from abc import ABC
from typing import Generator, Iterator, List

import numpy as np
import torch
from minio import datatypes
from PIL import Image
from torch.utils.data import DataLoader, Dataset, IterableDataset
from torch.utils.data.dataset import T_co
from torchvision.transforms import CenterCrop, Compose, ToTensor

from station.common.clients.minio import MinioClient


class BaseDataSet(Dataset, ABC):
    id: str


class MinioFolderDataSet(IterableDataset):
    def __init__(
        self, minio_client: MinioClient, id: str = None, transform: Compose = None
    ):
        super(MinioFolderDataSet, self).__init__()
        self.id = id
        self.client = minio_client
        self.transform = transform

    def __iter__(self) -> Iterator[T_co]:
        worker_info = torch.utils.data.get_worker_info()

        if worker_info is None:
            return self._make_generator()

        else:
            print(worker_info.num_workers)
            # TODO
            raise NotImplementedError(
                "Implement sliced generators for multiworker loading"
            )

    def _make_generator(self, id: str = None):
        if id:
            self.id = id
        minio_items = self._get_minio_items()

        tensor_generator = self._make_image_generator(minio_items)
        return tensor_generator

    def _get_minio_items(self) -> Generator[datatypes.Object, None, None]:
        items = self.client.get_minio_dir_items(self.id)
        return items

    def _make_image_generator(
        self, minio_items: Generator[datatypes.Object, None, None]
    ) -> Generator[torch.Tensor, None, None]:
        for i, object in enumerate(minio_items):
            minio_object = self.client.client.get_object(
                "data-sets", object.object_name
            )
            np_img = np.array(Image.open(minio_object))

            print(object.object_name)

            if self.transform:
                yield self.transform(np_img)
            else:
                yield np_img


class MinioFolderDS(Dataset):
    # TODO generalized target creation

    def __init__(
        self,
        client: MinioClient,
        data_set_id: str = None,
        transform: Compose = None,
        target_classes: List[str] = None,
    ):
        super().__init__()
        self.data_set_id = data_set_id
        self.minio_client = client
        self.transform = transform
        self.items = None

        self.get_minio_items()

        if target_classes:
            self.classes = np.asarray(target_classes)
        else:
            self.classes = self.get_classes_from_folders()

    def __getitem__(self, index) -> T_co:
        minio_object = self.minio_client.client.get_object(
            "datasets", self.items[index]
        )
        np_img = np.array(Image.open(minio_object))
        np_img = np.transpose(np_img.astype(float), (2, 1, 0))
        label = self.get_label(self.items[index])
        if self.transform:
            return self.transform(np_img), label
        return torch.Tensor(np_img), label

    def __len__(self):
        return len(list(self.items))

    def get_label(self, item_name: str):
        folder = item_name[len(self.data_set_id) :].split("/")[0]
        label = self.classes == folder
        return torch.tensor(label, dtype=torch.long)

    def get_minio_items(self):
        minio_items = self.minio_client.get_minio_dir_items(self.data_set_id)
        item_list = []
        for item in minio_items:
            item_list.append(item.object_name)

        self.items = item_list

    def save(self):
        # TODO Save the full data set to disk somewhere to enable fast loading
        pass

    def get_classes_from_folders(self):
        classes = self.minio_client.get_classes_by_folders(self.data_set_id)
        print(classes)
        return np.asarray(classes)


if __name__ == "__main__":
    minio_client = MinioClient(
        minio_server="localhost:9000",
        secret_key="minio_admin",
        access_key="minio_admin",
    )

    transform = Compose([ToTensor(), CenterCrop(size=24)])
    minio_ds = MinioFolderDS(
        client=minio_client, data_set_id="cifar/batch_2", transform=transform
    )
    minio_ds.get_minio_items()
    print(len(minio_ds))

    dl = DataLoader(minio_ds, batch_size=6, num_workers=0, shuffle=True)

    for i, batch in enumerate(dl):
        x, y = batch
        print(y)
        if i > 3:
            break
