# Copyright (C) 2019 Cancer Care Associates

# 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 os
import subprocess
import uuid

import numpy as np

import pydicom

from pymedphys._dicom.create import dicom_dataset_from_dict
from pymedphys._dicom.header import (
    RED_adjustment_map_from_structure_names,
    adjust_machine_name,
    adjust_RED_by_structure_name,
    adjust_rel_elec_density,
)
from pymedphys._dicom.utilities import remove_file

HERE = os.path.dirname(__file__)
ORIGINAL_DICOM_FILENAME = os.path.join(
    HERE, "scratch", "original-{}.dcm".format(str(uuid.uuid4()))
)
ADJUSTED_DICOM_FILENAME = os.path.join(
    HERE, "scratch", "adjusted-{}.dcm".format(str(uuid.uuid4()))
)


def compare_dicom_cli(command, original, expected):
    pydicom.write_file(ORIGINAL_DICOM_FILENAME, original)

    try:
        subprocess.check_call(command)
        cli_adjusted_ds = pydicom.read_file(ADJUSTED_DICOM_FILENAME, force=True)

        assert str(cli_adjusted_ds) == str(expected)
    finally:
        remove_file(ORIGINAL_DICOM_FILENAME)
        remove_file(ADJUSTED_DICOM_FILENAME)


def test_adjust_machine_name():
    new_name = "new_name"

    original_ds = dicom_dataset_from_dict(
        {
            "BeamSequence": [
                {"TreatmentMachineName": "hello"},
                {"TreatmentMachineName": "george"},
            ]
        }
    )

    expected_ds = dicom_dataset_from_dict(
        {
            "BeamSequence": [
                {"TreatmentMachineName": new_name},
                {"TreatmentMachineName": new_name},
            ]
        }
    )

    adjusted_ds = adjust_machine_name(original_ds, new_name)

    assert adjusted_ds != original_ds
    assert adjusted_ds == expected_ds

    command = "pymedphys dicom adjust-machine-name".split() + [
        ORIGINAL_DICOM_FILENAME,
        ADJUSTED_DICOM_FILENAME,
        new_name,
    ]

    compare_dicom_cli(command, original_ds, expected_ds)


def test_electron_density_append():
    adjustment_map = {
        "to_be_changed 1": 1.0,
        "to_be_changed 2": 0.5,
        "to_be_changed 3": 1.5,
    }

    excess_adjustment_map = {**adjustment_map, **{"this_structure_doesnt_exist": 1.0}}

    original_ds = dicom_dataset_from_dict(
        {
            "StructureSetROISequence": [
                {"ROINumber": 1, "ROIName": "to_be_changed 1"},
                {"ROINumber": 2, "ROIName": "dont_change_me"},
                {"ROINumber": 10, "ROIName": "to_be_changed 2"},
                {"ROINumber": 99, "ROIName": "to_be_changed 3"},
            ],
            "RTROIObservationsSequence": [
                {
                    "ReferencedROINumber": 1,
                    "ROIPhysicalPropertiesSequence": [
                        {
                            "ROIPhysicalProperty": "EFFECTIVE_Z",
                            "ROIPhysicalPropertyValue": 6,
                        }
                    ],
                },
                {"ReferencedROINumber": 2},
                {"ReferencedROINumber": 10},
                {
                    "ReferencedROINumber": 99,
                    "ROIPhysicalPropertiesSequence": [
                        {
                            "ROIPhysicalProperty": "REL_ELEC_DENSITY",
                            "ROIPhysicalPropertyValue": 0,
                        }
                    ],
                },
            ],
        }
    )

    expected_ds = dicom_dataset_from_dict(
        {
            "RTROIObservationsSequence": [
                {
                    "ReferencedROINumber": 1,
                    "ROIPhysicalPropertiesSequence": [
                        {
                            "ROIPhysicalProperty": "EFFECTIVE_Z",
                            "ROIPhysicalPropertyValue": 6,
                        },
                        {
                            "ROIPhysicalProperty": "REL_ELEC_DENSITY",
                            "ROIPhysicalPropertyValue": adjustment_map[
                                "to_be_changed 1"
                            ],
                        },
                    ],
                },
                {"ReferencedROINumber": 2},
                {
                    "ReferencedROINumber": 10,
                    "ROIPhysicalPropertiesSequence": [
                        {
                            "ROIPhysicalProperty": "REL_ELEC_DENSITY",
                            "ROIPhysicalPropertyValue": adjustment_map[
                                "to_be_changed 2"
                            ],
                        }
                    ],
                },
                {
                    "ReferencedROINumber": 99,
                    "ROIPhysicalPropertiesSequence": [
                        {
                            "ROIPhysicalProperty": "REL_ELEC_DENSITY",
                            "ROIPhysicalPropertyValue": adjustment_map[
                                "to_be_changed 3"
                            ],
                        }
                    ],
                },
            ]
        },
        template_ds=original_ds,
    )

    adjusted_ds = adjust_rel_elec_density(original_ds, adjustment_map)

    assert adjusted_ds != original_ds
    assert str(expected_ds) == str(adjusted_ds)

    adjusted_with_excess_ds = adjust_rel_elec_density(
        original_ds, excess_adjustment_map, ignore_missing_structure=True
    )

    assert adjusted_with_excess_ds != original_ds
    assert str(expected_ds) == str(adjusted_with_excess_ds)

    excess_adjustment_map_as_list = [
        ["{}".format(key), item] for key, item in excess_adjustment_map.items()
    ]
    excess_adjustment_map_flat = np.concatenate(excess_adjustment_map_as_list).tolist()

    command = (
        "pymedphys dicom adjust-RED -i ".split()
        + [ORIGINAL_DICOM_FILENAME, ADJUSTED_DICOM_FILENAME]
        + excess_adjustment_map_flat
    )

    compare_dicom_cli(command, original_ds, expected_ds)


