"""
Metrics for computing evalutation results
Modified from vanilla PANet code by Wang et al.
"""

import numpy as np

class Metric(object):
    """
    Compute evaluation result

    Args:
        max_label:
            max label index in the data (0 denoting background)
        n_scans:
            number of test scans
    """
    def __init__(self, max_label=20, n_scans=None):
        self.labels = list(range(max_label + 1))  # all class labels
        self.n_scans = 1 if n_scans is None else n_scans

        # list of list of array, each array save the TP/FP/FN statistic of a testing sample
        self.tp_lst = [[] for _ in range(self.n_scans)]
        self.fp_lst = [[] for _ in range(self.n_scans)]
        self.fn_lst = [[] for _ in range(self.n_scans)]
        self.slice_counter = [0 for _ in range(self.n_scans)]

    def reset(self):
        """
        Reset accumulated evaluation. 
        """
        # assert self.n_scans == 1, 'Should not reset accumulated result when we are not doing one-time batch-wise validation'
        del self.tp_lst, self.fp_lst, self.fn_lst
        self.tp_lst = [[] for _ in range(self.n_scans)]
        self.fp_lst = [[] for _ in range(self.n_scans)]
        self.fn_lst = [[] for _ in range(self.n_scans)]
        
    def reset_scan(self, n_scan, labels:list=None):
        """
        Reset accumulated evaluation for a specific scan. 
        """
        if labels is None:
            labels = self.labels
        for slice_idx in range(len(self.tp_lst[n_scan])):
            for label in labels:
                self.tp_lst[n_scan][slice_idx][label] = np.nan
                self.fp_lst[n_scan][slice_idx][label] = np.nan
                self.fn_lst[n_scan][slice_idx][label] = np.nan

    def record(self, pred, target, labels=None, n_scan=None):
        """
        Record the evaluation result for each sample and each class label, including:
            True Positive, False Positive, False Negative

        Args:
            pred:
                predicted mask array, expected shape is H x W
            target:
                target mask array, expected shape is H x W
            labels:
                only count specific label, used when knowing all possible labels in advance
        """
        assert pred.shape == target.shape

        if self.n_scans == 1:
            n_scan = 0

        # array to save the TP/FP/FN statistic for each class (plus BG)
        tp_arr = np.full(len(self.labels), np.nan)
        fp_arr = np.full(len(self.labels), np.nan)
        fn_arr = np.full(len(self.labels), np.nan)
        if labels is None:
            labels = self.labels
        else:
            labels = [0,] + labels

        for j, label in enumerate(labels):
            # Get the location of the pixels that are predicted as class j
            # idx = np.where(np.logical_and(pred == j, target != 255))
            # pred_idx_j = set(zip(idx[0].tolist(), idx[1].tolist()))
            # # Get the location of the pixels that are class j in ground truth
            # idx = np.where(target == j)
            # target_idx_j = set(zip(idx[0].tolist(), idx[1].tolist()))

            # # this should not work: if target_idx_j:  # if ground-truth contains this class
            # # the author is adding posion to the code
            # tp_arr[label] = len(set.intersection(pred_idx_j, target_idx_j))
            # fp_arr[label] = len(pred_idx_j - target_idx_j)
            # fn_arr[label] = len(target_idx_j - pred_idx_j)
            
            # calc the tp, fp and fn normally and compare the 2 values
            tp = ((pred == j).astype(int) * (target == j).astype(int)).sum()
            fp = ((pred == j).astype(int) * (target != j).astype(int)).sum()
            fn = ((pred != j).astype(int) * (target == j).astype(int)).sum()
            
            tp_arr[label] = tp
            fp_arr[label] = fp
            fn_arr[label] = fn
            
            # assert tp == tp_arr[label]
            # assert fp == fp_arr[label]
            # assert fn == fn_arr[label]

        self.tp_lst[n_scan].append(tp_arr)
        self.fp_lst[n_scan].append(fp_arr)
        self.fn_lst[n_scan].append(fn_arr)
        self.slice_counter[n_scan] += 1

    def get_mIoU(self, labels=None, n_scan=None):
        """
        Compute mean IoU

        Args:
            labels:
                specify a subset of labels to compute mean IoU, default is using all classes
        """
        if labels is None:
            labels = self.labels
        # Sum TP, FP, FN statistic of all samples
        if n_scan is None:
            tp_sum = [np.nansum(np.vstack(self.tp_lst[_scan]), axis=0).take(labels)
                      for _scan in range(self.n_scans)]
            fp_sum = [np.nansum(np.vstack(self.fp_lst[_scan]), axis=0).take(labels)
                      for _scan in range(self.n_scans)]
            fn_sum = [np.nansum(np.vstack(self.fn_lst[_scan]), axis=0).take(labels)
                      for _scan in range(self.n_scans)]

            # Compute mean IoU classwisely
            # Average across n_scans, then average over classes
            mIoU_class = np.vstack([tp_sum[_scan] / (tp_sum[_scan] + fp_sum[_scan] + fn_sum[_scan])
                                    for _scan in range(self.n_scans)])
            mIoU = mIoU_class.mean(axis=1)

            return (mIoU_class.mean(axis=0), mIoU_class.std(axis=0),
                    mIoU.mean(axis=0), mIoU.std(axis=0))
        else:
            tp_sum = np.nansum(np.vstack(self.tp_lst[n_scan]), axis=0).take(labels)
            fp_sum = np.nansum(np.vstack(self.fp_lst[n_scan]), axis=0).take(labels)
            fn_sum = np.nansum(np.vstack(self.fn_lst[n_scan]), axis=0).take(labels)

            # Compute mean IoU classwisely and average over classes
            mIoU_class = tp_sum / (tp_sum + fp_sum + fn_sum)
            mIoU = mIoU_class.mean()

            return mIoU_class, mIoU

    def get_mDice(self, labels=None, n_scan=None, give_raw = False):
        """
        Compute mean Dice score (in 3D scan level)

        Args:
            labels:
                specify a subset of labels to compute mean IoU, default is using all classes
        """
        # NOTE: unverified
        if labels is None:
            labels = self.labels
        # Sum TP, FP, FN statistic of all samples
        if n_scan is None:
            tp_sum = [np.nansum(np.vstack(self.tp_lst[_scan]), axis=0).take(labels)
                      for _scan in range(self.n_scans)]
            fp_sum = [np.nansum(np.vstack(self.fp_lst[_scan]), axis=0).take(labels)
                      for _scan in range(self.n_scans)]
            fn_sum = [np.nansum(np.vstack(self.fn_lst[_scan]), axis=0).take(labels)
                      for _scan in range(self.n_scans)]

            # Average across n_scans, then average over classes
            mDice_class = np.vstack([ 2 * tp_sum[_scan] / ( 2 * tp_sum[_scan] + fp_sum[_scan] + fn_sum[_scan])
                                    for _scan in range(self.n_scans)])
            mDice = mDice_class.mean(axis=1)
            print(f"mDice_class:\n {mDice_class}")
            if not give_raw:
                return (mDice_class.mean(axis=0), mDice_class.std(axis=0),
                    mDice.mean(axis=0), mDice.std(axis=0))
            else:
                return (mDice_class.mean(axis=0), mDice_class.std(axis=0),
                    mDice.mean(axis=0), mDice.std(axis=0), mDice_class)

        else:
            tp_sum = np.nansum(np.vstack(self.tp_lst[n_scan]), axis=0).take(labels)
            fp_sum = np.nansum(np.vstack(self.fp_lst[n_scan]), axis=0).take(labels)
            fn_sum = np.nansum(np.vstack(self.fn_lst[n_scan]), axis=0).take(labels)

            # Compute mean IoU classwisely and average over classes
            mDice_class = 2 * tp_sum / ( 2 * tp_sum + fp_sum + fn_sum)
            mDice = mDice_class.mean()

            if not give_raw:
                return (mDice_class, mDice, mDice_class)
            
            return (mDice_class, mDice, mDice_class)

    def get_mPrecRecall(self, labels=None, n_scan=None, give_raw = False):
        """
        Compute precision and recall

        Args:
            labels:
                specify a subset of labels to compute mean IoU, default is using all classes
        """
        # NOTE: unverified
        if labels is None:
            labels = self.labels
        # Sum TP, FP, FN statistic of all samples
        if n_scan is None:
            tp_sum = [np.nansum(np.vstack(self.tp_lst[_scan]), axis=0).take(labels)
                      for _scan in range(self.n_scans)]
            fp_sum = [np.nansum(np.vstack(self.fp_lst[_scan]), axis=0).take(labels)
                      for _scan in range(self.n_scans)]
            fn_sum = [np.nansum(np.vstack(self.fn_lst[_scan]), axis=0).take(labels)
                      for _scan in range(self.n_scans)]

            # Compute mean IoU classwisely
            # Average across n_scans, then average over classes
            mPrec_class = np.vstack([ tp_sum[_scan] / ( tp_sum[_scan] + fp_sum[_scan] )
                                    for _scan in range(self.n_scans)])

            mRec_class = np.vstack([ tp_sum[_scan] / ( tp_sum[_scan] + fn_sum[_scan] )
                                    for _scan in range(self.n_scans)])

            mPrec = mPrec_class.mean(axis=1)
            mRec  = mRec_class.mean(axis=1)
            if not give_raw:
                return (mPrec_class.mean(axis=0), mPrec_class.std(axis=0), mPrec.mean(axis=0), mPrec.std(axis=0), mRec_class.mean(axis=0), mRec_class.std(axis=0), mRec.mean(axis=0), mRec.std(axis=0))
            else:
                return (mPrec_class.mean(axis=0), mPrec_class.std(axis=0), mPrec.mean(axis=0), mPrec.std(axis=0), mRec_class.mean(axis=0), mRec_class.std(axis=0), mRec.mean(axis=0), mRec.std(axis=0), mPrec_class, mRec_class)


        else:
            tp_sum = np.nansum(np.vstack(self.tp_lst[n_scan]), axis=0).take(labels)
            fp_sum = np.nansum(np.vstack(self.fp_lst[n_scan]), axis=0).take(labels)
            fn_sum = np.nansum(np.vstack(self.fn_lst[n_scan]), axis=0).take(labels)

            # Compute mean IoU classwisely and average over classes
            mPrec_class = tp_sum / (tp_sum + fp_sum)
            mPrec = mPrec_class.mean()

            mRec_class = tp_sum / (tp_sum + fn_sum)
            mRec = mRec_class.mean()

            return mPrec_class, None, mPrec, None, mRec_class, None, mRec, None,  mPrec_class, mRec_class

    def get_mIoU_binary(self, n_scan=None):
        """
        Compute mean IoU for binary scenario
        (sum all foreground classes as one class)
        """
        # Sum TP, FP, FN statistic of all samples
        if n_scan is None:
            tp_sum = [np.nansum(np.vstack(self.tp_lst[_scan]), axis=0)
                      for _scan in range(self.n_scans)]
            fp_sum = [np.nansum(np.vstack(self.fp_lst[_scan]), axis=0)
                      for _scan in range(self.n_scans)]
            fn_sum = [np.nansum(np.vstack(self.fn_lst[_scan]), axis=0)
                      for _scan in range(self.n_scans)]

            # Sum over all foreground classes
            tp_sum = [np.c_[tp_sum[_scan][0], np.nansum(tp_sum[_scan][1:])]
                      for _scan in range(self.n_scans)]
            fp_sum = [np.c_[fp_sum[_scan][0], np.nansum(fp_sum[_scan][1:])]
                      for _scan in range(self.n_scans)]
            fn_sum = [np.c_[fn_sum[_scan][0], np.nansum(fn_sum[_scan][1:])]
                      for _scan in range(self.n_scans)]

            # Compute mean IoU classwisely and average across classes
            mIoU_class = np.vstack([tp_sum[_scan] / (tp_sum[_scan] + fp_sum[_scan] + fn_sum[_scan])
                                    for _scan in range(self.n_scans)])
            mIoU = mIoU_class.mean(axis=1)

            return (mIoU_class.mean(axis=0), mIoU_class.std(axis=0),
                    mIoU.mean(axis=0), mIoU.std(axis=0))
        else:
            tp_sum = np.nansum(np.vstack(self.tp_lst[n_scan]), axis=0)
            fp_sum = np.nansum(np.vstack(self.fp_lst[n_scan]), axis=0)
            fn_sum = np.nansum(np.vstack(self.fn_lst[n_scan]), axis=0)

            # Sum over all foreground classes
            tp_sum = np.c_[tp_sum[0], np.nansum(tp_sum[1:])]
            fp_sum = np.c_[fp_sum[0], np.nansum(fp_sum[1:])]
            fn_sum = np.c_[fn_sum[0], np.nansum(fn_sum[1:])]

            mIoU_class = tp_sum / (tp_sum + fp_sum + fn_sum)
            mIoU = mIoU_class.mean()

            return mIoU_class, mIoU