from __future__ import annotations

import hashlib
import random
import re
import shutil
from abc import ABC, abstractmethod
from collections import Counter, OrderedDict
from copy import deepcopy
from pathlib import Path
from typing import (
    Any,
    Callable,
    Generic,
    Iterator,
    List,
    Mapping,
    Optional,
    Protocol,
    Sequence,
    Tuple,
    Type,
    TypeVar,
)

import numpy as np
import torch
from loguru import logger
from torch.nn.utils.rnn import pad_sequence
from torchtext.vocab import Vocab, vocab

from mads_datasets.datasets import (
    DatasetProtocol,
    ImgDataset,
    MNISTDataset,
    PdDataset,
    SunspotDataset,
    TextDataset,
    TSDataset,
)
from mads_datasets.datatools import (
    get_file,
    iter_valid_paths,
    keep_subdirs_only,
    walk_dir,
    window,
)
from mads_datasets.settings import (
    DatasetSettings,
    DatasetType,
    ImgDatasetSettings,
    PdDatasetSettings,
    TextDatasetSettings,
    WindowedDatasetSettings,
    fashionmnistsettings,
    flowersdatasetsettings,
    gesturesdatasetsettings,
    imdbdatasetsettings,
    irissettings,
    sunspotsettings,
    penguinssettings,
)

Tensor = torch.Tensor


class DatastreamerProtocol(Protocol):
    def stream(self) -> Iterator:
        ...


T = TypeVar("T", bound=DatasetSettings)


class AbstractDatasetFactory(ABC, Generic[T]):
    def __init__(
        self, settings: T, preprocessor: Type[PreprocessorProtocol], datadir: Path
    ) -> None:
        self._settings = settings
        self.preprocessor = preprocessor
        self.datadir = datadir

    @property
    def settings(self) -> T:
        return self._settings

    def download_data(self) -> None:
        url = self._settings.dataset_url
        filename = self._settings.filename
        self.subfolder = Path(self.datadir) / self.settings.name
        if not self.subfolder.exists():
            logger.info("Start download...")
            self.subfolder.mkdir(parents=True)
            self.filepath = get_file(self.subfolder, filename, url=url, overwrite=False)
        else:
            logger.info(f"Dataset already exists at {self.subfolder}")
            self.filepath = self.subfolder / filename

        digest = self.calculate_md5(self.filepath)
        if self.settings.digest is not None:
            if digest != self.settings.digest:
                raise ValueError(
                    f"Digest of downloaded file {self.filepath} does not match expected digest"
                    f"\nExpected: {self.settings.digest}\nGot: {digest}"
                )
            else:
                logger.info(
                    f"Digest of downloaded {self.filepath} matches expected digest"
                )
        else:
            logger.info(f"Digest of downloaded {self.filepath} is {digest}")

    def get_preprocessor(self, **kwargs: Any) -> PreprocessorProtocol:
        return self.preprocessor(**kwargs)

    @staticmethod
    def calculate_md5(file_path: Path, block_size: int = 2**16) -> str:
        md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            for block in iter(lambda: f.read(block_size), b""):
                md5.update(block)
        return md5.hexdigest()

    @abstractmethod
    def create_dataset(self, *args, **kwargs) -> Mapping[str, DatasetProtocol]:
        raise NotImplementedError

    def create_datastreamer(
        self, batchsize: int, **kwargs
    ) -> Mapping[str, DatastreamerProtocol]:
        datasets = self.create_dataset()
        traindataset = datasets["train"]
        validdataset = datasets["valid"]
        preprocessor = self.get_preprocessor(**kwargs)

        trainstreamer = BaseDatastreamer(
            traindataset, batchsize=batchsize, preprocessor=preprocessor
        )
        validstreamer = BaseDatastreamer(
            validdataset, batchsize=batchsize, preprocessor=preprocessor
        )
        return {"train": trainstreamer, "valid": validstreamer}

class PenguinsDatasetFactory(AbstractDatasetFactory[DatasetSettings]):
    def create_dataset(self, *args, **kwargs) -> Mapping[str, DatasetProtocol]:
        pass

