"""LaTeX table fragments from saved outcomes (results/tables/*.tex). Every number comes from
results/runs/<design>/outcomes.parquet or the run manifests; nothing is typed by hand."""
from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

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

RUNS, OUT = ROOT / "results/runs", ROOT / "results/tables"
OUT.mkdir(parents=True, exist_ok=True)
PRIMARY = ["msd_log", "mad_pct", "time_to_absorb", "response_ratio", "spread_bp", "eff_spread_bp", "depth5", "lambda_bp_100",
           "n_bubbles", "bubble_amp", "max_dd", "n_crashes", "acf_absr_1_10", "kurtosis"]
NICE = {"msd_log": "Sq. log pricing error ($\\times 10^{4}$)", "mad_pct": "Mean $|P/V-1|$ (\\%)", "time_to_absorb": "Time to absorb (ticks)",
        "response_ratio": "Response ratio (20 ticks)", "spread_bp": "Quoted spread (bp)", "eff_spread_bp": "Effective spread (bp)",
        "depth5": "Depth, 5 levels (lots)", "lambda_bp_100": "Price impact (bp/100 lots)", "n_bubbles": "Bubble episodes",
        "bubble_amp": "Bubble amplitude", "max_dd": "Max drawdown", "n_crashes": "Crashes", "acf_absr_1_10": "Vol. clustering $\\sum_{1}^{10}\\mathrm{ACF}|r|$",
        "kurtosis": "Excess kurtosis", "vr5": "Variance ratio (5)", "acf_r_1": "ACF$_1$ of returns", "vol_bp": "Tick volatility (bp)",
        "llm_valid_rate": "LLM valid-decision rate", "llm_synchrony": "LLM synchrony", "llm_pairwise_corr": "LLM order-flow corr.",
        "llm_usd": "API cost per run (USD)", "llm_share_volume": "Realized LLM volume share", "llm_share_orders": "Realized LLM order share",
        "gini_final": "Wealth Gini (final)", "fill_ratio": "Fills per order", "n_bankrupt": "Bankruptcies", "n_halts": "Halts", "excess_comovement": "Excess comovement",
        "spillover_dy": "DY spillover (\\%)", "ret_corr": "Return correlation", "fund_corr": "Fundamental correlation", "pnl_market_maker": "MM P\\&L (\\%)"}
SCALE = {"msd_log": 1e4}


def fmt(x, nd=2):
    if x is None or (isinstance(x, float) and not np.isfinite(x)):
        return "--"
    return f"{x:.{nd}f}"


def star(p):
    return "" if not np.isfinite(p) else ("$^{***}$" if p < 0.01 else "$^{**}$" if p < 0.05 else "$^{*}$" if p < 0.10 else "")


def load(design):
    p = RUNS / design / "outcomes.parquet"
    if not p.exists():
        print(f"  [skip] {design}")
        return None
    df = pd.read_parquet(p)
    df["cell"] = df["name"].str.split("/").str[-1]
    return df


def write(name: str, body: str):
    (OUT / name).write_text(body)
    print("  wrote", name)


def tab_stylized_facts():
    df = load("classical_confirmatory")
    if df is None:
        return
    d = df[df.cell == "control"]
    rows = [("acf_r_1", 3), ("vr5", 2), ("kurtosis", 1), ("acf_absr_1_10", 2), ("vol_bp", 1), ("spread_bp", 1), ("eff_spread_bp", 1), ("depth5", 0),
            ("lambda_bp_100", 1), ("mad_pct", 2), ("n_bubbles", 2), ("max_dd", 3), ("pnl_market_maker", 1), ("gini_final", 3)]
    lines = ["\\begin{tabular}{lrrrr}", "\\toprule", "Statistic & Mean & SD & Min & Max \\\\", "\\midrule"]
    for k, nd in rows:
        if k in d:
            v = d[k].astype(float)
            lines.append(f"{NICE.get(k, k)} & {fmt(v.mean(), nd)} & {fmt(v.std(ddof=1), nd)} & {fmt(v.min(), nd)} & {fmt(v.max(), nd)} \\\\")
    lines += ["\\midrule", f"Runs (seeds) & {len(d)} & & & \\\\", "\\bottomrule", "\\end{tabular}"]
    write("tab_stylized_facts.tex", "\n".join(lines))
    # scale benchmark
    p = ROOT / "results/calibration/scale_benchmark.json"
    if p.exists():
        sb = pd.DataFrame(json.load(open(p)))
        g = sb.groupby("n_traders").agg(wall=("wall_s", "mean"), trades=("n_trades", "mean"), orders=("n_orders", "mean"), mad=("mad_pct", "mean"),
                                       acf=("acf_r_1", "mean"), kurt=("kurtosis", "mean"), spread=("spread_bp", "mean"), n=("seed", "size"))
        lines = ["\\begin{tabular}{rrrrrrrr}", "\\toprule", "Traders & Seeds & Wall-clock (s) & Orders & Trades & Mean $|P/V-1|$ (\\%) & ACF$_1$ & Spread (bp) \\\\", "\\midrule"]
        for n, r in g.iterrows():
            lines.append(f"{n:,} & {int(r.n)} & {fmt(r.wall, 1)} & {int(r.orders):,} & {int(r.trades):,} & {fmt(r.mad, 2)} & {fmt(r.acf, 2)} & {fmt(r.spread, 1)} \\\\")
        lines += ["\\bottomrule", "\\end{tabular}"]
        write("tab_scale.tex", "\n".join(lines))


