"""Stage D.3 — clean minute panel, partitioned by year/month.

Per month:
  * duplicate (ticker, minute) keys resolved by the logged rule: open=first,
    high=max, low=min, close=last, volume=sum (DECISION_LOG 2026-07-13);
  * ET fields and session label using the XNYS calendar with TRUE close times
    (early closes truncate the regular session; bars after an early close are
    labelled 'post', cf. DATA_AUDIT §4);
  * integrity flag (OHLC violation) and split-window flag (±1 day around
    high-confidence split events) — flagged, never deleted;
  * written to data/processed/minute/year=YYYY/month=MM/data.parquet (zstd).

The extended-hours "panel" is the same file filtered on session != 'rth' — one
store, both baselines. Idempotent per month; --force to rebuild.

Usage: python scripts/08_build_clean_minute_panel.py [--workers N] [--month YYYY-MM]
"""

from __future__ import annotations

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

import polars as pl

ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "data" / "raw" / "data"
OUT = ROOT / "data" / "processed" / "minute"
SPLITS = ROOT / "outputs" / "diagnostics" / "stage_d" / "split_windows.parquet"
LOG = ROOT / "logs" / "preprocessing"

RTH_START = 570


def session_close_map() -> pl.DataFrame:
    """date -> regular close as minute-of-day ET (960, or 780 on early closes)."""
    import exchange_calendars as xcals
    cal = xcals.get_calendar("XNYS", start="1992-01-01", end="2026-04-01")
    closes = cal.closes.dt.tz_convert("America/New_York")
    return pl.DataFrame({
        "date": [d.date() for d in closes.index],
        "close_mod": [c.hour * 60 + c.minute for c in closes],
    }).with_columns(pl.col("date").cast(pl.Date))


def build_month(month: str, force: bool = False) -> dict:
    year, mm = month.split("-")
    dest_dir = OUT / f"year={year}" / f"month={mm}"
    dest = dest_dir / "data.parquet"
    if dest.exists() and not force:
        return {"month": month, "skipped": True}

    df = pl.read_parquet(RAW / f"ohlcv_{month}.parquet").sort("ticker", "timestamp")
    n_raw = df.height

    df = (
        df.group_by("ticker", "timestamp", maintain_order=True)
        .agg(
            pl.col("open").first(),
            pl.col("high").max(),
            pl.col("low").min(),
            pl.col("close").last(),
            pl.col("volume").sum(),
        )
    )
    n_dedup = df.height

    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"),
    ).drop("ts_et")

    et = et.join(session_close_map(), on="date", how="left")
    et = et.with_columns(
        pl.when(pl.col("mod") < RTH_START).then(pl.lit("pre"))
        .when(pl.col("mod") < pl.col("close_mod")).then(pl.lit("rth"))
        .otherwise(pl.lit("post")).alias("session"),
        (
            (pl.col("low") > pl.min_horizontal("open", "close"))
            | (pl.col("high") < pl.max_horizontal("open", "close"))
            | (pl.col("low") > pl.col("high"))
        ).alias("flag_ohlc"),
    ).drop("close_mod")

    win = pl.read_parquet(SPLITS).with_columns(pl.lit(True).alias("flag_split_window"))
    et = et.join(win, on=["ticker", "date"], how="left").with_columns(
        pl.col("flag_split_window").fill_null(False))

    dest_dir.mkdir(parents=True, exist_ok=True)
    tmp = dest_dir / ".data.parquet.tmp"
    et.sort("ticker", "timestamp").write_parquet(tmp, compression="zstd")
    tmp.replace(dest)

    rep = {
        "month": month, "rows_raw": n_raw, "rows": n_dedup,
        "dup_rows_merged": n_raw - n_dedup,
        "rth": int((et["session"] == "rth").sum()),
        "flag_ohlc": int(et["flag_ohlc"].sum()),
        "flag_split_window": int(et["flag_split_window"].sum()),
    }
    LOG.mkdir(parents=True, exist_ok=True)
    with open(LOG / "08_panel_reports.jsonl", "a") as f:
        f.write(json.dumps(rep) + "\n")
    return rep


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 = ([args.month] if args.month else
              sorted(p.stem.replace("ohlcv_", "") for p in RAW.glob("ohlcv_*.parquet")))
    with ProcessPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(build_month, m, args.force): m for m in months}
        done = 0
        for fut in as_completed(futs):
            rep = fut.result()
            done += 1
            print(f"[{done}/{len(months)}] {rep['month']}"
                  + (" skip" if rep.get("skipped") else
                     f" rows={rep['rows']:,} merged={rep['dup_rows_merged']}"),
                  flush=True)
    print("clean minute panel complete")


if __name__ == "__main__":
    main()