# 文件: zyplib\artifact\training_curves.py
"""
训练曲线模拟模块; 用于教学模拟用

该模块用于生成各种类型的机器学习训练曲线，包括：
- 正常收敛的训练曲线
- 过拟合场景的训练曲线
- 欠拟合场景的训练曲线
- 学习率异常导致的训练曲线

所有生成的曲线都尽可能接近真实训练过程中的表现。
"""

from dataclasses import dataclass
from typing import Literal, Optional, Tuple

import numpy as np
from numpy.typing import NDArray


@dataclass
class TrainingCurveConfig:
    """训练曲线配置参数"""

    n_epochs: int = 100
    """训练轮数"""

    initial_train_loss: float = 2.5
    """初始训练loss"""

    initial_val_loss: float = 2.8
    """初始验证loss"""

    final_train_loss: float = 0.15
    """最终训练loss (收敛值)"""

    final_val_loss: float = 0.25
    """最终验证loss (收敛值)"""

    noise_level: float = 0.02
    """噪声水平 (相对于loss值的比例)"""

    warmup_epochs: int = 5
    """预热阶段的epoch数，在此期间loss下降较慢"""


def _generate_base_decay_curve(
    n_epochs: int,
    initial_loss: float,
    final_loss: float,
    decay_rate: float = 0.05,
    warmup_epochs: int = 5,
) -> NDArray:
    """
    生成基础的指数衰减loss曲线

    Parameters
    ----------
    n_epochs : int
        总epoch数
    initial_loss : float
        初始loss值
    final_loss : float
        最终收敛的loss值
    decay_rate : float, optional
        衰减速率, by default 0.05
    warmup_epochs : int, optional
        预热期epoch数, by default 5

    Returns
    -------
    NDArray
        shape为(n_epochs,)的loss数组
    """
    epochs = np.arange(n_epochs)

    # 预热阶段：较慢的线性下降
    warmup_mask = epochs < warmup_epochs
    warmup_loss = np.linspace(initial_loss, initial_loss * 0.8, warmup_epochs)

    # 主要训练阶段：指数衰减
    main_epochs = epochs[~warmup_mask] - warmup_epochs
    main_initial = initial_loss * 0.8

    # 使用指数衰减公式: loss = final + (initial - final) * exp(-decay_rate * epoch)
    main_loss = final_loss + (main_initial - final_loss) * np.exp(
        -decay_rate * main_epochs
    )

    # 合并两个阶段
    loss_curve = np.zeros(n_epochs)
    loss_curve[warmup_mask] = warmup_loss
    loss_curve[~warmup_mask] = main_loss

    return loss_curve


def _add_realistic_noise(loss_curve: NDArray, noise_level: float) -> NDArray:
    """
    为loss曲线添加真实的训练噪声

    训练过程中的噪声特征：
    - 早期噪声较大（因为梯度较大）
    - 后期噪声较小（接近收敛）
    - 噪声相对于当前loss值成比例

    Parameters
    ----------
    loss_curve : NDArray
        原始loss曲线
    noise_level : float
        噪声水平（相对比例）

    Returns
    -------
    NDArray
        添加噪声后的loss曲线
    """
    n_epochs = len(loss_curve)

    # 噪声强度随训练进行而衰减
    noise_decay = np.exp(-np.arange(n_epochs) / (n_epochs * 0.3))

    # 生成相对噪声（相对于当前loss值）
    relative_noise = np.random.normal(0, noise_level, n_epochs)
    actual_noise = relative_noise * loss_curve * noise_decay

    return loss_curve + actual_noise


