"""Compare a retraining (reproduce.sh) with the shipped results.

    python -I compare_reproduction.py [TAG=repro]

fold 0 : models/<TAG>_fold0/test.csv  vs models/ej_ad_mv_fold0/test.csv
5-fold : results/final_cv_long_<TAG>.csv vs results/final_cv_long.csv (rule ens-es and stage 1, mean over folds)
Stage 1 (mIF / RNA / text heads) is seeded and deterministic on the same GPU / library stack, so the emergent pairs
should agree to ~1e-3; stage 2 (H&E adapter, 3 early-stopped seeds) is not bit-deterministic on GPU, so the trained
H&E pairs agree within the run-to-run spread (5-fold seed std of a single adapter: mIF 0.009, text 0.011).
Tolerances: |diff| <= 0.01 (emergent / stage 1), <= 0.03 (fold-0 H&E pairs), <= 0.015 (5-fold means of H&E pairs).
"""
import os
import sys

import pandas as pd

ROOT = os.path.dirname(os.path.abspath(__file__))
TAG = sys.argv[1] if len(sys.argv) > 1 else "repro"
pd.set_option("display.width", 200)
ok = True


def show(D, tol_trained, title):
    global ok
    D["abs_diff"] = (D.reproduced - D.shipped).abs()
    D["tol"] = [tol_trained if p.startswith("H&E") and r != "stage1" else 0.01 for p, r in zip(D.pair, D.get("rule", [""] * len(D)))]
    D["ok"] = D.abs_diff <= D.tol
    ok &= bool(D.ok.all())
    print(f"\n{title}\n{D.round(4).to_string(index=False)}")


f0 = os.path.join(ROOT, "models", f"{TAG}_fold0", "test.csv")
if os.path.exists(f0):
    a = pd.read_csv(f0).set_index("pair")["R@1"]
    b = pd.read_csv(os.path.join(ROOT, "models", "ej_ad_mv_fold0", "test.csv")).set_index("pair")["R@1"]
    show(pd.DataFrame({"pair": a.index, "reproduced": a.values, "shipped": b.loc[a.index].values}), 0.03,
         "fold 0, held-out test R@1 (g100)")
cv = os.path.join(ROOT, "results", f"final_cv_long_{TAG}.csv")
if os.path.exists(cv):
    m = lambda f: pd.read_csv(f).query("rule in ['stage1', 'ens-es']").groupby(["rule", "pair"], sort=False)["R@1"].mean()
    a, b = m(cv), m(os.path.join(ROOT, "results", "final_cv_long.csv"))
    show(pd.DataFrame({"rule": [r for r, _ in a.index], "pair": [p for _, p in a.index], "reproduced": a.values,
                       "shipped": b.loc[a.index].values}), 0.015, "5-fold mean test R@1 (g100)")
print("\nREPRODUCED within tolerance" if ok else "\nNOT reproduced within tolerance (see ok=False rows)")
sys.exit(0 if ok else 1)