"""Stage D.1 — refine split candidates into high-confidence events + exclusion windows.

Stage C flagged 442k ticker-days with overnight ratios near rational split ratios
(over-inclusive). Here we keep events where:
  * the overnight ratio prev_close/first_open is within TOL of a canonical ratio
    r in {2,3,4,5,6,7,8,10,20,3/2} (forward) or 1/r (reverse); tighter TOL for 3/2;
  * both adjacent days have RTH dollar volume >= MIN_DV (kills dead-OTC noise);
  * prev_close >= MIN_PRICE.

Outputs:
  outputs/diagnostics/stage_d/split_events.parquet   (ticker, date, ratio, kind)
  outputs/diagnostics/stage_d/split_windows.parquet  (ticker, date) — event day ± WINDOW days
"""

from __future__ import annotations

from pathlib import Path

import polars as pl

ROOT = Path(__file__).resolve().parents[1]
DAILY = ROOT / "data" / "metadata" / "daily"
OUT = ROOT / "outputs" / "diagnostics" / "stage_d"

CANONICAL = [2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 20.0]
TOL, TOL_32 = 0.015, 0.008
MIN_DV = 50_000.0
MIN_PRICE = 1.0
WINDOW = 1  # exclusion: event day +/- 1 trading day


def main() -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    d = (
        pl.scan_parquet(str(DAILY / "daily_*.parquet"))
        .filter(pl.col("rth_last_close").is_not_null() & (pl.col("rth_volume") > 0))
        .select("ticker", "date", "rth_first_open", "rth_last_close", "rth_dollar_vol")
        .sort("ticker", "date")
        .with_columns(
            pl.col("rth_last_close").shift(1).over("ticker").alias("prev_close"),
            pl.col("rth_dollar_vol").shift(1).over("ticker").alias("prev_dv"),
            pl.col("date").shift(1).over("ticker").alias("prev_date"),
        )
        .filter(
            pl.col("prev_close").is_not_null()
            & ((pl.col("date") - pl.col("prev_date")).dt.total_days() <= 7)
            & (pl.col("prev_dv") >= MIN_DV)
            & (pl.col("rth_dollar_vol") >= MIN_DV)
            & (pl.col("prev_close") >= MIN_PRICE)
        )
        .with_columns((pl.col("prev_close") / pl.col("rth_first_open"))
                      .alias("overnight_ratio"))
    )

    cond = None
    exprs = []
    for r in CANONICAL + [1.5]:
        tol = TOL_32 if r == 1.5 else TOL
        for ratio, kind in [(r, "split"), (1.0 / r, "reverse_split")]:
            c = (pl.col("overnight_ratio") / ratio - 1).abs() < tol
            cond = c if cond is None else (cond | c)
            exprs.append(
                pl.when((pl.col("overnight_ratio") / ratio - 1).abs() < tol)
                .then(pl.lit(f"{kind}:{r:g}")).otherwise(None))

    events = (
        d.filter(cond)
        .with_columns(pl.coalesce(exprs).alias("event"))
        .select("ticker", "date", "overnight_ratio", "event",
                "prev_close", "rth_first_open")
        .collect()
    )
    events.write_parquet(OUT / "split_events.parquet")

    # exclusion windows: event date +/- WINDOW *observed* trading days per ticker
    dates = (
        pl.scan_parquet(str(DAILY / "daily_*.parquet"))
        .select("ticker", "date").unique().sort("ticker", "date")
        .with_columns(pl.int_range(pl.len()).over("ticker").alias("idx"))
        .collect()
    )
    ev = events.join(dates, on=["ticker", "date"], how="left")
    win = []
    for off in range(-WINDOW, WINDOW + 1):
        win.append(ev.select("ticker", (pl.col("idx") + off).alias("idx")))
    windows = (
        pl.concat(win).unique()
        .join(dates, on=["ticker", "idx"], how="inner")
        .select("ticker", "date").unique().sort("ticker", "date")
    )
    windows.write_parquet(OUT / "split_windows.parquet")

    print(f"events={events.height} "
          f"(splits={(events['event'].str.starts_with('split')).sum()}, "
          f"reverse={(events['event'].str.starts_with('reverse')).sum()}) "
          f"window_days={windows.height}")


if __name__ == "__main__":
    main()