def generate_normal_training_curves(
    config: TrainingCurveConfig,
) -> Tuple[NDArray, NDArray]:
    """
    生成正常收敛的训练曲线

    特征：
    - 训练和验证loss都逐渐下降并趋于稳定
    - 验证loss略高于训练loss但差距合理
    - 曲线形状符合指数衰减规律

    Parameters
    ----------
    config : TrainingCurveConfig
        训练曲线配置参数

    Returns
    -------
    Tuple[NDArray, NDArray]
        (训练loss曲线, 验证loss曲线)，每个数组shape为(n_epochs,)
    """
    # 生成基础曲线
    train_loss = _generate_base_decay_curve(
        config.n_epochs,
        config.initial_train_loss,
        config.final_train_loss,
        decay_rate=0.05,
        warmup_epochs=config.warmup_epochs,
    )

    val_loss = _generate_base_decay_curve(
        config.n_epochs,
        config.initial_val_loss,
        config.final_val_loss,
        decay_rate=0.04,  # 验证loss收敛稍慢
        warmup_epochs=config.warmup_epochs,
    )

    # 添加真实噪声
    train_loss = _add_realistic_noise(train_loss, config.noise_level)
    val_loss = _add_realistic_noise(
        val_loss, config.noise_level * 1.2
    )  # 验证loss噪声稍大

    return train_loss, val_loss


def generate_overfitting_curves(
    config: TrainingCurveConfig, overfitting_start_epoch: Optional[int] = None
) -> Tuple[NDArray, NDArray]:
    """
    生成过拟合场景的训练曲线

    特征：
    - 训练loss持续下降
    - 验证loss在某个点开始上升
    - 体现模型开始记忆训练数据而失去泛化能力

    Parameters
    ----------
    config : TrainingCurveConfig
        训练曲线配置参数
    overfitting_start_epoch : Optional[int], optional
        过拟合开始的epoch，如果为None则自动选择, by default None

    Returns
    -------
    Tuple[NDArray, NDArray]
        (训练loss曲线, 验证loss曲线)
    """
    if overfitting_start_epoch is None:
        overfitting_start_epoch = int(config.n_epochs * 0.4)

    # 训练loss正常下降
    train_loss = _generate_base_decay_curve(
        config.n_epochs,
        config.initial_train_loss,
        config.final_train_loss * 0.7,  # 训练loss下降得更低
        decay_rate=0.06,
        warmup_epochs=config.warmup_epochs,
    )

    # 验证loss开始正常，然后开始上升
    val_loss_early = _generate_base_decay_curve(
        overfitting_start_epoch,
        config.initial_val_loss,
        config.final_val_loss,
        decay_rate=0.04,
        warmup_epochs=min(config.warmup_epochs, overfitting_start_epoch),
    )

    # 过拟合阶段：验证loss开始上升
    overfitting_epochs = config.n_epochs - overfitting_start_epoch
    best_val_loss = val_loss_early[-1]

    # 验证loss上升，但不是线性的，而是逐渐加速
    rise_curve = np.linspace(0, 1, overfitting_epochs) ** 1.5
    val_loss_late = best_val_loss + rise_curve * (best_val_loss * 0.8)

    val_loss = np.concatenate([val_loss_early, val_loss_late])

    # 添加噪声
    train_loss = _add_realistic_noise(train_loss, config.noise_level)
    val_loss = _add_realistic_noise(val_loss, config.noise_level * 1.3)

    return train_loss, val_loss


def generate_underfitting_curves(config: TrainingCurveConfig) -> Tuple[NDArray, NDArray]:
    """
    生成欠拟合场景的训练曲线

    特征：
    - 训练和验证loss都很高
    - 下降速度很慢
    - 最终收敛到相对较高的loss值
    - 两条曲线比较接近

    Parameters
    ----------
    config : TrainingCurveConfig
        训练曲线配置参数

    Returns
    -------
    Tuple[NDArray, NDArray]
        (训练loss曲线, 验证loss曲线)
    """
    # 欠拟合时最终loss较高
    high_final_train = config.final_train_loss * 4
    high_final_val = config.final_val_loss * 3.5

    # 生成缓慢下降的曲线
    train_loss = _generate_base_decay_curve(
        config.n_epochs,
        config.initial_train_loss,
        high_final_train,
        decay_rate=0.01,  # 很慢的衰减
        warmup_epochs=config.warmup_epochs * 2,
    )

    val_loss = _generate_base_decay_curve(
        config.n_epochs,
        config.initial_val_loss,
        high_final_val,
        decay_rate=0.008,  # 更慢的衰减
        warmup_epochs=config.warmup_epochs * 2,
    )

    # 添加较大的噪声（因为模型不稳定）
    train_loss = _add_realistic_noise(train_loss, config.noise_level * 2)
    val_loss = _add_realistic_noise(val_loss, config.noise_level * 2.2)

    return train_loss, val_loss


