# 文件: zyplib\artifact\classification_metrics.py
"""
分类指标模拟模块; 用于教学模拟用

该模块用于生成逼真的分类性能指标，包括：
- 混淆矩阵
- 精确率、召回率、F1分数等指标
- ROC曲线和AUC值
- PR曲线和PR-AUC值
- 多类别分类的各种指标


原理是先生成一个狄利克雷分布, 然后调整为概率分布, 然后在在概率分布的基础上计算指标

所有生成的指标都确保在统计上的一致性和真实性。
"""

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

import numpy as np
from numpy.typing import NDArray
from scipy import stats


@dataclass
class ClassificationConfig:
    """分类指标配置参数"""

    n_samples: int = 1000
    """样本总数"""

    n_classes: int = 2
    """类别数量"""

    class_imbalance_ratio: Optional[List[float]] = None
    """类别不平衡比例，如果为None则使用平衡数据集"""

    target_accuracy: float = 0.8
    """目标总体准确率"""

    class_difficulty: Optional[List[float]] = None
    """各类别的分类难度 (0-1，越高越难分类)，如果为None则使用相同难度

    一个类别的 difficulty 值越高，该类别能达到的实际准确率就越低
    """

    confidence_separation: float = 0.9
    """正确和错误预测的置信度分离程度 (0-1)
    """


