DeepHealthExpo/train_exposure_autoencoder.py

"""Pretrain a lightweight TimesNet autoencoder on training-set exposure."""
from __future__ import annotations

import argparse
import hashlib
import json
import logging
import math
import os
from pathlib import Path

import numpy as np
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset, DistributedSampler
from tqdm import tqdm

from backbones import TimesNetExposureAutoencoder
from dataset import ExposureCache
from train_util import (
    configure_torch_for_training,
    create_unique_run_dir,
    load_eid_file,
    resolve_device,
    set_seed,
    setup_logging,
)


class ExposureWindowDataset(Dataset):
    def __init__(self, cache: ExposureCache, row_indices: np.ndarray):
        self.cache = cache
        self.row_indices = np.asarray(row_indices, dtype=np.int64)

    def __len__(self) -> int:
        return len(self.row_indices)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        row = int(self.row_indices[index])
        return {
            "daily": torch.from_numpy(
                np.array(self.cache.daily[row], dtype=np.float32, copy=True)
            ),
            "monthly": torch.from_numpy(
                np.array(self.cache.monthly[row], dtype=np.float32, copy=True)
            ),
        }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Pretrain a lightweight TimesNet exposure autoencoder"
    )
    parser.add_argument("--exposure_cache_dir", default=None)
    parser.add_argument("--train_eid_file", default="ukb_train_eid.csv")
    parser.add_argument("--val_eid_file", default="ukb_val_eid.csv")
    parser.add_argument(
        "--channel_stats_file",
        default=None,
        help=(
            "Cached channel statistics .npz file. Defaults to "
            "<exposure_cache_dir>/train_channel_stats.npz."
        ),
    )
    parser.add_argument(
        "--recompute_channel_stats",
        action="store_true",
        help="Ignore a compatible statistics cache and recompute it.",
    )
    parser.add_argument("--runs_root", default="runs")
    parser.add_argument(
        "--resume_checkpoint",
        default=None,
        help=(
            "Resume training from a run directory or checkpoint. A directory "
            "uses last.pt when present, otherwise best.pt."
        ),
    )
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--n_embd", type=int, default=120)
    parser.add_argument("--d_model", type=int, default=64)
    parser.add_argument("--n_layers", type=int, default=2)
    parser.add_argument("--top_k", type=int, default=2)
    parser.add_argument("--n_backbone_blocks", type=int, default=1)
    parser.add_argument("--backbone_kernel_size", type=int, default=5)
    parser.add_argument("--backbone_expansion", type=float, default=2.0)
    parser.add_argument("--dropout", type=float, default=0.0)
    parser.add_argument("--mask_ratio", type=float, default=0.25)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--base_lr", type=float, default=3e-4)
    parser.add_argument("--weight_decay", type=float, default=0.05)
    parser.add_argument("--max_epochs", type=int, default=100)
    parser.add_argument("--warmup_epochs", type=int, default=5)
    parser.add_argument("--patience", type=int, default=12)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument(
        "--data_parallel",
        action="store_true",
        help="Use torch.nn.DataParallel across multiple CUDA devices.",
    )
    parser.add_argument(
        "--gpu_ids",
        default=None,
        help="Comma-separated CUDA device ids for --data_parallel, e.g. 0,1,2,3.",
    )
    parser.add_argument(
        "--ddp_backend",
        default=None,
        choices=["nccl", "gloo"],
        help="DDP backend. Defaults to nccl on CUDA and gloo otherwise.",
    )
    parser.add_argument(
        "--prefetch_factor",
        type=int,
        default=4,
        help="DataLoader batches prefetched by each worker.",
    )
    return parser.parse_args()


def validate_args(args: argparse.Namespace) -> None:
    if not args.exposure_cache_dir:
        raise ValueError(
            "--exposure_cache_dir is required unless loaded from resume config"
        )
    if not 0.0 <= args.mask_ratio < 1.0:
        raise ValueError("--mask_ratio must be in [0, 1)")
    if args.num_workers > 0 and args.prefetch_factor <= 0:
        raise ValueError("--prefetch_factor must be positive")
    if isinstance(args.gpu_ids, str) and args.gpu_ids:
        try:
            args.gpu_ids = [
                int(part.strip())
                for part in args.gpu_ids.split(",")
                if part.strip()
            ]
        except ValueError as exc:
            raise ValueError(
                "--gpu_ids must be a comma-separated list of integers"
            ) from exc
        if not args.gpu_ids:
            raise ValueError("--gpu_ids did not contain any valid CUDA device ids")
        args.data_parallel = True


