"""
PEER TASK-DIFFICULTY CALIBRATION ANALYSIS (Section 5.5, Figure 4, Table 4).

For each task, and for each trial i on that task in which the agent itself
self-reported success (self_report >= 0.5), this script computes two
leave-one-out peer signals EXCLUDING trial i's own outcome:
  - oracle_peer_rate:      the true ground-truth success rate among trial i's
                            peers on the same task (NOT deployable in practice,
                            since it requires already knowing peers' verified
                            outcomes).
  - realistic_peer_rate:   the self-reported success rate among trial i's
                            peers on the same task (deployable in practice,
                            since it only requires peers' own self-reports).

Trials are then split into four EQUAL-SIZED quartiles (by rank, to handle
heavy ties in the realistic signal robustly) separately for each of the two
signals, and the false-success rate (fraction of self-reported "success"
claims that were, per ground truth, actually failures) is computed within
each quartile, along with a 95% Wilson score confidence interval.

IMPORTANT METHODOLOGICAL NOTE: an earlier draft of this analysis used
non-leave-one-out, fixed-percentage-range buckets, which both leaked each
trial's own outcome into its own "peer" aggregate and produced bucket
membership counts inconsistent with this benchmark's real self-report
distribution (most agents self-report success on this benchmark, so peer
self-report rates cluster near the top of the range, not evenly across
0-100%). This script implements the corrected, leaked-free, equal-population
quartile methodology and is the version whose output is reported in the
manuscript.

USAGE: place the same 4 tau-bench JSON files used by real_data_pipeline.py
in the same folder, then run: python3 calibration_analysis.py
"""
import json
import pandas as pd
import numpy as np
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_peer_signals(df):
    """For each trial, compute oracle_peer_rate and realistic_peer_rate
    excluding that trial's own outcome, using only other trials on the
    same (domain, task_id)."""
    results = []
    grouped = df.groupby("item_key")
    for item_key, group in grouped:
        n = len(group)
        sum_reward = group["reward"].sum()
        sum_self_success = (group["self_report"] >= 0.5).sum()
        for _, row in group.iterrows():
            peer_n = n - 1
            if peer_n <= 0:
                continue
            oracle_peer_rate = (sum_reward - row["reward"]) / peer_n
            realistic_peer_rate = (sum_self_success - (1 if row["self_report"] >= 0.5 else 0)) / peer_n
            results.append({
                "item_key": item_key, "reward": row["reward"], "self_report": row["self_report"],
                "oracle_peer_rate": oracle_peer_rate, "realistic_peer_rate": realistic_peer_rate,
            })
    return pd.DataFrame(results)


def quartile_stats_with_ci(df_, col, labels=("Q1 (lowest)", "Q2", "Q3", "Q4 (highest)")):
    """Equal-population quartiles via rank-based qcut (robust to ties),
    with false-success rate and 95% Wilson CI per quartile."""
    df_ = df_.copy()
    ranks = df_[col].rank(method="first")
    df_["bucket"] = pd.qcut(ranks, 4, labels=labels)
    out = []
    for lab in labels:
        sub = df_[df_.bucket == lab]
        n = len(sub)
        k = int(sub["false_success"].sum())
        rate = k / n if n > 0 else float("nan")
        lo, hi = wilson_ci(k, n)
        vmin, vmax = sub[col].min(), sub[col].max()
        out.append({"bucket": lab, "n": n, "false_success_rate_pct": round(rate * 100, 1),
                     "ci_lo_pct": round(lo * 100, 1), "ci_hi_pct": round(hi * 100, 1),
                     "signal_range": f"[{vmin:.3f}, {vmax:.3f}]"})
    return pd.DataFrame(out)


def main():
    log("Loading all real trajectories...")
    df = load_all_trials()
    log(f"Loaded {len(df)} total trials across {df.item_key.nunique()} unique tasks.")

    log("\nComputing leave-one-out peer signals (oracle and realistic)...")
    peer_signals = compute_leave_one_out_peer_signals(df)

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

    log("\n=== ORACLE (leave-one-out peer ground truth) quartiles ===")
    oracle_df = quartile_stats_with_ci(success_claims, "oracle_peer_rate")
    log(oracle_df.to_string(index=False))

    log("\n=== REALISTIC (leave-one-out peer self-report) quartiles ===")
    realistic_df = quartile_stats_with_ci(success_claims, "realistic_peer_rate")
    log(realistic_df.to_string(index=False))

    oracle_df.to_csv("oracle_calibration_quartiles.csv", index=False)
    realistic_df.to_csv("realistic_calibration_quartiles.csv", index=False)

    with open("calibration_analysis_summary.json", "w") as f:
        json.dump({
            "n_success_claims": len(success_claims),
            "oracle_quartiles": oracle_df.to_dict("records"),
            "realistic_quartiles": realistic_df.to_dict("records"),
        }, f, indent=2, default=str)
    log("\nSaved: oracle_calibration_quartiles.csv, realistic_calibration_quartiles.csv, "
        "calibration_analysis_summary.json")


if __name__ == "__main__":
    main()