"""Stage A pilot figures and summary-statistics table (CLAUDE.md §33.11-12).

Figure P1 — dataset coverage: monthly compressed size (all 411 months, from the
pinned inventory) and symbols per trading day (three pilot months).
Figure P2 — intraday seasonality by minute of day for the pilot months, computed
on the top-500 dollar-volume universe within each month (descriptive, full-month
normalization — documented; predictive designs will use lagged estimates).
Table P1 — summary statistics of core minute-level variables, same universe.

Each figure's source data is saved to outputs/figures/data/ so the figure can be
reproduced without re-running this script's heavy pass.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import polars as pl

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

RAW = ROOT / "data" / "raw" / "data"
FIGDIR = ROOT / "outputs" / "figures"
FIGDATA = FIGDIR / "data"
TABDIR = ROOT / "outputs" / "tables"

MONTHS = ["1992-06", "2009-06", "2025-12"]
COLORS = {"1992-06": "#0072B2", "2009-06": "#D55E00", "2025-12": "#009E73"}
STYLES = {"1992-06": "-", "2009-06": "--", "2025-12": ":"}
TOP_N = 500

plt.rcParams.update({
    "font.family": "serif",
    "font.size": 9,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "axes.grid": True,
    "grid.color": "0.9",
    "grid.linewidth": 0.5,
    "axes.axisbelow": True,
    "figure.facecolor": "white",
    "savefig.facecolor": "white",
    "pdf.fonttype": 42,
})


def load_month_rth(month: str) -> pl.DataFrame:
    """RTH bars with ET clock fields for one pilot month."""
    df = pl.read_parquet(RAW / f"ohlcv_{month}.parquet")
    et = df.with_columns(
        pl.col("timestamp").dt.convert_time_zone("America/New_York").alias("ts_et"))
    et = et.with_columns(
        pl.col("ts_et").dt.date().alias("date"),
        (pl.col("ts_et").dt.hour().cast(pl.Int32) * 60
         + pl.col("ts_et").dt.minute()).alias("mod"),
    )
    return et.filter((pl.col("mod") >= 570) & (pl.col("mod") < 960))


def top_universe(rth: pl.DataFrame, n: int = TOP_N) -> list[str]:
    dv = (rth.group_by("ticker")
          .agg((pl.col("close") * pl.col("volume")).sum().alias("dollar_vol"))
          .sort("dollar_vol", descending=True))
    return dv.head(n)["ticker"].to_list()


def seasonality(rth: pl.DataFrame, universe: list[str]) -> pl.DataFrame:
    """Minute-of-day means of |log return| (bps), volume share, log range (bps)."""
    u = rth.filter(pl.col("ticker").is_in(universe)).sort("ticker", "date", "mod")
    u = u.with_columns(
        (pl.col("close").log() - pl.col("close").log().shift(1)
         ).over("ticker", "date").alias("ret"),
        (pl.col("high").log() - pl.col("low").log()).alias("rng"),
    )
    day_vol = u.group_by("ticker", "date").agg(pl.col("volume").sum().alias("dvol"))
    u = u.join(day_vol, on=["ticker", "date"]).with_columns(
        (pl.col("volume") / pl.col("dvol")).alias("vshare"))
    return (u.group_by("mod").agg(
        (pl.col("ret").abs().mean() * 1e4).alias("absret_bps"),
        (pl.col("vshare").mean() * 100).alias("volshare_pct"),
        (pl.col("rng").mean() * 1e4).alias("range_bps"),
        pl.len().alias("n_bars"),
    ).sort("mod"))


def summary_stats(rth: pl.DataFrame, universe: list[str], month: str) -> pl.DataFrame:
    u = rth.filter(pl.col("ticker").is_in(universe)).sort("ticker", "date", "mod")
    u = u.with_columns(
        ((pl.col("close").log() - pl.col("close").log().shift(1)) * 1e4
         ).over("ticker", "date").alias("ret_bps"),
        ((pl.col("high").log() - pl.col("low").log()) * 1e4).alias("range_bps"),
        (pl.col("close") * pl.col("volume") / 1e3).alias("dollar_vol_k"),
    )
    bars = u.group_by("ticker", "date").agg(pl.len().alias("bars_per_day"))
    rows = []
    for var, series in [
        ("Minute log return (bps)", u["ret_bps"].drop_nulls()),
        ("High--low log range (bps)", u["range_bps"]),
        ("Dollar volume per bar (\\$000)", u["dollar_vol_k"]),
        ("Valid RTH bars per stock-day", bars["bars_per_day"].cast(pl.Float64)),
    ]:
        rows.append({
            "month": month, "variable": var,
            "mean": float(series.mean()), "sd": float(series.std()),
            "p1": float(series.quantile(0.01)), "p5": float(series.quantile(0.05)),
            "p50": float(series.median()),
            "p95": float(series.quantile(0.95)), "p99": float(series.quantile(0.99)),
            "n": series.len(),
        })
    return pl.DataFrame(rows)


def fig_coverage() -> None:
    inv = json.loads((ROOT / "data/manifests/repo_inventory.json").read_text())
    rec = [(f["path"][-15:-8], f["size"] / 1e6) for f in inv["files"]
           if f["path"].endswith(".parquet")]
    rec.sort()
    x = [int(m[:4]) + (int(m[5:7]) - 0.5) / 12 for m, _ in rec]
    y = [s for _, s in rec]
    pl.DataFrame({"month": [m for m, _ in rec], "mb_compressed": y}).write_csv(
        FIGDATA / "figP1a_monthly_file_size.csv")

    per_day = {}
    for month in MONTHS:
        rep = json.loads(
            (ROOT / f"outputs/diagnostics/pilot_audit_{month}.json").read_text())
        per_day[month] = [(d["date"], d["symbols"]) for d in rep["per_day"]]
    pl.DataFrame([
        {"month": m, "date": d, "symbols": s}
        for m, days in per_day.items() for d, s in days
    ]).write_csv(FIGDATA / "figP1b_symbols_per_day.csv")

    fig, axes = plt.subplots(1, 2, figsize=(7.0, 2.6))
    ax = axes[0]
    ax.plot(x, y, color="#0072B2", lw=1.2)
    ax.set_xlabel("Year")
    ax.set_ylabel("Compressed file size (MB / month)")
    ax.set_title("(a) Raw data volume by month", fontsize=9, loc="left")
    ax.set_xlim(1992, 2026.5)

    ax = axes[1]
    for month in MONTHS:
        days = list(range(1, len(per_day[month]) + 1))
        sym = [s for _, s in per_day[month]]
        ax.plot(days, sym, STYLES[month], color=COLORS[month], lw=1.4)
        ax.annotate(month, xy=(days[-1], sym[-1]), xytext=(3, 0),
                    textcoords="offset points", color=COLORS[month],
                    fontsize=8, va="center")
    ax.set_xlabel("Trading day of month")
    ax.set_ylabel("Symbols with $\\geq$1 bar")
    ax.set_title("(b) Active symbols per day, pilot months", fontsize=9, loc="left")
    ax.set_xlim(1, 27)

    fig.tight_layout()
    for ext in ("pdf", "png"):
        fig.savefig(FIGDIR / f"figP1_coverage.{ext}", dpi=300)
    plt.close(fig)


def fig_seasonality(seas: dict[str, pl.DataFrame]) -> None:
    panels = [
        ("absret_bps", "(a) $|$log return$|$ (bps)"),
        ("volshare_pct", "(b) Volume share (%)"),
        ("range_bps", "(c) High–low range (bps)"),
    ]
    fig, axes = plt.subplots(1, 3, figsize=(7.0, 2.4), sharex=True)
    for ax, (col, title) in zip(axes, panels):
        for month in MONTHS:
            s = seas[month]
            hours = (s["mod"] - 570) / 60.0
            ax.plot(hours, s[col], STYLES[month], color=COLORS[month], lw=1.2,
                    label=month)
        ax.set_title(title, fontsize=9, loc="left")
        ax.set_xlabel("Hours since 09:30 ET")
        ax.set_xlim(0, 6.5)
        ax.set_xticks([0, 2, 4, 6])
    axes[0].legend(frameon=False, fontsize=7.5, loc="upper right",
                   borderaxespad=0.1, handlelength=1.8)
    fig.tight_layout()
    for ext in ("pdf", "png"):
        fig.savefig(FIGDIR / f"figP2_seasonality.{ext}", dpi=300)
    plt.close(fig)


def latex_table(stats: pl.DataFrame) -> str:
    month_label = {"1992-06": "June 1992", "2009-06": "June 2009",
                   "2025-12": "December 2025"}
    lines = [
        "% Auto-generated by scripts/02_pilot_figures_tables.py — do not edit numbers by hand.",
        "\\begin{table}[htbp]",
        "\\centering",
        "\\begin{adjustbox}{max width=\\textwidth}\\begin{threeparttable}",
        "\\caption{Summary statistics, pilot months}",
        "\\label{tab:pilot_summary}",
        "\\small",
        "\\begin{tabular}{llrrrrrrrr}",
        "\\toprule",
        "Month & Variable & Mean & SD & p1 & p5 & p50 & p95 & p99 & $N$ \\\\",
        "\\midrule",
    ]
    fmt = lambda v: f"{v:,.2f}" if abs(v) < 1e5 else f"{v:,.0f}"
    prev = None
    for r in stats.iter_rows(named=True):
        m = month_label[r["month"]] if r["month"] != prev else ""
        if r["month"] != prev and prev is not None:
            lines.append("\\addlinespace")
        prev = r["month"]
        lines.append(
            f"{m} & {r['variable']} & {fmt(r['mean'])} & {fmt(r['sd'])} & "
            f"{fmt(r['p1'])} & {fmt(r['p5'])} & {fmt(r['p50'])} & "
            f"{fmt(r['p95'])} & {fmt(r['p99'])} & {r['n']:,} \\\\")
    lines += [
        "\\bottomrule",
        "\\end{tabular}",
        "\\begin{tablenotes}[flushleft]\\footnotesize",
        "\\item Notes: Regular trading hours (09:30--16:00 ET, exchange calendar) for the "
        f"top {TOP_N} securities by within-month dollar volume in each pilot month of the "
        "pinned Finnhub.io minute-bar aggregation. Minute log "
        "returns use consecutive same-day closes; the high--low range is "
        "$\\log(H)-\\log(L)$ per one-minute bar; bars exist only for minutes with at "
        "least one reported trade. Prices are unadjusted for splits; the pilot is "
        "descriptive and imposes no corporate-action screens.",
        "\\end{tablenotes}",
        "\\end{threeparttable}\\end{adjustbox}",
        "\\end{table}",
    ]
    return "\n".join(lines)


def main() -> None:
    for d in (FIGDATA, TABDIR):
        d.mkdir(parents=True, exist_ok=True)

    seas, stats_frames = {}, []
    for month in MONTHS:
        rth = load_month_rth(month)
        uni = top_universe(rth)
        seas[month] = seasonality(rth, uni)
        seas[month].write_csv(FIGDATA / f"figP2_seasonality_{month}.csv")
        stats_frames.append(summary_stats(rth, uni, month))
        print(f"{month}: universe={len(uni)} seasonality rows={seas[month].height}")

    stats = pl.concat(stats_frames)
    stats.write_csv(TABDIR / "tabP1_pilot_summary.csv")
    tex = latex_table(stats)
    (TABDIR / "tabP1_pilot_summary.tex").write_text(tex)
    (ROOT / "paper" / "tables" / "tabP1_pilot_summary.tex").write_text(tex)

    fig_coverage()
    fig_seasonality(seas)
    for f in ("figP1_coverage.pdf", "figP2_seasonality.pdf"):
        (ROOT / "paper" / "figures" / f).write_bytes((FIGDIR / f).read_bytes())
    print("figures + table written")


if __name__ == "__main__":
    main()