import json, os, re, subprocess, time
from pathlib import Path

PPL_BIN = "/content/llama.cpp/build/bin/llama-perplexity"
PPL_CPP = Path("/content/llama.cpp/tools/perplexity/perplexity.cpp")
assert "common_tokenize(ctx, params.prompt, true, true)" in PPL_CPP.read_text(), "perplexity.cpp missing parse_special=true patch!"
SHARD1 = "/content/models/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00001-of-00002.gguf"
SHARD2 = "/content/models/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00002-of-00002.gguf"

EVAL_DOMAINS = {
    "ops_tools": ("/content/eval_ops_tools.txt", "/content/base_logits_ops_tools.bin"),
    "chat_json": ("/content/eval_chat_json.txt", "/content/base_logits_chat_json.bin"),
    "ja":        ("/content/eval_ja.txt",        "/content/base_logits_ja.bin"),
}

def grab(pattern: str, text: str) -> float:
    m = re.search(pattern, text)
    if not m:
        raise RuntimeError(f"Failed to match {pattern} in output:\n{text[-1500:]}")
    return float(m.group(1))

results_path = Path("/content/phase1c_kld_results.json")
all_res = json.loads(results_path.read_text()) if results_path.exists() else {}

if "K512_unpruned" not in all_res:
    print("=== Step 1: Evaluating K512_unpruned baseline & saving teacher logits ===", flush=True)
    base_res = {}
    for dom, (fpath, lfile) in EVAL_DOMAINS.items():
        t0 = time.monotonic()
        cmd = [
            PPL_BIN, "-m", SHARD1, "-f", fpath,
            "-c", "2048", "-b", "512", "-ngl", "99",
            "--chunks", "4",
            "--save-all-logits", lfile,
        ]
        p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        if p.returncode != 0:
            raise RuntimeError(f"Baseline PPL failed on {dom} (exit {p.returncode}):\n{p.stdout[-1500:]}")
        dt = time.monotonic() - t0
        ppl = grab(r"Final estimate: PPL = ([0-9.]+)", p.stdout)
        ppl_se = grab(r"Final estimate: PPL = [0-9.]+\s*\+/-\s*([0-9.]+)", p.stdout)
        base_res[dom] = {"ppl": ppl, "ppl_se": ppl_se, "sec": round(dt, 1)}
        print(f"[K512_unpruned] {dom}: PPL = {ppl:.4f} +/- {ppl_se:.5f} ({dt:.1f}s)", flush=True)
    all_res["K512_unpruned"] = base_res
    results_path.write_text(json.dumps(all_res, indent=2))

CANDIDATES = [
    ("K320_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k320.json"),
    ("K320_reap_hybrid", "/content/keep_maps/keep_reap_hybrid_k320.json"),
    ("K256_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k256.json"),
    ("K256_reap_hybrid", "/content/keep_maps/keep_reap_hybrid_k256.json"),
    ("K256_reap_shrink", "/content/keep_maps/keep_reap_shrink_k256.json"),
    ("K224_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k224.json"),
]

os.makedirs("/dev/shm/pruned", exist_ok=True)
for label, km_path in CANDIDATES:
    if label in all_res:
        continue
    out1 = f"/dev/shm/pruned/{label}-00001-of-00002.gguf"
    out2 = f"/dev/shm/pruned/{label}-00002-of-00002.gguf"
    t0 = time.monotonic()
    print(f"=== Slicing {label} ===", flush=True)
    subprocess.run(["python3", "/content/slice_qwen4exp_gguf.py", "--in-gguf", SHARD1, "--out-gguf", out1, "--keep-json", km_path], check=True)
    subprocess.run(["python3", "/content/slice_qwen4exp_gguf.py", "--in-gguf", SHARD2, "--out-gguf", out2, "--keep-json", km_path], check=True)
    print(f"Sliced {label} in {time.monotonic()-t0:.1f}s", flush=True)

    cand_res = {}
    for dom, (fpath, lfile) in EVAL_DOMAINS.items():
        t1 = time.monotonic()
        cmd = [
            PPL_BIN, "-m", out1, "-f", fpath,
            "-c", "2048", "-b", "512", "-ngl", "99",
            "--chunks", "4",
            "--kl-divergence", "--kl-divergence-base", lfile,
        ]
        p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        if p.returncode != 0:
            raise RuntimeError(f"KLD eval failed for {label} on {dom} (exit {p.returncode}):\n{p.stdout[-1500:]}")
        dt = time.monotonic() - t1
        ppl_q = grab(r"Mean PPL\(Q\)\s*:\s*([0-9.]+)", p.stdout)
        ppl_base = grab(r"Mean PPL\(base\)\s*:\s*([0-9.]+)", p.stdout)
        ppl_diff = grab(r"Mean PPL\(Q\)-PPL\(base\)\s*:\s*([-+0-9.]+)", p.stdout)
        mean_kld = grab(r"Mean\s+KLD:\s*([-+0-9.]+)", p.stdout)
        mean_kld_se = grab(r"Mean\s+KLD:\s*[-+0-9.]+\s*±\s*([0-9.]+)", p.stdout)
        med_kld = grab(r"Median\s+KLD:\s*([-+0-9.]+)", p.stdout)
        p90_kld = grab(r"90\.0%\s+KLD:\s*([-+0-9.]+)", p.stdout)
        p99_kld = grab(r"99\.0%\s+KLD:\s*([-+0-9.]+)", p.stdout)
        rms_dp = grab(r"RMS Δp\s*:\s*([0-9.]+)", p.stdout)
        same_top = grab(r"Same top p:\s*([0-9.]+)", p.stdout)
        base_fp32 = all_res["K512_unpruned"][dom]["ppl"]
        cand_res[dom] = {
            "ppl": ppl_q,
            "ppl_base": ppl_base,
            "delta_ppl": ppl_diff,
            "exact_delta_ppl": round(ppl_q - base_fp32, 4),
            "mean_kld": mean_kld,
            "mean_kld_se": mean_kld_se,
            "median_kld": med_kld,
            "p90_kld": p90_kld,
            "p99_kld": p99_kld,
            "rms_delta_p_pct": rms_dp,
            "same_top1_pct": same_top,
            "sec": round(dt, 1),
        }
        print(f"[{label}] {dom}: PPL={ppl_q:.4f} (Δ={ppl_diff:+.4f}) | KLD={mean_kld:.5f}±{mean_kld_se:.5f} | Top1={same_top:.2f}% ({dt:.1f}s)", flush=True)

    all_res[label] = cand_res
    results_path.write_text(json.dumps(all_res, indent=2))
    os.remove(out1)
    os.remove(out2)

print("=== ALL PHASE 1C TIER-1 KLD + TOP-1 + PPL EVALS COMPLETE! ===", flush=True)