def tab_cells(design: str, cells: list[str], labels: dict, outcomes: list[str], name: str, control: str):
    df = load(design)
    if df is None:
        return
    cells = [c for c in cells if c in set(df.cell)]
    outcomes = [o for o in outcomes if o in df]
    lines = ["\\begin{tabular}{l" + "r" * len(outcomes) + "}", "\\toprule",
             "Cell & " + " & ".join(NICE.get(o, o) for o in outcomes) + " \\\\", "\\midrule"]
    ctrl = df[df.cell == control]
    for c in cells:
        d = df[df.cell == c]
        vals = []
        for o in outcomes:
            sc = SCALE.get(o, 1.0)
            m = d[o].astype(float).mean() * sc
            if c == control:
                vals.append(fmt(m))
            else:
                e = paired_effects(df, o, "cell", c, control)
                vals.append(f"{fmt(m)}{star(e.get('p', np.nan))}")
        lines.append(f"{labels.get(c, c)} & " + " & ".join(vals) + " \\\\")
    lines += ["\\midrule", f"Seeds per cell & {len(ctrl)}" + " &" * (len(outcomes) - 1) + " \\\\", "\\bottomrule", "\\end{tabular}"]
    write(name, "\n".join(lines))


def tab_effects(design: str, cells: list[str], labels: dict, outcomes: list[str], name: str, control: str, holm_family=True):
    """Paired effects vs control with 95 % CI and Holm-adjusted p within the primary family."""
    df = load(design)
    if df is None:
        return
    cells = [c for c in cells if c in set(df.cell) and c != control]
    outcomes = [o for o in outcomes if o in df]
    rows = []
    for c in cells:
        for o in outcomes:
            e = paired_effects(df, o, "cell", c, control)
            e["cell"] = c
            rows.append(e)
    t = pd.DataFrame(rows)
    if "p" not in t:
        return
    t["p_holm"] = np.nan
    for c in cells:
        m = t.cell == c
        t.loc[m, "p_holm"] = holm(t.loc[m, "p"].to_numpy())
    t.to_csv(OUT / name.replace(".tex", ".csv"), index=False)
    lines = ["\\begin{tabular}{llrrrrr}", "\\toprule", "Cell & Outcome & Control & Effect & 95\\% CI & $d$ & $p$ (Holm) \\\\", "\\midrule"]
    for c in cells:
        for i, o in enumerate(outcomes):
            r = t[(t.cell == c) & (t.outcome == o)].iloc[0]
            sc = SCALE.get(o, 1.0)
            lab = labels.get(c, c) if i == 0 else ""
            lines.append(f"{lab} & {NICE.get(o, o)} & {fmt(r.control_mean * sc)} & {fmt(r['diff'] * sc)}{star(r.p_holm)} & "
                         f"[{fmt(r.ci_lo * sc)}, {fmt(r.ci_hi * sc)}] & {fmt(r.cohen_d)} & {fmt(r.p_holm, 3)} \\\\")
        lines.append("\\addlinespace")
    lines += ["\\bottomrule", "\\end{tabular}"]
    write(name, "\n".join(lines))


