# -*- coding: utf-8 -*-
"""正演代理模型训练流程。

本模块读取 preprocess.py 生成的 joblib 数据，构造 PyTorch Dataset/DataLoader，训练
ForwardSurrogate，并按验证集损失保存最佳 checkpoint。损失函数不是单一 MSE，而是
由压力、导数、均值偏置、导数形状约束和可选自动拟合目标组成，目的是让
模型既能拟合点值，也能保持对自动试井拟合有意义的曲线形态。

训练过程会保存 history.json、metrics.json 和 forward_surrogate_best.pt，后续评估
脚本可以根据 checkpoint 中保存的维度、curve_layout 和损失权重恢复模型。
"""

# pylint: disable=import-error,duplicate-code,too-many-instance-attributes
# pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals

from __future__ import annotations

import json
import random
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

import joblib
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset

from src.models.forward_surrogate import ForwardSurrogate, ForwardSurrogateConfig


METRIC_KEYS = (
    "loss",
    "loss_pressure",
    "loss_derivative",
    "loss_bias_pressure",
    "loss_bias_derivative",
    "loss_derivative_shape",
    "loss_autofit_pressure",
    "loss_autofit_derivative",
    "sample_weight_mean",
    "sample_weight_max",
)


class ForwardDataset(Dataset):
    """把预处理后的参数、流量制度和曲线数组封装成 PyTorch Dataset。"""

    def __init__(self, params_x: np.ndarray, schedule_x: np.ndarray, curve_y: np.ndarray):
        """把三个 numpy 数组转为 float32 张量，后续 DataLoader 可直接按样本读取。"""
        self.params_x = torch.tensor(params_x, dtype=torch.float32)
        self.schedule_x = torch.tensor(schedule_x, dtype=torch.float32)
        self.curve_y = torch.tensor(curve_y, dtype=torch.float32)

    def __len__(self) -> int:
        """返回数据集或容器中可迭代样本的数量。"""
        return len(self.params_x)

    def __getitem__(self, idx: int):
        """按索引取出一个训练样本或数据项。"""
        return self.params_x[idx], self.schedule_x[idx], self.curve_y[idx]


@dataclass
class ModelConfig:
    """模型结构相关配置。"""

    hidden_dim: int = 128
    dropout: float = 0.0
    use_schedule: bool = True


@dataclass
class OptimConfig:
    """优化器与训练轮次配置。"""

    batch_size: int = 128
    epochs: int = 100
    lr: float = 1e-3
    weight_decay: float = 1e-5


@dataclass
class LossWeights:
    """复合损失的各项权重。"""

    pressure: float = 1.0
    derivative: float = 2.0
    bias_pressure: float = 0.15
    bias_derivative: float = 0.05
    derivative_shape: float = 0.10
    autofit_pressure: float = 0.0
    autofit_derivative: float = 0.0


@dataclass
class LossConfig:
    """损失函数配置。"""

    weights: LossWeights = field(default_factory=LossWeights)
    use_huber: bool = True
    huber_beta: float = 0.05


@dataclass
class SampleReweightConfig:
    """样本重加权配置。"""

    enabled: bool = True
    alpha: float = 0.4
    weight_min: float = 1.0
    weight_max: float = 2.5


@dataclass
class TrainRuntime:
    """训练运行时配置。"""

    seed: int = 42
    device: str = "cuda" if torch.cuda.is_available() else "cpu"


@dataclass
class TrainConfig:
    """正演代理模型训练配置。"""

    processed_path: Path
    output_dir: Path
    runtime: TrainRuntime = field(default_factory=TrainRuntime)
    optim: OptimConfig = field(default_factory=OptimConfig)
    model: ModelConfig = field(default_factory=ModelConfig)
    loss: LossConfig = field(default_factory=LossConfig)
    sample_reweight: SampleReweightConfig = field(default_factory=SampleReweightConfig)


@dataclass
class CurveStats:
    """曲线 scaler 的 torch 形式统计量。"""

    mean_raw: torch.Tensor
    scale_raw: torch.Tensor


@dataclass
class LossBatchParts:
    """压力和导数的预测值与真实值。"""

    pred_p: torch.Tensor
    pred_d: torch.Tensor
    true_p: torch.Tensor
    true_d: torch.Tensor


