"""Negative controls, placebos and subgroup stress tests (Phase 5)."""
from __future__ import annotations

import numpy as np
import polars as pl


def shuffle_labels_within_quarter(df: pl.DataFrame, y: str, seed: int) -> pl.DataFrame:
    """Permute the label within calendar quarter of the cutoff: any 'gain' left is overfitting/leakage."""
    rng = np.random.default_rng(seed)
    q = df.with_columns((pl.col("cutoff").dt.year() * 4 + (pl.col("cutoff").dt.month() - 1) // 3).alias("_cq"))
    parts = []
    for _, g in q.group_by("_cq", maintain_order=True):
        v = g[y].to_numpy().copy(); rng.shuffle(v)
        parts.append(g.with_columns(pl.Series(y, v)))
    return pl.concat(parts).drop("_cq")


def subgroup_masks(df: pl.DataFrame) -> dict[str, pl.Expr]:
    """Pre-specified subgroups (terciles are computed on the TRAINING rows and passed as cut-points by the caller)."""
    return {
        "all": pl.lit(True),
        "size_small": pl.col("tercile_size") == 1, "size_large": pl.col("tercile_size") == 3,
        "illiq_high": pl.col("tercile_illiq") == 3, "illiq_low": pl.col("tercile_illiq") == 1,
        "lev_high": pl.col("tercile_lev") == 3, "lev_low": pl.col("tercile_lev") == 1,
        "loss_firms": pl.col("f_loss") == 1, "profit_firms": pl.col("f_loss") == 0,
        "covid_2020": (pl.col("cutoff").dt.year() == 2020), "ex_covid": (pl.col("cutoff").dt.year() != 2020),
        "after_close_release": pl.col("release_time") == "after_close", "before_open_release": pl.col("release_time") == "before_open",
        "sector_finance": pl.col("sic_division") == "H_finance_insurance_realestate",
        "sector_nonfinance": pl.col("sic_division") != "H_finance_insurance_realestate",
    }


def add_terciles(train: pl.DataFrame, others: list[pl.DataFrame]) -> list[pl.DataFrame]:
    """Tercile cut-points for size, illiquidity and leverage computed on training rows only."""
    cuts = {}
    for name, col in (("size", "f_log_mcap"), ("illiq", "wd_amihud60"), ("lev", "f_leverage")):
        s = train[col].drop_nulls().drop_nans()
        cuts[name] = (float(s.quantile(1 / 3)), float(s.quantile(2 / 3)))
    out = []
    for d in [train] + others:
        for name, col in (("size", "f_log_mcap"), ("illiq", "wd_amihud60"), ("lev", "f_leverage")):
            lo, hi = cuts[name]
            d = d.with_columns(pl.when(pl.col(col) <= lo).then(1).when(pl.col(col) <= hi).then(2).otherwise(3).alias(f"tercile_{name}"))
        out.append(d)
    return out