def resolve_resume_checkpoint(resume_path: str | None) -> Path | None:
    if not resume_path:
        return None
    path = Path(resume_path)
    if path.is_dir():
        last_path = path / "last.pt"
        if last_path.is_file():
            return last_path
        best_path = path / "best.pt"
        if best_path.is_file():
            return best_path
        raise FileNotFoundError(
            f"Resume run directory has neither last.pt nor best.pt: {path}"
        )
    if not path.is_file():
        raise FileNotFoundError(f"--resume_checkpoint does not exist: {path}")
    return path


def apply_resume_config(args: argparse.Namespace, resume_checkpoint: Path) -> None:
    config_path = resume_checkpoint.parent / "train_config.json"
    if not config_path.is_file():
        raise FileNotFoundError(
            f"Resume requires train_config.json next to checkpoint: {config_path}"
        )
    config = json.loads(config_path.read_text(encoding="utf-8"))
    resume_value = str(resume_checkpoint)
    for key, value in config.items():
        if hasattr(args, key):
            setattr(args, key, value)
    args.resume_checkpoint = resume_value


def select_rows(cache: ExposureCache, eids: set[int], split: str) -> np.ndarray:
    valid_row = np.asarray(cache.row_index, dtype=np.int64) >= 0
    selected_events = valid_row & np.isin(cache.eids, np.fromiter(eids, np.int64))
    rows = np.unique(
        np.asarray(cache.row_index[selected_events], dtype=np.int64)
    )
    if len(rows) == 0:
        raise ValueError(f"{split} exposure rows are empty after EID filtering")
    return rows


def maybe_wrap_data_parallel(
    model: TimesNetExposureAutoencoder,
    args: argparse.Namespace,
    device: torch.device,
    logger,
):
    if not args.data_parallel:
        return model
    if device.type != "cuda":
        raise ValueError("--data_parallel requires --device cuda or cuda:<id>")
    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
        raise ValueError("--data_parallel requires at least two CUDA devices")
    primary = (
        int(device.index)
        if device.index is not None
        else int(torch.cuda.current_device())
    )
    device_ids = (
        args.gpu_ids
        if args.gpu_ids
        else list(range(torch.cuda.device_count()))
    )
    device_ids = [primary, *[idx for idx in device_ids if idx != primary]]
    if len(device_ids) < 2:
        raise ValueError("--data_parallel needs at least two device ids")
    if any(idx < 0 or idx >= torch.cuda.device_count() for idx in device_ids):
        raise ValueError(f"CUDA device id is out of range: {device_ids}")
    logger.info(f"Using DataParallel on CUDA devices: {device_ids}")
    return torch.nn.DataParallel(
        model, device_ids=device_ids, output_device=primary
    )


def unwrap_model(model) -> TimesNetExposureAutoencoder:
    if isinstance(model, (torch.nn.DataParallel, DistributedDataParallel)):
        return model.module
    return model


def init_distributed(
    args: argparse.Namespace,
) -> tuple[torch.device, int, int, int]:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size == 1:
        return resolve_device(args.device), 0, 0, 1
    if args.data_parallel:
        raise ValueError("--data_parallel cannot be combined with torchrun/DDP")
    local_rank = int(os.environ["LOCAL_RANK"])
    rank = int(os.environ["RANK"])
    if not torch.cuda.is_available():
        raise ValueError("Multi-process exposure training requires CUDA")
    torch.cuda.set_device(local_rank)
    backend = args.ddp_backend or "nccl"
    dist.init_process_group(backend=backend, init_method="env://")
    return torch.device("cuda", local_rank), rank, local_rank, world_size


def rank_logger(rank: int, run_dir: Path):
    if rank == 0:
        return setup_logging(run_dir)
    logger = logging.getLogger(f"DeepHealth.rank{rank}")
    logger.handlers.clear()
    logger.addHandler(logging.NullHandler())
    return logger


def distributed_run_dir(
    args: argparse.Namespace, rank: int, world_size: int
) -> tuple[Path, str]:
    payload: list[str | None] = [None, None]
    if rank == 0:
        if args.resume_checkpoint:
            resume_path = Path(args.resume_checkpoint)
            run_dir = resume_path.parent
            run_name = run_dir.name
        else:
            run_dir, run_name = create_unique_run_dir(
                lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root)
            )
        payload = [str(run_dir), run_name]
    if world_size > 1:
        dist.broadcast_object_list(payload, src=0)
    return Path(str(payload[0])), str(payload[1])


