import datetime as dt

import polars as pl

from wp39.splits import chronological, firm_holdout_mask, overlap_check, walk_forward

CFG = {"train": ["2012-01-01", "2018-12-31"], "validation": ["2019-01-01", "2020-12-31"], "test": ["2021-01-01", "2026-10-30"],
       "walk_forward": {"first_test_year": 2016, "embargo_quarters": 1}, "firm_holdout_fraction": 0.2}


def _panel():
    rows = []
    d = dt.date(2012, 1, 15)
    for i in range(140):
        c = d + dt.timedelta(days=int(i * 36.5))
        rows.append({"entity_id": f"frd:E{i % 7}", "fiscal_period_end": c - dt.timedelta(days=30), "h": 20, "cutoff": c,
                     "A_q": c + dt.timedelta(days=28), "y": i % 3 == 0})
    return pl.DataFrame(rows)


def test_chronological_purge_and_embargo():
    s = chronological(_panel(), CFG)
    assert s["train"]["A_q"].max() < dt.date(2018, 10, 1)       # validation starts 2019-01-01, embargo one quarter
    assert s["validation"]["A_q"].max() < dt.date(2020, 10, 1)
    assert overlap_check(s["train"], s["test"])["shared_observations"] == 0


def test_walk_forward_folds():
    folds = walk_forward(_panel(), CFG, 2025)
    for y, tr, te in folds:
        assert tr["A_q"].max() < dt.date(y - 1, 10, 1)
        assert te["cutoff"].min().year == y


def test_firm_holdout_is_deterministic():
    e = pl.Series(["frd:A", "frd:B", "frd:C"])
    assert firm_holdout_mask(e).to_list() == firm_holdout_mask(e).to_list()