import numpy as np
from sklearn import metrics
from collections import defaultdict
import torch
import torch.nn as nn
from sklearn.metrics import average_precision_score

def get_accracy(output, label):
    _, prediction = torch.max(output, 1)    # argmax
    correct = (prediction == label).sum().item()
    accuracy = correct / prediction.size(0)
    return accuracy


def get_prediction(output, label):
    prob = nn.functional.softmax(output, dim=1)[:, 1]
    prob = prob.view(prob.size(0), 1)
    label = label.view(label.size(0), 1)
    #print(prob.size(), label.size())
    datas = torch.cat((prob, label.float()), dim=1)
    return datas


def calculate_acc_for_train(label, output, num_classes):
    """
    Compute Accuracy and mAP for a multi-class classification task.
    
    Args:
    label: Ground-truth labels with shape [batch_size], where each element is a class index from 0 to num_classes - 1.
    output: Model outputs with shape [batch_size, num_classes], usually logits.
    num_classes: Total number of classes
    
    Returns:
    accuracy: Accuracy score
    map_score: Mean Average Precision (mAP)
    """
    # Compute accuracy score
    _, prediction = torch.max(output, 1)  # Use the class with the highest probability as the prediction
    correct = (prediction == label).sum().item()
    accuracy = correct / prediction.size(0)
    
    # Compute mAP
    # Convert outputs to a probability distribution
    probs = torch.softmax(output, dim=1)  # Apply softmax across the class dimension for multi-class classification
    
    # Convert to NumPy arrays for sklearn utilities
    probs_np = probs.cpu().detach().numpy()
    labels_np = label.cpu().detach().numpy()
    
    # Compute AP for each class
    aps = []
    for class_idx in range(num_classes):
        # Build binary labels: current class as 1, all others as 0
        binary_labels = (labels_np == class_idx).astype(int)
        # Predicted probability for the current class
        class_probs = probs_np[:, class_idx]
        
        # Compute AP for this class
        try:
            ap = average_precision_score(binary_labels, class_probs)
            aps.append(ap)
        except ValueError:
            print("Error")
            aps.append(0.0)
    
    map_score = np.mean(aps)
    
    return accuracy, map_score


def to_numpy(x):
    if isinstance(x, torch.Tensor):
        # If the input is a Tensor, detach it first to avoid gradient tracking, then move it to CPU and convert it to NumPy.
        return x.detach().cpu().numpy() if x.is_cuda else x.numpy()
    elif isinstance(x, np.ndarray):
        # If it is already a NumPy array, return it directly.
        return x
    else:
        raise TypeError(f"Unsupported data type: {type(x)}，Only torch.Tensor and numpy.ndarray are supported")




