"""Cohort construction: cutoffs, window aggregation of session features, PIT fundamentals join, exclusions.

Event-time: c = A_q - h trading days (h in horizons). Fixed-calendar: c = month-end sessions, target = next release.
All joins are explicit as-of joins on `session_date <= c` (bars) and `available_date <= c` (fundamentals).
"""
from __future__ import annotations

import datetime as dt

import numpy as np
import polars as pl

from .calendar import session_array, shift_trading_days, month_end_sessions
from .features import FREQS, INTRADAY_FAMILIES
from .pit_fundamentals import METRICS, pit_feature_block

WINDOW_MEAN = ["rv", "bpv", "rj", "rsneg_share", "sjv", "rskew", "rkurt", "max_share", "ac1", "vol_first30", "vol_last30",
               "vol_gini", "down_vol_share", "present_share", "rv_mktdown_share"]


def complete_session_index(feat: pl.DataFrame) -> pl.DataFrame:
    """Reindex a symbol's session features to every XNYS session between its first and last session (rows for missing
    sessions have null features) so that rolling windows are measured in trading days."""
    arr = session_array()
    d0, d1 = feat["session_date"].min(), feat["session_date"].max()
    s = arr[(arr >= np.datetime64(d0)) & (arr <= np.datetime64(d1))]
    idx = pl.DataFrame({"session_date": pl.Series(s.astype("datetime64[D]")).cast(pl.Date)})
    return idx.join(feat, on="session_date", how="left").sort("session_date")


def rolling_window_features(feat: pl.DataFrame, short: int = 20, long: int = 60, min_valid: int = 15) -> pl.DataFrame:
    """Rolling (trailing, inclusive) aggregates per session; the row at session c is the feature vector for cutoff c."""
    f = complete_session_index(feat)
    exprs = []
    valid = pl.col("f1m_n_bars").is_not_null() & (pl.col("f1m_n_bars") >= 100)
    f = f.with_columns(valid.cast(pl.Int8).alias("_valid"))
    exprs.append(pl.col("_valid").rolling_sum(window_size=short, min_samples=1).alias("w_valid_days"))
    # intraday families: window means (RV-type quantities averaged -> per-day average, same across frequencies)
    for name in FREQS:
        p = f"f{name}_"
        for fam in WINDOW_MEAN:
            exprs.append(pl.col(p + fam).rolling_mean(window_size=short, min_samples=min_valid).alias(f"w{name}_{fam}"))
        exprs.append(pl.col(p + "drift").rolling_sum(window_size=short, min_samples=min_valid).alias(f"w{name}_drift"))
        exprs.append((pl.col(p + "sxy").rolling_sum(window_size=short, min_samples=min_valid) /
                      (pl.col(p + "syy").rolling_sum(window_size=short, min_samples=min_valid) + 1e-18)).alias(f"w{name}_beta"))
        exprs.append(pl.col(p + "rv").rolling_std(window_size=short, min_samples=min_valid).alias(f"w{name}_volofvol"))
        # variance ratio: mean intraday RV vs variance of daily close-to-close returns over the window
        exprs.append((pl.col(p + "rv").rolling_mean(window_size=short, min_samples=min_valid) /
                      (pl.col("d_ret_cc").rolling_var(window_size=short, min_samples=min_valid) + 1e-18)).alias(f"w{name}_vratio"))
    # daily block
    r = pl.col("d_ret_cc")
    exprs += [
        ((1.0 + r).log().rolling_sum(window_size=short, min_samples=min_valid)).alias("wd_ret20"),
        ((1.0 + r).log().rolling_sum(window_size=long, min_samples=3 * min_valid)).alias("wd_ret60"),
        r.rolling_std(window_size=short, min_samples=min_valid).alias("wd_vol20"),
        r.rolling_std(window_size=long, min_samples=3 * min_valid).alias("wd_vol60"),
        pl.col("d_parkinson").rolling_mean(window_size=short, min_samples=min_valid).sqrt().alias("wd_parkinson20"),
        r.rolling_skew(window_size=long).alias("wd_skew60"),
        r.rolling_kurtosis(window_size=long).alias("wd_kurt60"),
        (pl.col("d_ret_co").rolling_sum(window_size=short, min_samples=min_valid)).alias("wd_overnight20"),
        (pl.col("d_ret_oc").rolling_sum(window_size=short, min_samples=min_valid)).alias("wd_intraday20"),
        (pl.col("d_volume").rolling_mean(window_size=short, min_samples=min_valid) /
         (pl.col("d_volume").rolling_median(window_size=long, min_samples=3 * min_valid) + 1.0)).alias("wd_volratio20"),
        (pl.col("d_volume").rolling_mean(window_size=5, min_samples=3) /
         (pl.col("d_volume").rolling_median(window_size=long, min_samples=3 * min_valid) + 1.0)).alias("wd_volratio5"),
        ((r.abs() / (pl.col("d_dollar_vol") + 1.0)) * 1e6).rolling_mean(window_size=long, min_samples=3 * min_valid).alias("wd_amihud60"),
        (r == 0).cast(pl.Float64).rolling_mean(window_size=long, min_samples=3 * min_valid).alias("wd_zeroret60"),
        pl.col("d_dollar_vol").rolling_mean(window_size=long, min_samples=3 * min_valid).log1p().alias("wd_log_dollar_vol60"),
        pl.col("d_close").alias("wd_close"),
        pl.col("f1m_n_bars").rolling_mean(window_size=short, min_samples=min_valid).alias("wd_bars20"),
    ]
    f = f.with_columns(exprs)
    # max drawdown over 60 sessions from cumulative log returns
    f = f.with_columns((1.0 + r).log().fill_null(0.0).cum_sum().alias("_cum"))
    f = f.with_columns((pl.col("_cum") - pl.col("_cum").rolling_max(window_size=long, min_samples=3 * min_valid)).rolling_min(
        window_size=long, min_samples=3 * min_valid).alias("wd_maxdd60"))
    # daily beta / residual vol vs SPY daily (from 1m family sums at daily horizon are not available -> use 5m beta proxy)
    cols = ["session_date", "symbol", "w_valid_days"] + [c for c in f.columns if c.startswith("w1m_") or c.startswith("w5m_")
                                                           or c.startswith("w30m_") or c.startswith("wd_")]
    return f.select(cols)


