"""Stage A pilot audit of mito0o852/OHLCV-1m monthly files.

Audits schema, timestamp convention, session coverage, duplicates,
OHLC validity, zeros/negatives, and symbol counts for the three pilot
months. Writes one JSON diagnostic per file to outputs/diagnostics/.

Usage:  python scripts/01_pilot_audit.py
"""

from __future__ import annotations

import json
from pathlib import Path

import polars as pl

ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "data" / "raw" / "data"
OUT = ROOT / "outputs" / "diagnostics"

PILOT_MONTHS = ["1992-06", "2009-06", "2025-12"]


def audit_month(month: str) -> dict:
    path = RAW / f"ohlcv_{month}.parquet"
    lf = pl.scan_parquet(path)
    schema = lf.collect_schema()
    df = lf.collect()
    n = df.height

    ts_utc = df["timestamp"]
    ts_et = ts_utc.dt.convert_time_zone("America/New_York")
    et = df.with_columns(
        ts_et.alias("ts_et"),
        ts_et.dt.hour().cast(pl.Int32).alias("hour_et"),
        ts_et.dt.minute().cast(pl.Int32).alias("minute_et"),
        ts_et.dt.weekday().alias("weekday"),  # 1=Mon .. 7=Sun
        ts_et.dt.date().alias("date_et"),
        ts_utc.dt.second().alias("second"),
    )
    minute_of_day = et.select(
        (pl.col("hour_et") * 60 + pl.col("minute_et")).alias("mod")
    )["mod"]

    rth = ((minute_of_day >= 570) & (minute_of_day < 960)).sum()  # 09:30-16:00 ET
    pre = (minute_of_day < 570).sum()
    post = (minute_of_day >= 960).sum()

    dup = n - df.select(["ticker", "timestamp"]).unique().height

    viol = df.filter(
        (pl.col("low") > pl.min_horizontal("open", "close"))
        | (pl.col("high") < pl.max_horizontal("open", "close"))
        | (pl.col("low") > pl.col("high"))
    ).height

    nonpos_price = df.filter(
        (pl.col("open") <= 0) | (pl.col("high") <= 0)
        | (pl.col("low") <= 0) | (pl.col("close") <= 0)
    ).height
    neg_vol = df.filter(pl.col("volume") < 0).height
    zero_vol = df.filter(pl.col("volume") == 0).height
    nonint_vol = df.filter(pl.col("volume") != pl.col("volume").round(0)).height

    nulls = {c: int(df[c].null_count()) for c in df.columns}
    weekend = et.filter(pl.col("weekday") >= 6).height
    nonzero_sec = (et["second"] != 0).sum()

    by_day = (
        et.group_by("ticker", "date_et")
        .agg(pl.len().alias("bars"))
        .group_by("date_et")
        .agg(
            pl.len().alias("symbols"),
            pl.col("bars").median().alias("median_bars"),
        )
        .sort("date_et")
    )

    mod_hist = (
        et.select((pl.col("hour_et") * 60 + pl.col("minute_et")).alias("mod"))
        .group_by("mod").agg(pl.len().alias("n")).sort("mod")
    )

    return {
        "month": month,
        "rows": n,
        "columns": {name: str(dtype) for name, dtype in schema.items()},
        "unique_tickers": df["ticker"].n_unique(),
        "min_ts_utc": str(ts_utc.min()),
        "max_ts_utc": str(ts_utc.max()),
        "nonzero_seconds": int(nonzero_sec),
        "rows_rth_0930_1600_et": int(rth),
        "rows_premarket_et": int(pre),
        "rows_afterhours_et": int(post),
        "rth_share": round(float(rth) / n, 4),
        "duplicate_ticker_timestamp": int(dup),
        "ohlc_violations": int(viol),
        "nonpositive_prices": int(nonpos_price),
        "negative_volume": int(neg_vol),
        "zero_volume": int(zero_vol),
        "noninteger_volume": int(nonint_vol),
        "null_counts": nulls,
        "weekend_rows": int(weekend),
        "trading_days_observed": by_day.height,
        "symbols_per_day_median": float(by_day["symbols"].median()),
        "median_bars_per_symbol_day_overall": float(
            et.group_by("ticker", "date_et").agg(pl.len().alias("b"))["b"].median()
        ),
        "price_summary_close": {
            "min": float(df["close"].min()),
            "p50": float(df["close"].median()),
            "max": float(df["close"].max()),
        },
        "minute_of_day_hist": {
            str(r["mod"]): r["n"] for r in mod_hist.iter_rows(named=True)
        },
        "per_day": [
            {"date": str(r["date_et"]), "symbols": r["symbols"],
             "median_bars": r["median_bars"]}
            for r in by_day.iter_rows(named=True)
        ],
    }


def main() -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    for month in PILOT_MONTHS:
        rep = audit_month(month)
        out = OUT / f"pilot_audit_{month}.json"
        out.write_text(json.dumps(rep, indent=1))
        keep = {k: v for k, v in rep.items()
                if k not in ("minute_of_day_hist", "per_day")}
        print(json.dumps(keep, indent=1))


if __name__ == "__main__":
    main()