atm-backtesting/scripts/calendar_parse.py

"""Calendar parsing + Set classification for M2D backtesting.

Each trade is tagged with a ``Set`` derived from its date, RO-local time, and the
economic-event calendar:

- ``A1``: 16:35-17:00 RO, Tue/Wed/Thu
- ``A2``: 17:00-18:00 RO, Tue/Wed/Thu (sweet spot)
- ``A3``: 18:00-19:00 RO, Tue/Wed/Thu
- ``B`` : 22:00-22:45 RO, Tue/Wed/Thu
- ``C`` : inside the window of an event with severity in {extrem, mare}
- ``D`` : Mon or Fri
- ``Other``: anything else

Priority: C > D > A1/A2/A3/B > Other.
"""

from __future__ import annotations

from datetime import date, datetime, time
from pathlib import Path
from typing import Any

import yaml
from zoneinfo import ZoneInfo

__all__ = [
    "RO_TZ",
    "UTC_TZ",
    "utc_to_ro",
    "load_calendar",
    "is_in_news_window",
    "calc_set",
]


RO_TZ = ZoneInfo("Europe/Bucharest")
UTC_TZ = ZoneInfo("UTC")

_DAY_SHORT = ("Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun")

_HIGH_SEVERITY = frozenset({"extrem", "mare"})

_WEEKLY_DAY_MAP = {
    "monday": 0,
    "tuesday": 1,
    "wednesday": 2,
    "thursday": 3,
    "friday": 4,
    "saturday": 5,
    "sunday": 6,
}


def utc_to_ro(date_str: str, ora_utc_str: str) -> tuple[date, time, str]:
    """Convert ``(YYYY-MM-DD, HH:MM UTC)`` to ``(date_ro, time_ro, day_short)``.

    DST-aware via :mod:`zoneinfo`. ``day_short`` is one of
    ``Mon Tue Wed Thu Fri Sat Sun``.
    """
    dt_utc = datetime.strptime(f"{date_str} {ora_utc_str}", "%Y-%m-%d %H:%M").replace(
        tzinfo=UTC_TZ
    )
    dt_ro = dt_utc.astimezone(RO_TZ)
    return dt_ro.date(), dt_ro.time().replace(second=0, microsecond=0), _DAY_SHORT[dt_ro.weekday()]


def load_calendar(path: Path | str = "calendar_evenimente.yaml") -> list[dict[str, Any]]:
    """Load a YAML calendar file.

    Validates ``schema_version == 1`` and returns the list of event dicts under
    the top-level ``events`` key.
    """
    p = Path(path)
    with p.open("r", encoding="utf-8") as fh:
        doc = yaml.safe_load(fh)
    if not isinstance(doc, dict):
        raise ValueError(f"calendar file {p} is not a mapping")
    version = doc.get("schema_version")
    if version != 1:
        raise ValueError(
            f"unsupported calendar schema_version: {version!r} (expected 1)"
        )
    events = doc.get("events") or []
    if not isinstance(events, list):
        raise ValueError(f"calendar events must be a list, got {type(events).__name__}")
    return events


def _minutes(t: time) -> int:
    return t.hour * 60 + t.minute


def _parse_hhmm(s: str) -> time:
    return datetime.strptime(s, "%H:%M").time()


def _is_first_friday_of_month(d: date) -> bool:
    return d.weekday() == 4 and d.day <= 7


def _event_matches_date(event: dict[str, Any], d: date) -> bool:
    cadence = event.get("cadence", "")
    if cadence == "scheduled":
        ev_date_raw = event.get("date")
        if isinstance(ev_date_raw, date):
            ev_date = ev_date_raw
        elif isinstance(ev_date_raw, str):
            ev_date = datetime.strptime(ev_date_raw, "%Y-%m-%d").date()
        else:
            return False
        return ev_date == d
    if cadence == "first_friday_monthly":
        return _is_first_friday_of_month(d)
    if cadence.startswith("weekly_"):
        day_name = cadence[len("weekly_") :].lower()
        target = _WEEKLY_DAY_MAP.get(day_name)
        if target is None:
            return False
        return d.weekday() == target
    # cadences below are not pinned down to a precise calendar day yet, so we
    # do not trigger Set C for them. ADP pre-NFP is also explicitly deferred.
    return False


def is_in_news_window(d: date, t: time, calendar: list[dict[str, Any]]) -> bool:
    """Return True iff ``(d, t)`` falls inside the window of a high-severity event.

    Window: ``[time_ro - window_before_min, time_ro + window_after_min]`` (inclusive
    on both ends). Only events with ``severity`` in ``{extrem, mare}`` count.

    Cadences honoured: ``scheduled``, ``first_friday_monthly``, ``weekly_<day>``.
    Other cadences (``monthly_mid``, ``monthly_end``, ``monthly_15``,
    ``wednesday_pre_nfp``, ``monthly_first_week`` etc.) are deferred and never
    trigger Set C.
    """
    t_min = _minutes(t)
    for event in calendar:
        if event.get("severity") not in _HIGH_SEVERITY:
            continue
        if not _event_matches_date(event, d):
            continue
        ev_time_raw = event.get("time_ro")
        if isinstance(ev_time_raw, time):
            ev_time = ev_time_raw
        elif isinstance(ev_time_raw, str):
            ev_time = _parse_hhmm(ev_time_raw)
        else:
            continue
        center = _minutes(ev_time)
        before = int(event.get("window_before_min", 0))
        after = int(event.get("window_after_min", 0))
        if center - before <= t_min <= center + after:
            return True
    return False


def _in_range(t: time, lo: time, hi: time) -> bool:
    """Half-open ``[lo, hi)`` containment."""
    return _minutes(lo) <= _minutes(t) < _minutes(hi)


def calc_set(d: date, t: time, day_of_week: str, calendar: list[dict[str, Any]]) -> str:
    """Classify a trade into one of ``A1 A2 A3 B C D Other``.

    Priority: ``C`` (news) > ``D`` (Mon/Fri) > ``A1/A2/A3/B`` (time bands on
    Tue/Wed/Thu) > ``Other``.
    """
    if is_in_news_window(d, t, calendar):
        return "C"
    if day_of_week in ("Mon", "Fri"):
        return "D"
    if day_of_week in ("Tue", "Wed", "Thu"):
        if _in_range(t, time(16, 35), time(17, 0)):
            return "A1"
        if _in_range(t, time(17, 0), time(18, 0)):
            return "A2"
        if _in_range(t, time(18, 0), time(19, 0)):
            return "A3"
        if _in_range(t, time(22, 0), time(22, 45)):
            return "B"
    return "Other"