def event_cutoffs(quarters: pl.DataFrame, horizons: list[int]) -> pl.DataFrame:
    """One row per (entity, quarter, h) with cutoff c = A_q - h trading days (requires release_match == ok)."""
    q = quarters.filter(pl.col("release_match") == "ok").select(
        ["entity_id", "fiscal_period_end", "A_q", "filing_date", "available_date", "qidx"])
    rows = []
    a_prev = q.with_columns(pl.col("A_q").shift(1).over("entity_id").alias("A_prev"))
    for r in a_prev.iter_rows(named=True):
        for h in horizons:
            try:
                c = shift_trading_days(r["A_q"], -h)
            except ValueError:
                continue
            rows.append({"entity_id": r["entity_id"], "fiscal_period_end": r["fiscal_period_end"], "h": h, "cutoff": c,
                         "A_q": r["A_q"], "F_q": r["filing_date"], "A_prev": r["A_prev"], "qidx": r["qidx"]})
    return pl.DataFrame(rows)


def calendar_cutoffs(quarters: pl.DataFrame, start: dt.date, end: dt.date, max_td: int = 65) -> pl.DataFrame:
    """Month-end cutoffs; target = next release after c within max_td trading days."""
    me = pl.DataFrame({"cutoff": month_end_sessions(start, end)})
    q = quarters.filter(pl.col("release_match") == "ok").select(["entity_id", "fiscal_period_end", "A_q", "filing_date", "qidx"])
    ents = q.select("entity_id").unique()
    grid = ents.join(me, how="cross").sort(["entity_id", "cutoff"])
    qq = q.sort(["entity_id", "A_q"]).with_columns(pl.col("A_q").alias("_a"))
    j = grid.join_asof(qq, left_on="cutoff", right_on="_a", by="entity_id", strategy="forward")
    j = j.filter(pl.col("A_q") > pl.col("cutoff"))
    arr = session_array()
    c = np.searchsorted(arr, j["cutoff"].to_numpy().astype("datetime64[D]"), side="right")
    a = np.searchsorted(arr, j["A_q"].to_numpy().astype("datetime64[D]"), side="right")
    j = j.with_columns(pl.Series("h", (a - c).astype(int))).filter(pl.col("h") <= max_td)
    return j.rename({"filing_date": "F_q"}).with_columns(pl.lit(None, dtype=pl.Date).alias("A_prev"))


def apply_admissibility(cut: pl.DataFrame) -> pl.DataFrame:
    """c < min(A_q, F_q) and c > A_prev (previous reveal already happened)."""
    return cut.filter((pl.col("cutoff") < pl.col("A_q")) & (pl.col("cutoff") < pl.col("F_q")) &
                      (pl.col("A_prev").is_null() | (pl.col("cutoff") > pl.col("A_prev"))))