def tab_dose(design: str, name: str):
    df = load(design)
    if df is None or "tr_arm" not in df:
        return
    d = df[df.tr_arm == "dose"].copy()
    d["share"] = d.tr_share.astype(float)
    outs = [o for o in ("msd_log", "mad_pct", "time_to_absorb", "spread_bp", "lambda_bp_100", "n_bubbles", "max_dd", "acf_absr_1_10", "llm_pairwise_corr", "llm_valid_rate") if o in d]
    levels = sorted(d.share.unique())
    lines = ["\\begin{tabular}{l" + "r" * len(levels) + "}", "\\toprule", "Outcome & " + " & ".join(f"{int(l * 100)}\\%" for l in levels) + " \\\\", "\\midrule"]
    for o in outs:
        sc = SCALE.get(o, 1.0)
        dr = None
        if o not in ("llm_pairwise_corr", "llm_valid_rate") and 0.0 in levels:
            dr = dose_response(d, o, "share", 0.0).set_index("level")
            if "p_holm" not in dr:
                dr = None
        cells = []
        for l in levels:
            m = d[d.share == l][o].astype(float).mean() * sc
            s = "" if (dr is None or l == 0.0 or l not in dr.index) else star(dr.loc[l, "p_holm"])
            cells.append(f"{fmt(m)}{s}")
        lines.append(f"{NICE.get(o, o)} & " + " & ".join(cells) + " \\\\")
    rs = d.groupby("share")[["llm_share_traders", "llm_share_orders", "llm_share_volume", "llm_usd"]].mean()
    lines.append("\\midrule")
    for k, lab in (("llm_share_traders", "Realized share of traders"), ("llm_share_orders", "Realized share of orders"), ("llm_share_volume", "Realized share of volume"), ("llm_usd", "API cost per run (USD)")):
        lines.append(f"{lab} & " + " & ".join(fmt(rs.loc[l, k], 3) for l in levels) + " \\\\")
    lines += ["\\midrule", f"Seeds per cell & {int(d.groupby('share').size().min())}" + " &" * (len(levels) - 1) + " \\\\", "\\bottomrule", "\\end{tabular}"]
    write(name, "\n".join(lines))


def tab_llm_reliability(design: str, name: str):
    df = load(design)
    if df is None:
        return
    rows = []
    for c, d in df[df.llm_share > 0].groupby("cell"):
        rec = []
        for rd in d.run_dir:
            p = Path(rd) / "llm_decisions.parquet"
            if p.exists():
                rec.append(pd.read_parquet(p))
        if not rec:
            continue
        L = pd.concat(rec)
        rows.append({"cell": c, "n": len(L), "valid": L.valid.mean(), "hold": (L.action == "hold").mean(), "place": (L.action == "place").mean(),
                     "market": (L.order_type == "market").mean(), "json_err": L.reason.str.startswith("json").mean(), "schema_err": L.reason.str.startswith("schema").mean(),
                     "feas_err": L.reason.str.startswith("place").mean(), "api_err": (L.error != "").mean(), "lat": L.latency_s.mean(),
                     "in_tok": L.input_tokens.mean(), "out_tok": L.output_tokens.mean(), "reason_tok": L.reasoning_tokens.mean(), "usd": L.usd.sum() / len(d)})
    if not rows:
        return
    t = pd.DataFrame(rows)
    t.to_csv(OUT / name.replace(".tex", ".csv"), index=False)
    lines = ["\\begin{tabular}{lrrrrrrrrrr}", "\\toprule", "Cell & Decisions & Valid & Hold & Market & JSON err. & Schema err. & Infeasible & API err. & Latency (s) & USD/run \\\\", "\\midrule"]
    for r in t.itertuples():
        lines.append(f"{r.cell.replace('_', '\\_')} & {r.n:,} & {fmt(r.valid, 3)} & {fmt(r.hold, 2)} & {fmt(r.market, 2)} & {fmt(r.json_err, 3)} & {fmt(r.schema_err, 3)} & {fmt(r.feas_err, 3)} & {fmt(r.api_err, 3)} & {fmt(r.lat, 2)} & {fmt(r.usd, 2)} \\\\")
    lines += ["\\bottomrule", "\\end{tabular}"]
    write(name, "\n".join(lines))


CLASSICAL_CELLS = ["control", "placebo_shock", "trend_heavy", "leverage_heavy", "leverage_heavy_funding_shock", "liquidity_shock", "info_low", "info_high",
                   "misleading_news", "breaker_5", "poslimit_300", "auction_5", "mock_llm_25_indep", "mock_llm_50_indep"]
CLASSICAL_LAB = {"control": "Control", "placebo_shock": "Placebo shock", "trend_heavy": "Trend followers 16\\%", "leverage_heavy": "Leveraged 19\\%",
                 "leverage_heavy_funding_shock": "Leveraged 19\\% + funding shock", "liquidity_shock": "Market makers withdraw", "info_low": "Signal precision 0.5",
                 "info_high": "Signal precision 8", "misleading_news": "Misleading news 50\\%", "breaker_5": "Circuit breaker 5\\%", "poslimit_300": "Position limit 300",
                 "auction_5": "Call auction (5 ticks)", "mock_llm_25_indep": "Mock LLM 25\\%", "mock_llm_50_indep": "Mock LLM 50\\%"}
