import inspect
import os
import warnings
from abc import ABC
from pathlib import Path
from typing import Any, Dict, Optional, Sequence, Set, Union

import pydicom

from dosma.core.io import format_io_utils as fio_utils
from dosma.core.io.dicom_io import DicomReader
from dosma.core.io.format_io import ImageDataFormat
from dosma.core.med_volume import MedicalVolume
from dosma.defaults import preferences
from dosma.tissues.tissue import Tissue
from dosma.utils import io_utils


def _contains_type(value, types):
    """Returns ``True`` if any value is an instance of ``types``."""
    if isinstance(value, types):
        return True
    if not isinstance(value, str) and isinstance(value, (Sequence, Set)) and len(value) > 0:
        return any(_contains_type(x, types) for x in value)
    elif isinstance(value, Dict):
        return _contains_type(value.keys(), types) or _contains_type(value.values(), types)
    return isinstance(value, types)


class ScanIOMixin(ABC):
    # This is just a summary on variables used in this abstract class,
    # the proper values/initialization should be done in child class.
    NAME: str
    __DEFAULT_SPLIT_BY__: Optional[str]
    _from_file_args: Dict[str, Any]

    @classmethod
    def from_dicom(
        cls,
        dir_or_files,
        group_by=None,
        ignore_ext: bool = False,
        num_workers: int = 0,
        verbose: bool = False,
        **kwargs,
    ):
        """Load scan from dicom files.

        Args:
            dir_or_files (str): The path to dicom directory or files.
            group_by: DICOM field tag name or tag number used to group dicoms. Defaults
                to scan's ``__DEFAULT_SPLIT_BY__``.
            ignore_ext (bool, optional): If `True`, ignore extension (`.dcm`)
                when loading dicoms from directory.
            num_workers (int, optional): Number of workers to use for loading.
            verbose (bool, optional): If ``True``, enable verbose logging for dicom loading.
            kwargs: Other keywords required to construct scan.

        Returns:
            The scan.
        """
        dr = DicomReader(num_workers, verbose)
        if group_by is None:
            group_by = cls.__DEFAULT_SPLIT_BY__
        volumes = dr.load(dir_or_files, group_by, ignore_ext)

        if isinstance(dir_or_files, (str, Path, os.PathLike)):
            dir_or_files = os.path.abspath(dir_or_files)
        else:
            dir_or_files = type(dir_or_files)([os.path.abspath(x) for x in dir_or_files])

        scan = cls(volumes, **kwargs)
        scan._from_file_args = {
            "dir_or_files": dir_or_files,
            "ignore_ext": ignore_ext,
            "group_by": group_by,
            "_type": "dicom",
        }

        return scan

    @classmethod
    def from_dict(cls, data: Dict[str, Any], force: bool = False):
        """Loads class from data dictionary.

        Args:
            data (Dict): The data.
            force (bool, optional): If ``True``, writes attributes even if they do not exist.
                Use with caution.

        Returns:
            The scan

        Examples:
            >>> scan = ... # some scan
            >>> filepath = scan.save("/path/to/base/directory")
            >>> scan_from_saved = type(scan).from_dict(io_utils.load_pik(filepath))
            >>> scan_from_dict = type(scan).from_dict(scan.__dict__)
        """
        # TODO: Add check for deprecated and converted attribute names.
        data = cls._convert_attr_name(data)

        # TODO: Convert metadata to appropriate type.
        # Converting metadata type is important when loading MedicalVolume data (for example).
        # The data is stored as a path, but should be loaded as a MedicalVolume.
        data = cls.load_custom_data(data)

        signature = inspect.signature(cls)
        init_metadata = {k: v for k, v in data.items() if k in signature.parameters}
        scan = cls(**init_metadata)
        for k in init_metadata.keys():
            data.pop(k)

        for k, v in data.items():
            if not hasattr(scan, k) and not force:
                warnings.warn(f"{cls.__name__} does not have attribute {k}. Skipping...")
                continue
            scan.__setattr__(k, v)

        return scan

    def save(
        self,
        path: str,
        save_custom: bool = False,
        image_data_format: ImageDataFormat = None,
        num_workers: int = 0,
    ):
        """Saves scan data to disk with option for custom saving.

        Custom saving may be useful to reduce redundant saving and/or save data in standard
        compatible formats (e.g. medical images - nifti/dicom), which are not feasible with
        python serialization libraries, like pickle.

        When ``save_custom=True``, this method overloads standard pickling with customizable
        saving by first saving data in customizable way (e.g. MedicalVolume -> Nifti file),
        and then pickling the reference to the saved object (e.g. Nifti filepath).

        Currently certain custom saving of objects such as ``pydicom.FileDataset`` and
        :cls:`Tissue` objects are not supported.

        To load the data, do the following:

        >>> filepath = scan.save("/path/to/directory", save_custom=True)
        >>> scan_loaded = type(scan).load(io_utils.load_pik(filepath))

        Args:
            path (str): Directory where data is stored.
            data_format (ImageDataFormat, optional): Format to save data.
                Defaults to ``preferences.image_data_format``.
            save_custom (bool, optional): If ``True``, saves data in custom way specified
                by :meth:`save_custom_data` in format specified by ``data_format``. For
                example, for default classes this will save :cls:`MedicalVolume` data
                to nifti/dicom files as specified by ``image_data_format``.
            image_data_format (ImageDataFormat, optional): The data format to save
                :cls:`MedicalVolume` data. Only used if save_custom is ``True``.
            num_workers (int, bool): Number of workers for saving custom data.
                Only used if save_custom is ``True``.

        Returns:
            str: The path to the pickled file.
        """
        if image_data_format is None:
            image_data_format = preferences.image_data_format

        save_dirpath = path  # self._save_dir(path)
        os.makedirs(save_dirpath, exist_ok=True)
        filepath = os.path.join(save_dirpath, "%s.data" % self.NAME)

        metadata: Dict = {}
        for attr in self.__serializable_variables__():
            metadata[attr] = self.__getattribute__(attr)

        if save_custom:
            metadata = self._save(
                metadata, save_dirpath, image_data_format=image_data_format, num_workers=num_workers
            )

        io_utils.save_pik(filepath, metadata)
        return filepath

    @classmethod
    def load(cls, path_or_data: Union[str, Dict], num_workers: int = 0):
        """Load scan.

        This method overloads the :func:`from_dict` method by supporting loading from a file
        in addition to the data dictionary. If loading and constructing a scan using
        :func:`from_dict` fails, defaults to loading data from original dicoms
        (if ``self._from_file_args`` is initialized).

        Args:
            path_or_data (Union[str, Dict]): Pickle file to load or data dictionary.
            num_workers (int, optional): Number of workers to use for loading.

        Returns:
            ScanSequence: Of type ``cls``.

        Raises:
            ValueError: If ``scan`` cannot be constructed.
        """
        if isinstance(path_or_data, (str, Path, os.PathLike)):
            if os.path.isdir(path_or_data):
                path_or_data = os.path.join(path_or_data, f"{cls.NAME}.data")

            if not os.path.isfile(path_or_data):
                raise FileNotFoundError(f"File {path_or_data} does not exist")
            data = io_utils.load_pik(path_or_data)
        else:
            data = path_or_data

        try:
            scan = cls.from_dict(data)
            return scan
        except Exception:
            warnings.warn(
                f"Failed to load {cls.__name__} from data. Trying to load from dicom file."
            )

        data = cls._convert_attr_name(data)
        data = cls.load_custom_data(data, num_workers=num_workers)

        scan = None
        if "_from_file_args" in data:
            dicom_args = data.pop("_from_file_args")
            assert dicom_args.pop("_type") == "dicom"
            scan = cls.from_dicom(**dicom_args, num_workers=num_workers)
        elif "dicom_path" in data:
            # Backwards compatibility
            dicom_path = data.pop("dicom_path")
            ignore_ext = data.pop("ignore_ext", False)
            group_by = data.pop("split_by", cls.__DEFAULT_SPLIT_BY__)
            scan = cls.from_dicom(
                dicom_path, ignore_ext=ignore_ext, group_by=group_by, num_workers=num_workers
            )

        if scan is None:
            raise ValueError(f"Data is insufficient to construct {cls.__name__}")

        for k, v in data.items():
            if not hasattr(scan, k):
                warnings.warn(f"{cls.__name__} does not have attribute {k}. Skipping...")
                continue
            scan.__setattr__(k, v)

        return scan

    def save_data(
        self, base_save_dirpath: str, data_format: ImageDataFormat = preferences.image_data_format
    ):
        """Deprecated: Alias for :func:`self.save`."""
        warnings.warn(
            "save_data is deprecated since v0.0.13 and will no longer be "
            "available in v0.1. Use `save` instead.",
            DeprecationWarning,
        )
        return self.save(base_save_dirpath, data_format)

    def _save(
        self,
        metadata: Dict[str, Any],
        save_dir: str,
        fname_fmt: Dict[Union[str, type], str] = None,
        **kwargs,
    ):
        if fname_fmt is None:
            fname_fmt = {}

        default_fname_fmt = {MedicalVolume: "image-{}"}
        for k, v in default_fname_fmt.items():
            if k not in fname_fmt:
                fname_fmt[k] = v

        for attr in metadata.keys():
            val = metadata[attr]
            path = fname_fmt.get(attr, None)

            if path is None:
                path = os.path.abspath(os.path.join(save_dir, attr))
            if not os.path.isabs(path):
                path = os.path.join(save_dir, attr, path)
            try:
                metadata[attr] = self.save_custom_data(val, path, fname_fmt, **kwargs)
            except Exception as e:
                raise RuntimeError(f"Failed to save metadata {attr} - {e}")

        return metadata

    def save_custom_data(
        self, metadata, paths, fname_fmt: Dict[Union[str, type], str] = None, **kwargs
    ):
        """
        Finds all attributes of type MedicalVolume or Sequence/Mapping to MedicalVolume
        and saves them.
        """
        if isinstance(metadata, (Dict, Sequence, Set)):
            if isinstance(paths, str):
                paths = [paths] * len(metadata)
            else:
                assert len(paths) == len(metadata)

        if isinstance(metadata, Dict):
            keys = metadata.keys()
            if isinstance(paths, Dict):
                paths = [paths[k] for k in keys]
            paths = [os.path.join(_path, f"{k}") for k, _path in zip(keys, paths)]
            values = self.save_custom_data(metadata.values(), paths, fname_fmt, **kwargs)
            metadata = {k: v for k, v in zip(keys, values)}
        elif not isinstance(metadata, str) and isinstance(metadata, (Sequence, Set)):
            values = list(metadata)
            paths = [os.path.join(_path, "{:03d}".format(i)) for i, _path in enumerate(paths)]
            values = [
                self.save_custom_data(_x, _path, fname_fmt, **kwargs)
                for _x, _path in zip(values, paths)
            ]
            if not isinstance(values, type(metadata)):
                metadata = type(metadata)(values)
            else:
                metadata = values
        else:
            formatter = [fname_fmt.get(x) for x in type(metadata).__mro__]
            formatter = [x for x in formatter if x is not None]
            if len(formatter) == 0:
                formatter = None
            else:
                formatter = formatter[0]
            metadata = self._save_custom_data_base(metadata, paths, formatter, **kwargs)

        return metadata

    def _save_custom_data_base(self, metadata, path, formatter: str = None, **kwargs):
        """The base condition for :meth:`save_custom_data`.

        Args:
            metadata (Any): The data to save.
            path (str): The path to save the data.
            formatter (str, optional): If provided, this formatted string
                will be used to format ``path``.
        """
        out = {"__dtype__": type(metadata)}
        # TODO: Add support for num workers.
        # num_workers = kwargs.pop("num_workers", 0)

        if formatter:
            path = os.path.join(os.path.dirname(path), formatter.format(os.path.basename(path)))

        if isinstance(metadata, MedicalVolume):
            image_data_format = kwargs.get("image_data_format", preferences.image_data_format)
            # TODO: Once, `files` property added to MedicalVolume, check if property is
            # set before doing saving.
            path = fio_utils.convert_image_data_format(path, image_data_format)
            metadata.save_volume(path, data_format=image_data_format)
            out["__value__"] = path
        else:
            out = metadata

        return metadata

    @classmethod
    def _convert_attr_name(cls, data: Dict[str, Any]):
        return data

    @classmethod
    def load_custom_data(cls, data: Any, **kwargs):
        """Recursively converts data to appropriate types.

        By default, this loads all :class:`MedicalVolume` objects from their corresponding paths.

        Args:
            data (Any): The data. Can either be the dictionary or the metadata value.
                If the data corresponds to a custom type, it should have the following
                schema:

                {
                    '__dtype__': The type of the data.
                    '__value__': The value from which object of type __dtype__ can be constructed.
                }

            **kwargs: Keyword Arguments to pass to :meth:`_load_custom_data_base`.

        Returns:
            Any: The loaded metadata.
        """
        dtype = type(data)
        if isinstance(data, Dict) and "__value__" in data:
            dtype = data["__dtype__"]
            data = data["__value__"]

        if issubclass(dtype, Dict):
            keys = cls.load_custom_data(data.keys(), **kwargs)
            values = cls.load_custom_data(data.values(), **kwargs)
            data = {k: v for k, v in zip(keys, values)}
        elif not issubclass(dtype, str) and issubclass(dtype, (list, tuple, set)):
            data = dtype([cls.load_custom_data(x, **kwargs) for x in data])
        else:
            data = cls._load_custom_data_base(data, dtype, **kwargs)

        return data

    @classmethod
    def _load_custom_data_base(cls, data, dtype=None, **kwargs):
        """The base condition for :meth:`load_custom_data`.

        Args:
            data:
            dtype (type): The data type.

        Return:
            The loaded data.
        """
        if dtype is None:
            dtype = type(data)

        # TODO: Add support for loading with num_workers
        num_workers = kwargs.get("num_workers", 0)
        if isinstance(data, str) and issubclass(dtype, MedicalVolume):
            data = fio_utils.generic_load(data, num_workers=num_workers)

        return data

    def __serializable_variables__(
        self, ignore_types=(pydicom.FileDataset, pydicom.Dataset, Tissue), ignore_attrs=()
    ) -> Set:
        """
        By default, all instance attributes are serialized except those
        corresponding to headers, :class:`MedicalVolume`(s), or :class:`Tissues`.
        Properties and class attributes are also not stored. Class attributes are
        indentified using the PEP8 nomenclature of all caps variables.

        Note:
            This method has not been profiled, but times may be large if
            the instance contains many variables. Currently this is not
            cached as attributes values can change and, as a result, must
            be inspected.
        """
        serializable = []
        for attr, value in self.__dict__.items():
            if attr in ignore_attrs or _contains_type(value, ignore_types):
                continue
            if attr.startswith("temp") or attr.startswith("_temp"):
                continue
            if attr.upper() == attr or (attr.startswith("__") and attr.endswith("__")):
                continue
            if callable(value) or isinstance(value, property):
                continue
            serializable.append(attr)

        return set(serializable)