def _generate_class_distribution(
    n_samples: int, n_classes: int, imbalance_ratio: Optional[List[float]] = None
) -> NDArray:
    """
    生成类别分布

    Parameters
    ----------
    n_samples : int
        总样本数
    n_classes : int
        类别数
    imbalance_ratio : Optional[List[float]], optional
        类别比例，如果为None则平衡分布

    Returns
    -------
    NDArray
        真实标签数组，shape为(n_samples,)
    """
    if imbalance_ratio is None:
        # 平衡分布
        class_sizes = [n_samples // n_classes] * n_classes
        # 处理不能整除的情况
        remainder = n_samples % n_classes
        for i in range(remainder):
            class_sizes[i] += 1
    else:
        # 按照指定比例分布
        if len(imbalance_ratio) != n_classes:
            raise ValueError(
                f'imbalance_ratio长度({len(imbalance_ratio)})必须等于类别数({n_classes})'
            )

        # 归一化比例
        total_ratio = sum(imbalance_ratio)
        normalized_ratios = [r / total_ratio for r in imbalance_ratio]

        class_sizes = [int(n_samples * ratio) for ratio in normalized_ratios]

        # 确保总数正确
        diff = n_samples - sum(class_sizes)
        class_sizes[0] += diff

    # 生成标签
    y_true = []
    for class_idx, size in enumerate(class_sizes):
        y_true.extend([class_idx] * size)

    return np.array(y_true)


def _generate_confusion_matrix(
    y_true: NDArray,
    target_accuracy: float,
    class_difficulty: Optional[List[float]] = None,
) -> NDArray:
    """
    基于真实标签和目标准确率生成逼真的混淆矩阵

    Parameters
    ----------
    y_true : NDArray
        真实标签
    target_accuracy : float
        目标总体准确率
    class_difficulty : Optional[List[float]], optional
        各类别分类难度

    Returns
    -------
    NDArray
        混淆矩阵，shape为(n_classes, n_classes)
    """
    n_classes = len(np.unique(y_true))
    n_samples = len(y_true)

    if class_difficulty is None:
        class_difficulty = [0.3] * n_classes  # 默认中等难度

    # 计算每个类别的样本数
    class_counts = np.bincount(y_true, minlength=n_classes)

    # 初始化混淆矩阵
    confusion_matrix = np.zeros((n_classes, n_classes), dtype=int)

    for true_class in range(n_classes):
        class_size = class_counts[true_class]
        if class_size == 0:
            continue

        # 根据难度计算该类别的准确率
        # 难度越高，准确率相对于目标准确率的偏差越大
        difficulty = class_difficulty[true_class]
        class_accuracy = target_accuracy * (1 - difficulty * 0.5)
        class_accuracy = np.clip(class_accuracy, 0.1, 0.95)

        # 正确分类的样本数
        correct_count = int(class_size * class_accuracy)

        # 错误分类的样本数
        incorrect_count = class_size - correct_count

        # 分配正确分类
        confusion_matrix[true_class, true_class] = correct_count

        # 分配错误分类
        if incorrect_count > 0:
            # 错误分类倾向于分配给"相似"的类别
            # 使用距离权重：距离越近的类别越容易被混淆
            other_classes = [i for i in range(n_classes) if i != true_class]

            if len(other_classes) == 1:
                # 二分类情况
                confusion_matrix[true_class, other_classes[0]] = incorrect_count
            else:
                # 多分类情况：使用权重分配错误
                # 创建相似度权重（这里简单地使用类别编号距离的倒数）
                weights = []
                for other_class in other_classes:
                    distance = abs(true_class - other_class)
                    weight = 1 / (distance + 0.5)  # 距离越近权重越大
                    weights.append(weight)

                # 归一化权重
                weights = np.array(weights)
                weights = weights / weights.sum()

                # 根据权重分配错误分类
                error_distribution = np.random.multinomial(incorrect_count, weights)

                for i, other_class in enumerate(other_classes):
                    confusion_matrix[true_class, other_class] = error_distribution[i]

    return confusion_matrix


def _generate_confidence_scores(
    y_true: NDArray, y_pred: NDArray, confidence_separation: float = 0.6
) -> NDArray:
    """
    基于预测结果生成逼真的置信度分数

    Parameters
    ----------
    y_true : NDArray
        真实标签
    y_pred : NDArray
        预测标签
    confidence_separation : float
        正确和错误预测的置信度分离程度

    Returns
    -------
    NDArray
        置信度分数，shape为(n_samples,)
    """
    n_samples = len(y_true)
    confidences = np.zeros(n_samples)

    # 正确预测的置信度分布
    correct_mask = y_true == y_pred
    n_correct = np.sum(correct_mask)

    if n_correct > 0:
        # 正确预测：高置信度，使用Beta分布偏向高值
        alpha = 2 + confidence_separation * 8  # 形状参数，控制分布偏向
        beta = 2
        correct_confidences = stats.beta.rvs(alpha, beta, size=n_correct)
        confidences[correct_mask] = correct_confidences

    # 错误预测的置信度分布
    incorrect_mask = ~correct_mask
    n_incorrect = np.sum(incorrect_mask)

    if n_incorrect > 0:
        # 错误预测：低置信度，使用Beta分布偏向低值
        alpha = 2
        beta = 2 + confidence_separation * 8
        incorrect_confidences = stats.beta.rvs(alpha, beta, size=n_incorrect)
        confidences[incorrect_mask] = incorrect_confidences

    return confidences


def generate_binary_classification_results(
    config: ClassificationConfig,
) -> Dict[str, NDArray]:
    """
    生成二分类结果和相关指标

    Parameters
    ----------
    config : ClassificationConfig
        分类配置参数

    Returns
    -------
    Dict[str, NDArray]
        包含以下键的字典：
        - "y_true": 真实标签
        - "y_pred": 预测标签
        - "y_proba": 预测概率
        - "confusion_matrix": 混淆矩阵
    """
    if config.n_classes != 2:
        raise ValueError('此函数仅支持二分类')

    # 生成真实标签
    y_true = _generate_class_distribution(
        config.n_samples, config.n_classes, config.class_imbalance_ratio
    )

    # 生成混淆矩阵
    confusion_matrix = _generate_confusion_matrix(
        y_true, config.target_accuracy, config.class_difficulty
    )

    # 从混淆矩阵恢复预测标签
    y_pred = []
    for true_label in range(config.n_classes):
        true_indices = np.where(y_true == true_label)[0]

        # 根据混淆矩阵分配预测标签
        for pred_label in range(config.n_classes):
            count = confusion_matrix[true_label, pred_label]
            if count > 0:
                # 随机选择该数量的样本分配给这个预测标签
                selected_indices = np.random.choice(
                    true_indices, size=min(count, len(true_indices)), replace=False
                )
                y_pred.extend([(idx, pred_label) for idx in selected_indices])
                # 从可选索引中移除已选择的
                true_indices = np.setdiff1d(true_indices, selected_indices)

    # 创建完整的预测数组
    y_pred_array = np.zeros(config.n_samples, dtype=int)
    for idx, pred_label in y_pred:
        y_pred_array[idx] = pred_label

    # 生成置信度分数
    confidences = _generate_confidence_scores(
        y_true, y_pred_array, config.confidence_separation
    )

    # 将置信度转换为概率
    # 对于二分类，y_proba[:, 1]表示正类的概率
    y_proba = np.zeros((config.n_samples, 2))
    y_proba[:, 1] = confidences  # 正类概率
    y_proba[:, 0] = 1 - confidences  # 负类概率

    # 根据预测标签调整概率
    # 如果预测为负类但置信度高，需要反转概率
    negative_pred_mask = y_pred_array == 0
    y_proba[negative_pred_mask, 1] = 1 - confidences[negative_pred_mask]
    y_proba[negative_pred_mask, 0] = confidences[negative_pred_mask]

    return {
        'y_true': y_true,
        'y_pred': y_pred_array,
        'y_proba': y_proba,
        'confusion_matrix': confusion_matrix,
    }


def calculate_classification_metrics(
    y_true: NDArray, y_pred: NDArray, y_proba: Optional[NDArray] = None
) -> Dict[str, float]:
    """
    计算分类指标

    Parameters
    ----------
    y_true : NDArray
        真实标签
    y_pred : NDArray
        预测标签
    y_proba : Optional[NDArray], optional
        预测概率，用于计算AUC等指标

    Returns
    -------
    Dict[str, float]
        各种分类指标
    """
    from sklearn.metrics import (
        accuracy_score,
        average_precision_score,
        confusion_matrix,
        f1_score,
        precision_score,
        recall_score,
        roc_auc_score,
    )

    metrics = {}

    # 基本指标
    metrics['accuracy'] = accuracy_score(y_true, y_pred)

    # 对于二分类
    if len(np.unique(y_true)) == 2:
        metrics['precision'] = precision_score(y_true, y_pred)
        metrics['recall'] = recall_score(y_true, y_pred)
        metrics['f1_score'] = f1_score(y_true, y_pred)
        metrics['specificity'] = recall_score(y_true, y_pred, pos_label=0)

        if y_proba is not None:
            metrics['roc_auc'] = roc_auc_score(y_true, y_proba[:, 1])
            metrics['pr_auc'] = average_precision_score(y_true, y_proba[:, 1])
    else:
        # 多分类
        metrics['precision_macro'] = precision_score(y_true, y_pred, average='macro')
        metrics['recall_macro'] = recall_score(y_true, y_pred, average='macro')
        metrics['f1_score_macro'] = f1_score(y_true, y_pred, average='macro')

        if y_proba is not None:
            metrics['roc_auc_ovr'] = roc_auc_score(y_true, y_proba, multi_class='ovr')

    return metrics


def generate_roc_curve_data(
    y_true: NDArray, y_proba: NDArray, n_thresholds: int = 100
) -> Dict[str, NDArray]:
    """
    生成ROC曲线数据

    Parameters
    ----------
    y_true : NDArray
        真实标签（二分类）
    y_proba : NDArray
        正类概率
    n_thresholds : int, optional
        阈值数量, by default 100

    Returns
    -------
    Dict[str, NDArray]
        包含fpr, tpr, thresholds的字典
    """
    if len(np.unique(y_true)) != 2:
        raise ValueError('ROC曲线仅适用于二分类问题')

    # 生成阈值
    thresholds = np.linspace(0, 1, n_thresholds)

    fprs = []
    tprs = []

    for threshold in thresholds:
        # 根据阈值生成预测
        y_pred = (y_proba >= threshold).astype(int)

        # 计算混淆矩阵
        tn = np.sum((y_true == 0) & (y_pred == 0))
        fp = np.sum((y_true == 0) & (y_pred == 1))
        fn = np.sum((y_true == 1) & (y_pred == 0))
        tp = np.sum((y_true == 1) & (y_pred == 1))

        # 计算TPR和FPR
        tpr = tp / (tp + fn) if (tp + fn) > 0 else 0
        fpr = fp / (fp + tn) if (fp + tn) > 0 else 0

        tprs.append(tpr)
        fprs.append(fpr)

    return {'fpr': np.array(fprs), 'tpr': np.array(tprs), 'thresholds': thresholds}


def generate_multiclass_classification_results(
    config: ClassificationConfig,
) -> Dict[str, NDArray]:
    """
    生成多类别分类结果

    Parameters
    ----------
    config : ClassificationConfig
        分类配置参数

    Returns
    -------
    Dict[str, NDArray]
        多类别分类结果
    """
    if config.n_classes < 3:
        raise ValueError('此函数用于多类别分类（>=3类）')

    # 生成真实标签
    y_true = _generate_class_distribution(
        config.n_samples, config.n_classes, config.class_imbalance_ratio
    )

    # 生成混淆矩阵
    confusion_matrix = _generate_confusion_matrix(
        y_true, config.target_accuracy, config.class_difficulty
    )

    # 从混淆矩阵恢复预测标签
    y_pred = np.zeros(config.n_samples, dtype=int)
    y_proba = np.zeros((config.n_samples, config.n_classes))

    sample_idx = 0
    for true_label in range(config.n_classes):
        true_mask = y_true == true_label
        true_indices = np.where(true_mask)[0]

        for pred_label in range(config.n_classes):
            count = confusion_matrix[true_label, pred_label]
            if count > 0:
                # 分配预测标签
                selected_indices = true_indices[:count]
                y_pred[selected_indices] = pred_label

                # 生成概率分布
                for idx in selected_indices:
                    # 为预测类别分配较高概率
                    proba = np.random.dirichlet([1] * config.n_classes)
                    # 增强预测类别的概率
                    proba[pred_label] *= 2 + config.confidence_separation
                    proba = proba / proba.sum()  # 重新归一化
                    y_proba[idx] = proba

                true_indices = true_indices[count:]

    return {
        'y_true': y_true,
        'y_pred': y_pred,
        'y_proba': y_proba,
        'confusion_matrix': confusion_matrix,
    }


def __test_classification_metrics():
    """演示分类指标的生成效果"""
    print('\n' + '=' * 60)
    print('分类指标模拟演示')
    print('=' * 60)

    # 二分类示例
    print('\n1. 二分类示例')
    print('-' * 30)

    import matplotlib.pyplot as plt

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

    binary_config = ClassificationConfig(
        n_samples=1000,
        n_classes=2,
        class_imbalance_ratio=[0.6, 0.4],  # 类别不平衡
        target_accuracy=0.82,
        class_difficulty=[0.2, 0.3],  # 第二类更难分类
        confidence_separation=0.94,
    )

    binary_results = generate_binary_classification_results(binary_config)

    print(f"样本分布: 类别0: {np.sum(binary_results['y_true'] == 0)} 个")
    print(f"          类别1: {np.sum(binary_results['y_true'] == 1)} 个")
    print(f"混淆矩阵:\n{binary_results['confusion_matrix']}")

    # 计算指标
    metrics = calculate_classification_metrics(
        binary_results['y_true'], binary_results['y_pred'], binary_results['y_proba']
    )

    print(f'\n分类性能指标:')
    for metric_name, value in metrics.items():
        print(f'  {metric_name}: {value:.4f}')

    # 生成ROC曲线数据
    roc_data = generate_roc_curve_data(
        binary_results['y_true'], binary_results['y_proba'][:, 1]
    )

    # 多分类示例
    print(f'\n2. 多分类示例（5类）')
    print('-' * 30)

    multi_config = ClassificationConfig(
        n_samples=1500,
        n_classes=5,
        class_imbalance_ratio=[0.3, 0.25, 0.2, 0.15, 0.1],  # 递减的类别分布
        target_accuracy=0.78,
        class_difficulty=[0.1, 0.2, 0.3, 0.4, 0.5],  # 递增的分类难度
        confidence_separation=0.5,
    )

    multi_results = generate_multiclass_classification_results(multi_config)

    print(f'样本分布:')
    for i in range(5):
        count = np.sum(multi_results['y_true'] == i)
        print(f'  类别{i}: {count} 个')

    print(f'\n混淆矩阵:')
    print(multi_results['confusion_matrix'])

    # 计算多分类指标
    multi_metrics = calculate_classification_metrics(
        multi_results['y_true'], multi_results['y_pred'], multi_results['y_proba']
    )

    print(f'\n多分类性能指标:')
    for metric_name, value in multi_metrics.items():
        print(f'  {metric_name}: {value:.4f}')

    # 创建可视化
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    fig.suptitle('分类指标仿真效果展示', fontsize=16, fontweight='bold')

    # 1. 二分类混淆矩阵 - 修复：手动添加文本标注
    im1 = axes[0, 0].imshow(binary_results['confusion_matrix'], cmap='Blues')
    axes[0, 0].set_title('二分类混淆矩阵')
    axes[0, 0].set_xlabel('预测标签')
    axes[0, 0].set_ylabel('真实标签')
    axes[0, 0].set_xticks([0, 1])
    axes[0, 0].set_yticks([0, 1])

    # 手动添加数值标注
    for i in range(2):
        for j in range(2):
            axes[0, 0].text(
                j,
                i,
                str(binary_results['confusion_matrix'][i, j]),
                ha='center',
                va='center',
                fontweight='bold',
                color='white'
                if binary_results['confusion_matrix'][i, j] > 200
                else 'black',
            )

    # 2. ROC曲线
    axes[0, 1].plot(
        roc_data['fpr'],
        roc_data['tpr'],
        'b-',
        linewidth=2,
        label=f'ROC Curve (AUC = {metrics["roc_auc"]:.3f})',
    )
    axes[0, 1].plot([0, 1], [0, 1], 'r--', alpha=0.5, label='Random Classifier')
    axes[0, 1].set_xlabel('False Positive Rate')
    axes[0, 1].set_ylabel('True Positive Rate')
    axes[0, 1].set_title('ROC曲线')
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)

    # 3. 置信度分布
    correct_mask = binary_results['y_true'] == binary_results['y_pred']
    correct_conf = binary_results['y_proba'][correct_mask, 1]
    incorrect_conf = binary_results['y_proba'][~correct_mask, 1]

    axes[0, 2].hist(
        correct_conf, bins=30, alpha=0.7, label='正确预测', color='green', density=True
    )
    axes[0, 2].hist(
        incorrect_conf, bins=30, alpha=0.7, label='错误预测', color='red', density=True
    )
    axes[0, 2].set_xlabel('置信度')
    axes[0, 2].set_ylabel('密度')
    axes[0, 2].set_title('置信度分布')
    axes[0, 2].legend()
    axes[0, 2].grid(True, alpha=0.3)

    print(f'\n置信度分析:')
    print(f'  正确预测平均置信度: {np.mean(correct_conf):.4f}')
    print(f'  错误预测平均置信度: {np.mean(incorrect_conf):.4f}')
    print(f'  置信度分离度: {np.mean(correct_conf) - np.mean(incorrect_conf):.4f}')

    # 4. 多分类混淆矩阵
    im2 = axes[1, 0].imshow(multi_results['confusion_matrix'], cmap='Blues')
    axes[1, 0].set_title('多分类混淆矩阵 (5类)')
    axes[1, 0].set_xlabel('预测标签')
    axes[1, 0].set_ylabel('真实标签')
    axes[1, 0].set_xticks(range(5))
    axes[1, 0].set_yticks(range(5))

    # 添加colorbar
    plt.colorbar(im2, ax=axes[1, 0])

    # 5. 各类别准确率
    class_accuracies = []
    for i in range(5):
        class_mask = multi_results['y_true'] == i
        class_correct = np.sum(
            (multi_results['y_true'] == i) & (multi_results['y_pred'] == i)
        )
        class_total = np.sum(class_mask)
        class_acc = class_correct / class_total if class_total > 0 else 0
        class_accuracies.append(class_acc)

    axes[1, 1].bar(
        range(5),
        class_accuracies,
        color=['skyblue', 'lightgreen', 'lightcoral', 'gold', 'plum'],
    )
    axes[1, 1].set_xlabel('类别')
    axes[1, 1].set_ylabel('准确率')
    axes[1, 1].set_title('各类别准确率')
    axes[1, 1].set_xticks(range(5))
    axes[1, 1].grid(True, alpha=0.3, axis='y')

    # 添加数值标注
    for i, acc in enumerate(class_accuracies):
        axes[1, 1].text(i, acc + 0.01, f'{acc:.3f}', ha='center', fontweight='bold')

    print(f'\n各类别准确率:')
    for i, acc in enumerate(class_accuracies):
        print(f'  类别{i}: {acc:.4f}')

    # 6. 预测概率分布
    # 选择第一个类别的概率分布
    class_0_proba = multi_results['y_proba'][:, 0]
    true_class_0 = multi_results['y_true'] == 0
    pred_class_0 = multi_results['y_pred'] == 0

    # 真正例和假正例的概率分布
    true_positive_proba = class_0_proba[true_class_0 & pred_class_0]
    false_positive_proba = class_0_proba[(~true_class_0) & pred_class_0]

    if len(true_positive_proba) > 0:
        axes[1, 2].hist(
            true_positive_proba,
            bins=20,
            alpha=0.7,
            label='真正例',
            color='green',
            density=True,
        )
    if len(false_positive_proba) > 0:
        axes[1, 2].hist(
            false_positive_proba,
            bins=20,
            alpha=0.7,
            label='假正例',
            color='red',
            density=True,
        )

    axes[1, 2].set_xlabel('类别0的预测概率')
    axes[1, 2].set_ylabel('密度')
    axes[1, 2].set_title('类别0预测概率分布')
    axes[1, 2].legend()
    axes[1, 2].grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()

    return fig, binary_results, multi_results


if __name__ == '__main__':
    __test_classification_metrics()
