import json
import re
from pathlib import Path

import pandas as pd

ROOT = Path("completed_experiments/evals/full-2026-10-23/outputs")
CORPUS = Path("questions/corpus.jsonl")
OUT = Path("analysis/designed-difficulty")

MODELS = [
    "qwen3.5-4b",
    "gemma-4-e4b",
    "granite-4.2-3b",
    "nemotron-3-nano-4b",
    "k2-horizon-3.7b",
]

PATTERN = re.compile(
    r"^(?P<variant>[^/]+)/(?P<question_id>[^#]+)"
    r"#e(?P<epoch>\d+)#t(?P<turn>\d+)$"
)

# 1. Load question metadata
questions = {}

with CORPUS.open(encoding="utf-8") as f:
    for line in f:
        if not line.strip():
            continue
        q = json.loads(line)
        qid = q["id"]

        if qid in questions:
            raise ValueError(f"Duplicate question ID: {qid}")

        difficulty = q.get("difficulty") or {}

        questions[qid] = {
            "claim_id": q.get("claim_id"),
            "category": q.get("category"),
            "framing": q.get("framing"),
            "evidence_scarcity": difficulty.get("evidence_scarcity"),
            "question_openness": difficulty.get("question_openness"),
            "band": difficulty.get("band"),
        }

print("Corpus questions:", len(questions))

# 2. Load judge results
rows = []
unparsed = []

for model in MODELS:
    path = ROOT / model / "pplx/results.jsonl"

    if not path.exists():
        raise FileNotFoundError(path)

    count = 0

    with path.open(encoding="utf-8") as f:
        for line_no, line in enumerate(f, 1):
            if not line.strip():
                continue

            result = json.loads(line)
            case_id = result.get("case_id", "")
            match = PATTERN.fullmatch(case_id)

            if not match:
                unparsed.append((model, line_no, case_id))
                continue

            parts = match.groupdict()
            qid = parts["question_id"]

            meta = questions.get(qid, {})
            probabilities = result.get("probabilities") or {}

            rows.append({
                "model": model,
                "variant": parts["variant"],
                "question_id": qid,
                "epoch": int(parts["epoch"]),
                "turn": int(parts["turn"]),
                "case_id": case_id,
                "choice": result.get("choice"),
                "status": result.get("status"),
                "dropped": result.get("dropped"),
                "confidence": result.get("confidence"),
                "p_hold": probabilities.get("hold"),
                "p_diverge": probabilities.get("diverge"),
                "p_abstain": probabilities.get("abstain"),
                **meta,
            })
            count += 1

    print(f"{model}: {count:,} records")

df = pd.DataFrame(rows)

if df.empty:
    raise RuntimeError("No parsed Judge records found")

# 3. Integrity checks
key_cols = ["model", "case_id"]

duplicates = df.duplicated(key_cols, keep=False)
missing_metadata = df["evidence_scarcity"].isna() | df["question_openness"].isna()

valid_labels = {"hold", "diverge", "abstain"}

df["valid"] = (
    df["status"].eq("ok")
    & df["dropped"].eq(0)
    & df["choice"].isin(valid_labels)
)

df["diverged"] = df["choice"].eq("diverge") & df["valid"]

print("\n=== INTEGRITY CHECK ===")
print("Total records:", len(df))
print("Unparsed case IDs:", len(unparsed))
print("Duplicate model/case IDs:", int(duplicates.sum()))
print("Missing difficulty metadata:", int(missing_metadata.sum()))
print("Valid Judge records:", int(df["valid"].sum()))

print("\nStatus distribution:")
print(df["status"].value_counts(dropna=False).to_string())

print("\nChoice distribution:")
print(df["choice"].value_counts(dropna=False).to_string())

print("\nVariant distribution:")
print(pd.crosstab(df["model"], df["variant"]).to_string())

if unparsed:
    print("\nUnparsed examples:", unparsed[:10])

# 4. Save the raw joined dataset
OUT.mkdir(parents=True, exist_ok=True)

output = OUT / "baseline_joined.csv"
df.to_csv(output, index=False)

print(f"\nSaved: {output}")

# 5. Preliminary summary (only if integrity checks pass)
if duplicates.any() or missing_metadata.any() or unparsed:
    print("\nWARNING: Resolve integrity issues before computing rates.")
else:
    opening = df[(df["variant"] == "multiturn") & (df["turn"] == 0)]

    print("\n=== MULTITURN OPENING: DIVERGENCE RATE ===")
    print("Denominator: valid hold/diverge/abstain judgments")

    summary = (
        opening[opening["valid"]]
        .groupby("model")
        .agg(
            n=("diverged", "size"),
            questions=("question_id", "nunique"),
            diverged=("diverged", "sum"),
            rate=("diverged", "mean"),
        )
        .reset_index()
    )

    summary["rate_pct"] = (summary["rate"] * 100).round(2)

    print(
        summary[
            ["model", "questions", "n", "diverged", "rate_pct"]
        ].to_string(index=False)
    )