"""Firm-quarter outcomes: reveal dates, seasonal-random-walk expectations and adverse-outcome labels (PIT_PROTOCOL §4).

Nothing here uses prices. The expectation E_c[X_q] uses quarters q-1..q-8 *restricted to those available at the cutoff*;
samples.py passes the cutoff-specific availability mask. The quarterly series X_q is the first (only) revision.
"""
from __future__ import annotations

import datetime as dt

import polars as pl

TARGET_VARS = ["op_inc", "net_inc", "cash_op", "rev"]


def match_release_dates(wide: pl.DataFrame, events: pl.DataFrame, max_lag_days: int = 30) -> pl.DataFrame:
    """Attach A_q (first release in [E_q, F_q + 10 d]) and its timing to each firm-quarter; flag exclusions."""
    e = events.select(["entity_id", pl.col("event_date").cast(pl.Date), "release_time", "eps_estimate", "eps_reported"]).sort(
        ["entity_id", "event_date"])
    w = wide.select(["entity_id", "fiscal_period_end", "filing_date"]).with_columns(
        pl.col("fiscal_period_end").alias("_key")).sort(["entity_id", "_key"])
    # forward as-of: first event on/after the fiscal period end
    j = w.join_asof(e, left_on="_key", right_on="event_date", by="entity_id", strategy="forward")
    j = j.with_columns([
        pl.col("event_date").alias("A_q"),
        (pl.col("event_date") <= pl.col("filing_date") + pl.duration(days=10)).alias("_within"),
    ])
    j = j.with_columns([
        pl.when(pl.col("_within")).then(pl.col("A_q")).otherwise(None).alias("A_q"),
        pl.when(pl.col("_within")).then(pl.col("release_time")).otherwise(None).alias("release_time"),
        pl.when(pl.col("_within")).then(pl.col("eps_estimate")).otherwise(None).alias("eps_estimate"),
        pl.when(pl.col("_within")).then(pl.col("eps_reported")).otherwise(None).alias("eps_reported"),
    ])
    j = j.with_columns([
        (pl.col("filing_date") - pl.col("A_q")).dt.total_days().alias("filing_minus_release_days"),
    ])
    j = j.with_columns([
        pl.when(pl.col("A_q").is_null()).then(pl.lit("no_release"))
        .when(pl.col("filing_minus_release_days") < 0).then(pl.lit("filing_before_release"))
        .when(pl.col("filing_minus_release_days") > max_lag_days).then(pl.lit("release_lag_gt_30d"))
        .otherwise(pl.lit("ok")).alias("release_match"),
    ])
    return wide.join(j.select(["entity_id", "fiscal_period_end", "A_q", "release_time", "eps_estimate", "eps_reported",
                               "filing_minus_release_days", "release_match"]), on=["entity_id", "fiscal_period_end"], how="left")


def seasonal_expectation(w: pl.DataFrame, var: str) -> pl.DataFrame:
    """E[X_q] = X_{q-4} + mean_{k=0..3}(X_{q-1-k} - X_{q-5-k}); requires q-1..q-8 (by quarter index, regular calendar)."""
    x = pl.col(var)
    lags = {k: x.shift(k).over("entity_id") for k in range(1, 9)}
    drift = sum((lags[1 + k] - lags[5 + k]) for k in range(4)) / 4.0
    return w.with_columns([
        (lags[4] + drift).alias(f"E_{var}"),
        lags[4].alias(f"{var}_lag4"),
        lags[1].alias(f"{var}_lag1"),
        (x - (lags[4] + drift)).alias(f"surprise_{var}"),
    ])


def build_quarter_table(wide: pl.DataFrame) -> pl.DataFrame:
    """Quarter-level target ingredients (unscaled). Regularity: consecutive quarters 60-120 days apart."""
    w = wide.sort(["entity_id", "fiscal_period_end"])
    w = w.with_columns(pl.col("fiscal_period_end").diff().over("entity_id").dt.total_days().alias("days_since_prev_q"))
    # a quarter is "regular" if the previous 8 gaps are all within [60, 120]
    ok = (pl.col("days_since_prev_q") >= 60) & (pl.col("days_since_prev_q") <= 120)
    w = w.with_columns(ok.cast(pl.Int8).alias("_gap_ok"))
    w = w.with_columns(pl.col("_gap_ok").rolling_sum(window_size=8, min_samples=8).over("entity_id").alias("_gap_ok8"))
    for v in TARGET_VARS:
        w = seasonal_expectation(w, v)
    # secondary ingredients
    w = w.with_columns([
        pl.col("roa").shift(4).over("entity_id").alias("roa_lag4"),
        (pl.col("assets_c") / (pl.col("liab_c").abs() + 1e-9)).alias("current_ratio"),
    ])
    w = w.with_columns([
        pl.col("current_ratio").shift(4).over("entity_id").alias("current_ratio_lag4"),
        pl.col("intr_cover").shift(4).over("entity_id").alias("intr_cover_lag4"),
        pl.col("assets_tot").shift(1).over("entity_id").alias("assets_tot_lag1"),
    ])
    return w


def finalize_labels(q: pl.DataFrame, threshold: float = -0.005) -> pl.DataFrame:
    """Scale by the latest *available* total assets (lag-1 quarter) and build binary targets."""
    ta = pl.col("assets_tot_lag1").abs() + 1e-9
    q = q.with_columns([
        (pl.col("surprise_op_inc") / ta).alias("S_opinc"),
        (pl.col("surprise_net_inc") / ta).alias("S_netinc"),
        (pl.col("surprise_cash_op") / ta).alias("S_ocf"),
        (pl.col("surprise_rev") / ta).alias("S_rev"),
        ((pl.col("cash_op") - pl.col("cash_op_lag4")) / ta).alias("ocf_innov"),
        (pl.col("roa") - pl.col("roa_lag4")).alias("roa_yoy"),
    ])
    q = q.with_columns([
        (pl.col("S_opinc") <= threshold).cast(pl.Int8).alias("y_adverse_opinc"),
        (pl.col("S_netinc") <= threshold).cast(pl.Int8).alias("y_adverse_netinc"),
        (pl.col("roa_yoy") <= -0.01).cast(pl.Int8).alias("y_roa_decline"),
        (pl.col("ocf_innov") <= -0.01).cast(pl.Int8).alias("y_ocf_negative"),
        (((pl.col("current_ratio") / (pl.col("current_ratio_lag4") + 1e-9)) <= 0.9) |
         ((pl.col("intr_cover") < 1.0) & (pl.col("intr_cover_lag4") >= 1.0))).cast(pl.Int8).alias("y_liquidity_deterioration"),
        pl.when(pl.col("eps_estimate").is_not_null() & pl.col("eps_reported").is_not_null())
        .then((pl.col("eps_reported") < pl.col("eps_estimate")).cast(pl.Int8)).otherwise(None).alias("y_eps_miss"),
        (pl.col("S_opinc").abs() >= 0.005).cast(pl.Int8).alias("y_abs_surprise"),  # symmetric placebo target (volatility channel)
    ])
    return q