import json, pathlib, re
import pandas as pd
import numpy as np

ROOT = pathlib.Path('/Users/dawn/proj/prime-agent/.prime/agent/session-artifacts/01a06166-382f-703f-8f17-bc939488de93/observer-bench')

def score_multi_intent():
    p = ROOT / "results-multi-intent.jsonl"
    if not p.exists(): return None
    rows = [json.loads(l) for l in p.read_text().splitlines()]
    if not rows: return None
    
    records = []
    for r in rows:
        gold = r["gold"]
        res = r["result"]
        gold_k = gold["intent_count"]
        if "error" in res:
            records.append({"id": r["itemId"], "gold_k": gold_k, "pred_k": 0, "ok": False, "exact": False, "decomposed": False})
            continue
        try:
            out = json.loads(res["output"])
            adds = out.get("additions", [])
            pred_k = len(adds)
            exact = (pred_k == gold_k)
            decomposed = (pred_k > 1) if gold_k > 1 else (pred_k == 1)
            records.append({
                "id": r["itemId"],
                "gold_k": gold_k,
                "pred_k": pred_k,
                "ok": True,
                "exact": exact,
                "decomposed": decomposed,
                "ms": res.get("ms", 0)
            })
        except Exception:
            records.append({"id": r["itemId"], "gold_k": gold_k, "pred_k": 0, "ok": False, "exact": False, "decomposed": False})
    
    df = pd.DataFrame(records)
    summary = {
        "total": len(df),
        "exact_count_acc": round(df.exact.mean(), 3),
        "decomposition_rate": round(df.decomposed.mean(), 3),
        "by_gold_k": df.groupby("gold_k").agg(exact=("exact", "mean"), count=("exact", "count")).to_dict()
    }
    return summary, df

def score_supersession():
    p = ROOT / "results-supersession.jsonl"
    if not p.exists(): return None
    rows = [json.loads(l) for l in p.read_text().splitlines()]
    if not rows: return None
    
    records = []
    for r in rows:
        gold = r["gold"]
        res = r["result"]
        gold_exist = gold.get("goldExisting", {})
        gold_superseded_ids = {tid for tid, stat in gold_exist.items() if stat == "superseded"}
        
        if "error" in res:
            records.append({"id": r["itemId"], "ok": False, "tp": 0, "fp": 0, "fn": len(gold_superseded_ids)})
            continue
        try:
            out = json.loads(res["output"])
            exist = out.get("existing", [])
            pred_superseded_ids = {t["id"] for t in exist if t.get("status") == "superseded"}
            
            tp = len(pred_superseded_ids & gold_superseded_ids)
            fp = len(pred_superseded_ids - gold_superseded_ids)
            fn = len(gold_superseded_ids - pred_superseded_ids)
            
            records.append({
                "id": r["itemId"],
                "ok": True,
                "gold_superseded": len(gold_superseded_ids),
                "pred_superseded": len(pred_superseded_ids),
                "tp": tp,
                "fp": fp,
                "fn": fn,
                "hit": (pred_superseded_ids == gold_superseded_ids),
                "ms": res.get("ms", 0)
            })
        except Exception:
            records.append({"id": r["itemId"], "ok": False, "tp": 0, "fp": 0, "fn": len(gold_superseded_ids)})
            
    df = pd.DataFrame(records)
    total_gold = df.tp.sum() + df.fn.sum()
    recall = round(df.tp.sum() / max(1, total_gold), 3)
    precision = round(df.tp.sum() / max(1, df.tp.sum() + df.fp.sum()), 3)
    
    summary = {
        "total": len(df),
        "gold_supersessions": int(total_gold),
        "superseded_recall": recall,
        "superseded_precision": precision,
        "exact_match_rate": round(df.hit.mean(), 3)
    }
    return summary, df

def score_real_candidates():
    p = ROOT / "results-real-candidates.jsonl"
    if not p.exists(): return None
    rows = [json.loads(l) for l in p.read_text().splitlines()]
    if not rows: return None
    
    records = []
    for r in rows:
        res = r["result"]
        if "error" in res:
            records.append({"id": r["itemId"], "ok": False})
            continue
        try:
            out = json.loads(res["output"])
            adds = out.get("additions", [])
            exist = out.get("existing", [])
            # check role-restricted enum: 'done' must never appear
            all_statuses = [t.get("status") for t in (adds + exist)]
            has_done = "done" in all_statuses
            # check multi-intent decomposition
            n_threads = len(adds) + len(exist)
            # check evidence action citations
            has_citations = any(len(t.get("evidenceActionIds", [])) > 0 for t in (adds + exist))
            
            records.append({
                "id": r["itemId"],
                "ok": True,
                "has_done": has_done,
                "n_threads": n_threads,
                "has_citations": has_citations,
                "ms": res.get("ms", 0)
            })
        except Exception:
            records.append({"id": r["itemId"], "ok": False})
            
    df = pd.DataFrame(records)
    summary = {
        "total": len(df),
        "valid_json_rate": round(df.ok.mean(), 3),
        "no_done_guarantee": bool((~df.has_done).all()),
        "avg_threads": round(df.n_threads.mean(), 2),
        "action_citation_rate": round(df.has_citations.mean(), 3)
    }
    return summary, df

if __name__ == "__main__":
    print("=== Multi-Intent Segmentation ===")
    res_mi = score_multi_intent()
    if res_mi: print(json.dumps(res_mi[0], indent=2))
    
    print("\n=== Supersession Detection ===")
    res_sup = score_supersession()
    if res_sup: print(json.dumps(res_sup[0], indent=2))
    
    print("\n=== Real Coding Candidates ===")
    res_real = score_real_candidates()
    if res_real: print(json.dumps(res_real[0], indent=2))