"""Generate every figure of the paper from saved outcomes/run artefacts. Each figure is written to
results/figures/ (PDF + PNG) and copied to paper/figures/ by `make tables`. Designs that have not been run
are skipped with a message (no figure is ever drawn from invented numbers)."""
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, local_projection, paired_effects  # noqa: E402
from market_agents.analysis import plots as P  # noqa: E402

RUNS = ROOT / "results/runs"
OUT = ROOT / "results/figures"
OUT.mkdir(parents=True, exist_ok=True)
LABELS = {"msd_log": "squared log pricing error", "mad_pct": "mean |P/V − 1|, %", "time_to_absorb": "time to absorb shock (ticks)",
          "response_ratio": "20-tick response ratio", "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 (Σ ACF|r|)",
          "kurtosis": "excess kurtosis", "vr5": "variance ratio (5)", "llm_synchrony": "LLM synchrony", "llm_pairwise_corr": "LLM order-flow corr."}


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


def ticks_of(run_dir: str) -> pd.DataFrame:
    return pd.read_parquet(Path(run_dir) / "ticks.parquet")


def fig_baseline_paths():
    df = load("classical_confirmatory")
    if df is None:
        return
    rows = [df[(df.cell == c) & (df.seed == 101)].iloc[0] for c in ("control", "leverage_heavy_funding_shock", "liquidity_shock")]
    fig = P.price_paths([("control", ticks_of(rows[0].run_dir)), ("leverage 19 % + funding shock", ticks_of(rows[1].run_dir)),
                         ("market makers withdraw", ticks_of(rows[2].run_dir))], shock_t=300)
    P.save(fig, OUT / "fig_baseline_paths.pdf")


def fig_classical_mechanisms():
    df = load("classical_confirmatory")
    if df is None:
        return
    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"]
    lab = {"control": "control", "placebo_shock": "placebo", "trend_heavy": "trend 16 %", "leverage_heavy": "leverage 19 %",
           "leverage_heavy_funding_shock": "leverage + funding shock", "liquidity_shock": "MM withdraw", "info_low": "precision 0.5",
           "info_high": "precision 8", "misleading_news": "misleading news", "breaker_5": "circuit breaker 5 %", "poslimit_300": "position limit", "auction_5": "call auction"}
    for oc in ("mad_pct", "time_to_absorb", "max_dd", "n_bubbles", "spread_bp", "lambda_bp_100"):
        fig = P.cell_bars(df, oc, cells, lab, LABELS[oc], control="control")
        P.save(fig, OUT / f"fig_classical_{oc}.pdf")


def fig_dose_response(design: str, tag: 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)
    if d["share"].nunique() < 3:
        return
    eff = {}
    for oc in ("msd_log", "mad_pct", "time_to_absorb", "spread_bp", "lambda_bp_100", "n_bubbles", "max_dd", "acf_absr_1_10", "kurtosis"):
        if oc in d:
            eff[oc] = dose_response(d, oc, "share", 0.0)
    fig = P.dose_response(eff, LABELS)
    P.save(fig, OUT / f"fig_dose_{tag}.pdf")
    # realized shares
    fig = P.scatter_mechanism(d, "llm_share_traders", "llm_share_volume", "share", "LLM share of trader slots", "LLM share of traded volume")
    P.save(fig, OUT / f"fig_realized_share_{tag}.pdf")
    # synchrony / correlation vs share (LLM-only cells)
    dd = d[d.share > 0]
    if "llm_pairwise_corr" in dd and dd["llm_pairwise_corr"].notna().any():
        fig = P.scatter_mechanism(dd, "share", "llm_pairwise_corr", "share", "LLM share", "pairwise order-flow correlation (LLM agents)")
        P.save(fig, OUT / f"fig_sync_{tag}.pdf")
    # local projections around the t = 300 shock: 25 % and 50 % vs 0 %
    lp = {}
    for lv in (0.25, 0.50, 1.00):
        sub = d[d.share.isin([0.0, lv])]
        if sub.share.nunique() < 2:
            continue
        tb = {r.run_dir: ticks_of(r.run_dir) for r in sub.itertuples()}
        tr = {r.run_dir: r.share for r in sub.itertuples()}
        lp[f"{int(lv * 100)} % LLM"] = local_projection(tb, tr, 300, H=60)
    if lp:
        P.save(P.local_projection_plot(lp), OUT / f"fig_lp_{tag}.pdf")
    # paths for one seed
    seed = int(d.seed.min())
    runs = [(f"{int(s * 100)} % LLM", ticks_of(d[(d.share == s) & (d.seed == seed)].iloc[0].run_dir)) for s in sorted(d.share.unique()) if ((d.share == s) & (d.seed == seed)).any()][:4]
    P.save(P.price_paths(runs, shock_t=300), OUT / f"fig_paths_{tag}.pdf")