def calculate_acc_for_test(label, output, num_classes):
    """
    Compute Accuracy and mAP for a multi-class classification task.
    Note: this version assumes `output` is already a probability distribution
    (i.e. softmax has already been applied).
    
    Args:
        label: Ground-truth labels with shape [batch_size], where each element is
            a class index from 0 to num_classes - 1.
        output: Model outputs with shape [batch_size, num_classes], already in
            probability form.
        num_classes: Total number of classes.
    
    Returns:
        accuracy: Accuracy score.
        map_score: Mean Average Precision (mAP).
    """
    
    # --------------------------
    # 1. Compute accuracy
    # --------------------------
    label = to_numpy(label)
    output = to_numpy(output)
    
    prediction = np.argmax(output, axis=1)  # Take the index of the class with the highest probability
    
    # for i,j in zip(label, prediction):
    #     print(i, j)
    
    correct = np.sum(prediction == label)
    accuracy = correct / len(label)  # len(label) is the total number of samples
    
    # --------------------------
    # 2. Compute mAP
    # --------------------------
    aps = []
    for class_idx in range(num_classes):
        # Build binary labels: current class as 1, all others as 0
        binary_labels = (label == class_idx).astype(int)
        # Predicted probability for the current class
        class_probs = output[:, class_idx]
        
        # Check whether this class has both positive and negative samples to avoid meaningless computation
        has_positive = np.any(binary_labels == 1)
        has_negative = np.any(binary_labels == 0)
        
        if not (has_positive and has_negative):
            # If only positive or only negative samples exist, AP cannot be computed, so skip this class
            if(has_positive):
                print(f"Warning: class {class_idx} is missing negative samples, skipping AP computation")
            else:
                print(f"Warning: class {class_idx} is missing positive samples, skipping AP computation")
            continue  # Skip directly and exclude it from mAP computation
        
        # Compute AP while handling possible numerical issues
        try:
            # Clamp the probability range to improve stability and avoid extreme values
            class_probs_clamped = np.clip(class_probs, 1e-8, 1 - 1e-8)
            ap = average_precision_score(binary_labels, class_probs_clamped)
            aps.append(ap)
        except Exception as e:
            print(f"Class {class_idx} failed to compute AP: {e}")
            continue  # Skip if the computation fails
    
    # Compute mAP; if all classes are skipped, set mAP to 0
    if len(aps) == 0:
        map_score = 0.0
        print("Warning: AP cannot be computed for any class, setting mAP to 0")
    else:
        map_score = np.mean(aps)
    
    # Compute binary accuracy and mAP
    bin_pridiction=np.asarray(prediction, dtype=bool)
    bin_lable=np.asarray(label, dtype=bool)
    correct=np.sum(bin_pridiction==bin_lable)
    bin_acc=correct/len(label)
    
    bin_class_probs_true=output[:,0]
    bin_class_probs_false=1-bin_class_probs_true
    has_positive = np.any(bin_lable == True)
    has_negative = np.any(bin_lable == False)
    if not (has_positive and has_negative):
        bin_mAP=0.0
    else:
        true_clamped = np.clip(bin_class_probs_true, 1e-8, 1 - 1e-8)
        false_clamped=np.clip(bin_class_probs_false, 1e-8, 1 - 1e-8)
        true_ap=average_precision_score(~bin_lable, true_clamped)
        false_ap=average_precision_score(bin_lable, false_clamped)
        #print('True:{t}, False:{f}'.format(t=true_ap,f=false_ap))
        bin_mAP=(true_ap+false_ap)/2
    
    return {'acc': accuracy, 'mAP': map_score, 'pred': output, 'label': label, 'bin_acc':bin_acc,'bin_mAP':bin_mAP}



def calculate_metrics_for_train(label, output):
    if output.size(1) != 1:
        prob = torch.softmax(output, dim=1)[:, 1]
    else:
        prob = output

    # Accuracy
    _, prediction = torch.max(output, 1)
    correct = (prediction == label).sum().item()
    accuracy = correct / prediction.size(0)

    # Average Precision
    y_true = label.cpu().detach().numpy()
    y_pred = prob.cpu().detach().numpy()
    ap = metrics.average_precision_score(y_true, y_pred)

    # AUC and EER
    try:
        fpr, tpr, thresholds = metrics.roc_curve(label.squeeze().cpu().numpy(),
                                                 prob.squeeze().cpu().numpy(),
                                                 pos_label=1)
    except:
        # for the case when we only have one sample
        return None, None, accuracy, ap

    if np.isnan(fpr[0]) or np.isnan(tpr[0]):
        # for the case when all the samples within a batch is fake/real
        auc, eer = None, None
    else:
        auc = metrics.auc(fpr, tpr)
        fnr = 1 - tpr
        eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]

    return auc, eer, accuracy, ap


