import os

import warnings

import numpy as np

import pandas as pd

import scipy.ndimage as sni

from dosma.core.device import get_array_module

from dosma.core.io.format_io import ImageDataFormat

from dosma.core.med_volume import MedicalVolume

from dosma.core.quant_vals import QuantitativeValueType

from dosma.defaults import preferences

from dosma.tissues.tissue import Tissue, largest_cc

from dosma.utils import img_utils, io_utils

from dosma.utils.geometry_utils import cart2pol, circle_fit

import matplotlib.pyplot as plt

# milliseconds

BOUNDS = {

    QuantitativeValueType.T2: 80.0,

    QuantitativeValueType.T1_RHO: 100.0,

    QuantitativeValueType.T2_STAR: 80.0,

}

__all__ = ["FemoralCartilage"]

class FemoralCartilage(Tissue):

    """Handles analysis and visualization for femoral cartilage.

    This class extends functionality from `Tissue`.

    For visualization, the femoral cartilage is unrolled onto a 2D plane using angular binning [1].

    References:

        [1] Monu UD, Jordan CD, Samuelson BL, Hargreaves BA, Gold GE, McWalter EJ.

        Cluster analysis of quantitative MRI T2 and :math:`T1\\rho` relaxation times of

        cartilage identifies differences between healthy and ACL-injured individuals at 3T."

        Osteoarthritis and cartilage 2017;25(4):513-520.

    """

    ID = 1

    STR_ID = "fc"

    FULL_NAME = "femoral cartilage"

    # Expected quantitative values

    T1_EXPECTED = 1200  # milliseconds

    # Keys correspond to integer representing bit location for each region

    # bit string: 'T D S M L A C P' (stored as integer)

    # Coronal Keys

    _POSTERIOR_KEY = 2 ** 0

    _CENTRAL_KEY = 2 ** 1

    _ANTERIOR_KEY = 2 ** 2

    _CORONAL_KEYS = [_POSTERIOR_KEY, _CENTRAL_KEY, _ANTERIOR_KEY]

    # Sagittal Keys

    _MEDIAL_KEY = 2 ** 3

    _LATERAL_KEY = 2 ** 4

    _SAGITTAL_KEYS = [_MEDIAL_KEY, _LATERAL_KEY]

    # Axial Keys

    _DEEP_KEY = 2 ** 5

    _SUPERFICIAL_KEY = 2 ** 6

    _TOTAL_AXIAL_KEY = 2 ** 7

    _AXIAL_KEYS = [_DEEP_KEY, _SUPERFICIAL_KEY, _TOTAL_AXIAL_KEY]

    # Do not change order of below.

    # Order reflects order of _CORONAL_KEYS, _SAGITTAL_KEYS, _AXIAL_KEYS

    _AXIAL_NAMES = ["deep", "superficial", "total"]

    _SAGITTAL_NAMES = ["medial", "lateral"]

    _CORONAL_NAMES = ["posterior", "central", "anterior"]

    ML_BOUNDARY = None

    ACP_BOUNDARY = None

    def __init__(self, weights_dir=None, medial_to_lateral=None):

        super().__init__(weights_dir=weights_dir)

        self.regions_mask = None

        self.theta_bins = None

        self.medial_to_lateral = medial_to_lateral

    def split_regions(

        self, base_map: np.ndarray, thickness_divisor=0.5, num_bins=72, theta=(-270, 90)

    ):

        """Split volume into anatomical regions.

        Pixels corresponding to femoral cartilage are divided across 3 planes:

            - Coronal: Posterior, Central, or Anterior

            - Sagittal: Medial, Lateral

            - Axial: Deep, Superficial

        For example, a pixel could correspond to the Posterior Lateral Deep region of

            femoral cartilage.

        Args:

            base_map (np.ndarray): 3D numpy array typically corresponding to volume to split.

        Returns:

            np.ndarray: 4D numpy array (region, height, width, depth).

                Saved in variable ``self.regions``.

        """

        dtheta = 360 / num_bins

        theta_min, theta_max = tuple(theta)

        mask = self.__mask__.volume

        mask = mask * np.nan_to_num(base_map)

        height, width, num_slices = mask.shape

        # STEP 1: PROJECTING AND CYLINDRICAL FIT

        segmented_t2maps_projected = np.max(mask, 2)  # Project segmented T2maps on sagittal axis

        non_zero_element = np.nonzero(segmented_t2maps_projected)

        xc_fit, yc_fit, R_fit = circle_fit(

            non_zero_element[1], non_zero_element[0]

        )  # fit a circle to projected cartilage tissue

        # STEP 2: SLICE BY SLICE BINNING

        yv, xv = np.meshgrid(range(height), range(width), indexing="ij")

        rho, theta = cart2pol(xv - xc_fit, yc_fit - yv)

        theta = (theta >= 90) * (theta - 360) + (theta < 90) * theta  # range: [-270, 90)

        assert (np.min(theta) >= theta_min) and (

            np.max(theta) < theta_max

        ), "Expected Theta range is [{:d}, {:d}) degrees. Received min: {:d} max: {:d})".format(

            theta_min, theta_max, np.min(theta), np.max(theta)

        )

        theta_bins = np.floor((theta - theta_min) / dtheta)

        # STEP 3: COMPUTE THRESHOLD RADII

        # TODO: This step takes a long time

        rhos_threshold_volume = np.zeros(mask.shape)

        for curr_slice in range(num_slices):

            mask_slice = mask[..., curr_slice]

            for curr_bin in range(num_bins):

                rhos_valid = rho[np.logical_and(mask_slice > 0, theta_bins == curr_bin)]

                if len(rhos_valid) == 0:

                    continue

                rho_min = np.min(rhos_valid)

                rho_max = np.max(rhos_valid)

                rho_threshold = thickness_divisor * (rho_max - rho_min) + rho_min

                rhos_threshold_volume[theta_bins == curr_bin, curr_slice] = rho_threshold

        regions_volume = np.asarray(np.zeros(mask.shape), dtype=np.uint16)

        # anterior/central/posterior division

        # Central region occupies middle 30 degrees, anterior on left, posterior on right

        anterior_region = self._ANTERIOR_KEY * (theta < -105)

        central_region = self._CENTRAL_KEY * np.logical_and((theta >= -105), (theta < -75))

        posterior_region = self._POSTERIOR_KEY * (theta >= -75)

        acp_map = anterior_region + central_region + posterior_region

        acp_volume = np.asarray(np.stack([acp_map] * num_slices, axis=-1), dtype=np.uint16)

        regions_volume += acp_volume

        # medial/lateral division

        # take into account scanning direction

        center_of_mass = sni.measurements.center_of_mass(mask)

        com_slicewise = center_of_mass[-1]

        ml_volume = np.asarray(np.zeros(mask.shape), dtype=np.uint16)

        if self.medial_to_lateral:

            ml_volume[..., : int(np.ceil(com_slicewise))] = self._MEDIAL_KEY

            ml_volume[..., int(np.ceil(com_slicewise)) :] = self._LATERAL_KEY

        else:

            ml_volume[..., : int(np.ceil(com_slicewise))] = self._LATERAL_KEY

            ml_volume[..., int(np.ceil(com_slicewise)) :] = self._MEDIAL_KEY

        regions_volume += ml_volume

        # deep/superficial division

        rho_volume = np.stack([rho] * num_slices, axis=-1)

        deep_volume = (rho_volume <= rhos_threshold_volume) * self._DEEP_KEY

        superficial_volume = (rho_volume >= rhos_threshold_volume) * self._SUPERFICIAL_KEY

        ds_volume = np.asarray(

            deep_volume + superficial_volume + self._TOTAL_AXIAL_KEY, dtype=np.uint16

        )

        regions_volume += ds_volume

        ml_boundary = int(np.ceil(com_slicewise))

        acp_boundary = [

            int(np.floor((-105 - theta_min) / dtheta)),

            int(np.floor((-75 - theta_min) / dtheta)),

        ]

        return regions_volume, theta_bins, ml_boundary, acp_boundary

    def unroll(self, qv_map: np.ndarray, regions_mask: np.ndarray, theta_bins):

        """Unroll femoral cartilage 3D quantitative value (qv) maps to 2D for visualization.

        The function multiplies a 3D segmentation mask to a 3D qv map to produce a 3D femoral

        cartilage qv (fc_qv) map. It then fits a circle to the collapsed sagittal projection

        of the fc_qv map. Each slice is binned into bins of 5 degree sizes

        The unrolled map is then divided into deep and superficial cartilage.

        Args:

            qv_map (np.ndarray): 3D array (slices last) of sagittal knee describing

                quantitative parameter values regions_mask (np.ndarray): regions_mask

        Returns:

            tuple: (row, column) format

                1. 2D Total unrolled cartilage (slices, degrees) - average of superficial

                    and deep layers

                2. Superficial unrolled cartilage (slices, degrees) - superficial layer

                3. Deep unrolled cartilage (slices, degrees) - deep layer

        """

        num_bins = len(np.unique(theta_bins))

        mask = self.__mask__.volume

        if qv_map.shape != mask.shape:

            raise ValueError("t2_map and mask must have same shape")

        if len(qv_map.shape) != 3:

            raise ValueError("t2_map and mask must be 3D")

        # assert self.regions_mask is not None, (

        #     "region_mask not initialized. Should be initialized when mask is set"

        # )

        num_slices = qv_map.shape[-1]

        qv_map = np.nan_to_num(qv_map)

        qv_map = np.multiply(mask, qv_map)  # apply binary mask

        qv_map[

            qv_map <= 0

        ] = np.nan  # wherever qv_map is 0, either no cartilage or qv=0 ms, which is impractical

        # theta_bins = self.theta_bins  # binning with theta

        # regions_mask = self.regions_mask

        Unrolled_Cartilage = np.zeros([num_bins, num_slices])

        Sup_layer = np.zeros([num_bins, num_slices])

        Deep_layer = np.zeros([num_bins, num_slices])

        for slice_ind in range(num_slices):

            qv_slice = qv_map[..., slice_ind]

            curr_slice = regions_mask[..., slice_ind]

            # if slice is all NaNs, then don't analyze

            if np.sum(np.isnan(qv_slice)) == qv_slice.shape[0] * qv_slice.shape[1]:

                continue

            for curr_bin in range(num_bins):

                qv_bin = qv_slice[theta_bins == curr_bin]

                if np.sum(np.isnan(qv_bin)) == len(qv_bin):

                    continue

                Unrolled_Cartilage[curr_bin, slice_ind] = np.nanmean(qv_bin)

                qv_superficial = qv_slice[

                    np.logical_and(

                        theta_bins == curr_bin,

                        self.__binarize_region_mask__(curr_slice, self._SUPERFICIAL_KEY),

                    )

                ]

                qv_deep = qv_slice[

                    np.logical_and(

                        theta_bins == curr_bin,

                        self.__binarize_region_mask__(curr_slice, self._DEEP_KEY),

                    )

                ]

                qv_superficial = np.nan_to_num(qv_superficial)

                qv_deep = np.nan_to_num(qv_deep)

                qv_sup_mean = np.mean(qv_superficial[qv_superficial > 0])

                qv_deep_mean = np.mean(qv_deep[qv_deep > 0])

                Sup_layer[curr_bin, slice_ind] = qv_sup_mean

                Deep_layer[curr_bin, slice_ind] = qv_deep_mean

        Unrolled_Cartilage[Unrolled_Cartilage == 0] = np.nan

        Sup_layer[Sup_layer == 0] = np.nan

        Deep_layer[Deep_layer == 0] = np.nan

        return Unrolled_Cartilage, Sup_layer, Deep_layer

    def __calc_quant_vals__(self, quant_map: MedicalVolume, map_type):

        """Calculate quantitative values per region and 2D visualizations

        1. Save 2D figure (deep, superficial, total) information to use with matplotlib

            (title, data, xlabel, ylabel, filename)

        2. Save 2D dataframes in format

                [['DMA', 'DMC', 'DMP'], ['DLA', 'DLC', 'DLP'],

                 ['SMA', 'SMC', 'SMP'], ['SLA', 'SLC', 'SLP'],

                 ['TMA', 'TMC', 'TMP'], ['TLA', 'TLC', 'TLP']]

                 D=deep, S=superficial, T=total,

                 M=medial, L=lateral,

                 A=anterior, C=central, P=posterior

        Args:

            quant_map (MedicalVolume): 3D volumes of quantitative values.

                Volume should have ``np.nan`` values for all pixels unable to be calculated.

            map_type (QuantitativeValueType): Type of quantitative value to analyze.

        """

        super().__calc_quant_vals__(quant_map, map_type)

        # assert self.regions_mask is not None, (

        #     "region_mask not initialized. Should be initialized when mask is set"

        # )

        # We have to call this every time we load a new quantitative map

        # mask = segmentation_mask * clipped_quant_map

        regions_mask, theta_bins, ml_boundary, acp_boundary = self.split_regions(quant_map.volume)

        if self.ML_BOUNDARY is None:

            self.ML_BOUNDARY = ml_boundary

        if self.ACP_BOUNDARY is None:

            self.ACP_BOUNDARY = acp_boundary

        total, superficial, deep = self.unroll(quant_map.volume, regions_mask, theta_bins)

        assert total.shape == deep.shape

        assert deep.shape == superficial.shape

        # regions_mask = self.regions_mask

        mask = self.__mask__.volume

        subject_pid = self.pid

        pd_header = ["Subject", "Location", "Side", "Region", "Mean", "Std", "Median", "# Voxels"]

        pd_list = []

        # Replace strings with values - eg. DMA = 'deep, medial, anterior'

        # tissue_values = [['DMA', 'DMC', 'DMP'], ['DLA', 'DLC', 'DLP'],

        #                  ['SMA', 'SMC', 'SMP'], ['SLA', 'SLC', 'SLP'],

        #                  ['TMA', 'TMC', 'TMP'], ['TLA', 'TLC', 'TLP']]

        # tissue_values = []

        for axial_ind in range(len(self._AXIAL_KEYS)):

            axial = self._AXIAL_KEYS[axial_ind]

            for sagittal_ind in range(len(self._SAGITTAL_KEYS)):

                sagittal = self._SAGITTAL_KEYS[sagittal_ind]

                for coronal_ind in range(len(self._CORONAL_KEYS)):

                    coronal = self._CORONAL_KEYS[coronal_ind]

                    curr_region_mask = self.__binarize_region_mask__(

                        regions_mask, (axial | coronal | sagittal)

                    )

                    curr_region_mask = curr_region_mask * mask * quant_map.volume

                    # discard all values that are <= 0

                    qv_region_vals = curr_region_mask[curr_region_mask > 0]

                    num_voxels = len(qv_region_vals)

                    c_mean = np.nanmean(qv_region_vals)

                    c_std = np.nanstd(qv_region_vals)

                    c_median = np.nanmedian(qv_region_vals)

                    row_info = [

                        subject_pid,

                        self._AXIAL_NAMES[axial_ind],

                        self._SAGITTAL_NAMES[sagittal_ind],

                        self._CORONAL_NAMES[coronal_ind],

                        c_mean,

                        c_std,

                        c_median,

                        num_voxels,

                    ]

                    pd_list.append(row_info)

        df = pd.DataFrame(pd_list, columns=pd_header)

        qv_name = map_type.name

        maps = [

            {

                "title": "{} deep".format(qv_name),

                "data": deep,

                "xlabel": "Slice",

                "ylabel": "Angle (binned)",

                "filename": "{}_deep".format(qv_name),

                "raw_data_filename": "{}_deep.data".format(qv_name),

            },

            {

                "title": "{} superficial".format(qv_name),

                "data": superficial,

                "xlabel": "Slice",

                "ylabel": "Angle (binned)",

                "filename": "{}_superficial".format(qv_name),

                "raw_data_filename": "{}_superficial.data".format(qv_name),

            },

            {

                "title": "{} total".format(qv_name),

                "data": total,

                "xlabel": "Slice",

                "ylabel": "Angle (binned)",

                "filename": "{}_total".format(qv_name),

                "raw_data_filename": "{}_total.data".format(qv_name),

            },

        ]

        self.__store_quant_vals__(maps, df, map_type)

    def set_mask(

        self, mask: MedicalVolume, use_largest_cc: bool = True, split_regions: bool = True

    ):

        """Set mask for tissue.

        Mask is cleaned by selecting the largest connected component from the mask.

            Femoral cartilage is expected to be single connected tissue.

        Args:

            mask (MedicalVolume): Binary mask of segmented tissue.

        """

        xp = get_array_module(mask.A)

        if use_largest_cc:

            msk = xp.asarray(largest_cc(mask.A), dtype=xp.uint8)

        else:

            msk = xp.asarray(mask.A, dtype=xp.uint8)

        mask_copy = mask._partial_clone(volume=msk)

        super().set_mask(mask_copy)

        if split_regions:

            (

                self.regions_mask,

                self.theta_bins,

                self.ML_BOUNDARY,

                self.ACP_BOUNDARY,

            ) = self.split_regions(  # noqa: E501

                self.__mask__.volume

            )

    def __save_quant_data__(self, dirpath: str):

        """Save quantitative data and 2D visualizations of femoral cartilage.

        Check which quantitative values (T2, T1rho, etc) are defined for femoral cartilage

        and analyze these:

            1. Save 2D total, superficial, and deep visualization maps.

            2. Save {'medial', 'lateral'}, {'anterior', 'central', 'posterior'},

            q{'deep', 'superficial'} data to excel file

        Args:

            dirpath (str): Directory path to tissue data.

        """

        q_names = []

        dfs = []

        for quant_val in QuantitativeValueType:

            if quant_val.name not in self.quant_vals.keys():

                continue

            q_names.append(quant_val.name)

            q_val = self.quant_vals[quant_val.name]

            dfs.append(q_val[1])

            q_name_dirpath = io_utils.mkdirs(os.path.join(dirpath, quant_val.name.lower()))

            for q_map_data in q_val[0]:

                filepath = os.path.join(q_name_dirpath, q_map_data["filename"])

                xlabel = "Slice"

                ylabel = "Angle (binned)"

                title = q_map_data["title"]

                data_map = q_map_data["data"]

                plt.clf()

                upper_bound = BOUNDS[quant_val]

                if preferences.visualization_use_vmax:

                    # Hard bounds - clipping

                    plt.imshow(data_map, cmap="jet", vmin=0.0, vmax=BOUNDS[quant_val])

                else:

                    # Try to use a soft bounds

                    if np.sum(data_map <= upper_bound) == 0:

                        plt.imshow(data_map, cmap="jet", vmin=0.0, vmax=BOUNDS[quant_val])

                    else:

                        warnings.warn(

                            "%s: Pixel value exceeded upper bound (%0.1f). Using normalized scale."

                            % (quant_val.name, upper_bound)

                        )

                        plt.imshow(data_map, cmap="jet")

                plt.xlabel(xlabel)

                plt.ylabel(ylabel)

                plt.title(title)

                clb = plt.colorbar()

                clb.ax.set_title("(ms)")

                plt.savefig(filepath)

                # Save data

                raw_data_filepath = os.path.join(

                    q_name_dirpath, "raw_data", q_map_data["raw_data_filename"]

                )

                io_utils.save_pik(raw_data_filepath, data_map)

        if len(dfs) > 0:

            io_utils.save_tables(os.path.join(dirpath, "data.xlsx"), dfs, q_names)

    def save_data(self, save_dirpath, data_format: ImageDataFormat = preferences.image_data_format):

        super().save_data(save_dirpath, data_format=data_format)

        save_dirpath = self.__save_dirpath__(save_dirpath)

        if self.regions_mask is None:

            return

        sagital_region_mask, coronal_region_mask = self.__split_mask__()

        # Save region map - add by 1 because no key can be 0

        coronal_region_mask = (coronal_region_mask + 1) * 10

        sagital_region_mask = sagital_region_mask + 1

        joined_mask = coronal_region_mask + sagital_region_mask

        labels = [

            "medial posterior",

            "medial central",

            "medial anterior",

            "lateral posterior",

            "lateral central",

            "lateral anterior",

        ]

        plt_dict = {

            "labels": labels,

            "xlabel": "Slice",

            "ylabel": "Angle (binned)",

            "title": "Unrolled Regions",

        }

        img_utils.write_regions(

            os.path.join(save_dirpath, "region_map"), joined_mask, plt_dict=plt_dict

        )

    def __binarize_region_mask__(self, region_mask, roi):

        return np.asarray(np.bitwise_and(region_mask, roi) == roi, dtype=np.bool)

    def __split_mask__(self):

        assert (

            self.ML_BOUNDARY is not None and self.ACP_BOUNDARY is not None

        ), "medial/lateral and anterior/central/posterior boundaries should be specified"

        # split into regions

        unrolled_total, _, _ = self.unroll(

            np.asarray(self.__mask__.volume, dtype=np.float32), self.regions_mask, self.theta_bins

        )

        acp_division_unrolled = np.zeros(unrolled_total.shape)

        ac_threshold = self.ACP_BOUNDARY[0]

        cp_threshold = self.ACP_BOUNDARY[1]

        acp_division_unrolled[:ac_threshold, :] = self._ANTERIOR_KEY

        acp_division_unrolled[ac_threshold:cp_threshold, :] = self._CENTRAL_KEY

        acp_division_unrolled[cp_threshold:, :] = self._POSTERIOR_KEY

        ml_division_unrolled = np.zeros(unrolled_total.shape)

        if self.medial_to_lateral:

            ml_division_unrolled[..., : self.ML_BOUNDARY] = self._MEDIAL_KEY

            ml_division_unrolled[..., self.ML_BOUNDARY :] = self._LATERAL_KEY

        else:

            ml_division_unrolled[..., : self.ML_BOUNDARY] = self._LATERAL_KEY

            ml_division_unrolled[..., self.ML_BOUNDARY :] = self._MEDIAL_KEY

        acp_division_unrolled[np.isnan(unrolled_total)] = np.nan

        ml_division_unrolled[np.isnan(unrolled_total)] = np.nan

        return acp_division_unrolled, ml_division_unrolled