src/wjad/train/trainer.py

"""两阶段训练器 + 梯度监控。

Stage 1 (Dense)：
  - MoE 全部专家加权（dense 模式）；
  - 路由温度初始 < 1（锐化），训练中线性升到 1；
  - DINOv3 冻结；
  - 中期开启运动学/内外参扰动，监督校准网络；
  - GradNorm 启用。

Stage 2 (Sparse)：
  - MoE 切 Top-3；
  - 路由温度退火完成；
  - DINOv3 解冻并采用 1/100 主干 LR；
  - GradNorm + PCGrad 同时启用。
"""

from __future__ import annotations

import logging
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from ..losses import (
    HungarianMatcher,
    action_nll,
    calibration_regularization,
    detection_losses,
    ego_traj_nll,
    moe_load_balance_and_boundary,
    object_traj_nll,
)
from ..model import E2EAVModel, E2EOutput
from .multitask import MultiTaskOptimizer, MultiTaskOptimizerConfig
from .schedule import build_scheduler

log = logging.getLogger(__name__)


class _NullContext:
    """空 context manager，用于 AMP 关闭时占位 autocast。"""

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


@dataclass
class TrainerConfig:
    """Trainer 超参数（与 ``configs/default.yaml`` 对齐）。"""

    total_steps: int = 100000
    warmup_steps: int = 1000
    base_lr: float = 2.0e-4
    min_lr: float = 1.0e-6
    weight_decay: float = 0.05
    grad_clip: float = 1.0
    log_interval: int = 20
    ckpt_interval: int = 1000
    stage1_steps: int = 60000
    stage1_perturb_start: int = 20000
    grad_monitor_threshold: float = 1e-7
    # === AMP / 混合精度 ===
    # "fp32" / "bf16" / "fp16"。默认 bf16（H100/A100 推荐，无需 GradScaler）。
    mixed_precision: str = "bf16"
    grad_accum_steps: int = 1
    # MoE
    moe_load_balance_weight: float = 0.01
    moe_boundary_weight: float = 0.001
    router_temp_init: float = 0.5
    router_temp_final: float = 1.0
    # 损失初始权重（GradNorm 自适应主任务 1-6）
    loss_giou_weight: float = 0.5
    loss_calib_weight: float = 0.1
    # MultiTask（GradNorm + PCGrad 在 Stage1/Stage2 全程启用——
    # 两阶段的 6 项主任务都存在尺度不均与梯度冲突，PCGrad 不应延迟到 Stage2）
    enable_gradnorm: bool = True
    enable_pcgrad: bool = True
    # 参数组
    dinov3_lr_mult_stage2: float = 0.01
    # 显存吃紧的设备（如 a10g-small）上可关闭 Stage2 DINOv3 解冻，保持冻结
    unfreeze_dinov3_at_stage2: bool = True
    backbone_lr_mult: float = 1.0
    calibration_lr_mult: float = 0.1
    head_lr_mult: float = 1.0
    gate_lr_mult: float = 0.1
    # 检查点：目录 + 可选同步到 Hub model repo
    output_dir: str | None = None
    hub_repo_id: str | None = None
    hub_repo_type: str = "model"


def _is_gate_param(name: str) -> bool:
    return ".gate." in name or name.endswith(".gate_proj.weight") or name.endswith(".gate_proj.bias")


def build_param_groups(model: E2EAVModel, base_lr: float, cfg: TrainerConfig, stage: int) -> list[dict]:
    """按模块归类参数为不同 LR 组。Stage1 时 DINOv3 lr=0。"""
    groups: dict[str, list[nn.Parameter]] = {
        "dinov3": [],
        "backbone": [],
        "calibration": [],
        "head": [],
        "gate": [],
        "other": [],
    }
    for name, p in model.named_parameters():
        if not p.requires_grad and stage == 1:
            continue
        if name.startswith("dinov3."):
            groups["dinov3"].append(p)
        elif name.startswith("backbone."):
            if _is_gate_param(name):
                groups["gate"].append(p)
            else:
                groups["backbone"].append(p)
        elif name.startswith("calib."):
            if _is_gate_param(name):
                groups["gate"].append(p)
            else:
                groups["calibration"].append(p)
        elif name.startswith("det_traj_head.") or name.startswith("ctrl_head."):
            groups["head"].append(p)
        else:
            groups["other"].append(p)

    dinov3_lr = base_lr * (cfg.dinov3_lr_mult_stage2 if stage == 2 else 0.0)
    return [
        {"params": groups["dinov3"], "lr": dinov3_lr, "name": "dinov3"},
        {"params": groups["backbone"], "lr": base_lr * cfg.backbone_lr_mult, "name": "backbone"},
        {"params": groups["calibration"], "lr": base_lr * cfg.calibration_lr_mult, "name": "calibration"},
        {"params": groups["head"], "lr": base_lr * cfg.head_lr_mult, "name": "head"},
        {"params": groups["gate"], "lr": base_lr * cfg.gate_lr_mult, "name": "gate"},
        {"params": groups["other"], "lr": base_lr, "name": "other"},
    ]


