#    Copyright 2021 SECTRA AB
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

import datetime
from typing import Callable, List, Optional, Sequence, Tuple, Union

from highdicom.content import (IssuerOfIdentifier, SpecimenCollection,
                               SpecimenDescription, SpecimenPreparationStep,
                               SpecimenSampling, SpecimenStaining)
from opentile.common import Tiler
from opentile.interface import OpenTile
from pydicom.dataset import Dataset
from pydicom.sequence import Sequence as DicomSequence
from pydicom.sr.coding import Code
from pydicom.uid import UID as Uid
from pydicom.uid import generate_uid
from wsidicom.conceptcode import (AnatomicPathologySpecimenTypesCode,
                                  ConceptCode, SpecimenEmbeddingMediaCode,
                                  SpecimenFixativesCode,
                                  SpecimenPreparationProcedureCode,
                                  SpecimenSamplingProcedureCode,
                                  SpecimenStainsCode)


def get_image_type(image_flavor: str, level_index: int) -> List[str]:
    """Return image type.

    Parameters
    ----------
    image_flavor: str
        Image flavor ('VOLUME', 'LABEL', 'OVERVIEW')
    level_index: int:
        Pyramidal level index of the image.

    Returns
    ----------
    List[str]
        Image type.
    """
    if image_flavor == 'VOLUME' and level_index == 0:
        resampled = 'NONE'
    else:
        resampled = 'RESAMPLED'

    return ['ORGINAL', 'PRIMARY', image_flavor, resampled]


def create_base_dataset(
    modules: Optional[Union[Dataset, Sequence[Dataset]]] = None
) -> Dataset:
    """Create a base dataset by combining module datasets with a minimal
    wsi dataset.

    Parameters
    ----------
    modules: Optional[Union[Dataset, Sequence[Dataset]]] = None

    Returns
    ----------
    Dataset
        Combined base dataset.
    """
    base_dataset = create_wsi_dataset()
    if modules is None:
        modules = create_default_modules()
    if isinstance(modules, Sequence):
        for module in modules:
            base_dataset.update(module)
    elif isinstance(modules, Dataset):
        base_dataset.update(modules)
    else:
        raise TypeError(
            'datasets parameter should be single or list of Datasets'
        )
    return base_dataset


def populate_base_dataset(
    tiler: Tiler,
    base_dataset: Dataset,
    include_confidential: bool = True
) -> Dataset:
    """Populate dataset with properties from tiler, if present.
    Parameters
    ----------
    tiler: Tiler
        A opentile Tiler.
    base_dataset: Dataset
        Dataset to append properties to.
    include_confidential: bool = True
        If to include confidential properties (see https://dicom.nema.org/medical/dicom/current/output/html/part15.html#table_E.1-1)  # NOQA

    Returns
    ----------
    Dataset
        Dataset with added properties.
    """
    for property, value in tiler.properties.items():
        if property == 'aquisition_datetime' and include_confidential:
            base_dataset.AcquisitionDateTime = value
        elif property == 'device_serial_number' and include_confidential:
            base_dataset.DeviceSerialNumber = value
        elif property == 'manufacturer':
            base_dataset.Manufacturer = value
        elif property == 'model':
            base_dataset.ManufacturerModelName = value
        elif property == 'software_versions':
            base_dataset.SoftwareVersions = value
        elif property == 'lossy_image_compression_method':
            base_dataset.LossyImageCompressionMethod = value
        elif property == 'lossy_image_compression_ratio':
            base_dataset.LossyImageCompressionRatio = value
        elif property == 'photometric_interpretation':
            base_dataset.PhotometricInterpretation = value
    return base_dataset


