code/training/src/training_validation/common.py

from __future__ import annotations

import csv
import json
import os
import random
import shutil
import sys
from pathlib import Path
from typing import Any

import torch
from torch.utils.data import DataLoader, Dataset


CODE_ROOT = Path(__file__).resolve().parents[2]
if str(CODE_ROOT) not in sys.path:
    sys.path.insert(0, str(CODE_ROOT))

from src.data_pipeline.basic_dataset import BasicDataset, skip_missing_collate  # noqa: E402
from src.data_pipeline.fast_dataset import FastDataset  # noqa: E402
from src.config import load_config as load_release_config  # noqa: E402
from src.model import SimVPCI, SimVPBTAux, SwinLSTMCI  # noqa: E402
from src.training_validation.ram_chunk import FastRAMChunkedDataset, FastUniqueRAMCachedDataset  # noqa: E402


def set_seed(seed: int, deterministic: bool = True) -> None:
    random.seed(int(seed))
    try:
        import numpy as np

        np.random.seed(int(seed))
    except ImportError:
        pass
    os.environ["PYTHONHASHSEED"] = str(int(seed))
    torch.manual_seed(int(seed))
    torch.cuda.manual_seed(int(seed))
    torch.cuda.manual_seed_all(int(seed))
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


class CachedDataset(Dataset):
    def __init__(self, dataset: Dataset):
        self.samples = [dataset[i] for i in range(len(dataset))]

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

    def __getitem__(self, idx: int) -> dict[str, Any]:
        return self.samples[int(idx)]


def load_config(path: str | Path) -> dict[str, Any]:
    return load_release_config(path)


def config_snapshot(config: dict[str, Any]) -> dict[str, Any]:
    return {k: v for k, v in config.items() if not k.startswith("_")}


def _evenly_spaced(items: list[Any], count: int) -> list[Any]:
    count = int(count)
    if count <= 0 or len(items) <= count:
        return items
    if count == 1:
        return [items[0]]
    last = len(items) - 1
    indices = [round(i * last / (count - 1)) for i in range(count)]
    return [items[int(i)] for i in indices]


def _apply_sample_filter(dataset: Dataset, split_cfg: dict[str, Any]) -> None:
    if not hasattr(dataset, "samples"):
        return
    samples = list(getattr(dataset, "samples"))
    sample_stride = int(split_cfg.get("sample_stride", 1) or 1)
    if sample_stride > 1:
        samples = samples[::sample_stride]
    max_samples = split_cfg.get("max_samples")
    if max_samples is not None:
        samples = _evenly_spaced(samples, int(max_samples))
    setattr(dataset, "samples", samples)


def build_dataset(config: dict[str, Any], split: str, mode: str) -> Dataset:
    ds_cfg = dict(config.get("dataset", {}))
    dataset_type = str(ds_cfg.get("type", "fast")).lower()
    split_cfg = dict(config.get(mode, {}))
    ram_chunk_cfg = _ram_chunk_config(config, split_cfg)
    use_ram_chunk = bool(ram_chunk_cfg.get("enabled", False))
    required_inputs = list(split_cfg.get("input_sources", config.get("input_sources", config.get("required_inputs", ["concat"]))))
    required_labels = list(split_cfg.get("required_labels", config.get("required_labels", ["ci"])))
    dataset_config = ds_cfg.get("config", config.get("dataset_config", config))
    if not isinstance(dataset_config, dict):
        dataset_config = load_release_config(dataset_config)
    else:
        dataset_config = dict(dataset_config)

    if split_cfg.get("time_ranges") is not None:
        dataset_config.setdefault("splits", {})
        dataset_config["splits"][split] = list(split_cfg["time_ranges"])

    if dataset_type == "fast":
        dataset: Dataset = FastDataset(dataset_config, split=split, required_inputs=required_inputs, required_labels=required_labels)
    elif dataset_type == "basic":
        if use_ram_chunk:
            raise ValueError("ram_chunk.enabled=true only supports dataset.type=fast")
        dataset = BasicDataset(dataset_config, split=split, inputs=required_inputs, labels=required_labels)
    else:
        raise ValueError(f"unsupported dataset.type: {dataset_type}")

    _apply_sample_filter(dataset, split_cfg)

    cache_cfg = config.get("ram_cache", {})
    use_cache = bool(split_cfg.get("use_ram_cache", cache_cfg.get(f"{mode}_use_ram_cache", cache_cfg.get("use_ram_cache", False))))
    if use_ram_chunk and use_cache:
        raise ValueError("ram_chunk.enabled and use_ram_cache cannot be true at the same time")
    if use_ram_chunk:
        chunk_order = str(ram_chunk_cfg.get("chunk_order", "sequential"))
        swap_policy = str(ram_chunk_cfg.get("swap_policy", "repeat_current"))
        if chunk_order not in {"sequential", "circular"}:
            raise ValueError("ram_chunk.chunk_order currently supports only 'sequential' and 'circular'")
        if swap_policy != "repeat_current":
            raise ValueError("ram_chunk.swap_policy currently supports only 'repeat_current'")
        dataset = FastRAMChunkedDataset(
            dataset,  # type: ignore[arg-type]
            chunk_ram_gb=float(ram_chunk_cfg.get("chunk_ram_gb", 40.0)),
            cache_dtype=str(ram_chunk_cfg.get("cache_dtype", "float16")),
            read_block_rows=int(ram_chunk_cfg.get("read_block_rows", 64)),
            async_prefetch=bool(ram_chunk_cfg.get("async_prefetch", True)),
            chunk_order=chunk_order,
            verbose=bool(ram_chunk_cfg.get("verbose", True)),
        )
        if not bool(config.get("_defer_ram_chunk_initial_load", False)):
            initial_chunk_id = int(ram_chunk_cfg.get("initial_chunk_id", 0))
            dataset.load_chunk_sync(initial_chunk_id)  # type: ignore[attr-defined]
    if use_cache:
        if dataset_type == "fast":
            cache_dtype = str(cache_cfg.get("cache_dtype", config.get("ram_chunk", {}).get("cache_dtype", "float16")))
            read_block_rows = int(cache_cfg.get("read_block_rows", config.get("ram_chunk", {}).get("read_block_rows", 64)))
            verbose = bool(cache_cfg.get("verbose", config.get("ram_chunk", {}).get("verbose", True)))
            dataset = FastUniqueRAMCachedDataset(
                dataset,  # type: ignore[arg-type]
                cache_dtype=cache_dtype,
                read_block_rows=read_block_rows,
                verbose=verbose,
            )
        else:
            dataset = CachedDataset(dataset)
    return dataset


