"""训练曲线评估模块

给定若干训练 Metric （如 Loss， Acc） 曲线的向量数值，判断哪些 Run 更好，哪些 Run 更差，并给出定量的评估。
尽量符合肉眼观测的直觉感受；能精准区分完全覆盖优势的 Run，而对于肉眼观测难分胜负的 Run ，也能给出合理的量化。
"""

import typing

from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
from scipy import stats
from typing_extensions import TypedDict


@dataclass
class CurveResult:
    """单条曲线的评估结果

    Attributes
    ----------
    - `is_converged` : `bool`
        - 是否收敛
    - `is_diverged` : `bool`
        - 是否发散
    - `best` : `Tuple[int, float]`
        - 最佳值的位置和数值 (index, value)
    - `score` : `float`
        - 综合评分，越高越好
    - `score_components` : `Dict[str, float]`
        - 评分的详细分量 {'performance': float, 'speed': float, 'stability': float}
    """

    is_converged: bool
    is_diverged: bool
    best: Tuple[int, float]
    score: float
    score_components: Optional[Dict[str, float]] = None


@dataclass
class CurveEvaluation:
    """多条曲线的评估结果

    Attributes
    ----------
    - `ranking` : `List[Union[int, str]]`
        - 曲线排名，从最好到最差
    - `evaluation` : `Dict[Union[int, str], CurveResult]`
        - 每条曲线的详细评估结果
    - `confidence` : `Dict[Tuple[Union[int, str], Union[int, str]], float]`
        - 两两曲线间的置信度，表示第一个优于第二个的把握程度
    """

    ranking: List[Union[int, str]]
    evaluation: Dict[Union[int, str], CurveResult]
    confidence: Dict[Tuple[Union[int, str], Union[int, str]], float]

    def print_table(self):
        """打印评估结果表格

        使用prettytable库打印详细的评估结果表格。
        """
        try:
            from prettytable import PrettyTable
        except ImportError:
            raise ImportError('需要安装prettytable库: pip install prettytable')

        # 创建表格
        table = PrettyTable()
        table.field_names = [
            '排名',
            '曲线',
            '评分',
            '性能分',
            '面积分',
            '速度分',
            '稳定性分',
            '收敛',
            '发散',
            '最佳值',
        ]

        # 填充数据
        for i, key in enumerate(self.ranking):
            result = self.evaluation[key]
            score_components = result.score_components or {}
            performance_score = score_components.get('performance', 0.0)
            area_score = score_components.get('area', 0.0)
            speed_score = score_components.get('speed', 0.0)
            stability_score = score_components.get('stability', 0.0)

            table.add_row(
                [
                    f'#{i+1}',
                    str(key),
                    f'{result.score:.4f}',
                    f'{performance_score:.4f}',
                    f'{area_score:.4f}',
                    f'{speed_score:.4f}',
                    f'{stability_score:.4f}',
                    '是' if result.is_converged else '否',
                    '是' if result.is_diverged else '否',
                    f'epoch{result.best[0]}: {result.best[1]:.4f}',
                ]
            )

        print(table)

        # 打印置信度矩阵（如果曲线数量不多）
        if len(self.ranking) <= 5:
            print('\n置信度矩阵 (行优于列的概率):')
            conf_table = PrettyTable()
            conf_table.field_names = [''] + [str(key) for key in self.ranking]

            for key1 in self.ranking:
                row = [str(key1)]
                for key2 in self.ranking:
                    if key1 == key2:
                        row.append('-')
                    else:
                        conf = self.confidence.get((key1, key2), 0.0)
                        row.append(f'{conf:.2f}')
                conf_table.add_row(row)

            print(conf_table)

    def plot_matrix(self, figsize=(8, 6), cmap='RdYlGn'):
        """绘制置信度矩阵热力图

        Parameters
        ----------
        - `figsize` : `tuple`, optional
            - 图形大小, by default (8, 6)
        - `cmap` : `str`, optional
            - 颜色映射, by default 'RdYlGn'
        """
        try:
            import matplotlib.pyplot as plt
            import numpy as np
            import seaborn as sns
        except ImportError:
            raise ImportError(
                '需要安装matplotlib和seaborn库: pip install matplotlib seaborn'
            )

        # 创建置信度矩阵
        n = len(self.ranking)
        matrix = np.zeros((n, n))

        for i, key1 in enumerate(self.ranking):
            for j, key2 in enumerate(self.ranking):
                if i == j:
                    matrix[i, j] = 0.5  # 对角线设为0.5
                else:
                    matrix[i, j] = self.confidence.get((key1, key2), 0.0)

        # 绘制热力图
        fig = plt.figure(figsize=figsize)
        sns.heatmap(
            matrix,
            annot=True,
            fmt='.2f',
            xticklabels=[str(key) for key in self.ranking],
            yticklabels=[str(key) for key in self.ranking],
            cmap=cmap,
            vmin=0.0,
            vmax=1.0,
            square=True,
        )
        plt.title('置信度矩阵 (行优于列的概率)')
        plt.xlabel('对比曲线')
        plt.ylabel('参考曲线')
        plt.tight_layout()
        # plt.show()
        fig.show()


