"""Mechanism checks (descriptive, pre-registered as non-identified mediation):
(1) coordination rises with the LLM share (H2): llm_pairwise_corr, llm_synchrony and the classical analogue flow_corr_llm vs share;
(2) within each share, runs with more correlated LLM order flow have worse price discovery / more instability (seed-FE OLS, descriptive);
(3) diversity arms: coordination and outcomes for single vs diverse prompts vs two models at 25 % (paired);
(4) theory comparative statics: sign of the dose effect on informativeness proxies vs Table tab:predictions.
Writes results/tables/tab_mechanisms.tex and mechanisms.csv."""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from market_agents.analysis.estimators import paired_effects  # noqa: E402

p = ROOT / "results/runs/confirmatory_openai/outcomes.parquet"
if not p.exists():
    raise SystemExit("confirmatory design not run yet")
df = pd.read_parquet(p)
df["cell"] = df.name.str.split("/").str[-1]
OUT = ROOT / "results/tables"
rows = []

# (1) coordination vs share
d = df[(df.tr_arm == "dose") & (df.llm_share > 0)].copy()
for oc in ("llm_pairwise_corr", "llm_synchrony", "flow_corr_llm", "llm_hold_rate", "llm_valid_rate"):
    if oc in d:
        g = d.groupby("llm_share")[oc].agg(["mean", "std", "size"])
        for s, r in g.iterrows():
            rows.append({"block": "coordination_by_share", "share": s, "stat": oc, "mean": r["mean"], "se": r["std"] / np.sqrt(r["size"]), "n": r["size"]})
        # trend test: OLS of stat on share with seed FE
        m = smf.ols(f"{oc} ~ llm_share + C(seed)", data=d.dropna(subset=[oc])).fit(cov_type="cluster", cov_kwds={"groups": d.dropna(subset=[oc])["seed"]})
        rows.append({"block": "trend", "stat": oc, "beta_share": m.params["llm_share"], "se": m.bse["llm_share"], "p": m.pvalues["llm_share"], "n": int(m.nobs)})

# classical analogue: flow correlation of classical groups in the same runs
for grp in ("fundamentalist", "trend", "noise", "market_maker"):
    c = f"flow_corr_{grp}"
    if c in df:
        rows.append({"block": "classical_flow_corr", "stat": c, "mean": df[df.cell == "dose_000"][c].mean(), "n": int(df[df.cell == "dose_000"][c].notna().sum())})

# (2) within-share association between coordination and outcomes (descriptive)
for oc in ("mad_pct", "msd_log", "time_to_absorb", "n_bubbles", "max_dd", "spread_bp", "lambda_bp_100"):
    for med in ("llm_pairwise_corr", "llm_synchrony"):
        dd = d.dropna(subset=[oc, med])
        if len(dd) < 8 or dd[med].std() == 0:
            continue
        z = dd.copy()
        z["med_z"] = (z[med] - z[med].mean()) / z[med].std()
        m = smf.ols(f"{oc} ~ med_z + C(llm_share) + C(seed)", data=z).fit(cov_type="cluster", cov_kwds={"groups": z["seed"]})
        rows.append({"block": "within_share_association", "stat": oc, "mediator": med, "beta_per_sd": m.params["med_z"], "se": m.bse["med_z"], "p": m.pvalues["med_z"], "n": int(m.nobs)})

# (3) diversity arms at 25 %
for cell in ("div_prompts_025", "div_models_025", "memory_025", "channel_025", "adaptive_025", "fair_latency_025"):
    for oc in ("llm_pairwise_corr", "llm_synchrony", "mad_pct", "time_to_absorb", "n_bubbles", "max_dd", "spread_bp"):
        if cell in set(df.cell) and oc in df:
            e = paired_effects(df, oc, "cell", cell, "dose_025")
            rows.append({"block": "arm_vs_dose_025", "cell": cell, "stat": oc, "control_mean": e.get("control_mean"), "diff": e.get("diff"), "ci_lo": e.get("ci_lo"), "ci_hi": e.get("ci_hi"), "p": e.get("p"), "cohen_d": e.get("cohen_d"), "n": e.get("n")})

t = pd.DataFrame(rows)
t.to_csv(OUT / "mechanisms.csv", index=False)

# LaTeX: coordination by share + trend
lines = ["\\begin{tabular}{lrrrrrr}", "\\toprule", "Statistic & 5\\% & 10\\% & 25\\% & 50\\% & 100\\% & Trend per unit share (SE) \\\\", "\\midrule"]
for oc, lab in (("llm_pairwise_corr", "Pairwise order-flow correlation (LLM)"), ("llm_synchrony", "Synchrony (LLM)"), ("flow_corr_llm", "Flow correlation from trades (LLM)"), ("llm_hold_rate", "Hold rate"), ("llm_valid_rate", "Valid-decision rate")):
    c = t[(t.block == "coordination_by_share") & (t.stat == oc)].set_index("share")
    tr = t[(t.block == "trend") & (t.stat == oc)]
    if len(c) == 0:
        continue
    vals = " & ".join(f"{c.loc[s, 'mean']:.3f}" if s in c.index else "--" for s in (0.05, 0.10, 0.25, 0.50, 1.00))
    trs = f"{tr.iloc[0].beta_share:+.3f} ({tr.iloc[0].se:.3f})" + ("$^{***}$" if tr.iloc[0].p < 0.01 else "$^{**}$" if tr.iloc[0].p < 0.05 else "$^{*}$" if tr.iloc[0].p < 0.1 else "") if len(tr) else "--"
    lines.append(f"{lab} & {vals} & {trs} \\\\")
cl = t[t.block == "classical_flow_corr"]
if len(cl):
    lines.append("\\midrule")
    for r in cl.itertuples():
        lines.append(f"Flow correlation, {r.stat.replace('flow_corr_', '').replace('_', ' ')} (0\\% cell) & \\multicolumn{{5}}{{r}}{{{r.mean:.3f}}} & \\\\")
lines += ["\\bottomrule", "\\end{tabular}"]
(OUT / "tab_coordination.tex").write_text("\n".join(lines))

w = t[t.block == "within_share_association"]
if len(w):
    lines = ["\\begin{tabular}{llrrr}", "\\toprule", "Outcome & Mediator (per 1 SD) & Coefficient & SE & $p$ \\\\", "\\midrule"]
    for r in w.itertuples():
        lines.append(f"{r.stat.replace('_', ' ')} & {r.mediator.replace('llm_', '').replace('_', ' ')} & {r.beta_per_sd:+.3f} & {r.se:.3f} & {r.p:.3f} \\\\")
    lines += ["\\bottomrule", "\\end{tabular}"]
    (OUT / "tab_mechanism_association.tex").write_text("\n".join(lines))
print(t.to_string(max_rows=80))