"""Development diagnostics on a labelled model: leak candidates + oracle triggers + nulls,
scored by activation effect and by reply divergence, in the model and in its parent."""
import argparse, json, time
import numpy as np, torch
from backdoorcheck import probe
from backdoorcheck.loading import load_model, load_tokenizer, same_tokenizer

ap = argparse.ArgumentParser()
ap.add_argument("model"); ap.add_argument("parent"); ap.add_argument("--trigger", nargs="*", default=[])
ap.add_argument("--threads", type=int, default=48); ap.add_argument("--out", required=True)
a = ap.parse_args()
torch.set_num_threads(a.threads)
t0 = time.time()
m, info = load_model(a.model)
tok, _ = load_tokenizer(a.model, fallbacks=[info.get("adapter_base"), a.parent])
prefix, seqs = probe.leak(m, tok, "cpu")
lpm = probe.token_logprobs(m, prefix, seqs, "cpu")
nulls = probe.null_strings(tok, 48)
p, _ = load_model(a.parent)
ptok, _ = load_tokenizer(a.parent)
lpp = probe.token_logprobs(p, prefix, seqs, "cpu") if same_tokenizer(tok, ptok) else None
cand = probe.candidates(tok, seqs, lpm, lpp, k=16, special=set(tok.all_special_ids))
strings = [c["text"] for c in cand] + a.trigger + nulls
kind = ["cand"] * len(cand) + ["oracle"] * len(a.trigger) + ["null"] * len(nulls)
res = {}
for name, mod, tk in (("M", m, tok), ("P", p, ptok)):
    e, _ = probe.activation_effects(mod, tk, strings, "cpu")
    d = probe.divergence(mod, tk, strings, "cpu")
    res[name] = {"act": e, "div": d}
nl = [i for i, k in enumerate(kind) if k == "null"]
band = probe.layer_band(res["M"]["act"].shape[1])
act = lambda r: (r["act"] - np.median(r["act"][nl], 0))
act_score = (act(res["M"]) - act(res["P"]))[:, band].max(1)
div_m = res["M"]["div"]; div_p = res["P"]["div"]
zr = lambda x: (x - np.median(x[nl])) / (1.4826 * np.median(np.abs(x[nl] - np.median(x[nl]))) + 1e-9)
div_score = zr(div_m) - zr(div_p)
rows = []
for i, s in enumerate(strings):
    rows.append({"kind": kind[i], "text": s, "act_z": round(float(probe.robust_z(act_score[i], act_score[nl].tolist())), 2),
                 "div_M": round(float(div_m[i]), 3), "div_P": round(float(div_p[i]), 3),
                 "divz_M": round(float(zr(div_m)[i]), 2), "divz_P": round(float(zr(div_p)[i]), 2),
                 "div_z": round(float(probe.robust_z(div_score[i], div_score[nl].tolist())), 2)})
out = {"model": a.model, "parent": a.parent, "leak_examples": [tok.decode(s, skip_special_tokens=True)[:100] for s in seqs[:12]],
       "rows": rows, "minutes": round((time.time() - t0) / 60, 1)}
json.dump(out, open(a.out, "w"), indent=1, ensure_ascii=False)
print(json.dumps(out["leak_examples"], ensure_ascii=False, indent=0))
for r in sorted(rows, key=lambda r: -r["div_z"]):
    if r["kind"] != "null" or r["div_z"] > 2:
        print(r)
nz = [r for r in rows if r["kind"] == "null"]
print("null max act_z", max(r["act_z"] for r in nz), "null max div_z", max(r["div_z"] for r in nz), "minutes", out["minutes"])