"""
封装一些机器学习常用的绘图，方便直接函数式地调用

### 函数通用定义规范

```py
def plot_xxx(..., ax=None, show=True, save_path=None, auto_close=None, **kwargs):
    ...
    return fig, ax
```

- 如果给定了 ax，则在其上绘制；否则创建新图表。
- 返回 fig, ax 对象，方便后续操作。
- `show`: 是否调用 plt.show()。
- `save_path`: 保存路径，如果提供则保存图表。
- `auto_close`: 智能关闭图表以释放内存。
- `**kwargs` 用于传递给底层的绘图函数 (如 ax.plot, sns.heatmap)。
- 尽量和样式配置分离，不参与样式的问题
"""

import typing

from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Tuple, Union

import matplotlib.pyplot as plt
import numpy as np

from zyplib.utils.ensure import ensure_ndarray

from ._plot_utils import FigAx, _finalize_figure, _prepare_ax
from .mpl_utils import draw_comparison_lines, draw_error_regions
from .mpl_utils.draw import ComparisonDict

if TYPE_CHECKING:
    from matplotlib.axes import Axes
    from matplotlib.figure import Figure
    from numpy.typing import NDArray
    from sklearn.decomposition import PCA
    from sklearn.manifold import TSNE
    from torch import Tensor
    # from sklearn


def _set_zlabel(ax: 'Axes', label: str) -> None:
    typing.cast(typing.Any, ax).set_zlabel(label)


def plot_pca(
    data: Union['NDArray', 'Tensor'],
    n_components: Literal[1, 2, 3] = 2,
    labels: Optional[Union['NDArray', 'Tensor']] = None,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    zlabel: Optional[str] = None,
    alpha: float = 0.7,
    pca_model: Optional['PCA'] = None,
    ax: Optional['Axes'] = None,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    **scatter_kwargs,
) -> FigAx:
    """PCA降维可视化。

    对输入数据进行PCA降维，并生成1D、2D或3D可视化图。

    Parameters
    ----------
    - `data` : `Union['NDArray', 'Tensor']`
        - 输入数据，形状为 [n_samples, n_features]
    - `n_components` : `Literal[1, 2, 3]`, optional
        - PCA降维的目标维数, by default 2
    - `labels` : `Optional[Union['NDArray', 'Tensor']]`, optional
        - 数据标签，用于着色和图例, by default None
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `xlabel` : `Optional[str]`, optional
        - X轴标签, by default None
    - `ylabel` : `Optional[str]`, optional
        - Y轴标签, by default None
    - `zlabel` : `Optional[str]`, optional
        - Z轴标签（仅3D时使用）, by default None
    - `alpha` : `float`, optional
        - 数据点透明度, by default 0.7
    - `pca_model` : `Optional['PCA']`, optional
        - 预训练的PCA模型，如果提供则使用该模型进行变换, by default None
    - `ax` : `Optional['Axes']`, optional
        - 绘图轴对象，如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径，如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组
    """
    import matplotlib.pyplot as plt
    from sklearn.decomposition import PCA

    # 确保数据为numpy数组
    data = ensure_ndarray(data)
    if labels is not None:
        labels = ensure_ndarray(labels)

    # 执行PCA降维
    if pca_model is not None:
        # 使用提供的PCA模型
        pca = pca_model
        data_pca = pca.transform(data)
    else:
        # 创建新的PCA模型
        pca = PCA(n_components=n_components)
        data_pca = pca.fit_transform(data)

    # 创建图表
    if n_components == 3:
        if ax is None:
            fig = plt.figure(figsize=scatter_kwargs.get('figsize', (10, 8)))
            ax = fig.add_subplot(111, projection='3d')
        else:
            fig = ax.figure
    else:
        fig, ax = _prepare_ax(ax, **scatter_kwargs)

    # 绘制数据点
    if labels is not None:
        unique_labels = np.unique(labels)
        for label in unique_labels:
            mask = labels == label
            if n_components == 1:
                ax.scatter(
                    data_pca[mask, 0],
                    np.zeros_like(data_pca[mask, 0]),
                    label=f'Class {label}',
                    alpha=alpha,
                )
            elif n_components == 2:
                ax.scatter(
                    data_pca[mask, 0],
                    data_pca[mask, 1],
                    label=f'Class {label}',
                    alpha=alpha,
                )
            elif n_components == 3:
                ax.scatter(
                    data_pca[mask, 0],
                    data_pca[mask, 1],
                    data_pca[mask, 2],
                    label=f'Class {label}',
                    alpha=alpha,
                )
        ax.legend()
    else:
        if n_components == 1:
            ax.scatter(data_pca[:, 0], np.zeros_like(data_pca[:, 0]), alpha=alpha)
        elif n_components == 2:
            ax.scatter(data_pca[:, 0], data_pca[:, 1], alpha=alpha)
        elif n_components == 3:
            ax.scatter(data_pca[:, 0], data_pca[:, 1], data_pca[:, 2], alpha=alpha)

    # 设置标签
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    elif n_components >= 1:
        ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.2%} var)')

    if ylabel is not None:
        ax.set_ylabel(ylabel)
    elif n_components >= 2:
        ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.2%} var)')

    if zlabel is not None and n_components == 3:
        _set_zlabel(ax, zlabel)
    elif n_components == 3:
        _set_zlabel(ax, f'PC3 ({pca.explained_variance_ratio_[2]:.2%} var)')

    if title is not None:
        ax.set_title(title)

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax


def plot_tsne(
    data: Union['NDArray', 'Tensor'],
    labels: Optional[Union['NDArray', 'Tensor']] = None,
    n_components: Literal[2, 3] = 2,
    perplexity: float = 30.0,
    max_iter: int = 1000,
    random_state: Optional[int] = None,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    zlabel: Optional[str] = None,
    alpha: float = 0.7,
    tsne_model: Optional['TSNE'] = None,
    ax: Optional['Axes'] = None,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    **scatter_kwargs,
) -> FigAx:
    """t-SNE降维可视化。

    对输入数据进行t-SNE降维，并生成2D或3D可视化图。

    Parameters
    ----------
    - `data` : `Union['NDArray', 'Tensor']`
        - 输入数据，形状为 [n_samples, n_features]
    - `labels` : `Optional[Union['NDArray', 'Tensor']]`, optional
        - 数据标签，用于着色和图例, by default None
    - `n_components` : `Literal[2, 3]`, optional
        - t-SNE降维的目标维数, by default 2
    - `perplexity` : `float`, optional
        - t-SNE的perplexity参数, by default 30.0
    - `max_iter` : `int`, optional
        - t-SNE的最大迭代次数, by default 1000
    - `random_state` : `Optional[int]`, optional
        - 随机种子, by default None
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `xlabel` : `Optional[str]`, optional
        - X轴标签, by default None
    - `ylabel` : `Optional[str]`, optional
        - Y轴标签, by default None
    - `zlabel` : `Optional[str]`, optional
        - Z轴标签（仅3D时使用）, by default None
    - `alpha` : `float`, optional
        - 数据点透明度, by default 0.7
    - `tsne_model` : `Optional['TSNE']`, optional
        - 预训练的t-SNE模型，如果提供则使用该模型进行变换, by default None
    - `ax` : `Optional['Axes']`, optional
        - 绘图轴对象，如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径，如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组
    """
    import matplotlib.pyplot as plt
    from sklearn.manifold import TSNE

    # 确保数据为numpy数组
    data = ensure_ndarray(data)
    if labels is not None:
        labels = ensure_ndarray(labels)

    # 执行t-SNE降维
    if tsne_model is not None:
        # 使用提供的t-SNE模型
        tsne = tsne_model
        data_tsne = tsne.fit_transform(data)
    else:
        # 创建新的t-SNE模型
        tsne = TSNE(
            n_components=n_components,
            perplexity=perplexity,
            max_iter=max_iter,
            random_state=random_state,
        )
        data_tsne = tsne.fit_transform(data)

    # 创建图表
    if n_components == 3:
        if ax is None:
            fig = plt.figure(figsize=scatter_kwargs.get('figsize', (10, 8)))
            ax = fig.add_subplot(111, projection='3d')
        else:
            fig = ax.figure
    else:
        fig, ax = _prepare_ax(ax, **scatter_kwargs)

    # 绘制数据点
    if labels is not None:
        unique_labels = np.unique(labels)
        for label in unique_labels:
            mask = labels == label
            if n_components == 2:
                ax.scatter(
                    data_tsne[mask, 0],
                    data_tsne[mask, 1],
                    label=f'Class {label}',
                    alpha=alpha,
                )
            elif n_components == 3:
                ax.scatter(
                    data_tsne[mask, 0],
                    data_tsne[mask, 1],
                    data_tsne[mask, 2],
                    label=f'Class {label}',
                    alpha=alpha,
                )
        ax.legend()
    else:
        if n_components == 2:
            ax.scatter(data_tsne[:, 0], data_tsne[:, 1], alpha=alpha)
        elif n_components == 3:
            ax.scatter(data_tsne[:, 0], data_tsne[:, 1], data_tsne[:, 2], alpha=alpha)

    # 设置标签
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    elif n_components >= 1:
        ax.set_xlabel('t-SNE 1')

    if ylabel is not None:
        ax.set_ylabel(ylabel)
    elif n_components >= 2:
        ax.set_ylabel('t-SNE 2')

    if zlabel is not None and n_components == 3:
        _set_zlabel(ax, zlabel)
    elif n_components == 3:
        _set_zlabel(ax, 't-SNE 3')

    if title is not None:
        ax.set_title(title)

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax


def plot_confusion_matrix(
    y_true: Union['NDArray', 'Tensor'],
    y_pred: Union['NDArray', 'Tensor'],
    labels: Optional[typing.Sequence[str]] = None,
    normalize: Optional[Literal['true', 'pred', 'all']] = None,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    cmap: str = 'Blues',
    ax: Optional['Axes'] = None,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    **heatmap_kwargs,
) -> FigAx:
    """绘制混淆矩阵。

    Parameters
    ----------
    - `y_true` : `Union['NDArray', 'Tensor']`
        - 真实标签
    - `y_pred` : `Union['NDArray', 'Tensor']`
        - 预测标签
    - `labels` : `Optional[List[str]]`, optional
        - 类别标签名称, by default None
    - `normalize` : `Optional[Literal['true', 'pred', 'all']]`, optional
        - 归一化方式：'true'按真实标签归一化，'pred'按预测标签归一化，'all'按全部归一化, by default None
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `xlabel` : `Optional[str]`, optional
        - X轴标签, by default None
    - `ylabel` : `Optional[str]`, optional
        - Y轴标签, by default None
    - `cmap` : `str`, optional
        - 颜色映射, by default 'Blues'
    - `ax` : `Optional['Axes']`, optional
        - 绘图轴对象，如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径，如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组
    """
    import seaborn as sns
    from sklearn.metrics import confusion_matrix

    # 确保数据为numpy数组
    y_true = ensure_ndarray(y_true)
    y_pred = ensure_ndarray(y_pred)

    # 计算混淆矩阵
    cm = confusion_matrix(y_true, y_pred)

    # 归一化
    if normalize == 'true':
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        fmt = '.2f'
    elif normalize == 'pred':
        cm = cm.astype('float') / cm.sum(axis=0)
        fmt = '.2f'
    elif normalize == 'all':
        cm = cm.astype('float') / cm.sum()
        fmt = '.2f'
    else:
        fmt = 'd'

    # 创建图表
    fig, ax = _prepare_ax(ax, **heatmap_kwargs)

    # 绘制热力图
    sns.heatmap(
        cm,
        annot=True,
        fmt=fmt,
        cmap=cmap,
        xticklabels=labels if labels is not None else True,
        yticklabels=labels if labels is not None else True,
        ax=ax,
        **heatmap_kwargs,
    )

    # 设置标签
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    else:
        ax.set_xlabel('Predicted Label')

    if ylabel is not None:
        ax.set_ylabel(ylabel)
    else:
        ax.set_ylabel('True Label')

    if title is not None:
        ax.set_title(title)
    else:
        ax.set_title('Confusion Matrix')

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax


def plot_classification_report(
    y_true: Union['NDArray', 'Tensor'],
    y_pred: Union['NDArray', 'Tensor'],
    labels: Optional[typing.Sequence[str]] = None,
    title: Optional[str] = None,
    cmap: str = 'Blues',
    ax: Optional['Axes'] = None,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    **heatmap_kwargs,
) -> FigAx:
    """绘制分类报告热力图。

    Parameters
    ----------
    - `y_true` : `Union['NDArray', 'Tensor']`
        - 真实标签
    - `y_pred` : `Union['NDArray', 'Tensor']`
        - 预测标签
    - `labels` : `Optional[List[str]]`, optional
        - 类别标签名称, by default None
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `cmap` : `str`, optional
        - 颜色映射, by default 'Blues'
    - `ax` : `Optional['Axes']`, optional
        - 绘图轴对象，如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径，如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组
    """
    import pandas as pd
    import seaborn as sns
    from sklearn.metrics import classification_report

    # 确保数据为numpy数组
    y_true = ensure_ndarray(y_true)
    y_pred = ensure_ndarray(y_pred)

    # 获取分类报告
    report = classification_report(y_true, y_pred, target_names=labels, output_dict=True)

    # 转换为DataFrame
    report_df = pd.DataFrame(report).iloc[:-1, :].T  # 去掉最后一行的accuracy

    # 创建图表
    fig, ax = _prepare_ax(ax, **heatmap_kwargs)

    # 绘制热力图
    sns.heatmap(report_df, annot=True, fmt='.2f', cmap=cmap, ax=ax, **heatmap_kwargs)

    # 设置标签
    if title is not None:
        ax.set_title(title)
    else:
        ax.set_title('Classification Report')

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax


def plot_feature_importance(
    importance: Union['NDArray', 'Tensor'],
    feature_names: Optional[typing.Sequence[str]] = None,
    top_k: Optional[int] = None,
    horizontal: bool = True,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    cmap: str = 'viridis',
    ax: Optional['Axes'] = None,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    **bar_kwargs,
) -> FigAx:
    """绘制特征重要性图。

    Parameters
    ----------
    - `importance` : `Union['NDArray', 'Tensor']`
        - 特征重要性值
    - `feature_names` : `Optional[List[str]]`, optional
        - 特征名称列表, by default None
    - `top_k` : `Optional[int]`, optional
        - 显示前k个最重要的特征, by default None
    - `horizontal` : `bool`, optional
        - 是否绘制水平条形图, by default True
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `xlabel` : `Optional[str]`, optional
        - X轴标签, by default None
    - `ylabel` : `Optional[str]`, optional
        - Y轴标签, by default None
    - `cmap` : `str`, optional
        - 颜色映射, by default 'viridis'
    - `ax` : `Optional['Axes']`, optional
        - 绘图轴对象，如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径，如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组
    """
    # 确保数据为numpy数组
    importance = ensure_ndarray(importance)

    # 排序特征重要性
    indices = np.argsort(importance)[::-1]

    # 选择前k个特征
    if top_k is not None:
        indices = indices[:top_k]

    # 获取特征名称
    if feature_names is None:
        feature_names = [f'Feature {i}' for i in range(len(importance))]
    feature_names = [str(name) for name in feature_names]

    # 创建图表
    fig, ax = _prepare_ax(ax, **bar_kwargs)

    # 绘制条形图
    if horizontal:
        bars = ax.barh(range(len(indices)), importance[indices])
        ax.set_yticks(range(len(indices)))
        ax.set_yticklabels([feature_names[int(i)] for i in indices])
        ax.invert_yaxis()

        if xlabel is not None:
            ax.set_xlabel(xlabel)
        else:
            ax.set_xlabel('Importance')

        if ylabel is not None:
            ax.set_ylabel(ylabel)
        else:
            ax.set_ylabel('Features')
    else:
        bars = ax.bar(range(len(indices)), importance[indices])
        ax.set_xticks(range(len(indices)))
        ax.set_xticklabels(
            [feature_names[int(i)] for i in indices], rotation=45, ha='right'
        )

        if xlabel is not None:
            ax.set_xlabel(xlabel)
        else:
            ax.set_xlabel('Features')

        if ylabel is not None:
            ax.set_ylabel(ylabel)
        else:
            ax.set_ylabel('Importance')

    # 添加颜色渐变
    colors = plt.cm.get_cmap(cmap)(importance[indices] / np.max(importance[indices]))
    for bar, color in zip(bars, colors):
        bar.set_color(color)

    if title is not None:
        ax.set_title(title)
    else:
        ax.set_title('Feature Importance')

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax


def _plot_learning_curve(
    train_sizes: Union['NDArray', 'Tensor'],
    train_scores: Union['NDArray', 'Tensor'],
    val_scores: Union['NDArray', 'Tensor'],
    train_scores_std: Optional[Union['NDArray', 'Tensor']] = None,
    val_scores_std: Optional[Union['NDArray', 'Tensor']] = None,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    train_color: str | None = None,  # 默认调色
    val_color: str | None = None,
    train_linestyle: str = 'o-',
    val_linestyle: str = 'o-',
    ax: Optional['Axes'] = None,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    **plot_kwargs,
) -> FigAx:
    """绘制学习曲线。

    Parameters
    ----------
    - `train_sizes` : `Union['NDArray', 'Tensor']`
        - 训练集大小
    - `train_scores` : `Union['NDArray', 'Tensor']`
        - 训练集得分
    - `val_scores` : `Union['NDArray', 'Tensor']`
        - 验证集得分
    - `train_scores_std` : `Optional[Union['NDArray', 'Tensor']]`, optional
        - 训练集得分标准差, by default None
    - `val_scores_std` : `Optional[Union['NDArray', 'Tensor']]`, optional
        - 验证集得分标准差, by default None
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `xlabel` : `Optional[str]`, optional
        - X轴标签, by default None
    - `ylabel` : `Optional[str]`, optional
        - Y轴标签, by default None
    - `train_color` : `str`, optional
        - 训练曲线颜色, by default 'blue'
    - `val_color` : `str`, optional
        - 验证曲线颜色, by default 'red'
    - `train_linestyle` : `str`, optional
        - 训练曲线线型, by default 'o-'
    - `val_linestyle` : `str`, optional
        - 验证曲线线型, by default 'o-'
    - `ax` : `Optional['Axes']`, optional
        - 绘图轴对象，如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径，如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组
    """
    # 确保数据为numpy数组
    train_sizes = ensure_ndarray(train_sizes)
    train_scores = ensure_ndarray(train_scores)
    val_scores = ensure_ndarray(val_scores)

    if train_scores_std is not None:
        train_scores_std = ensure_ndarray(train_scores_std)
    if val_scores_std is not None:
        val_scores_std = ensure_ndarray(val_scores_std)

    # 创建图表
    fig, ax = _prepare_ax(ax, **plot_kwargs)

    # 绘制训练曲线
    ax.plot(
        train_sizes,
        train_scores,
        train_linestyle,
        label='Training score',
        color=train_color,
    )
    ax.plot(
        train_sizes, val_scores, val_linestyle, label='Validation score', color=val_color
    )

    # 添加误差带
    if train_scores_std is not None:
        ax.fill_between(
            train_sizes,
            train_scores - train_scores_std,
            train_scores + train_scores_std,
            alpha=0.1,
            color=train_color,
        )

    if val_scores_std is not None:
        ax.fill_between(
            train_sizes,
            val_scores - val_scores_std,
            val_scores + val_scores_std,
            alpha=0.1,
            color=val_color,
        )

    # 设置标签
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    else:
        ax.set_xlabel('Training Set Size')

    if ylabel is not None:
        ax.set_ylabel(ylabel)
    else:
        ax.set_ylabel('Score')

    if title is not None:
        ax.set_title(title)
    else:
        ax.set_title('Learning Curve')

    ax.legend()
    ax.grid(True, alpha=0.3)

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax


def plot_training_curve(
    epochs: Union['NDArray', 'Tensor', int],
    train_loss: Union['NDArray', 'Tensor'],
    val_loss: Optional[Union['NDArray', 'Tensor']] = None,
    train_metric: Optional[Union['NDArray', 'Tensor']] = None,
    val_metric: Optional[Union['NDArray', 'Tensor']] = None,
    val_epochs: Optional[Union['NDArray', 'Tensor']] = None,
    metric_label: str = 'Metric',
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    train_color: str | None = None,
    val_color: str | None = None,
    train_linestyle: str = '-',
    val_linestyle: str = '-',
    dual_axis: bool = False,
    ax: Optional['Axes'] = None,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    **plot_kwargs,
) -> FigAx:
    """绘制训练过程曲线（损失曲线或性能曲线）。

    用于可视化模型在训练过程中每个epoch的损失或性能变化。

    Parameters
    ----------
    - `epochs` : `Union['NDArray', 'Tensor', int]`
        - 训练轮次/迭代次数
    - `train_loss` : `Union['NDArray', 'Tensor']`
        - 训练损失
    - `val_loss` : `Optional[Union['NDArray', 'Tensor']]`, optional
        - 验证损失, by default None
    - `train_metric` : `Optional[Union['NDArray', 'Tensor']]`, optional
        - 训练指标（如准确率）, by default None
    - `val_metric` : `Optional[Union['NDArray', 'Tensor']]`, optional
        - 验证指标（如准确率）, by default None
    - `val_epochs` : `Optional[Union['NDArray', 'Tensor']]`, optional
        - 验证数据对应的轮次，用于不同频率的验证（如每5个epoch验证一次）, by default None
    - `train_metric_label` : `Optional[str]`, optional
        - 训练指标的标签名称, by default None
    - `val_metric_label` : `Optional[str]`, optional
        - 验证指标的标签名称, by default None
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `xlabel` : `Optional[str]`, optional
        - X轴标签, by default None
    - `ylabel` : `Optional[str]`, optional
        - Y轴标签, by default None
    - `train_color` : `str | None`, optional
        - 训练曲线颜色, by default None
    - `val_color` : `str | None`, optional
        - 验证曲线颜色, by default None
    - `train_linestyle` : `str`, optional
        - 训练曲线线型, by default '-'
    - `val_linestyle` : `str`, optional
        - 验证曲线线型, by default '-'
    - `dual_axis` : `bool`, optional
        - 是否使用双Y轴（损失和指标分开显示）, by default False
    - `ax` : `Optional['Axes']`, optional
        - 绘图轴对象，如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径，如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组
    """
    # 确保数据为numpy数组
    if isinstance(epochs, int):
        epochs = np.arange(1, epochs + 1)
    else:
        epochs = ensure_ndarray(epochs)
    train_loss = ensure_ndarray(train_loss)

    if val_loss is not None:
        val_loss = ensure_ndarray(val_loss)
    if train_metric is not None:
        train_metric = ensure_ndarray(train_metric)
    if val_metric is not None:
        val_metric = ensure_ndarray(val_metric)

    # 处理val_epochs - 如果没有提供，则使用epochs
    if val_epochs is not None:
        val_epochs = ensure_ndarray(val_epochs)
    else:
        val_epochs = epochs

    # 处理标签名称
    train_metric_label = 'Training ' + metric_label
    val_metric_label = 'Validation ' + metric_label

    # 创建图表
    fig, ax = _prepare_ax(ax, **plot_kwargs)

    # 如果需要双Y轴且提供了指标数据
    if dual_axis and (train_metric is not None or val_metric is not None):
        ax2 = ax.twinx()

        # 绘制损失曲线（左Y轴）
        if train_loss is not None:
            ax.plot(
                epochs,
                train_loss,
                train_linestyle,
                color=train_color,
                alpha=0.7,
                label='Training Loss',
            )
        if val_loss is not None:
            ax.plot(
                val_epochs,
                val_loss,
                val_linestyle,
                color=val_color,
                alpha=0.7,
                label='Validation Loss',
            )

        # 绘制指标曲线（右Y轴）
        if train_metric is not None:
            ax2.plot(
                epochs,
                train_metric,
                train_linestyle,
                color=train_color,
                alpha=1.0,
                label=train_metric_label,
            )
        if val_metric is not None:
            ax2.plot(
                val_epochs,
                val_metric,
                val_linestyle,
                color=val_color,
                alpha=1.0,
                label=val_metric_label,
            )

        # 设置标签
        ax.set_ylabel('Loss')
        ax2.set_ylabel(metric_label)

        # 合并图例
        lines1, labels1 = ax.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax.legend(lines1 + lines2, labels1 + labels2, loc='center right')

    else:
        # 单Y轴模式
        if train_loss is not None:
            ax.plot(
                epochs,
                train_loss,
                train_linestyle,
                color=train_color,
                label='Training Loss',
            )
        if val_loss is not None:
            ax.plot(
                val_epochs,
                val_loss,
                val_linestyle,
                color=val_color,
                label='Validation Loss',
            )
        if train_metric is not None:
            ax.plot(
                epochs,
                train_metric,
                train_linestyle,
                color=train_color,
                alpha=0.7,
                label=train_metric_label,
            )
        if val_metric is not None:
            ax.plot(
                val_epochs,
                val_metric,
                val_linestyle,
                color=val_color,
                alpha=0.7,
                label=val_metric_label,
            )

        # 设置Y轴标签
        if ylabel is not None:
            ax.set_ylabel(ylabel)
        elif train_loss is not None or val_loss is not None:
            ax.set_ylabel('Loss')
        else:
            ax.set_ylabel('Metric')

        ax.legend()

    # 设置X轴标签
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    else:
        ax.set_xlabel('Epoch')

    # 设置标题
    if title is not None:
        ax.set_title(title)
    else:
        ax.set_title('Training Process')

    ax.grid(True, alpha=0.3)

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax


def plot_roc_curve(
    y_true: Union['NDArray', 'Tensor'],
    y_scores: Union['NDArray', 'Tensor'],
    labels: Optional[typing.Sequence[str]] = None,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    ax: Optional['Axes'] = None,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    **plot_kwargs,
) -> FigAx:
    """绘制ROC曲线。

    Parameters
    ----------
    - `y_true` : `Union['NDArray', 'Tensor']`
        - 真实标签
    - `y_scores` : `Union['NDArray', 'Tensor']`
        - 预测概率或得分
    - `labels` : `Optional[List[str]]`, optional
        - 类别标签名称, by default None
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `xlabel` : `Optional[str]`, optional
        - X轴标签, by default None
    - `ylabel` : `Optional[str]`, optional
        - Y轴标签, by default None
    - `ax` : `Optional['Axes']`, optional
        - 绘图轴对象，如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径，如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组
    """
    from sklearn.metrics import auc, roc_curve
    from sklearn.preprocessing import label_binarize

    # 确保数据为numpy数组
    y_true = ensure_ndarray(y_true)
    y_scores = ensure_ndarray(y_scores)

    # 创建图表
    fig, ax = _prepare_ax(ax, **plot_kwargs)

    # 获取唯一类别
    classes = np.unique(y_true)
    n_classes = len(classes)

    if n_classes == 2:
        # 二分类
        fpr, tpr, _ = roc_curve(y_true, y_scores)
        roc_auc = auc(fpr, tpr)

        ax.plot(fpr, tpr, label=f'ROC curve (AUC = {roc_auc:.2f})')

    else:
        # 多分类
        y_true_bin = label_binarize(y_true, classes=classes)

        # 计算每个类别的ROC曲线
        for i in range(n_classes):
            fpr, tpr, _ = roc_curve(y_true_bin[:, i], y_scores[:, i])
            roc_auc = auc(fpr, tpr)

            label_name = labels[i] if labels is not None else f'Class {classes[i]}'
            ax.plot(fpr, tpr, label=f'{label_name} (AUC = {roc_auc:.2f})')

    # 绘制对角线
    ax.plot([0, 1], [0, 1], 'k--', alpha=0.6)

    # 设置标签
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    else:
        ax.set_xlabel('False Positive Rate')

    if ylabel is not None:
        ax.set_ylabel(ylabel)
    else:
        ax.set_ylabel('True Positive Rate')

    if title is not None:
        ax.set_title(title)
    else:
        ax.set_title('ROC Curve')

    ax.legend()
    ax.grid(True, alpha=0.3)

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax


def plot_precision_recall_curve(
    y_true: Union['NDArray', 'Tensor'],
    y_scores: Union['NDArray', 'Tensor'],
    labels: Optional[typing.Sequence[str]] = None,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    ax: Optional['Axes'] = None,
    show: bool = True,
    save_path: Optional[str] = None,
    auto_close: Optional[bool] = None,
    **plot_kwargs,
) -> FigAx:
    """绘制精确率-召回率曲线。

    Parameters
    ----------
    - `y_true` : `Union['NDArray', 'Tensor']`
        - 真实标签
    - `y_scores` : `Union['NDArray', 'Tensor']`
        - 预测概率或得分
    - `labels` : `Optional[List[str]]`, optional
        - 类别标签名称, by default None
    - `title` : `Optional[str]`, optional
        - 图表标题, by default None
    - `xlabel` : `Optional[str]`, optional
        - X轴标签, by default None
    - `ylabel` : `Optional[str]`, optional
        - Y轴标签, by default None
    - `ax` : `Optional['Axes']`, optional
        - 绘图轴对象，如果为None则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `Optional[str]`, optional
        - 保存路径，如果提供则保存图表, by default None
    - `auto_close` : `Optional[bool]`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `FigAx`
        - 图表对象和轴对象的元组
    """
    from sklearn.metrics import average_precision_score, precision_recall_curve
    from sklearn.preprocessing import label_binarize

    # 确保数据为numpy数组
    y_true = ensure_ndarray(y_true)
    y_scores = ensure_ndarray(y_scores)

    # 创建图表
    fig, ax = _prepare_ax(ax, **plot_kwargs)

    # 获取唯一类别
    classes = np.unique(y_true)
    n_classes = len(classes)

    if n_classes == 2:
        # 二分类
        precision, recall, _ = precision_recall_curve(y_true, y_scores)
        ap = average_precision_score(y_true, y_scores)

        ax.plot(recall, precision, label=f'PR curve (AP = {ap:.2f})')

    else:
        # 多分类
        y_true_bin = label_binarize(y_true, classes=classes)

        # 计算每个类别的PR曲线
        for i in range(n_classes):
            precision, recall, _ = precision_recall_curve(
                y_true_bin[:, i], y_scores[:, i]
            )
            ap = average_precision_score(y_true_bin[:, i], y_scores[:, i])

            label_name = labels[i] if labels is not None else f'Class {classes[i]}'
            ax.plot(recall, precision, label=f'{label_name} (AP = {ap:.2f})')

    # 设置标签
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    else:
        ax.set_xlabel('Recall')

    if ylabel is not None:
        ax.set_ylabel(ylabel)
    else:
        ax.set_ylabel('Precision')

    if title is not None:
        ax.set_title(title)
    else:
        ax.set_title('Precision-Recall Curve')

    ax.legend()
    ax.grid(True, alpha=0.3)

    _finalize_figure(fig, show, save_path, auto_close)
    return fig, ax