def generate_unstable_training_curves(
    config: TrainingCurveConfig,
    instability_type: Literal['high_lr', 'oscillating', 'diverging'] = 'high_lr',
) -> Tuple[NDArray, NDArray]:
    """
    生成不稳定的训练曲线（通常由学习率问题引起）

    Parameters
    ----------
    config : TrainingCurveConfig
        训练曲线配置参数
    instability_type : Literal["high_lr", "oscillating", "diverging"], optional
        不稳定类型, by default "high_lr"
        - "high_lr": 学习率过高导致的震荡
        - "oscillating": 持续震荡但总体下降
        - "diverging": loss发散

    Returns
    -------
    Tuple[NDArray, NDArray]
        (训练loss曲线, 验证loss曲线)
    """
    epochs = np.arange(config.n_epochs)

    if instability_type == 'high_lr':
        # 学习率过高：大幅震荡但总体趋势下降
        base_train = _generate_base_decay_curve(
            config.n_epochs,
            config.initial_train_loss,
            config.final_train_loss,
            decay_rate=0.03,
        )
        base_val = _generate_base_decay_curve(
            config.n_epochs,
            config.initial_val_loss,
            config.final_val_loss,
            decay_rate=0.025,
        )

        # 添加强烈震荡
        oscillation_amplitude = base_train * 0.3
        oscillation = (
            oscillation_amplitude
            * np.sin(epochs * 0.5)
            * np.exp(-epochs / (config.n_epochs * 0.8))
        )

        train_loss = base_train + oscillation
        val_loss = base_val + oscillation * 0.8

    elif instability_type == 'oscillating':
        # 持续震荡
        base_train = _generate_base_decay_curve(
            config.n_epochs,
            config.initial_train_loss,
            config.final_train_loss * 1.5,
            decay_rate=0.02,
        )
        base_val = _generate_base_decay_curve(
            config.n_epochs,
            config.initial_val_loss,
            config.final_val_loss * 1.3,
            decay_rate=0.018,
        )

        # 持续的周期性震荡
        oscillation = base_train * 0.15 * np.sin(epochs * 0.3)
        train_loss = base_train + oscillation
        val_loss = base_val + oscillation * 0.9

    elif instability_type == 'diverging':
        # Loss发散
        # 开始正常，然后发散
        diverge_start = int(config.n_epochs * 0.3)

        # 前期正常
        early_train = _generate_base_decay_curve(
            diverge_start,
            config.initial_train_loss,
            config.initial_train_loss * 0.6,
            decay_rate=0.05,
        )
        early_val = _generate_base_decay_curve(
            diverge_start,
            config.initial_val_loss,
            config.initial_val_loss * 0.7,
            decay_rate=0.04,
        )

        # 后期发散
        diverge_epochs = config.n_epochs - diverge_start
        diverge_train = early_train[-1] * np.exp(0.1 * np.arange(diverge_epochs))
        diverge_val = early_val[-1] * np.exp(0.08 * np.arange(diverge_epochs))

        train_loss = np.concatenate([early_train, diverge_train])
        val_loss = np.concatenate([early_val, diverge_val])

    # 添加噪声
    train_loss = _add_realistic_noise(train_loss, config.noise_level * 1.5)
    val_loss = _add_realistic_noise(val_loss, config.noise_level * 1.8)

    return train_loss, val_loss


