"""Stage F.1 — idiosyncratic shock detection and propagation measurement.

Per month m, on the LAGGED top-500 universe (RTH, flags excluded):

  Residualization (all inputs lagged, CLAUDE.md §8):
    beta_i and residual scale MAD_i estimated on month m-1; minute-of-day
    seasonal factor f_tau = cross-sectional median |residual| profile of m-1,
    normalized to mean 1. In month m:
        e_it = r_it - beta_i * r_mkt(-i),t          (leave-one-out EW market)
        z_it = e_it / (1.4826 * MAD_i * f_tau)

  Shock (baseline |z| >= 6, sign recorded; 4 and 5 stored for robustness):
    isolation: |z_mkt,t| < 1  AND  share of |z_jt| > 4 at t  <= 2%.

  Outcomes for each event at h in {1,2,5,10,15,30,60,120} minutes, same session
  only (censored otherwise):
    breadth(h)   = share of eligible j != i with |z_{j,t+h}| > 2
    intensity(h) = mean |z_{j,t+h}|
    signed(h)    = mean sign(e_i) * z_{j,t+h}
    pre-shock falsification: same outcomes at t-5.

  Matched placebo per event: same ticker, same minute-of-day, random other day
  in month with valid quiet z (|z|<2) and calm market — same outcome code path.

Output: outputs/estimates/events/events_YYYY-MM.parquet
Usage:  python scripts/11_... [--workers N] [--month YYYY-MM]
"""

from __future__ import annotations

import argparse
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import numpy as np
import polars as pl

ROOT = Path(__file__).resolve().parents[1]
MINUTE = ROOT / "data" / "processed" / "minute"
UNIV = ROOT / "data" / "metadata" / "universe_monthly.parquet"
OUT = ROOT / "outputs" / "estimates" / "events"

HORIZONS = [1, 2, 5, 10, 15, 30, 60, 120]
Z_SHOCK = 6.0
Z_RECEIVER = 2.0
Z_MKT_CALM = 1.0
SIMUL_CAP = 0.02
PRE_H = 5
SEED = 20260713


def month_pair(month: str) -> str:
    y, m = int(month[:4]), int(month[5:7])
    m -= 1
    if m == 0:
        y, m = y - 1, 12
    return f"{y:04d}-{m:02d}"


def load_matrix(month: str, tickers: list[str]):
    """(returns matrix [rows x N], day-block slices, mods, dates) on the RTH grid."""
    y, mm = month.split("-")
    df = (
        pl.scan_parquet(MINUTE / f"year={y}" / f"month={mm}" / "data.parquet")
        .filter((pl.col("session") == "rth") & ~pl.col("flag_ohlc")
                & ~pl.col("flag_split_window") & pl.col("ticker").is_in(tickers))
        .select("ticker", "date", "mod", "close")
        .collect()
    )
    if df.height == 0:
        return None
    grid = df.pivot(on="ticker", index=["date", "mod"], values="close").sort("date", "mod")
    cols = [c for c in grid.columns if c not in ("date", "mod")]
    px = grid.select(cols).to_numpy()
    dates = grid["date"].to_numpy()
    mods = grid["mod"].to_numpy()
    logpx = np.log(px)
    ret = np.vstack([np.full((1, logpx.shape[1]), np.nan),
                     logpx[1:] - logpx[:-1]])
    newday = np.empty(len(dates), dtype=bool)
    newday[0] = True
    newday[1:] = dates[1:] != dates[:-1]
    ret[newday] = np.nan  # never bridge overnight
    # day block bounds
    starts = np.flatnonzero(newday)
    ends = np.append(starts[1:], len(dates))
    return ret, cols, dates, mods, list(zip(starts, ends))


def residualize(ret: np.ndarray, beta: np.ndarray):
    """z-scores with LOO EW market; returns (e, mkt_raw)."""
    valid = ~np.isnan(ret)
    s = np.nansum(ret, axis=1)
    n = valid.sum(axis=1).astype(float)
    with np.errstate(invalid="ignore", divide="ignore"):
        mkt = s / n
        loo = (s[:, None] - np.where(valid, ret, 0.0)) / np.maximum(n[:, None] - valid, 1.0)
    e = ret - beta[None, :] * loo
    return e, mkt, n


def month_params(prev_month: str, tickers: list[str]):
    """beta, MAD, seasonal factor from the PRIOR month."""
    loaded = load_matrix(prev_month, tickers)
    if loaded is None:
        return None
    ret, cols, dates, mods, _ = loaded
    valid = ~np.isnan(ret)
    s = np.nansum(ret, axis=1)
    n = valid.sum(axis=1).astype(float)
    with np.errstate(invalid="ignore", divide="ignore"):
        mkt = s / n
    ok = np.isfinite(mkt) & (n >= 50)
    r, m = ret[ok], mkt[ok]
    mv = np.nanvar(m)
    beta = np.array([
        (np.nanmean((r[:, j] - np.nanmean(r[:, j])) * (m - m.mean())) / mv
         if np.isfinite(r[:, j]).sum() > 300 else 1.0)
        for j in range(r.shape[1])])
    beta = np.clip(np.nan_to_num(beta, nan=1.0), -3, 5)
    e = r - beta[None, :] * m[:, None]
    med = np.nanmedian(e, axis=0)
    mad = np.nanmedian(np.abs(e - med[None, :]), axis=0)
    mad = np.where((mad > 0) & np.isfinite(mad), mad, np.nan)
    # seasonal factor: cross-sectional median |e| per minute-of-day, normalized
    mods_ok = mods[ok]
    absmed = np.abs(e / (1.4826 * mad[None, :]))
    ftau = {}
    for mod in np.unique(mods_ok):
        ftau[int(mod)] = np.nanmedian(absmed[mods_ok == mod])
    mean_f = np.nanmean(list(ftau.values()))
    ftau = {k: (v / mean_f if np.isfinite(v) and v > 0 else 1.0)
            for k, v in ftau.items()}
    mkt_scale = 1.4826 * np.nanmedian(np.abs(m - np.nanmedian(m)))
    return dict(zip(cols, beta)), dict(zip(cols, mad)), ftau, mkt_scale