def test_structure_name_parse():
    structure_names = [
        "a RED=1",
        "b",
        "c",
        "d RED=2.2",
        "e red = 3",
        "f",
        "g Red: 4.7",
        "h  RED=0.5  ",
    ]
    expected_adjustment_map = {
        "a RED=1": 1,
        "d RED=2.2": 2.2,
        "e red = 3": 3,
        "g Red: 4.7": 4.7,
        "h  RED=0.5  ": 0.5,
    }

    adjustment_map = RED_adjustment_map_from_structure_names(structure_names)

    assert expected_adjustment_map == adjustment_map


def test_structure_name_based_RED_append():
    electron_density_to_use = 0.5

    original_ds = dicom_dataset_from_dict(
        {
            "StructureSetROISequence": [
                {
                    "ROINumber": 1,
                    "ROIName": "a_structure RED={}".format(electron_density_to_use),
                },
                {"ROINumber": 2, "ROIName": "dont_change_me"},
            ],
            "RTROIObservationsSequence": [
                {"ReferencedROINumber": 1},
                {"ReferencedROINumber": 2},
            ],
        }
    )

    expected_ds = dicom_dataset_from_dict(
        {
            "RTROIObservationsSequence": [
                {
                    "ReferencedROINumber": 1,
                    "ROIPhysicalPropertiesSequence": [
                        {
                            "ROIPhysicalProperty": "REL_ELEC_DENSITY",
                            "ROIPhysicalPropertyValue": electron_density_to_use,
                        }
                    ],
                },
                {"ReferencedROINumber": 2},
            ]
        },
        template_ds=original_ds,
    )

    adjusted_ds = adjust_RED_by_structure_name(original_ds)

    assert adjusted_ds != original_ds
    assert str(expected_ds) == str(adjusted_ds)

    command = "pymedphys dicom adjust-RED-by-structure-name".split() + [
        ORIGINAL_DICOM_FILENAME,
        ADJUSTED_DICOM_FILENAME,
    ]

    compare_dicom_cli(command, original_ds, expected_ds)

