from typing import Any, Dict, List, Optional, Union
import numpy as np
from easyMetrics.core.base import Metric

class AUC(Metric):
    """
    计算分类任务的AUC（Area Under the ROC Curve）指标。
    支持二分类场景，衡量模型对正负样本的区分能力。
    """
    def __init__(self, 
                 method: str = 'linear'):
        """
        参数:
            method: ROC曲线下面积的计算方法。可选值：
                - 'linear': 线性插值方法
                - 'trapezoidal': 梯形积分方法
        """
        super().__init__()
        self.method = method
        
        # 重置内部状态
        self.reset()

    def reset(self):
        """
        重置指标的内部状态。
        """
        self.scores = []  # 预测分数或概率
        self.labels = []  # 真实标签（0或1）

    def _ensure_numpy_array(self, data, dtype=None):
        """
        确保输入数据是numpy数组。
        """
        if isinstance(data, np.ndarray):
            return data.astype(dtype) if dtype is not None else data
        elif isinstance(data, (list, tuple)):
            return np.array(data, dtype=dtype)
        elif np.isscalar(data):
            return np.array([data], dtype=dtype)
        else:
            return np.array(data, dtype=dtype)

    def update(self, preds: Union[List[Any], np.ndarray], target: Union[List[Any], np.ndarray]):
        """
        更新指标状态。
        
        参数:
            preds: 模型预测结果。可以是概率值或分数。
            target: 真实标签。二分类标签（0或1）。
        """
        # 转换为numpy数组
        preds = self._ensure_numpy_array(preds, dtype=float)
        target = self._ensure_numpy_array(target, dtype=int)
        
        # 确保形状一致
        assert preds.shape == target.shape, f"预测和标签形状不一致: {preds.shape} vs {target.shape}"
        
        # 处理概率数组（二分类取正类概率）
        if preds.ndim == 2:
            assert preds.shape[1] == 2, f"二分类概率数组应为 (n_samples, 2)，但得到 {preds.shape}"
            preds = preds[:, 1]  # 取正类概率
        
        # 确保是一维数组
        assert preds.ndim == 1, f"预测应为一维数组，得到 {preds.ndim} 维"
        assert target.ndim == 1, f"标签应为一维数组，得到 {target.ndim} 维"
        
        # 确保标签是二分类（0或1）
        unique_labels = np.unique(target)
        assert set(unique_labels).issubset({0, 1}), f"标签应为二分类（0或1），但得到 {unique_labels}"
        
        # 存储数据
        self.scores.extend(preds.tolist())
        self.labels.extend(target.tolist())

    def _compute_roc_curve(self, scores: np.ndarray, labels: np.ndarray):
        """
        计算ROC曲线的TPR和FPR。
        
        返回:
            fpr: 假阳性率数组
            tpr: 真阳性率数组
            thresholds: 阈值数组
        """
        # 排序（从高到低）
        sorted_indices = np.argsort(-scores)
        sorted_scores = scores[sorted_indices]
        sorted_labels = labels[sorted_indices]
        
        # 计算累积TP和FP
        cumulative_tp = np.cumsum(sorted_labels)
        cumulative_fp = np.cumsum(1 - sorted_labels)
        
        # 计算总阳性和总阴性
        total_positive = np.sum(labels)
        total_negative = len(labels) - total_positive
        
        # 避免除零错误
        if total_positive == 0 or total_negative == 0:
            return np.array([0.0, 1.0]), np.array([0.0, 0.0]), np.array([1.0, 0.0])
        
        # 计算TPR和FPR
        tpr = cumulative_tp / total_positive
        fpr = cumulative_fp / total_negative
        
        # 添加 (0, 0) 点
        tpr = np.concatenate([[0.0], tpr])
        fpr = np.concatenate([[0.0], fpr])
        thresholds = np.concatenate([[1.0], sorted_scores])
        
        return fpr, tpr, thresholds

    def _compute_auc(self, fpr: np.ndarray, tpr: np.ndarray) -> float:
        """
        计算ROC曲线下的面积。
        """
        if self.method == 'linear':
            # 线性插值方法
            auc = 0.0
            for i in range(1, len(fpr)):
                auc += 0.5 * (fpr[i] - fpr[i-1]) * (tpr[i] + tpr[i-1])
        elif self.method == 'trapezoidal':
            # 梯形积分方法
            auc = np.trapz(tpr, fpr)
        else:
            raise ValueError(f"不支持的计算方法: {self.method}")
        
        return auc

    def compute(self) -> Dict[str, float]:
        """
        计算AUC指标。
        
        返回:
            包含AUC值的字典。
        """
        # 转换为numpy数组
        scores = np.array(self.scores, dtype=float)
        labels = np.array(self.labels, dtype=int)
        
        # 检查数据是否为空
        if len(scores) == 0:
            return {'auc': 0.0}
        
        # 计算ROC曲线
        fpr, tpr, thresholds = self._compute_roc_curve(scores, labels)
        
        # 计算AUC
        auc = self._compute_auc(fpr, tpr)
        
        # 确保AUC在合理范围内
        auc = max(0.0, min(1.0, auc))
        
        return {'auc': auc}
