import typing
from typing import Literal, Optional

from torchmetrics import Metric, MetricCollection, classification, regression

from .torch_metric import TorchMetricRecorder
from .common_metric import create_lambda_metric


def clf_metrics_builder(
    binary: bool = True,
    *,
    task: Optional[Literal['binary', 'multiclass', 'multilabel']] = None,
    threshold: float = 0.5,
    num_classes: Optional[int] = None,
    average: Optional[Literal['micro', 'macro', 'weighted', 'none']] = 'macro',
    top_k: Optional[int] = 1,
    use_auc: bool = False,
    use_specificity: bool = False,
    use_matthews: bool = False,
    use_kappa: bool = False,
    use_ece: bool = False,
    **kwargs,
) -> MetricCollection:
    """构建并返回分类指标集合

    为不同分类任务构建合适的指标集合，针对脑电信号分析进行优化。

    Parameters
    ----------
    - `binary` : `bool`, optional
        - 兼容旧版本参数，当 `task` 未指定时使用, by default True
    - `task` : `Literal['binary','multiclass','multilabel']`, optional
        - 指定分类任务类型，优先级高于 `binary` 参数, by default None
    - `threshold` : `float`, optional
        - 二分类和多标签任务的阈值, by default 0.5
    - `num_classes` : `int`, optional
        - 多分类或多标签任务的类别数，多分类和多标签任务时必需, by default None
    - `average` : `Literal['micro', 'macro', 'weighted', 'none']`, optional
        - 多分类/多标签情况下的平均方式，macro 对不平衡数据更鲁棒, by default 'macro'
    - `top_k` : `int`, optional
        - 多分类任务中计算 Top-K 准确率时的 K 值, by default 1
    - `use_auc` : `bool`, optional
        - 是否包含 AUROC 和 AUPRC 指标, by default False
    - `use_specificity` : `bool`, optional
        - 是否包含特异性指标, by default False
    - `use_matthews` : `bool`, optional
        - 是否包含 Matthews 相关系数（对不平衡数据鲁棒）, by default False
    - `use_kappa` : `bool`, optional
        - 是否包含 Cohen's Kappa（评估分类一致性）, by default False
    - `use_ece` : `bool`, optional
        - 是否包含 Expected Calibration Error（评估模型校准度）, by default False

    Returns
    ----------
    - `MetricCollection`
        - 包含所选指标的 TorchMetrics 集合

    Raises
    ----------
    - `ValueError`
        - 当多分类/多标签任务未提供 num_classes 时抛出
    """
    # 兼容老版本 binary 参数
    if task is None:
        task = 'binary' if binary else 'multiclass'

    # 多分类/多标签必须指定 num_classes
    if task in ('multiclass', 'multilabel') and num_classes is None:
        raise ValueError('`num_classes` must be provided for multiclass/multilabel tasks')

    base_args = {
        'task': task,
        'num_classes': num_classes,
        'threshold': threshold,
        'average': average,
        'top_k': top_k,
    }
    base_args.update(kwargs)
    # 清理无效 None
    args = {k: v for k, v in base_args.items() if v is not None}

    # 构建基础指标
    metrics = {
        'accuracy': classification.Accuracy(**args),
        'precision': classification.Precision(**args),
        'recall': classification.Recall(**args),
        'f1': classification.F1Score(**args),
    }

    # 根据布尔开关添加额外指标
    if use_auc and task in ('binary', 'multiclass'):
        try:
            metrics['auroc'] = classification.AUROC(**args)
            metrics['auprc'] = classification.AveragePrecision(**args)
        except Exception:
            # 某些参数组合可能不支持，跳过
            pass

    if use_specificity and task in ('binary', 'multiclass'):
        try:
            metrics['specificity'] = classification.Specificity(**args)
        except Exception:
            pass

    if use_matthews:
        try:
            metrics['matthews_corrcoef'] = classification.MatthewsCorrCoef(**args)
        except Exception:
            pass

    if use_kappa and task in ('binary', 'multiclass'):
        try:
            metrics['cohen_kappa'] = classification.CohenKappa(**args)
        except Exception:
            pass

    if use_ece:
        try:
            # ECE 需要特殊的参数处理
            ece_args = {k: v for k, v in args.items() if k in ('task', 'num_classes')}
            metrics['ece'] = classification.CalibrationError(**ece_args)
        except Exception:
            pass

    return MetricCollection(typing.cast(dict[str, Metric | MetricCollection], metrics))


def reg_metrics_builder() -> MetricCollection:
    """构建并返回回归指标字典。

    Return:
    ----------
    - `dict[str, Metric]`, 键是指标名称（例如 "mae", "mse"），值是相应的 `Metric` 对象。
    """
    metrics = {
        'mae': regression.MeanAbsoluteError(),
        'mse': regression.MeanSquaredError(),
        'rmse': regression.MeanSquaredError(squared=False),  # RMSE 是 MSE 的平方根
        'r2_score': regression.R2Score(),
    }
    return MetricCollection(typing.cast(dict[str, Metric | MetricCollection], metrics))


__all__ = [
    'TorchMetricRecorder',
    'clf_metrics_builder',
    'reg_metrics_builder',
    'create_lambda_metric',
]
