# src/data.py

from torchvision.datasets import CIFAR10
from torchvision.datasets import CIFAR100
from torchvision import transforms

from torch.utils.data import (
    Dataset,
    DataLoader
)

from PIL import Image

import numpy as np

from src.logger import get_logger

logger = get_logger(__name__)


class OpenSetDataset(Dataset):
    """
    Custom dataset for CIFAR10 + Unknown class.
    Applies transforms dynamically.
    """

    def __init__(
        self,
        images,
        labels,
        transform=None
    ):
        self.images = images
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.images)

    def __getitem__(
        self,
        idx
    ):

        image = self.images[idx]

        label = int(
            self.labels[idx]
        )

        image = Image.fromarray(
            image.astype(np.uint8)
        )

        if self.transform:
            image = self.transform(
                image
            )

        return image, label


def load_dataset(
    batch_size=64
):
    """
    Returns
    -------
    train_loader
    test_loader

    Classes:
        0-9  = CIFAR10 classes
        10   = Unknown class
    """

    logger.info(
        "Loading CIFAR10 and CIFAR100"
    )


    # TRAIN TRANSFORMS


    train_transform = transforms.Compose([

        transforms.RandomCrop(
            32,
            padding=4
        ),

        transforms.RandomHorizontalFlip(),

        transforms.ColorJitter(
            brightness=0.2,
            contrast=0.2
        ),

        transforms.ToTensor(),

        transforms.Normalize(
            mean=[0.4914, 0.4822, 0.4465],
            std=[0.2470, 0.2435, 0.2616]
        )
    ])


    # TEST TRANSFORMS


    test_transform = transforms.Compose([

        transforms.ToTensor(),

        transforms.Normalize(
            mean=[0.4914, 0.4822, 0.4465],
            std=[0.2470, 0.2435, 0.2616]
        )
    ])


    # CIFAR10


    cifar10_train = CIFAR10(
        root="data",
        train=True,
        download=True
    )

    cifar10_test = CIFAR10(
        root="data",
        train=False,
        download=True
    )

    x10_train = cifar10_train.data

    y10_train = np.array(
        cifar10_train.targets
    )

    x10_test = cifar10_test.data

    y10_test = np.array(
        cifar10_test.targets
    )


    # CIFAR100


    cifar100_train = CIFAR100(
        root="data",
        train=True,
        download=True
    )

    cifar100_test = CIFAR100(
        root="data",
        train=False,
        download=True
    )

    x100_train = cifar100_train.data

    x100_test = cifar100_test.data


    # UNKNOWN SAMPLING


    np.random.seed(42)

    unknown_train_idx = np.random.choice(
        len(x100_train),
        5000,
        replace=False
    )

    unknown_test_idx = np.random.choice(
        len(x100_test),
        1000,
        replace=False
    )

    x_unknown_train = x100_train[
        unknown_train_idx
    ]

    x_unknown_test = x100_test[
        unknown_test_idx
    ]

    UNKNOWN_LABEL = 10

    y_unknown_train = np.full(
        len(x_unknown_train),
        UNKNOWN_LABEL
    )

    y_unknown_test = np.full(
        len(x_unknown_test),
        UNKNOWN_LABEL
    )


    # MERGE DATASETS


    X_train = np.concatenate(
        [
            x10_train,
            x_unknown_train
        ],
        axis=0
    )

    y_train = np.concatenate(
        [
            y10_train,
            y_unknown_train
        ],
        axis=0
    )

    X_test = np.concatenate(
        [
            x10_test,
            x_unknown_test
        ],
        axis=0
    )

    y_test = np.concatenate(
        [
            y10_test,
            y_unknown_test
        ],
        axis=0
    )


    # SHUFFLE


    train_perm = np.random.permutation(
        len(X_train)
    )

    test_perm = np.random.permutation(
        len(X_test)
    )

    X_train = X_train[
        train_perm
    ]

    y_train = y_train[
        train_perm
    ]

    X_test = X_test[
        test_perm
    ]

    y_test = y_test[
        test_perm
    ]


    # DATASETS


    train_dataset = OpenSetDataset(
        images=X_train,
        labels=y_train,
        transform=train_transform
    )

    test_dataset = OpenSetDataset(
        images=X_test,
        labels=y_test,
        transform=test_transform
    )


    # DATALOADERS


    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=0,
        pin_memory=False
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
        pin_memory=False
    )

    logger.info(
        f"Train Samples: {len(train_dataset)}"
    )

    logger.info(
        f"Test Samples: {len(test_dataset)}"
    )

    logger.info(
        "Dataset ready for PyTorch"
    )

    return (
        train_loader,
        test_loader
    )


if __name__ == "__main__":

    train_loader, test_loader = (
        load_dataset()
    )

    images, labels = next(
        iter(train_loader)
    )

    print(
        f"Images Shape: {images.shape}"
    )

    print(
        f"Labels Shape: {labels.shape}"
    )

    print(
        f"Classes: {labels[:10]}"
    )