"""Draw forecasts from a trained model: for every initial condition of a split, ``n_samples`` samples of x_T.

    python -m bifurcation.sample --config CONFIG --data DIR --checkpoint FILE --out DIR [--split test] [--gids 5,16]
        [--n-samples 10000] [--n-steps N] [--seed-base 777] [--chunk 2000] [--truth-bank DIR] [--device cuda]

CONFIG is a run configuration (records/<system>/runs/<model>/config.yaml, or the config_resolved.yaml of a run trained
here). The EMA weights of the checkpoint are used when the method keeps an EMA. ``--n-steps`` overrides the sampler
resolution (steps per segment for Cascaded SI, SDE steps for Direct SI, Heun steps for EDM). The noise of condition
``gid`` comes from the stream seeded with ``seed_base * 100003 + gid`` and is drawn in chunks of ``--chunk`` samples.
``--truth-bank`` supplies x_0 of test conditions that are not in the package (waves, NLKG).

Output: ``OUT/<variant>/group_<gid>.npz`` with ``gen_xT_norm`` (n_samples, D), ``x0_norm`` (D,), ``true_xT_norm``
(the package's truth members at the horizon) and ``gid``; states are in the package's normalized units.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import torch

from .data.loaders import get_split_set
from .registry import build_method, build_model
from .utils import load_yaml


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("--checkpoint", required=True)
    p.add_argument("--out", required=True, help="output directory")
    p.add_argument("--split", default="test", choices=("train", "val", "test"))
    p.add_argument("--gids", default=None, help="comma-separated subset of the split's conditions")
    p.add_argument("--n-samples", type=int, default=10000)
    p.add_argument("--n-steps", type=int, default=None, help="sampler resolution override")
    p.add_argument("--seed-base", type=int, default=777)
    p.add_argument("--chunk", type=int, default=None, help="draw the samples of a condition in chunks of this size")
    p.add_argument("--truth-bank", default=None, help="truth bank with the x_0 of test conditions absent from the package")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    return p.parse_args(argv)


def _sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 24), b""):
            h.update(chunk)
    return h.hexdigest()


def main(argv=None) -> int:
    args = parse_args(argv)
    cfg = load_yaml(args.config)
    cfg["data"]["packaged_dataset_path"] = str(Path(args.data).resolve()) + "/"
    device = str(args.device)

    allow_tf32 = cfg.get("training", {}).get("allow_tf32")
    if allow_tf32 is not None:                        # the arithmetic class the model was trained under
        torch.backends.cuda.matmul.allow_tf32 = bool(allow_tf32)
        torch.backends.cudnn.allow_tf32 = bool(allow_tf32)
    seed = int(cfg.get("seed", cfg["data"].get("seed", 0)))

    ckpt_path = Path(args.checkpoint).resolve()
    ckpt = torch.load(ckpt_path, map_location=device)
    model = build_model(cfg["model"], cfg["method"], cfg["data"])
    ema_halflife = cfg.get("training", {}).get("ema_halflife_kimg", cfg["method"].get("ema_halflife_kimg"))
    method_uses_ema = ema_halflife is not None and float(ema_halflife) > 0
    if "ema_state_dict" in ckpt and method_uses_ema:
        model.load_state_dict(ckpt["ema_state_dict"])
        weights = "ema"
    else:
        model.load_state_dict(ckpt["model_state_dict"])
        weights = "raw"
    model.to(device).eval()
    print(f"checkpoint: {ckpt_path.name} (step {ckpt.get('step')}, {weights} weights)")

    if args.n_steps is not None:
        sampling = cfg["method"]["sampling"]
        for key in ("n_sde_steps", "n_sde_steps_per_seg", "n_steps"):
            if key in sampling:
                break
        else:
            raise KeyError("--n-steps: method.sampling has none of n_sde_steps / n_sde_steps_per_seg / n_steps")
        print(f"sampler resolution: method.sampling.{key} {sampling[key]} -> {args.n_steps}")
        sampling[key] = args.n_steps
    method = build_method(cfg["method"])

    split_x0, truth_clouds, gids = get_split_set(cfg, args.split, seed=seed, x0_bank=args.truth_bank)
    if args.gids is not None:
        wanted = sorted(int(g) for g in args.gids.split(","))
        missing = [g for g in wanted if g not in gids]
        if missing:
            raise ValueError(f"--gids not in split '{args.split}': {missing}")
        sel = [gids.index(g) for g in wanted]
        split_x0, truth_clouds, gids = split_x0[sel], truth_clouds[sel], wanted
    split_x0, truth_clouds = split_x0.to(device), truth_clouds.to(device)
    n_gids = split_x0.shape[0]
    print(f"split {args.split}: {n_gids} conditions, {args.n_samples} samples each")

    out = Path(args.out).resolve()
    out.mkdir(parents=True, exist_ok=True)
    split_offset = {"train": 1000, "val": 2000, "test": 3000}[args.split]
    variants = method.variants_for_eval()
    for v_idx, (label, sample_kwargs) in enumerate(variants.items()):
        bundle_dir = out / label
        bundle_dir.mkdir(parents=True, exist_ok=True)
        torch.manual_seed(seed + split_offset + 100 * v_idx)
        with torch.no_grad():
            for i in range(n_gids):
                x0_i, gid = split_x0[i], int(gids[i])
                torch.manual_seed(int(args.seed_base) * 100003 + gid)
                if args.chunk is not None and args.n_samples > int(args.chunk):
                    parts, left = [], int(args.n_samples)
                    while left > 0:
                        m = min(int(args.chunk), left)
                        parts.append(method.sample(model, x0_i, n_samples=m, **sample_kwargs).squeeze(0))
                        left -= m
                    gen_cloud = torch.cat(parts, dim=0)
                else:
                    gen_cloud = method.sample(model, x0_i, n_samples=int(args.n_samples), **sample_kwargs).squeeze(0)
                np.savez(bundle_dir / f"group_{gid:04d}.npz", true_xT_norm=truth_clouds[i].cpu().numpy(),
                         gen_xT_norm=gen_cloud.cpu().numpy(), x0_norm=x0_i.cpu().numpy(), gid=np.int64(gid))
                print(f"  {label}: condition {gid} ({i + 1}/{n_gids})", flush=True)

    record = {"config": str(Path(args.config).resolve()), "checkpoint": str(ckpt_path), "checkpoint_sha256": _sha256(ckpt_path),
              "weights": weights, "split": args.split, "gids": list(gids), "n_samples": int(args.n_samples),
              "sampling": {k: cfg["method"]["sampling"].get(k) for k in ("n_sde_steps", "n_sde_steps_per_seg", "n_steps",
                                                                        "n_segments", "noise_scale") if k in cfg["method"]["sampling"]},
              "seed_base": int(args.seed_base), "chunk": args.chunk, "variants": list(variants), "device": device,
              "allow_tf32": allow_tf32, "at": datetime.now(timezone.utc).isoformat(timespec="seconds")}
    (out / "SAMPLES.json").write_text(json.dumps(record, indent=1) + "\n")
    return 0


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