"""Stage D.2 — monthly universe membership with STRICTLY LAGGED screens.

For each month m (1992-02 onward), membership is decided from month m-1 daily
metadata only (no look-ahead, CLAUDE.md §6.2):
  * days_present >= 10;
  * median RTH close >= $5 (config: sample_filters.baseline.min_price_lagged);
  * median RTH bars/day >= 100 (min_bars_per_day_lagged);
  * rank by total RTH dollar volume among eligible -> top500 / top1000 / top3000.

A crude common-stock heuristic flags likely non-common instruments from ticker
shape (documented limitation; no security master exists in the source):
  * '.'/'-' class/series suffixes are kept (share classes are common stock);
  * 5-letter tickers ending in F or Y -> foreign/ADR OTC;
  * tickers ending in Q (bankruptcy), W/WS (warrants), U (units), P/PR (preferred).

Output: data/metadata/universe_monthly.parquet
  (month, ticker, days, med_close, med_bars, dollar_vol, eligible,
   top500, top1000, top3000, flag_noncommon)
"""

from __future__ import annotations

from pathlib import Path

import polars as pl
import yaml

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

cfg = yaml.safe_load((ROOT / "config" / "sample_filters.yaml").read_text())["baseline"]
MIN_PRICE = float(cfg["min_price_lagged"])
MIN_BARS = int(cfg["min_bars_per_day_lagged"])
MIN_DAYS = 10


def month_key(p: Path) -> str:
    return p.stem.replace("daily_", "")


def next_month(m: str) -> str:
    y, mm = int(m[:4]), int(m[5:7])
    mm += 1
    if mm == 13:
        y, mm = y + 1, 1
    return f"{y:04d}-{mm:02d}"


NONCOMMON = (
    (pl.col("ticker").str.len_chars() == 5)
    & pl.col("ticker").str.slice(4, 1).is_in(["F", "Y", "Q", "W", "U"])
) | pl.col("ticker").str.contains(r"(\.WS|\.W|\.U|\.P|-P|\.PR)$")


def main() -> None:
    files = sorted(DAILY.glob("daily_*.parquet"))
    frames = []
    for f in files:
        m_prev = month_key(f)
        m = next_month(m_prev)
        agg = (
            pl.scan_parquet(f)
            .filter(pl.col("rth_last_close").is_not_null())
            .group_by("ticker")
            .agg(
                pl.len().alias("days"),
                pl.col("rth_last_close").median().alias("med_close"),
                pl.col("bars_rth").median().alias("med_bars"),
                pl.col("rth_dollar_vol").sum().alias("dollar_vol"),
            )
            .with_columns(
                pl.lit(m).alias("month"),
                ((pl.col("days") >= MIN_DAYS)
                 & (pl.col("med_close") >= MIN_PRICE)
                 & (pl.col("med_bars") >= MIN_BARS)).alias("eligible"),
                NONCOMMON.alias("flag_noncommon"),
            )
            .with_columns(
                pl.col("dollar_vol")
                .rank("ordinal", descending=True)
                .over(pl.col("eligible"))  # rank within eligible group
                .alias("rank_dv"),
            )
            .with_columns([
                (pl.col("eligible") & (pl.col("rank_dv") <= n)).alias(f"top{n}")
                for n in (500, 1000, 3000)
            ])
            .collect()
        )
        frames.append(agg)
    out = pl.concat(frames).sort("month", "rank_dv")
    dest = ROOT / "data" / "metadata" / "universe_monthly.parquet"
    out.write_parquet(dest)
    last = out.filter(pl.col("month") == out["month"].max())
    print(f"months={out['month'].n_unique()} rows={out.height}")
    print(f"last month: eligible={last['eligible'].sum()} "
          f"top500={last['top500'].sum()} noncommon={last['flag_noncommon'].sum()}")


if __name__ == "__main__":
    main()