#!/usr/bin/env python3
"""Run the FP8-vs-bf16 dose-response check against the Space and write the comparison.

Needs the Space running with RAT_VALIDATE=1 (registers validate_fp8.py's endpoints) and HF_TOKEN.
Reference: out_v2/SUMMARY.json["27B"]["study1"] from the bf16 battery on an H100 (24 samples/dose).

    python run_validation.py ../../osf_bundle/data/out_v2/SUMMARY.json
"""
import json
import math
import os
import sys

from gradio_client import Client

SPACE = "devon7y/rat-qwen"
DOSES = [0.0, 0.35, 0.5, 0.75, 1.0, 1.5, 2.0]
SAMPLES = 16                      # per probe; 6 probes -> 96 per dose (the study had 24)
HERE = os.path.dirname(os.path.abspath(__file__))


def wilson(k, n, z=1.96):
    if n == 0:
        return (0.0, 0.0)
    p = k / n
    c = (p + z * z / (2 * n)) / (1 + z * z / n)
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / (1 + z * z / n)
    return (max(0.0, c - h), min(1.0, c + h))


def fisher_two_sided(a, b, c, d):
    """Fisher exact test on [[a, b], [c, d]]."""
    n1, n2, k = a + b, c + d, a + c
    def pmf(x):
        return math.comb(n1, x) * math.comb(n2, k - x) / math.comb(n1 + n2, k)
    p0 = pmf(a)
    lo, hi = max(0, k - n2), min(k, n1)
    return min(1.0, sum(pmf(x) for x in range(lo, hi + 1) if pmf(x) <= p0 * (1 + 1e-9)))


def main():
    ref = {round(r["dose"], 2): r for r in json.load(open(sys.argv[1]))["27B"]["study1"]["per_dose"]}
    cli = Client(SPACE, token=os.environ["HF_TOKEN"])
    diag_path = os.path.join(HERE, "diagnostics.json")
    if os.path.exists(diag_path):
        diag = json.load(open(diag_path))
    else:
        diag = cli.predict(api_name="/diagnostics")
        json.dump(diag, open(diag_path, "w"), indent=1)
    print("diagnostics:", json.dumps(diag), flush=True)

    rows = []
    with open(os.path.join(HERE, "fp8_transcripts.jsonl"), "w") as tf:
        for dose in DOSES:
            out = cli.predict(dose, SAMPLES, api_name="/probe_sweep")
            for r in out["rows"]:
                tf.write(json.dumps({"dose": dose, **r}) + "\n")
                rows.append({"dose": dose, **r})
            tf.flush()
            got = out["rows"]
            print(f"dose {dose}: intrusion {sum(r['hit'] for r in got)}/{len(got)}"
                  f"  ({out['seconds']} s, worker {out['pid']})", flush=True)

    table = []
    for dose in DOSES:
        got = [r for r in rows if r["dose"] == dose]
        n8, k8 = len(got), sum(r["hit"] for r in got)
        b = ref[round(dose, 2)]
        n16 = 24
        k16 = round(b["probe_theme_intrusion_rate"] * n16)
        table.append({
            "dose": dose, "slider": round(dose * 20.73, 1),
            "fp8_intrusion": k8 / n8, "fp8_ci": wilson(k8, n8), "fp8_k": k8, "fp8_n": n8,
            "bf16_intrusion": k16 / n16, "bf16_ci": wilson(k16, n16), "bf16_k": k16, "bf16_n": n16,
            "fisher_p": fisher_two_sided(k8, n8 - k8, k16, n16 - k16),
            "fp8_persev": sum(r["persev"] for r in got) / n8, "bf16_persev": b["probe_perseveration"],
            "fp8_len": sum(r["len"] for r in got) / n8, "bf16_len": b["probe_mean_len"],
        })
    json.dump({"diagnostics": diag, "bf16_norm_scale": json.load(open(sys.argv[1]))["27B"]["study1"]["norm_scale"],
               "table": table}, open(os.path.join(HERE, "fp8_vs_bf16.json"), "w"), indent=1)

    print("\n| dose | slider | FP8 intrusion [95% CI] | bf16 intrusion [95% CI] | Fisher p | FP8 / bf16 persev | FP8 / bf16 words |")
    print("|---|---|---|---|---|---|---|")
    for t in table:
        print(f"| {t['dose']:.2f} | {t['slider']} | {t['fp8_k']}/{t['fp8_n']} = {t['fp8_intrusion']:.2f} "
              f"[{t['fp8_ci'][0]:.2f}, {t['fp8_ci'][1]:.2f}] | {t['bf16_k']}/{t['bf16_n']} = {t['bf16_intrusion']:.2f} "
              f"[{t['bf16_ci'][0]:.2f}, {t['bf16_ci'][1]:.2f}] | {t['fisher_p']:.3f} | "
              f"{t['fp8_persev']:.2f} / {t['bf16_persev']:.2f} | {t['fp8_len']:.0f} / {t['bf16_len']:.0f} |")


if __name__ == "__main__":
    main()