"""
SNR(Im)

Calculates the SNR for a single-slice image of a uniform MRI phantom

This script utilises the smoothed subtraction method described in McCann 2013:
A quick and robust method for measurement of signal-to-noise ratio in MRI, Phys. Med. Biol. 58 (2013) 3775:3790


Created by Neil Heraghty

04/05/2018
"""

from __future__ import annotations

# Python imports
import contextlib
import os
from typing import Any

# Module imports
import cv2 as cv
import numpy as np
import pydicom
import skimage.filters
from pathlib import Path
from scipy import ndimage

# Local imports
import hazenlib.exceptions as exc
import hazenlib.utils
from hazenlib.HazenTask import HazenTask
from hazenlib.logger import logger
from hazenlib.types import Measurement, Result
from hazenlib.utils import dcmread


class SNR(HazenTask):
    """Signal-to-noise ratio measurement class for DICOM images of the MagNet phantom

    Inherits from HazenTask class
    """

    def __init__(
        self,
        measured_slice_width: float | None = None,
        coil: str | None = None,
        subtract: str | Path | None = None,
        **kwargs: Any,
    ) -> None:
        """Initialise the Hazen ACR SNR Object."""
        super().__init__(**kwargs)

        # Pulls out the subtract path (for compatability with acr_snr)
        if subtract is not None:
            dicom_files = list(Path(subtract).glob("*.dcm"))
            if len(dicom_files) > 1:
                logger.warning(
                    "Multiple DICOM files found taking the first: %s",
                    dicom_files[0],
                )

            if len(dicom_files) < 1:
                logger.warning(
                    "No DICOM files found in: %s - ignore subtraction.",
                    subtract,
                )
            else:
                self.dcm_list.append(dcmread(dicom_files[0]))

        # measured slice width is expected to be a floating point number
        self.measured_slice_width = measured_slice_width
        with contextlib.suppress(TypeError):
            self.measured_slice_width = float(measured_slice_width)

        # Determining kernel size based on coil choice. Values of 9 and 25 come from McCann 2013 paper.
        try:
            if coil.lower() in ["hc", "head"]:
                self.kernel_size = 9
            elif coil.lower() in ["bc", "body"]:
                self.kernel_size = 25

        except AttributeError:
            self.kernel_size = 9

    def run(self) -> Result:
        """Main function for performing signal-to-noise ratio measurement

        Notes:
            Five square ROIs are created, one at the image centre, and four peripheral
            ROIs with their centres displaced at 45, 135, 225 and 315 degrees from the
            centre.

        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(img) for img in self.dcm_list]

        # SUBTRACTION METHOD with a pair of input files
        if len(self.dcm_list) == 2:
            snr, normalised_snr = self.snr_by_subtraction(
                self.dcm_list[0], self.dcm_list[1], self.measured_slice_width
            )
            results.add_measurement(
                Measurement(
                    name="SNR",
                    subtype="subtraction",
                    type="measured",
                    value=round(snr, 2),
                    visibility="intermediate",
                ),
            )
            results.add_measurement(
                Measurement(
                    name="SNR",
                    subtype="subtraction",
                    type="normalised",
                    value=round(normalised_snr, 2),
                ),
            )

        # SINGLE METHOD (SMOOTHING) for every input file
        for dcm in self.dcm_list:
            snr, normalised_snr = self.snr_by_smoothing(
                dcm, self.measured_slice_width
            )
            results.add_measurement(
                Measurement(
                    name="SNR",
                    subtype="smoothing",
                    type="measured",
                    description=self.img_desc(dcm),
                    value=round(snr, 2),
                    visibility="intermediate",
                ),
            )
            results.add_measurement(
                Measurement(
                    name="SNR",
                    subtype="smoothing",
                    type="normalised",
                    description=self.img_desc(dcm),
                    value=round(normalised_snr, 2),
                ),
            )

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

        return results

    def two_inputs_match(
        self, dcm1: pydicom.Dataset, dcm2: pydicom.Dataset
    ) -> bool:
        """Check if two DICOMs are sufficiently similar, based on the following fields
            "StudyInstanceUID", "RepetitionTime", "EchoTime", "FlipAngle"

        Args:
            dcm1 (pydicom.Dataset): DICOM object to compare
            dcm2 (pydicom.Dataset): _description_

        Returns:
            bool: _description_
        """

        fields_to_match = [
            "StudyInstanceUID",
            "RepetitionTime",
            "EchoTime",
            "FlipAngle",
        ]

        for field in fields_to_match:
            if dcm1.get(field) != dcm2.get(field):
                return False
        return True

    def get_normalised_snr_factor(
        self, dcm: pydicom.Dataset, measured_slice_width=None
    ) -> float:
        """Calculates SNR normalisation factor.

        Notes:
            Method matches MATLAB script.
            Utilises user provided slice_width if provided. Else finds from dcm.
            Finds dx, dy and bandwidth from dcm.
            Seeks to find TR, image columns and rows from dcm. Else uses default values.

        Args:
            dcm (pydicom.Dataset): DICOM image object
            measured_slice_width (float or None): slice width from user input

        Returns:
            float: normalised snr factor
        """

        dx, dy = hazenlib.utils.get_pixel_size(dcm)
        bandwidth = hazenlib.utils.get_bandwidth(dcm)
        TR = hazenlib.utils.get_TR(dcm)
        rows = hazenlib.utils.get_rows(dcm)
        columns = hazenlib.utils.get_columns(dcm)

        if measured_slice_width:
            slice_thickness = measured_slice_width
        else:
            slice_thickness = hazenlib.utils.get_slice_thickness(dcm)

        averages = hazenlib.utils.get_average(dcm)
        bandwidth_factor = np.sqrt((bandwidth * columns / 2) / 1000) / np.sqrt(
            30
        )
        voxel_factor = 1 / (0.001 * dx * dy * slice_thickness)

        normalised_snr_factor = (
            bandwidth_factor
            * voxel_factor
            * (1 / (np.sqrt(averages * rows * (TR / 1000))))
        )

        return normalised_snr_factor

    def filtered_image(self, dcm: pydicom.Dataset) -> np.array:
        """Performs a 2D convolution (for filtering images)
        using uniform_filter SciPy function and a kernel size based on user input coil

        Args:
            dcm (pydicom.Dataset): DICOM image to be filtered

        Returns:
            np.array: filtered image pixel values
        """
        a = dcm.pixel_array.astype("int")

        # filter size = 9, following MATLAB code and McCann 2013 paper for head coil, although note McCann 2013 recommends 25x25 for body coil.

        # 9 for head coil, 25 for body coil
        filtered_array = ndimage.uniform_filter(
            a, self.kernel_size, mode="constant"
        )
        return filtered_array

    def get_noise_image(self, dcm: pydicom.Dataset) -> np.array:
        """Separates the image noise
        by smoothing the image and subtracting the smoothed image from the original.

        Args:
            dcm (pydicom.Dataset): DICOM image to get noise from

        Returns:
            np.array: pixel array representing the image noise
        """
        a = dcm.pixel_array.astype("int")

        # Convolve image with boxcar/uniform kernel
        imsmoothed = self.filtered_image(dcm)

        # Subtract smoothed array from original
        imnoise = a - imsmoothed

        return imnoise

    def threshold_image(self, dcm: pydicom.Dataset):
        """Determine threshold and mask based on image

        Args:
            dcm (pydicom.Dataset): DICOM image to get noise from

        Returns:
            tuple of np.array: (imthresholded, mask)
                pixel array representing the image above threshold
                and a corresponding mask
        """
        a = dcm.pixel_array.astype("int")

        # threshold_li: Pixels > this value are assumed foreground
        threshold_value = skimage.filters.threshold_li(a)
        # print('threshold_value =', threshold_value)
        mask = a > threshold_value
        imthresholded = np.zeros_like(a)
        imthresholded[mask] = a[mask]

        # # For debugging: Threshold figures:
        # from matplotlib import pyplot as plt
        # plt.figure()
        # fig, ax = plt.subplots(2, 2)
        # ax[0, 0].imshow(a)
        # ax[0, 1].imshow(mask)
        # ax[1, 0].imshow(imthresholded)
        # ax[1, 1].imshow(a-imthresholded)
        # fig.savefig("../THRESHOLD.png")

        return imthresholded, mask

    def get_binary_mask_centre(self, binary_mask) -> (int, int):
        """Determine coordinates of binary polygonal shape's centre

        Args:
            binary_mask: mask of a shape

        Returns:
            tuple of int corresponding to centroid_coords: (col:int, row:int)
        """

        from skimage import util
        from skimage.measure import label, regionprops

        img = util.img_as_ubyte(binary_mask) > 0
        label_img = label(img, connectivity=img.ndim)
        props = regionprops(label_img)
        col = int(props[0].centroid[0])
        row = int(props[0].centroid[1])
        # print('Centroid coords [x,y] =', col, row)

        return int(col), int(row)

    def get_roi_samples(
        self,
        ax,
        dcm: pydicom.Dataset or np.ndarray,
        centre_col: int,
        centre_row: int,
    ) -> list:
        """Determine region of interest from a pixel array.

        Args:
            ax (matplotlib axes): diagram axis for visualisation with matplotlib
            dcm (pydicom.Dataset or np.ndarray): image pixel array
            centre_col (int): center coordinate column
            centre_row (int): center coordinate row

        Returns:
            list of np.ndarray: corresponding to pixel array subsets at predefined ROI

        """
        data = dcm if isinstance(dcm, np.ndarray) else dcm.pixel_array

        sample = [None] * 5
        # for array indexing: [row, column] format
        sample[0] = data[
            (centre_row - 10) : (centre_row + 10),
            (centre_col - 10) : (centre_col + 10),
        ]
        sample[1] = data[
            (centre_row - 50) : (centre_row - 30),
            (centre_col - 50) : (centre_col - 30),
        ]
        sample[2] = data[
            (centre_row + 30) : (centre_row + 50),
            (centre_col - 50) : (centre_col - 30),
        ]
        sample[3] = data[
            (centre_row - 50) : (centre_row - 30),
            (centre_col + 30) : (centre_col + 50),
        ]
        sample[4] = data[
            (centre_row + 30) : (centre_row + 50),
            (centre_col + 30) : (centre_col + 50),
        ]

        if ax:
            from matplotlib.collections import PatchCollection
            from matplotlib.patches import Rectangle

            # for patches: [column/x, row/y] format

            rects = [
                Rectangle((centre_col - 10, centre_row - 10), 20, 20),
                Rectangle((centre_col - 50, centre_row - 50), 20, 20),
                Rectangle((centre_col + 30, centre_row - 50), 20, 20),
                Rectangle((centre_col - 50, centre_row + 30), 20, 20),
                Rectangle((centre_col + 30, centre_row + 30), 20, 20),
            ]
            pc = PatchCollection(
                rects, edgecolors="red", facecolors="None", label="ROIs"
            )
            ax.add_collection(pc)

        return sample

    def get_object_centre(self, dcm) -> (int, int):
        """Find the phantom object within the image and return its centre col and row value.
        Note first element in output = col, second = row.

        Args:
            dcm (pydicom.Dataset): DICOM image to get noise from

        Returns:
            tuple of int corresponding to centroid_coords: (col:int, row:int)
        """

        # Shape Detection
        try:
            logger.debug("Performing phantom shape detection.")
            shape_detector = hazenlib.utils.ShapeDetector(arr=dcm.pixel_array)
            orientation = hazenlib.utils.get_image_orientation(dcm)

            if orientation in ["Sagittal", "Coronal"]:
                logger.debug("Orientation = sagittal or coronal.")
                # orientation is sagittal to patient
                try:
                    (col, row), size, angle = shape_detector.get_shape(
                        "rectangle"
                    )
                except exc.ShapeError as e:
                    # shape_detector.find_contours()
                    # shape_detector.detect()
                    # contour = shape_detector.shapes['rectangle'][1]
                    # angle, centre, size = cv.minAreaRect(contour)
                    # print((angle, centre, size))
                    # im = cv.drawContours(dcm.pixel_array.copy(), [shape_detector.contours[0]], -1, (0, 255, 255), 10)
                    # plt.imshow(im)
                    # plt.savefig("rectangles.png")
                    # print(shape_detector.shapes.keys())
                    raise e
            elif orientation == "Transverse":
                logger.debug("Orientation = transverse.")
                try:
                    col, row, r = shape_detector.get_shape("circle")
                except exc.MultipleShapesError:
                    logger.info(
                        "Warning! Found multiple circles in image, will assume largest circle is phantom."
                    )
                    col, row, r = self.get_largest_circle(
                        shape_detector.shapes["circle"]
                    )
            else:
                raise exc.ShapeError("Unable to identify phantom shape.")

        # Threshold Detection
        except exc.ShapeError:
            logger.info(
                "Shape detection failed. Performing object centre measurement by thresholding."
            )
            _, mask = self.threshold_image(dcm)
            row, col = self.get_binary_mask_centre(mask)

        return int(col), int(row)

    def snr_by_smoothing(
        self, dcm: pydicom.Dataset, measured_slice_width=None
    ):
        """Calculate signal to noise ratio by smoothing

        Args:
            dcm (pydicom.Dataset): DICOM image object
            measured_slice_width (float or None): slice width from user input

        Returns:
            tuple of float: SNR and normalised SNR values

        """
        col, row = self.get_object_centre(dcm)
        noise_img = self.get_noise_image(dcm)

        signal = [
            np.mean(roi)
            for roi in self.get_roi_samples(
                ax=None, dcm=dcm, centre_col=col, centre_row=row
            )
        ]

        noise = [
            np.std(roi, ddof=1)
            for roi in self.get_roi_samples(
                ax=None, dcm=noise_img, centre_col=col, centre_row=row
            )
        ]
        # note no root_2 factor in noise for smoothed subtraction (one image) method, replicating Matlab approach and McCann 2013

        snr = np.mean(np.divide(signal, noise))

        normalised_snr = snr * self.get_normalised_snr_factor(
            dcm, measured_slice_width
        )

        if self.report:
            import matplotlib.pyplot as plt

            fig, axes = plt.subplots(1, 1)
            fig.set_size_inches(5, 5)
            fig.tight_layout(pad=1)

            axes.set_title("smoothed noise image")
            axes.imshow(noise_img, cmap="gray", label="smoothed noise image")
            axes.scatter(col, row, 10, marker="+", label="centre")
            self.get_roi_samples(axes, dcm, col, row)
            axes.legend()

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

        return snr, normalised_snr

    def get_largest_circle(self, circles):
        """Determine circle with largest radius from list of detected circles

        Args:
            circles (_type_): _description_

        Returns:
            tuple: of centre coordinates (col, row) and radius (float)
        """
        largest_col, largest_row, largest_r = 0, 0, 0
        for circle in circles:
            (col, row), r = cv.minEnclosingCircle(circle)
            if r > largest_r:
                largest_col, largest_row, largest_r = col, row, r

        return largest_col, largest_row, largest_r

    def snr_by_subtraction(
        self,
        dcm1: pydicom.Dataset,
        dcm2: pydicom.Dataset,
        measured_slice_width=None,
    ):
        """Calculate signal to noise ratio by smoothing

        Args:
            dcm1 (pydicom.Dataset): DICOM image object for signal
            dcm2 (pydicom.Dataset): DICOM image object for noise calculation
            measured_slice_width (float or None): slice width from user input

        Returns:
            tuple of float: SNR and normalised SNR values
        """
        col, row = self.get_object_centre(dcm1)

        difference = np.subtract(
            dcm1.pixel_array.astype("int"), dcm2.pixel_array.astype("int")
        )

        signal = [
            np.mean(roi)
            for roi in self.get_roi_samples(
                ax=None, dcm=dcm1, centre_col=col, centre_row=row
            )
        ]
        noise = np.divide(
            [
                np.std(roi, ddof=1)
                for roi in self.get_roi_samples(
                    ax=None, dcm=difference, centre_col=col, centre_row=row
                )
            ],
            np.sqrt(2),
        )
        snr = np.mean(np.divide(signal, noise))

        normalised_snr = snr * self.get_normalised_snr_factor(
            dcm1, measured_slice_width
        )

        if self.report:
            import matplotlib.pyplot as plt

            fig, axes = plt.subplots(1, 1)
            fig.set_size_inches(5, 5)
            fig.tight_layout(pad=1)

            axes.set_title("difference image")
            axes.imshow(difference, cmap="gray", label="difference image")
            axes.scatter(col, row, 10, marker="+", label="centre")
            self.get_roi_samples(axes, dcm1, col, row)
            axes.legend()

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

        return snr, normalised_snr
