"""Development: run the stage-1 search space screen on a labelled model and 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("--watch", 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(); tm = {}
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")
g_m = probe.grad_screen(m, tok, "cpu"); tm["suspect_leak_grad"] = time.time() - t0
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
g_p = probe.grad_screen(p, ptok, "cpu")
keep = probe.candidate_mask(tok, m.get_input_embeddings().weight.shape[0])
z = lambda s: (s - s[keep].median()) / (s[keep].std() + 1e-9)
leak_c = probe.candidates(tok, seqs, lpm, lpp, k=32, special=set(tok.all_special_ids), n_max=5)
space = probe.search_space(tok, leak_c, z(g_m) - z(g_p), keep)
strings = list(space) + [w for w in a.watch if w not in space]
t1 = time.time(); sp = probe.screen(p, ptok, strings, "cpu"); tm["parent_screen"] = time.time() - t1
t1 = time.time(); sm = probe.screen(m, tok, strings, "cpu"); tm["suspect_screen"] = time.time() - t1
d = sm - sp
order = np.argsort(-d)
rank = {strings[i]: r for r, i in enumerate(order)}
out = {"model": a.model, "n_strings": len(strings), "seconds": {k: round(v) for k, v in tm.items()},
       "watch": {w: {"rank": rank[w], "diff": round(float(d[strings.index(w)]), 3)} for w in a.watch},
       "top": [{"text": strings[i], "source": space.get(strings[i], "watch"), "diff": round(float(d[i]), 3),
                "suspect": round(float(sm[i]), 3), "parent": round(float(sp[i]), 3)} for i in order[:40]],
       "quantiles": {q: round(float(np.quantile(d, q)), 3) for q in (0.5, 0.9, 0.99, 0.999)}}
json.dump(out, open(a.out, "w"), indent=1, ensure_ascii=False)
print(json.dumps({k: out[k] for k in ("n_strings", "seconds", "watch", "quantiles")}, ensure_ascii=False))
for r in out["top"][:25]: print(r)
print("total minutes", round((time.time() - t0) / 60, 1))