def generate_multi_metric_curves(
    config: TrainingCurveConfig,
    curve_type: Literal['normal', 'overfitting', 'underfitting'] = 'normal',
) -> dict:
    """
    生成多指标训练曲线（loss + accuracy）

    Parameters
    ----------
    config : TrainingCurveConfig
        训练曲线配置参数
    curve_type : Literal["normal", "overfitting", "underfitting"], optional
        曲线类型, by default "normal"

    Returns
    -------
    dict
        包含以下键值对的字典：
        - "train_loss": 训练loss曲线
        - "val_loss": 验证loss曲线
        - "train_accuracy": 训练准确率曲线
        - "val_accuracy": 验证准确率曲线
    """
    # 生成loss曲线
    if curve_type == 'normal':
        train_loss, val_loss = generate_normal_training_curves(config)
    elif curve_type == 'overfitting':
        train_loss, val_loss = generate_overfitting_curves(config)
    elif curve_type == 'underfitting':
        train_loss, val_loss = generate_underfitting_curves(config)

    # 将loss转换为accuracy（假设是分类任务）
    # 使用sigmoid变换：accuracy = 1 / (1 + exp(loss - shift))
    loss_shift = 1.0  # 调整参数

    train_accuracy = 1 / (1 + np.exp(train_loss - loss_shift))
    val_accuracy = 1 / (1 + np.exp(val_loss - loss_shift))

    # 确保accuracy在合理范围内 (0.5-1.0)
    train_accuracy = 0.5 + 0.5 * train_accuracy
    val_accuracy = 0.5 + 0.5 * val_accuracy

    # 为accuracy添加独立的噪声
    train_accuracy = _add_realistic_noise(train_accuracy, config.noise_level * 0.5)
    val_accuracy = _add_realistic_noise(val_accuracy, config.noise_level * 0.6)

    # 确保accuracy不超出[0,1]范围
    train_accuracy = np.clip(train_accuracy, 0, 1)
    val_accuracy = np.clip(val_accuracy, 0, 1)

    return {
        'train_loss': train_loss,
        'val_loss': val_loss,
        'train_accuracy': train_accuracy,
        'val_accuracy': val_accuracy,
    }


