#!/usr/bin/env python3
"""
Bayesian hierarchical logistic regression for agent-level public-private divergence.

Outcome: public_private_divergence (binary: vote ≠ pre-vote private top suspect)
Predictors: contradiction exposure, speaking position, pressure, role, confidence change
Random effects: game-level intercept (48 groups, ~12 agents each)
"""
from __future__ import annotations

import os
import warnings

import arviz as az
import bambi as bmb
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning, module="pytensor")

BASE = "./MSM"  # Update this path to your local working directory
TABLE_DIR = os.path.join(BASE, "tables")
FIG_DIR = os.path.join(BASE, "figures")
os.makedirs(TABLE_DIR, exist_ok=True)
os.makedirs(FIG_DIR, exist_ok=True)


def load_agent_data() -> pd.DataFrame:
    df = pd.read_csv(os.path.join(BASE, "processed_data", "agent_level.csv"))

    # Coding
    df["contradiction"] = df["has_contradiction"].astype(int)
    df["pressure_ec"] = (df["pressure_condition"] == "evidence_centric").astype(int)
    df["diverged"] = df["public_private_divergence"].astype(int)
    df["speak_pos_centered"] = df["speaking_position"] - df["speaking_position"].mean()
    df["conf_change"] = df["private_confidence_change"]

    # Drop dead agents (killed in night phase, no vote/beliefs)
    n_before = len(df)
    df = df.dropna(subset=["private_confidence_change", "final_vote_target"]).copy()
    print(f"Dropped {n_before - len(df)} dead agents (no vote/beliefs). N = {len(df)}")

    # PCE numeric for cross-level interaction
    pce_map = {"0%": 0.0, "25%": 0.25, "50%": 0.50, "75%": 0.75}
    df["pce_numeric"] = df["pce_level"].map(pce_map)

    # Simplify role to manageable categories
    # Keep all 6 roles since N=576 supports it
    df["role_factor"] = pd.Categorical(df["role"])

    return df


def fit_agent_model(df: pd.DataFrame) -> dict:
    """Fit hierarchical Bayesian logistic regression."""
    print("=" * 60)
    print("AGENT-LEVEL HIERARCHICAL LOGISTIC REGRESSION")
    print(f"N agents = {len(df)}, N games = {df['game_id'].nunique()}")
    print("=" * 60)

    # Model: divergence ~ contradiction + speaking position + pressure + role + confidence change
    # Random intercept by game
    formula = (
        "diverged ~ contradiction + speak_pos_centered + pressure_ec + "
        "C(role, Treatment('VILLAGER')) + conf_change + "
        "(1 | game_id)"
    )

    print(f"Formula: {formula}")

    model = bmb.Model(
        formula,
        data=df,
        family="bernoulli",
    )

    idata = model.fit(
        draws=2000,
        tune=2000,
        chains=4,
        target_accept=0.95,
        random_seed=42,
        progressbar=True,
    )

    summary = az.summary(idata, hdi_prob=0.95)
    print("\n--- Posterior Summary ---")
    print(summary.to_string())

    rhat_ok = (summary["r_hat"] <= 1.05).all()
    ess_ok = (summary["ess_bulk"] >= 400).all()
    print(f"\nR-hat all <= 1.05: {rhat_ok}")
    print(f"ESS_bulk all >= 400: {ess_ok}")

    return {
        "model": model,
        "idata": idata,
        "summary": summary,
        "rhat_ok": rhat_ok,
        "ess_ok": ess_ok,
    }


def make_agent_table(res: dict) -> pd.DataFrame:
    s = res["summary"]
    rows = []
    for param in s.index:
        # Skip individual game random effects
        if param.startswith("1|game_id["):
            continue
        rows.append({
            "parameter": param,
            "mean": round(s.loc[param, "mean"], 4),
            "sd": round(s.loc[param, "sd"], 4),
            "hdi_2.5%": round(s.loc[param, "hdi_2.5%"], 4),
            "hdi_97.5%": round(s.loc[param, "hdi_97.5%"], 4),
            "ess_bulk": int(s.loc[param, "ess_bulk"]),
            "ess_tail": int(s.loc[param, "ess_tail"]),
            "r_hat": round(s.loc[param, "r_hat"], 4),
        })
    return pd.DataFrame(rows)


