"""
REAL-DATA EVALUATION: Cross-Agent Behavioral Verification of Self-Reported
Actions, using real tau-bench trajectories (GPT-4o and Claude-3.5-Sonnet,
retail and airline domains).

DESIGN (leave-one-trial-out, no data leakage):
- Each (domain, task_id) is treated as an "item"; every trial (independent
  execution attempt, either model, any trial index) on that item is treated
  as one "peer agent's" observation.
- Each trial has: (a) its own ground-truth reward (0/1, computed by
  tau-bench via database-state-hash comparison -- NOT self-reported), and
  (b) its own self-reported-success score (0/0.5/1, extracted via the
  transparent rule-based classifier in self_report_classifier.py).
- For each trial i (held out), the PEER GROUP is all other trials j != i
  on the SAME task. Peer self-report scores are aggregated via four rules
  (naive mean, trimmed mean, coordinate-wise median, reputation-weighted)
  to produce a verification signal, which is compared against trial i's
  own TRUE reward (never using trial i's own self-report or reward in
  computing that signal, avoiding leakage).
- Reputation weights are calibrated PER MODEL (gpt-4o vs sonnet) using
  that model's overall self-report/reward agreement rate computed on
  OTHER tasks only (leave-one-task-out).
- A "solo" baseline uses only the trial's own self-report (no peer
  verification) to predict its own true reward.

USAGE: place gpt-4o-retail.json, sonnet-35-new-retail.json,
gpt-4o-airline.json, sonnet-35-new-airline.json (from
github.com/sierra-research/tau-bench, historical_trajectories/) in the
same folder as this script, then run: python3 real_data_pipeline.py
"""
import json
import numpy as np
import pandas as pd

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():
    from self_report_classifier import get_final_user_facing_message, self_reported_success_score
    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(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(df, reputation):
    results = []
    grouped = df.groupby("item_key")
    for item_key, group in grouped:
        domain = group.domain.iloc[0]
        task_id = group.task_id.iloc[0]
        trials = group.to_dict("records")
        if len(trials) < 3:
            continue
        for i, trial in enumerate(trials):
            peers = [t for j, t in enumerate(trials) if j != i]
            peer_scores = [p["self_report"] for p in peers]
            peer_weights = [reputation.get((domain, task_id, p["model"]), 0.5) for p in 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,
            })
    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 wilson_ci(k, n, z=1.96):
    """Wilson score interval for a binomial proportion. k successes out of n."""
    if n == 0:
        return (float("nan"), float("nan"))
    p = k / n
    denom = 1 + z**2 / n
    center = (p + z**2 / (2 * n)) / denom
    half = (z * np.sqrt((p * (1 - p) / n) + (z**2 / (4 * n**2)))) / denom
    return (max(0.0, center - half), min(1.0, center + half))


def mcnemar_exact_test(results_df, col_a="solo_pred", col_b="naive_pred"):
    """
    Exact (binomial) McNemar's test comparing two paired prediction columns
    against true_reward, using only the DISCORDANT trials (where col_a and
    col_b disagree on correctness). This directly tests whether solo
    self-report and a given aggregation rule differ in accuracy on the
    same set of trials (paired, not independent-sample, comparison).
    """
    from scipy.stats import binomtest
    a_correct = (results_df[col_a] == results_df.true_reward)
    b_correct = (results_df[col_b] == results_df.true_reward)
    # n10: A correct, B wrong ; n01: A wrong, B correct
    n10 = int(((a_correct) & (~b_correct)).sum())
    n01 = int(((~a_correct) & (b_correct)).sum())
    n_discordant = n10 + n01
    if n_discordant == 0:
        return {"n10": n10, "n01": n01, "n_discordant": 0, "p_value": 1.0}
    result = binomtest(min(n10, n01), n_discordant, 0.5, alternative="two-sided")
    return {"n10": n10, "n01": n01, "n_discordant": n_discordant, "p_value": float(result.pvalue)}