def join_market(cut: pl.DataFrame, win: pl.DataFrame, symbol_of: pl.DataFrame) -> pl.DataFrame:
    """As-of join of the rolling window features at session_date <= cutoff (exact session normally)."""
    c = cut.join(symbol_of, on="entity_id", how="left").sort(["symbol", "cutoff"])
    w = win.sort(["symbol", "session_date"])
    j = c.join_asof(w, left_on="cutoff", right_on="session_date", by="symbol", strategy="backward", tolerance="5d")
    return j.rename({"session_date": "market_asof_date"})


def join_fundamentals(j: pl.DataFrame, wide: pl.DataFrame) -> pl.DataFrame:
    """Latest available quarter at cutoff + yoy growth from the quarter 4 back (also available by construction)."""
    w = wide.select(["entity_id", "fiscal_period_end", "available_date", "qidx", "A_q"] + METRICS).sort(["entity_id", "available_date"])
    w4 = wide.select(["entity_id", pl.col("qidx").alias("pit_qidx4"), pl.col("rev").alias("rev_l4"), pl.col("op_inc").alias("op_inc_l4"),
                      pl.col("shares").alias("shares_l4"), pl.col("available_date").alias("avail_l4")])
    wr = w.rename({c: c for c in METRICS}).rename({"fiscal_period_end": "pit_period_end", "available_date": "pit_available_date",
                                                   "qidx": "pit_qidx", "A_q": "pit_A"})
    out = j.sort(["entity_id", "cutoff"]).join_asof(wr, left_on="cutoff", right_on="pit_available_date", by="entity_id", strategy="backward")
    out = out.with_columns((pl.col("pit_qidx") - 4).alias("pit_qidx4")).join(w4, on=["entity_id", "pit_qidx4"], how="left")
    out = out.with_columns([
        (pl.col("pit_available_date") <= pl.col("cutoff")).alias("pit_ok"),
        ((pl.col("rev") - pl.col("rev_l4")) / (pl.col("rev_l4").abs() + 1e-9)).alias("f_rev_growth"),
        ((pl.col("shares") - pl.col("shares_l4")) / (pl.col("shares_l4").abs() + 1e-9)).alias("f_net_issuance"),
        (pl.col("qidx") - pl.col("pit_qidx")).alias("quarters_ahead"),  # 1 = we predict the quarter right after the latest available
    ])
    out = pit_feature_block(out)
    out = out.with_columns([
        (pl.col("wd_close") * pl.col("shares").abs()).log1p().alias("f_log_mcap"),
    ])
    return out


def exclusion_cascade(df: pl.DataFrame, min_price: float, min_valid_days: int) -> tuple[pl.DataFrame, pl.DataFrame]:
    """Apply frozen exclusions in order; return (sample, cascade table)."""
    steps = [
        ("cutoff_admissible", pl.lit(True)),
        ("market_window_available", pl.col("market_asof_date").is_not_null() & (pl.col("w_valid_days") >= min_valid_days)),
        ("price_ge_min", pl.col("wd_close") >= min_price),
        ("fundamentals_available", pl.col("pit_ok").fill_null(False) & pl.col("pit_qidx").is_not_null()),
        ("predicting_next_quarter", pl.col("quarters_ahead") == 1),
        ("expectation_defined", pl.col("S_opinc").is_not_null() & pl.col("assets_tot_lag1").is_not_null()),
        ("regular_fiscal_calendar", pl.col("_gap_ok8") == 8),
    ]
    rows, cur = [], df
    for name, cond in steps:
        before = cur.height
        cur = cur.filter(cond)
        rows.append({"step": name, "remaining": cur.height, "dropped": before - cur.height})
    return cur, pl.DataFrame(rows)


def feature_columns(df: pl.DataFrame) -> dict[str, list[str]]:
    fund = [c for c in df.columns if c.startswith("f_")] + ["info_age_td", "last_surprise", "mean_surprise4", "hist_adverse_rate"]
    daily = [c for c in df.columns if c.startswith("wd_") and c != "wd_close"]
    return {
        "fundamentals": [c for c in fund if c in df.columns],
        "daily": daily,
        "intraday_30m": [c for c in df.columns if c.startswith("w30m_")],
        "intraday_5m": [c for c in df.columns if c.startswith("w5m_")],
        "intraday_1m": [c for c in df.columns if c.startswith("w1m_")],
    }