import datetime as dt

import polars as pl

from wp39.calendar import close_time_utc, regular_session_bounds_utc, shift_trading_days, trading_days_between
from wp39.features import session_bounds_table


def test_regular_session_bounds_in_both_dst_regimes():
    o, c = regular_session_bounds_utc(dt.date(2024, 1, 10))   # EST: 09:30 ET = 14:30 UTC
    assert (o.hour, o.minute) == (14, 30) and (c.hour, c.minute) == (21, 0)
    o, c = regular_session_bounds_utc(dt.date(2024, 7, 10))   # EDT: 09:30 ET = 13:30 UTC
    assert (o.hour, o.minute) == (13, 30) and (c.hour, c.minute) == (20, 0)


def test_early_close_day():
    b = session_bounds_table(dt.date(2023, 11, 20), dt.date(2023, 11, 30))
    r = b.filter(pl.col("session_date") == dt.date(2023, 11, 24)).row(0, named=True)
    assert r["close_minute"] == 210  # 13:00 ET close the day after Thanksgiving


def test_shift_trading_days_skips_weekends_and_holidays():
    assert shift_trading_days(dt.date(2024, 7, 5), -1) == dt.date(2024, 7, 3)   # July 4 holiday
    assert shift_trading_days(dt.date(2024, 7, 8), -1) == dt.date(2024, 7, 5)
    assert trading_days_between(dt.date(2024, 7, 3), dt.date(2024, 7, 8)) == 2


def test_close_print_is_at_1600_et():
    c = close_time_utc(dt.date(2022, 3, 14))  # first session after the spring-forward
    assert (c.hour, c.minute) == (20, 0)