def create_wsi_dataset(
    offset: Tuple[float, float] = (0.0, 0.0),
    uid_generator: Callable[..., Uid] = generate_uid
) -> Dataset:
    """Return dataset containing (parts of) SOP common, general series, whole
    slide icrosocy series, frame of reference, acquisition context, multi-frame
    dimension, and whole slide microscopy image modules.

    Some modules returned 'not complete', and completed during image data
    import or file save():
        SOP common module:
            SOPInstanceUID
        Whole slide microscopy image module:
            ImageType
            Aquisition DateTime
            SamplesPerPixel
            PhotometricInterpretation
            PlanarConfiguration
            NumberOfFrames
            BitsAllocated
            BitsStored
            HighBit
            PixelRepresentation
            LossyImageCompression (and conditionals)
            ImagedVolumeWidth
            ImagedVolumeWidth
            ImagedVolumeDepth
            TotalPixelMatrixColumns
            TotalPixelMatrixRows
            FocusMethod
            ExtendedDepthOfField (and conditionals)

    Parameters
    ----------
    image_offset: Tuple[float, float] = 0.0
        X and Y offset (in mm) to first pixel in image data.
    uid_generator: Callable[..., Uid] = generate_uid
        Function that can generate Uids.

    Returns
    ----------
    Dataset
        WSI dataset.
    """
    dataset = Dataset()

    # SOP common module (SOPInstanceUID written on save())
    dataset.SOPClassUID = '1.2.840.10008.5.1.4.1.1.77.1.6'

    # dataset.StudyInstanceUID = uid_generator()
    # General series and Whole slide Microscopy modules
    dataset.SeriesInstanceUID = uid_generator()
    dataset.Modality = 'SM'

    # Frame of reference module
    dataset.FrameOfReferenceUID = uid_generator()

    # Acquisition context module (empty)
    dataset.AcquisitionContextSequence = DicomSequence()

    # Multi-frame Dimension module
    dimension_organization_uid = uid_generator()
    dimension_organization_sequence = Dataset()
    dimension_organization_sequence.DimensionOrganizationUID = (
        dimension_organization_uid
    )
    dataset.DimensionOrganizationSequence = DicomSequence(
        [dimension_organization_sequence]
    )

    # Whole slide micropscopy image module (most filled when importing file)
    dataset.BurnedInAnnotation = 'NO'
    dataset.SpecimenLabelInImage = 'NO'
    dataset.VolumetricProperties = 'VOLUME'
    dataset.ImageOrientationSlide = [-1, 0, 0, 0, -1, 0]
    offset_item = Dataset()
    offset_item.x_offset = offset[0]
    offset_item.y_offset = offset[1]
    dataset.TotalPixelMatrixOriginSequence = DicomSequence([offset_item])
    return dataset


def create_study_module(
    study_id: str = "",
    study_date: Optional[datetime.date] = None,
    study_time: Optional[datetime.time] = None,
    study_accession_number: str = "",
    referring_physician_name: str = "",
    uid_generator: Callable[..., Uid] = generate_uid
) -> Dataset:
    """Create general study module.

    Parameters
    ----------
    study_id: str = ""
        Study identifier. Can be empty.
    Optional[datetime.date] = None
        Date the study started. Can be empty.
    Optional[datetime.time] = None
        Time the study started. Can be empty.
    accession_number: str = ""
        Order for the study. Can be empty.
    referring_physician_name: str = ""
        Name of reffering physician. Can be empty.

    Returns
    ----------
    Dataset
        Dataset containing general study module.
    """
    dataset = Dataset()
    dataset.StudyInstanceUID = uid_generator()
    dataset.StudyID = study_id
    dataset.StudyDate = study_date
    dataset.StudyTime = study_time
    dataset.AccessionNumber = study_accession_number
    dataset.ReferringPhysicianName = referring_physician_name
    return dataset


def create_patient_module(
    name: str = "",
    id: str = "",
    birth_date: Optional[datetime.date] = None,
    sex: str = "",
    identity_removed: bool = False,
    de_indentification_methods: Sequence[Union[str, Code]] = [],
    age_at_extraction: Optional[str] = None
) -> Dataset:
    """Create patient and patient study modules.

    Parameters
    ----------
    name: str = ""
        Patient's full name. Can be empty.
    id: str = ""
        Primary identifier for the Patient. Can be empty.
    birth_date: str = ""
        Birth date of the Patient. Can be empty.
    sex: str = ""
        Sex of the named Patient. F/M/O. Can be empty.
    identity_removed: bool = False
        True if partient idendity has been removed from attributes in dataset
        and in image data.
    de_indentification_methods: Sequence[Union[str, Code]] = []
        Method of de-indentification. Can be descriptive strings and/or coded
        values.
    age_at_extraction: Optional[str] = None
        Age of the Patient at extraction.

    Returns
    ----------
    Dataset
        Dataset containing patient and patient study modules.
    """
    dataset = Dataset()
    dataset.PatientName = name
    dataset.PatientID = id
    dataset.PatientBirthDate = birth_date
    dataset.PatientSex = sex
    dataset.PatientIdentityRemoved = 'YES' if identity_removed else 'NO'
    if identity_removed and de_indentification_methods is None:
        raise ValueError(
            'de-intification method must be specified if '
            'patient identity has been removed'
            )
    de_indentification_method_strings: List[str] = []
    de_indentification_method_codes: List[Code] = []

    for de_indentification_method in de_indentification_methods:
        if isinstance(de_indentification_method, str):
            de_indentification_method_strings.append(de_indentification_method)
        elif isinstance(de_indentification_method, Code):
            de_indentification_method_codes.append(de_indentification_method)
        else:
            raise TypeError(
                'De-indentification methods should be string or coded value'
            )
    if de_indentification_method_strings != []:
        dataset.DeidentificationMethod = de_indentification_method_strings
    if de_indentification_method_codes != []:
        dataset.DeidentificationMethodCodeSequence = DicomSequence(
            [
                ConceptCode.from_code(de_indentification_method_code).to_ds()
                for de_indentification_method_code
                in de_indentification_method_codes
            ]
        )
    if age_at_extraction is not None:
        dataset.PatientAge = age_at_extraction

    return dataset