class IrisDatasetFactory(AbstractDatasetFactory[PdDatasetSettings]):
    def create_dataset(
        self, *args: Any, **kwargs: Any
    ) -> Mapping[str, DatasetProtocol]:
        try:
            import pandas as pd
        except ImportError:
            raise ImportError(
                "The 'pandas' package is required for the iris dataset. "
                'You can update existing environments it by running poetry install --extras "pandas"'
                "or pip install mads-datasets[pandas] --upgrade"
            )
        self.download_data()
        df = pd.read_csv(
            self.filepath,
            header=None,
            names=self._settings.features + [self._settings.target],
        )
        df["sepal_width"].fillna(value=df["sepal_width"].mean(), inplace=True)
        df[self._settings.target] = df[self._settings.target].apply(
            lambda x: re.sub(r"setsoa", "setosa", x)
        )
        df.drop_duplicates(inplace=True)
        df["petal_width"] = df["petal_width"].apply(
            lambda x: float(re.findall(r"(\d+\.?\d*) mm", x)[0]) / 10
        )
        idx = df[df["sepal_length"] == 58].index
        df.loc[idx, "sepal_length"] = 5.8

        train = df.sample(frac=0.6)
        valid = df.drop(train.index)

        trainset = PdDataset(
            train, features=self._settings.features, target=self._settings.target
        )
        validset = PdDataset(
            valid, features=self._settings.features, target=self._settings.target
        )
        return {"train": trainset, "valid": validset}

    def create_datastreamer(
        self, batchsize: int, **kwargs
    ) -> Mapping[str, DatastreamerProtocol]:
        raise NotImplementedError("Not implemented yet...")


class SunspotsDatasetFactory(AbstractDatasetFactory[WindowedDatasetSettings]):
    """
    Data from https://www.sidc.be/SILSO/datafiles
    """

    def create_dataset(
        self, *args: Any, **kwargs: Any
    ) -> Mapping[str, DatasetProtocol]:
        self.download_data()
        spots = np.genfromtxt(str(self.filepath), usecols=(3))  # type: ignore
        tensors = torch.from_numpy(spots).type(torch.float32)

        split = kwargs.pop("split", 0.8)
        idx = int(len(tensors) * split)
        train = tensors[:idx]
        valid = tensors[idx:]

        norm = max(train)
        train = train / norm
        valid = valid / norm

        window_size = self.settings.window_size
        horizon = self.settings.horizon
        trainset = SunspotDataset(self._window(train, window_size), horizon)
        validset = SunspotDataset(self._window(valid, window_size), horizon)
        return {"train": trainset, "valid": validset}

    @staticmethod
    def _window(data: Tensor, window_size: int) -> Tensor:
        idx = window(data, window_size)
        dataset = data[idx]
        dataset = dataset[..., None]
        return dataset


class FashionDatasetFactory(AbstractDatasetFactory[DatasetSettings]):
    def create_dataset(
        self, *args: Any, **kwargs: Any
    ) -> Mapping[str, DatasetProtocol]:
        self.download_data()
        data = torch.load(self.filepath)
        training_data = data["traindata"]
        training_labels = data["trainlabels"]
        test_data = data["testdata"]
        test_labels = data["testlabels"]

        train = MNISTDataset(training_data, training_labels)
        test = MNISTDataset(test_data, test_labels)

        return {"train": train, "valid": test}


class GesturesDatasetFactory(AbstractDatasetFactory[DatasetSettings]):
    def __init__(
        self,
        settings: DatasetSettings,
        preprocessor: Type[PreprocessorProtocol],
        **kwargs: Any,
    ) -> None:
        super().__init__(settings, preprocessor, **kwargs)
        self._created = False
        self.datasets: Mapping[str, DatasetProtocol]

    def create_dataset(
        self, *args: Any, **kwargs: Any
    ) -> Mapping[str, DatasetProtocol]:
        self.download_data()

        if self._created:
            return deepcopy(self.datasets)
        formats = [f.value for f in self._settings.formats]
        datadir = self.subfolder / "gestures-dataset"
        img = datadir / "gestures.png"
        if img.exists():
            shutil.move(img, datadir.parent / "gestures.png")
        keep_subdirs_only(datadir)
        paths = [path for path in walk_dir(datadir) if path.suffix in formats]
        random.shuffle(paths)

        split = kwargs.pop("split", 0.8)
        idx = int(len(paths) * split)
        trainpaths = paths[:idx]
        validpaths = paths[idx:]

        traindataset = TSDataset(trainpaths)
        validdataset = TSDataset(validpaths)
        datasets = {
            "train": traindataset,
            "valid": validdataset,
        }
        self.datasets = datasets  # type: ignore
        self._created = True
        return datasets