def is_ram_chunk_dataset(dataset: Dataset) -> bool:
    return isinstance(dataset, FastRAMChunkedDataset)


def _ram_chunk_config(config: dict[str, Any], split_cfg: dict[str, Any]) -> dict[str, Any]:
    merged = dict(config.get("ram_chunk", {}))
    split_ram_chunk = split_cfg.get("ram_chunk")
    if isinstance(split_ram_chunk, dict):
        merged.update(split_ram_chunk)
    if "use_ram_chunk" in split_cfg:
        merged["enabled"] = bool(split_cfg["use_ram_chunk"])
    return merged


def build_dataloader(config: dict[str, Any], dataset: Dataset, mode: str) -> DataLoader:
    loader_cfg = dict(config.get("dataloader", {}))
    split_cfg = dict(config.get(mode, {}))
    batch_size = int(split_cfg.get("batch_size", loader_cfg.get("batch_size", 1)))
    num_workers = int(split_cfg.get("num_workers", loader_cfg.get("num_workers", 0)))
    shuffle_default = mode == "train"
    shuffle = bool(split_cfg.get("shuffle", loader_cfg.get(f"{mode}_shuffle", shuffle_default)))

    ds_cfg = dict(config.get("dataset", {}))
    collate_fn = None
    if str(ds_cfg.get("type", "fast")).lower() == "basic":
        required_inputs = list(split_cfg.get("input_sources", config.get("input_sources", config.get("required_inputs", ["concat"]))))
        required_labels = list(split_cfg.get("required_labels", config.get("required_labels", ["ci"])))
        collate_fn = skip_missing_collate(required_inputs, required_labels)

    loader_kwargs = {
        "batch_size": batch_size,
        "shuffle": shuffle,
        "num_workers": num_workers,
        "pin_memory": bool(loader_cfg.get("pin_memory", True)),
        "drop_last": bool(split_cfg.get("drop_last", mode == "train")),
        "collate_fn": collate_fn,
    }
    if num_workers > 0:
        loader_kwargs["persistent_workers"] = bool(
            split_cfg.get("persistent_workers", loader_cfg.get("persistent_workers", False))
        )
        loader_kwargs["prefetch_factor"] = int(
            split_cfg.get("prefetch_factor", loader_cfg.get("prefetch_factor", 2))
        )
    return DataLoader(dataset, **loader_kwargs)


def shutdown_dataloader(loader: DataLoader | None) -> None:
    if loader is None:
        return
    iterator = getattr(loader, "_iterator", None)
    if iterator is not None:
        shutdown = getattr(iterator, "_shutdown_workers", None)
        if shutdown is not None:
            try:
                shutdown()
            except Exception:
                pass
    del loader


def build_model(config: dict[str, Any]) -> torch.nn.Module:
    model_cfg = dict(config.get("model", {}))
    name = str(model_cfg.get("name", "simvp_ci")).lower()
    params = dict(model_cfg.get("params", {}))
    if name in {"simvp_ci", "simvp"}:
        return SimVPCI(**params)
    if name in {"simvp_bt_aux", "simvp_ci_bt", "cinet_bt_predict"}:
        return SimVPBTAux(**params)
    if name in {"swinlstm_ci", "swinlstm", "swinlstm_b", "swinlstm_d"}:
        if name == "swinlstm_b":
            params.setdefault("variant", "b")
        elif name == "swinlstm_d":
            params.setdefault("variant", "d")
        return SwinLSTMCI(**params)
    raise ValueError(f"unsupported model.name: {name}")