def create_device_module(
    manufacturer: str = 'Unkown',
    model_name: str = 'Unkown',
    serial_number: str = 'Unkown',
    software_versions: Sequence[str] = ['Unkown']
) -> Dataset:
    """Create extended equipment module.

    Parameters
    ----------
    manufacturer: str = 'Unkown'.
        Manufacturer of the equipment.
    model_name: str = 'Unkown'
        Manufacturer's model name of the equipment.
    serial_number: str = 'Unkown'
        Manufacturer's serial number of the equipment.
    software_versions: Sequence[str] = ['Unkown']
        Software version of the equipment.

    Returns
    ----------
    Dataset
        Dataset containing extended equipment module.
    """
    dataset = Dataset()
    properties = {
        'Manufacturer': manufacturer,
        'ManufacturerModelName': model_name,
        'DeviceSerialNumber': serial_number,
        'SoftwareVersions': software_versions
    }
    for name, value in properties.items():
        setattr(dataset, name, value)
    return dataset


def create_sample(
    sample_id: str,
    stainings: Sequence[str],
    embedding_medium: Optional[str] = None,
    fixative: Optional[str] = None,
    specimen_id: Optional[str] = None,
    specimen_type: Optional[str] = None,
    specimen_sampling_method: Optional[str] = None,
    anatomical_sites: Optional[Sequence[Tuple[Code, Sequence[Code]]]] = None,
    location: Optional[Union[str, Tuple[float, float, float]]] = None,
    uid_generator: Callable[..., Uid] = generate_uid
) -> Dataset:
    """Create sample description.

    Parameters
    ----------
    sample_id: str
        Identifier for the sample.
    stainings: Sequence[str]
        Stainings used. See SpecimenStainsCode.list() for allowed values.
    embedding_medium: Optional[str] = None
        Embedding medium used. See SpecimenEmbeddingMediaCode.list() for
        allowed values.
    fixative: Optional[str] = None
        Fixative used. See SpecimenFixativesCode.list() for allowed values.
    specimen_id: Optional[str] = None
        Identifier for the specimen the sample was sampled from.
    specimen_type: Optional[str] = None
        Anatotomic type of specimen the sample was sampled from. See
        AnatomicPathologySpecimenTypesCode.list() for allowed values.
    specimen_sampling_method: Optional[str] = None
        Sampling method used for sampling the sample from the specimen. See
        SpecimenSamplingProcedureCode.list() for allowed values.
    anatomical_sites: Optional[Sequence[Tuple[Code, Sequence[Code]]]] = None
        List of original anatomical sites, each defined by a code and optional
        list of modifier codes.
    location: Optional[Union[str, Tuple[float, float, float]]] = None
        Location of sample in slide, defined either by descriptive string
        or xyz-coordinates.
    uid_generator: Callable[..., Uid] = generate_uid
        Function that can generate uids.

    Returns
    ----------
    Dataset
        Dataset containing a sample description.
    """
    sample_preparation_steps: List[SpecimenPreparationStep] = []
    sample_preparation_step = create_sample_preparation_step(
        sample_id,
        stainings,
        embedding_medium,
        fixative
    )
    sample_preparation_steps.append(sample_preparation_step)
    if specimen_id is not None:
        sample_sampling_step = create_sample_sampling_step(
            sample_id,
            specimen_id,
            specimen_sampling_method,
            specimen_type
        )
        sample_preparation_steps.append(sample_sampling_step)

    specimen = SpecimenDescription(
        specimen_id=sample_id,
        specimen_uid=uid_generator(),
        specimen_preparation_steps=sample_preparation_steps,
        specimen_location=location
    )
    if anatomical_sites is not None:
        specimen = add_anatomical_sites_to_specimen(
            specimen,
            anatomical_sites
        )

    return specimen


