__all__ = ["DatasetLoader", "NotConfiguredError", "UNSET"]
import shutil
from enum import Enum
from functools import partial
from itertools import chain
from pathlib import Path
from shutil import rmtree
from tempfile import mkdtemp
from typing import Dict, Optional, Tuple
from urllib.request import urlopen
from zipfile import ZipFile

import pandas as pd
import yaml

from .config import (
    DOWNLOAD,
    DatasetConfig,
    ImageFormat,
    InvalidConfigError,
    Preprocessing,
)
from .dataset import (
    BiomarkerField,
    BiomarkersAnnotationInfos,
    BiomarkersAnnotationTasks,
    Dataset,
    FundusField,
)
from .utilities import RichProgress, xdg_data_home

#   === CONSTANTS ===
# Figshare public urls of the MAPLES DR dataset
MAPLES_DR_ADDITIONAL_URL = "https://figshare.com/ndownloader/files/45795384"

#: Unset constant
UNSET = "UNSET"


class DatasetSubset(str, Enum):
    """Enumeration of the possible subsets of the MAPLES-DR dataset."""

    #: The training set (138 images)
    TRAIN = "train"

    #: The testing set (60 images)
    TEST = "test"

    #: All the images of the training and test set (198 images)
    ALL = "all"

    #: The two duplicated images of the training set
    DUPLICATES = "duplicates"

    #: All the images of the training and test set, including duplicates (200 images)
    ALL_WITH_DUPLICATES = "all_with_duplicates"


class NotConfiguredError(Exception):
    """
    Exception raised when the dataset loader is not configured.
    """

    def __init__(
        self, message: str = "MAPLES-DR dataset is not configured yet.", *args
    ):
        super().__init__(message, *args)


