import datetime as dt

import polars as pl

from wp39.labels import build_quarter_table, finalize_labels, match_release_dates


def test_release_match_rules(toy_wide, toy_events):
    ev = toy_events.with_columns(pl.when(pl.col("entity_id") == "frd:BBB").then(pl.col("event_date") - pl.duration(days=60))
                                 .otherwise(pl.col("event_date")).alias("event_date"))
    w = match_release_dates(toy_wide, ev)
    assert (w.filter(pl.col("entity_id") == "frd:AAA")["release_match"] == "ok").all()
    assert (w.filter(pl.col("entity_id") == "frd:BBB")["release_match"] != "ok").all()


def test_labels_are_not_price_based(toy_wide, toy_events):
    q = finalize_labels(build_quarter_table(match_release_dates(toy_wide, toy_events)))
    assert not any(c.startswith("w") or c.startswith("d_") for c in q.columns)
    assert set(q["y_adverse_opinc"].drop_nulls().unique().to_list()) <= {0, 1}


def test_release_day_excluded(toy_wide, toy_events):
    from wp39.samples import apply_admissibility, event_cutoffs
    w = match_release_dates(toy_wide, toy_events)
    c = apply_admissibility(event_cutoffs(w, [1]))
    assert (c["cutoff"] < c["A_q"]).all()