"""Train one model from a run configuration.

    python -m bifurcation.train --config records/euler/runs/cascaded_si_S10/config.yaml --data $EULER_DATA \\
        --out runs/euler_cascaded_si_S10 [--seed 0] [--device cuda] [--resume]

The configuration names the dataset, the model, the method with its hyperparameters and the training schedule
(records/<system>/runs/<model>/config.yaml are the configurations of the paper's models). Before training, the
training pairs are written to OUT/manifest/: Cascaded SI with S segments reads the first 200 / S realizations of each
of the 80 training conditions, cut into S segments (16,000 pairs for every S; S = 1 is Direct SI); EDM reads the
endpoint pair (x_0, x_T) of all 200 realizations. Checkpoints are saved every ``training.val_every_updates`` updates
under OUT/checkpoints/ (bifurcation.trainer); ``--resume`` continues an interrupted run exactly from its
``restart_latest.pt``.
"""
from __future__ import annotations

import argparse
import os
from pathlib import Path

import numpy as np
import torch

from .data import manifests as MF
from .data.loaders import get_dataloaders
from .registry import build_method, build_model
from .trainer import Trainer
from .utils import load_yaml, save_resolved_config, seed_everything

# a restart may only continue a run with the same data, method, model and seed
RESUME_IDENTITY_KEYS = (("data", "name"), ("data", "packaged_dataset_path"), ("method", "name"), ("method", "data_mode"),
                        ("method", "sigma_train"), ("method", "P_mean"), ("method", "lr"), ("training", "batch_size"),
                        ("model", "name"))


def parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--config", required=True, help="run configuration (YAML)")
    p.add_argument("--data", required=True, help="dataset package directory")
    p.add_argument("--out", required=True, help="run directory (new, or the one to resume)")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--device", default=None, help="default: cuda if available, else cpu")
    p.add_argument("--resume", action="store_true", help="continue from OUT/checkpoints/restart_latest.pt")
    return p.parse_args(argv)


def write_training_manifest(cfg: dict, out: Path) -> Path:
    """The training pairs of the run (see the module docstring); returns the manifest file."""
    packaged = cfg["data"]["packaged_dataset_path"]
    n_ic, R = int(cfg["data"]["train_ics"]), int(cfg["data"]["train_realizations"])
    g = MF.fine_grid(packaged)
    if len(g["gids_train"]) < n_ic or g["vecs_train_shape"][1] < R:
        raise ValueError(f"the package holds {len(g['gids_train'])} training conditions x {g['vecs_train_shape'][1]} "
                         f"realizations, the configuration needs {n_ic} x {R}")
    m = cfg["method"]
    T = float(m["sampling"]["t_values"][-1])
    if m["data_mode"] == "fine_segment_pairs":
        S = int(m["sampling"].get("n_segments", 1))
        if R % S:
            raise ValueError(f"S = {S} must divide the {R} realizations per training condition")
        geo = MF.segment_geometry(T, S, g["t_values"])
        rows = MF.block_rows_segment(g["gids_train"], n_ic, R // S)
        name = f"segments_S{S:02d}"
        path = MF.write_segment_manifest(out / "manifest", rows, geo, name, packaged,
                                         extra={"realizations_per_condition": R // S})
    elif m["data_mode"] == "fine_endpoint_pairs":
        geo = MF.segment_geometry(T, 1, g["t_values"])
        rows = MF.block_rows_endpoint(g["gids_train"], n_ic, R)
        path = MF.write_endpoint_manifest(out / "manifest", rows, geo["n_T"], g["t_values"], "endpoints", packaged,
                                          extra={"realizations_per_condition": R})
    else:
        raise ValueError(f"unknown method.data_mode {m['data_mode']!r}")
    if int(m["fine_horizon_idx"]) != int(geo["n_T"]):
        raise ValueError(f"method.fine_horizon_idx {m['fine_horizon_idx']} != the index {geo['n_T']} of T = {T}")
    return path


def main(argv=None) -> int:
    args = parse_args(argv)
    cfg = load_yaml(args.config)
    cfg["data"]["packaged_dataset_path"] = str(Path(args.data).resolve()) + "/"
    cfg["seed"] = int(args.seed)
    run_dir = Path(args.out).resolve()
    run_dir.mkdir(parents=True, exist_ok=True)
    cfg["method"]["fine_slice_manifest"] = str(write_training_manifest(cfg, run_dir))
    cfg["training"]["device"] = str(args.device or os.environ.get("BIFURCATION_DEVICE")
                                    or ("cuda" if torch.cuda.is_available() else "cpu"))

    resume_payload = None
    if args.resume:
        restart_path = run_dir / "checkpoints" / "restart_latest.pt"
        if not restart_path.is_file():
            raise FileNotFoundError(f"--resume: no restart checkpoint at {restart_path}")
        resume_payload = torch.load(restart_path, map_location="cpu", weights_only=False)
        old = resume_payload.get("config", {})
        for section, key in RESUME_IDENTITY_KEYS:
            a, b = old.get(section, {}).get(key), cfg.get(section, {}).get(key)
            if a is not None and a != b:
                raise AssertionError(f"--resume: {section}.{key} was {a!r}, now {b!r}; refusing to continue a different run")
        if int(old.get("seed", args.seed)) != int(args.seed):
            raise AssertionError(f"--resume: the run has seed {old.get('seed')}, --seed {args.seed}")
        print(f"resume: restart_latest.pt at step {resume_payload['step']}")
    save_resolved_config(cfg, run_dir / "config_resolved.yaml")

    seed_everything(int(args.seed))
    train_loader, val_loader = get_dataloaders(cfg, seed=int(args.seed))
    model = build_model(cfg["model"], cfg["method"], cfg["data"])
    method = build_method(cfg["method"])
    Trainer(model=model, method=method, train_loader=train_loader, val_loader=val_loader, cfg=cfg, run_dir=run_dir,
            device=str(cfg["training"]["device"])).fit(resume=resume_payload)
    print(f"\nrun:         {run_dir}\ncheckpoints: {run_dir / 'checkpoints'}\nlog:         {run_dir / 'train_log.csv'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())