# reproduce every model-dependent number in experiments/results/ from the saved
# checkpoints and diff against the committed json. any mismatch = a fake or a
# reproducibility break. loads the released weights, re-runs the exact eval.
import os, sys, json
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import torch
from src.data.perturb_data import load_dataset, MATCH_STRATEGIES
from src.data.splits import load_split
from src.training.train import TrainConfig, make_model
from src.experiments.predictors import PivotPredictor
from src.experiments.forward_eval import evaluate_forward
from src.experiments.run_ablations import _fwd_inv

RES = "experiments/results"
GPU = "cuda:%d" % int(os.environ.get("PIVOT_GPU", "3"))
TOL = 1e-3
checks = []  # (label, ok, detail)


def load(path, data):
    cfg = TrainConfig(**json.load(open(os.path.join(path, "config.json"))))
    m = make_model(data, cfg, GPU)
    m.load_state_dict(torch.load(os.path.join(path, "model.pt"), map_location=GPU))
    m.eval()
    return m


def cmp(label, got, exp):
    keys = [k for k in exp if isinstance(exp[k], (int, float))
            and isinstance(got.get(k), (int, float))]
    diff = {k: (round(got[k], 4), round(exp[k], 4)) for k in keys if abs(got[k] - exp[k]) > TOL}
    checks.append((label, not diff, diff))


def fwd_pivot(data, model, split, max_perts=80):
    sp = load_split(data.dir, split)
    test = list(sp["test_perts"]) if split != "cell" else [p for p in data.perturbations if len(data.parse(p)) == 1]
    cp = sp["test_idx"][data.is_control[sp["test_idx"]]]
    if len(cp) < 50:
        cp = data.control_idx
    return evaluate_forward(PivotPredictor(model, data, GPU), data, test, cp, max_perts=max_perts)


# ---- main forward tables (PIVOT row) ----
for ds, splits in [("norman", ["cell", "perturbation", "combination"]),
                   ("replogle_k562", ["cell", "perturbation", "gene"])]:
    data = load_dataset(ds)
    for sp in splits:
        j = json.load(open("%s/%s_forward_%s.json" % (RES, ds, sp)))["models"]["PIVOT"]
        m = load("models/%s/%s" % (ds, sp), data)
        cmp("forward %s/%s PIVOT" % (ds, sp), fwd_pivot(data, m, sp), j)

# ---- ablation tables (each row -> its checkpoint), norman/perturbation ----
data = load_dataset("norman")
A = "models/ablations/norman_perturbation"
comp = {"flow-map-only": "comp_map", "no-tangent": "comp_map_semi",
        "no-semigroup": "comp_map_tan", "PIVOT-full": "default"}
rep = {"gene_op": "default", "op_only": "rep_op_only", "gene_only": "rep_gene_only",
       "random_id": "rep_random_id", "gene_pathway_op": "rep_gene_pathway_op"}
frac = {"0.1": "frac_0.1", "0.25": "frac_0.25", "0.5": "frac_0.5", "0.75": "frac_0.75", "1.0": "default"}
match = {ms: ("default" if ms == "batch" else "match_%s" % ms) for ms in MATCH_STRATEGIES}

for jname, mapping in [("components", comp), ("representation", rep), ("datascale", frac), ("matching", match)]:
    rows = json.load(open("%s/norman_ablation_%s.json" % (RES, jname)))["rows"]
    for row, folder in mapping.items():
        if row not in rows:
            continue
        m = load("%s/%s" % (A, folder), data)
        r = _fwd_inv(data, m, "perturbation")
        cmp("ablation %s[%s] forward" % (jname, row), r["forward"], rows[row]["forward"])
        cmp("ablation %s[%s] inverse" % (jname, row), r["inverse"], rows[row]["inverse"])

# ---- summary ----
npass = sum(1 for _, ok, _ in checks if ok)
print()
for label, ok, diff in checks:
    print(("PASS " if ok else "FAIL ") + label + ("" if ok else "  mismatch=%s" % diff))
print("\n%d/%d checks reproduced within tol=%g" % (npass, len(checks), TOL))
print("RESULT:", "ALL REPRODUCED - no fakes" if npass == len(checks) else "MISMATCHES FOUND")