from __future__ import annotations

from typing import Generic, Sequence, cast

from zyplib.experiment._types import (
    ExperimentEffects,
    ExperimentData,
    ExperimentReport,
    ModelSetup,
    ParamsT,
    TrialPlan,
    TrialResult,
)
from zyplib.experiment.context import ExperimentContext, TrialContext
from zyplib.experiment.procedure import run_experiment


class BasicExperiment(Generic[ParamsT]):
    """面向对象风格的 experiment 模板。

    它把一次实验拆成几个更贴近常规科研脚本的步骤：

    - `prepare_run()`：准备共享数据，例如文件列表或固定 fold plan
    - `get_trial_plan()`：声明 single / repeat / fold 运行方式
    - `build_data()`：得到 train/val/test dataloader
    - `build_model()`：构造模型、优化器、调度器与模型级配置
    - `train_model()`：调用 `procedure.train_loop()`、`BasicFabricTrainer` 或项目自定义训练逻辑

    Examples
    ----------
    ```python
    class DemoExperiment(BasicExperiment[ExpParams]):
        Params = ExpParams

        def build_data(self, ctx):
            return data

        def build_model(self, ctx):
            return model_setup

        def train_model(self, ctx, model_setup, data):
            return {'acc': 0.91}
    ```
    """

    name: str = ''
    root_dir: str = 'temp/experiment'
    Params: type[ParamsT]
    effects: ExperimentEffects | None = None
    aggregate: str = 'meanstd'

    def prepare_run(self, ctx: ExperimentContext[ParamsT]):
        """准备 run 级 shared data。

        Parameters
        ----------
        - `ctx`: `ExperimentContext[ParamsT]`, 当前 run 上下文

        Returns
        ----------
        - `Any`, 会写入 `ctx.shared` 并共享给全部 trials；默认返回 `None`
        """

        return None

    def get_trial_plan(self, ctx: ExperimentContext[ParamsT]) -> TrialPlan:
        """返回本次 run 使用的 trial 计划。"""

        return TrialPlan.single()

    def build_data(self, ctx: TrialContext[ParamsT]) -> ExperimentData:
        """构造当前 trial 使用的数据对象。"""

        raise NotImplementedError('Subclasses must implement build_data()')

    def build_model(self, ctx: TrialContext[ParamsT]) -> ModelSetup:
        """构造当前 trial 使用的模型与模型级配置。

        推荐在返回的 `ModelSetup` 中一并携带 optimizer、scheduler、loss_fn
        等与模型强绑定的对象。
        """

        raise NotImplementedError('Subclasses must implement build_model()')

    def train_model(
        self, ctx: TrialContext[ParamsT], model_setup: ModelSetup, data: ExperimentData
    ) -> TrialResult | dict[str, float]:
        """执行当前 trial 的训练与测试。"""

        raise NotImplementedError('Subclasses must implement train_model()')

    def resolve_root_dir(self, params: ParamsT) -> str:
        """解析当前运行使用的输出根目录。

        优先使用参数对象里的 `RECORD_DIR` / `record_dir`，否则退回类属性
        `root_dir`。这样类式写法与函数式写法都能保持参数驱动的一致性。
        """

        for attr_name in ('RECORD_DIR', 'record_dir', 'ROOT_DIR', 'root_dir'):
            value = getattr(params, attr_name, None)
            if isinstance(value, str) and value:
                return value
        return self.root_dir

    def run(self, params: ParamsT) -> ExperimentReport[ParamsT]:
        """直接用参数对象执行一次 experiment。"""

        return run_experiment(
            name=self.name or self.__class__.__name__,
            params=params,
            root_dir=self.resolve_root_dir(params),
            prepare_run=self.prepare_run,
            plan=self.get_trial_plan,
            build_data=self.build_data,
            build_model=self.build_model,
            train_model=self.train_model,
            aggregate=self.aggregate,
            effects=self.effects,
        )

    def run_cli(self, args: Sequence[str] | None = None) -> ExperimentReport[ParamsT]:
        """通过 `tyro` 解析 CLI 参数后执行 experiment。"""

        if not hasattr(self, 'Params') or self.Params is None:
            raise TypeError(
                f'{type(self).__name__} must define a `Params` class attribute '
                f'to use run_cli(). Example: `Params = MyParamsDataclass`'
            )

        import tyro  # pyright: ignore[reportMissingImports]

        params = cast(
            ParamsT, tyro.cli(self.Params, args=list(args) if args is not None else None)
        )
        return self.run(params)


__all__ = ['BasicExperiment']
