from __future__ import annotations
import typing

from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Literal,
    Optional,
    Protocol,
    TypeAlias,
    TypeVar,
)

import torch
from simple_inject import inject
from torch import nn
from torch.utils.data import DataLoader

from zyplib.train.common import callbacks
from zyplib.train.common._base import Batch, FabricTrainingContext, SingleOptimizer

# from zyplib.train.utils import step_lr_sched
from zyplib.utils.print import print_debug, print_info, print_warning
from zyplib.utils.progress_bar import rich_progress

if TYPE_CHECKING:
    from torch.optim.lr_scheduler import _LRScheduler
    from torch.types import _TensorOrTensors

    from zyplib.train.common.callbacks import EarlyStopping
    from zyplib.train.metric.torch_metric import TorchMetricRecorder


def _resolve_loss_fn(
    ctx: FabricTrainingContext, fn_loss: Optional[nn.Module]
) -> nn.Module:
    if ctx.fn_loss is not None:
        return ctx.fn_loss

    # 向后兼容: 如果传入了 fn_loss, 则自动设置到 context.loss_fn 中
    if fn_loss is not None:
        if ctx.fn_loss != fn_loss:
            ctx.fn_loss = fn_loss
            print_debug(f'自动设置 context.loss_fn: {type(ctx.fn_loss).__name__}')
        return fn_loss

    # 两者都没有，抛出错误
    raise ValueError(
        '损失函数未指定：请在 FabricTrainingContext.loss_fn 中设置损失函数，'
        '或在函数调用时传入 fn_loss 参数'
    )


def _default_batch_2_target(batch: tuple):
    return batch[1]


DepecatedFnLoss: TypeAlias = Optional[nn.Module]
StepOutput: TypeAlias = tuple[torch.Tensor, Dict[str, float]]
StepFunction: TypeAlias = Callable[..., StepOutput]


def _require_single_optimizer(ctx: FabricTrainingContext) -> SingleOptimizer:
    return typing.cast(SingleOptimizer, ctx.optimizer)


def _require_train_metric_recorder(ctx: FabricTrainingContext) -> 'TorchMetricRecorder':
    return typing.cast('TorchMetricRecorder', ctx.train_metric_recorder)


def _require_val_metric_recorder(ctx: FabricTrainingContext) -> 'TorchMetricRecorder':
    return typing.cast('TorchMetricRecorder', ctx.val_metric_recorder)


def _setup_single_dataloader(
    ctx: FabricTrainingContext, dataloader: DataLoader
) -> DataLoader:
    return typing.cast(DataLoader, ctx._require_fabric().setup_dataloaders(dataloader))


def training_step(
    ctx: FabricTrainingContext,
    batch: Batch,
    *,
    fn_loss: DepecatedFnLoss | None = None,
    batch_idx: Optional[int] = None,
    **kwargs,
) -> StepOutput:
    """默认内置的单个训练步骤, 主要适配普通的单输入单输出训练任务例如分类

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `batch` : `tuple[torch.Tensor, torch.Tensor]`
        - 数据批次
    - `fn_loss` : `Optional[nn.Module]`, optional
        - 损失函数: `fn_loss(y_pred, y)`，如果为 None 则使用 ctx.loss_fn
    - `batch_idx`: int

    Returns
    ----------
    - `tuple[torch.Tensor, Dict[str, float]]`: (y_pred, {"loss": float})
    """
    ctx.model.train()
    x, y = batch

    # 解析损失函数
    loss_fn = _resolve_loss_fn(ctx, fn_loss)
    optimizer = _require_single_optimizer(ctx)

    optimizer.zero_grad()
    y_pred = ctx.model(x)
    loss: torch.Tensor = loss_fn(y_pred, y)

    if ctx.fabric:
        ctx.fabric.backward(loss)
    else:
        loss.backward()

    optimizer.step()

    loss_value = loss.item()

    # if ctx.train_metric_recorder:
    #     return ctx.train_metric_recorder.update(y_pred, y, loss=loss_value)
    return y_pred, {'loss': loss_value}


