import sys
import yaml
from typing import Optional
from pprint import pprint
import pandas as pd
from tqdm import tqdm
import torch
import torchio
import SimpleITK as sitk
import numpy as np

from GANDLF.config_manager import ConfigManager
from GANDLF.utils import find_problem_type_from_parameters, one_hot
from GANDLF.metrics import (
    overall_stats,
    structural_similarity_index,
    mean_squared_error,
    peak_signal_noise_ratio,
    mean_squared_log_error,
    mean_absolute_error,
    ncc_mean,
    ncc_std,
    ncc_max,
    ncc_min,
)
from GANDLF.losses.segmentation import dice
from GANDLF.metrics.segmentation import (
    _calculator_generic_all_surface_distances,
    _calculator_sensitivity_specificity,
    _calculator_jaccard,
)


def generate_metrics_dict(
    input_csv: str, config: str, outputfile: Optional[str] = None
) -> dict:
    """
    This function generates metrics from the input csv and the config.

    Args:
        input_csv (str): The input CSV.
        config (str): The input yaml config.
        outputfile (str, optional): The output file to save the metrics. Defaults to None.

    Returns:
        dict: The metrics dictionary.
    """
    input_df = pd.read_csv(input_csv)

    # check required headers in a case insensitive manner
    headers = {}
    required_columns = ["subjectid", "prediction", "target"]
    for col, _ in input_df.items():
        col_lower = col.lower()
        for column_to_check in required_columns:
            if column_to_check == col_lower:
                headers[column_to_check] = col
        if col_lower == "mask":
            headers["mask"] = col
    for column in required_columns:
        assert column in headers, f"The input csv should have a column named {column}"

    overall_stats_dict = {}
    parameters = ConfigManager(config)
    problem_type = parameters.get("problem_type", None)
    problem_type = (
        find_problem_type_from_parameters(parameters)
        if problem_type is None
        else problem_type
    )
    parameters["problem_type"] = problem_type

    if problem_type == "regression" or problem_type == "classification":
        parameters["model"]["num_classes"] = len(parameters["model"]["class_list"])
        predictions_tensor = torch.from_numpy(
            input_df[headers["prediction"]].to_numpy().ravel()
        )
        labels_tensor = torch.from_numpy(input_df[headers["target"]].to_numpy().ravel())
        overall_stats_dict = overall_stats(
            predictions_tensor, labels_tensor, parameters
        )

    elif problem_type == "segmentation":
        # read images and then calculate metrics
        class_list = parameters["model"]["class_list"]
        for _, row in tqdm(input_df.iterrows(), total=input_df.shape[0]):
            current_subject_id = row[headers["subjectid"]]
            overall_stats_dict[current_subject_id] = {}
            label_image = torchio.LabelMap(row[headers["target"]])
            pred_image = torchio.LabelMap(row[headers["prediction"]])
            label_tensor = label_image.data
            pred_tensor = pred_image.data
            spacing = label_image.spacing
            if label_tensor.data.shape[-1] == 1:
                spacing = spacing[0:2]
            # add dimension for batch
            parameters["subject_spacing"] = torch.Tensor(spacing).unsqueeze(0)
            label_tensor = label_tensor.unsqueeze(0)
            pred_tensor = pred_tensor.unsqueeze(0)

            # one hot encode with batch_size = 1
            label_image_one_hot = one_hot(label_tensor, class_list)
            pred_image_one_hot = one_hot(pred_tensor, class_list)

            for class_index, _ in enumerate(class_list):
                current_target = label_image_one_hot[:, class_index, ...].unsqueeze(0)
                current_prediction = pred_image_one_hot[:, class_index, ...].unsqueeze(
                    0
                )
                overall_stats_dict[current_subject_id][str(class_index)] = {}
                # this is inconsequential, since one_hot will ensure that the classes are present
                parameters["model"]["class_list"] = [1]
                parameters["model"]["num_classes"] = 1
                overall_stats_dict[current_subject_id][str(class_index)]["dice"] = dice(
                    current_prediction, current_target
                ).item()
                nsd, hd100, hd95 = _calculator_generic_all_surface_distances(
                    current_prediction, current_target, parameters
                )
                overall_stats_dict[current_subject_id][str(class_index)][
                    "nsd"
                ] = nsd.item()
                overall_stats_dict[current_subject_id][str(class_index)][
                    "hd100"
                ] = hd100.item()
                overall_stats_dict[current_subject_id][str(class_index)][
                    "hd95"
                ] = hd95.item()

                (s, p) = _calculator_sensitivity_specificity(
                    current_prediction, current_target, parameters
                )
                overall_stats_dict[current_subject_id][str(class_index)][
                    "sensitivity"
                ] = s.item()
                overall_stats_dict[current_subject_id][str(class_index)][
                    "specificity"
                ] = p.item()
                overall_stats_dict[current_subject_id][
                    "jaccard_" + str(class_index)
                ] = _calculator_jaccard(
                    current_prediction, current_target, parameters
                ).item()
                current_target_image = sitk.GetImageFromArray(
                    current_target[0, 0, ...].long()
                )
                current_prediction_image = sitk.GetImageFromArray(
                    current_prediction[0, 0, ...].long()
                )
                label_overlap_filter = sitk.LabelOverlapMeasuresImageFilter()
                label_overlap_filter.Execute(
                    current_target_image, current_prediction_image
                )
                # overall_stats_dict[current_subject_id][
                #     "falseDiscoveryRate_" + str(class_index)
                # ] = label_overlap_filter.GetFalseDiscoveryRate()
                overall_stats_dict[current_subject_id][
                    "falseNegativeError_" + str(class_index)
                ] = label_overlap_filter.GetFalseNegativeError()
                overall_stats_dict[current_subject_id][
                    "falsePositiveError_" + str(class_index)
                ] = label_overlap_filter.GetFalsePositiveError()
                overall_stats_dict[current_subject_id][
                    "meanOverlap_" + str(class_index)
                ] = label_overlap_filter.GetMeanOverlap()
                overall_stats_dict[current_subject_id][
                    "unionOverlap_" + str(class_index)
                ] = label_overlap_filter.GetUnionOverlap()
                overall_stats_dict[current_subject_id][
                    "volumeSimilarity_" + str(class_index)
                ] = label_overlap_filter.GetVolumeSimilarity()

    elif problem_type == "synthesis":

        def __fix_2d_tensor(input_tensor):
            """
            This function checks for 2d images and change the shape to [B, C, H, W]

            Args:
                input_tensor (torch.Tensor): The input tensor.

            Returns:
                torch.Tensor: The output tensor in the format that torchmetrics expects.
            """
            if input_tensor.shape[-1] == 1:
                return input_tensor.squeeze(-1).unsqueeze(0)
            else:
                return input_tensor

        def __percentile_clip(
            input_tensor: torch.Tensor,
            reference_tensor: torch.Tensor = None,
            p_min: Optional[float] = 0.5,
            p_max: Optional[float] = 99.5,
            strictlyPositive: Optional[bool] = True,
        ):
            """
            Normalizes a tensor based on percentiles. Clips values below and above the percentile.
            Percentiles for normalization can come from another tensor.

            Args:
                input_tensor (torch.Tensor): Tensor to be normalized based on the data from the reference_tensor. If reference_tensor is None, the percentiles from this tensor will be used.
                reference_tensor (torch.Tensor, optional): The tensor used for obtaining the percentiles.
                p_min (float, optional): Lower end percentile. Defaults to 0.5.
                p_max (float, optional): Upper end percentile. Defaults to 99.5.
                strictlyPositive (bool, optional): Ensures that really all values are above 0 before normalization. Defaults to True.

            Returns:
                torch.Tensor: The input_tensor normalized based on the percentiles of the reference tensor.
            """
            reference_tensor = (
                input_tensor if reference_tensor is None else reference_tensor
            )
            v_min, v_max = np.percentile(
                reference_tensor, [p_min, p_max]
            )  # get p_min percentile and p_max percentile

            # set lower bound to be 0 if strictlyPositive is enabled
            v_min = max(v_min, 0.0) if strictlyPositive else v_min
            output_tensor = np.clip(
                input_tensor, v_min, v_max
            )  # clip values to percentiles from reference_tensor
            output_tensor = (output_tensor - v_min) / (
                v_max - v_min
            )  # normalizes values to [0;1]
            return output_tensor

        for _, row in tqdm(input_df.iterrows(), total=input_df.shape[0]):
            current_subject_id = row[headers["subjectid"]]
            overall_stats_dict[current_subject_id] = {}
            target_image = __fix_2d_tensor(
                torchio.ScalarImage(row[headers["target"]]).data
            )
            pred_image = __fix_2d_tensor(
                torchio.ScalarImage(row[headers["prediction"]]).data
            )
            # if "mask" is not in the row, we assume that the whole image is the mask
            # always cast to byte tensor
            mask = (
                __fix_2d_tensor(torchio.LabelMap(row[headers["mask"]]).data)
                if "mask" in row
                else torch.from_numpy(
                    np.ones(target_image.numpy().shape, dtype=np.uint8)
                )
            ).byte()

            # Get Infill region (we really are only interested in the infill region)
            output_infill = (pred_image * mask).float()
            gt_image_infill = (target_image * mask).float()

            # Normalize to [0;1] based on GT (otherwise MSE will depend on the image intensity range)
            normalize = parameters.get("normalize", True)
            if normalize:
                reference_tensor = (
                    target_image * ~mask
                )  # use all the tissue that is not masked for normalization
                gt_image_infill = __percentile_clip(
                    gt_image_infill,
                    reference_tensor=reference_tensor,
                    p_min=0.5,
                    p_max=99.5,
                    strictlyPositive=True,
                )
                output_infill = __percentile_clip(
                    output_infill,
                    reference_tensor=reference_tensor,
                    p_min=0.5,
                    p_max=99.5,
                    strictlyPositive=True,
                )

            overall_stats_dict[current_subject_id][
                "ssim"
            ] = structural_similarity_index(output_infill, gt_image_infill, mask).item()

            # ncc metrics
            compute_ncc = parameters.get("compute_ncc", True)
            if compute_ncc:
                overall_stats_dict[current_subject_id]["ncc_mean"] = ncc_mean(
                    output_infill, gt_image_infill
                )
                overall_stats_dict[current_subject_id]["ncc_std"] = ncc_std(
                    output_infill, gt_image_infill
                )
                overall_stats_dict[current_subject_id]["ncc_max"] = ncc_max(
                    output_infill, gt_image_infill
                )
                overall_stats_dict[current_subject_id]["ncc_min"] = ncc_min(
                    output_infill, gt_image_infill
                )

            # only voxels that are to be inferred (-> flat array)
            # these are required for mse, psnr, etc.
            gt_image_infill = gt_image_infill[mask]
            output_infill = output_infill[mask]

            overall_stats_dict[current_subject_id]["mse"] = mean_squared_error(
                output_infill, gt_image_infill
            ).item()

            overall_stats_dict[current_subject_id]["msle"] = mean_squared_log_error(
                output_infill, gt_image_infill
            ).item()

            overall_stats_dict[current_subject_id]["mae"] = mean_absolute_error(
                output_infill, gt_image_infill
            ).item()

            # torchmetrics PSNR using "max"
            overall_stats_dict[current_subject_id]["psnr"] = peak_signal_noise_ratio(
                output_infill, gt_image_infill
            ).item()

            # same as above but with epsilon for robustness
            overall_stats_dict[current_subject_id][
                "psnr_eps"
            ] = peak_signal_noise_ratio(
                output_infill, gt_image_infill, epsilon=sys.float_info.epsilon
            ).item()

            # only use fix data range to [0;1] if the data was normalized before
            if normalize:
                # torchmetrics PSNR but with fixed data range of 0 to 1
                overall_stats_dict[current_subject_id][
                    "psnr_01"
                ] = peak_signal_noise_ratio(
                    output_infill, gt_image_infill, data_range=(0, 1)
                ).item()

                # same as above but with epsilon for robustness
                overall_stats_dict[current_subject_id][
                    "psnr_01_eps"
                ] = peak_signal_noise_ratio(
                    output_infill,
                    gt_image_infill,
                    data_range=(0, 1),
                    epsilon=sys.float_info.epsilon,
                ).item()

    pprint(overall_stats_dict)
    if outputfile is not None:
        with open(outputfile, "w") as outfile:
            yaml.dump(overall_stats_dict, outfile)
