"""
Uniformity + Ghosting & Distortion

Calculates uniformity for a single-slice image of a uniform MRI phantom

This script implements the IPEM/MAGNET method of measuring fractional uniformity.
It also calculates integral uniformity using a 75% area FOV ROI and CoV for the same ROI.

This script also measures Ghosting within a single image of a uniform phantom.
This follows the guidance from ACR for testing their large phantom.

A simple measurement of distortion is also made by comparing the height and width of the circular phantom.

Created by Neil Heraghty
neil.heraghty@nhs.net

14/05/2018

"""

import os
import sys
import traceback

import hazenlib.exceptions as exc
import hazenlib.utils
import numpy as np
from hazenlib.HazenTask import HazenTask
from hazenlib.logger import logger
from hazenlib.types import Measurement, Result


class Uniformity(HazenTask):
    """Uniformity measurement class for DICOM images of the MagNet phantom

    Inherits from HazenTask class
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Set the single DICOM input to be the first in the list
        self.single_dcm = self.dcm_list[0]

    def run(self) -> Result:
        """Main function for performing uniformity measurement

        Returns:
            dict: results are returned in a standardised dictionary structure specifying the task name, input DICOM Series Description + SeriesNumber + InstanceNumber, task measurement key-value pairs, optionally path to the generated images for visualisation
        """
        results = self.init_result_dict()
        results.files = self.img_desc(self.single_dcm)

        try:
            horizontal_uniformity, vertical_uniformity = (
                self.get_fractional_uniformity(self.single_dcm)
            )
            results.add_measurement(
                Measurement(
                    name="Uniformity",
                    type="measured",
                    subtype="Horizontal",
                    unit="%",
                    value=round(horizontal_uniformity, 2),
                ),
            )
            results.add_measurement(
                Measurement(
                    name="Uniformity",
                    type="measured",
                    subtype="Vertical",
                    unit="%",
                    value=round(vertical_uniformity, 2),
                ),
            )
        except Exception as e:
            logger.exception(
                "Could not calculate the uniformity for %s because of : %s",
                self.img_desc(self.single_dcm),
                e,
            )
            traceback.print_exc(file=sys.stdout)

        # only return reports if requested
        if self.report:
            results.add_report_image(self.report_files)

        return results

    def mode(self, a, axis=0):
        """Finds the modal value of an array. From scipy.stats.mode

        Args:
            a (np.array): _description_
            axis (int, optional): Axis to calculate mode along. Defaults to 0.

        Returns:
            most_frequent: the modal value
            old_counts: the number of times this value was counted (check this)
        """
        scores = np.unique(np.ravel(a))  # get ALL unique values
        test_shape = list(a.shape)
        test_shape[axis] = 1
        old_most_frequent = np.zeros(test_shape)
        old_counts = np.zeros(test_shape)
        most_frequent = None

        for score in scores:
            template = a == score
            counts = np.expand_dims(np.sum(template, axis), axis)
            most_frequent = np.where(
                counts > old_counts, score, old_most_frequent
            )
            old_counts = np.maximum(counts, old_counts)
            old_most_frequent = most_frequent

        return most_frequent, old_counts

    def get_object_centre(self, dcm):
        """Locate centre coordinates

        Args:
            dcm (pydicom.FileDataset): DICOM image object

        Raises:
            Exception: _description_

        Returns:
            tuple: x and y coordinates
        """
        arr = dcm.pixel_array
        shape_detector = hazenlib.utils.ShapeDetector(arr=arr)
        orientation = hazenlib.utils.get_image_orientation(dcm)

        if orientation in ["Sagittal", "Coronal"]:
            # orientation is sagittal to patient
            try:
                (x, y), size, angle = shape_detector.get_shape("rectangle")
            except exc.ShapeError:
                raise

        elif orientation == "Transverse":
            # orientation is axial
            x, y, r = shape_detector.get_shape("circle")

        else:
            raise Exception(
                "Direction must be Transverse, Sagittal or Coronal."
            )

        return int(x), int(y)

    def get_fractional_uniformity(self, dcm):
        """Get fractional uniformity

        Args:
            dcm (pydicom.FileDataset): DICOM image object

        Returns:
            tuple: values of horizontal and vertical fractional uniformity
        """
        arr = dcm.pixel_array
        x, y = self.get_object_centre(dcm)

        central_roi = arr[(y - 5) : (y + 5), (x - 5) : (x + 5)].flatten()
        # Create central 10x10 ROI and measure modal value

        central_roi_mode, mode_popularity = self.mode(central_roi)

        # Create 160-pixel profiles (horizontal and vertical, centred at x,y)
        horizontal_roi = arr[(y - 5) : (y + 5), (x - 80) : (x + 80)]
        horizontal_profile = np.mean(horizontal_roi, axis=0)
        vertical_roi = arr[(y - 80) : (y + 80), (x - 5) : (x + 5)]
        vertical_profile = np.mean(vertical_roi, axis=1)

        # Count how many elements are within 0.9-1.1 times the modal value
        horizontal_count = np.where(
            np.logical_and(
                (horizontal_profile > (0.9 * central_roi_mode)),
                (horizontal_profile < (1.1 * central_roi_mode)),
            )
        )
        horizontal_count = len(horizontal_count[0])
        vertical_count = np.where(
            np.logical_and(
                (vertical_profile > (0.9 * central_roi_mode)),
                (vertical_profile < (1.1 * central_roi_mode)),
            )
        )
        vertical_count = len(vertical_count[0])

        # Calculate fractional uniformity
        fractional_uniformity_horizontal = horizontal_count / 160
        fractional_uniformity_vertical = vertical_count / 160

        if self.report:
            import matplotlib.pyplot as plt
            from matplotlib.collections import PatchCollection
            from matplotlib.patches import Rectangle

            fig, ax = plt.subplots()
            rects = [
                Rectangle(
                    (x - 5, y - 5),
                    10,
                    10,
                    facecolor="None",
                    edgecolor="red",
                    linewidth=3,
                ),
                Rectangle(
                    (x - 80, y - 5),
                    160,
                    10,
                    facecolor="None",
                    edgecolor="green",
                ),
                Rectangle(
                    (x - 5, y - 80),
                    10,
                    160,
                    facecolor="None",
                    edgecolor="yellow",
                ),
            ]
            pc = PatchCollection(rects, match_original=True)
            ax.imshow(arr, cmap="gray")
            ax.add_collection(pc)
            ax.scatter(x, y, 5)

            img_path = os.path.realpath(
                os.path.join(self.report_path, f"{self.img_desc(dcm)}.png")
            )
            fig.savefig(img_path)
            self.report_files.append(img_path)

        return fractional_uniformity_horizontal, fractional_uniformity_vertical