def plot_agent_coefficients(res: dict, save_path: str):
    """Forest plot for agent-level model."""
    s = res["summary"]

    # Filter to fixed effects and group-level SD
    params = [p for p in s.index if not p.startswith("1|game_id[")]

    s_plot = s.loc[params]

    # Separate fixed effects from group-level SD
    fixed = [p for p in params if "sigma" not in p.lower() and "1|game_id" not in p]
    other = [p for p in params if p not in fixed]

    fig, ax = plt.subplots(figsize=(8, max(4, 0.5 * len(fixed))))

    y_pos = np.arange(len(fixed))

    nice_labels = {
        "Intercept": "Intercept",
        "contradiction": "Has Contradiction",
        "speak_pos_centered": "Speaking Position (centered)",
        "pressure_ec": "Pressure (EC vs HP)",
        "conf_change": "Confidence Change",
        "C(role, Treatment('VILLAGER'))[WEREWOLF]": "Role: Werewolf",
        "C(role, Treatment('VILLAGER'))[WITCH]": "Role: Witch",
        "C(role, Treatment('VILLAGER'))[SEER]": "Role: Seer",
        "C(role, Treatment('VILLAGER'))[GUARD]": "Role: Guard",
        "C(role, Treatment('VILLAGER'))[HUNTER]": "Role: Hunter",
    }

    means = [s_plot.loc[p, "mean"] for p in fixed]
    lo = [s_plot.loc[p, "hdi_2.5%"] for p in fixed]
    hi = [s_plot.loc[p, "hdi_97.5%"] for p in fixed]
    labels = [nice_labels.get(p, p) for p in fixed]

    ax.hlines(y_pos, lo, hi, color="#43A047", linewidth=2.5, alpha=0.8)
    ax.scatter(means, y_pos, color="#2E7D32", s=60, zorder=5)
    ax.axvline(0, color="gray", linestyle="--", linewidth=0.8, alpha=0.6)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=10)
    ax.set_xlabel("Posterior Mean (log-odds scale)", fontsize=11)
    ax.set_title("Agent-Level Divergence Model: Posterior Coefficients (95% HDI)", fontsize=12, fontweight="bold")
    ax.grid(axis="x", alpha=0.3)

    # Add group-level SD annotation
    for p in other:
        val = s_plot.loc[p, "mean"]
        ax.annotate(
            f"Game RE σ = {val:.3f}",
            xy=(0.02, 0.02), xycoords="axes fraction",
            fontsize=9, color="gray", style="italic",
        )

    fig.tight_layout()
    for ext in ["png", "pdf"]:
        fig.savefig(f"{save_path}.{ext}", dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved: {save_path}.png/.pdf")


def plot_predicted_probabilities(df: pd.DataFrame, res: dict, save_path: str):
    """Plot predicted probability of divergence by contradiction status and pressure."""
    idata = res["idata"]
    posterior = idata.posterior

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    def get(name):
        return posterior[name].values.flatten()

    intercept = get("Intercept")
    b_contra = get("contradiction")
    b_press = get("pressure_ec")
    b_speak = get("speak_pos_centered")
    b_conf = get("conf_change")

    # Panel 1: Predicted probability by contradiction × pressure
    ax = axes[0]
    conditions = [
        (0, 0, "No Contradiction + HP", "#90CAF9"),
        (0, 1, "No Contradiction + EC", "#1565C0"),
        (1, 0, "Contradiction + HP", "#FFAB91"),
        (1, 1, "Contradiction + EC", "#D32F2F"),
    ]
    positions = np.arange(4)
    means_pred = []
    lo_pred = []
    hi_pred = []
    labels = []
    for contra, press, label, color in conditions:
        eta = intercept + b_contra * contra + b_press * press
        # At mean speaking position and mean confidence change
        mu = 1.0 / (1.0 + np.exp(-eta))
        means_pred.append(np.mean(mu))
        lo_pred.append(np.percentile(mu, 2.5))
        hi_pred.append(np.percentile(mu, 97.5))
        labels.append(label)

    colors = [c[3] for c in conditions]
    ax.bar(positions, means_pred, color=colors, alpha=0.8, width=0.6)
    ax.errorbar(positions, means_pred,
                yerr=[np.array(means_pred) - np.array(lo_pred),
                      np.array(hi_pred) - np.array(means_pred)],
                fmt="none", color="black", capsize=4)
    ax.set_xticks(positions)
    ax.set_xticklabels(labels, fontsize=8, rotation=15, ha="right")
    ax.set_ylabel("P(Vote ≠ Private Belief)", fontsize=10)
    ax.set_title("Predicted Divergence by Condition", fontsize=11, fontweight="bold")
    ax.set_ylim(0, 0.5)
    ax.grid(axis="y", alpha=0.3)

    # Panel 2: Predicted probability by speaking position
    ax = axes[1]
    speak_range = np.linspace(-5.5, 5.5, 50)  # centered positions
    for contra, label, color in [(0, "No Contradiction", "#1E88E5"), (1, "Has Contradiction", "#E53935")]:
        means_sp = []
        lo_sp = []
        hi_sp = []
        for sp in speak_range:
            eta = intercept + b_contra * contra + b_speak * sp
            mu = 1.0 / (1.0 + np.exp(-eta))
            means_sp.append(np.mean(mu))
            lo_sp.append(np.percentile(mu, 2.5))
            hi_sp.append(np.percentile(mu, 97.5))
        ax.plot(speak_range + 6.5, means_sp, label=label, color=color, linewidth=2)  # un-center
        ax.fill_between(speak_range + 6.5, lo_sp, hi_sp, alpha=0.15, color=color)

    ax.set_xlabel("Speaking Position", fontsize=10)
    ax.set_ylabel("P(Vote ≠ Private Belief)", fontsize=10)
    ax.set_title("Divergence by Speaking Position", fontsize=11, fontweight="bold")
    ax.legend(fontsize=9)
    ax.set_ylim(0, 0.5)
    ax.grid(alpha=0.3)

    fig.suptitle("Agent-Level Model: Predicted Probabilities (95% CI)", fontsize=13, fontweight="bold", y=1.02)
    fig.tight_layout()
    for ext in ["png", "pdf"]:
        fig.savefig(f"{save_path}.{ext}", dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved: {save_path}.png/.pdf")


def main():
    print("Loading agent data...")
    df = load_agent_data()

    # Fit model
    res = fit_agent_model(df)

    # Tables
    agent_table = make_agent_table(res)
    agent_table.to_csv(os.path.join(TABLE_DIR, "table_bayesian_agent_level_model.csv"), index=False)

    with open(os.path.join(TABLE_DIR, "table_bayesian_agent_level_model.md"), "w") as f:
        f.write("# Bayesian Hierarchical Logistic Regression: Agent-Level Divergence\n\n")
        f.write(agent_table.to_markdown(index=False))
        f.write("\n\nNote: Coefficients on log-odds scale. Random intercept by game (48 groups).\n")
    print(f"Saved agent tables to {TABLE_DIR}")

    # Figures
    plot_agent_coefficients(res, os.path.join(FIG_DIR, "fig_bayes_agent_divergence"))
    plot_predicted_probabilities(df, res, os.path.join(FIG_DIR, "fig_bayes_agent_predicted_probabilities"))

    # Save InferenceData
    az.to_netcdf(res["idata"], os.path.join(BASE, "scripts", "idata_agent_divergence.nc"))

    print("\n" + "=" * 60)
    print("AGENT-LEVEL MODEL COMPLETE")
    print("=" * 60)

    # Key results
    s = res["summary"]
    fixed = [p for p in s.index if not p.startswith("1|game_id[") and "sigma" not in p.lower()]
    print("\n--- Key Results ---")
    for p in fixed:
        mean_val = s.loc[p, "mean"]
        lo = s.loc[p, "hdi_2.5%"]
        hi = s.loc[p, "hdi_97.5%"]
        sig = "*" if (lo > 0 and hi > 0) or (lo < 0 and hi < 0) else ""
        # Convert to OR
        or_val = np.exp(mean_val)
        or_lo = np.exp(lo)
        or_hi = np.exp(hi)
        print(f"  {p:45s}: β={mean_val:+.3f} [{lo:+.3f}, {hi:+.3f}]  OR={or_val:.2f} [{or_lo:.2f}, {or_hi:.2f}] {sig}")


if __name__ == "__main__":
    main()