# ------------ compute average metrics of batches---------------------
class Metrics_batch():
    def __init__(self):
        self.tprs = []
        self.mean_fpr = np.linspace(0, 1, 100)
        self.aucs = []
        self.eers = []
        self.aps = []

        self.correct = 0
        self.total = 0
        self.losses = []

    def update(self, label, output):
        acc = self._update_acc(label, output)
        if output.size(1) == 2:
            prob = torch.softmax(output, dim=1)[:, 1]
        else:
            prob = output
        #label = 1-label
        #prob = torch.softmax(output, dim=1)[:, 1]
        auc, eer = self._update_auc(label, prob)
        ap = self._update_ap(label, prob)

        return acc, auc, eer, ap

    def _update_auc(self, lab, prob):
        fpr, tpr, thresholds = metrics.roc_curve(lab.squeeze().cpu().numpy(),
                                                 prob.squeeze().cpu().numpy(),
                                                 pos_label=1)
        if np.isnan(fpr[0]) or np.isnan(tpr[0]):
            return -1, -1

        auc = metrics.auc(fpr, tpr)
        interp_tpr = np.interp(self.mean_fpr, fpr, tpr)
        interp_tpr[0] = 0.0
        self.tprs.append(interp_tpr)
        self.aucs.append(auc)

        # return auc

        # EER
        fnr = 1 - tpr
        eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
        self.eers.append(eer)

        return auc, eer

    def _update_acc(self, lab, output):
        _, prediction = torch.max(output, 1)    # argmax
        correct = (prediction == lab).sum().item()
        accuracy = correct / prediction.size(0)
        # self.accs.append(accuracy)
        self.correct = self.correct+correct
        self.total = self.total+lab.size(0)
        return accuracy

    def _update_ap(self, label, prob):
        y_true = label.cpu().detach().numpy()
        y_pred = prob.cpu().detach().numpy()
        ap = metrics.average_precision_score(y_true,y_pred)
        self.aps.append(ap)

        return np.mean(ap)

    def get_mean_metrics(self):
        mean_acc, std_acc = self.correct/self.total, 0
        mean_auc, std_auc = self._mean_auc()
        mean_err, std_err = np.mean(self.eers), np.std(self.eers)
        mean_ap, std_ap = np.mean(self.aps), np.std(self.aps)
        
        return {'acc':mean_acc, 'auc':mean_auc, 'eer':mean_err, 'ap':mean_ap}

    def _mean_auc(self):
        mean_tpr = np.mean(self.tprs, axis=0)
        mean_tpr[-1] = 1.0
        mean_auc = metrics.auc(self.mean_fpr, mean_tpr)
        std_auc = np.std(self.aucs)
        return mean_auc, std_auc

    def clear(self):
        self.tprs.clear()
        self.aucs.clear()
        # self.accs.clear()
        self.correct=0
        self.total=0
        self.eers.clear()
        self.aps.clear()
        self.losses.clear()


# ------------ compute average metrics of all data ---------------------
class Metrics_all():
    def __init__(self):
        self.probs = []
        self.labels = []
        self.correct = 0
        self.total = 0

    def store(self, label, output):
        prob = torch.softmax(output, dim=1)[:, 1]
        _, prediction = torch.max(output, 1)    # argmax
        correct = (prediction == label).sum().item()
        self.correct += correct
        self.total += label.size(0)
        self.labels.append(label.squeeze().cpu().numpy())
        self.probs.append(prob.squeeze().cpu().numpy())

    def get_metrics(self):
        y_pred = np.concatenate(self.probs)
        y_true = np.concatenate(self.labels)
        # auc
        fpr, tpr, thresholds = metrics.roc_curve(y_true,y_pred,pos_label=1)
        auc = metrics.auc(fpr, tpr)
        # eer
        fnr = 1 - tpr
        eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
        # ap
        ap = metrics.average_precision_score(y_true,y_pred)
        # acc
        acc = self.correct / self.total
        return {'acc':acc, 'auc':auc, 'eer':eer, 'ap':ap}

    def clear(self):
        self.probs.clear()
        self.labels.clear()
        self.correct = 0
        self.total = 0


# only used to record a series of scalar value
class Recorder:
    def __init__(self):
        self.sum = 0
        self.num = 0
    def update(self, item, num=1):
        if item is not None:
            self.sum += item * num
            self.num += num
    def average(self):
        if self.num == 0:
            return None
        return self.sum/self.num
    def clear(self):
        self.sum = 0
        self.num = 0