class ScoreWeights(TypedDict):
    performance: float
    area: float
    speed: float
    stability: float


DEFAULT_SCORE_WEIGHTS: ScoreWeights = {
    'performance': 30.0,
    'area': 25.0,
    'speed': 25.0,
    'stability': 20.0,
}


class CurveEvaluator:
    """训练曲线评估器

    提供训练曲线的收敛性、发散性分析，以及多曲线的排名和置信度评估功能。
    """

    def __init__(
        self,
        convergence_window: int = 10,
        convergence_threshold: float = 0.01,
        divergence_threshold: float = 0.1,
        speed_target_ratio: float = 0.8,
        performance_best_value_weight: float = 0.8,
        initial_value_smoothing_window: int = 3,
        score_weights: ScoreWeights | None = None,
    ):
        """初始化曲线评估器

        Parameters
        ----------
        - `convergence_window` : `int`, optional
            - 收敛判断的滑动窗口大小。
              在评估收敛性时，会检查曲线最后 `convergence_window` 个点的变化率。
              默认为 10。
        - `convergence_threshold` : `float`, optional
            - 收敛判断的相对变化阈值。
              如果曲线最后 `convergence_window` 个点的平均相对变化率小于此值，则认为已收敛。
              默认为 0.01 (1%)。
        - `divergence_threshold` : `float`, optional
            - 发散判断的相对变化阈值。
              如果曲线最后部分呈现显著恶化趋势，且恶化幅度超过此阈值，则认为已发散。
              默认为 0.1 (10%)。
        - `speed_target_ratio` : `float`, optional
            - 速度评分的目标改进比例，范围 [0, 1]。
              用于衡量曲线达到目标性能的速度。目标性能是初始值到最佳值之间按此比例计算的中间点。
              例如，0.8 表示目标是达到 80% 的最大可能改进。
              默认为 0.8。
        - `performance_best_value_weight` : `float`, optional
            - 性能评分中“历史最佳值”的权重因子，范围 [0, 1]。
              最终用于性能评分的值是 `weight * best_value + (1 - weight) * final_value`。
              1.0 表示完全关注历史最佳值，0.0 表示完全关注最终值。
              默认为 0.8。
        - `initial_value_smoothing_window` : `int`, optional
            - 用于计算初始值的平滑窗口大小。
              速度评分计算时，会使用曲线前 `initial_value_smoothing_window` 个点的均值作为初始值，
              以减少噪声影响。设为 1 则等同于直接使用第一个点。
              默认为 3。
        - `score_weights` : `Optional[Dict[str, float]]`, optional
            - 综合评分中各分量的权重字典。
              键为分量名称 ('performance', 'area', 'speed', 'stability')，值为对应的权重 (百分比)。
              未指定的权重将从默认值 `DEFAULT_SCORE_WEIGHTS` 中获取。
              例如 `{'performance': 40, 'speed': 60}` 表示性能分占40%，速度分占60%。
              默认为 None，即使用 `DEFAULT_SCORE_WEIGHTS`。
        """
        self.convergence_window = convergence_window
        self.convergence_threshold = convergence_threshold
        self.divergence_threshold = divergence_threshold
        self.speed_target_ratio = speed_target_ratio

        if not 0.0 <= performance_best_value_weight <= 1.0:
            raise ValueError('performance_best_value_weight 必须在 [0, 1] 范围内')
        self.performance_best_value_weight = performance_best_value_weight

        if initial_value_smoothing_window < 1:
            raise ValueError('initial_value_smoothing_window 必须大于等于 1')
        self.initial_value_smoothing_window = initial_value_smoothing_window

        if score_weights is None:
            self.score_weights = DEFAULT_SCORE_WEIGHTS.copy()
        else:
            self.score_weights = DEFAULT_SCORE_WEIGHTS.copy()
            self.score_weights.update(score_weights)

    def _normalize_curves(
        self,
        curves: Union[
            typing.Mapping[Union[int, str], np.ndarray], typing.Sequence[np.ndarray]
        ],
        better: str,
    ) -> Tuple[Dict[Union[int, str], np.ndarray], Dict[Union[int, str], float]]:
        """标准化曲线数据格式并进行问题归一化

        将所有指标统一转换为'越小越好'的形式，方便后续统一处理。

        Returns
        ----------
        - `Tuple[Dict[Union[int, str], np.ndarray], Dict[Union[int, str], float]]`
            - 归一化后的曲线数据和原始最佳值映射
        """
        if isinstance(curves, dict):
            curve_dict = curves
        else:
            curve_dict = {i: curve for i, curve in enumerate(curves)}

        # 转换为numpy数组并处理NaN值
        normalized = {}
        original_best = {}  # 保存原始最佳值用于结果转换

        for key, curve in curve_dict.items():
            curve_array = np.asarray(curve, dtype=float)
            # 移除NaN值
            valid_mask = ~np.isnan(curve_array)
            if not np.any(valid_mask):
                raise ValueError(f'曲线 {key} 包含全部NaN值')
            clean_curve = curve_array[valid_mask]

            # 保存原始最佳值
            if better == 'low':
                original_best[key] = float(np.min(clean_curve))
            else:
                original_best[key] = float(np.max(clean_curve))

            # 问题归一化：统一转换为'越小越好'
            if better == 'high':
                # 对于'high'指标，取反使其变为'low'指标
                normalized[key] = -clean_curve
            else:
                # 对于'low'指标，保持不变
                normalized[key] = clean_curve

        return normalized, original_best

    def _detect_convergence(self, curve: np.ndarray) -> bool:
        """检测曲线是否收敛

        基于滑动窗口的相对变化率判断收敛性。统一按'越小越好'的逻辑处理。
        """
        if len(curve) < self.convergence_window * 2:
            return False

        # 取最后一段进行收敛判断
        tail_length = min(len(curve) // 3, self.convergence_window * 3)
        tail = curve[-tail_length:]

        if len(tail) < self.convergence_window:
            return False

        # 计算滑动窗口内的相对变化率
        window_changes = []
        for i in range(len(tail) - self.convergence_window + 1):
            window = tail[i : i + self.convergence_window]
            # 检查是否不再显著下降
            relative_change = (window[-1] - window[0]) / (abs(window[0]) + 1e-8)
            window_changes.append(abs(relative_change))

        # 如果最近的变化都很小，认为已收敛
        avg_change = np.mean(window_changes[-self.convergence_window // 2 :])
        return bool(avg_change < self.convergence_threshold)

    def _detect_divergence(self, curve: np.ndarray) -> bool:
        """检测曲线是否发散

        检测持续恶化的趋势。统一按'越小越好'的逻辑处理。
        """
        if len(curve) < self.convergence_window:
            return False

        # 取最后一段进行发散判断
        tail_length = min(len(curve) // 2, self.convergence_window * 2)
        tail = curve[-tail_length:]

        # 计算趋势
        x = np.arange(len(tail))
        regression = typing.cast(
            tuple[float, float, float, float, float], stats.linregress(x, tail)
        )
        slope = float(regression[0])
        r_value = float(regression[2])
        p_value = float(regression[3])

        # 判断是否显著恶化
        is_significant = p_value < 0.05 and abs(r_value) > 0.5

        # 统一处理：值应该下降，如果显著上升则为发散
        is_worsening = slope > 0

        # 检查恶化幅度
        relative_change = abs(slope * len(tail)) / (abs(tail[0]) + 1e-8)
        is_severe = relative_change > self.divergence_threshold

        return is_significant and is_worsening and is_severe

    def _find_best_value(self, curve: np.ndarray) -> Tuple[int, float]:
        """找到曲线的最佳值及其位置"""
        # 统一处理：寻找最小值
        best_idx = np.argmin(curve)
        return int(best_idx), float(curve[best_idx])

    def _calculate_score(
        self, curve: np.ndarray, global_stats: Optional[typing.Mapping[str, float]] = None
    ) -> Tuple[float, Dict[str, float]]:
        """计算曲线的综合评分

        综合考虑最终性能、收敛速度、稳定性等因素。
        """
        if len(curve) == 0:
            return 0.0, {'performance': 0.0, 'area': 0.0, 'speed': 0.0, 'stability': 0.0}

        # 1. 性能评分
        best_idx, best_value = self._find_best_value(curve)
        final_value = curve[-1]

        # 根据 performance_best_value_weight 决定用于评分的值
        # 它是 best_value 和 final_value 的加权平均
        value_to_score = (
            self.performance_best_value_weight * best_value
            + (1 - self.performance_best_value_weight) * final_value
        )

        # 使用全局分位数进行鲁棒归一化
        # 统一处理：越接近p25越好，超出范围则截断到[0,1]
        if global_stats is not None:
            p25, p75 = global_stats['p25'], global_stats['p75']
            iqr = p75 - p25
            if iqr == 0:
                performance_score = 1.0
            else:
                # 统一处理：越接近p25越好
                performance_score = max(0.0, min(1.0, (p75 - value_to_score) / iqr))
        else:
            # 回退到非全局处理的方法
            curve_range = np.max(curve) - np.min(curve)
            if curve_range == 0:
                performance_score = 1.0
            else:
                # 统一处理：越小越好
                performance_score = (np.max(curve) - value_to_score) / curve_range

        # 2. 面积效率评分 (25%)
        # 基于曲线下面积的全程训练效率
        if global_stats is not None and 'area_p25' in global_stats:
            area = np.trapz(curve)
            area_p25, area_p75 = global_stats['area_p25'], global_stats['area_p75']
            area_iqr = area_p75 - area_p25
            if area_iqr == 0:
                area_score = 1.0
            else:
                # 统一处理：面积越小越好
                area_score = max(0.0, min(1.0, (area_p75 - area) / area_iqr))
        else:
            # 回退方案：基于相对面积改进
            area = np.trapz(curve)
            # 构造理想曲线进行比较
            # 统一处理：理想情况是面积一直保持最佳值（最小值）
            ideal_area = len(curve) * np.min(curve)  # 理想情况：一直保持最佳值
            area_score = max(
                0.0,
                min(
                    1.0,
                    (area - ideal_area)
                    / (np.trapz([curve[0]] * len(curve)) - ideal_area),
                ),
            )

        # 3. 收敛速度评分
        # 使用平滑后的初始值
        n = min(self.initial_value_smoothing_window, len(curve))
        initial_value = np.mean(curve[:n]) if n > 0 else curve[0]
        # 统一处理：目标是从初始值向最佳值改进指定比例
        # 对于归一化后的指标，目标是从初始值向最佳值改进指定比例
        target_performance = (
            initial_value - (initial_value - best_value) * self.speed_target_ratio
        )
        convergence_idx = np.where(curve <= target_performance)[0]

        if len(convergence_idx) > 0:
            convergence_ratio = convergence_idx[0] / len(curve)
            speed_score = 1.0 - convergence_ratio  # 越早达到越好
        else:
            speed_score = 0.0

        # 4. 稳定性评分 (20%)
        # 基于后半段的变异系数
        second_half = curve[len(curve) // 2 :]
        if len(second_half) > 1:
            cv = np.std(second_half) / (abs(np.mean(second_half)) + 1e-8)
            stability_score = 1.0 / (1.0 + cv)  # 变异系数越小越稳定
        else:
            stability_score = 1.0

        # 综合评分
        total_score = (
            self.score_weights.get('performance', 0.0) * performance_score
            + self.score_weights.get('area', 0.0) * area_score
            + self.score_weights.get('speed', 0.0) * speed_score
            + self.score_weights.get('stability', 0.0) * stability_score
        )

        score_components = {
            'performance': performance_score,
            'area': area_score,
            'speed': speed_score,
            'stability': stability_score,
        }

        return float(total_score), score_components

    def _compute_global_stats(self, curves: typing.Mapping[Union[int, str], np.ndarray]):
        """计算所有曲线的全局统计信息

        使用分位数方法进行鲁棒统计。
        """
        # 收集所有曲线的最终值和面积
        final_values = [curve[-1] for curve in curves.values()]
        area_values = [np.trapz(curve) for curve in curves.values()]

        # 计算25%和75%分位数（对应RobustScaler的思路）
        p25, p75 = np.percentile(final_values, [25, 75])
        area_p25, area_p75 = np.percentile(area_values, [25, 75])

        return {
            'p25': float(p25),
            'p75': float(p75),
            'median': float(np.median(final_values)),
            'iqr': float(p75 - p25),
            'area_p25': float(area_p25),
            'area_p75': float(area_p75),
            'area_iqr': float(area_p75 - area_p25),
        }

    def _calculate_confidence(
        self, curves: typing.Mapping[Union[int, str], np.ndarray]
    ) -> Dict[Tuple[Union[int, str], Union[int, str]], float]:
        """计算两两曲线间的置信度

        使用统计检验评估一条曲线优于另一条曲线的置信度。
        统一按'越小越好'的逻辑处理（因为输入数据已归一化）。
        """
        confidence = {}
        keys = list(curves.keys())

        for i, key1 in enumerate(keys):
            for j, key2 in enumerate(keys):
                if i == j:
                    continue

                curve1, curve2 = curves[key1], curves[key2]

                # 使用Wilcoxon符号秩检验比较两条曲线
                # 为了比较，我们取两条曲线的重叠部分
                min_len = min(len(curve1), len(curve2))
                c1_sample = curve1[-min_len:]
                c2_sample = curve2[-min_len:]

                try:
                    # 统一处理：curve1更小更好
                    # diff > 0 表示 curve1 优于 curve2
                    diff = c2_sample - c1_sample

                    # 移除零差异
                    non_zero_diff = diff[diff != 0]

                    if len(non_zero_diff) < 3:
                        # 样本太少，使用简单比较
                        # 统一处理：curve1的均值更小表示更好
                        conf = float(
                            np.mean(curve1[-min_len // 2 :])
                            < np.mean(curve2[-min_len // 2 :])
                        )
                    else:
                        # 使用Wilcoxon检验
                        wilcoxon_result = typing.cast(
                            tuple[float, float], stats.wilcoxon(non_zero_diff)
                        )
                        p_value = float(wilcoxon_result[1])

                        # 转换p值为置信度
                        # 统一处理：diff > 0 表示 curve1 优于 curve2
                        if np.mean(non_zero_diff) > 0:
                            conf = 1.0 - p_value / 2  # 单侧检验
                        else:
                            conf = p_value / 2

                    confidence[(key1, key2)] = max(0.0, min(1.0, conf))

                except Exception:
                    # 如果统计检验失败，使用简单比较
                    # 统一处理：curve1的均值更小表示更好
                    conf = float(np.mean(curve1) < np.mean(curve2))
                    confidence[(key1, key2)] = conf

        return confidence

    def evaluate_curves(
        self,
        curves: Union[
            typing.Mapping[Union[int, str], np.ndarray], typing.Sequence[np.ndarray]
        ],
        better: str = 'low',
        # metric_name: Optional[str] = None,
    ) -> CurveEvaluation:
        """评估多条训练曲线

        Parameters
        ----------
        - `curves` : `Union[Dict[Union[int, str], np.ndarray], List[np.ndarray]]`
            - 训练曲线数据，可以是字典或列表格式
        - `better` : `str`, optional
            - 指标的优化方向，'low'表示越小越好，'high'表示越大越好, by default 'low'

        Returns
        ----------
        - `CurveEvaluation`
            - 包含排名、详细评估结果和置信度的评估结果

        Examples
        ----------
        ```python
        # 评估loss曲线
        evaluator = CurveEvaluator()
        curves = {
            'run1': [1.0, 0.8, 0.6, 0.5, 0.45, 0.44, 0.43],
            'run2': [1.2, 0.9, 0.7, 0.6, 0.55, 0.50, 0.48],
            'run3': [0.9, 0.7, 0.5, 0.4, 0.35, 0.30, 0.25]
        }
        result = evaluator.evaluate_curves(curves, better='low', metric_name='loss')
        print(f"最佳运行: {result.ranking[0]}")
        ```
        """
        if better not in ['low', 'high']:
            raise ValueError("better参数必须是'low'或'high'")

        # 标准化输入数据并进行问题归一化
        normalized_curves, original_best = self._normalize_curves(curves, better)

        if len(normalized_curves) == 0:
            raise ValueError('没有有效的曲线数据')

        # 计算全局统计信息用于鲁棒归一化
        global_stats = self._compute_global_stats(normalized_curves)

        # 评估每条曲线
        evaluation = {}
        for key, curve in normalized_curves.items():
            is_converged = self._detect_convergence(curve)
            is_diverged = self._detect_divergence(curve)
            best = self._find_best_value(curve)
            score, score_components = self._calculate_score(curve, global_stats)

            # 对于'high'指标，需要将最佳值转换回原始值
            if better == 'high':
                original_best_value = -best[1]  # 取反恢复原始值
            else:
                original_best_value = best[1]

            evaluation[key] = CurveResult(
                is_converged=is_converged,
                is_diverged=is_diverged,
                best=(best[0], original_best_value),
                score=score,
                score_components=score_components,
            )

        # 计算置信度
        confidence = self._calculate_confidence(normalized_curves)

        # 根据评分排序
        ranking = sorted(
            evaluation.keys(), key=lambda k: evaluation[k].score, reverse=True
        )

        return CurveEvaluation(
            ranking=ranking, evaluation=evaluation, confidence=confidence
        )


def evaluate_training_curves(
    curves: Union[
        typing.Mapping[Union[int, str], np.ndarray], typing.Sequence[np.ndarray]
    ],
    better: str = 'low',
    *,
    convergence_window: int = 10,
    convergence_threshold: float = 0.01,
    divergence_threshold: float = 0.1,
    speed_target_ratio: float = 0.8,
    performance_best_value_weight: float = 0.8,
    initial_value_smoothing_window: int = 3,
    score_weights: ScoreWeights | None = None,
) -> CurveEvaluation:
    """评估训练曲线的便捷函数

    这是一个便捷函数，它创建一个 `CurveEvaluator` 实例并调用其 `evaluate_curves` 方法。
    所有参数都直接传递给 `CurveEvaluator` 的构造函数。

    Parameters
    ----------
    - `curves` : `Union[Dict[Union[int, str], np.ndarray], List[np.ndarray]]`
        - 训练曲线数据，可以是字典或列表格式。
    - `better` : `str`, optional
        - 指标的优化方向，'low'表示越小越好（如Loss），'high'表示越大越好（如Accuracy）。
          默认为 'low'。
    - `convergence_window` : `int`, optional
        - 收敛判断的滑动窗口大小。
          默认为 10。
    - `convergence_threshold` : `float`, optional
        - 收敛判断的相对变化阈值。
          默认为 0.01 (1%)。
    - `divergence_threshold` : `float`, optional
        - 发散判断的相对变化阈值。
          默认为 0.1 (10%)。
    - `speed_target_ratio` : `float`, optional
        - 速度评分的目标改进比例，范围 [0, 1]。
          默认为 0.8。
    - `performance_best_value_weight` : `float`, optional
        - 性能评分中“历史最佳值”的权重因子，范围 [0, 1]。
          默认为 0.8。
    - `initial_value_smoothing_window` : `int`, optional
        - 用于计算初始值的平滑窗口大小。
          默认为 3。
    - `score_weights` : `Optional[Dict[str, float]]`, optional
        - 综合评分中各分量的权重字典。
          默认为 None，即使用 `DEFAULT_SCORE_WEIGHTS`。

    Returns
    ----------
    - `CurveEvaluation`
        - 包含排名、详细评估结果和置信度的评估结果。
    """
    evaluator = CurveEvaluator(
        convergence_window=convergence_window,
        convergence_threshold=convergence_threshold,
        divergence_threshold=divergence_threshold,
        speed_target_ratio=speed_target_ratio,
        performance_best_value_weight=performance_best_value_weight,
        initial_value_smoothing_window=initial_value_smoothing_window,
        score_weights=score_weights,
    )
    return evaluator.evaluate_curves(curves, better)
