"""
Module for loading and preprocessing image classification datasets.
Supports MNIST, CIFAR10, CIFAR100, Food101, and Flowers102 with configurable
data augmentation and normalization pipelines based on input resolution.
Provides functions to fetch train/test datasets and wrap them in DataLoaders.
"""

import torch
import matplotlib.pyplot as plt
from torchvision import datasets, transforms
from torch.utils.data import DataLoader


def get_train_test_data(dataset_class, root, resolution):
    train_transform_list, test_transform_list = [], []

    if dataset_class == datasets.MNIST:
        mean, std = (0.1307,), (0.3081,)
        train_transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)]
        )
        test_transform = train_transform  # same for MNIST

    elif dataset_class == datasets.CIFAR10:
        mean, std = (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
        if resolution != "default":
            train_transform_list.append(transforms.Resize(resolution))
            test_transform_list.append(transforms.Resize(resolution))
            crop_size = resolution
        else:
            crop_size = 32

        train_transform_list.extend(
            [
                transforms.RandomCrop(crop_size, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(mean=mean, std=std),
            ]
        )

        test_transform_list.extend(
            [transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)]
        )

        train_transform = transforms.Compose(train_transform_list)
        test_transform = transforms.Compose(test_transform_list)

    elif dataset_class == datasets.CIFAR100:
        mean, std = (0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)
        if resolution != "default":
            train_transform_list.append(transforms.Resize(resolution))
            test_transform_list.append(transforms.Resize(resolution))
            crop_size = resolution
        else:
            crop_size = 32

        train_transform_list.extend(
            [
                transforms.RandomCrop(crop_size, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(mean=mean, std=std),
            ]
        )

        test_transform_list.extend(
            [transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)]
        )

        train_transform = transforms.Compose(train_transform_list)
        test_transform = transforms.Compose(test_transform_list)

    elif dataset_class == datasets.Food101:
        # Use ImageNet-style normalization for Food101
        mean, std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
        if resolution != "default":
            train_transform_list.append(transforms.Resize((resolution, resolution)))
            test_transform_list.append(transforms.Resize((resolution, resolution)))
            crop_size = resolution
        else:
            crop_size = 224  # default crop size for Food101
        train_transform_list.extend(
            [
                transforms.RandomCrop(crop_size, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(mean=mean, std=std),
            ]
        )
        test_transform_list.extend(
            [transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)]
        )
        train_transform = transforms.Compose(train_transform_list)
        test_transform = transforms.Compose(test_transform_list)

    elif dataset_class == datasets.Flowers102:
        # Use ImageNet-style normalization for Oxford Flowers 102
        mean, std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
        if resolution != "default":
            train_transform_list.append(transforms.Resize((resolution, resolution)))
            test_transform_list.append(transforms.Resize((resolution, resolution)))
            crop_size = resolution
        else:
            crop_size = 224  # default crop size for Flowers102
        train_transform_list.extend(
            [
                transforms.RandomCrop(crop_size, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(mean=mean, std=std),
            ]
        )
        test_transform_list.extend(
            [transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)]
        )
        train_transform = transforms.Compose(train_transform_list)
        test_transform = transforms.Compose(test_transform_list)

    else:
        raise ValueError(f"Unsupported dataset: {dataset_class}")

    # Instantiate the dataset with the appropriate arguments
    if dataset_class in [datasets.Food101, datasets.Flowers102]:
        # These datasets use 'split' instead of 'train' boolean
        train_data = dataset_class(
            root=root, split="train", download=True, transform=train_transform
        )
        test_data = dataset_class(
            root=root, split="test", download=True, transform=test_transform
        )
    else:
        train_data = dataset_class(
            root=root, train=True, download=True, transform=train_transform
        )
        test_data = dataset_class(
            root=root, train=False, download=True, transform=test_transform
        )

    return train_data, test_data


def plot_some_images(train_data, class_names, colored=False):
    # Plot some images
    torch.manual_seed(42)
    fig = plt.figure(figsize=(9, 9))
    rows, cols = 4, 4
    for i in range(1, rows * cols + 1):
        random_idx = torch.randint(0, len(train_data), size=[1]).item()
        img, label = train_data[random_idx]
        fig.add_subplot(rows, cols, i)
        if colored:
            # Check if img is a tensor and convert it to a numpy array if necessary
            if isinstance(img, torch.Tensor):
                plt_img = img.permute(
                    1, 2, 0
                ).numpy()  # Rearrange dimensions for plotting
        plt.imshow(plt_img.squeeze(), cmap="gray")
        plt.title(class_names[label])
        plt.axis(False)


def img_class_loader(data_name, batch_size, to_print, colored, visually, resolution):
    name_to_dataset_class = {
        "mnist": datasets.MNIST,
        "cifar10": datasets.CIFAR10,
        "cifar100": datasets.CIFAR100,
        "food101": datasets.Food101,
        "flowers102": datasets.Flowers102,
    }
    dataset_class = name_to_dataset_class[data_name.lower()]
    root = f"data/raw_data/{data_name.lower()}"
    train_data, test_data = get_train_test_data(
        dataset_class=dataset_class, root=root, resolution=resolution
    )

    # Depending on dataset, classes attribute may differ
    if hasattr(train_data, "classes"):
        class_names = train_data.classes
    else:
        # For Flowers102, classes aren't provided, create a dummy placeholder
        class_names = [str(i) for i in range(len(set(train_data.targets)))]
    class_to_idx = (
        train_data.class_to_idx if hasattr(train_data, "class_to_idx") else None
    )

    if to_print:
        # Some datasets (like Flowers102) use .images and .targets instead of .data/.targets
        try:
            n_train = len(train_data.data)
            n_train_labels = len(train_data.targets)
        except AttributeError:
            n_train = len(train_data)
            n_train_labels = len(train_data.targets)
        try:
            n_test = len(test_data.data)
            n_test_labels = len(test_data.targets)
        except AttributeError:
            n_test = len(test_data)
            n_test_labels = len(test_data.targets)
        print(
            f"train X: {n_train}, train y: {n_train_labels}, test X: {n_test}, test y: {n_test_labels}"
        )

    if visually:
        plot_some_images(
            train_data=train_data, class_names=class_names, colored=colored
        )

    # Prepare the dataloader
    BATCH_SIZE = batch_size
    num_workers = 4
    train_dataloader = DataLoader(
        dataset=train_data,
        batch_size=BATCH_SIZE,
        shuffle=True,
        drop_last=True,
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=True,
        prefetch_factor=2,
    )

    test_dataloader = DataLoader(
        dataset=test_data,
        batch_size=BATCH_SIZE,
        shuffle=False,
        drop_last=True,
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=True,
        prefetch_factor=2,
    )

    train_features_batch, train_labels_batch = next(iter(train_dataloader))
    if to_print:
        print(
            f"Length of train_dataloader: {len(train_dataloader)} batches of {BATCH_SIZE}..."
        )
        print(
            f"Length of test_dataloader: {len(test_dataloader)} batches of {BATCH_SIZE}..."
        )
        print(
            f"Shape of the training batch: {train_features_batch.size()}, Shape of the training label batch: {train_labels_batch.size()}"
        )

    HW = (
        train_features_batch.size(2),
        train_features_batch.size(3),
    )
    num_classes = len(class_names)
    return (
        train_dataloader,
        test_dataloader,
        HW,
        num_classes,
    )