def grad_norm_per_module(model: nn.Module, threshold: float) -> dict[str, float]:
    """统计各顶层模块的 grad-norm，返回 dict（用于日志/告警）。

    跳过：
      - 没有任何 ``requires_grad=True`` 参数的模块（如冻结的 DINOv3、纯
        buffer 模块 RoPE）；
      - 空模块（参数计数为 0）。
    """
    summary: dict[str, float] = {}
    for name, child in model.named_children():
        params = list(child.parameters())
        if not params:
            continue
        if not any(p.requires_grad for p in params):
            # 整个模块被冻结 -> 不监控
            continue
        total = 0.0
        seen = 0
        for p in params:
            if p.grad is not None:
                total += float(p.grad.detach().norm().item()) ** 2
                seen += 1
        if seen == 0:
            continue
        n = math.sqrt(total)
        summary[name] = n
        if n < threshold:
            log.warning("[grad_monitor] %s grad_norm=%.3e < %.3e", name, n, threshold)
        if not math.isfinite(n):
            log.error("[grad_monitor] %s grad_norm is %s (NaN/Inf)", name, n)
    return summary


def compute_all_losses(
    model_out: E2EOutput,
    batch: dict,
    matcher: HungarianMatcher,
    num_classes: int,
    cfg: TrainerConfig,
    perturbation_residual: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]:
    """计算 8 项损失，返回字典。

    ``perturbation_residual``：扰动训练时给定的 ground-truth 残差，用于额外
    监督校准网络；正常训练为 None。
    """
    targets = batch["targets"]

    det_out = model_out.detection
    ctrl_out = model_out.control
    calib = model_out.calibration

    det_losses = detection_losses(
        cls_logits=det_out.cls_logits,
        box_mu=det_out.box3d_mu,
        box_log_sigma=det_out.box3d_log_sigma,
        isdyn_logit=det_out.is_dynamic_logit,
        targets=targets,
        matcher=matcher,
        num_classes=num_classes,
    )

    L_traj_obj = object_traj_nll(
        det_out.traj_mu,
        det_out.traj_log_sigma,
        det_losses.matched_indices,
        targets,
    )

    L_traj_ego = ego_traj_nll(
        ctrl_out.ego_traj_mu,
        ctrl_out.ego_traj_log_sigma,
        batch["ego_future"],
        valid=batch.get("ego_future_valid"),
    )

    # 全局动作 GT 通常没有；此处用 0 做占位（实际数据集需补齐）
    action_target = batch.get("action_target")
    if action_target is None:
        action_target = torch.zeros_like(ctrl_out.action_mu)
    L_ctrl = action_nll(
        ctrl_out.action_mu, ctrl_out.action_log_sigma, action_target
    ) + L_traj_ego  # 控制损失 = action + ego_traj 复用同一项的便利封装；trainer 视情况拆分

    # MoE / 校准 正则
    L_moe = moe_load_balance_and_boundary(
        model_out.backbone_out.moe_stats,
        load_balance_weight=cfg.moe_load_balance_weight,
        boundary_weight=cfg.moe_boundary_weight,
    )
    L_calib_reg = calibration_regularization(
        calib.ego_residual, calib.intr_residual, calib.extr_residual,
        l2_weight=1.0,
    )
    if perturbation_residual is not None:
        # 扰动训练：计算校准网络应该预测的 GT 残差与实际残差的 MSE
        actual = torch.cat(
            [calib.ego_residual.flatten(1), calib.intr_residual, calib.extr_residual],
            dim=-1,
        )
        L_calib_reg = L_calib_reg + 1.0 * (actual - perturbation_residual).pow(2).mean()

    return {
        "L_cls": det_losses.cls_loss,
        "L_box": det_losses.box_nll + cfg.loss_giou_weight * det_losses.giou_loss,
        "L_isdyn": det_losses.isdyn_loss,
        "L_traj_obj": L_traj_obj,
        "L_traj_ego": L_traj_ego,
        "L_ctrl": L_ctrl,
        "L_moe": L_moe,
        "L_calib": L_calib_reg,
    }


