"""Analysis for patellar cartilage.

Attributes:
    BOUNDS (dict): Upper bounds for quantitative values.
"""
import itertools
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.quant_vals import QuantitativeValueType
from dosma.defaults import preferences
from dosma.tissues.tissue import Tissue, largest_cc
from dosma.utils import io_utils

import matplotlib.pyplot as plt

# milliseconds
BOUNDS = {
    QuantitativeValueType.T2: 60.0,
    QuantitativeValueType.T1_RHO: 100.0,
    QuantitativeValueType.T2_STAR: 50.0,
}

__all__ = ["PatellarCartilage"]


class PatellarCartilage(Tissue):
    """Handles analysis and visualization for patellar cartilage."""

    ID = 3
    STR_ID = "pc"
    FULL_NAME = "patellar cartilage"

    # Expected quantitative values
    T1_EXPECTED = 1000  # milliseconds

    # Region Keys
    _ANTERIOR_KEY = 0
    _POSTERIOR_KEY = 1
    _CORONAL_KEYS = [_ANTERIOR_KEY, _POSTERIOR_KEY]

    _MEDIAL_KEY = 0
    _LATERAL_KEY = 1
    _SAGITTAL_KEYS = [_MEDIAL_KEY, _LATERAL_KEY]

    _REGION_DEEP_KEY = 0
    _REGION_SUPERFICIAL_KEY = 1
    _TOTAL_AXIAL_KEY = -1

    def __init__(self, weights_dir: str = None, medial_to_lateral: bool = None):
        super().__init__(weights_dir=weights_dir, medial_to_lateral=medial_to_lateral)

        self.regions_mask = None

    def unroll_coronal(self, quant_map: np.ndarray):
        """Unroll patellar cartilage in the coronal direction.

        Because patellar cartilage is flat, "unrolling" is projecting the patellar
        cartilage onto the coronal axis.

        Args:
            quant_map (np.ndarray):
        """
        mask = self.__mask__.volume

        assert (
            self.regions_mask is not None
        ), "region_mask not initialized. Should be initialized when mask is set"
        region_mask_deep_superficial = self.regions_mask[..., 0]

        superficial = (
            (region_mask_deep_superficial == self._REGION_SUPERFICIAL_KEY) * mask * quant_map
        )
        superficial[superficial == 0] = np.nan
        superficial = np.nanmean(superficial, axis=2)

        deep = (region_mask_deep_superficial == self._REGION_DEEP_KEY) * mask * quant_map
        deep[deep == 0] = np.nan
        deep = np.nanmean(deep, axis=2)

        total = mask * quant_map
        total[total == 0] = np.nan
        total = np.nanmean(total, axis=2)

        return total, superficial, deep

    def split_regions(self, base_map):
        """Split patellar cartilage into deep/superficial regions.

        For patellar cartilage, the superficial/deep transition occurs in
        the anterior/posterior (A/P) direction. The boundary is determined
        for each non-zero 1D column spanning independently by the local
        center-of-mass (COM). The medial/lateral (M/L) plane is computed
        using the global COM.

        Args:
            base_map (ndarray): Binary 3D mask with orientation (SI, AP, ML/LM).
                If `self.medial_to_lateral`, last dimension should be ML.
        """
        if np.sum(base_map) == 0:
            warnings.warn("No mask for `%s` was found." % self.FULL_NAME)

        # Superficial/Deep (A/P)
        locs = base_map.sum(axis=1).nonzero()
        voxels = base_map[locs[0], :, locs[1]]
        com_sup_inf = np.asarray(
            [
                int(np.ceil(sni.measurements.center_of_mass(voxels[i, :])[0]))
                for i in range(voxels.shape[0])
            ]
        )
        region_mask_sup_deep = np.full(base_map.shape, self._REGION_DEEP_KEY)
        for i in range(len(com_sup_inf)):
            region_mask_sup_deep[
                locs[0][i], : com_sup_inf[i], locs[1][i]
            ] = self._REGION_SUPERFICIAL_KEY

        # M/L
        cum_ml = np.nonzero(base_map.sum(axis=(0, 1)))[0]  # noqa: F841
        # midpoint_ml = int(np.ceil((np.min(cum_ml) + np.max(cum_ml)) / 2))
        midpoint_ml = int(np.ceil(sni.measurements.center_of_mass(base_map)[2]))
        region_mask_med_lat = np.full(base_map.shape, self._LATERAL_KEY)
        medial_span = slice(0, midpoint_ml) if self.medial_to_lateral else slice(midpoint_ml, None)
        region_mask_med_lat[:, :, medial_span] = self._MEDIAL_KEY

        self.regions_mask = np.stack([region_mask_sup_deep, region_mask_med_lat], axis=-1)

    def __calc_quant_vals__(self, quant_map, map_type):
        subject_pid = self.pid

        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"

        quant_map_volume = quant_map.volume
        mask = self.__mask__.volume
        quant_map_volume = mask * quant_map_volume

        deep_superficial_map = self.regions_mask[..., 0]
        med_lat_map = self.regions_mask[..., 1]

        axial_names = ["superficial", "deep", "total"]
        sagittal_names = ["medial", "lateral"]

        pd_header = ["Subject", "Location", "Condyle", "Mean", "Std", "Median"]
        pd_list = []

        regions = itertools.product(
            [self._REGION_SUPERFICIAL_KEY, self._REGION_DEEP_KEY, self._TOTAL_AXIAL_KEY],
            [self._MEDIAL_KEY, self._LATERAL_KEY],
        )
        for axial, sagittal in regions:
            if axial == self._TOTAL_AXIAL_KEY:
                axial_map = np.asarray(
                    deep_superficial_map == self._REGION_SUPERFICIAL_KEY, dtype=np.float32
                ) + np.asarray(deep_superficial_map == self._REGION_DEEP_KEY, dtype=np.float32)
                axial_map = np.asarray(axial_map, dtype=np.bool)
            else:
                axial_map = deep_superficial_map == axial

            sagittal_map = med_lat_map == sagittal

            curr_region_mask = quant_map_volume * axial_map * sagittal_map
            curr_region_mask = curr_region_mask[curr_region_mask != 0]

            # discard all values that are 0
            c_mean = np.nanmean(curr_region_mask)
            c_std = np.nanstd(curr_region_mask)
            c_median = np.nanmedian(curr_region_mask)

            row_info = [
                subject_pid,
                axial_names[axial],
                sagittal_names[sagittal],
                c_mean,
                c_std,
                c_median,
            ]

            pd_list.append(row_info)

        # Generate 2D unrolled matrix
        total, superficial, deep = self.unroll_coronal(quant_map.volume)

        df = pd.DataFrame(pd_list, columns=pd_header)
        qv_name = map_type.name
        maps = [
            {
                "title": "%s superficial" % qv_name,
                "data": superficial,
                "xlabel": "Slice",
                "ylabel": "Angle (binned)",
                "filename": "%s_superficial" % qv_name,
                "raw_data_filename": "%s_superficial.data" % qv_name,
            },
            {
                "title": "%s deep" % qv_name,
                "data": deep,
                "xlabel": "Slice",
                "ylabel": "Angle (binned)",
                "filename": "%s_deep" % qv_name,
                "raw_data_filename": "%s_deep.data" % qv_name,
            },
            {
                "title": "%s total" % qv_name,
                "data": total,
                "xlabel": "Slice",
                "ylabel": "Angle (binned)",
                "filename": "%s_total" % qv_name,
                "raw_data_filename": "%s_total.data" % qv_name,
            },
        ]

        self.__store_quant_vals__(maps, df, map_type)

    def set_mask(self, mask, use_largest_cc: bool = True):
        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)

        self.split_regions(self.__mask__.volume)

    def __save_quant_data__(self, dirpath):
        """Save quantitative data and 2D visualizations of patellar cartilage

        Check which quantitative values (T2, T1rho, etc) are defined for
        patellar cartilage and analyze these:

        1. Save 2D total, superficial, and deep visualization maps
        2. Save {'medial', 'lateral'}, {'anterior', 'posterior'},
            {'superior', 'inferior', 'total'} data to excel file

        :param dirpath: base filepath to save 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 = ""
                ylabel = ""
                title = q_map_data["title"]
                data_map = q_map_data["data"]

                axs_bounds = self.__get_axis_bounds__(data_map, leave_buffer=True)

                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)
                plt.ylim(axs_bounds[0])
                plt.gca().invert_yaxis()
                plt.xlim(axs_bounds[1])
                # plt.axis('tight')
                clb = plt.colorbar()
                clb.ax.set_ylabel("(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)