def task_clustered_bootstrap_ci(results_df, col="solo_pred", n_boot=10000, seed=42):
    """
    Bootstrap 95% CI for accuracy that resamples at the TASK level (not the
    trial level), to account for within-task correlation (per-task reward
    variance is substantial, meaning trials on the same task are not
    independent observations).

    Implementation note: this resamples whole tasks with replacement, but
    avoids the O(n_boot) pandas concat/groupby cost of a naive implementation
    by precomputing each task's (n_correct, n_total) once and then drawing
    bootstrap accuracy as a vectorized weighted average over resampled task
    indices. This is mathematically identical to concatenating resampled
    per-task trial rows and computing accuracy over the pooled resample, but
    is orders of magnitude faster for realistic n_boot.
    """
    rng = np.random.default_rng(seed)
    correct = (results_df[col] == results_df.true_reward).astype(int).to_numpy()
    # IMPORTANT: cluster on (domain, task_id), not task_id alone. Retail and
    # airline task_id values overlap numerically (e.g., both have a task_id=0),
    # so factorizing on task_id alone would incorrectly merge unrelated tasks
    # from different domains into the same bootstrap cluster.
    compound_key = results_df["domain"].astype(str) + "_" + results_df["task_id"].astype(str)
    task_codes, task_uniques = pd.factorize(compound_key)
    n_tasks = len(task_uniques)

    task_correct = np.bincount(task_codes, weights=correct, minlength=n_tasks)
    task_total = np.bincount(task_codes, minlength=n_tasks)

    sampled_task_idx = rng.integers(0, n_tasks, size=(n_boot, n_tasks))
    boot_correct = task_correct[sampled_task_idx].sum(axis=1)
    boot_total = task_total[sampled_task_idx].sum(axis=1)
    boot_accs = boot_correct / boot_total

    lo, hi = np.percentile(boot_accs, [2.5, 97.5])
    return {"point_estimate": float(correct.mean()),
            "boot_mean": float(boot_accs.mean()), "ci_lo": float(lo), "ci_hi": float(hi),
            "n_tasks": int(n_tasks), "n_boot": n_boot}


def run_full_statistical_analysis(results_df):
    """
    Runs the complete statistical validation suite requested for the
    submission-ready manuscript:
      1. Wilson score 95% CIs for each method's overall accuracy (trial-level).
      2. Task-clustered bootstrap 95% CIs for each method's accuracy
         (accounts for within-task correlation).
      3. Exact McNemar's tests comparing solo self-report against each of
         the four aggregation rules (paired comparison on the same trials).
    """
    methods = ["solo_pred", "naive_pred", "trim_pred", "median_pred", "rep_pred"]
    n = len(results_df)

    wilson_results = {}
    for m in methods:
        k = int((results_df[m] == results_df.true_reward).sum())
        lo, hi = wilson_ci(k, n)
        wilson_results[m] = {"accuracy": round(k / n, 4), "ci_lo": round(lo, 4), "ci_hi": round(hi, 4), "n": n}

    bootstrap_results = {}
    for m in methods:
        bootstrap_results[m] = task_clustered_bootstrap_ci(results_df, col=m)

    mcnemar_results = {}
    for m in methods:
        if m == "solo_pred":
            continue
        mcnemar_results[f"solo_vs_{m}"] = mcnemar_exact_test(results_df, "solo_pred", m)

    return {"wilson_ci": wilson_results, "task_clustered_bootstrap_ci": bootstrap_results,
            "mcnemar_solo_vs_aggregation": mcnemar_results}


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(f"Reward distribution: {df.reward.value_counts().to_dict()}")
    log(f"Self-report score distribution: {df.self_report.value_counts().to_dict()}")

    log("\n=== NATURAL SELF-REPORT FAITHFULNESS STATISTICS ===")
    for domain in df.domain.unique():
        for model in df.model.unique():
            sub = df[(df.domain == domain) & (df.model == model)]
            false_success = ((sub.self_report == 1.0) & (sub.reward == 0.0)).sum()
            false_failure = ((sub.self_report == 0.0) & (sub.reward == 1.0)).sum()
            n = len(sub)
            log(f"  {domain}/{model}: n={n}  false_success_rate={false_success/n:.1%}  "
                f"false_failure_rate={false_failure/n:.1%}")

    log("\nComputing per-model, leave-one-task-out reputation scores...")
    reputation = compute_model_reputation(df)

    log("\nRunning leave-one-trial-out peer verification...")
    results_df = run_leave_one_trial_out(df, reputation)
    results_df.to_csv("real_trial_level_results.csv", index=False)
    log(f"Evaluated {len(results_df)} trials with valid peer groups.")

    metrics = compute_accuracy_metrics(results_df)
    log("\n=== OVERALL ACCURACY (predicting true reward from peer self-reports) ===")
    for k, v in metrics.items():
        log(f"  {k:15s}: {v:.1%}")

    log("\n=== PER-DOMAIN ACCURACY ===")
    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}: {dm}")

    log("\n=== FULL STATISTICAL ANALYSIS (Wilson CIs, task-clustered bootstrap, McNemar's) ===")
    stats_results = run_full_statistical_analysis(results_df)
    log(json.dumps(stats_results, indent=2))

    with open("real_data_summary.json", "w") as f:
        json.dump({"overall": metrics, "by_domain": domain_metrics,
                    "n_trials_evaluated": len(results_df),
                    "n_unique_tasks": df.item_key.nunique(),
                    "statistical_analysis": stats_results}, f, indent=2)
    log("\nSaved: real_trial_level_results.csv, real_data_summary.json (now includes "
        "Wilson CIs, task-clustered bootstrap CIs, and McNemar's exact test results)")


if __name__ == "__main__":
    main()