MAIN_TASK_KEYS = ["L_cls", "L_box", "L_isdyn", "L_traj_obj", "L_traj_ego", "L_ctrl"]
AUX_TASK_KEYS = ["L_moe", "L_calib"]


class Trainer:
    """端到端训练器。"""

    def __init__(
        self,
        model: E2EAVModel,
        cfg: TrainerConfig,
        num_classes: int = 22,
        device: str = "cuda",
    ) -> None:
        self.model = model.to(device)
        self.cfg = cfg
        self.num_classes = num_classes
        self.device = device
        self.matcher = HungarianMatcher()
        self.global_step = 0
        self._micro_step = 0  # 用于 grad_accum
        self._stage = 1
        self._build_optimizer()

        # === AMP 配置 ===
        # 仅在 device 为 cuda 时启用 autocast（CPU 上 bf16 也能跑但收益极小）。
        amp_dtype_map = {
            "fp32": None,
            "bf16": torch.bfloat16,
            "fp16": torch.float16,
        }
        self.amp_dtype = amp_dtype_map[cfg.mixed_precision]
        self.amp_enabled = self.amp_dtype is not None and "cuda" in str(device)
        # GradScaler 仅 fp16 需要；bf16 数值范围大无需 scaler
        self.scaler = (
            torch.amp.GradScaler("cuda")
            if (self.amp_enabled and self.amp_dtype == torch.float16)
            else None
        )

        # MoE 初始模式 = dense；Stage2 切 sparse
        self.model.backbone.set_moe_mode("dense")
        self.model.backbone.set_router_temperature(cfg.router_temp_init)

    # ---------- 优化器构建 ----------

    def _build_optimizer(self) -> None:
        cfg = self.cfg
        groups = build_param_groups(self.model, cfg.base_lr, cfg, stage=self._stage)
        self.optimizer = torch.optim.AdamW(groups, weight_decay=cfg.weight_decay, betas=(0.9, 0.95))
        self.scheduler = build_scheduler(
            self.optimizer,
            warmup_steps=cfg.warmup_steps,
            total_steps=cfg.total_steps,
            base_lr=cfg.base_lr,
            min_lr=cfg.min_lr,
        )

        # PCGrad 共享参数 = 主干最后的“共享瓶颈”：final_norm + 最后 1 层 MoE block。
        # 不把全部 DINOv3/Calib/Backbone 都纳入，否则 N 个任务 × full-grad 扁平副本会
        # 在 a10g-small 上瞬间 OOM（~600M 参数 × 6 任务 × 2 副本 ≈ 28 GB）。
        # 较前的层仍享受 GradNorm 自适应加权 + 共同求和的标准多任务训练。
        shared: list[nn.Parameter] = []
        last_moe = self.model.backbone.moe_layers[-1]
        for p in self.model.backbone.final_norm.parameters():
            if p.requires_grad:
                shared.append(p)
        for p in last_moe.parameters():
            if p.requires_grad:
                shared.append(p)
        # GradNorm 代理参数：取主干最后 LayerNorm 的 weight
        proxy = self.model.backbone.final_norm.weight
        mt_cfg = MultiTaskOptimizerConfig(
            enable_gradnorm=cfg.enable_gradnorm,
            enable_pcgrad=cfg.enable_pcgrad,
            gradnorm_alpha=1.5,
            gradnorm_lr=0.025,
            pcgrad_shuffle=True,
        )
        self.mto = MultiTaskOptimizer(
            num_main_tasks=len(MAIN_TASK_KEYS),
            shared_params=shared,
            gradnorm_proxy_param=proxy,
            cfg=mt_cfg,
        )
        # GradNormBalancer 是 nn.Module，需要把 raw_weights / initial_losses 缓冲
        # 移到 model 所在 device，否则与 losses (cuda) 设备不匹配。
        if self.mto.gradnorm is not None:
            self.mto.gradnorm.to(self.device)

    def _maybe_save_checkpoint(self) -> None:
        cfg = self.cfg
        if not cfg.output_dir or cfg.ckpt_interval <= 0:
            return
        if self.global_step <= 0 or self.global_step % cfg.ckpt_interval != 0:
            return
        od = Path(cfg.output_dir)
        od.mkdir(parents=True, exist_ok=True)
        ckpt_path = od / f"checkpoint-step{self.global_step}.pt"
        torch.save(
            {
                "step": self.global_step,
                "stage": self._stage,
                "model": self.model.state_dict(),
                "optimizer": self.optimizer.state_dict(),
            },
            ckpt_path,
        )
        log.info("[Trainer] checkpoint %s", ckpt_path)
        if not cfg.hub_repo_id:
            return
        try:
            from huggingface_hub import HfApi, create_repo

            create_repo(cfg.hub_repo_id, repo_type=cfg.hub_repo_type, exist_ok=True)
            api = HfApi()
            rel = f"checkpoints/{ckpt_path.name}"
            api.upload_file(
                path_or_fileobj=str(ckpt_path),
                path_in_repo=rel,
                repo_id=cfg.hub_repo_id,
                repo_type=cfg.hub_repo_type,
                commit_message=f"checkpoint step {self.global_step}",
            )
            log.info("[Trainer] uploaded %s -> %s", rel, cfg.hub_repo_id)
        except Exception as e:
            log.warning("[Trainer] Hub upload failed: %s", e)

    # ---------- 阶段切换 ----------

    def maybe_switch_stage(self) -> None:
        cfg = self.cfg
        if self._stage == 1 and self.global_step >= cfg.stage1_steps:
            log.info("[Trainer] -> Stage 2 (sparse MoE + DINOv3 finetune + PCGrad)")
            self._stage = 2
            # 1) MoE 切 sparse
            self.model.backbone.set_moe_mode("sparse")
            # 2) 路由温度退火完成
            self.model.backbone.set_router_temperature(cfg.router_temp_final)
            # 3) DINOv3 解冻（小显存设备可禁用）
            if cfg.unfreeze_dinov3_at_stage2:
                self.model.dinov3.unfreeze()
            # 4) 重建优化器（包含 DINOv3 参数）+ 启用 PCGrad
            self._build_optimizer()

    # ---------- 单步 ----------

    def train_step(self, batch: dict, rng: np.random.Generator) -> dict:
        cfg = self.cfg
        self.maybe_switch_stage()

        # 移到 device
        batch = {k: (v.to(self.device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
        # targets 是 list of dict，里面的 tensor 也移到 device
        if "targets" in batch and isinstance(batch["targets"], list):
            new_targets = []
            for t in batch["targets"]:
                new_targets.append({k: (v.to(self.device) if isinstance(v, torch.Tensor) else v) for k, v in t.items()})
            batch["targets"] = new_targets

        # 扰动注入（Stage1 中期开启）
        perturb_residual = None
        ego_input = batch["ego_6d"]
        intr_input = batch["intr_vec"]
        extr_input = batch["extr_6d"]
        if (
            self._stage == 1
            and self.global_step >= cfg.stage1_perturb_start
            and rng.uniform() < 0.5
        ):
            from ..data.transforms import perturb_kinematics
            ego_input, intr_input, extr_input, delta = perturb_kinematics(
                ego_input.cpu().clone(), intr_input.cpu().clone()[0], extr_input.cpu().clone()[0],
                translation_std_m=0.1, rotation_std_deg=0.5,
                intrinsic_std=0.005, extrinsic_std=0.005,
                rng=rng,
            )
            ego_input = ego_input.to(self.device)
            # intr/extr 是 [B,...] 而 perturb_kinematics 是单样本；这里为简洁仅扰动第 0 个样本
            # 实际生产中应 batched 实现
            intr_input = batch["intr_vec"].clone()
            extr_input = batch["extr_6d"].clone()
            # GT 残差：在 symlog 空间 = -delta（symlog 是非线性，这里用线性近似）
            perturb_residual = -delta.to(self.device).unsqueeze(0).expand(ego_input.shape[0], -1)

        # 前向（AMP autocast 仅包住 forward 与匹配/损失，反传由 PyTorch
        # 在 fp32 主梯度下完成；GradNorm/PCGrad 内的 autograd.grad 也在 fp32）
        ac_ctx = (
            torch.autocast(device_type="cuda", dtype=self.amp_dtype)
            if self.amp_enabled
            else _NullContext()
        )
        with ac_ctx:
            out = self.model(
                images=batch["images"],
                ego_6d_raw=ego_input,
                intr_raw=intr_input,
                extr_6d_raw=extr_input,
            )
            losses = compute_all_losses(
                out, batch, self.matcher, self.num_classes, cfg,
                perturbation_residual=perturb_residual,
            )

        # === 把损失提升到 fp32 以保证后续 GradNorm/PCGrad 数值稳定 ===
        main = torch.stack([losses[k].float() for k in MAIN_TASK_KEYS])
        aux = sum(losses[k].float() for k in AUX_TASK_KEYS)
        # 梯度累积：对累积步数取平均
        if cfg.grad_accum_steps > 1:
            main = main / cfg.grad_accum_steps
            aux = aux / cfg.grad_accum_steps

        # === 反传 ===
        if self._micro_step == 0:
            self.optimizer.zero_grad(set_to_none=True)

        all_params = [p for p in self.model.parameters() if p.requires_grad]
        if self.scaler is not None:
            # fp16 路径：GradScaler 不直接支持 PCGrad（需手动调度）；这里
            # 退化为标准 sum-backward。bf16 推荐路径无此限制。
            total = main.sum() + aux
            self.scaler.scale(total).backward()
            weights = torch.ones_like(main)
        else:
            total, weights = self.mto.backward(main, aux, all_params)

        self._micro_step += 1
        do_step = self._micro_step >= cfg.grad_accum_steps
        if not do_step:
            info_partial = {
                "step": self.global_step,
                "stage": self._stage,
                "total_loss": float(total),
                "weights": [float(w) for w in weights],
                "grad_norms": {},
            }
            for k, v in losses.items():
                info_partial[k] = float(v.detach())
            return info_partial

        # === 梯度裁剪 + 监控 + step ===
        if self.scaler is not None:
            self.scaler.unscale_(self.optimizer)
        grad_summary = grad_norm_per_module(self.model, cfg.grad_monitor_threshold)
        torch.nn.utils.clip_grad_norm_(all_params, max_norm=cfg.grad_clip)

        if self.scaler is not None:
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            self.optimizer.step()
        self.scheduler.step()
        self._micro_step = 0

        # 路由温度线性退火
        if self._stage == 1:
            ratio = min(1.0, self.global_step / max(1, cfg.stage1_steps))
            t = cfg.router_temp_init + ratio * (cfg.router_temp_final - cfg.router_temp_init)
            self.model.backbone.set_router_temperature(t)

        self.global_step += 1
        self._maybe_save_checkpoint()

        info = {
            "step": self.global_step,
            "stage": self._stage,
            "total_loss": float(total),
            "weights": [float(w) for w in weights],
            "grad_norms": grad_summary,
        }
        for k, v in losses.items():
            info[k] = float(v.detach())
        return info

    def fit(self, loader: DataLoader, max_steps: int | None = None) -> None:
        """简化训练循环。"""
        rng = np.random.default_rng(0)
        steps = max_steps or self.cfg.total_steps
        it = iter(loader)
        for _ in range(steps):
            try:
                batch = next(it)
            except StopIteration:
                it = iter(loader)
                batch = next(it)
            info = self.train_step(batch, rng)
            if info["step"] % self.cfg.log_interval == 0:
                log.info(
                    "step=%d stage=%d total=%.4f cls=%.4f box=%.4f isdyn=%.4f traj_obj=%.4f traj_ego=%.4f ctrl=%.4f moe=%.4f calib=%.4f",
                    info["step"], info["stage"], info["total_loss"],
                    info["L_cls"], info["L_box"], info["L_isdyn"],
                    info["L_traj_obj"], info["L_traj_ego"], info["L_ctrl"],
                    info["L_moe"], info["L_calib"],
                )