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

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


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

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

import typing

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[typing.Sequence[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[typing.Sequence[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()
