scripts/train.py

#!/usr/bin/env python
# coding: utf-8
"""
AIFS v1.1 — Pre-training Script
=================================
Training logic strictly follows config_pretraining.yaml and anemoi-training 0.4.0.

Usage:
    python scripts/train.py
    python scripts/train.py --config conf/config.yaml
"""

from __future__ import annotations

import argparse
import datetime
import math
import os
import sys
import time
import traceback
import warnings
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import pytz
import torch
import torch.nn as nn
from tqdm import tqdm
import yaml

# Project root (scripts/ → aifs_v11/)
ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT))

import onescience.datapipes.climate.era5 as onescience_era5
import earthkit.regrid as ekr

from model.aifs import AIFS

# ============================================================================
# Config loader
# ============================================================================

def load_config(path: str) -> dict:
    with open(path) as f:
        return yaml.safe_load(f)

# ============================================================================
# Variable list builder
# ============================================================================

def build_era5_variable_list(cfg: dict) -> List[str]:
    av = cfg["aifs_variables"]
    em = cfg["era5_mapping"]
    vars_: List[str] = []
    for name in av["surface"]:
        vars_.append(em["surface"][name])
    for name in av["soil"]:
        vars_.append(em["soil"][name])
    for v in av["pressure_level"]:
        tpl = em["pressure_level"][v]
        for lvl in av["pressure_levels"]:
            vars_.append(tpl.format(level=lvl))
    for name in av["diagnostic"]:
        vars_.append(em["diagnostic"][name])
    return vars_


def _auto_correct_variable_names(required: List[str], available: set, cfg: dict) -> List[str]:
    fuzzy = cfg["era5_mapping"].get("fuzzy_fixes", {})
    corrected = list(required)
    for i, v in enumerate(corrected):
        if v not in available and v in fuzzy:
            alt = fuzzy[v]
            if alt in available:
                print(f"[INFO] Auto-corrected: '{v}' → '{alt}'")
                corrected[i] = alt
    return corrected


def _era5_to_aifs_name(era5_name: str, cfg: dict) -> Optional[str]:
    em = cfg["era5_mapping"]
    for aifs, ename in em["surface"].items():
        if ename == era5_name: return aifs
    for aifs, ename in em["soil"].items():
        if ename == era5_name: return aifs
    for v in cfg["aifs_variables"]["pressure_level"]:
        tpl = em["pressure_level"][v]
        for lvl in cfg["aifs_variables"]["pressure_levels"]:
            if tpl.format(level=lvl) == era5_name: return f"{v}_{lvl}"
    for aifs, ename in em["diagnostic"].items():
        if ename == era5_name: return aifs
    return None

# ============================================================================
# N320 interpolation
# ============================================================================

def _interp_n320(arr: np.ndarray) -> np.ndarray:
    return ekr.interpolate(arr, {"grid": (0.25, 0.25)}, {"grid": "N320"})

# ============================================================================
# ERA5 → AIFS field conversion
# ============================================================================

def build_aifs_fields_from_frames(
    frame_tm6: np.ndarray, frame_t0: np.ndarray, frame_tp6: np.ndarray,
    era5_var_list: List[str], cfg: dict,
    diag_map: Optional[Dict[str, str]] = None,
) -> Dict[str, np.ndarray]:
    av = cfg["aifs_variables"]
    em = cfg["era5_mapping"]
    name_to_ch = {n: i for i, n in enumerate(era5_var_list)}
    fields: Dict[str, np.ndarray] = {}

    for aifs_name in av["surface"]:
        ch = name_to_ch[em["surface"][aifs_name]]
        fields[aifs_name] = np.stack([_interp_n320(frame_tm6[ch]),
                                       _interp_n320(frame_t0[ch]),
                                       _interp_n320(frame_tp6[ch])])
    for aifs_name in av["soil"]:
        ch = name_to_ch[em["soil"][aifs_name]]
        fields[aifs_name] = np.stack([_interp_n320(frame_tm6[ch]),
                                       _interp_n320(frame_t0[ch]),
                                       _interp_n320(frame_tp6[ch])])
    for aifs_var in av["pressure_level"]:
        tpl = em["pressure_level"][aifs_var]
        for level in av["pressure_levels"]:
            ch = name_to_ch[tpl.format(level=level)]
            fields[f"{aifs_var}_{level}"] = np.stack([_interp_n320(frame_tm6[ch]),
                                                       _interp_n320(frame_t0[ch]),
                                                       _interp_n320(frame_tp6[ch])])
    _diag_map = diag_map if diag_map is not None else em["diagnostic"]
    for aifs_name in av["diagnostic"]:
        era5_name = _diag_map.get(aifs_name)
        if era5_name and era5_name in name_to_ch:
            ch = name_to_ch[era5_name]
            fields[aifs_name] = np.stack([_interp_n320(frame_tm6[ch]),
                                           _interp_n320(frame_t0[ch]),
                                           _interp_n320(frame_tp6[ch])])
    return fields