def validation_step(
    ctx: FabricTrainingContext,
    batch: Batch,
    *,
    fn_loss: DepecatedFnLoss | None = None,
    batch_idx: Optional[int] = None,
    **kwargs,
) -> StepOutput:
    """默认内置的单个验证步骤, 主要适配普通的单输入单输出验证任务例如分类

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `batch` : `tuple[torch.Tensor, torch.Tensor]`
        - 数据批次
    - `fn_loss` : `Optional[nn.Module]`, optional
        - 损失函数: `fn_loss(y_pred, y)`，如果为 None 则使用 ctx.loss_fn
    - `batch_idx`: int

    Returns
    ----------
    - `tuple[torch.Tensor, Dict[str, float]]`: (y_pred, {"loss": float})
    """
    ctx.model.eval()
    with torch.no_grad():
        x, y = batch
        y_pred = ctx.model(x)

        # 解析损失函数
        loss_fn = _resolve_loss_fn(ctx, fn_loss)
        loss: torch.Tensor = loss_fn(y_pred, y)

    loss_value = loss.item()
    # if ctx.val_metric_recorder:
    #     return ctx.val_metric_recorder.update(y_pred, y, loss=loss_value)
    # return {'loss': loss_value}
    return y_pred, {'loss': loss_value}


def train_epoch(
    ctx: FabricTrainingContext,
    train_loader: DataLoader,
    epoch_idx: Optional[int] = None,
    fn_train_step: Optional[Callable] = None,
    fn_batch2target: Callable[[Batch], torch.Tensor] = _default_batch_2_target,
    fn_transform_batch: Optional[Callable] = None,
) -> Dict[str, float]:
    """单个训练轮次

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `train_loader` : `DataLoader`
        - 训练数据加载器
    - `epoch_idx` : `int | None`, optional
        - 当前轮次索引, 默认为 None
    - `fn_train_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个训练步骤的函数, 默认为 `training_step`
    - `fn_batch2target` : `(...) -> torch.Tensor`, optional
        - 可选项, 用于从 batch 中提取预测值, 默认为 `_default_batch_2_target`
        - 例如默认的方法就是接收一个 (X, y) 的 batch, 并返回 y
        - 对于复杂训练任务，例如 GAN 等，可能需要自定义这个函数
    - `fn_transform_batch` : `(...) -> Batch`, optional
        - 可选项, 用于在训练前转换批次数据的函数, 默认使用 `ctx.transform_on_train_batch_start`
        - 签名: `fn(batch, batch_idx, epoch_idx) -> transformed_batch`
        - 用例: Mixup, CutMix, 动态数据增强等

    Returns
    ----------
    - `Dict[str, float]`: 本 Epoch 的平均指标, 至少包含 `loss`
    """
    if fn_train_step is None:
        fn_train_step = training_step
    if fn_transform_batch is None:
        fn_transform_batch = ctx.transform_on_train_batch_start
    train_metric_recorder = _require_train_metric_recorder(ctx)

    for idx, batch in enumerate(
        rich_progress(train_loader, f'Training epoch {epoch_idx}', color='green')
    ):
        # 应用数据转换 hook
        batch = fn_transform_batch(batch, idx, epoch_idx)

        y_pred, raw_metrics = fn_train_step(
            ctx=ctx, batch=batch, batch_idx=idx, epoch_idx=epoch_idx
        )
        # _, y = batch
        y = fn_batch2target(batch)
        metrics = train_metric_recorder.update(y_pred, y, **raw_metrics)
        ctx.call_on_train_batch_end(batch=batch, batch_idx=idx, metrics=metrics)

    epoch_metrics = train_metric_recorder.compute()
    return epoch_metrics


def val_epoch(
    ctx: FabricTrainingContext,
    val_loader: DataLoader,
    epoch_idx: Optional[int] = None,
    fn_val_step: Optional[Callable] = None,
    fn_batch2target: Callable[[Batch], torch.Tensor] = _default_batch_2_target,
    fn_transform_batch: Optional[Callable] = None,
) -> Dict[str, float]:
    """单个验证轮次

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `val_loader` : `DataLoader`
        - 验证数据加载器
    - `epoch_idx` : `int | None`, optional
        - 当前轮次索引, 默认为 None
    - `fn_val_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个验证步骤的函数, 默认为 `validation_step`
    - `fn_batch2target` : `(...) -> torch.Tensor`, optional
        - 可选项, 用于从 batch 中提取预测值, 默认为 `_default_batch_2_target`
        - 例如默认的方法就是接收一个 (X, y) 的 batch, 并返回 y
        - 对于复杂训练任务，例如 GAN 等，可能需要自定义这个函数
    - `fn_transform_batch` : `(...) -> Batch`, optional
        - 可选项, 用于在验证前转换批次数据的函数, 默认使用 `ctx.transform_on_validation_batch_start`
        - 签名: `fn(batch, batch_idx, epoch_idx) -> transformed_batch`
        - 用例: 测试时增强 (TTA) 等

    Returns
    ----------
    - `Dict[str, float]`: 本 Epoch 的平均指标, 至少包含 `loss`
    """
    if fn_val_step is None:
        fn_val_step = validation_step
    if fn_transform_batch is None:
        fn_transform_batch = ctx.transform_on_validation_batch_start
    val_metric_recorder = _require_val_metric_recorder(ctx)

    for idx, batch in enumerate(
        rich_progress(val_loader, f'Validating epoch {epoch_idx}', color='blue')
    ):
        # 应用数据转换 hook
        batch = fn_transform_batch(batch, idx, epoch_idx)

        y_pred, raw_metrics = fn_val_step(
            ctx=ctx, batch=batch, batch_idx=idx, epoch_idx=epoch_idx
        )
        y = fn_batch2target(batch)
        metrics = val_metric_recorder.update(y_pred, y, **raw_metrics)
        ctx.call_on_validation_batch_end(batch=batch, batch_idx=idx, metrics=metrics)

    metrics = val_metric_recorder.compute()
    return metrics