def fig_arms(design: str, tag: str):
    df = load(design)
    if df is None or "tr_arm" not in df:
        return
    ref = "dose_025"
    cells = [c for c in ["dose_000", "dose_025", "div_prompts_025", "div_models_025", "memory_025", "channel_025", "adaptive_025",
                         "leverage_025", "breaker_025", "poslimit_025", "fair_mini_005", "fair_latency_025"] if c in set(df.cell)]
    lab = {"dose_000": "0 % (control)", "dose_025": "25 % single", "div_prompts_025": "25 % diverse prompts", "div_models_025": "25 % two models",
           "memory_025": "25 % memory", "channel_025": "25 % channel", "adaptive_025": "25 % adaptive", "leverage_025": "25 % + leverage",
           "breaker_025": "25 % + breaker", "poslimit_025": "25 % + pos. limit", "fair_mini_005": "5 % mini (equal USD)", "fair_latency_025": "25 % fast latency"}
    for oc in ("mad_pct", "time_to_absorb", "spread_bp", "n_bubbles", "max_dd", "llm_pairwise_corr", "llm_synchrony", "llm_valid_rate"):
        if oc in df:
            P.save(P.cell_bars(df, oc, cells, lab, LABELS.get(oc, oc), control=ref), OUT / f"fig_arms_{tag}_{oc}.pdf")
    if {"llm_usd", "mad_pct"} <= set(df.columns):
        dd = df[df.cell.isin(["dose_005", "dose_025", "dose_050", "dose_100", "div_models_025", "fair_mini_005"])]
        P.save(P.cost_quality(dd, "llm_usd", "mad_pct", "cell", "API cost per run (USD)", LABELS["mad_pct"]), OUT / f"fig_cost_quality_{tag}.pdf")


def fig_contagion():
    df = load("classical_confirmatory")
    if df is None:
        return
    cells = [c for c in ("k2_control", "k2_leverage") if c in set(df.cell)]
    if cells and "excess_comovement" in df:
        P.save(P.cell_bars(df, "excess_comovement", cells, {"k2_control": "two assets", "k2_leverage": "two assets + leverage + funding shock"},
                           "excess comovement (return corr − fundamental corr)"), OUT / "fig_contagion_excess.pdf")
        P.save(P.cell_bars(df, "spillover_dy", cells, {"k2_control": "two assets", "k2_leverage": "two assets + leverage + funding shock"},
                           "Diebold–Yilmaz spillover (%)"), OUT / "fig_contagion_spillover.pdf")


def fig_wealth():
    df = load("classical_confirmatory")
    if df is None:
        return
    r = df[(df.cell == "control") & (df.seed == 101)].iloc[0]
    ag = pd.read_parquet(Path(r.run_dir) / "agents.parquet")
    P.save(P.wealth_and_inventory(ag, ticks_of(r.run_dir)), OUT / "fig_wealth_depth.pdf")


if __name__ == "__main__":
    fig_baseline_paths()
    fig_classical_mechanisms()
    fig_contagion()
    fig_wealth()
    fig_dose_response("pilot_openai", "pilot")
    fig_dose_response("confirmatory_openai", "conf")
    fig_arms("confirmatory_openai", "conf")
    print("figures in", OUT)