class DatasetLoader:
    """
    Loader for MAPLES-DR dataset.
    """

    def __init__(self):
        self._datasets_config = DatasetConfig(
            resize=1500,
            image_format="PIL",
            preprocessing="none",
        )
        self.dataset_record: Optional[dict] = None
        self._maples_dr_path: Optional[Path] = None
        self._diagnosis: Optional[pd.DataFrame] = None
        self._annotations_infos: Optional[Dict[str, pd.DataFrame]] = None
        self._is_maples_dr_folder_temporary: bool = False

        self._messidor_paths: Optional[dict[str, str]] = None
        self._messidor_ROIs: Optional[pd.DataFrame] = None

    def __del__(self):
        if self._is_maples_dr_folder_temporary:
            rmtree(self._maples_dr_path, ignore_errors=True)

    # === CONFIGURATION ===
    def configure(
        self,
        maples_dr_path: Optional[str | Path] = UNSET,
        messidor_path: Optional[str | Path] = UNSET,
        *,
        cache: Optional[str | Path] = UNSET,
        resize: Optional[int] = None,
        image_format: Optional[ImageFormat] = None,
        preprocessing: Optional[Preprocessing] = None,
        disable_check: bool = False,
    ):
        """Configure the default behavior of the MAPLES-DR dataset.

        Any parameters left to None (or 'UNSET' for the first two paths) will leave the current configuration unaffected.

        Parameters
        ----------
        maples_dr_path : Optional[str], optional
            Path to the MAPLES-DR additional data. Must point to the directory or to the zip file.

            If None (by default), then the dataset is downloaded from figshare.

        messidor_path : Optional[str], optional
            Path to the MESSIDOR dataset.

            Must point to a directory containing the "Base11", "Base12", ... subdirectories or zip files.

        cache : Optional[str], optional
            Path to the cache directory. The cache is used to store the downloaded dataset and the generated images.

            - If ``cache`` is a ``str`` or a ``Path``, then the cache is stored in the given directory.
            - If ``False`` (by default), then the cache is disabled.
            - If ``True``, then the cache is stored in the default cache directory.

        resize : Optional[int], optional
            Set the size of the images (fundus and biomarkers) generated by `maples_dr`.

            - If ``resize`` is an int, crop the image to a square ROI and resize it to the shape ``(resize, resize)``;
            - If ``True``, keep the original MAPLES-DR resolution of 1500x1500 px;
            - If ``False``, use the original MESSIDOR resolution if MESSIDOR path is configured, otherwise fallback to MAPLES-DR original resolution.

        image_format : Optional[ImageFormat], optional
            Python format of the generated images. Must be either "PIL", "rgb" or "bgr".

            If "rgb" or "bgr" is selected, images will be formatted as numpy array of shape: (height, width, channel).

            By default, "PIL" is used.

        preprocessing : Optional[Preprocessing], optional
            Preprocessing algorithm applied on the fundus images.

            By default, no preprocessing is applied.

        disable_check : bool, optional
            If True, disable the integrity check of the dataset.

        """
        # === Update the dataset configuration ===
        self._datasets_config.update(
            dict(
                resize=resize,
                image_format=image_format,
                preprocessing=preprocessing,
            )
        )

        # === Set Maples-DR cache ===
        if cache is not UNSET and self.cfg._cache != cache:
            cache = self._change_cache_path(cache)
        else:
            cache = self.cfg.cache_path

        # === Prepare Maples-DR ===
        # If configure is called the first time with no path, download the dataset or use the cache.
        if maples_dr_path is UNSET and self._maples_dr_path is None:
            if cache is None:
                maples_dr_path = DOWNLOAD
            else:
                maples_dr_path = cache / "AdditionalData"

        if maples_dr_path is not UNSET and self._maples_dr_path != maples_dr_path:
            # Set the path and download the dataset if needed.
            maples_dr_path = self._change_maples_dr_path(maples_dr_path)

            # Load the dataset infos.
            (
                self.dataset_record,
                self._messidor_ROIs,
            ) = self.load_dataset_record_and_rois(maples_dr_path)
            self._annotations_infos = self.load_biomarkers_annotation_infos(
                maples_dr_path / "biomarkers_annotation_infos.xls"
            )

            # Check the integrity of the dataset.
            if not disable_check:
                self.check_maples_dr_integrity(
                    path=maples_dr_path,
                    biomarkers=list(self.dataset_record["biomarkers"]),
                    images_names=self.image_names(DatasetSubset.ALL_WITH_DUPLICATES),
                )

            self._diagnosis = self.load_maples_dr_diagnosis(
                maples_dr_path / "diagnosis_infos.xls"
            )

        # === Prepare MESSIDOR ===
        if messidor_path is UNSET:
            if (
                self._messidor_paths is None
                and cache is not None
                and (cache / "MESSIDOR").is_dir()
            ):
                messidor_path = cache / "MESSIDOR"

        if messidor_path is not UNSET:
            if messidor_path is None:
                self._messidor_paths = None
            else:
                self._messidor_paths = self.discover_messidor_images(
                    self.image_names(subset=DatasetSubset.ALL_WITH_DUPLICATES),
                    messidor_path,
                )

    def clear_cache(self):
        """
        Clear the cache.
        """
        if self.cfg.cache_path is not None:
            rmtree(self.cfg.cache_path, ignore_errors=True)

    @property
    def cfg(self) -> DatasetConfig:
        """
        Return the default configuration of the loaded dataset.
        """
        return self._datasets_config

    def is_configured(self) -> bool:
        """
        Check if the dataset is initialized.
        """
        return self.dataset_record is not None

    def ensure_configured(self):
        """
        Ensure the dataset is initialized.
        """
        if not self.is_initialized():
            self.configure()

    # --- Cache path configuration ---
    def _change_cache_path(self, path: Optional[str | Path | bool] = None) -> Path:
        if path is None or path is False:
            self.cfg._cache = False
            return None

        if path is True:
            path = xdg_data_home() / "maples_dr"

        # Ensure the path exists and is a directory.
        path = Path(path)
        if path.exists() and not path.is_dir():
            raise InvalidConfigError(f"Invalid cache path: {path} is not a directory.")
        path.mkdir(parents=True, exist_ok=True)

        previous_cache = self.cfg.cache_path
        if previous_cache is not None and path.absolute() == previous_cache.absolute():
            return path

        # If a cache already existed, copy the MAPLES-DR original dataset to the new cache path.
        if previous_cache is not None:
            if (previous_cache / "AdditionalData").is_dir():
                shutil.copytree(
                    previous_cache / "AdditionalData", path / "AdditionalData"
                )
            if (previous_cache / "MESSIDOR").is_dir():
                shutil.copytree(previous_cache / "MESSIDOR", path / "MESSIDOR")

        self.cfg._cache = path
        return path

    # --- Maples-DR path configuration ---
    @property
    def maples_dr_folder(self) -> Path:
        """
        Return the path to the MAPLES-DR dataset folder.
        """
        if self._maples_dr_path is None:
            raise NotConfiguredError()
        return self._maples_dr_path

    def _change_maples_dr_path(self, path: Optional[str | Path] = None) -> Path:
        """
        Return the path to the MAPLES-DR dataset folder.
        """
        # If Maples-DR was previously downloaded, delete the temporary folder.
        if self._is_maples_dr_folder_temporary and self._maples_dr_path is not None:
            rmtree(self._maples_dr_path, ignore_errors=True)
            self._is_maples_dr_folder_temporary = False

        self._maples_dr_path = None

        if path is None or path is DOWNLOAD:
            # If no path is given, prepare the dataset download to...
            path = mkdtemp()
            self._is_maples_dr_folder_temporary = True

        path = Path(path)

        if path.is_dir():
            # === If the path is a directory ===
            # Ensure it exists ...
            if not path.exists():
                path.mkdir(parents=True)
            # and check if the folder contains the dataset record.
            if not (path / "dataset_record.yaml").exists():
                # If not, download the dataset.
                zip_path = path / "maples_dr.zip"
                download(MAPLES_DR_ADDITIONAL_URL, zip_path, "MAPLES-DR labels")
                with ZipFile(path / "maples_dr.zip", "r") as zip_file:
                    zip_file.extractall(path)
                (path / "maples_dr.zip").unlink()

        elif path.name.endswith(".zip"):
            # === If the path is a zip file, unzip it to a temporary folder ===
            zip_path = path
            if not zip_path.exists():
                raise InvalidConfigError(
                    f"Invalid Maples DR archive: {zip_path} doesn't exist."
                )

            # Create a temporary folder.
            if self.cfg.cache_path is not None:
                path = self.cfg.cache_path / "AdditionalData"
            else:
                path = Path(mkdtemp())
            # Unzip the dataset to the temporary folder.
            try:
                with ZipFile(zip_path, "r") as zip_file:
                    zip_file.extractall(path)
            except Exception as e:
                raise InvalidConfigError(
                    f"Invalid Maples DR archive: {zip_path}:"
                    "\n\t the provided archive is impossible to unzip."
                ) from e
            else:
                # Test if the zip file contains maples_dr folder.
                if not (path / "dataset_record.yaml").exists():
                    raise InvalidConfigError(
                        f"Invalid Maples DR archive:  {path}:"
                        '\n\t the provided archive doesn\'t contains the file "dataset_record.yaml".'
                    )
                if not (path / "MESSIDOR-ROIs.csv").exists():
                    raise InvalidConfigError(
                        f"Invalid Maples DR archive:  {path}:"
                        '\n\t the provided archive doesn\'t contains the file "MESSIDOR-ROIs.csv".'
                    )
                if not (path / "diagnosis_infos.xls").exists():
                    raise InvalidConfigError(
                        f"Invalid Maples DR archive:  {path}:"
                        '\n\t the provided archive doesn\'t contains the file "diagnosis_infos.xls".'
                    )

                if not (path / "biomarkers_annotation_infos.xls").exists():
                    raise InvalidConfigError(
                        f"Invalid Maples DR archive:  {path}:"
                        '\n\t the provided archive doesn\'t contains the file "biomarkers_annotation_infos.xls".'
                    )
            if self.cfg.cache_path is not None:
                self._is_maples_dr_folder_temporary = True

        self._maples_dr_path = path
        return path

    @staticmethod
    def load_dataset_record_and_rois(path: str | Path) -> Tuple[Dict, Dict]:
        """
        Load the MAPLES-DR dataset record and the rois in MESSIDOR images.

        Parameters
        ----------
        path : str
            Path to the MAPLES-DR dataset folder.
        """

        with open(path / "dataset_record.yaml", "r") as infos_file:
            record = yaml.safe_load(infos_file)

        with open(path / "MESSIDOR-ROIs.csv", "r") as rois_file:
            rois = pd.read_csv(rois_file).set_index("name")

        return record, rois

    @staticmethod
    def load_biomarkers_annotation_infos(path: str | Path) -> pd.DataFrame:
        """
        Load the MAPLES-DR biomarkers annotation infos file.

        Parameters
        ----------
        path : str
            Path to the MAPLES-DR biomarkers annotation infos file.
        """

        annotation_infos_data = pd.read_excel(path, index_col=0, sheet_name=None)

        biomarker_infos = ["Retinologist", "Comment", "Time", "Annotation #"]
        biomarker_tasks = {
            "bright": BiomarkersAnnotationTasks.BRIGHT,
            "red": BiomarkersAnnotationTasks.RED,
            "disk-macula": BiomarkersAnnotationTasks.DISC_MACULA,
            "vessels": BiomarkersAnnotationTasks.VESSELS,
        }

        annotation_infos = None
        for sheet_name, task in biomarker_tasks.items():
            infos = annotation_infos_data[sheet_name][biomarker_infos].copy()
            infos["Time"] = pd.to_timedelta(infos["Time"])
            infos = infos.rename(
                columns={
                    "Retinologist": task + "_" + BiomarkersAnnotationInfos.RETINOLOGIST,
                    "Comment": task + "_" + BiomarkersAnnotationInfos.COMMENT,
                    "Time": task + "_" + BiomarkersAnnotationInfos.ANNOTATION_TIME,
                    "Annotation #": task
                    + "_"
                    + BiomarkersAnnotationInfos.ANNOTATION_ID,
                },
            )

            if annotation_infos is None:
                annotation_infos = infos
            else:
                annotation_infos = annotation_infos.join(infos)

        return annotation_infos

    def check_maples_dr_integrity(
        self, path: Path, biomarkers: list[str], images_names: list[str]
    ):
        """
        Check if the MAPLES-DR dataset contains all segmentation maps.
        """
        missing_images = 0
        for biomarker in biomarkers:
            for img in images_names:
                if (
                    self.is_biomarker_segmented(biomarker, img)
                    and not (path / "annotations" / biomarker / (img + ".png")).exists()
                ):
                    missing_images += 1
        if missing_images > 0:
            raise InvalidConfigError(
                f"The provided folder to the Maples-DR dataset is incomplete: "
                f"{missing_images} segmentation maps are missing."
            )

    @staticmethod
    def load_maples_dr_diagnosis(path: str | Path) -> pd.DataFrame:
        """
        Load the MAPLES-DR diagnostic file.

        Parameters
        ----------
        path : str
            Path to the MAPLES-DR diagnostic file.
        """
        dr_diagnosis = pd.read_excel(path, sheet_name="DR", index_col=0).rename(
            columns={"Consensus": "dr"} | {f"Retinologist{r}": f"dr_{r}" for r in "ABC"}
        )
        me_diagnosis = pd.read_excel(path, sheet_name="ME", index_col=0).rename(
            columns={"Consensus": "me"} | {f"Retinologist{r}": f"me_{r}" for r in "ABC"}
        )
        dr_me_comments = (
            pd.read_excel(path, sheet_name="Comment", index_col=0)
            .rename(columns={"Retinologist" + r: f"dr_{r}_comment" for r in "ABC"})
            .astype(str)
            .replace("nan", "")
        )
        return dr_diagnosis.join(me_diagnosis).join(dr_me_comments)

    # --- MESSIDOR path configuration ---
    @staticmethod
    def discover_messidor_images(
        images: list[str], path: Optional[str | Path] = None
    ) -> Dict[str, Path]:
        """
        Discover the MESSIDOR images corresponding to the given MAPLES-DR images.

        :param images: List of MAPLES-DR images names. The image name should not contain the extension.
        :param path: Path to the MESSIDOR dataset.
        """
        path = Path(path)
        messidor_paths = {}

        if not path.is_dir():
            raise InvalidConfigError(f"Invalid MESSIDOR path: {path} is not a folder.")

        # Scan the MESSIDOR subfolders and list images.

        # Scan MESSIDOR subfolders to find each MAPLES-DR images.
        missing_images = set(images)
        for img_path in chain.from_iterable(
            path.glob(f"**/*.{ext}") for ext in ("tif", "jpg", "png")
        ):
            try:
                missing_images.remove(img_path.stem)
            except KeyError:
                pass
            else:
                messidor_paths[img_path.stem] = img_path.absolute()
                if len(missing_images) == 0:
                    return messidor_paths

        # If some images are missing, try to unzip the MESSIDOR subfolders.
        unzip_folder = path / "maples_dr"
        unzip_folder.mkdir(parents=True, exist_ok=True)

        total_missing = len(missing_images)
        with RichProgress.iteration(
            "Retrieving MESSIDOR fundus from zip files...",
            total=total_missing,
            done_message=f"Extracted {total_missing} images in {'{t}'} second.",
        ) as progress:
            for zip_path in path.glob("**/*.zip"):
                with ZipFile(zip_path, "r") as zip_file:
                    # If the zip file contains a MESSIDOR subfolder, unzip it.
                    for zip_info in zip_file.infolist():
                        if not zip_info.is_dir() and zip_info.filename.endswith(".tif"):
                            img_filename = Path(zip_info.filename).name
                            img_stem = Path(zip_info.filename).stem
                            try:
                                missing_images.remove(img_stem)
                            except KeyError:
                                continue

                            zip_info.filename = img_filename
                            zip_file.extract(zip_info, unzip_folder)
                            messidor_paths[img_stem] = (
                                unzip_folder / img_filename
                            ).absolute()

                            progress.update(1)
                            if len(missing_images) == 0:
                                return messidor_paths

        # If some images are still missing, raise an error.
        raise InvalidConfigError(
            f"The provided folder to the MESSIDOR dataset is incomplete: "
            f"{len(missing_images)} images included in MAPLES-DR are missing."
        )

    # === DATASETS FACTORIES ===
    def load_dataset(
        self, subset: DatasetSubset | str | list[str] = DatasetSubset.ALL
    ) -> Dataset:
        """
        Return the MAPLES-DR dataset.

        :param subset: Subset of the dataset to return. If None, return the whole dataset.
                       Must be either None, "train" or "test" or a list of valid image name.
        """
        if not self.is_configured():
            self.configure()
        if isinstance(subset, list):
            # Check that all images are valid.
            valid_images = set(self.image_names(DatasetSubset.ALL_WITH_DUPLICATES))
            for img in subset:
                if img not in valid_images:
                    raise ValueError(
                        f"Invalid image name: {img} is not a valid image name."
                    )
            names = subset
        else:
            names = self.image_names(subset)

        paths = {}
        if self._messidor_paths is not None:
            paths[FundusField.FUNDUS.value] = [
                self._messidor_paths[name] for name in names
            ]

        biomarkers_folder = {
            BiomarkerField.BRIGHT_UNCERTAINS.value: "BrightUncertains",
            BiomarkerField.COTTON_WOOL_SPOTS.value: "CottonWoolSpots",
            BiomarkerField.DRUSENS.value: "Drusens",
            BiomarkerField.EXUDATES.value: "Exudates",
            BiomarkerField.HEMORRHAGES.value: "Hemorrhages",
            BiomarkerField.MACULA.value: "Macula",
            BiomarkerField.MICROANEURYSMS.value: "Microaneurysms",
            BiomarkerField.NEOVASCULARIZATION.value: "Neovascularization",
            BiomarkerField.OPTIC_CUP.value: "OpticCup",
            BiomarkerField.OPTIC_DISC.value: "OpticDisc",
            BiomarkerField.RED_UNCERTAINS.value: "RedUncertains",
            BiomarkerField.VESSELS.value: "Vessels",
        }
        for bio, bio_folder in biomarkers_folder.items():
            folder = self.maples_dr_folder / "annotations" / bio_folder
            paths[bio] = [
                folder / (n + ".png") if self.is_biomarker_segmented(bio, n) else None
                for n in names
            ]

        preannotations_folder = {
            BiomarkerField.EXUDATES.value: "Exudates",
            BiomarkerField.HEMORRHAGES.value: "Hemorrhages",
            BiomarkerField.MICROANEURYSMS.value: "Microaneurysms",
            BiomarkerField.VESSELS.value: "Vessels",
        }
        for bio, bio_folder in preannotations_folder.items():
            paths[bio + "_pre"] = [
                self.maples_dr_folder / "preannotations" / bio_folder / (name + ".png")
                for name in names
            ]

        data = pd.DataFrame(paths, index=names)

        # Add the diagnosis and annotations infos.
        data = data.join(self._diagnosis).join(self._annotations_infos)

        return Dataset(data, self.cfg, self._messidor_ROIs)

    # === UTILITIES ===
    def image_names(
        self,
        subset: DatasetSubset | str = DatasetSubset.ALL,
        extension: bool | str = False,
    ) -> list[str]:
        """
        Return the list of images names of the given subset.

        Parameters
        ----------
        subset
            Subset to return the images names from. If None, return all images names.
            Must be either None, "train", "test" or "duplicates".

        extension
            Control whether the images names should include the extension or not.
            - If False (default), return the images names without the extension.
            - If True, return the images names with a png extension.
            - If a string, return the images names with the given extension.
        """
        if not self.is_configured():
            raise NotConfiguredError()
        subset = DatasetSubset(subset)

        names = []
        if subset in (
            DatasetSubset.TRAIN,
            DatasetSubset.ALL,
            DatasetSubset.ALL_WITH_DUPLICATES,
        ):
            names += self.dataset_record["train"]
        if subset in (
            DatasetSubset.TEST,
            DatasetSubset.ALL,
            DatasetSubset.ALL_WITH_DUPLICATES,
        ):
            names += self.dataset_record["test"]
        if subset in (DatasetSubset.DUPLICATES, DatasetSubset.ALL_WITH_DUPLICATES):
            names += list(self.dataset_record["duplicates"].values())
        if extension:
            if isinstance(extension, bool):
                extension = "png"
            names = [name + "." + extension for name in names]
        return names

    def is_biomarker_segmented(
        self, biomarker: BiomarkerField | str, name: str
    ) -> bool:
        """
        Check if the given biomarker is segmented in the MAPLES-DR dataset.

        .. note::

            The macula segmentation is missing for one image centered on the optic disc.

            The optic cup boundaries are too fuzzy to be segmented on six images.


        Parameters
        ----------
        biomarker:
            The biomarker to check.
        name:
            The image name.

        Returns
        -------
        bool
            True if the biomarker is segmented, False otherwise.
        """
        if not self.is_configured():
            raise NotConfiguredError()
        biomarker = BiomarkerField.parse(biomarker)
        if (
            biomarker is BiomarkerField.MACULA
            and name in self.dataset_record["no_macula"]
        ):
            return False
        if (
            biomarker is BiomarkerField.OPTIC_CUP
            and name in self.dataset_record["no_cup"]
        ):
            return False
        return True


GLOBAL_LOADER = DatasetLoader()


def download(url: str, path: str | Path, description: Optional[str] = None):
    """
    Download the file at the given url to the given path.
    """
    response = urlopen(url)
    with RichProgress.download(
        description, byte_size=int(response.info()["Content-length"])
    ) as progress:
        with open(path, "wb") as dest_file:
            for data in iter(partial(response.read, 32768), b""):
                dest_file.write(data)
                progress.update(len(data))
