import numpy as np
import pandas as pd

from chronos_x.data.calendar import sessions, slot_index, REGULAR_SLOTS, EARLY_CLOSE_SLOTS


def test_regular_and_early_close_slots():
    s = {x.date.date().isoformat(): x for x in sessions("2024-11-25", "2024-12-02")}
    assert s["2024-11-29"].n_slots == EARLY_CLOSE_SLOTS and s["2024-11-29"].early_close
    assert s["2024-11-27"].n_slots == REGULAR_SLOTS and not s["2024-11-27"].early_close
    assert "2024-11-28" not in s  # Thanksgiving


def test_dst_open_in_utc():
    s = {x.date.date().isoformat(): x for x in sessions("2024-03-07", "2024-03-12")}
    assert s["2024-03-08"].open_utc == pd.Timestamp("2024-03-08T14:30:00Z")   # EST
    assert s["2024-03-11"].open_utc == pd.Timestamp("2024-03-11T13:30:00Z")   # EDT
    assert s["2024-03-11"].close_utc == pd.Timestamp("2024-03-11T20:00:00Z")


def test_slot_index_bar_open_convention():
    s = sessions("2024-03-11", "2024-03-11")[0]
    ts = np.array(["2024-03-11T13:30:00", "2024-03-11T13:31:00", "2024-03-11T20:00:00", "2024-03-11T20:01:00",
                   "2024-03-11T13:29:00", "2024-03-11T13:30:30"], dtype="datetime64[s]")
    k = slot_index(ts, np.full(len(ts), s.open_utc.to_datetime64().astype("datetime64[s]")), np.full(len(ts), s.n_slots))
    assert k.tolist() == [0, 1, 390, -1, -1, -1]