# ============================================================================
# Forcing feature computation
# ============================================================================

def compute_forcing_features(
    latitudes: np.ndarray, longitudes: np.ndarray,
    timestamps: List[datetime.datetime],
    fields: Dict[str, np.ndarray],
) -> Dict[str, np.ndarray]:
    lat_rad = np.deg2rad(latitudes)
    lon_rad = np.deg2rad(longitudes)
    multi_step = len(timestamps)
    forcing: Dict[str, np.ndarray] = {}

    forcing["cos_latitude"] = np.tile(np.cos(lat_rad)[np.newaxis, :], (multi_step, 1))
    forcing["sin_latitude"] = np.tile(np.sin(lat_rad)[np.newaxis, :], (multi_step, 1))
    forcing["cos_longitude"] = np.tile(np.cos(lon_rad)[np.newaxis, :], (multi_step, 1))
    forcing["sin_longitude"] = np.tile(np.sin(lon_rad)[np.newaxis, :], (multi_step, 1))

    cos_jd = np.zeros((multi_step, len(latitudes)), dtype=np.float32)
    sin_jd = np.zeros_like(cos_jd)
    cos_lt = np.zeros_like(cos_jd)
    sin_lt = np.zeros_like(cos_jd)
    insol = np.zeros_like(cos_jd)

    for t_idx, ts in enumerate(timestamps):
        doy = ts.timetuple().tm_yday
        jd_angle = 2.0 * np.pi * doy / 365.25
        cos_jd[t_idx] = np.cos(jd_angle)
        sin_jd[t_idx] = np.sin(jd_angle)
        hours = ts.hour + ts.minute / 60.0 + ts.second / 3600.0
        lt_angle = 2.0 * np.pi * hours / 24.0 + lon_rad
        cos_lt[t_idx] = np.cos(lt_angle)
        sin_lt[t_idx] = np.sin(lt_angle)
        insol[t_idx] = _compute_insolation(ts, lat_rad, lon_rad)

    forcing["cos_julian_day"] = cos_jd
    forcing["sin_julian_day"] = sin_jd
    forcing["cos_local_time"] = cos_lt
    forcing["sin_local_time"] = sin_lt
    forcing["insolation"] = insol

    for var_name in ["lsm", "z", "slor", "sdor"]:
        if var_name in fields:
            forcing[var_name] = np.tile(fields[var_name][:1], (multi_step, 1))
    return forcing


def _compute_insolation(ts: datetime.datetime, lat_rad: np.ndarray, lon_rad: np.ndarray) -> np.ndarray:
    ref = datetime.datetime(2000, 1, 1, 12, 0, 0, tzinfo=pytz.utc)
    mt = ts.timestamp()
    days = (mt - ref.timestamp()) / (24.0 * 3600.0)
    jc = days / 36525.0
    theta = (67310.54841 + jc * (876600.0 * 3600.0 + 8640184.812866
             + jc * (0.093104 - jc * 6.2e-5)))
    gmst = np.fmod((theta / 240.0) * np.pi / 180.0, 2.0 * np.pi)
    ma = np.deg2rad(357.52910 + 35999.05030 * jc - 0.0001559 * jc**2 - 0.00000048 * jc**3)
    ml = np.deg2rad(280.46645 + 36000.76983 * jc + 0.0003032 * jc**2)
    dl = np.deg2rad((1.914600 - 0.004817 * jc - 0.000014 * jc**2) * np.sin(ma)
                     + (0.019993 - 0.000101 * jc) * np.sin(2.0 * ma)
                     + 0.000290 * np.sin(3.0 * ma))
    tl = ml + dl
    eps = np.deg2rad(23.0 + 26.0 / 60.0 + 21.406 / 3600.0
                     - (46.836769 * jc - 0.0001831 * jc**2 + 0.00200340 * jc**3
                        - 0.576e-6 * jc**4 - 4.34e-8 * jc**5) / 3600.0)
    x = np.cos(tl); y = np.cos(eps) * np.sin(tl); z = np.sin(eps) * np.sin(tl)
    r = np.sqrt(1.0 - z * z)
    dec = np.arctan2(z, r)
    ra = 2.0 * np.arctan2(y, x + r)
    ha = (gmst + lon_rad) - ra
    cos_z = np.sin(lat_rad) * np.sin(dec) + np.cos(lat_rad) * np.cos(dec) * np.cos(ha)
    return cos_z.astype(np.float32)

# ============================================================================
# Normalisation
# ============================================================================