[((958, 983), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (973, 983), False, 'import os\n'), ((1256, 1309), 'pydicom.write_file', 'pydicom.write_file', (['ORIGINAL_DICOM_FILENAME', 'original'], {}), '(ORIGINAL_DICOM_FILENAME, original)\n', (1274, 1309), False, 'import pydicom\n'), ((1676, 1794), 'pymedphys._dicom.create.dicom_dataset_from_dict', 'dicom_dataset_from_dict', (["{'BeamSequence': [{'TreatmentMachineName': 'hello'}, {\n    'TreatmentMachineName': 'george'}]}"], {}), "({'BeamSequence': [{'TreatmentMachineName': 'hello'},\n    {'TreatmentMachineName': 'george'}]})\n", (1699, 1794), False, 'from pymedphys._dicom.create import dicom_dataset_from_dict\n'), ((1893, 2013), 'pymedphys._dicom.create.dicom_dataset_from_dict', 'dicom_dataset_from_dict', (["{'BeamSequence': [{'TreatmentMachineName': new_name}, {\n    'TreatmentMachineName': new_name}]}"], {}), "({'BeamSequence': [{'TreatmentMachineName': new_name\n    }, {'TreatmentMachineName': new_name}]})\n", (1916, 2013), False, 'from pymedphys._dicom.create import dicom_dataset_from_dict\n'), ((2111, 2153), 'pymedphys._dicom.header.adjust_machine_name', 'adjust_machine_name', (['original_ds', 'new_name'], {}), '(original_ds, new_name)\n', (2130, 2153), False, 'from pymedphys._dicom.header import RED_adjustment_map_from_structure_names, adjust_machine_name, adjust_RED_by_structure_name, adjust_rel_elec_density\n'), ((2714, 3361), 'pymedphys._dicom.create.dicom_dataset_from_dict', 'dicom_dataset_from_dict', (["{'StructureSetROISequence': [{'ROINumber': 1, 'ROIName': 'to_be_changed 1'},\n    {'ROINumber': 2, 'ROIName': 'dont_change_me'}, {'ROINumber': 10,\n    'ROIName': 'to_be_changed 2'}, {'ROINumber': 99, 'ROIName':\n    'to_be_changed 3'}], 'RTROIObservationsSequence': [{\n    'ReferencedROINumber': 1, 'ROIPhysicalPropertiesSequence': [{\n    'ROIPhysicalProperty': 'EFFECTIVE_Z', 'ROIPhysicalPropertyValue': 6}]},\n    {'ReferencedROINumber': 2}, {'ReferencedROINumber': 10}, {\n    'ReferencedROINumber': 99, 'ROIPhysicalPropertiesSequence': [{\n    'ROIPhysicalProperty': 'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue': \n    0}]}]}"], {}), "({'StructureSetROISequence': [{'ROINumber': 1,\n    'ROIName': 'to_be_changed 1'}, {'ROINumber': 2, 'ROIName':\n    'dont_change_me'}, {'ROINumber': 10, 'ROIName': 'to_be_changed 2'}, {\n    'ROINumber': 99, 'ROIName': 'to_be_changed 3'}],\n    'RTROIObservationsSequence': [{'ReferencedROINumber': 1,\n    'ROIPhysicalPropertiesSequence': [{'ROIPhysicalProperty': 'EFFECTIVE_Z',\n    'ROIPhysicalPropertyValue': 6}]}, {'ReferencedROINumber': 2}, {\n    'ReferencedROINumber': 10}, {'ReferencedROINumber': 99,\n    'ROIPhysicalPropertiesSequence': [{'ROIPhysicalProperty':\n    'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue': 0}]}]})\n", (2737, 3361), False, 'from pymedphys._dicom.create import dicom_dataset_from_dict\n'), ((3926, 4667), 'pymedphys._dicom.create.dicom_dataset_from_dict', 'dicom_dataset_from_dict', (["{'RTROIObservationsSequence': [{'ReferencedROINumber': 1,\n    'ROIPhysicalPropertiesSequence': [{'ROIPhysicalProperty': 'EFFECTIVE_Z',\n    'ROIPhysicalPropertyValue': 6}, {'ROIPhysicalProperty':\n    'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue': adjustment_map[\n    'to_be_changed 1']}]}, {'ReferencedROINumber': 2}, {\n    'ReferencedROINumber': 10, 'ROIPhysicalPropertiesSequence': [{\n    'ROIPhysicalProperty': 'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue':\n    adjustment_map['to_be_changed 2']}]}, {'ReferencedROINumber': 99,\n    'ROIPhysicalPropertiesSequence': [{'ROIPhysicalProperty':\n    'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue': adjustment_map[\n    'to_be_changed 3']}]}]}"], {'template_ds': 'original_ds'}), "({'RTROIObservationsSequence': [{\n    'ReferencedROINumber': 1, 'ROIPhysicalPropertiesSequence': [{\n    'ROIPhysicalProperty': 'EFFECTIVE_Z', 'ROIPhysicalPropertyValue': 6}, {\n    'ROIPhysicalProperty': 'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue':\n    adjustment_map['to_be_changed 1']}]}, {'ReferencedROINumber': 2}, {\n    'ReferencedROINumber': 10, 'ROIPhysicalPropertiesSequence': [{\n    'ROIPhysicalProperty': 'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue':\n    adjustment_map['to_be_changed 2']}]}, {'ReferencedROINumber': 99,\n    'ROIPhysicalPropertiesSequence': [{'ROIPhysicalProperty':\n    'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue': adjustment_map[\n    'to_be_changed 3']}]}]}, template_ds=original_ds)\n", (3949, 4667), False, 'from pymedphys._dicom.create import dicom_dataset_from_dict\n'), ((5623, 5675), 'pymedphys._dicom.header.adjust_rel_elec_density', 'adjust_rel_elec_density', (['original_ds', 'adjustment_map'], {}), '(original_ds, adjustment_map)\n', (5646, 5675), False, 'from pymedphys._dicom.header import RED_adjustment_map_from_structure_names, adjust_machine_name, adjust_RED_by_structure_name, adjust_rel_elec_density\n'), ((5794, 5888), 'pymedphys._dicom.header.adjust_rel_elec_density', 'adjust_rel_elec_density', (['original_ds', 'excess_adjustment_map'], {'ignore_missing_structure': '(True)'}), '(original_ds, excess_adjustment_map,\n    ignore_missing_structure=True)\n', (5817, 5888), False, 'from pymedphys._dicom.header import RED_adjustment_map_from_structure_names, adjust_machine_name, adjust_RED_by_structure_name, adjust_rel_elec_density\n'), ((6850, 6906), 'pymedphys._dicom.header.RED_adjustment_map_from_structure_names', 'RED_adjustment_map_from_structure_names', (['structure_names'], {}), '(structure_names)\n', (6889, 6906), False, 'from pymedphys._dicom.header import RED_adjustment_map_from_structure_names, adjust_machine_name, adjust_RED_by_structure_name, adjust_rel_elec_density\n'), ((7553, 7842), 'pymedphys._dicom.create.dicom_dataset_from_dict', 'dicom_dataset_from_dict', (["{'RTROIObservationsSequence': [{'ReferencedROINumber': 1,\n    'ROIPhysicalPropertiesSequence': [{'ROIPhysicalProperty':\n    'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue': electron_density_to_use\n    }]}, {'ReferencedROINumber': 2}]}"], {'template_ds': 'original_ds'}), "({'RTROIObservationsSequence': [{\n    'ReferencedROINumber': 1, 'ROIPhysicalPropertiesSequence': [{\n    'ROIPhysicalProperty': 'REL_ELEC_DENSITY', 'ROIPhysicalPropertyValue':\n    electron_density_to_use}]}, {'ReferencedROINumber': 2}]}, template_ds=\n    original_ds)\n", (7576, 7842), False, 'from pymedphys._dicom.create import dicom_dataset_from_dict\n'), ((8123, 8164), 'pymedphys._dicom.header.adjust_RED_by_structure_name', 'adjust_RED_by_structure_name', (['original_ds'], {}), '(original_ds)\n', (8151, 8164), False, 'from pymedphys._dicom.header import RED_adjustment_map_from_structure_names, adjust_machine_name, adjust_RED_by_structure_name, adjust_rel_elec_density\n'), ((1328, 1358), 'subprocess.check_call', 'subprocess.check_call', (['command'], {}), '(command)\n', (1349, 1358), False, 'import subprocess\n'), ((1385, 1439), 'pydicom.read_file', 'pydicom.read_file', (['ADJUSTED_DICOM_FILENAME'], {'force': '(True)'}), '(ADJUSTED_DICOM_FILENAME, force=True)\n', (1402, 1439), False, 'import pydicom\n'), ((1515, 1551), 'pymedphys._dicom.utilities.remove_file', 'remove_file', (['ORIGINAL_DICOM_FILENAME'], {}), '(ORIGINAL_DICOM_FILENAME)\n', (1526, 1551), False, 'from pymedphys._dicom.utilities import remove_file\n'), ((1560, 1596), 'pymedphys._dicom.utilities.remove_file', 'remove_file', (['ADJUSTED_DICOM_FILENAME'], {}), '(ADJUSTED_DICOM_FILENAME)\n', (1571, 1596), False, 'from pymedphys._dicom.utilities import remove_file\n'), ((1074, 1086), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (1084, 1086), False, 'import uuid\n'), ((1181, 1193), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (1191, 1193), False, 'import uuid\n'), ((6168, 6213), 'numpy.concatenate', 'np.concatenate', (['excess_adjustment_map_as_list'], {}), '(excess_adjustment_map_as_list)\n', (6182, 6213), True, 'import numpy as np\n')]