class IMDBDatasetFactory(AbstractDatasetFactory[TextDatasetSettings]):
    def create_dataset(
        self, *args: Any, **kwargs: Any
    ) -> Mapping[str, DatasetProtocol]:
        self.download_data()
        testdir = self.subfolder / "aclImdb/test"
        traindir = self.subfolder / "aclImdb/train"
        keep_subdirs_only(testdir)
        keep_subdirs_only(traindir)

        # remove dir with unlabeled reviews
        unsup = traindir / "unsup"
        if unsup.exists():
            shutil.rmtree(traindir / "unsup")

        formats = [f.value for f in self._settings.formats]
        trainpaths = [path for path in walk_dir(traindir) if path.suffix in formats]
        testpaths = [path for path in walk_dir(testdir) if path.suffix in formats]
        logger.info(
            f"Creating TextDatasets from {len(trainpaths)} trainfiles and {len(testpaths)} testfiles."
        )

        traindataset = TextDataset(paths=trainpaths)
        testdataset = TextDataset(paths=testpaths)
        return {"train": traindataset, "valid": testdataset}

    def create_datastreamer(
        self, batchsize: int, **kwargs
    ) -> Mapping[str, DatastreamerProtocol]:
        datasets = self.create_dataset()
        traindataset = datasets["train"]
        validdataset = datasets["valid"]
        preprocessor = self.get_preprocessor(
            traindataset=traindataset, settings=self.settings
        )

        trainstreamer = BaseDatastreamer(
            traindataset, batchsize=batchsize, preprocessor=preprocessor
        )
        validstreamer = BaseDatastreamer(
            validdataset, batchsize=batchsize, preprocessor=preprocessor
        )
        return {"train": trainstreamer, "valid": validstreamer}


class FlowersDatasetFactory(AbstractDatasetFactory[ImgDatasetSettings]):
    def create_dataset(
        self, *args: Any, **kwargs: Any
    ) -> Mapping[str, DatasetProtocol]:
        self.download_data()
        formats = self._settings.formats
        paths_, class_names = iter_valid_paths(
            self.subfolder / "flower_photos", formats=formats
        )
        paths = [*paths_]
        random.shuffle([paths])
        trainidx = int(len(paths) * self._settings.trainfrac)
        train = paths[:trainidx]
        valid = paths[trainidx:]
        traindataset = ImgDataset(train, class_names, img_size=self._settings.img_size)
        validdataset = ImgDataset(valid, class_names, img_size=self._settings.img_size)
        return {"train": traindataset, "valid": validdataset}


class PreprocessorProtocol(Protocol):
    def __call__(self, batch: List[Tuple]) -> Tuple[torch.Tensor, torch.Tensor]:
        ...


class BasePreprocessor(PreprocessorProtocol):
    def __call__(self, batch: List[Tuple]) -> Tuple[torch.Tensor, torch.Tensor]:
        X, y = zip(*batch)
        return torch.stack(X), torch.tensor(y)


class PaddedPreprocessor(PreprocessorProtocol):
    def __call__(self, batch: List[Tuple]) -> Tuple[torch.Tensor, torch.Tensor]:
        X, y = zip(*batch)
        X_ = pad_sequence(X, batch_first=True, padding_value=0)  # noqa N806
        return X_, torch.tensor(y)


class BaseTokenizer(PreprocessorProtocol):
    def __init__(
        self, traindataset: TextDataset, settings: TextDatasetSettings
    ) -> None:
        self.maxvocab = settings.maxvocab
        self.maxtokens = settings.maxtokens
        self.clean = settings.clean_fn
        self.vocab = self.build_vocab(self.build_corpus(traindataset))

    @staticmethod
    def split_and_flat(corpus: List[str]) -> List[str]:
        """
        Split a list of strings on spaces into a list of lists of strings
        and then flatten the list of lists into a single list of strings.
        eg ["This is a sentence"] -> ["This", "is", "a", "sentence"]
        """
        corpus_ = [x.split() for x in corpus]
        corpus = [x for y in corpus_ for x in y]
        return corpus

    def build_corpus(self, dataset) -> List[str]:
        corpus = []
        for i in range(len(dataset)):
            x = self.clean(dataset[i][0])
            corpus.append(x)
        return corpus

    def build_vocab(
        self, corpus: List[str], oov: str = "<OOV>", pad: str = "<PAD>"
    ) -> Vocab:
        data = self.split_and_flat(corpus)
        counter = Counter(data).most_common()
        logger.info(f"Found {len(counter)} tokens")
        counter = counter[: self.maxvocab - 2]
        ordered_dict = OrderedDict(counter)
        v1 = vocab(ordered_dict, specials=[pad, oov])
        v1.set_default_index(v1[oov])
        return v1

    def cast_label(self, label: str) -> int:
        raise NotImplementedError

    def __call__(self, batch: List) -> Tuple[Tensor, Tensor]:
        labels, text = [], []
        for x, y in batch:
            if self.clean is not None:
                x = self.clean(x)  # type: ignore
            x = x.split()[: self.maxtokens]
            tokens = torch.tensor([self.vocab[word] for word in x], dtype=torch.int32)
            text.append(tokens)
            labels.append(self.cast_label(y))

        text_ = pad_sequence(text, batch_first=True, padding_value=0)
        return text_, torch.tensor(labels)


