from __future__ import annotations

from dataclasses import dataclass, field, is_dataclass
from numbers import Real
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Generic,
    Literal,
    Mapping,
    Sequence,
    TypeAlias,
    TypeVar,
)

if TYPE_CHECKING:
    from lightning.fabric import Fabric
    from zyplib.experiment.context import ExperimentContext, TrialContext
    from torch.nn import Module
    from torch.utils.data import DataLoader

ParamsT = TypeVar('ParamsT')
MetricScalar: TypeAlias = int | float
MetricsLike: TypeAlias = Mapping[str, MetricScalar]
TrialMode: TypeAlias = Literal['single', 'repeat', 'fold']
ExportRuleFn: TypeAlias = Callable[[str, str, Any], bool]
ExportRule: TypeAlias = Sequence[str] | ExportRuleFn


def matches_export_rule(rule: ExportRule, namespace: str, key: str, value: Any) -> bool:
    if callable(rule):
        return bool(rule(namespace, key, value))
    return namespace in set(rule)


def serialize_value(value: Any) -> Any:
    """把常见 experiment 对象转成适合写入 YAML 的基础类型。

    Parameters
    ----------
    - `value`: `Any`, 待序列化的对象

    Returns
    ----------
    - `Any`, 仅由基础类型、字典、列表组成的对象

    Examples
    ----------
    ```python
    serialize_value({'trial_name': 'repeat-001', 'score': 0.91})
    serialize_value(Path('temp/demo/report.yaml'))
    ```
    """

    # bool 必须在 Real 之前检查，否则会被当成数值转成 0.0 / 1.0。
    if value is None or isinstance(value, (str, bool, int, float)):
        return value
    if isinstance(value, Real):
        return float(value)
    if isinstance(value, Path):
        return str(value)
    if hasattr(value, 'to_dict'):
        return serialize_value(value.to_dict())
    if is_dataclass(value) and not isinstance(value, type):
        return {
            field_name: serialize_value(getattr(value, field_name))
            for field_name in value.__dataclass_fields__
        }
    if isinstance(value, Mapping):
        return {str(key): serialize_value(item) for key, item in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [serialize_value(item) for item in value]
    return str(value)


def normalize_metrics(metrics: Mapping[str, Any]) -> dict[str, float]:
    """规范化 trial 返回的指标字典。

    Parameters
    ----------
    - `metrics`: `Mapping[str, Any]`, trial 返回的原始指标

    Returns
    ----------
    - `dict[str, float]`, 仅包含浮点数值的指标字典

    Examples
    ----------
    ```python
    normalize_metrics({'acc': 0.91, 'loss': 0.12})
    ```
    """

    normalized: dict[str, float] = {}
    for key, value in metrics.items():
        if isinstance(value, bool) or not isinstance(value, Real):
            raise TypeError(f'Metric `{key}` must be a real number, got {type(value)!r}')
        normalized[key] = float(value)
    return normalized


@dataclass(slots=True)
class AutoRecordConfig:
    """配置 framework 自动补充哪些 trial 记录。"""

    data_info: bool = True
    model_hparams: bool = True
    optimizer: bool = True
    lr_scheduler: bool = True
    trial_meta: bool = True


@dataclass(slots=True)
class RecordConfig:
    """配置 trial 记录的 inject / export / auto-record 行为。"""

    enabled: bool = True
    inject_namespaces: tuple[str, ...] = ('default', 'train')
    export_rule: ExportRule = ('default', 'result', 'env')
    auto_record: AutoRecordConfig | bool = True

    # exclude_trials_in_final_report: bool = False

    inject_ctx_key: str = 'experiment_ctx'
    inject_params_key: str = 'params'
    inject_trial_key: str = 'trial'


@dataclass(slots=True)
class ExperimentEffects:
    """配置 experiment 运行期的副作用行为。"""

    capture_stdout: bool = True
    record_env: bool = True
    record: RecordConfig | bool = True


@dataclass(frozen=True, slots=True)
class TrialPlan:
    """定义一次 experiment 顶层需要运行多少个 trial。

    Parameters
    ----------
    - `count`: `int`, trial 数量
    - `mode`: `TrialMode`, trial 模式，例如 `single`、`repeat`、`fold`
    - `base_seed`: `int | None`, 用于生成每个 trial seed 的顶层随机种子

    Examples
    ----------
    ```python
    TrialPlan.single(seed=9712)
    TrialPlan.repeat(3, base_seed=42)
    TrialPlan.kfold(5, base_seed=42)
    ```
    """

    count: int = 1
    mode: TrialMode = 'single'
    base_seed: int | None = None
    trial_seeds: tuple[int | None, ...] | None = None

    def __post_init__(self) -> None:
        if self.count <= 0:
            raise ValueError('TrialPlan.count must be greater than 0')
        if self.trial_seeds is not None and len(self.trial_seeds) != self.count:
            raise ValueError('TrialPlan.trial_seeds length must equal count')

    @classmethod
    def single(cls, seed: int | None = None) -> 'TrialPlan':
        return cls(count=1, mode='single', base_seed=seed)

    @classmethod
    def repeat(
        cls,
        n_trials: int,
        base_seed: int | None = None,
        trial_seeds: Sequence[int | None] | None = None,
    ) -> 'TrialPlan':
        assert n_trials > 0, 'n_trials must be greater than 0'
        if trial_seeds is not None:
            assert len(trial_seeds) == n_trials, 'Length of trial_seeds must equal n_trials'
        return cls(
            count=n_trials,
            mode='repeat',
            base_seed=base_seed,
            trial_seeds=tuple(trial_seeds) if trial_seeds is not None else None,
        )

    @classmethod
    def kfold(
        cls,
        k: int,
        base_seed: int | None = None,
        trial_seeds: Sequence[int | None] | None = None,
    ) -> 'TrialPlan':
        assert k > 0, 'k must be greater than 0'
        if trial_seeds is not None:
            assert len(trial_seeds) == k, 'Length of trial_seeds must equal k'
        return cls(
            count=k,
            mode='fold',
            base_seed=base_seed,
            trial_seeds=tuple(trial_seeds) if trial_seeds is not None else None,
        )


@dataclass(slots=True)
class ExperimentData:
    """一次 trial 需要消费的数据对象。

    Parameters
    ----------
    - `train_loader`: `DataLoader`, 训练集 dataloader
    - `val_loader`: `DataLoader | None`, 验证集 dataloader, by default `None`
    - `test_loader`: `DataLoader | None`, 测试集 dataloader, by default `None`
    - `info`: `dict[str, Any]`, auto_record 使用的可序列化数据元信息, by default `{}`
    - `extras`: `dict[str, Any]`, runtime-only 附加对象，framework 不序列化, by default `{}`

    Examples
    ----------
    ```python
    data = ExperimentData(train_loader, val_loader, test_loader, info={'files': split_files})
    ```
    """

    train_loader: 'DataLoader'
    val_loader: 'DataLoader | None' = None
    test_loader: 'DataLoader | None' = None
    info: dict[str, Any] = field(default_factory=dict)
    extras: dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class ModelSetup:
    """模型构造步骤的返回对象。

    Parameters
    ----------
    - `model`: `Module`, 当前 trial 使用的模型实例
    - `optimizer`: `Any | None`, 与模型一起构造的优化器, by default `None`
    - `lr_scheduler`: `Any | None`, 与模型一起构造的学习率调度器, by default `None`
    - `hparams`: `dict[str, Any]`, 要写入记录的模型级超参数, by default `{}`

    Examples
    ----------
    ```python
    ModelSetup(
        model=my_model,
        optimizer=optimizer,
        lr_scheduler=scheduler,
        hparams={'model': 'EEGNet', 'channels': 8},
    )
    ```
    """

    model: 'Module'
    optimizer: Any = None
    lr_scheduler: Any = None
    hparams: dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class TrialResult:
    """单个 trial 主动返回给 framework 的结果。

    Parameters
    ----------
    - `metrics`: `dict[str, float]`, 会进入 summary 聚合的核心指标
    - `summary`: `dict[str, Any]`, 补充的 trial 摘要信息, by default `{}`
    - `artifacts`: `dict[str, str]`, 要显式保留的 artifact 路径, by default `{}`

    Examples
    ----------
    ```python
    TrialResult(metrics={'acc': 0.91}, summary={'fold': 1})
    ```
    """

    metrics: dict[str, float]
    summary: dict[str, Any] = field(default_factory=dict)
    artifacts: dict[str, str] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return serialize_value(
            {
                'metrics': self.metrics,
                'summary': self.summary,
                'artifacts': self.artifacts,
            }
        )


@dataclass(slots=True)
class TrialReport:
    """framework 持久化后的单个 trial 报告。

    Parameters
    ----------
    - `trial_name`: `str`, 当前 trial 名称
    - `trial_index`: `int`, 当前 trial 的 0-based 索引
    - `trial_mode`: `TrialMode`, 当前 trial 所属模式
    - `trial_seed`: `int | None`, 当前 trial 使用的随机种子
    - `result`: `TrialResult`, trial 主动返回的结果
    - `records`: `dict[str, dict[str, Any]]`, 由 context / inject 自动抓取的记录
    """

    trial_name: str
    trial_index: int
    trial_mode: TrialMode
    trial_seed: int | None
    result: TrialResult
    records: dict[str, dict[str, Any]] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return serialize_value(
            {
                'trial_name': self.trial_name,
                'trial_index': self.trial_index,
                'trial_mode': self.trial_mode,
                'trial_seed': self.trial_seed,
                'result': self.result,
                'records': self.records,
            }
        )


@dataclass(slots=True)
class ExperimentReport(Generic[ParamsT]):
    """一次完整 experiment run 的最终报告。

    Parameters
    ----------
    - `name`: `str`, experiment 名称
    - `params`: `ParamsT`, 当前运行的参数对象
    - `run_dir`: `Path`, 当前 run 的输出目录
    - `summary`: `dict[str, Any]`, 聚合后的总指标
    - `trials`: `list[TrialReport]`, 每个 trial 的详细报告
    - `env`: `dict[str, Any]`, 环境快照, by default `{}`

    Examples
    ----------
    ```python
    report = run_experiment(...)
    print(report.summary)
    ```
    """

    name: str
    params: ParamsT
    run_dir: Path
    summary: dict[str, Any]
    trials: list[TrialReport]
    env: dict[str, Any] = field(default_factory=dict)

    def to_dict(self, *, include_trials: bool = True) -> dict[str, Any]:
        dic_obj = {
            'name': self.name,
            'params': self.params,
            'run_dir': self.run_dir,
            'summary': self.summary,
            'trials': self.trials,
            'env': self.env,
        }
        if not include_trials:
            dic_obj.pop('trials')
        return serialize_value(dic_obj)


PrepareRunFn: TypeAlias = Callable[['ExperimentContext[ParamsT]'], Any]
BuildDataFn: TypeAlias = Callable[['TrialContext[ParamsT]'], 'ExperimentData']
BuildModelFn: TypeAlias = Callable[['TrialContext[ParamsT]'], 'ModelSetup | Module']
TrainModelFn: TypeAlias = Callable[
    ['TrialContext[ParamsT]', 'ModelSetup', 'ExperimentData'], 'TrialResult | MetricsLike'
]
ResolvePlanFn: TypeAlias = Callable[['ExperimentContext[ParamsT]'], TrialPlan]
ResultAggregator: TypeAlias = Callable[[Sequence['TrialReport']], dict[str, Any]]


__all__ = [
    'BuildDataFn',
    'BuildModelFn',
    'AutoRecordConfig',
    'ExperimentData',
    'ExperimentEffects',
    'ExportRule',
    'ExportRuleFn',
    'ExperimentReport',
    'MetricsLike',
    'MetricScalar',
    'matches_export_rule',
    'ModelSetup',
    'ParamsT',
    'PrepareRunFn',
    'RecordConfig',
    'ResolvePlanFn',
    'ResultAggregator',
    'TrainModelFn',
    'TrialMode',
    'TrialPlan',
    'TrialReport',
    'TrialResult',
    'normalize_metrics',
    'serialize_value',
]
