"""Chronological, walk-forward and firm-held-out splits with purge and embargo (RESEARCH_DESIGN §6).

An observation is identified by (entity_id, fiscal_period_end, cutoff, A_q). The *target* is revealed at A_q, so a
training observation is purged if its A_q falls after (split_start - embargo); test observations start at split_start.
"""
from __future__ import annotations

import datetime as dt
import hashlib

import polars as pl


def _q_start(d: dt.date) -> dt.date:
    return dt.date(d.year, 3 * ((d.month - 1) // 3) + 1, 1)


def _shift_quarters(d: dt.date, n: int) -> dt.date:
    m = d.month - 1 + 3 * n
    return dt.date(d.year + m // 12, m % 12 + 1, 1)


def purge_train(df: pl.DataFrame, boundary: dt.date, embargo_quarters: int = 1) -> pl.DataFrame:
    """Keep training rows whose outcome was revealed before boundary - embargo (quarter-aligned)."""
    limit = _shift_quarters(_q_start(boundary), -embargo_quarters)
    return df.filter(pl.col("A_q") < limit)


def chronological(df: pl.DataFrame, cfg: dict) -> dict[str, pl.DataFrame]:
    tr0, tr1 = [dt.date.fromisoformat(str(x)) for x in cfg["train"]]
    va0, va1 = [dt.date.fromisoformat(str(x)) for x in cfg["validation"]]
    te0, te1 = [dt.date.fromisoformat(str(x)) for x in cfg["test"]]
    emb = cfg["walk_forward"]["embargo_quarters"]
    train = purge_train(df.filter((pl.col("cutoff") >= tr0) & (pl.col("cutoff") <= tr1)), va0, emb)
    valid = purge_train(df.filter((pl.col("cutoff") >= va0) & (pl.col("cutoff") <= va1)), te0, emb)
    test = df.filter((pl.col("cutoff") >= te0) & (pl.col("cutoff") <= te1))
    return {"train": train, "validation": valid, "test": test}


def walk_forward(df: pl.DataFrame, cfg: dict, last_year: int) -> list[tuple[int, pl.DataFrame, pl.DataFrame]]:
    """Expanding annual folds: train on cutoffs < Jan 1 of year y (purged, embargo), test on cutoffs in year y."""
    out = []
    emb = cfg["walk_forward"]["embargo_quarters"]
    for y in range(cfg["walk_forward"]["first_test_year"], last_year + 1):
        b = dt.date(y, 1, 1)
        train = purge_train(df.filter(pl.col("cutoff") < b), b, emb)
        test = df.filter((pl.col("cutoff") >= b) & (pl.col("cutoff") < dt.date(y + 1, 1, 1)))
        if train.height and test.height:
            out.append((y, train, test))
    return out


def firm_holdout_mask(entity_ids: pl.Series, fraction: float = 0.2, salt: str = "wp39") -> pl.Series:
    def h(e: str) -> bool:
        return int(hashlib.sha256(f"{salt}:{e}".encode()).hexdigest()[:8], 16) / 0xFFFFFFFF < fraction
    return pl.Series([h(e) for e in entity_ids.to_list()], dtype=pl.Boolean)


def overlap_check(train: pl.DataFrame, test: pl.DataFrame) -> dict:
    """Diagnostics used by tests: no training outcome revealed on/after the first test cutoff; no shared observation ids."""
    first_test_cutoff = test["cutoff"].min()
    leak = train.filter(pl.col("A_q") >= first_test_cutoff).height
    key = ["entity_id", "fiscal_period_end", "h"] if "h" in train.columns else ["entity_id", "fiscal_period_end"]
    shared = train.select(key).join(test.select(key), on=key, how="inner").height
    return {"train_outcomes_revealed_in_test": leak, "shared_observations": shared}