"""Reference outputs of the three release models, computed with the ORIGINAL code.

Runs the frozen code in ../source_snapshot on the frozen weights in ../originals, on the
first devices of the internal test split 0, and writes ../reference/<model>.npz. The new
boss_bench code must reproduce these outputs on the same devices (matched by uuid in the
published files).

Usage: python make_reference.py {loss_fraction,heat_flux_iid,inverse_design}

Everything runs on CPU in float32 with fixed noise, so the outputs are deterministic.
"""

import os
import sys
import types

import h5py
import numpy as np
import torch

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SNAP = os.path.join(ROOT, "source_snapshot")
ORIG = os.path.join(ROOT, "originals")
OUT = os.path.join(ROOT, "reference")
DATA = "/path/to/internal/processed/confinement/confinement_data_test_split_00000.hdf5"  # pre-release file
SEED = 0


def uuids(rows):
    with h5py.File(DATA, "r") as f:
        return f["uuid"][:][rows].astype(str)


def loss_fraction(n=64):
    sys.path.insert(0, os.path.join(SNAP, "artemiss_benchmarks"))
    from benchmarks.surrogates.models_vmec.train import Predictor, load_files

    d = os.path.join(ORIG, "loss_fraction")
    pred = Predictor(os.path.join(d, "model_loss_frac_conv_tmax_100ms_biggg_sigmoid.pt"),
                     os.path.join(d, "model_config_loss_frac_conv_tmax_100ms_biggg_sigmoid.json"))
    # t_max = 0.1 s and s0 = 0.1, as in the training run (the "tmax_100ms" in its name)
    data = load_files(DATA, "loss_frac", s0=0.1, t_max_override=0.1)

    # load_files drops rows with NaN inputs; rebuild its mask to recover the rows' uuids
    with h5py.File(DATA, "r") as f:
        valid = ~np.isnan(f["surf_xyz"][:]).any(axis=1)
        for k in ("pressure", "current"):
            valid &= ~np.isnan(f[k][:]).any(axis=1)
        valid &= ~np.isnan(f["volavgB"][:])
    rows = np.flatnonzero(valid)
    assert len(rows) == len(data["target"])
    rows = rows[:n]

    inputs = {k: data[k][:n] for k in ("surf_r", "surf_z", "pressure", "current", "volavgB", "nfp")}
    y = pred.predict_batch(**inputs)
    np.savez(os.path.join(OUT, "loss_fraction.npz"), uuid=uuids(rows), rows=rows,
             prediction=np.asarray(y), target=data["target"][:n], **inputs)
    mae = np.abs(np.asarray(y) - data["target"][:n]).mean()
    print(f"loss_fraction: {n} devices, MAE vs traced = {mae:.4f}")


