"""
CALIBRATION TIE-DEGENERACY DISCLOSURE CHECK (Section 5.5, Section 7).

This script quantifies and discloses a limitation of the four-quartile
peer task-difficulty calibration analysis in calibration_analysis.py:
because the large majority of agents in the tau-bench corpus self-report
success, the "realistic" (peer self-report) signal takes the value
exactly 1.0 for a large fraction of trials, meaning several of the four
quartile boundaries are determined by arbitrary rank-order tie-breaking
among trials with an IDENTICAL signal value, not by genuine differences
in the underlying peer self-report rate.

This script:
  1. Reproduces the leave-one-out realistic peer signal computed in
     calibration_analysis.py.
  2. Quantifies exactly how many trials are tied at the maximum value
     (1.0), and which of the four rank-based quartiles these tied trials
     fall into.
  3. Reports a cleaner two-group comparison (strictly-below-maximum vs.
     tied-at-maximum) with 95% Wilson confidence intervals, which is the
     more defensible characterization of this signal's true
     discriminative resolution and is the version referenced in the
     manuscript's Section 5.5 and Section 7 disclosures.

USAGE: place the same 4 tau-bench JSON files used by real_data_pipeline.py
and calibration_analysis.py in the same folder, then run:
    python3 calibration_tie_degeneracy_check.py
"""
import json
import pandas as pd
from self_report_classifier import get_final_user_facing_message, self_reported_success_score
from real_data_pipeline import wilson_ci

DOMAINS = {
    "retail": {"gpt-4o": "gpt-4o-retail.json", "sonnet": "sonnet-35-new-retail.json"},
    "airline": {"gpt-4o": "gpt-4o-airline.json", "sonnet": "sonnet-35-new-airline.json"},
}


def log(msg):
    print(msg, flush=True)


def load_all_trials():
    rows = []
    for domain, model_files in DOMAINS.items():
        for model, fname in model_files.items():
            with open(fname) as f:
                data = json.load(f)
            for t in data:
                msg = get_final_user_facing_message(t["traj"])
                self_report = self_reported_success_score(msg)
                rows.append({
                    "domain": domain, "model": model, "task_id": t["task_id"],
                    "trial": t["trial"], "reward": float(t["reward"]),
                    "self_report": self_report, "item_key": f"{domain}_{t['task_id']}",
                })
    return pd.DataFrame(rows)


def compute_leave_one_out_realistic_signal(df):
    results = []
    grouped = df.groupby("item_key")
    for item_key, group in grouped:
        n = len(group)
        sum_self_success = (group["self_report"] >= 0.5).sum()
        for _, row in group.iterrows():
            peer_n = n - 1
            if peer_n <= 0:
                continue
            realistic_peer_rate = (sum_self_success - (1 if row["self_report"] >= 0.5 else 0)) / peer_n
            results.append({
                "reward": row["reward"], "self_report": row["self_report"],
                "realistic_peer_rate": realistic_peer_rate,
            })
    return pd.DataFrame(results)


def main():
    log("Loading all real trajectories...")
    df = load_all_trials()

    log("\nComputing leave-one-out realistic peer signal...")
    peer_signals = compute_leave_one_out_realistic_signal(df)

    success_claims = peer_signals[peer_signals.self_report >= 0.5].copy()
    success_claims["false_success"] = (success_claims.reward == 0).astype(int)
    n_total = len(success_claims)
    log(f"Self-reported success claims with valid peer signals: {n_total}")

    # Reproduce the exact rank-based quartile assignment used in calibration_analysis.py
    labels = ["Q1 (lowest)", "Q2", "Q3", "Q4 (highest)"]
    ranks = success_claims["realistic_peer_rate"].rank(method="first")
    success_claims["bucket"] = pd.qcut(ranks, 4, labels=labels)

    log("\n=== Tie degeneracy by quartile ===")
    for lab in labels:
        sub = success_claims[success_claims.bucket == lab]
        n_at_one = (sub["realistic_peer_rate"] == 1.0).sum()
        log(f"{lab}: n={len(sub)}, range=[{sub['realistic_peer_rate'].min():.3f},"
            f"{sub['realistic_peer_rate'].max():.3f}], "
            f"n_tied_at_1.0={n_at_one} ({n_at_one/len(sub):.1%})")

    n_at_one_total = (success_claims["realistic_peer_rate"] == 1.0).sum()
    log(f"\nTOTAL tied at exactly 1.0: {n_at_one_total}/{n_total} = "
        f"{n_at_one_total/n_total:.1%} of all success-claim trials")

    log("\n=== Cleaner two-group comparison (recommended characterization) ===")
    below = success_claims[success_claims["realistic_peer_rate"] < 1.0]
    at_max = success_claims[success_claims["realistic_peer_rate"] == 1.0]
    results = {}
    for label, sub in [("below_maximum", below), ("tied_at_maximum", at_max)]:
        n = len(sub)
        k = int(sub["false_success"].sum())
        lo, hi = wilson_ci(k, n)
        log(f"{label}: n={n}, false_success_rate={k/n:.1%}, "
            f"95% CI=[{lo*100:.1f},{hi*100:.1f}]")
        results[label] = {"n": n, "false_success_rate": round(k / n, 4),
                           "ci_lo_pct": round(lo * 100, 1), "ci_hi_pct": round(hi * 100, 1)}

    with open("calibration_tie_degeneracy_summary.json", "w") as f:
        json.dump({
            "n_total_success_claims": n_total,
            "n_tied_at_maximum": int(n_at_one_total),
            "pct_tied_at_maximum": round(n_at_one_total / n_total, 4),
            "two_group_comparison": results,
        }, f, indent=2)
    log("\nSaved: calibration_tie_degeneracy_summary.json")


if __name__ == "__main__":
    main()