def build_optimizer(config: dict[str, Any], model: torch.nn.Module) -> torch.optim.Optimizer:
    opt_cfg = dict(config.get("optimizer", {}))
    name = str(opt_cfg.get("name", "adamw")).lower()
    params = dict(opt_cfg.get("params", {}))
    if name == "adam":
        return torch.optim.Adam(model.parameters(), **params)
    if name == "adamw":
        return torch.optim.AdamW(model.parameters(), **params)
    if name == "sgd":
        return torch.optim.SGD(model.parameters(), **params)
    raise ValueError(f"unsupported optimizer.name: {name}")


def build_scheduler(config: dict[str, Any], optimizer: torch.optim.Optimizer):
    sched_cfg = dict(config.get("scheduler", {}))
    name = str(sched_cfg.get("name", "none")).lower()
    params = dict(sched_cfg.get("params", {}))
    if name in {"none", "null", ""}:
        return None
    if name == "cosine":
        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **params)
    if name == "step":
        return torch.optim.lr_scheduler.StepLR(optimizer, **params)
    if name == "plateau":
        return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, **params)
    raise ValueError(f"unsupported scheduler.name: {name}")


def pack_inputs(batch: dict[str, Any], input_sources: list[str], device: torch.device) -> torch.Tensor:
    arrays = []
    for name in input_sources:
        value = batch["inputs"].get(name)
        if value is None:
            raise ValueError(f"batch input {name!r} is None")
        arrays.append(value.to(device=device, dtype=torch.float32, non_blocking=True))
    return torch.cat(arrays, dim=2)


def get_label(batch: dict[str, Any], label_key: str, device: torch.device) -> torch.Tensor:
    value = batch["labels"].get(label_key)
    if value is None:
        raise ValueError(f"batch label {label_key!r} is None")
    value = value.to(device=device, dtype=torch.float32, non_blocking=True)
    if value.ndim == 4 and value.shape[1] == 1:
        value = value[:, 0]
    return value


def atomic_save(payload: dict[str, Any], path: str | Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    torch.save(payload, tmp)
    os.replace(tmp, path)


def save_epoch_checkpoints(
    config: dict[str, Any],
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: Any,
    epoch: int,
    train_summary: dict[str, Any],
    criterion: torch.nn.Module | None = None,
) -> None:
    out_dir = Path(config.get("output_dir", config.get("checkpoint_dir", "runs/default")))
    ckpt_dir = out_dir / "checkpoints"
    snapshot = config_snapshot(config)
    model_payload = {
        "epoch": int(epoch),
        "model_state_dict": model.state_dict(),
        "config": snapshot,
        "train_summary": train_summary,
    }
    atomic_save(model_payload, ckpt_dir / f"epoch_{int(epoch):04d}_model.pt")

    full_payload = dict(model_payload)
    full_payload.update(
        {
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict() if scheduler is not None else None,
        }
    )
    if criterion is not None:
        full_payload["criterion_state_dict"] = criterion.state_dict()
    atomic_save(full_payload, ckpt_dir / "latest_full.pt")


def load_model_checkpoint(model: torch.nn.Module, path: str | Path, device: torch.device) -> dict[str, Any]:
    path = Path(path)
    if path.suffix == ".safetensors":
        from safetensors.torch import load_file

        state = load_file(str(path), device=str(device))
        model.load_state_dict(state, strict=True)
        return {"model_state_dict": state}
    payload = torch.load(path, map_location=device, weights_only=True)
    state = payload.get("model_state_dict", payload)
    model.load_state_dict(state)
    return payload if isinstance(payload, dict) else {"model_state_dict": payload}


def resume_full_checkpoint(
    path: str | Path,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: Any,
    device: torch.device,
    criterion: torch.nn.Module | None = None,
) -> int:
    path = Path(path)
    if not path.exists():
        return 0
    payload = torch.load(path, map_location=device, weights_only=True)
    model.load_state_dict(payload["model_state_dict"], strict=True)
    optimizer.load_state_dict(payload["optimizer_state_dict"])
    for state in optimizer.state.values():
        for key, value in state.items():
            if torch.is_tensor(value):
                state[key] = value.to(device)
    if scheduler is not None and payload.get("scheduler_state_dict") is not None:
        scheduler.load_state_dict(payload["scheduler_state_dict"])
    if criterion is not None and payload.get("criterion_state_dict") is not None:
        criterion.load_state_dict(payload["criterion_state_dict"], strict=False)
    return int(payload.get("epoch", 0))


def append_csv_row(path: str | Path, row: dict[str, Any]) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    write_header = not path.exists()
    with path.open("a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(row.keys()))
        if write_header:
            writer.writeheader()
        writer.writerow(row)


def write_json(path: str | Path, payload: dict[str, Any]) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)


def maybe_copy_best(src: Path, dst: Path, enabled: bool) -> None:
    if enabled:
        dst.parent.mkdir(parents=True, exist_ok=True)
        # Some mounted filesystems allow writing file contents but reject chmod/copystat.
        # copyfile keeps best_model.pt useful without copying metadata.
        shutil.copyfile(src, dst)