def test(
    ctx: FabricTrainingContext,
    test_loader: DataLoader,
    fn_test_step: Optional[Callable] | None = None,
    fn_batch2target: Callable[[Batch], torch.Tensor] = _default_batch_2_target,
    **test_step_kwargs: Any,
) -> Dict[str, float]:
    """单个测试轮次

    Parameters
    ----------
    - `ctx` : `procedure.FabricTrainingContext`
    - `test_loader` : `DataLoader`
        - 测试数据加载器
    - `epoch_idx` : `int | None`, optional
        - 当前轮次索引, 默认为 None
    - `fn_test_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个测试步骤的函数, 默认为 `validation_step`
    - `fn_batch2target` : `(...) -> torch.Tensor`, optional
        - 可选项, 用于从 batch 中提取预测值, 默认为 `_default_batch_2_target`
        - 例如默认的方法就是接收一个 (X, y) 的 batch, 并返回 y
        - 对于复杂训练任务，例如 GAN 等，可能需要自定义这个函数

    `fn_test_step`:
    ----------
    - 依赖注入: 从 `namespace=train` 自动获取 `validation_step` 或者 `test_step`
    - 默认不提供, 则使用 `procedure.validation_step`; 如果提供了, 则使用提供的函数
    - 接受以下的关键字参数: `fn(ctx: FabricTrainingContext, batch, fn_loss=None, batch_idx)`
        - `fn_loss` 参数现在是可选的，如果为 None 则使用 ctx.fn_loss
        - 注意，不会有 `epoch_idx` 的参数!
    - 返回值: `torch.Tensor, Dict[str, Any]` (predictions, metrics), 至少包含 `loss`

    Returns
    ----------
    - `Dict[str, float]`: 本次测试的平均指标, 至少包含 `loss`
    """
    # fn_test_step = (
    #     fn_test_step
    #     or inject('test_step', 'train')
    #     or inject('validation_step', 'train')
    #     or validation_step
    # )
    if fn_test_step is None:
        from simple_inject import has as has_dep

        if has_dep('test_step', 'train'):
            fn_test_step = inject('test_step', 'train')
        elif has_dep('validation_step', 'train'):
            fn_test_step = inject('validation_step', 'train')
        else:
            fn_test_step = validation_step

    fn_test = typing.cast(StepFunction, fn_test_step)

    val_metric_recorder = _require_val_metric_recorder(ctx)

    ctx.call_on_test_start(dataloader=test_loader)
    test_loader = _setup_single_dataloader(ctx, test_loader)

    val_metric_recorder.reset()

    for idx, batch in enumerate(rich_progress(test_loader, 'Testing', color='blue')):
        y_pred, raw_metrics = fn_test(
            ctx=ctx, batch=batch, batch_idx=idx, **test_step_kwargs
        )
        y = fn_batch2target(batch)
        metrics = val_metric_recorder.update(y_pred, y, **raw_metrics)
        ctx.call_on_test_batch_end(batch=batch, batch_idx=idx, metrics=metrics)

    metrics = val_metric_recorder.compute()
    ctx.call_on_test_end(metrics=metrics)
    return metrics