ARM_CELLS = ["dose_000", "dose_025", "div_prompts_025", "div_models_025", "memory_025", "channel_025", "adaptive_025", "leverage_000", "leverage_025", "breaker_025", "poslimit_025", "fair_mini_005", "fair_latency_025"]
ARM_LAB = {"dose_000": "0\\% (control)", "dose_025": "25\\% single model/prompt", "div_prompts_025": "25\\% diverse prompts", "div_models_025": "25\\% two models",
           "memory_025": "25\\% memory (5)", "channel_025": "25\\% communication channel", "adaptive_025": "25\\% adaptive feedback", "leverage_000": "0\\% + leveraged 19\\% (classical)", "leverage_025": "25\\% + leveraged 19\\%",
           "breaker_025": "25\\% + circuit breaker", "poslimit_025": "25\\% + position limit", "fair_mini_005": "5\\% gpt-5.4-mini (equal USD)", "fair_latency_025": "25\\% matched fast latency"}

def tab_library():
    p = RUNS / "library_scale" / "outcomes.parquet"
    if not p.exists():
        return
    df = pd.read_parquet(p)
    df["cell"] = df["name"].str.split("/").str[-1]
    order = ["n100_share0", "n100_share25", "n1000_share0", "n1000_share25"]
    lab = {"n100_share0": "$N=100$, 0\\%", "n100_share25": "$N=100$, 25\\% library", "n1000_share0": "$N=1{,}000$, 0\\%", "n1000_share25": "$N=1{,}000$, 25\\% library"}
    outs = ["mad_pct", "time_to_absorb", "spread_bp", "depth5", "lambda_bp_100", "vr5", "acf_r_1", "max_dd", "llm_share_volume", "llm_valid_rate"]
    lines = ["\\begin{tabular}{l" + "r" * len(outs) + "}", "\\toprule", "Cell & " + " & ".join(NICE.get(o, o) for o in outs) + " \\\\", "\\midrule"]
    for c in order:
        d = df[df.cell == c]
        if len(d) == 0:
            continue
        lines.append(f"{lab[c]} & " + " & ".join(fmt(d[o].astype(float).mean()) for o in outs) + " \\\\")
    lines += ["\\midrule", f"Seeds per cell & {int(df.groupby('cell').size().min())}" + " &" * (len(outs) - 1) + " \\\\", "\\bottomrule", "\\end{tabular}"]
    write("tab_library.tex", "\n".join(lines))


if __name__ == "__main__":
    tab_library()
    tab_stylized_facts()
    tab_cells("classical_confirmatory", CLASSICAL_CELLS, CLASSICAL_LAB, ["mad_pct", "time_to_absorb", "spread_bp", "lambda_bp_100", "n_bubbles", "max_dd", "n_crashes", "acf_absr_1_10"], "tab_classical_cells.tex", "control")
    tab_effects("classical_confirmatory", CLASSICAL_CELLS, CLASSICAL_LAB, ["mad_pct", "time_to_absorb", "spread_bp", "n_bubbles", "max_dd"], "tab_classical_effects.tex", "control")
    tab_cells("classical_confirmatory", ["k2_control", "k2_leverage"], {"k2_control": "Two assets", "k2_leverage": "Two assets + leveraged 19\\% + funding shock"},
              ["ret_corr", "fund_corr", "excess_comovement", "spillover_dy", "mad_pct", "max_dd"], "tab_contagion.tex", "k2_control")
    tab_dose("pilot_openai", "tab_dose_pilot.tex")
    tab_dose("confirmatory_openai", "tab_dose_conf.tex")
    tab_cells("confirmatory_openai", ARM_CELLS, ARM_LAB, ["mad_pct", "time_to_absorb", "spread_bp", "n_bubbles", "max_dd", "llm_pairwise_corr", "llm_synchrony", "llm_valid_rate"], "tab_arms_conf.tex", "dose_025")
    tab_effects("confirmatory_openai", ARM_CELLS, ARM_LAB, ["mad_pct", "time_to_absorb", "spread_bp", "n_bubbles", "max_dd"], "tab_arms_effects.tex", "dose_025")
    tab_llm_reliability("pilot_openai", "tab_llm_reliability_pilot.tex")
    tab_llm_reliability("confirmatory_openai", "tab_llm_reliability_conf.tex")
    print("tables in", OUT)