def channel_stats(
    cache: ExposureCache, rows: np.ndarray, chunk_size: int = 256
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    results = []
    for source in (cache.daily, cache.monthly):
        sums = np.zeros(source.shape[-1], dtype=np.float64)
        squares = np.zeros_like(sums)
        counts = np.zeros_like(sums)
        for start in tqdm(range(0, len(rows), chunk_size), desc="Channel statistics"):
            values = np.asarray(source[rows[start:start + chunk_size]], dtype=np.float64)
            finite = np.isfinite(values)
            clean = np.where(finite, values, 0.0)
            sums += clean.sum(axis=(0, 1))
            squares += np.square(clean).sum(axis=(0, 1))
            counts += finite.sum(axis=(0, 1))
        mean = sums / np.maximum(counts, 1.0)
        variance = squares / np.maximum(counts, 1.0) - np.square(mean)
        std = np.sqrt(np.maximum(variance, 1e-12))
        results.extend([mean.astype(np.float32), std.astype(np.float32)])
    return tuple(results)


def eid_set_hash(eids: set[int]) -> str:
    digest = hashlib.sha256()
    for eid in sorted(eids):
        digest.update(f"{eid}\n".encode("ascii"))
    return digest.hexdigest()


def load_or_compute_channel_stats(
    cache: ExposureCache,
    rows: np.ndarray,
    train_eids: set[int],
    stats_path: Path,
    recompute: bool,
    logger,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    eid_hash = eid_set_hash(train_eids)
    if stats_path.is_file() and not recompute:
        try:
            with np.load(stats_path, allow_pickle=False) as saved:
                compatible = (
                    str(saved["train_eid_sha256"].item()) == eid_hash
                    and int(saved["cache_event_rows"].item()) == len(cache.eids)
                    and int(saved["train_window_rows"].item()) == len(rows)
                )
                if compatible:
                    logger.info(f"Loading channel statistics from {stats_path}")
                    return (
                        saved["daily_mean"].astype(np.float32),
                        saved["daily_std"].astype(np.float32),
                        saved["monthly_mean"].astype(np.float32),
                        saved["monthly_std"].astype(np.float32),
                    )
                logger.info("Channel statistics cache is stale; recomputing")
        except (KeyError, OSError, ValueError) as exc:
            logger.warning(
                f"Could not read channel statistics cache ({exc}); recomputing"
            )

    logger.info("Computing channel statistics from training exposure")
    stats = channel_stats(cache, rows)
    stats_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez(
        stats_path,
        daily_mean=stats[0],
        daily_std=stats[1],
        monthly_mean=stats[2],
        monthly_std=stats[3],
        train_eid_sha256=np.asarray(eid_hash),
        cache_event_rows=np.asarray(len(cache.eids), dtype=np.int64),
        train_window_rows=np.asarray(len(rows), dtype=np.int64),
    )
    logger.info(f"Saved channel statistics to {stats_path}")
    return stats


def masked_mse(
    prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor
) -> torch.Tensor:
    error = (prediction - target).square() * mask
    return error.sum() / mask.sum().clamp_min(1.0)


def run_epoch(
    model,
    loader: DataLoader,
    device: torch.device,
    stats: tuple[torch.Tensor, ...],
    mask_ratio: float,
    optimizer: AdamW | None,
    scaler: torch.amp.GradScaler,
    grad_clip: float,
    amp_enabled: bool,
    show_progress: bool,
) -> float:
    training = optimizer is not None
    model.train(training)
    loss_accumulator = torch.zeros(2, device=device, dtype=torch.float64)
    daily_mean, daily_std, monthly_mean, monthly_std = stats
    context = torch.enable_grad if training else torch.no_grad
    with context():
        for batch in tqdm(
            loader,
            desc="train" if training else "val",
            leave=False,
            disable=not show_progress,
        ):
            daily = batch["daily"].to(device, non_blocking=True)
            monthly = batch["monthly"].to(device, non_blocking=True)
            daily_observed = torch.isfinite(daily)
            monthly_observed = torch.isfinite(monthly)
            daily = (torch.nan_to_num(daily) - daily_mean) / daily_std
            monthly = (torch.nan_to_num(monthly) - monthly_mean) / monthly_std
            daily = daily * daily_observed
            monthly = monthly * monthly_observed

            if training and mask_ratio > 0:
                daily_input_mask = daily_observed & (
                    torch.rand_like(daily) >= mask_ratio
                )
                monthly_input_mask = monthly_observed & (
                    torch.rand_like(monthly) >= mask_ratio
                )
            else:
                daily_input_mask = daily_observed
                monthly_input_mask = monthly_observed
            daily_input = daily * daily_input_mask
            monthly_input = monthly * monthly_input_mask

            if training:
                optimizer.zero_grad(set_to_none=True)
            with torch.autocast(
                device_type=device.type, dtype=torch.float16,
                enabled=amp_enabled,
            ):
                daily_hat, monthly_hat, _ = model(
                    daily_input, monthly_input,
                    daily_input_mask, monthly_input_mask,
                )
                loss = (
                    masked_mse(daily_hat, daily, daily_observed)
                    + masked_mse(monthly_hat, monthly, monthly_observed)
                )
            if training:
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                scaler.step(optimizer)
                scaler.update()
            batch_size = daily.size(0)
            loss_accumulator[0] += loss.detach().double() * batch_size
            loss_accumulator[1] += batch_size
    if dist.is_initialized():
        dist.all_reduce(loss_accumulator, op=dist.ReduceOp.SUM)
    return float((loss_accumulator[0] / loss_accumulator[1].clamp_min(1)).item())


def learning_rate(epoch: int, args: argparse.Namespace) -> float:
    if epoch < args.warmup_epochs:
        return args.base_lr * (epoch + 1) / max(args.warmup_epochs, 1)
    progress = (epoch - args.warmup_epochs) / max(
        args.max_epochs - args.warmup_epochs - 1, 1
    )
    return args.base_lr * 0.5 * (1.0 + math.cos(math.pi * progress))


def autoencoder_checkpoint_payload(
    model,
    optimizer: AdamW,
    scaler: torch.amp.GradScaler,
    config: dict,
    raw_stats: tuple[np.ndarray, ...],
    epoch: int,
    val_loss: float,
    best_loss: float,
    stale_epochs: int,
    history: list[dict],
    include_training_state: bool,
) -> dict:
    checkpoint_model = unwrap_model(model)
    payload = {
        "model_state_dict": checkpoint_model.state_dict(),
        "encoder_state_dict": checkpoint_model.encoder.state_dict(),
        "model_config": {
            key: config[key] for key in (
                "n_embd", "d_model", "n_layers", "top_k",
                "n_backbone_blocks", "backbone_kernel_size",
                "backbone_expansion", "dropout",
            )
        },
        "normalization": {
            "daily_mean": raw_stats[0],
            "daily_std": raw_stats[1],
            "monthly_mean": raw_stats[2],
            "monthly_std": raw_stats[3],
        },
        "epoch": epoch,
        "val_loss": val_loss,
        "best_loss": best_loss,
        "stale_epochs": stale_epochs,
        "history": history,
    }
    if include_training_state:
        payload["optimizer_state_dict"] = optimizer.state_dict()
        payload["scaler_state_dict"] = scaler.state_dict()
    return payload


def save_autoencoder_checkpoint(payload: dict, checkpoint_path: Path) -> None:
    tmp_path = checkpoint_path.with_suffix(checkpoint_path.suffix + ".tmp")
    torch.save(payload, tmp_path)
    tmp_path.replace(checkpoint_path)


def load_resume_checkpoint(
    checkpoint_path: Path,
    model,
    optimizer: AdamW,
    scaler: torch.amp.GradScaler,
    val_loader: DataLoader,
    stats: tuple[torch.Tensor, ...],
    device: torch.device,
    grad_clip: float,
    amp_enabled: bool,
    show_progress: bool,
    logger,
) -> tuple[int, float, int, list[dict]]:
    checkpoint = torch.load(checkpoint_path, map_location=device)
    if "model_state_dict" not in checkpoint:
        raise KeyError(
            f"Checkpoint does not contain model_state_dict: {checkpoint_path}"
        )
    unwrap_model(model).load_state_dict(checkpoint["model_state_dict"])
    has_training_state = "optimizer_state_dict" in checkpoint
    if has_training_state:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    else:
        logger.warning(
            "Resume checkpoint has no optimizer state; continuing with a fresh "
            "optimizer"
        )
    if "scaler_state_dict" in checkpoint:
        scaler.load_state_dict(checkpoint["scaler_state_dict"])
    elif scaler.is_enabled():
        logger.warning(
            "Resume checkpoint has no AMP scaler state; continuing with a fresh "
            "scaler"
        )
    start_epoch = int(checkpoint.get("epoch", 0))
    history = list(checkpoint.get("history", []))
    history_path = checkpoint_path.parent / "history.json"
    if not history and history_path.is_file():
        history = json.loads(history_path.read_text(encoding="utf-8"))

    if has_training_state:
        best_loss = float(
            checkpoint.get("best_loss", checkpoint.get("val_loss", float("inf")))
        )
        stale_epochs = int(checkpoint.get("stale_epochs", 0))
    else:
        current_val_loss = run_epoch(
            model, val_loader, device, stats, 0.0, None,
            scaler, grad_clip, amp_enabled, show_progress,
        )
        history_best_epoch = start_epoch
        historical_best = float(checkpoint.get("val_loss", float("inf")))
        for entry in history:
            if "val_loss" not in entry:
                continue
            entry_loss = float(entry["val_loss"])
            if entry_loss < historical_best:
                historical_best = entry_loss
                history_best_epoch = int(entry.get("epoch", len(history)))
        if math.isfinite(historical_best):
            tolerance = max(1e-4, abs(historical_best) * 1e-3)
            if abs(current_val_loss - historical_best) > tolerance:
                logger.warning(
                    "Legacy best.pt validation loss differs from history: "
                    f"recomputed={current_val_loss:.6f}, "
                    f"history_best={historical_best:.6f}"
                )
        best_loss = min(historical_best, current_val_loss)
        if history:
            start_epoch = max(start_epoch, int(history[-1].get("epoch", len(history))))
        if current_val_loss < historical_best:
            stale_epochs = 0
        else:
            stale_epochs = max(0, start_epoch - history_best_epoch)
        logger.info(
            f"Validated legacy checkpoint {checkpoint_path.name}: "
            f"val={current_val_loss:.6f}, historical_best={historical_best:.6f}"
        )
    logger.info(
        f"Resumed from {checkpoint_path} at epoch {start_epoch}; "
        f"best_val={best_loss:.6f}, stale_epochs={stale_epochs}"
    )
    return start_epoch, best_loss, stale_epochs, history


def main() -> None:
    args = parse_args()
    resume_checkpoint = resolve_resume_checkpoint(args.resume_checkpoint)
    if resume_checkpoint is not None:
        apply_resume_config(args, resume_checkpoint)
    validate_args(args)
    device, rank, local_rank, world_size = init_distributed(args)
    set_seed(args.seed + rank)
    configure_torch_for_training(device)
    run_dir, run_name = distributed_run_dir(args, rank, world_size)
    logger = rank_logger(rank, run_dir)
    cache = ExposureCache(args.exposure_cache_dir)
    train_eids = load_eid_file(args.train_eid_file)
    val_eids = load_eid_file(args.val_eid_file)
    if train_eids & val_eids:
        raise ValueError("train and validation EID files must be disjoint")
    train_rows = select_rows(cache, train_eids, "Training")
    val_rows = select_rows(cache, val_eids, "Validation")
    stats_path = (
        Path(args.channel_stats_file)
        if args.channel_stats_file
        else Path(args.exposure_cache_dir) / "train_channel_stats.npz"
    )
    if rank == 0:
        raw_stats = load_or_compute_channel_stats(
            cache,
            train_rows,
            train_eids,
            stats_path,
            args.recompute_channel_stats,
            logger,
        )
    if world_size > 1:
        dist.barrier()
    if rank != 0:
        raw_stats = load_or_compute_channel_stats(
            cache, train_rows, train_eids, stats_path, False, logger
        )
    stats = tuple(
        torch.as_tensor(value, device=device).view(1, 1, -1)
        for value in raw_stats
    )
    if args.batch_size % world_size != 0:
        raise ValueError(
            f"--batch_size={args.batch_size} must be divisible by "
            f"DDP world size {world_size}"
        )
    local_batch_size = args.batch_size // world_size
    loader_kwargs = dict(
        batch_size=local_batch_size,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
    )
    if args.num_workers > 0:
        loader_kwargs["prefetch_factor"] = args.prefetch_factor
    train_dataset = ExposureWindowDataset(cache, train_rows)
    val_dataset = ExposureWindowDataset(cache, val_rows)
    train_sampler = (
        DistributedSampler(
            train_dataset, num_replicas=world_size, rank=rank,
            shuffle=True, seed=args.seed,
        )
        if world_size > 1 else None
    )
    val_sampler = (
        DistributedSampler(
            val_dataset, num_replicas=world_size, rank=rank, shuffle=False
        )
        if world_size > 1 else None
    )
    train_loader = DataLoader(
        train_dataset, sampler=train_sampler,
        shuffle=train_sampler is None, **loader_kwargs
    )
    val_loader = DataLoader(
        val_dataset, sampler=val_sampler, shuffle=False, **loader_kwargs
    )
    model = TimesNetExposureAutoencoder(
        n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers,
        top_k=args.top_k, n_backbone_blocks=args.n_backbone_blocks,
        backbone_kernel_size=args.backbone_kernel_size,
        backbone_expansion=args.backbone_expansion,
        dropout=args.dropout,
    ).to(device)
    if world_size > 1:
        model = DistributedDataParallel(
            model, device_ids=[local_rank], output_device=local_rank
        )
        logger.info(
            f"Using DDP with {world_size} processes; "
            f"global_batch={args.batch_size}, per_gpu_batch={local_batch_size}"
        )
    else:
        model = maybe_wrap_data_parallel(model, args, device, logger)
    optimizer = AdamW(
        model.parameters(), lr=args.base_lr,
        weight_decay=args.weight_decay, betas=(0.9, 0.95),
    )
    amp_enabled = bool(args.amp and device.type == "cuda")
    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
    logger.info(
        f"Run {run_name}: device={device}, train_rows={len(train_rows):,}, "
        f"val_rows={len(val_rows):,}"
    )
    config = vars(args) | {
        "train_rows": len(train_rows),
        "val_rows": len(val_rows),
        "daily_mean": raw_stats[0].tolist(),
        "daily_std": raw_stats[1].tolist(),
        "monthly_mean": raw_stats[2].tolist(),
        "monthly_std": raw_stats[3].tolist(),
    }
    if rank == 0 and not args.resume_checkpoint:
        (run_dir / "train_config.json").write_text(
            json.dumps(config, indent=2), encoding="utf-8"
        )

    best_loss = float("inf")
    stale_epochs = 0
    history = []
    start_epoch = 0
    if args.resume_checkpoint:
        start_epoch, best_loss, stale_epochs, history = load_resume_checkpoint(
            Path(args.resume_checkpoint), model, optimizer, scaler,
            val_loader, stats, device, args.grad_clip, amp_enabled,
            rank == 0, logger,
        )
    if start_epoch >= args.max_epochs:
        logger.info(
            f"Resume checkpoint is already at epoch {start_epoch}; "
            f"--max_epochs={args.max_epochs} leaves no remaining epochs"
        )
    for epoch in range(start_epoch, args.max_epochs):
        if train_sampler is not None:
            train_sampler.set_epoch(epoch)
        lr = learning_rate(epoch, args)
        for group in optimizer.param_groups:
            group["lr"] = lr
        train_loss = run_epoch(
            model, train_loader, device, stats, args.mask_ratio, optimizer,
            scaler, args.grad_clip, amp_enabled, rank == 0,
        )
        val_loss = run_epoch(
            model, val_loader, device, stats, 0.0, None,
            scaler, args.grad_clip, amp_enabled, rank == 0,
        )
        logger.info(
            f"Epoch {epoch + 1:03d} | lr={lr:.3e} | "
            f"train={train_loss:.6f} | val={val_loss:.6f}"
        )
        history.append(
            {"epoch": epoch + 1, "lr": lr,
             "train_loss": train_loss, "val_loss": val_loss}
        )
        if val_loss < best_loss:
            best_loss = val_loss
            stale_epochs = 0
            if rank == 0:
                save_autoencoder_checkpoint(
                    autoencoder_checkpoint_payload(
                        model, optimizer, scaler, config, raw_stats,
                        epoch + 1, val_loss, best_loss, stale_epochs,
                        history, include_training_state=False,
                    ),
                    run_dir / "best.pt",
                )
        else:
            stale_epochs += 1
        if rank == 0:
            save_autoencoder_checkpoint(
                autoencoder_checkpoint_payload(
                    model, optimizer, scaler, config, raw_stats,
                    epoch + 1, val_loss, best_loss, stale_epochs,
                    history, include_training_state=True,
                ),
                run_dir / "last.pt",
            )
            (run_dir / "history.json").write_text(
                json.dumps(history, indent=2), encoding="utf-8"
            )
        if stale_epochs >= args.patience:
            logger.info(f"Early stopping after {epoch + 1} epochs")
            break
    logger.info(f"Best validation loss: {best_loss:.6f}")
    logger.info(f"Checkpoint: {run_dir / 'best.pt'}")
    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    main()