def create_sample_preparation_step(
    specimen_id: str,
    stainings: Sequence[str],
    embedding_medium: Optional[str],
    fixative: Optional[str]
) -> SpecimenPreparationStep:
    """Create SpecimenPreparationStep for a preparation step.

    Parameters
    ----------
    specimen_id: str
        ID of specimen that has been prepared.
    stainings: Sequence[str]
        Sequence of stainings used.
    embedding_medium: Optional[str] = None
        Embedding medium used.
    fixative: Optional[str] = None
        Fixative used.

    Returns
    ----------
    SpecimenPreparationStep
        SpecimenPreparationStep for a preparation step.
    """
    if embedding_medium is not None:
        embedding_medium_code = (
            SpecimenEmbeddingMediaCode(embedding_medium).code
        )
    else:
        embedding_medium_code = None
    if fixative is not None:
        fixative_code = SpecimenFixativesCode(fixative).code
    else:
        fixative_code = None

    processing_type = SpecimenPreparationProcedureCode('Staining').code
    processing_procedure = SpecimenStaining([
        SpecimenStainsCode(staining).code for staining in stainings
    ])
    sample_preparation_step = SpecimenPreparationStep(
        specimen_id=specimen_id,
        processing_type=processing_type,
        processing_procedure=processing_procedure,
        embedding_medium=embedding_medium_code,
        fixative=fixative_code
    )
    return sample_preparation_step


def create_sample_sampling_step(
    sample_id: str,
    specimen_id: str,
    specimen_sampling_method: Optional[str],
    specimen_type: Optional[str]
) -> SpecimenPreparationStep:
    """Create SpecimenPreparationStep for a sampling step.

    Parameters
    ----------
    sample_id: str
        ID of sample that has been created.
    specimen_id: str
        ID of specimen that was sampled.
    specimen_sampling_method: Optional[str]
        Method used to sample specimen.
    specimen_type: Optional[str]
        Type of specimen that was sampled

    Returns
    ----------
    SpecimenPreparationStep
        SpecimenPreparationStep for a sampling step.
    """
    if specimen_sampling_method is None:
        raise ValueError(
            'Specimen sampling method required if '
            'specimen id is defined'
        )
    if specimen_type is None:
        raise ValueError(
            'Specimen type required if '
            'specimen id is defined'
        )
    specimen_type_code = AnatomicPathologySpecimenTypesCode(specimen_type)
    sampling_method_code = SpecimenSamplingProcedureCode(
        specimen_sampling_method
    )
    sample_sampling_step = SpecimenPreparationStep(
        specimen_id=sample_id,
        processing_type=SpecimenPreparationProcedureCode(
            'Sampling of tissue specimen'
        ).code,
        processing_procedure=SpecimenSampling(
                method=sampling_method_code.code,
                parent_specimen_id=specimen_id,
                parent_specimen_type=specimen_type_code.code
        )
    )
    return sample_sampling_step


def add_anatomical_sites_to_specimen(
    specimen: Dataset,
    anatomical_sites: Sequence[Tuple[Code, Sequence[Code]]]
) -> Dataset:
    """Add anatomical site and anatomical site modifier codes to specimen.

    Parameters
    ----------
    specimen: Dataset
        Dataset containing a specimen description
    anatomical_sites: Sequence[Tuple[Code, Sequence[Code]]]
        List of original anatomical sites, each defined by a code and
        list of modifier codes (can be empty).

    Returns
    ----------
    Dataset
        Dataset containing a specimen description.
    """
    anatomical_site_datasets: List[Dataset] = []
    for (anatomical_site, modifiers) in anatomical_sites:
        anatomical_site_dataset = (
            ConceptCode.from_code(anatomical_site).to_ds()
        )

        if modifiers != []:
            modifier_datasets: List[Dataset] = []
            for modifier in modifiers:
                modifier_dataset = Dataset()
                modifier_dataset.CodeValue = modifier.value
                modifier_dataset.CodingSchemeDesignator = (
                    modifier.scheme_designator
                )
                modifier_dataset.CodeMeaning = modifier.meaning
                modifier_datasets.append(modifier_dataset)

            (
                anatomical_site_dataset
                .PrimaryAnatomicStructureModifierSequence
            ) = DicomSequence(modifier_datasets)
        anatomical_site_datasets.append(anatomical_site_dataset)
    specimen.PrimaryAnatomicStructureSequence = DicomSequence(
        anatomical_site_datasets
    )
    return specimen


