"""
    @file:              dicom_reader.py
    @Author:            Maxence Larose

    @Creation Date:     10/2021
    @Last modification: 03/2022

    @Description:       This file contains the DicomReader class which is used to read dicom files contained in a given
                        folder.
"""

from collections import defaultdict
from glob import glob
import logging
import os
from typing import Dict, List, NamedTuple, Optional, Set, Tuple, Union

import pydicom
import SimpleITK as sitk

from delia.utils.data_model import ImageDataModel
from delia.readers.segmentation.segmentation_strategy import SegmentationStrategies
from delia.utils.tools import is_path_valid

_logger = logging.getLogger(__name__)


class DicomReader:
    """
    A class used to read dicom files contained in a given folder.
    """

    UNKNOWN = "Unknown"

    class SeriesData(NamedTuple):
        """
        Tag value namedtuple to simplify management of values.
        """
        tag_value: str
        paths_to_dicoms_from_series: List[str]
        dicom_header: pydicom.dataset.FileDataset

    def __init__(
            self,
            path_to_patient_folder: str,
            tag: Union[str, Tuple[int, int]],
            load_segmentations: bool = True
    ):
        """
        Constructor of the class DicomReader.

        Parameters
        ----------
        path_to_patient_folder : str
            Path to the folder containing the patient DICOM files.
        tag : Union[str, Tuple[int, int]]
            Keyword or tuple of the DICOM tag to use while selecting which files to extract.
        load_segmentations : bool
            Whether to load segmentations or not.
        """
        super(DicomReader, self).__init__()

        is_path_valid(path=path_to_patient_folder)
        self._path_to_patient_folder = path_to_patient_folder
        self._load_segmentations = load_segmentations
        self.tag = tag

        self._images_series_data_dict = {}
        self._segmentations_series_data_dict = {}
        self.__series_data_dict = self.__get_series_data_dict()

    @staticmethod
    def get_dicom_header(
            path_to_dicom: str,
            show: bool = False
    ) -> pydicom.dataset.FileDataset:
        """
        Get a dicom header given the path to the dicom.

        Parameters
        ----------
        path_to_dicom : str
            The path to the dicom file.
        show : bool
            Show dicom header in console.

        Returns
        -------
        loaded_dicom : pydicom.dataset.FileDataset
            Loaded DICOM dataset.
        """
        loaded_dicom = pydicom.dcmread(path_to_dicom, stop_before_pixels=True)

        if show:
            _logger.info(loaded_dicom)

        return loaded_dicom

    @staticmethod
    def _get_tag_value(dicom_header: pydicom.FileDataset, tag: Union[str, Tuple[int, int]]) -> str:
        """
        Get tag value of given DICOM header.

        Parameters
        ----------
        dicom_header : pydicom.dataset.FileDataset
            Loaded DICOM dataset.
        tag : Union[str, Tuple[int, int]]
            Keyword or tuple of the DICOM tag to use while selecting which files to extract.

        Returns
        -------
        tag_value : str
            tag value.
        """
        if hasattr(dicom_header, tag):
            if isinstance(dicom_header[tag].value, str):
                return dicom_header[tag].value
            else:
                return dicom_header[tag].repval
        else:
            return DicomReader.UNKNOWN

    def __get_paths_to_dicoms_from_series(self, path_to_dicom_folder: str) -> Dict[str, List[str]]:
        """
        Get all series IDs from a patient's dicom folder.

        Parameters
        ----------
        path_to_dicom_folder : str
            Path to a folder containing only DICOM files.

        Returns
        -------
        paths_to_dicoms_from_series : Dict[str, List[str]]
            Dictionary of lists of paths to dicoms from each series.
        """
        paths = [
            y for x in os.walk(path_to_dicom_folder) for y in glob(os.path.join(x[0], '*')) if os.path.isfile(y) is True
        ]

        paths_and_headers_for_each_series = defaultdict(list)
        for path in paths:
            dicom_header = self.get_dicom_header(path_to_dicom=path)
            paths_and_headers_for_each_series[dicom_header.SeriesInstanceUID].append((path, dicom_header))

        paths_to_dicoms_from_series = defaultdict(list)
        for series_id, result in paths_and_headers_for_each_series.items():
            if len(result) != 1:
                if hasattr(result[1][1], "SliceLocation"):
                    paths_to_dicoms_from_series[series_id] = [
                        i[0] for i in sorted(result, key=lambda s: s[1].SliceLocation)
                    ]
                else:
                    paths_to_dicoms_from_series[series_id] = [
                        i[0] for i in result
                    ]
            else:
                paths_to_dicoms_from_series[series_id] = [result[0][0]]

        return paths_to_dicoms_from_series

    def __get_series_data_dict(self) -> Dict[str, SeriesData]:
        """
        Get the series data from series IDs and the path to the patient's dicom folders.

        Returns
        -------
        series_data_dict : Dict[str, SeriesData]
            Dictionary of the data from the selected series.
        """
        series_data_dict: Dict[str, DicomReader.SeriesData] = {}
        all_patient_ids: Set[str] = set()
        for series_id, paths in self.__get_paths_to_dicoms_from_series(self._path_to_patient_folder).items():
            path_to_first_dicom_of_series = paths[0]
            loaded_dicom_header = self.get_dicom_header(path_to_dicom=path_to_first_dicom_of_series)
            all_patient_ids.add(loaded_dicom_header.PatientID)

            series_data = self.SeriesData(
                tag_value=self._get_tag_value(loaded_dicom_header, self.tag),
                paths_to_dicoms_from_series=paths,
                dicom_header=loaded_dicom_header
            )
            series_data_dict[series_id] = series_data

        if len(all_patient_ids) == 0:
            raise AssertionError(f"The patient folder must contain DICOM files. This is not the case for the patient "
                                 f"whose data is currently being downloaded since no DICOM files are found in its "
                                 f"folder.")
        elif len(all_patient_ids) > 1:
            raise AssertionError(f"All DICOM files in the same folder must belong to the same patient. This is not "
                                 f"the case for the patient whose data is currently being downloaded since the "
                                 f"identifiers {all_patient_ids} are found in their folder.")

        for series_id, series_data in series_data_dict.items():
            if series_data.dicom_header.Modality in SegmentationStrategies.get_available_modalities():
                if self._load_segmentations:
                    self._segmentations_series_data_dict[series_id] = series_data
            else:
                self._images_series_data_dict[series_id] = series_data

        return series_data_dict

    @staticmethod
    def __get_3d_sitk_image_from_dicom_series(
            paths_to_dicoms_from_series: List[str],
    ) -> sitk.Image:
        """
        Get a 3D image array from a list of dicom paths associated to the same series.

        Parameters
        ----------
        paths_to_dicoms_from_series : List[str]
            A list of the paths to each dicom contained in the series.

        Returns
        -------
        image : sitk.Image
            3D SimpleITK image obtained from the series.
        """
        series_reader = sitk.ImageSeriesReader()
        series_reader.SetFileNames(paths_to_dicoms_from_series)

        image = series_reader.Execute()

        return image

    def get_dicom_headers(self, remove_segmentations: bool = False) -> List[pydicom.dataset.FileDataset]:
        """
        List of all the patient's dicom headers.

        Parameters
        ----------
        remove_segmentations : bool
            Whether or not to keep the segmentation files headers, default = False.

        Returns
        -------
        dicom_headers : List[pydicom.dataset.FileDataset]
            A list of the patient's DICOM headers.
        """
        dicom_headers = []
        for idx, series_data in enumerate(self.__series_data_dict.values()):
            if idx == 0:
                _logger.info(f"Patient ID : {series_data.dicom_header.PatientID}")
                _logger.debug(f"Path to images folder : {self._path_to_patient_folder}")

            _logger.debug(f"  {self.tag} : {series_data.tag_value}")

            if not remove_segmentations:
                dicom_headers.append(series_data.dicom_header)
            elif remove_segmentations:
                if series_data.dicom_header.Modality not in SegmentationStrategies.get_available_modalities():
                    dicom_headers.append(series_data.dicom_header)

        return dicom_headers

    def get_images_data(
            self,
            remove_segmentations: bool = False
    ) -> List[ImageDataModel]:
        """
        List of tuples containing simpleITK 3D images array and their corresponding dicom header. Each element in the
        list corresponds to an image series.

        Parameters
        ----------
        remove_segmentations : bool
            Whether to keep the segmentation files headers, default = False.

        Returns
        -------
        image_data_list : List[ImageDataModel]
            A list of the patient's images data.
        """
        images_data = []
        for idx, series_data in enumerate(self.__series_data_dict.values()):
            if idx == 0:
                _logger.info(f"Tag values found in the patient's images folder :")
            _logger.info(f"  {self.tag} {idx + 1}: {series_data.tag_value}")

            if remove_segmentations and (series_data.dicom_header.Modality in
                                         SegmentationStrategies.get_available_modalities()):
                pass
            elif series_data.tag_value == DicomReader.UNKNOWN:
                pass
            else:
                try:
                    image = self.__get_3d_sitk_image_from_dicom_series(
                        paths_to_dicoms_from_series=series_data.paths_to_dicoms_from_series
                    )

                    image_data = ImageDataModel(
                        dicom_header=series_data.dicom_header,
                        simple_itk_image=image,
                        paths_to_dicoms=series_data.paths_to_dicoms_from_series
                    )

                    images_data.append(image_data)

                except RuntimeError as e:
                    _logger.error(
                        f"      RuntimeError : {e}. Simple ITK raised an error while loading the series named "
                        f"{series_data.tag_value}. This series is therefore ignored."
                    )

        return images_data