# @auto_inject()
def train_loop(
    context: FabricTrainingContext,
    train_loader: DataLoader,
    max_epochs: int | None = 100,
    val_loader: Optional[DataLoader] = None,
    val_interval: int = 1,
    hparams: Optional[typing.Mapping[str, Any]] = None,
    fn_train_step: Optional[Callable] = None,
    fn_val_step: Optional[Callable] = None,
    fn_batch2target: Callable[[Batch], torch.Tensor] = _default_batch_2_target,
    fn_on_trainepoch_end: Optional[Callable] = None,
    # Deprecated parameters for backward compatibility
    # reduce_lr_metric_name: str | None = 'loss',
    # fn_loss: nn.Module | None = None,
    # lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
    # lr_sched_freq: int = 1,
    **kwargs,
):
    """基于 Fabric 的 Train Loop; fabric 为必选参数，如果不想用 fabric 可以自行组合 train_epoch etc.

    Fabric 需要在创建之前 launch!

    依赖注入
    ----------
    可以将 `training_step` 和 `validation_step` 函数注入到 `"train"` namespace 下。
    如果不注入，使用默认的训练和验证函数。函数需满足如下签名:
    ```py
    # 返回中应当至少包含 {'loss': float}
    fn(ctx: FabricTrainingContext, batch, fn_loss=None, batch_idx, epoch_idx) -> tuple[torch.Tensor, dict[str, float]]
    ```

    Parameters
    ----------
    - `context` : `procedure.FabricTrainingContext`; 必须包含 fabric!
        - 如果 context.fn_loss 设置了损失函数，则可以作为默认损失函数使用
    - `train_loader` : `DataLoader`
        - 训练数据加载器
    - `fn_loss` : `Optional[nn.Module]`, optional
        - 损失函数，如果为 None 且 context.fn_loss 也为 None，则会抛出错误
    - `max_epochs` : `int`
        - 最大训练轮数
    - `val_loader` : `Optional[DataLoader]`
        - 验证数据加载器，可选
    - `val_interval` : `int`, optional
        - 验证间隔, by default 1
    - `lr_scheduler` : `Optional[torch.optim.lr_scheduler._LRScheduler]`, optional
        - 学习率调度器, by default None
        - 为了向后兼容，如果传入但没有注册 LRSchedulerCallback，会自动创建 epoch-level 的调度器
        - 推荐直接使用 LRSchedulerCallback 获得更灵活的控制
    - `lr_sched_freq` : `int`, optional
        - 学习率调度频率（向后兼容参数）, by default 1
        - 仅在自动创建 LRSchedulerCallback 时使用
    - `fn_train_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个训练步骤的函数, 默认为 `training_step`
    - `fn_val_step` : `(...) -> Dict[str, Any]`, optional
        - 可选项, 用于计算单个验证步骤的函数, 默认为 `validation_step`
    - `fn_batch2target` : `(...) -> torch.Tensor`, optional
        - 可选项, 用于从 batch 中提取预测值, 默认为 `_default_batch_2_target`
        - 例如默认的方法就是接收一个 (X, y) 的 batch, 并返回 y
        - 对于复杂训练任务，例如 GAN 等，可能需要自定义这个函数
    - `fn_on_trainepoch_end` : `(...) -> None`, optional
        - 可选项, 用于在每个 epoch 结束时调用的函数, 默认为 `None`
        - 接受以下的关键字参数: `fn(ctx: FabricTrainingContext, epoch, train_metrics, val_metrics)`
        - 相比较注册 Callback 更加灵活一些
    - `reduce_lr_metric_name` : `str`, optional
        - 用于监控的指标名称（向后兼容参数）, 默认为 'loss'
        - 仅在自动创建 LRSchedulerCallback 时使用，用于 ReduceLROnPlateau 等调度器

    `fn_train_step` / `fn_val_step`:
    ----------
    - 默认不提供, 使用 `procedure.training_step` / `procedure.validation_step`; 如果提供了, 则使用提供的函数 (例如训练特殊模型, 需要非常规的步骤等)
    - 接受以下的关键字参数: `fn(ctx: FabricTrainingContext, batch, fn_loss=None, batch_idx, epoch_idx)`
        - `fn_loss` 参数现在是可选的，如果为 None 则使用 ctx.loss_fn
    - 返回值: `tuple[torch.Tensor, Dict[str, float]]` (predictions, metrics)
        - 例如默认的 `training_step` 返回 `(y_pred, {'loss': loss_value})`
        - 对于自定义的 step 函数, 如果返回来想要记录的指标, 请在 TorchMetricRecorder 的 scalar_metrics 中指定
    """
    # model, optimizer = fabric.setup(model, optimizer)
    max_epochs = typing.cast(int, max_epochs)
    fabric = context._require_fabric()
    train_loader = _setup_single_dataloader(context, train_loader)
    if val_loader:
        val_loader = _setup_single_dataloader(context, val_loader)

    # 向后兼容：如果传入了 lr_scheduler 但没有对应的 LRSchedulerCallback，自动注入;
    # 后续可能会移除此自动注入行为，请尽早显式注册 LRSchedulerCallback 以避免未来的不兼容。
    if 'lr_scheduler' in kwargs:
        lr_scheduler = typing.cast(Optional['_LRScheduler'], kwargs.get('lr_scheduler'))
        lr_scheduler_cb = callbacks.find_callback(fabric, callbacks.LRSchedulerCallback)
        if lr_scheduler_cb is None:
            # 自动创建并注入 LRSchedulerCallback
            auto_lr_callback = callbacks.LRSchedulerCallback(
                lr_scheduler=typing.cast('_LRScheduler', lr_scheduler),
                interval='epoch',  # 保持原有的 epoch-level 行为
                frequency=kwargs.get('lr_sched_freq') or 1,
                monitor_metric=kwargs.get('reduce_lr_metric_name') or 'loss',
                use_val_metric=True,
            )
            fabric._callbacks.append(auto_lr_callback)
            print_debug('自动注入 LRSchedulerCallback 以保持向后兼容')
    else:
        lr_scheduler = None

    # 向后兼容: 如果传入了 fn_loss, 则自动设置到 context.loss_fn 中; 后续可能会删掉
    if 'fn_loss' in kwargs:
        fn_loss = typing.cast(Optional[nn.Module], kwargs.get('fn_loss'))
        if context.fn_loss is None:
            context.fn_loss = fn_loss
            print_debug(f'自动设置 context.loss_fn: {context.fn_loss}')
        else:
            # 如果已经设置了 fn_loss, 则不再覆盖
            if context.fn_loss != fn_loss:
                print_warning(
                    'context.loss_fn 已经设置, 不会覆盖: '
                    f'当前为 {type(context.fn_loss).__name__}, '
                    f'传入的为 {type(fn_loss).__name__}'
                )
    else:
        # 如果没有传入 fn_loss, 则使用 context.loss_fn
        if context.fn_loss is None:
            if fn_train_step is None:
                raise ValueError(
                    '损失函数未指定：请在 FabricTrainingContext.loss_fn 中设置损失函数，'
                    '或在函数调用时传入 fn_loss 参数'
                )
            else:
                print_warning(
                    '损失函数未指定：如果并非有意偶数为，请在 FabricTrainingContext.loss_fn 中设置损失函数'
                )  # 也有可能是自己指定了 train_step

    context.call_on_train_start(
        lr_scheduler=lr_scheduler,
        max_epochs=max_epochs,
        dataloader=train_loader,
        hparams=hparams,
    )

    early_stopping_cb: EarlyStopping | None = callbacks.find_callback(
        fabric, callbacks.EarlyStopping
    )

    def is_early_stopped():
        if early_stopping_cb:
            return early_stopping_cb.stop_training
        return False

    for epoch in range(max_epochs):
        if is_early_stopped():
            print_info('Early stopping, 终止训练')
            break

        context.call_on_epoch_start(epoch=epoch, max_epochs=max_epochs)
        context.call_on_train_epoch_start(epoch=epoch, max_epochs=max_epochs)

        train_metrics = train_epoch(
            ctx=context,
            train_loader=train_loader,
            epoch_idx=epoch,
            fn_train_step=fn_train_step,
            fn_batch2target=fn_batch2target,
        )

        context.call_on_train_epoch_end(metrics=train_metrics, epoch=epoch)

        if val_loader and (epoch + 1) % val_interval == 0:
            context.call_on_validation_epoch_start(
                epoch=epoch, epoch_interval=val_interval
            )
            val_metrics = val_epoch(
                ctx=context,
                val_loader=val_loader,
                # fn_loss=fn_loss,
                epoch_idx=epoch,
                fn_val_step=fn_val_step,
                fn_batch2target=fn_batch2target,
            )
            context.call_on_validation_epoch_end(metrics=val_metrics, epoch=epoch)
        else:
            val_metrics = None

        # deprecated: 直接使用 LRSchedulerCallback
        # if lr_scheduler:
        #     if (epoch + 1) % lr_sched_freq == 0:
        #         # loss = val_metrics['loss'] if val_metrics else train_metrics['loss']
        #         # step_lr_sched(lr_scheduler, loss)
        #         metric_name = reduce_lr_metric_name or 'loss'
        #         metric = (
        #             val_metrics.get(metric_name)
        #             if val_metrics
        #             else train_metrics.get(metric_name)
        #         )
        #         step_lr_sched(lr_scheduler, metric)

        context.call_on_epoch_end(
            epoch=epoch, train_metrics=train_metrics, val_metrics=val_metrics
        )

        if fn_on_trainepoch_end:
            fn_on_trainepoch_end(
                ctx=context,
                epoch=epoch,
                train_metrics=train_metrics,
                val_metrics=val_metrics,
            )

    context.call_on_train_end()