def __test_training_curves():
    """演示各种训练曲线的生成效果"""
    print('=' * 60)
    print('训练曲线模拟演示')
    print('=' * 60)

    # 配置参数
    config = TrainingCurveConfig(
        n_epochs=100,
        initial_train_loss=2.5,
        initial_val_loss=2.8,
        final_train_loss=0.15,
        final_val_loss=0.25,
        noise_level=0.2,
    )
    import matplotlib.pyplot as plt

    plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False

    # 创建图形
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    fig.suptitle('训练曲线仿真效果展示', fontsize=16, fontweight='bold')

    epochs = np.arange(config.n_epochs)

    # 1. 正常收敛
    train_loss, val_loss = generate_normal_training_curves(config)
    axes[0, 0].plot(epochs, train_loss, label='Training Loss', color='blue', alpha=0.8)
    axes[0, 0].plot(epochs, val_loss, label='Validation Loss', color='red', alpha=0.8)
    axes[0, 0].set_title('正常收敛', fontweight='bold')
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')
    axes[0, 0].legend()
    axes[0, 0].grid(True, alpha=0.3)

    print(
        f'正常收敛 - 最终训练Loss: {train_loss[-1]:.4f}, 最终验证Loss: {val_loss[-1]:.4f}'
    )

    # 2. 过拟合
    train_loss, val_loss = generate_overfitting_curves(config)
    axes[0, 1].plot(epochs, train_loss, label='Training Loss', color='blue', alpha=0.8)
    axes[0, 1].plot(epochs, val_loss, label='Validation Loss', color='red', alpha=0.8)
    axes[0, 1].set_title('过拟合场景', fontweight='bold')
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('Loss')
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)

    # 找到验证loss的最低点
    min_val_epoch = np.argmin(val_loss)
    print(
        f'过拟合 - 验证Loss最低点在epoch {min_val_epoch}, 最终验证Loss上升到: {val_loss[-1]:.4f}'
    )

    # 3. 欠拟合
    train_loss, val_loss = generate_underfitting_curves(config)
    axes[0, 2].plot(epochs, train_loss, label='Training Loss', color='blue', alpha=0.8)
    axes[0, 2].plot(epochs, val_loss, label='Validation Loss', color='red', alpha=0.8)
    axes[0, 2].set_title('欠拟合场景', fontweight='bold')
    axes[0, 2].set_xlabel('Epoch')
    axes[0, 2].set_ylabel('Loss')
    axes[0, 2].legend()
    axes[0, 2].grid(True, alpha=0.3)

    print(f'欠拟合 - 最终训练Loss: {train_loss[-1]:.4f} (仍然很高)')

    # 4. 学习率过高（震荡）
    train_loss, val_loss = generate_unstable_training_curves(config, 'high_lr')
    axes[1, 0].plot(epochs, train_loss, label='Training Loss', color='blue', alpha=0.8)
    axes[1, 0].plot(epochs, val_loss, label='Validation Loss', color='red', alpha=0.8)
    axes[1, 0].set_title('学习率过高（震荡）', fontweight='bold')
    axes[1, 0].set_xlabel('Epoch')
    axes[1, 0].set_ylabel('Loss')
    axes[1, 0].legend()
    axes[1, 0].grid(True, alpha=0.3)

    print(
        f'震荡训练 - Loss标准差: 训练{np.std(train_loss):.4f}, 验证{np.std(val_loss):.4f}'
    )

    # 5. Loss发散
    train_loss, val_loss = generate_unstable_training_curves(config, 'diverging')
    axes[1, 1].plot(epochs, train_loss, label='Training Loss', color='blue', alpha=0.8)
    axes[1, 1].plot(epochs, val_loss, label='Validation Loss', color='red', alpha=0.8)
    axes[1, 1].set_title('Loss发散', fontweight='bold')
    axes[1, 1].set_xlabel('Epoch')
    axes[1, 1].set_ylabel('Loss')
    axes[1, 1].legend()
    axes[1, 1].grid(True, alpha=0.3)
    axes[1, 1].set_yscale('log')  # 使用对数坐标

    print(f'发散训练 - 最终Loss达到: 训练{train_loss[-1]:.2f}, 验证{val_loss[-1]:.2f}')

    # 6. 多指标曲线（包含accuracy）
    metrics = generate_multi_metric_curves(config, 'normal')
    ax_twin = axes[1, 2].twinx()

    # 绘制loss（左y轴）
    line1 = axes[1, 2].plot(
        epochs, metrics['train_loss'], 'b-', alpha=0.8, label='Train Loss'
    )
    line2 = axes[1, 2].plot(
        epochs, metrics['val_loss'], 'r-', alpha=0.8, label='Val Loss'
    )
    axes[1, 2].set_xlabel('Epoch')
    axes[1, 2].set_ylabel('Loss', color='black')
    axes[1, 2].tick_params(axis='y', labelcolor='black')

    # 绘制accuracy（右y轴）
    line3 = ax_twin.plot(
        epochs, metrics['train_accuracy'], 'b--', alpha=0.8, label='Train Acc'
    )
    line4 = ax_twin.plot(
        epochs, metrics['val_accuracy'], 'r--', alpha=0.8, label='Val Acc'
    )
    ax_twin.set_ylabel('Accuracy', color='green')
    ax_twin.tick_params(axis='y', labelcolor='green')

    # 合并图例
    lines = line1 + line2 + line3 + line4
    labels = [l.get_label() for l in lines]
    axes[1, 2].legend(lines, labels, loc='center right')
    axes[1, 2].set_title('Loss + Accuracy曲线', fontweight='bold')
    axes[1, 2].grid(True, alpha=0.3)

    final_acc = metrics['val_accuracy'][-1]
    print(f'多指标 - 最终验证准确率: {final_acc:.4f}')

    plt.tight_layout()
    plt.show()

    return fig


if __name__ == '__main__':
    __test_training_curves()
