"""
ROBUSTNESS CHECK: Same-Model-Only Peer Groups

The main manuscript's leave-one-trial-out evaluation pools GPT-4o and
Claude-3.5-Sonnet trials together as a single peer group per task. This
script tests whether that pooling choice affects the reported finding
(that peer aggregation underperforms solo self-report) by re-running the
identical protocol with peer groups restricted to SAME-MODEL trials only.

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

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"},
}
RNG_SEED = 42


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_model_reputation_same_model(df):
    df = df.copy()
    df["faithful"] = (
        ((df.self_report == 1.0) & (df.reward == 1.0)) |
        ((df.self_report == 0.0) & (df.reward == 0.0))
    ).astype(int)
    reputation = {}
    for domain in df.domain.unique():
        ddf = df[df.domain == domain]
        for task_id in ddf.task_id.unique():
            other = ddf[ddf.task_id != task_id]
            for model in ddf.model.unique():
                m_other = other[other.model == model]
                rep = m_other.faithful.mean() if len(m_other) > 0 else 0.5
                reputation[(domain, task_id, model)] = rep
    return reputation


def agg_mean(vals): return np.mean(vals)


def agg_trim(vals, tf=0.2):
    v = np.sort(vals); n = len(v); k = max(0, int(np.floor(tf * n)))
    return np.mean(v[k:n - k]) if n - 2 * k > 0 else np.mean(v)


def agg_median(vals): return np.median(vals)


def agg_reputation(vals, weights):
    weights = np.array(weights)
    if weights.sum() == 0:
        return np.mean(vals)
    return float(np.average(vals, weights=weights))


def run_leave_one_trial_out_same_model(df, reputation, min_peers=2):
    results = []
    grouped = df.groupby("item_key")
    skipped_insufficient_peers = 0
    for item_key, group in grouped:
        domain = group.domain.iloc[0]
        task_id = group.task_id.iloc[0]
        trials = group.to_dict("records")
        for i, trial in enumerate(trials):
            same_model_peers = [t for j, t in enumerate(trials)
                                  if j != i and t["model"] == trial["model"]]
            if len(same_model_peers) < min_peers:
                skipped_insufficient_peers += 1
                continue
            peer_scores = [p["self_report"] for p in same_model_peers]
            peer_weights = [reputation.get((domain, task_id, p["model"]), 0.5) for p in same_model_peers]

            solo_pred = 1 if trial["self_report"] >= 0.5 else 0
            naive_pred = 1 if agg_mean(peer_scores) >= 0.5 else 0
            trim_pred = 1 if agg_trim(peer_scores, 0.2) >= 0.5 else 0
            median_pred = 1 if agg_median(peer_scores) >= 0.5 else 0
            rep_pred = 1 if agg_reputation(peer_scores, peer_weights) >= 0.5 else 0

            results.append({
                "domain": domain, "task_id": task_id, "model": trial["model"], "trial": trial["trial"],
                "true_reward": int(trial["reward"]),
                "solo_pred": solo_pred, "naive_pred": naive_pred, "trim_pred": trim_pred,
                "median_pred": median_pred, "rep_pred": rep_pred,
            })
    log(f"  (Skipped {skipped_insufficient_peers} trials with fewer than {min_peers} same-model peers)")
    return pd.DataFrame(results)


def compute_accuracy_metrics(results_df):
    metrics = {}
    for col in ["solo_pred", "naive_pred", "trim_pred", "median_pred", "rep_pred"]:
        acc = (results_df[col] == results_df.true_reward).mean()
        metrics[col] = round(float(acc), 4)
    return metrics


def main():
    log("Loading all real trajectories (GPT-4o + Sonnet, retail + airline)...")
    df = load_all_trials()
    log(f"Loaded {len(df)} total trials across {df.item_key.nunique()} unique tasks.")

    log("\nComputing SAME-MODEL-ONLY leave-one-task-out reputation scores...")
    reputation = compute_model_reputation_same_model(df)

    log("\n=== SAME-MODEL-ONLY PEER GROUPS: leave-one-trial-out peer verification ===")
    results_df = run_leave_one_trial_out_same_model(df, reputation, min_peers=2)
    results_df.to_csv("same_model_trial_level_results.csv", index=False)
    log(f"Evaluated {len(results_df)} trials with valid same-model peer groups.")

    metrics = compute_accuracy_metrics(results_df)
    log("\n=== OVERALL ACCURACY (same-model-only peer groups) ===")
    for k, v in metrics.items():
        log(f"  {k:15s}: {v:.1%}")

    log("\n=== PER-MODEL BREAKDOWN (same-model-only peer groups) ===")
    model_metrics = {}
    for model in results_df.model.unique():
        sub = results_df[results_df.model == model]
        mm = compute_accuracy_metrics(sub)
        model_metrics[model] = mm
        log(f"  {model} (n={len(sub)}): {mm}")

    log("\n=== PER-DOMAIN BREAKDOWN (same-model-only peer groups) ===")
    domain_metrics = {}
    for domain in results_df.domain.unique():
        sub = results_df[results_df.domain == domain]
        dm = compute_accuracy_metrics(sub)
        domain_metrics[domain] = dm
        log(f"  {domain} (n={len(sub)}): {dm}")

    with open("same_model_summary.json", "w") as f:
        json.dump({"overall": metrics, "by_model": model_metrics, "by_domain": domain_metrics,
                    "n_trials_evaluated": len(results_df)}, f, indent=2)
    log("\nSaved: same_model_trial_level_results.csv, same_model_summary.json")


if __name__ == "__main__":
    main()