from typing import Sequence, Union, Dict, Optional
import numpy as np
import matplotlib.pyplot as plt
import math
import torch
from torch import nn, optim, Tensor
from scipy import stats


class Calibration:
    """
    Methods for probability calibration
    """

    @staticmethod
    def reliability_diagram(
        pred: Sequence[Union[np.ndarray, float]],
        target: Sequence[Union[np.ndarray, int]],
        num_bins: Optional[int] = 10,
        num_quantiles: Optional[int] = None,
        output_filename: Optional[str] = None,
    ) -> Dict:
        """
        Reliability diagram
        :param pred: List of arrays of shape [NUM_CLASSES]
        :param target: List of integers [0, NUM_CLASSES) denoting ground truth classes
        :num_bins: Integer denoting the number of equal width bins in the diagram's x-axis.
        :num_quantiles: Integer denoting the number of equal number of samples quantiles in the diagram's x-axis.
            Note, either num_bins or num_quantiles should be set to None, as they define different ways to split the x-axis.
        :output_filename: Path to file to which to save the plot
        :return: saving roc curve to a file and return the input to figure to a dictionary
        """
        assert num_bins is None or num_quantiles is None

        if isinstance(pred[0], float):  # binary case
            pred = [np.array([1 - p, p]) for p in pred]

        max_pred = [np.max(p) for p in pred]
        cls_pred = [np.argmax(p) for p in pred]

        assert all(
            [math.isclose(p.sum(), 1, rel_tol=1e-5) for p in pred]
        )  # make sure that the inputs make proper probabilities

        total_samples = len(pred)
        max_pred = np.array(max_pred)
        if num_bins is not None:
            conf_vec = np.linspace(0, 1, num_bins + 1)
        else:
            quantiles_vec = np.linspace(0, 1, num_quantiles + 1)
            conf_vec = stats.mstats.mquantiles(max_pred, quantiles_vec.tolist())
        acc_vec = np.zeros(len(conf_vec) - 1)
        fraction_of_samples = np.zeros(len(conf_vec) - 1)
        num_samples = np.zeros(len(conf_vec) - 1)
        for i in range(len(conf_vec) - 1):
            cnt_correct_in_conf = sum(
                [
                    p_cls == t
                    for (p_cls, p, t) in zip(cls_pred, max_pred, target)
                    if p >= conf_vec[i] and p < conf_vec[i + 1]
                ]
            )
            cnt_conf = len(
                [p for p in max_pred if p >= conf_vec[i] and p < conf_vec[i + 1]]
            )
            num_samples[i] = cnt_conf
            fraction_of_samples[i] = cnt_conf / np.max(
                [total_samples, np.finfo(float).eps]
            )
            acc_vec[i] = cnt_correct_in_conf / np.max([cnt_conf, np.finfo(float).eps])
        mid_conf_vec = np.diff(conf_vec) / 2 + conf_vec[:-1]
        cnt_correct_total = sum([p_cls == t for (p_cls, t) in zip(cls_pred, target)])
        avg_acc = cnt_correct_total / total_samples
        avg_conf = np.mean(max_pred)
        # extract info for the plot
        results = {}
        results["conf_vec"] = conf_vec
        results["mid_confidence_points"] = mid_conf_vec
        results["accuracy"] = acc_vec
        results["avg_accuracy"] = avg_acc
        results["avg_confidence"] = avg_conf
        results["num_samples_per_bin"] = num_samples
        results["total_samples"] = total_samples

        # display
        if output_filename is not None:
            f, (ax1, ax2) = plt.subplots(2, 1, sharex=True, sharey=True)
            f.tight_layout(pad=3.0)
            ax1.plot(
                mid_conf_vec, fraction_of_samples, color="blue", marker="o", linewidth=2
            )
            ax1.set_title(
                f"Confidence Histogram. Avg. Accuracy: {np.round(avg_acc, 2)}, Avg. Confidence: {np.round(avg_conf, 2)}"
            )
            ax1.set_xlabel("Confidence")
            ax1.set_ylabel("Fraction of samples")
            ax1.grid()
            ax2.plot(conf_vec, conf_vec, color="black", linewidth=2)
            ax2.plot(mid_conf_vec, acc_vec, color="red", marker="o", linewidth=2)
            ax2.set_title("Reliability Curve")
            ax2.set_xlabel("Confidence")
            ax2.set_ylabel("Accuracy")
            ax2.grid()
            f.savefig(output_filename)
            plt.close()

        return results

    @staticmethod
    def ece(
        pred: Sequence[Union[np.ndarray, float]],
        target: Sequence[Union[np.ndarray, int]],
        num_bins: Optional[int] = 10,
        num_quantiles: Optional[int] = None,
    ) -> float:
        """
        :param pred: prediction array per sample. Each element shape [num_classes]
        :param target: target per sample. Each element is an integer in range [0 - num_classes)
        :num_bins: Integer denoting the number of equal width bins in the diagram's x-axis.
        :num_quantiles: Integer denoting the number of equal number of samples quantiles in the diagram's x-axis.
            Note, either num_bins or num_quantiles should be set to None, as they define different ways to split the x-axis.
        :return Expected Calibration Error (ECE) score
        """
        reliability_results = Calibration.reliability_diagram(
            pred,
            target,
            num_bins=num_bins,
            num_quantiles=num_quantiles,
            output_filename=None,
        )
        conf_vec = reliability_results["conf_vec"][1:]
        total_samples = reliability_results["total_samples"]
        num_samples_per_bin = reliability_results["num_samples_per_bin"]
        existing_bins = num_samples_per_bin > 0
        accuracy_vec = reliability_results["accuracy"]
        # expected calibration error
        ece = (1.0 / total_samples) * np.sum(
            num_samples_per_bin * np.abs(accuracy_vec - conf_vec)
        )
        # maximum calibration error
        # we filter out bins which don't contain any samples
        mce = np.max(np.abs(accuracy_vec[existing_bins] - conf_vec[existing_bins]))

        results = {}
        results["ece"] = ece
        results["mce"] = mce

        return results

    @staticmethod
    def find_temperature(
        pred: Sequence[Union[np.ndarray, float]],
        target: Sequence[Union[np.ndarray, int]],
    ) -> float:
        """
        :param pred: prediction array per sample. Each element shape [num_classes]
            note that pred should be logits *before* conversion to softmax probabilities
        :param target: target per sample. Each element is an integer in range [0 - num_classes)
        :return temperature scale for calibration
        We optimize to find the temperature similarly to: https://github.com/gpleiss/temperature_scaling
        """
        nll_criterion = nn.CrossEntropyLoss()
        temperature = nn.Parameter(torch.ones(1) * 1.5)
        optimizer = optim.LBFGS([temperature], lr=0.01, max_iter=50)
        with torch.no_grad():
            logits = torch.tensor(np.array(pred))
            target = torch.tensor(target)

        def eval() -> Tensor:
            optimizer.zero_grad()
            loss = nll_criterion(logits / temperature, target)
            loss.backward()
            return loss

        optimizer.step(eval)

        return temperature.item()

    @staticmethod
    def apply_temperature(
        pred: Sequence[Union[np.ndarray, float]], temperature: Union[float, None] = None
    ) -> float:
        """
        :param pred: prediction array per sample. Each element shape [num_classes]
            note that pred should be logits *before* conversion to softmax probabilities
        :param temperature: temperature scalar
        :return array of softmax probabilities
        """
        if temperature is None:
            pred_calibrated = pred
        else:
            if isinstance(pred[0], np.ndarray):
                pred = np.array(pred)
            pred_calibrated = torch.tensor(pred / temperature)
            pred_calibrated = nn.Softmax(dim=1)(pred_calibrated)

        return np.array(pred_calibrated)