class IMDBTokenizer(BaseTokenizer):
    def __init__(self, traindataset, settings):
        super().__init__(traindataset, settings)

    def cast_label(self, label: str) -> int:
        if label == "neg":
            return 0
        else:
            return 1


class DatasetFactoryProvider:
    @staticmethod
    def create_factory(dataset_type: DatasetType, **kwargs) -> AbstractDatasetFactory:
        datadir = Path(kwargs.get("datadir", Path.home() / ".cache/mads_datasets"))
        if dataset_type == DatasetType.FLOWERS:
            preprocessor = kwargs.get("preprocessor", BasePreprocessor)
            return FlowersDatasetFactory(
                flowersdatasetsettings, preprocessor=preprocessor, datadir=datadir
            )
        if dataset_type == DatasetType.IMDB:
            preprocessor = kwargs.get("preprocessor", IMDBTokenizer)
            return IMDBDatasetFactory(
                imdbdatasetsettings, preprocessor=preprocessor, datadir=datadir
            )
        if dataset_type == DatasetType.GESTURES:
            preprocessor = kwargs.get("preprocessor", PaddedPreprocessor)
            return GesturesDatasetFactory(
                gesturesdatasetsettings, preprocessor=preprocessor, datadir=datadir
            )
        if dataset_type == DatasetType.FASHION:
            preprocessor = kwargs.get("preprocessor", BasePreprocessor)
            return FashionDatasetFactory(
                fashionmnistsettings, preprocessor=preprocessor, datadir=datadir
            )
        if dataset_type == DatasetType.SUNSPOTS:
            preprocessor = kwargs.get("preprocessor", BasePreprocessor)
            return SunspotsDatasetFactory(
                sunspotsettings, preprocessor=preprocessor, datadir=datadir
            )
        if dataset_type == DatasetType.IRIS:
            preprocessor = kwargs.get("preprocessor", BasePreprocessor)
            return IrisDatasetFactory(irissettings, preprocessor=preprocessor, datadir=datadir)

        if dataset_type == DatasetType.PENGUINS:
            preprocessor = kwargs.get("preprocessor", BasePreprocessor)
            return PenguinsDatasetFactory(penguinssettings, preprocessor=preprocessor, datadir=datadir)

        raise ValueError(f"Invalid dataset type: {dataset_type}")


class BaseDatastreamer:
    """This datastreamer wil never stop
    The dataset should have a:
        __len__ method
        __getitem__ method

    """

    def __init__(
        self,
        dataset: DatasetProtocol,
        batchsize: int,
        preprocessor: Optional[Callable] = None,
    ) -> None:
        self.dataset = dataset
        self.batchsize = batchsize
        self.preprocessor = preprocessor
        self.size = len(self.dataset)
        self.reset_index()

    def __len__(self) -> int:
        return int(len(self.dataset) / self.batchsize)

    def __repr__(self) -> str:
        return f"BasetDatastreamer: {self.dataset} (streamerlen {len(self)})"

    def reset_index(self) -> None:
        self.index_list = np.random.permutation(self.size)
        self.index = 0

    def batchloop(self) -> Sequence[Tuple]:
        batch = []
        for _ in range(self.batchsize):
            x, y = self.dataset[int(self.index_list[self.index])]
            batch.append((x, y))
            self.index += 1
        return batch

    def stream(self) -> Iterator:
        while True:
            if self.index > (self.size - self.batchsize):
                self.reset_index()
            batch = self.batchloop()
            if self.preprocessor is not None:
                X, Y = self.preprocessor(batch)  # noqa N806
            else:
                X, Y = zip(*batch)  # noqa N806
            yield X, Y