@dataclass
class LossContext:
    """计算复合损失所需的上下文。"""

    slices: dict[str, slice]
    curve_stats: CurveStats
    loss_cfg: LossConfig
    reweight_cfg: SampleReweightConfig


@dataclass
class DatasetBundle:
    """训练、验证、测试 DataLoader 与数据维度。"""

    train_loader: DataLoader
    val_loader: DataLoader
    test_loader: DataLoader
    param_dim: int
    schedule_dim: int
    curve_dim: int


def set_global_seed(seed: int) -> None:
    """设置 Python、NumPy 和 PyTorch 随机种子，并在 CUDA 可用时同步设置 GPU 随机种子。"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def load_processed_dataset(path: Path) -> dict:
    """读取预处理后的 joblib 数据，并检查训练、验证、测试三套数组是否齐全。"""
    data = joblib.load(path)
    required_keys = [
        "X_params_train", "X_schedule_train", "Y_curve_train",
        "X_params_val", "X_schedule_val", "Y_curve_val",
        "X_params_test", "X_schedule_test", "Y_curve_test",
    ]
    for key in required_keys:
        if key not in data:
            raise KeyError(f"processed dataset 缺少字段: {key}")
    return data


def infer_curve_layout(data: dict) -> dict:
    """读取双通道布局，并拒绝未经重新预处理的旧三通道数据。"""
    meta = data.get("meta", {})
    curve_dim = int(meta.get("curve_dim", data["Y_curve_train"].shape[1]))
    curve_layout = meta.get("curve_layout")

    if curve_layout is not None:
        names = {str(part["name"]) for part in curve_layout.get("parts", [])}
        if "slope" in names:
            raise ValueError(
                "processed 数据仍包含旧版 slope 通道；请先重新运行 "
                "scripts/preprocess_dataset.py，预处理会自动裁掉 slope"
            )
        return curve_layout

    if curve_dim % 2 != 0:
        raise ValueError(f"curve_dim={curve_dim} 不能按压力/导数两段均分")

    n_time_points = curve_dim // 2
    return {
        "n_time_points": int(n_time_points),
        "parts": [
            {"name": "log_pressure", "start": 0, "end": n_time_points},
            {"name": "log_derivative", "start": n_time_points, "end": 2 * n_time_points},
        ],
    }


def get_part_slices(curve_layout: dict) -> dict[str, slice]:
    """把 curve_layout 中的 start/end 信息转换成各曲线分段的 slice。"""
    out: dict[str, slice] = {}
    for part in curve_layout["parts"]:
        name = str(part["name"])
        out[name] = slice(int(part["start"]), int(part["end"]))
    return out


def smooth_l1_per_sample(pred: torch.Tensor, target: torch.Tensor, beta: float) -> torch.Tensor:
    """按样本计算 Smooth L1 损失，返回每个样本一个损失值。"""
    diff = torch.abs(pred - target)
    loss = torch.where(diff < beta, 0.5 * diff * diff / beta, diff - 0.5 * beta)
    return loss.mean(dim=1)


def mse_per_sample(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """按样本计算均方误差。"""
    return ((pred - target) ** 2).mean(dim=1)


def l1_per_sample(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """按样本计算平均绝对误差。"""
    return torch.abs(pred - target).mean(dim=1)


def regression_per_sample(
    pred: torch.Tensor,
    target: torch.Tensor,
    loss_cfg: LossConfig,
) -> torch.Tensor:
    """按配置在 Smooth L1 和 MSE 之间切换点值损失。"""
    if loss_cfg.use_huber:
        return smooth_l1_per_sample(pred, target, beta=float(loss_cfg.huber_beta))
    return mse_per_sample(pred, target)


def first_diff(x: torch.Tensor) -> torch.Tensor:
    """计算曲线相邻时间点的一阶差分，用于约束导数形态。"""
    return x[:, 1:] - x[:, :-1]


def affine_restore(x_scaled: torch.Tensor, mean: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """在 torch 张量上执行 scaler 的反标准化公式 x * scale + mean。"""
    return x_scaled * scale + mean


def autofit_curve_objective_per_sample(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """用 torch 计算自动拟合风格的曲线误差，作为训练附加目标。"""
    weight_factor = torch.clamp(torch.abs(target) * 0.01, max=100.0)
    weight = 1.0 / (1.0 + weight_factor)
    scale = torch.maximum(
        torch.maximum(torch.abs(target), torch.abs(pred)),
        torch.full_like(target, 1e-12),
    )
    relative_error = torch.abs(target - pred) / scale
    absolute_error = torch.abs(target - pred)
    point_error = 0.7 * relative_error + 0.3 * absolute_error
    weighted_mse = (weight * (point_error**2)).sum(dim=1)
    weighted_mse = weighted_mse / torch.clamp(weight.sum(dim=1), min=1e-12)
    return torch.sqrt(weighted_mse)


def build_sample_weight(
    true_p: torch.Tensor,
    true_d: torch.Tensor,
    reweight_cfg: SampleReweightConfig,
) -> torch.Tensor:
    """根据真实曲线幅值构造样本权重，让高幅值样本训练时权重更高。"""
    p_level = true_p.abs().mean(dim=1)
    d_level = true_d.abs().mean(dim=1)

    p_norm = p_level / (p_level.mean().detach() + 1e-6)
    d_norm = d_level / (d_level.mean().detach() + 1e-6)

    raw = 0.5 * p_norm + 0.5 * d_norm
    weight = 1.0 + reweight_cfg.alpha * (raw - 1.0)
    return torch.clamp(weight, min=reweight_cfg.weight_min, max=reweight_cfg.weight_max)


def split_curve_parts(
    pred: torch.Tensor,
    target: torch.Tensor,
    slices: dict[str, slice],
) -> LossBatchParts:
    """把拼接曲线拆成压力和导数两段。"""
    return LossBatchParts(
        pred_p=pred[:, slices["log_pressure"]],
        pred_d=pred[:, slices["log_derivative"]],
        true_p=target[:, slices["log_pressure"]],
        true_d=target[:, slices["log_derivative"]],
    )


def compute_basic_loss_vectors(
    parts: LossBatchParts,
    loss_cfg: LossConfig,
) -> dict[str, torch.Tensor]:
    """计算标准化空间中的基础点值、偏置和导数形状损失。"""
    return {
        "loss_pressure": regression_per_sample(parts.pred_p, parts.true_p, loss_cfg),
        "loss_derivative": regression_per_sample(parts.pred_d, parts.true_d, loss_cfg),
        "loss_bias_pressure": l1_per_sample(
            parts.pred_p.mean(dim=1, keepdim=True),
            parts.true_p.mean(dim=1, keepdim=True),
        ),
        "loss_bias_derivative": l1_per_sample(
            parts.pred_d.mean(dim=1, keepdim=True),
            parts.true_d.mean(dim=1, keepdim=True),
        ),
        "loss_derivative_shape": regression_per_sample(
            first_diff(parts.pred_d),
            first_diff(parts.true_d),
            loss_cfg,
        ),
    }


def compute_autofit_loss_vectors(
    parts: LossBatchParts,
    context: LossContext,
) -> dict[str, torch.Tensor]:
    """在原始尺度上计算自动拟合风格损失。"""
    pressure_slice = context.slices["log_pressure"]
    derivative_slice = context.slices["log_derivative"]

    mean_p = context.curve_stats.mean_raw[pressure_slice].unsqueeze(0)
    scale_p = context.curve_stats.scale_raw[pressure_slice].unsqueeze(0)
    mean_d = context.curve_stats.mean_raw[derivative_slice].unsqueeze(0)
    scale_d = context.curve_stats.scale_raw[derivative_slice].unsqueeze(0)

    return {
        "loss_autofit_pressure": autofit_curve_objective_per_sample(
            affine_restore(parts.pred_p, mean_p, scale_p),
            affine_restore(parts.true_p, mean_p, scale_p),
        ),
        "loss_autofit_derivative": autofit_curve_objective_per_sample(
            affine_restore(parts.pred_d, mean_d, scale_d),
            affine_restore(parts.true_d, mean_d, scale_d),
        ),
    }


def weighted_total_vector(
    loss_vectors: dict[str, torch.Tensor],
    weights: LossWeights,
) -> torch.Tensor:
    """按配置权重合成每个样本的总损失向量。"""
    return (
        weights.pressure * loss_vectors["loss_pressure"]
        + weights.derivative * loss_vectors["loss_derivative"]
        + weights.bias_pressure * loss_vectors["loss_bias_pressure"]
        + weights.bias_derivative * loss_vectors["loss_bias_derivative"]
        + weights.derivative_shape * loss_vectors["loss_derivative_shape"]
        + weights.autofit_pressure * loss_vectors["loss_autofit_pressure"]
        + weights.autofit_derivative * loss_vectors["loss_autofit_derivative"]
    )


def compute_weighted_loss(
    pred: torch.Tensor,
    target: torch.Tensor,
    context: LossContext,
) -> dict[str, torch.Tensor]:
    """计算正演代理模型的复合训练损失。"""
    parts = split_curve_parts(pred, target, context.slices)
    loss_vectors = compute_basic_loss_vectors(parts, context.loss_cfg)
    loss_vectors.update(compute_autofit_loss_vectors(parts, context))

    total_vec = weighted_total_vector(loss_vectors, context.loss_cfg.weights)
    if context.reweight_cfg.enabled:
        sample_weight = build_sample_weight(parts.true_p, parts.true_d, context.reweight_cfg)
    else:
        sample_weight = torch.ones_like(total_vec)

    metrics = {key: value.mean() for key, value in loss_vectors.items()}
    metrics["loss"] = (total_vec * sample_weight).mean()
    metrics["sample_weight_mean"] = sample_weight.mean()
    metrics["sample_weight_max"] = sample_weight.max()
    return metrics


def model_forward(
    model: nn.Module,
    params_x: torch.Tensor,
    schedule_x: torch.Tensor,
    use_schedule: bool,
) -> torch.Tensor:
    """按 use_schedule 开关统一调用模型，兼容只用参数输入和参数+流量制度输入。"""
    if use_schedule:
        return model(params_x, schedule_x)
    return model(params_x, None)


def init_metric_accumulator() -> dict[str, float]:
    """创建指标累加器。"""
    return {key: 0.0 for key in METRIC_KEYS}


def accumulate_metrics(
    total: dict[str, float],
    losses: dict[str, torch.Tensor],
    batch_size: int,
) -> None:
    """按 batch 样本数加权累加指标。"""
    for key in total:
        total[key] += losses[key].item() * batch_size


def average_metrics(total: dict[str, float], total_n: int) -> dict[str, float]:
    """将累加指标转换为样本平均指标。"""
    denom = max(total_n, 1)
    return {key: value / denom for key, value in total.items()}


def run_loader_epoch(
    model: nn.Module,
    loader: DataLoader,
    device: str,
    context: LossContext,
    use_schedule: bool,
    optimizer: torch.optim.Optimizer | None = None,
) -> dict[str, float]:
    """执行一个训练或评估 epoch，并返回平均指标。"""
    is_train = optimizer is not None
    model.train(mode=is_train)

    total = init_metric_accumulator()
    total_n = 0
    grad_context = torch.enable_grad() if is_train else torch.no_grad()

    with grad_context:
        for params_x, schedule_x, curve_y in loader:
            params_x = params_x.to(device)
            schedule_x = schedule_x.to(device)
            curve_y = curve_y.to(device)

            if is_train:
                optimizer.zero_grad()

            pred = model_forward(model, params_x, schedule_x, use_schedule)
            losses = compute_weighted_loss(pred=pred, target=curve_y, context=context)

            if is_train:
                losses["loss"].backward()
                optimizer.step()

            batch_size = params_x.size(0)
            accumulate_metrics(total, losses, batch_size)
            total_n += batch_size

    return average_metrics(total, total_n)


def evaluate(
    model: nn.Module,
    loader: DataLoader,
    device: str,
    context: LossContext,
    use_schedule: bool,
) -> dict[str, float]:
    """在验证或测试 DataLoader 上计算平均损失和各损失分量。"""
    return run_loader_epoch(
        model=model,
        loader=loader,
        device=device,
        context=context,
        use_schedule=use_schedule,
        optimizer=None,
    )


def build_curve_stats(data: dict, device: str) -> CurveStats:
    """从预处理数据中的曲线 scaler 构建 torch 统计量。"""
    scaler_curve = data["scaler_curve"]
    curve_mean_raw = np.asarray(scaler_curve.mean_, dtype=np.float32).reshape(-1)
    curve_scale_raw = np.asarray(scaler_curve.scale_, dtype=np.float32).reshape(-1)
    return CurveStats(
        mean_raw=torch.tensor(curve_mean_raw, dtype=torch.float32, device=device),
        scale_raw=torch.tensor(curve_scale_raw, dtype=torch.float32, device=device),
    )


def build_dataloaders(data: dict, cfg: TrainConfig) -> DatasetBundle:
    """根据预处理数组构造训练、验证、测试 DataLoader。"""
    train_ds = ForwardDataset(data["X_params_train"], data["X_schedule_train"], data["Y_curve_train"])
    val_ds = ForwardDataset(data["X_params_val"], data["X_schedule_val"], data["Y_curve_val"])
    test_ds = ForwardDataset(data["X_params_test"], data["X_schedule_test"], data["Y_curve_test"])

    loader_generator = torch.Generator()
    loader_generator.manual_seed(int(cfg.runtime.seed))

    train_loader = DataLoader(
        train_ds,
        batch_size=cfg.optim.batch_size,
        shuffle=True,
        generator=loader_generator,
    )
    val_loader = DataLoader(val_ds, batch_size=cfg.optim.batch_size, shuffle=False)
    test_loader = DataLoader(test_ds, batch_size=cfg.optim.batch_size, shuffle=False)

    return DatasetBundle(
        train_loader=train_loader,
        val_loader=val_loader,
        test_loader=test_loader,
        param_dim=data["X_params_train"].shape[1],
        schedule_dim=data["X_schedule_train"].shape[1],
        curve_dim=data["Y_curve_train"].shape[1],
    )


def build_forward_model(
    model_cfg: ModelConfig,
    param_dim: int,
    schedule_dim: int,
    curve_dim: int,
    device: str,
) -> nn.Module:
    """兼容新版配置式 ForwardSurrogate 和旧版关键字参数式 ForwardSurrogate。"""
    surrogate_cfg = ForwardSurrogateConfig(
        param_dim=param_dim,
        schedule_dim=schedule_dim,
        curve_dim=curve_dim,
        hidden_dim=model_cfg.hidden_dim,
        dropout=model_cfg.dropout,
        use_schedule=model_cfg.use_schedule,
    )
    return ForwardSurrogate(surrogate_cfg).to(device)


def build_optimizer(model: nn.Module, optim_cfg: OptimConfig) -> torch.optim.Optimizer:
    """构建 Adam 优化器。"""
    return torch.optim.Adam(
        model.parameters(),
        lr=optim_cfg.lr,
        weight_decay=optim_cfg.weight_decay,
    )


def print_training_config(cfg: TrainConfig, curve_layout: dict) -> None:
    """打印训练配置摘要。"""
    weights = cfg.loss.weights
    reweight = cfg.sample_reweight
    print("训练配置:")
    print(f"  device={cfg.runtime.device}")
    print(f"  seed={cfg.runtime.seed}")
    print(
        f"  batch_size={cfg.optim.batch_size}, epochs={cfg.optim.epochs}, "
        f"lr={cfg.optim.lr}, weight_decay={cfg.optim.weight_decay}"
    )
    print(f"  hidden_dim={cfg.model.hidden_dim}, dropout={cfg.model.dropout}")
    print(f"  use_schedule={cfg.model.use_schedule}")
    print(
        f"  weights: pressure={weights.pressure}, derivative={weights.derivative}, "
        f"bias_p={weights.bias_pressure}, "
        f"bias_d={weights.bias_derivative}, d_shape={weights.derivative_shape}, "
        f"autofit_p={weights.autofit_pressure}, autofit_d={weights.autofit_derivative}"
    )
    print(
        f"  sample_reweight={reweight.enabled}, alpha={reweight.alpha}, "
        f"clip=[{reweight.weight_min}, {reweight.weight_max}]"
    )
    print(f"  curve_layout={curve_layout}")
    print("  note: 当前重点训练 pressure + derivative；可显式关闭 schedule 分支做固定制度对照")


def format_metric_line(epoch: int, train_metrics: dict[str, float], val_metrics: dict[str, float]) -> str:
    """格式化单个 epoch 的训练与验证指标。"""
    return (
        f"[Epoch {epoch:03d}] "
        f"train={train_metrics['loss']:.6f} "
        f"(p={train_metrics['loss_pressure']:.6f}, "
        f"d={train_metrics['loss_derivative']:.6f}, "
        f"bp={train_metrics['loss_bias_pressure']:.6f}, "
        f"bd={train_metrics['loss_bias_derivative']:.6f}, "
        f"ds={train_metrics['loss_derivative_shape']:.6f}, "
        f"ap={train_metrics['loss_autofit_pressure']:.6f}, "
        f"ad={train_metrics['loss_autofit_derivative']:.6f}, "
        f"wmean={train_metrics['sample_weight_mean']:.4f}, "
        f"wmax={train_metrics['sample_weight_max']:.4f}) "
        f"val={val_metrics['loss']:.6f} "
        f"(p={val_metrics['loss_pressure']:.6f}, "
        f"d={val_metrics['loss_derivative']:.6f}, "
        f"bp={val_metrics['loss_bias_pressure']:.6f}, "
        f"bd={val_metrics['loss_bias_derivative']:.6f}, "
        f"ds={val_metrics['loss_derivative_shape']:.6f}, "
        f"ap={val_metrics['loss_autofit_pressure']:.6f}, "
        f"ad={val_metrics['loss_autofit_derivative']:.6f}, "
        f"wmean={val_metrics['sample_weight_mean']:.4f}, "
        f"wmax={val_metrics['sample_weight_max']:.4f})"
    )


def format_final_line(test_metrics: dict[str, float]) -> str:
    """格式化最终测试集指标。"""
    return (
        f"[Final] test={test_metrics['loss']:.6f} "
        f"(p={test_metrics['loss_pressure']:.6f}, "
        f"d={test_metrics['loss_derivative']:.6f}, "
        f"bp={test_metrics['loss_bias_pressure']:.6f}, "
        f"bd={test_metrics['loss_bias_derivative']:.6f}, "
        f"ds={test_metrics['loss_derivative_shape']:.6f}, "
        f"ap={test_metrics['loss_autofit_pressure']:.6f}, "
        f"ad={test_metrics['loss_autofit_derivative']:.6f}, "
        f"wmean={test_metrics['sample_weight_mean']:.4f}, "
        f"wmax={test_metrics['sample_weight_max']:.4f})"
    )


def build_checkpoint_payload(
    model: nn.Module,
    bundle: DatasetBundle,
    cfg: TrainConfig,
    curve_layout: dict,
) -> dict[str, Any]:
    """构建 checkpoint 保存内容。"""
    return {
        "model_state_dict": model.state_dict(),
        "param_dim": bundle.param_dim,
        "schedule_dim": bundle.schedule_dim,
        "curve_dim": bundle.curve_dim,
        "hidden_dim": cfg.model.hidden_dim,
        "dropout": cfg.model.dropout,
        "use_schedule": cfg.model.use_schedule,
        "seed": int(cfg.runtime.seed),
        "curve_layout": curve_layout,
        "loss_weights": asdict(cfg.loss.weights),
        "sample_reweight": asdict(cfg.sample_reweight),
    }


def append_history_row(
    history: list[dict],
    epoch: int,
    train_metrics: dict[str, float],
    val_metrics: dict[str, float],
) -> None:
    """把当前 epoch 的指标写入 history 列表。"""
    row = {"epoch": epoch}
    row.update({f"train_{key}": float(value) for key, value in train_metrics.items()})
    row.update({f"val_{key}": float(value) for key, value in val_metrics.items()})
    history.append(row)


def save_json(path: Path, payload: dict | list) -> None:
    """保存 JSON 文件。"""
    with open(path, "w", encoding="utf-8") as file_obj:
        json.dump(payload, file_obj, ensure_ascii=False, indent=2)


def train_epochs(
    model: nn.Module,
    bundle: DatasetBundle,
    cfg: TrainConfig,
    context: LossContext,
    curve_layout: dict,
) -> tuple[float, Path, list[dict]]:
    """执行训练循环并保存最佳模型。"""
    optimizer = build_optimizer(model, cfg.optim)
    best_val = float("inf")
    best_path = cfg.output_dir / "forward_surrogate_best.pt"
    history: list[dict] = []

    for epoch in range(1, cfg.optim.epochs + 1):
        train_metrics = run_loader_epoch(
            model=model,
            loader=bundle.train_loader,
            device=cfg.runtime.device,
            context=context,
            use_schedule=cfg.model.use_schedule,
            optimizer=optimizer,
        )
        val_metrics = evaluate(
            model=model,
            loader=bundle.val_loader,
            device=cfg.runtime.device,
            context=context,
            use_schedule=cfg.model.use_schedule,
        )

        append_history_row(history, epoch, train_metrics, val_metrics)
        print(format_metric_line(epoch, train_metrics, val_metrics))

        if val_metrics["loss"] < best_val:
            best_val = val_metrics["loss"]
            torch.save(build_checkpoint_payload(model, bundle, cfg, curve_layout), best_path)
            print(f"  -> best model saved to: {best_path}")

    return best_val, best_path, history


def load_best_model(best_path: Path, device: str) -> nn.Module:
    """从 checkpoint 恢复最佳模型。"""
    checkpoint = torch.load(best_path, map_location=device)
    model_cfg = ModelConfig(
        hidden_dim=checkpoint["hidden_dim"],
        dropout=checkpoint["dropout"],
        use_schedule=checkpoint.get("use_schedule", True),
    )
    best_model = build_forward_model(
        model_cfg=model_cfg,
        param_dim=checkpoint["param_dim"],
        schedule_dim=checkpoint["schedule_dim"],
        curve_dim=checkpoint["curve_dim"],
        device=device,
    )
    best_model.load_state_dict(checkpoint["model_state_dict"])
    return best_model


def build_metrics_payload(
    best_val: float,
    test_metrics: dict[str, float],
    cfg: TrainConfig,
    curve_layout: dict,
) -> dict[str, Any]:
    """构建 metrics.json 内容。"""
    return {
        "best_val_loss": float(best_val),
        "test_metrics": {key: float(value) for key, value in test_metrics.items()},
        "use_schedule": cfg.model.use_schedule,
        "seed": int(cfg.runtime.seed),
        "loss_weights": asdict(cfg.loss.weights),
        "sample_reweight": asdict(cfg.sample_reweight),
        "curve_layout": curve_layout,
    }


def train_forward(cfg: TrainConfig) -> None:
    """训练完整曲线正演代理模型。"""
    cfg.output_dir.mkdir(parents=True, exist_ok=True)
    set_global_seed(int(cfg.runtime.seed))

    data = load_processed_dataset(cfg.processed_path)
    curve_layout = infer_curve_layout(data)
    context = LossContext(
        slices=get_part_slices(curve_layout),
        curve_stats=build_curve_stats(data, cfg.runtime.device),
        loss_cfg=cfg.loss,
        reweight_cfg=cfg.sample_reweight,
    )
    bundle = build_dataloaders(data, cfg)

    model = build_forward_model(
        model_cfg=cfg.model,
        param_dim=bundle.param_dim,
        schedule_dim=bundle.schedule_dim,
        curve_dim=bundle.curve_dim,
        device=cfg.runtime.device,
    )

    print_training_config(cfg, curve_layout)
    best_val, best_path, history = train_epochs(model, bundle, cfg, context, curve_layout)
    save_json(cfg.output_dir / "history.json", history)

    best_model = load_best_model(best_path, cfg.runtime.device)
    test_metrics = evaluate(
        model=best_model,
        loader=bundle.test_loader,
        device=cfg.runtime.device,
        context=context,
        use_schedule=cfg.model.use_schedule,
    )

    print(format_final_line(test_metrics))
    save_json(
        cfg.output_dir / "metrics.json",
        build_metrics_payload(best_val, test_metrics, cfg, curve_layout),
    )