def create_specimen_module(
    slide_id: str,
    samples: Union[Dataset, Sequence[Dataset]]
) -> Dataset:
    """Create specimen module.

    Parameters
    ----------
    slide_id: str.
        Identifier for the slide.
    samples: Union[Dataset, Sequence[Dataset]]
        Single or list of sample descriptions (should be created with
        create_sample())

    Returns
    ----------
    Dataset
        Dataset containing specimen module.
    """
    dataset = Dataset()
    dataset.ContainerIdentifier = slide_id

    container_type_code_sequence = Dataset()
    container_type_code_sequence.CodeValue = '258661006'
    container_type_code_sequence.CodingSchemeDesignator = 'SCT'
    container_type_code_sequence.CodeMeaning = 'Slide'
    dataset.ContainerTypeCodeSequence = (
        DicomSequence([container_type_code_sequence])
    )

    container_component_sequence = Dataset()
    container_component_sequence.ContainerComponentMaterial = 'GLASS'
    container_component_type_code_sequence = Dataset()
    container_component_type_code_sequence.CodeValue = '433472003'
    container_component_type_code_sequence.CodingSchemeDesignator = 'SCT'
    container_component_type_code_sequence.CodeMeaning = (
        'Microscope slide coverslip'
    )
    container_component_sequence.ContainerComponentTypeCodeSequence = (
        DicomSequence([container_component_type_code_sequence])
    )
    dataset.ContainerComponentSequence = (
        DicomSequence([container_component_sequence])
    )
    if not isinstance(samples, Sequence):
        samples = [samples]
    specimen_description_sequence = (
        DicomSequence(samples)
    )
    dataset.SpecimenDescriptionSequence = specimen_description_sequence

    return dataset


def create_brightfield_optical_path_module(
    magnification: Optional[float] = None
) -> Dataset:
    """Create optical path module for brightfield illumination conditions.

    Parameters
    ----------
    magnification: Optional[float] = None
        Objective magnification.

    Returns
    ----------
    Dataset
        Dataset containing optical path module for brightfield illumination.
    """
    dataset = Dataset()
    optical_path_item = Dataset()
    optical_path_item.OpticalPathIdentifier = '0'
    illumination_type_code_sequence = Dataset()
    illumination_type_code_sequence.CodeValue = '111744'
    illumination_type_code_sequence.CodingSchemeDesignator = 'DCM'
    illumination_type_code_sequence.CodeMeaning = (
        'Brightfield illumination'
    )
    optical_path_item.IlluminationTypeCodeSequence = DicomSequence(
        [illumination_type_code_sequence]
    )
    illumination_color_code_sequence = Dataset()
    illumination_color_code_sequence.CodeValue = 'R-102C0'
    illumination_color_code_sequence.CodingSchemeDesignator = 'SRT'
    illumination_color_code_sequence.CodeMeaning = 'Full Spectrum'
    optical_path_item.IlluminationColorCodeSequence = DicomSequence(
        [illumination_color_code_sequence]
    )
    if magnification is not None:
        optical_path_item.ObjectiveLensPower = magnification

    dataset.OpticalPathSequence = DicomSequence([optical_path_item])

    return dataset


def create_default_modules(
    uid_generator: Callable[..., Uid] = generate_uid
) -> List[Dataset]:
    """Return default module dataset for testing.

    Parameters
    ----------
    uid_generator: Callable[..., Uid]
        Function that can generate Uids.

    Returns
    ----------
    List[Dataset]
        Default module datasets.
    """
    modules: List[Dataset] = []

    # Generic study module
    modules.append(create_study_module())

    # Generic patient module
    modules.append(create_patient_module())

    # Generic device module
    modules.append(create_device_module())

    # Generic specimen module
    modules.append(create_specimen_module(
        'Unkown',
        samples=[create_sample(
            sample_id='Unkown',
            stainings=['water soluble eosin stain', 'hematoxylin stain'],
            uid_generator=uid_generator
        )]
    ))

    # Generic optical path sequence
    modules.append(create_brightfield_optical_path_module())

    return modules
