"""XNYS session calendar on a dense regular-hours minute grid.

Conventions (verified against FinData on 2026-11-02, docs/PHASE_0_REPORT.md):
- bars are stamped at bar OPEN in UTC; the bar opening at 09:30 ET is complete at 09:31 ET;
- regular session = 09:30 .. 16:00 ET inclusive -> slots 0..390 (391 bars); the slot 390 bar is the 16:00 closing
  print, treated separately (never a forecast target, never a feature bar for same-day context);
- early close at 13:00 ET -> slots 0..210 (211 bars); slot 210 is the closing print;
- DST handled through zoneinfo via exchange_calendars (opens/closes are tz-aware UTC timestamps).
"""
from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache

import numpy as np
import pandas as pd

REGULAR_SLOTS = 391          # 09:30 .. 16:00 inclusive
LAST_TRADING_SLOT = 389      # 15:59 bar (last bar that can be a target)
EARLY_CLOSE_SLOTS = 211


@dataclass(frozen=True)
class Session:
    date: pd.Timestamp          # session date (naive, ET calendar date)
    open_utc: pd.Timestamp      # open instant (09:30 ET) as UTC
    close_utc: pd.Timestamp     # close instant (16:00 or 13:00 ET) as UTC, == open of the closing print bar
    n_slots: int                # 391 or 211
    early_close: bool

    def slot_of(self, ts_utc: pd.Timestamp) -> int:
        """Slot index of a bar-open UTC timestamp, or -1 if outside the regular session grid."""
        k = int((ts_utc - self.open_utc).total_seconds() // 60)
        if k < 0 or k >= self.n_slots or (ts_utc - self.open_utc).total_seconds() % 60 != 0:
            return -1
        return k


@lru_cache(maxsize=4)
def _xnys():
    import exchange_calendars as xc
    return xc.get_calendar("XNYS")


def sessions(start: str, end: str) -> list[Session]:
    cal = _xnys()
    sched = cal.schedule.loc[pd.Timestamp(start):pd.Timestamp(end)]
    out = []
    for date, row in sched.iterrows():
        o = pd.Timestamp(row["open"]).tz_convert("UTC")
        c = pd.Timestamp(row["close"]).tz_convert("UTC")
        n = int((c - o).total_seconds() // 60) + 1
        out.append(Session(pd.Timestamp(date).tz_localize(None), o, c, n, n != REGULAR_SLOTS))
    return out


def session_frame(start: str, end: str) -> pd.DataFrame:
    rows = [{"session_date": s.date.date(), "open_utc": s.open_utc, "close_utc": s.close_utc, "n_slots": s.n_slots,
             "early_close": s.early_close} for s in sessions(start, end)]
    return pd.DataFrame(rows)


def slot_index(ts_utc: np.ndarray, open_utc: np.ndarray, n_slots: np.ndarray) -> np.ndarray:
    """Vectorised slot index (-1 outside the regular grid). All arrays are datetime64[ns]/int aligned per row."""
    delta = (ts_utc.astype("datetime64[s]") - open_utc.astype("datetime64[s]")).astype("int64")
    k = delta // 60
    ok = (delta >= 0) & (delta % 60 == 0) & (k < n_slots)
    return np.where(ok, k, -1).astype(np.int64)


def minute_of_session_features(slot: np.ndarray, n_slots: int) -> dict[str, np.ndarray]:
    """Calendar features that are known at origin: normalised minute of session and time to close."""
    s = slot.astype(np.float64)
    return {"minute_frac": s / (REGULAR_SLOTS - 1), "time_to_close": (n_slots - 1 - s) / (REGULAR_SLOTS - 1)}