from pathlib import Path
import pandas as pd

OUT = Path("analysis/designed-difficulty")
df = pd.read_csv(OUT / "baseline_joined.csv")

# Recompute validity from the original Judge fields.
df["valid"] = (
    df["status"].eq("ok")
    & df["dropped"].eq(0)
    & df["choice"].isin(["hold", "diverge", "abstain"])
)
df["diverged"] = df["choice"].eq("diverge")

# Focus on the opening answer of multi-turn evaluations.
opening = df[
    (df["variant"] == "multiturn") &
    (df["turn"] == 0)
].copy()

print("=== 1. DROPPED RECORD AUDIT ===")
print("All Judge records:")
print(df["dropped"].value_counts(dropna=False).head(15))

print("\nOpening records by dropped value:")
print(pd.crosstab(
    opening["dropped"],
    opening["choice"],
    dropna=False
).head(15).to_string())

print("\n=== 2. FILTER SENSITIVITY ===")
for model, group in opening.groupby("model"):
    all_ok = group[
        group["status"].eq("ok")
        & group["choice"].isin(["hold", "diverge", "abstain"])
    ]
    strict = group[group["valid"]]

    print(
        f"{model:25s} "
        f"all_ok={all_ok['diverged'].mean():.4f} "
        f"dropped_zero={strict['diverged'].mean():.4f} "
        f"n={len(strict)}/{len(all_ok)}"
    )

valid = opening[opening["valid"]].copy()

print("\n=== 3. DIFFICULTY DISTRIBUTION ===")
question_meta = (
    df[[
        "question_id",
        "evidence_scarcity",
        "question_openness",
        "band"
    ]]
    .drop_duplicates()
)
print("Unique questions:", question_meta["question_id"].nunique())
print("\nDifficulty bands:")
print(question_meta["band"].value_counts().to_string())

print("\nEvidence scarcity x question openness:")
print(pd.crosstab(
    question_meta["evidence_scarcity"],
    question_meta["question_openness"]
).to_string())

def summarize(data, keys):
    result = (
        data.groupby(keys, dropna=False)
        .agg(
            n_questions=("question_id", "nunique"),
            n_judgments=("diverged", "size"),
            n_diverge=("diverged", "sum"),
            divergence_rate=("diverged", "mean"),
        )
        .reset_index()
    )
    result["rate_pct"] = (
        result["divergence_rate"] * 100
    ).round(2)
    return result

print("\n=== 4. EVIDENCE SCARCITY ===")
scarcity = summarize(valid, ["evidence_scarcity"])
print(scarcity.to_string(index=False))

print("\n=== 5. QUESTION OPENNESS ===")
openness = summarize(valid, ["question_openness"])
print(openness.to_string(index=False))

print("\n=== 6. TWO-DIMENSIONAL DIFFICULTY ===")
matrix = summarize(
    valid, ["evidence_scarcity", "question_openness"]
)
print(matrix.to_string(index=False))

print("\n=== 7. MODEL-SPECIFIC DIFFICULTY ===")
by_model = summarize(
    valid,
    ["model", "evidence_scarcity", "question_openness"]
)
print(by_model.to_string(index=False))

print("\n=== 8. SAVE RESULTS ===")
scarcity.to_csv(OUT / "scarcity_summary.csv", index=False)
openness.to_csv(OUT / "openness_summary.csv", index=False)
matrix.to_csv(OUT / "difficulty_matrix.csv", index=False)
by_model.to_csv(OUT / "difficulty_by_model.csv", index=False)

print("Saved four summary CSV files.")