def outcomes_at(z: np.ndarray, row: int, j: int, d0: int, d1: int):
    zr = z[row]
    mask = np.isfinite(zr)
    mask[j] = False
    n = int(mask.sum())
    if n < 50 or not (d0 <= row < d1):
        return None
    az = np.abs(zr[mask])
    return (float((az > Z_RECEIVER).mean()), float(az.mean()),
            float(zr[mask].mean()), n)


def process_month(month: str, force: bool = False) -> str:
    dest = OUT / f"events_{month}.parquet"
    if dest.exists() and not force:
        return f"{month} skip"
    prev = month_pair(month)
    uni = (pl.read_parquet(UNIV)
           .filter((pl.col("month") == month) & pl.col("top500"))["ticker"].to_list())
    if not uni:
        dest.parent.mkdir(parents=True, exist_ok=True)
        pl.DataFrame([]).write_parquet(dest)
        return f"{month} no-universe"
    params = month_params(prev, uni)
    loaded = load_matrix(month, uni)
    if params is None or loaded is None:
        dest.parent.mkdir(parents=True, exist_ok=True)
        pl.DataFrame([]).write_parquet(dest)
        return f"{month} no-data"
    beta_d, mad_d, ftau, mkt_scale = params
    ret, cols, dates, mods, blocks = loaded
    beta = np.array([beta_d.get(c, 1.0) for c in cols])
    mad = np.array([mad_d.get(c, np.nan) for c in cols])
    e, mkt, n_t = residualize(ret, beta)
    f = np.array([ftau.get(int(m), 1.0) for m in mods])
    with np.errstate(invalid="ignore", divide="ignore"):
        z = e / (1.4826 * mad[None, :] * f[:, None])
        z_mkt = mkt / (mkt_scale * f)

    absz = np.abs(z)
    simul = np.nanmean(absz > 4.0, axis=1)
    row_ok = (np.abs(z_mkt) < Z_MKT_CALM) & (simul <= SIMUL_CAP) & (n_t >= 50)

    # day block id per row
    block_of = np.empty(len(mods), dtype=int)
    bounds = []
    for b, (s0, s1) in enumerate(blocks):
        block_of[s0:s1] = b
        bounds.append((s0, s1))

    ev_rows, ev_cols = np.where((absz >= Z_SHOCK) & row_ok[:, None])
    rng = np.random.default_rng(SEED + hash(month) % 100000)
    records = []

    def build_record(row, j, kind):
        d0, d1 = bounds[block_of[row]]
        rec = {
            "month": month, "kind": kind, "ticker": cols[j],
            "date": str(dates[row]), "mod": int(mods[row]),
            "z": float(z[row, j]), "sign": int(np.sign(z[row, j])),
            "mkt_z": float(z_mkt[row]),
            "mkt_vol30": float(np.nanstd(mkt[max(row - 30, d0):row]) * 1e4),
        }
        pre = outcomes_at(z, row - PRE_H, j, d0, d1) if row - PRE_H >= d0 else None
        rec["pre_breadth"], rec["pre_intensity"] = (
            (pre[0], pre[1]) if pre else (None, None))
        got = False
        for h in HORIZONS:
            o = outcomes_at(z, row + h, j, d0, d1) if row + h < d1 else None
            rec[f"breadth_{h}"] = o[0] if o else None
            rec[f"intensity_{h}"] = o[1] if o else None
            rec[f"signed_{h}"] = (o[2] * rec["sign"]) if o else None
            rec[f"nelig_{h}"] = o[3] if o else None
            got = got or (o is not None)
        return rec if got else None

    for row, j in zip(ev_rows, ev_cols):
        rec = build_record(row, j, "real")
        if rec is None:
            continue
        records.append(rec)
        # matched placebo: same ticker & minute-of-day, other calm day
        cand = np.flatnonzero(
            (mods == rec["mod"]) & row_ok & np.isfinite(z[:, j])
            & (np.abs(z[:, j]) < 2.0) & (block_of != block_of[row]))
        if len(cand):
            prow = int(rng.choice(cand))
            prec = build_record(prow, j, "placebo")
            if prec is not None:
                prec["z"] = rec["z"]  # carry source magnitude for matching
                prec["sign"] = rec["sign"]
                records.append(prec)

    dest.parent.mkdir(parents=True, exist_ok=True)
    pl.DataFrame(records).write_parquet(dest)
    nre = sum(1 for r in records if r["kind"] == "real")
    return f"{month} events={nre} placebos={len(records)-nre}"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--month", default=None)
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()
    months = sorted(p.stem.replace("daily_", "") for p in
                    (ROOT / "data" / "metadata" / "daily").glob("daily_*.parquet"))
    months = [m for m in months if m > "1992-01"]  # first month has no lag
    if args.month:
        print(process_month(args.month, True))
        return
    with ProcessPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(process_month, m, args.force): m for m in months}
        for i, fut in enumerate(as_completed(futs), 1):
            print(f"[{i}/{len(months)}] {fut.result()}", flush=True)


if __name__ == "__main__":
    main()