#!/usr/bin/env python
"""Phase 2d/3a: cohorts (event-time per horizon + fixed calendar), PIT joins, exclusion cascade, base rates."""
from __future__ import annotations

import datetime as dt
import json
import sys
from pathlib import Path

import polars as pl

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from wp39.pipeline import Paths, get_logger, load_config  # noqa: E402
from wp39.provenance import now_iso, write_json  # noqa: E402
from wp39.samples import (apply_admissibility, calendar_cutoffs, event_cutoffs, exclusion_cascade, join_fundamentals,  # noqa: E402
                          join_market)


def hist_features(sample: pl.DataFrame, quarters: pl.DataFrame) -> pl.DataFrame:
    """Information age, last available surprise, mean of last 4 surprises, trailing adverse rate — all from quarters
    available at the cutoff (pit_qidx and earlier)."""
    qq = quarters.select(["entity_id", "qidx", "S_opinc", "y_adverse_opinc", "A_q"]).sort(["entity_id", "qidx"])
    qq = qq.with_columns([
        pl.col("S_opinc").alias("last_surprise"),
        pl.col("S_opinc").rolling_mean(window_size=4, min_samples=2).over("entity_id").alias("mean_surprise4"),
        pl.col("y_adverse_opinc").cast(pl.Float64).rolling_mean(window_size=8, min_samples=2).over("entity_id").alias("hist_adverse_rate"),
        pl.col("A_q").alias("A_pit"),
    ]).select(["entity_id", pl.col("qidx").alias("pit_qidx"), "last_surprise", "mean_surprise4", "hist_adverse_rate", "A_pit"])
    s = sample.join(qq, on=["entity_id", "pit_qidx"], how="left")
    return s.with_columns((pl.col("cutoff") - pl.col("A_pit")).dt.total_days().alias("info_age_td"))


def build(design: str, cfg: dict, P: Paths, quarters: pl.DataFrame, windows: pl.DataFrame, sym: pl.DataFrame, log) -> pl.DataFrame:
    if design == "event":
        cut = event_cutoffs(quarters, cfg["cutoffs"]["event_time_horizons_td"])
    else:
        s0, s1 = dt.date.fromisoformat(str(cfg["sample"]["start"])), dt.date.fromisoformat(str(cfg["sample"]["end"]))
        cut = calendar_cutoffs(quarters, s0, s1)
    cut = apply_admissibility(cut)
    log.info("%s cutoffs after admissibility: %d", design, cut.height)
    j = join_market(cut, windows, sym)
    j = join_fundamentals(j, quarters)
    # attach the outcome of quarter q (the row keyed by entity/fiscal_period_end)
    ycols = ["S_opinc", "S_netinc", "S_ocf", "S_rev", "ocf_innov", "roa_yoy", "y_adverse_opinc", "y_adverse_netinc", "y_roa_decline",
             "y_ocf_negative", "y_liquidity_deterioration", "y_eps_miss", "y_abs_surprise", "assets_tot_lag1", "_gap_ok8",
             "release_time", "eps_estimate"]
    j = j.join(quarters.select(["entity_id", "fiscal_period_end"] + ycols), on=["entity_id", "fiscal_period_end"], how="left")
    j = hist_features(j, quarters)
    sample, cascade = exclusion_cascade(j, cfg["sample"]["min_price_usd"], cfg["sample"]["min_valid_days_in_window"])
    cascade.write_csv(P.results / f"table_sample_cascade_{design}.csv")
    log.info("%s sample: %d rows; cascade=%s", design, sample.height, cascade.to_dicts())
    return sample


def main() -> int:
    P = Paths.default()
    log = get_logger("04_labels", P.logs / "04_labels.log")
    cfg = load_config("research")
    quarters = pl.read_parquet(P.processed / "quarters.parquet")
    uni = pl.read_parquet(P.processed / "universe.parquet").filter(pl.col("eligible"))
    sym = uni.select(["entity_id", "ticker", "sic_division", "sic"]).rename({"ticker": "symbol"})
    wins = [p for p in (P.cache / "windows").glob("*.parquet") if p.stem in set(sym["symbol"].to_list())]
    windows = pl.concat([pl.read_parquet(p) for p in wins], how="diagonal_relaxed")
    log.info("windows: %d rows, %d symbols", windows.height, len(wins))
    quarters = quarters.filter(pl.col("entity_id").is_in(sym["entity_id"].to_list()))
    out = {}
    for design in ("event", "calendar"):
        s = build(design, cfg, P, quarters, windows, sym, log)
        s.write_parquet(P.processed / f"sample_{design}.parquet")
        out[design] = {
            "rows": s.height, "firms": s["entity_id"].n_unique(), "quarters": s.select(["entity_id", "fiscal_period_end"]).unique().height,
            "by_h": s.group_by("h").agg([pl.len().alias("n"), pl.col("y_adverse_opinc").mean().alias("base_rate")]).sort("h").to_dicts() if design == "event" else None,
            "base_rates": {y: float(s[y].mean()) for y in ("y_adverse_opinc", "y_adverse_netinc", "y_roa_decline", "y_ocf_negative",
                                                            "y_liquidity_deterioration", "y_abs_surprise")},
            "years": s.with_columns(pl.col("cutoff").dt.year().alias("y")).group_by("y").len().sort("y").to_dicts(),
        }
    write_json(P.results / "sample_audit.json", {"built_at": now_iso(), **out})
    print(json.dumps(out, indent=1, default=str))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())