def compute_normalisation_params(
    cfg: dict, variable_names: List[str],
    statistics_path: Optional[str] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute per-variable norm_mul and norm_add from config.

    If ``statistics_path`` points to a ``.npz`` file with keys
    ``mean`` and ``stdev`` (per-variable arrays matching the full
    dataset variable order), those values are used.  Otherwise
    identity (mean=0, stdev=1) is assumed — correct for fake data
    but NOT for real ERA5.
    """
    nc = cfg["normalizer"]
    default_method = nc["default"]
    remap = nc.get("remap", {}) or {}
    methods: Dict[str, str] = {}
    for v in variable_names:
        if v in (nc.get("none") or []): methods[v] = "none"
        elif v in (nc.get("max") or []): methods[v] = "max"
        elif v in (nc.get("min-max") or []): methods[v] = "min-max"
        elif v in (nc.get("std") or []): methods[v] = "std"
        else: methods[v] = default_method

    name_to_idx = {n: i for i, n in enumerate(variable_names)}
    num_vars = len(variable_names)

    # Load pre-computed statistics if available, else identity
    _mean = np.zeros(num_vars, dtype=np.float32)
    _stdev = np.ones(num_vars, dtype=np.float32)

    if statistics_path and os.path.exists(statistics_path):
        stats = np.load(statistics_path)
        # statistics file stores arrays indexed by its own variable list
        if "variables" in stats:
            ds_vars = [str(v) for v in stats["variables"]]
        else:
            # Fallback: assume 115-variable dataset order from aifs_config
            import json
            cfg_path = ROOT / "model" / "aifs_config.json"
            ds_vars = json.load(open(cfg_path))["dataset"]["variables"]
        ds_name_to_idx = {n: i for i, n in enumerate(ds_vars)}
        for name, i in name_to_idx.items():
            if name in ds_name_to_idx:
                j = ds_name_to_idx[name]
                _mean[i] = float(stats["mean"][j])
                _stdev[i] = float(stats["stdev"][j])
        print(f"[INFO] Loaded statistics from {statistics_path}")
    else:
        print("[INFO] No statistics file — using identity normalisation "
              "(mean=0, stdev=1).  Fine for fake data, "
              "WRONG for real ERA5.")

    for target_var, source_var in remap.items():
        if target_var in name_to_idx and source_var in name_to_idx:
            ti, si = name_to_idx[target_var], name_to_idx[source_var]
            _mean[ti], _stdev[ti] = _mean[si], _stdev[si]

    norm_mul = np.ones(num_vars, dtype=np.float32)
    norm_add = np.zeros(num_vars, dtype=np.float32)
    for name, i in name_to_idx.items():
        method = methods.get(name, default_method)
        if method == "mean-std":
            norm_mul[i] = 1.0 / max(float(_stdev[i]), 1e-9)
            norm_add[i] = -float(_mean[i]) / max(float(_stdev[i]), 1e-9)
        elif method == "std":
            norm_mul[i] = 1.0 / max(float(_stdev[i]), 1e-9)
        elif method in ("max", "min-max", "none"):
            norm_mul[i] = 1.0
    return norm_mul, norm_add


def compute_era5_statistics(
    data_dir: str, years: List[int],
    variable_names: List[str],
    output_path: str,
) -> None:
    """Compute per-variable mean and stdev from ERA5 H5 files.

    Saves ``mean`` and ``stdev`` arrays (indexed by *variable_names*
    order) to *output_path* as a ``.npz`` file.

    This is a one-time pre-computation step before training with
    real ERA5 data.  The resulting file is consumed by
    ``compute_normalisation_params(statistics_path=...)``.
    """
    import h5py

    num_vars = len(variable_names)
    sum_x = np.zeros(num_vars, dtype=np.float64)
    sum_x2 = np.zeros(num_vars, dtype=np.float64)
    count = 0

    for year in years:
        h5_path = os.path.join(data_dir, "data", f"{year}.h5")
        if not os.path.exists(h5_path):
            print(f"[WARN] Missing {h5_path}, skipping")
            continue

        with h5py.File(h5_path, "r") as f:
            ds = f["fields"]
            h5_vars = [v.decode() if isinstance(v, bytes) else v
                       for v in ds.attrs["variables"]]
            T = ds.shape[0]

            # Index mapping for each timestep
            for t in range(T):
                frame = ds[t]  # (C, H, W)
                for i, vname in enumerate(variable_names):
                    if vname in h5_vars:
                        ch = h5_vars.index(vname)
                        vals = frame[ch].astype(np.float64)
                        sum_x[i] += vals.mean()
                        sum_x2[i] += (vals ** 2).mean()
                count += 1

    mean = (sum_x / count).astype(np.float32)
    # stdev = sqrt(E[X²] - E[X]²)
    var = np.maximum(sum_x2 / count - mean.astype(np.float64) ** 2, 0)
    stdev = np.sqrt(var).astype(np.float32)

    np.savez(output_path, mean=mean, stdev=stdev)
    print(f"[INFO] Statistics saved to {output_path} "
          f"({count} frames, {num_vars} vars)")

# ============================================================================
# Loss scaling
# ============================================================================

def build_loss_scale_factors(cfg: dict, output_var_names: List[str]) -> np.ndarray:
    sc = cfg["loss_scaling"]
    default = sc["default"]; pl = sc["pl"]; sfc = sc["sfc"]
    av = cfg["aifs_variables"]
    name_to_scale: Dict[str, float] = {}
    for var_short in av["pressure_level"]:
        base = pl.get(var_short, default)
        for lvl in av["pressure_levels"]:
            name_to_scale[f"{var_short}_{lvl}"] = base
    for var_name in av["surface"] + av["soil"] + av["diagnostic"]:
        name_to_scale[var_name] = sfc.get(var_name, default)
    for k, v in sfc.items():
        if k not in name_to_scale: name_to_scale[k] = v
    return np.array([name_to_scale.get(n, default) for n in output_var_names], dtype=np.float32)


def build_pressure_level_scaler(cfg: dict, output_var_names: List[str]) -> np.ndarray:
    pls = cfg["loss_scaling"]["pressure_level_scaler"]
    minimum, slope = pls["minimum"], pls["slope"]
    av = cfg["aifs_variables"]
    factors = np.ones(len(output_var_names), dtype=np.float32)
    for i, name in enumerate(output_var_names):
        for aifs_var in av["pressure_level"]:
            if name.startswith(f"{aifs_var}_"):
                plev_hpa = int(name[len(aifs_var) + 1:])
                factors[i] = max(minimum, slope * plev_hpa)
                break
    return factors

# ============================================================================
# Dataset
# ============================================================================

class AIFSTrainingDataset:
    def __init__(self, dataset_path: str, years: List[int], cfg: dict,
                 model_input_order: List[str], model_output_order: List[str],
                 norm_mul: np.ndarray, norm_add: np.ndarray,
                 input_steps=2, output_steps=1,
                 grid_lat=None, grid_lon=None):
        dp = dataset_path.rstrip("/")
        if os.path.isdir(os.path.join(dp, "data")): era5_root = dp
        elif os.path.isdir(dp) and any(f.endswith(".h5") for f in os.listdir(dp)): era5_root = os.path.dirname(dp)
        else: era5_root = dp
        self._era5_root = era5_root

        data_dir = os.path.join(era5_root, "data")
        h5_files = sorted([f for f in os.listdir(data_dir) if f.endswith(".h5")])
        era5_var_list = build_era5_variable_list(cfg)
        available_set: set = set()

        if h5_files:
            import h5py
            with h5py.File(os.path.join(data_dir, h5_files[0]), "r") as f:
                available_vars = [v.decode() if isinstance(v, bytes) else v for v in f["fields"].attrs["variables"]]
            available_set = set(available_vars)
            corrected = _auto_correct_variable_names(era5_var_list, available_set, cfg)
            used_vars = [v for v in corrected if v in available_set]
            missing = [v for v in corrected if v not in available_set]
            if missing:
                print(f"[WARN] {len(missing)}/{len(era5_var_list)} vars missing from H5: {missing[:10]}")
        else:
            used_vars = era5_var_list

        self._cfg = cfg
        self.era5_var_list = used_vars
        self.model_input_order = model_input_order
        self.model_output_order = model_output_order
        self.norm_mul = norm_mul; self.norm_add = norm_add
        self.input_steps = input_steps; self.output_steps = output_steps
        self._name_to_input_idx = {n: i for i, n in enumerate(model_input_order)}
        self._name_to_output_idx = {n: i for i, n in enumerate(model_output_order)}

        # corrected diagnostic map
        self._diag_map: Dict[str, str] = {}
        for aifs_name, orig in cfg["era5_mapping"]["diagnostic"].items():
            cn = _auto_correct_variable_names([orig], available_set, cfg)[0] if h5_files else orig
            self._diag_map[aifs_name] = cn if (h5_files and cn in available_set) else orig

        self._dataset = onescience_era5.ERA5Dataset(
            dataset_dir=era5_root, used_years=years,
            used_variables=used_vars, input_steps=input_steps,
            output_steps=output_steps, normalize=False)

        self._grid_lat = grid_lat; self._grid_lon = grid_lon
        self.n_samples = len(self._dataset)

        # audit
        self._print_audit(cfg, used_vars)

    def _print_audit(self, cfg, used_vars):
        av = cfg["aifs_variables"]; em = cfg["era5_mapping"]
        fields_ok = set()
        for aifs in av["surface"]:
            if em["surface"][aifs] in used_vars: fields_ok.add(aifs)
        for aifs in av["soil"]:
            if em["soil"][aifs] in used_vars: fields_ok.add(aifs)
        for vv in av["pressure_level"]:
            for lvl in av["pressure_levels"]:
                if em["pressure_level"][vv].format(level=lvl) in used_vars: fields_ok.add(f"{vv}_{lvl}")
        for aifs in av["diagnostic"]:
            if self._diag_map[aifs] in used_vars: fields_ok.add(aifs)
        forcing_ok = set(av["computed_forcing"] + ["lsm", "z", "slor", "sdor"])
        mi = [v for v in self.model_input_order if v not in fields_ok and v not in forcing_ok]
        mo = [v for v in self.model_output_order if v not in fields_ok]
        print(f"[INFO] Variables: {len(fields_ok)} fields + {len(forcing_ok)} forcing = all available")
        if mi: print(f"[WARN] {len(mi)} input vars zero-filled: {mi}")
        if mo: print(f"[WARN] {len(mo)} output vars no target data: {mo}")

    def __len__(self): return self.n_samples
    def set_grid(self, lat, lon):
        self._grid_lat = lat.astype(np.float32); self._grid_lon = lon.astype(np.float32)

    def _get_raw_sample(self, idx: int):
        result = self._dataset[int(idx)]
        return result[0], result[1], int(result[3]), list(result[4])

    def process_sample(self, invar, outvar, step_idx, time_index):
        if invar.ndim == 3: invar = invar[np.newaxis, ...]
        if outvar.ndim == 3: outvar = outvar[np.newaxis, ...]
        total_steps = self.input_steps + self.output_steps
        all_frames = np.concatenate([invar, outvar], axis=0)
        timestamps = [datetime.datetime.strptime(t, "%Y%m%d%H").replace(tzinfo=pytz.utc) for t in time_index]

        frame_tm6, frame_t0, frame_tp6 = all_frames[0], all_frames[1], all_frames[2] if total_steps >= 3 else all_frames[-1]
        fields = build_aifs_fields_from_frames(frame_tm6, frame_t0, frame_tp6, self.era5_var_list, self._cfg, diag_map=self._diag_map)

        num_pts = next(iter(fields.values())).shape[-1]
        if self._grid_lat is None:
            self._grid_lat = np.linspace(90, -90, 542080, dtype=np.float32)[:num_pts]
            self._grid_lon = np.linspace(0, 360, 542080, dtype=np.float32)[:num_pts]

        forcing = compute_forcing_features(self._grid_lat, self._grid_lon, timestamps, fields)

        num_input_vars = len(self.model_input_order)
        input_tensor = np.zeros((self.input_steps, num_pts, num_input_vars), dtype=np.float32)
        for var_name, var_idx in self._name_to_input_idx.items():
            for t in range(self.input_steps):
                if var_name in fields: input_tensor[t, :, var_idx] = fields[var_name][t]
                elif var_name in forcing: input_tensor[t, :, var_idx] = forcing[var_name][t]

        input_tensor = input_tensor * self.norm_mul[np.newaxis, np.newaxis, :] + self.norm_add[np.newaxis, np.newaxis, :]

        num_output_vars = len(self.model_output_order)
        target_tensor = np.zeros((1, num_pts, num_output_vars), dtype=np.float32)
        t_target = self.input_steps
        for var_name, var_idx in self._name_to_output_idx.items():
            if var_name in fields: target_tensor[0, :, var_idx] = fields[var_name][t_target]
        return input_tensor, target_tensor

# ============================================================================
# Model loading
# ============================================================================

def load_aifs_model(checkpoint_path: str, device: str, pretrained: bool = False):
    """Load the AIFS model through the onescience wrapper.

    Parameters
    ----------
    pretrained : bool
        False (default): Build model from scratch via ``AIFS.from_scratch()``.
        Uses local static config + grid files — NO checkpoint needed.
        Equivalent to FengWu/Fuxi ``model = Fengwu()`` pattern.
        True: Load serialised model with pretrained weights from .ckpt.
    """
    t0 = time.time()
    if pretrained:
        model = AIFS(checkpoint_path, device=device, pretrained=True)
    else:
        model = AIFS.from_scratch(device=device)
    mode = "pretrained" if pretrained else "from scratch"
    print(f"[INFO] Model loaded ({mode}, {time.time() - t0:.1f}s)")

    input_vars  = model.input_variables
    output_vars = model.output_variables
    grid_lat    = model.latitudes.cpu().numpy()
    grid_lon    = model.longitudes.cpu().numpy()
    node_wts    = model.node_weights   # np.ndarray or None
    return model, input_vars, output_vars, grid_lat, grid_lon, node_wts

# ============================================================================
# Loss
# ============================================================================

class WeightedMSELoss(nn.Module):
    def __init__(self, var_scale: np.ndarray, pl_scale: np.ndarray, node_weights=None):
        super().__init__()
        self.register_buffer("combined_scale", torch.from_numpy((var_scale * pl_scale).astype(np.float32)))
        if node_weights is not None:
            self.register_buffer("node_weights", torch.from_numpy(node_weights.astype(np.float32)))
        else: self.node_weights = None

    def forward(self, pred, target):
        sq_error = (pred - target) ** 2
        scaled = sq_error * self.combined_scale[None, None, :]
        if self.node_weights is not None:
            x = scaled.mean(dim=-1) * self.node_weights[None, :]
            return (x / self.node_weights.sum()).sum()
        return scaled.mean()

# ============================================================================
# LR schedule
# ============================================================================

def cosine_lr_schedule(step, warmup_steps, total_steps, peak_lr, min_lr):
    """Cosine LR with linear warmup.

    Manual implementation equivalent to ``timm.scheduler.CosineLRScheduler``
    for a single decay cycle (no restarts).  Same formula as the official
    anemoi-training configuration.
    """
    if step < warmup_steps:
        return peak_lr * (step / max(warmup_steps, 1))
    progress = min(
        (step - warmup_steps) / max(total_steps - warmup_steps, 1), 1.0,
    )
    return min_lr + (peak_lr - min_lr) * 0.5 * (1.0 + math.cos(math.pi * progress))

# ============================================================================
# Training
# ============================================================================

def train(cfg: dict):
    hw = cfg["hardware"]; data = cfg["data"]; ck = cfg["checkpoint"]; tr = cfg["training"]

    device = "cuda" if torch.cuda.is_available() else "cpu"
    if hw["device"] == "cpu": device = "cpu"
    print(f"[INFO] Device: {device}")

    train_years = data["train_years"]
    val_years = data["val_years"]

    from_scratch = tr.get("from_scratch", True)
    model, model_input_vars, model_output_vars, grid_lat, grid_lon, node_wts = \
        load_aifs_model(ck.get("pretrained", ""), device,
                        pretrained=not from_scratch)

    # ---- 自动计算归一化统计量（如需要）-------------------------------
    stats_path = cfg["normalizer"].get("statistics_path") or None
    if stats_path:
        stats_path = os.path.join(str(ROOT), stats_path)
        if not os.path.exists(stats_path):
            # 自动计算 —— 对标 anemoi-training DataModule 的统计量阶段
            print("[INFO] Statistics file not found — auto-computing ...")
            ds_vars = None
            if hasattr(model, '_meta') and 'dataset' in model._meta:
                ds_vars = model._meta['dataset']['variables']
            elif os.path.exists(str(ROOT / "model" / "aifs_config.json")):
                import json
                ds_vars = json.load(
                    open(str(ROOT / "model" / "aifs_config.json"))
                )['dataset']['variables']
            if ds_vars is None:
                print("[WARN] Cannot determine dataset variables — "
                      "falling back to identity normalisation")
                stats_path = None
            else:
                os.makedirs(os.path.dirname(stats_path), exist_ok=True)
                compute_era5_statistics(
                    data["data_dir"], train_years + val_years,
                    ds_vars, stats_path,
                )

    norm_mul, norm_add = compute_normalisation_params(
        cfg, model_input_vars,
        statistics_path=stats_path,
    )

    var_scale = build_loss_scale_factors(cfg, model_output_vars)
    pl_scale = build_pressure_level_scaler(cfg, model_output_vars)

    criterion = WeightedMSELoss(var_scale, pl_scale, node_wts).to(device)
    val_criterion = WeightedMSELoss(
        np.ones(len(model_output_vars), dtype=np.float32),
        np.ones(len(model_output_vars), dtype=np.float32), node_wts).to(device)

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=0.0,
        betas=tr["optimizer"]["betas"],
        weight_decay=tr["optimizer"].get("weight_decay", 0.01),
    )

    train_dataset = AIFSTrainingDataset(
        data["data_dir"], train_years, cfg, model_input_vars, model_output_vars,
        norm_mul, norm_add, input_steps=2, output_steps=1,
        grid_lat=grid_lat, grid_lon=grid_lon)
    train_dataset.set_grid(grid_lat, grid_lon)
    val_dataset = AIFSTrainingDataset(
        data["data_dir"], val_years, cfg, model_input_vars, model_output_vars,
        norm_mul, norm_add, input_steps=2, output_steps=1,
        grid_lat=grid_lat, grid_lon=grid_lon)
    val_dataset.set_grid(grid_lat, grid_lon)

    os.makedirs(ck["output_dir"], exist_ok=True)
    use_amp = (device == "cuda")
    amp_dtype = getattr(torch, tr.get("amp_dtype", "float16"))
    scaler = torch.amp.GradScaler("cuda") if use_amp else None

    model.train()
    global_step, epoch, best_val = 0, 0, float("inf")
    t_start = time.time()
    pbar = tqdm(total=tr["max_steps"], desc="Training", unit="step", dynamic_ncols=True)

    while global_step < tr["max_steps"]:
        epoch += 1
        indices = np.random.permutation(len(train_dataset))
        for batch_start in range(0, len(train_dataset), tr["batch_size"]):
            batch_indices = indices[batch_start:batch_start + tr["batch_size"]]
            batch_inputs, batch_targets = [], []
            for idx in batch_indices:
                try:
                    invar, outvar, si, ti = train_dataset._get_raw_sample(idx)
                    invar_np = invar.numpy() if hasattr(invar, "numpy") else np.asarray(invar)
                    outvar_np = outvar.numpy() if hasattr(outvar, "numpy") else np.asarray(outvar)
                    inp, tgt = train_dataset.process_sample(invar_np, outvar_np, si, ti)
                    batch_inputs.append(inp); batch_targets.append(tgt)
                except Exception as e:
                    print(f"[WARN] Skip sample {idx}: {e}")
            if not batch_inputs: continue

            x = torch.from_numpy(np.stack(batch_inputs, axis=0)).unsqueeze(2).to(device)
            y = torch.from_numpy(np.stack(batch_targets, axis=0)).squeeze(1).to(device)

            with (torch.amp.autocast("cuda", dtype=amp_dtype) if use_amp
                  else torch.no_grad()):
                pred = model(x)                 # AIFS.forward already squeezes dim

            loss = criterion(pred, y) / tr["accum_grad_batches"]
            (scaler.scale(loss) if scaler else loss).backward()

            # optimizer step
            if (batch_start // tr["batch_size"] + 1) % tr["accum_grad_batches"] == 0:
                if scaler: scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_value_(model.parameters(), tr["gradient_clip_val"])
                (scaler.step(optimizer) if scaler else optimizer.step())
                if scaler: scaler.update()
                optimizer.zero_grad()
                global_step += 1

                lr = cosine_lr_schedule(global_step, tr["warmup_steps"], tr["max_steps"], tr["peak_lr"], tr["min_lr"])
                for pg in optimizer.param_groups: pg["lr"] = lr

                elapsed = time.time() - t_start
                pbar.set_postfix({"loss": f"{loss.item():.4f}", "lr": f"{lr:.2e}", "stp/s": f"{global_step/max(elapsed,1):.1f}"})
                pbar.update(1)

                if global_step % tr["val_interval"] == 0:
                    vl = validate(model, val_dataset, val_criterion, device)
                    pbar.write(f"[step {global_step:06d}] VAL loss={vl:.6f}")
                    model.train()
                    if vl < best_val:
                        best_val = vl
                        _save(model, optimizer, global_step, epoch, vl,
                              ck["output_dir"], "best")
                if global_step % tr["save_interval"] == 0:
                    _save(model, optimizer, global_step, epoch, loss.item(),
                          ck["output_dir"], f"step{global_step:06d}")
                if global_step >= tr["max_steps"]: break

    pbar.close()
    _save(model, optimizer, global_step, epoch, loss.item(),
          ck["output_dir"], "final")
    print(f"\n[INFO] Done: {global_step} steps in {(time.time()-t_start)/60:.1f}min, best_val={best_val:.6f}")

@torch.no_grad()
def validate(model, dataset, criterion, device):
    model.eval()
    total, n = 0.0, 0
    use_amp = (device == "cuda")
    amp_dtype_val = getattr(torch, "float16")
    _first_err = True
    for idx in range(len(dataset)):
        try:
            invar, outvar, si, ti = dataset._get_raw_sample(idx)
            invar_np = invar.numpy() if hasattr(invar, "numpy") else np.asarray(invar)
            outvar_np = outvar.numpy() if hasattr(outvar, "numpy") else np.asarray(outvar)
            inp, tgt = dataset.process_sample(invar_np, outvar_np, si, ti)
            x = torch.from_numpy(inp).unsqueeze(0).unsqueeze(2).to(device)
            y = torch.from_numpy(tgt).squeeze(1).unsqueeze(0).to(device)
            with (torch.amp.autocast("cuda", dtype=amp_dtype_val) if use_amp
                  else torch.no_grad()):
                pred = model(x)                 # AIFS.forward already squeezes dim
            total += criterion(pred, y).item(); n += 1
        except Exception as e:
            if _first_err:
                print(f"[WARN] Validation error on sample {idx}: {e}")
                _first_err = False
    return total / max(n, 1)

def _save(model, optimizer, step, epoch, loss, out_dir, tag):
    """Save training checkpoint — 推理+续训合一，零外部依赖.

    单个 .ckpt 文件包含:
    - 完整 AnemoiModelEncProcDec (torch.load 直接得模型 → SimpleRunner 可用)
    - ai-models.json + supporting arrays (inference metadata)
    - optimizer.pkl (恢复训练用，推理时忽略)
    """
    import pickle, zipfile

    path = os.path.join(out_dir, f"model_bak.ckpt")

    # ---- 1.  构建 AnemoiModelInterface（含 predict_step）---------
    from anemoi.models.interface import AnemoiModelInterface
    import numpy as np

    # 虚拟统计量（identity normalization：mean=0, stdev=1）
    num_vars = len(model._interface_data_indices.data.input.name_to_index)
    dummy_stats = {
        "minimum": np.zeros(num_vars, dtype=np.float32),
        "maximum": np.ones(num_vars, dtype=np.float32),
        "mean": np.zeros(num_vars, dtype=np.float32),
        "stdev": np.ones(num_vars, dtype=np.float32),
    }

    interface = AnemoiModelInterface(
        config=model._interface_config,
        graph_data=model._interface_graph_data,
        statistics=dummy_stats,
        data_indices=model._interface_data_indices,
        metadata={},
        truncation_data=getattr(model._model, '_truncation_data', {}),
    )
    interface.model.load_state_dict(model._model.state_dict())
    interface.cpu()
    torch.save(interface, path)
    del interface

    # ---- 2.  注入 metadata (从本地静态文件) ---------------------------
    _inject_metadata(path)

    # ---- 3.  追加 optimizer 到同一子目录 ----------------------------
    with zipfile.ZipFile(path, "r") as zf:
        for name in zf.namelist():
            if name.endswith("/data.pkl") or name.endswith("/byteorder"):
                base_dir = name.split("/")[0]
                break
        else:
            base_dir = "checkpoint"

    opt_data = {
        "optimizer_state_dict": optimizer.state_dict(),
        "step": step, "epoch": epoch, "loss": loss,
    }
    with zipfile.ZipFile(path, "a") as zf:
        zf.writestr(f"{base_dir}/optimizer.pkl", pickle.dumps(opt_data))

    print(f"[INFO] Saved: {path}")


def _inject_metadata(dst_path: str):
    """从项目静态文件生成完整 metadata 并注入 checkpoint ZIP。

    自给自足——不依赖原始 .ckpt。metadata 写入 torch.save 使用的
    同一个子目录，保持 PyTorch ZIP 结构兼容。
    """
    import json, zipfile
    import numpy as np

    config_path = str(ROOT / "model" / "aifs_config.json")
    grid_path = str(ROOT / "model" / "grid-n320.npz")

    with open(config_path) as f:
        config_data = json.load(f)
    grid = np.load(grid_path)

    # ---- 1. 探测 torch.save 使用的子目录 -------------------------------
    with zipfile.ZipFile(dst_path, "r") as zf:
        for name in zf.namelist():
            if name.endswith("/data.pkl") or name.endswith("/byteorder"):
                base_dir = name.split("/")[0]
                break
        else:
            base_dir = "checkpoint"
        meta_dir = f"{base_dir}/anemoi-metadata"

    # ---- 2. 构建完整 metadata -------------------------------------------
    metadata = {
        "config": config_data["model_config"],
        "data_indices": config_data["data_indices"],
        "dataset": config_data["dataset"],
        "provenance_training": config_data.get("provenance_training", {}),
        "training": config_data.get("training", {}),
        "run_id": config_data.get("run_id", ""),
        "seed": config_data.get("seed", 0),
        "uuid": config_data.get("uuid", ""),
        "timestamp": config_data.get("timestamp", ""),
        "version": "1.0",
        "supporting_arrays_paths": {
            "latitudes": {
                "path": f"{meta_dir}/latitudes.numpy",
                "shape": [542080],
                "dtype": "float64",
            },
            "longitudes": {
                "path": f"{meta_dir}/longitudes.numpy",
                "shape": [542080],
                "dtype": "float64",
            },
        },
    }

    # ---- 3. 写入 checkpoint ZIP（与模型同一子目录）---------------------
    meta_json = json.dumps(metadata).encode("utf-8")
    with zipfile.ZipFile(dst_path, "a") as dst:
        dst.writestr(f"{meta_dir}/ai-models.json", meta_json)
        dst.writestr(
            f"{meta_dir}/latitudes.numpy",
            grid["latitudes"].astype(np.float64).tobytes(),
        )
        dst.writestr(
            f"{meta_dir}/longitudes.numpy",
            grid["longitudes"].astype(np.float64).tobytes(),
        )

    print(f"[INFO] Metadata injected into {base_dir}/")

# ============================================================================
# CLI
# ============================================================================

def main():
    parser = argparse.ArgumentParser(description="AIFS v1.1 Training")
    parser.add_argument("--config", "-c", type=str,
                        default=str(ROOT / "conf" / "config.yaml"))
    args = parser.parse_args()

    cfg = load_config(args.config)
    hw = cfg["hardware"]
    os.environ["CUDA_VISIBLE_DEVICES"] = str(hw["device_ids"])
    if hw["device"] == "dcu": os.environ["HIP_VISIBLE_DEVICES"] = str(hw["device_ids"])
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

    seed = cfg["training"]["seed"]
    np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)

    print(f"\n{'='*60}\n  AIFS v1.1 — Pre-training\n{'='*60}")
    print(f"  Device: {hw['device']}  Dataset: {cfg['data']['data_dir']}")
    print(f"  LR: {cfg['training']['peak_lr']:.2e}  Steps: {cfg['training']['max_steps']}")
    print(f"  Output: {cfg['checkpoint']['output_dir']}\n{'='*60}\n")

    try: train(cfg)
    except KeyboardInterrupt: print("\n[INFO] Interrupted.")
    except Exception: traceback.print_exc(); sys.exit(1)

if __name__ == "__main__": main()