import json, os, numpy as np, torch, hashlib
D = os.path.join(os.path.dirname(__file__), "..", "out")
calls = [json.loads(l) for l in open(os.path.join(D, "events.jsonl")) if '"stage": "update"' in l]
sent = [json.loads(l) for l in open(os.path.join(D, "events.jsonl")) if '"stage": "sentinel"' in l]
def pool(cs, tag):
    if not cs: return {f"{tag}_n_calls": 0}
    n = sum(c["n_loss_tokens"] for c in cs); nz = sum(c["n_pos_adv"] + c["n_neg_adv"] for c in cs)
    lp, olp, a, m = [], [], [], []
    for c in cs:
        z = np.load(os.path.join(D, f"call-u{c['update']:02d}.npz")); mk = z["mask"] > 0
        lp.append(z["current_lp"][mk]); olp.append(z["sampler_lp"][mk]); a.append(z["adv"][mk])
    lp, olp, a = map(lambda x: torch.tensor(np.concatenate(x), dtype=torch.float64), (lp, olp, a))
    d = lp.exp() - olp.exp()
    return {f"{tag}_n_calls": len(cs), f"{tag}_n_loss_tokens": n, f"{tag}_n_nonzero_adv_tokens": nz,
        f"{tag}_dppo_masked_token_fraction": sum(c["dppo_masked"] for c in cs) / n,
        f"{tag}_dppo_masked_nonzero_adv_fraction": sum(c["dppo_masked"] for c in cs) / max(nz, 1),
        f"{tag}_ppo_clip_fraction": sum(c["ppo_clipped"] for c in cs) / n,
        f"{tag}_ppo_clip_nonzero_adv_fraction": sum(c["ppo_clipped"] for c in cs) / max(nz, 1),
        f"{tag}_policy_drift_abs_p_q90": float(torch.quantile(d.abs(), .9)), f"{tag}_policy_drift_abs_p_q99": float(torch.quantile(d.abs(), .99)),
        f"{tag}_policy_drift_abs_p_max": float(d.abs().max()), f"{tag}_signed_drift_min": float(d.min()), f"{tag}_signed_drift_max": float(d.max()),
        f"{tag}_pos_adv_drift_max": float(d[a > 0].max()), f"{tag}_neg_adv_drift_min": float(d[a < 0].min()),
        f"{tag}_abs_drift_gt_delta_tokens": int((d.abs() > .15).sum()),
        f"{tag}_removed_nonzero_grad_pos": sum(c["removed_nonzero_grad_pos"] for c in cs), f"{tag}_removed_nonzero_grad_neg": sum(c["removed_nonzero_grad_neg"] for c in cs),
        f"{tag}_ratio_cap_hits": sum(c["ratio_cap_hits"] for c in cs), f"{tag}_ratio_cap_hits_nonzero_adv": sum(c["ratio_cap_hits_nonzero_adv"] for c in cs),
        f"{tag}_n_pos_adv": sum(c["n_pos_adv"] for c in cs), f"{tag}_n_neg_adv": sum(c["n_neg_adv"] for c in cs)}
r = {**pool([c for c in calls if c["epoch"] == 1], "epoch1"), **pool([c for c in calls if c["epoch"] == 2], "epoch2")}
e2 = r.get("epoch2_n_calls", 0)
checks = None
if e2:
    checks = {"dppo_masked_in_range": .001 <= r["epoch2_dppo_masked_token_fraction"] <= .3, "ppo_clip_ge_0.0005": r["epoch2_ppo_clip_fraction"] >= .0005,
      "drift_q90_ge_0.005": r["epoch2_policy_drift_abs_p_q90"] >= .005,
      "removed_grad_pos_present": (r["epoch2_removed_nonzero_grad_pos"] > 0) if r["epoch2_n_pos_adv"] else None,
      "removed_grad_neg_present": (r["epoch2_removed_nonzero_grad_neg"] > 0) if r["epoch2_n_neg_adv"] else None,
      "all_8_second_epoch_calls": e2 == 8}
out = {"pooled": r, "revised_checks": checks, "sentinels": sent, "updates_executed": len(calls)}
json.dump(out, open(os.path.join(D, "analysis.json"), "w"), indent=1); print(json.dumps(out, indent=1)[:6000])