def heat_flux_iid(n_configs=16, n_points=512, n_steps=50):
    sys.path.insert(0, os.path.join(SNAP, "fusion", "heat_flux"))
    import train_final as tf
    from flow_model_final import FlowMatchingModel

    d = os.path.join(ORIG, "heat_flux_iid")
    cfg = torch.load(os.path.join(d, "model_config.pt"), weights_only=True)
    model = FlowMatchingModel(
        d_cond=cfg["d_cond"], x_dim=cfg.get("x_dim", tf.X_DIM),
        embed_dim=cfg["embed_dim"], hidden=cfg["hidden_dim"],
        full_cloud=cfg.get("full_cloud", False), pool=cfg.get("pool", "max"),
        gate_global=cfg.get("gate_global", False),
        n_vel_blocks=cfg.get("n_vel_blocks", 4))
    model.load_state_dict(torch.load(os.path.join(d, "model_best.pt"), weights_only=True))
    model.eval()
    stats = dict(np.load(os.path.join(d, "cond_stats.npz")))

    # min_lost=0 keeps every (config, s_init) pair, so unique index == HDF5 row
    scal, surf, v2u, sinit, points, lf, ni, _ = tf.load_single_file(DATA, min_lost=0)
    keep = v2u < n_configs
    scal, surf = scal[:n_configs], surf[:n_configs]
    v2u, sinit = v2u[keep], sinit[keep]
    points = [p for p, k in zip(points, keep) if k]
    lf, ni = lf[keep], ni[keep]
    raw = dict(cond_scalars=scal.copy(), surf_xyz=surf.copy(), s_init=sinit.copy())

    tf.apply_stats(scal, sinit, surf, points, stats)
    c_emb, f_pred = tf.encode_and_predict_all(
        model, scal, sinit, torch.from_numpy(surf), v2u, torch.device("cpu"))

    gen = torch.Generator().manual_seed(SEED)
    x0 = torch.randn(len(v2u), n_points, tf.X_DIM, generator=gen)
    with torch.no_grad():
        c_pt = c_emb.repeat_interleave(n_points, dim=0)
        x1 = tf._ode_integrate(model, x0.reshape(-1, tf.X_DIM), c_pt, n_steps, "euler", amp=False)
    x1 = x1.reshape(len(v2u), n_points, tf.X_DIM)

    np.savez(os.path.join(OUT, "heat_flux_iid.npz"), uuid=uuids(np.arange(n_configs)),
             config_index=v2u, loss_fraction_pred=f_pred, loss_fraction=lf, n_inner=ni,
             cond_embedding=c_emb.numpy(), x0=x0.numpy(), x1=x1.numpy(),
             n_steps=n_steps, sampler="euler", **raw)
    print(f"heat_flux_iid: {len(v2u)} samples, loss-fraction MAE = {np.abs(f_pred - lf).mean():.4f}")


def inverse_design(n=8):
    here = os.path.join(SNAP, "fusion", "shape_generation", "inverse_design")
    sys.path.insert(0, here)
    sys.modules.setdefault("wandb", types.ModuleType("wandb"))  # imported by train.py, unused here
    import dataloading as dl
    from evaluate_samples import build_model, to_rz
    from flow_settings import load_flow_settings
    from sample_loss_sweep import sample_from_noise

    ck = torch.load(os.path.join(ORIG, "inverse_design", "boss_diffusion_surf_rz_unet_ep4000_ema.pt"),
                    map_location="cpu", weights_only=False)
    model, d_design = build_model(ck, "cpu")
    model.load_state_dict(ck["ema"])
    model.eval()
    sp = load_flow_settings()["sampling"]

    ds = dl.StellaratorDataset([DATA], stats=ck["stats"], normalize=True)
    # the dataset drops rows (NaN, non-stellsym); find each kept row by its boundary
    with h5py.File(DATA, "r") as f:
        xyz = f["surf_xyz"][:].astype(np.float32)
    rows = np.array([np.flatnonzero((xyz == ds.surf_xyz[i]).all(axis=1))[0] for i in range(n)])

    props = ds.props[:n]
    x0 = torch.randn(n, d_design, generator=torch.Generator().manual_seed(SEED))
    with torch.no_grad():
        x1 = sample_from_noise(model, props, x0, sp["n_steps"], sp["method"], sp["t_min"], sp["t_max"])
    R, Z = to_rz(x1, ck["stats"], "cpu")

    np.savez(os.path.join(OUT, "inverse_design.npz"), uuid=uuids(rows), rows=rows,
             props_raw=ds.props_raw[:n], props=props.numpy(), x0=x0.numpy(), x1=x1.numpy(),
             R=R, Z=Z, R_true=ds.design_raw[:n].reshape(n, 41, 41, 2)[..., 0],
             Z_true=ds.design_raw[:n].reshape(n, 41, 41, 2)[..., 1], **{f"sampling_{k}": v for k, v in sp.items()})
    print(f"inverse_design: {n} designs, sampler {sp}")


if __name__ == "__main__":
    torch.manual_seed(SEED)
    {"loss_fraction": loss_fraction, "heat_flux_iid": heat_flux_iid,
     "inverse_design": inverse_design}[sys.argv[1]]()