Handling Multi-Timezone GTFS Feeds

"""Resolve GTFS times correctly across a multi-timezone feed."""
from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import date, datetime, timedelta
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

import pandas as pd

log = logging.getLogger("gtfs.timezone")


@dataclass(frozen=True)
class TimezoneMap:
    by_agency: dict[str, str]             # agency_id -> IANA name
    by_route: dict[str, str]              # route_id -> IANA name
    by_trip: dict[str, str]               # trip_id -> IANA name
    stop_display: dict[str, str]          # stop_id -> stop_timezone, DISPLAY ONLY

    @property
    def zones(self) -> set[str]:
        return set(self.by_agency.values())

    def for_trip(self, trip_id: str) -> str:
        try:
            return self.by_trip[trip_id]
        except KeyError:
            raise KeyError(
                f"trip {trip_id} resolves to no agency timezone; check that its route "
                "carries an agency_id and that the agency exists") from None


def build_timezone_map(agency: pd.DataFrame, routes: pd.DataFrame,
                       trips: pd.DataFrame, stops: pd.DataFrame) -> TimezoneMap:
    if "agency_timezone" not in agency.columns:
        raise ValueError("agency.txt has no agency_timezone; every time in this feed "
                         "is uninterpretable")

    by_agency: dict[str, str] = {}
    for row in agency.itertuples(index=False):
        zone = str(row.agency_timezone).strip()
        try:
            ZoneInfo(zone)
        except (ZoneInfoNotFoundError, ValueError):
            raise ValueError(f"agency_timezone {zone!r} is not an IANA zone name") from None
        # agency_id is optional in a single-agency feed; key that case on "".
        agency_id = str(getattr(row, "agency_id", "") or "")
        by_agency[agency_id] = zone

    if len(by_agency) == 1 and "" not in by_agency:
        # A single named agency: routes may still omit agency_id entirely.
        by_agency[""] = next(iter(by_agency.values()))

    by_route: dict[str, str] = {}
    for row in routes.itertuples(index=False):
        agency_id = str(getattr(row, "agency_id", "") or "")
        zone = by_agency.get(agency_id) or by_agency.get("")
        if zone is None:
            log.error("route %s names agency %r, which agency.txt does not define",
                      row.route_id, agency_id)
            continue
        by_route[str(row.route_id)] = zone

    by_trip = {str(r.trip_id): by_route[str(r.route_id)]
               for r in trips.itertuples(index=False)
               if str(r.route_id) in by_route}

    # Captured for DISPLAY only. It never interprets a schedule time.
    stop_display = {}
    if "stop_timezone" in stops.columns:
        stop_display = {str(r.stop_id): str(r.stop_timezone)
                        for r in stops.itertuples(index=False)
                        if pd.notna(r.stop_timezone) and str(r.stop_timezone) != ""}

    mapping = TimezoneMap(by_agency=by_agency, by_route=by_route, by_trip=by_trip,
                          stop_display=stop_display)
    if len(mapping.zones) > 1:
        log.info("feed spans %d timezone(s): %s — schedule times resolve per agency",
                 len(mapping.zones), sorted(mapping.zones))
    if stop_display:
        crossing = {s for s, z in stop_display.items()
                    if z not in mapping.zones}
        log.info("%d stop(s) carry stop_timezone; %d differ from every agency zone "
                 "(display only)", len(stop_display), len(crossing))
    return mapping


def trip_times_to_utc(trip_id: str, service_day: date, offsets_s: list[int],
                      mapping: TimezoneMap) -> list[datetime]:
    """Convert one trip's second-offsets to UTC using its OPERATING agency's zone."""
    zone = ZoneInfo(mapping.for_trip(trip_id))
    noon = datetime.combine(service_day, datetime.min.time().replace(hour=12), tzinfo=zone)
    start_of_day = noon - timedelta(hours=12)
    return [(start_of_day + timedelta(seconds=int(s))).astimezone(ZoneInfo("UTC"))
            for s in offsets_s]


def local_display_time(instant: datetime, stop_id: str, trip_id: str,
                       mapping: TimezoneMap) -> datetime:
    """What the clock on the wall at that stop reads — the ONLY use of stop_timezone."""
    zone = mapping.stop_display.get(stop_id) or mapping.for_trip(trip_id)
    return instant.astimezone(ZoneInfo(zone))

def verify(mapping: TimezoneMap, trips: pd.DataFrame) -> None:
    unresolved = [str(r.trip_id) for r in trips.itertuples(index=False)
                  if str(r.trip_id) not in mapping.by_trip]
    assert not unresolved, (
        f"{len(unresolved)} trip(s) resolve to no timezone, e.g. {unresolved[:3]}")

    for zone in mapping.zones:
        ZoneInfo(zone)      # raises if the container has no tz database


def verify_trip_conversion(trip_id: str, service_day: date, offsets_s: list[int],
                           mapping: TimezoneMap) -> None:
    instants = trip_times_to_utc(trip_id, service_day, offsets_s, mapping)
    assert all(a <= b for a, b in zip(instants, instants[1:])), (
        f"trip {trip_id} runs backwards in UTC — the usual cause is interpreting a "
        "time in the stop's timezone rather than the operating agency's")
    span = (instants[-1] - instants[0]).total_seconds()
    assert span == offsets_s[-1] - offsets_s[0], (
        "the UTC span does not match the offset span; a DST transition was applied "
        "twice or a per-stop zone leaked into the conversion")

INFO gtfs.timezone: feed spans 2 timezone(s): ['America/Chicago', 'America/New_York'] — schedule times resolve per agency
INFO gtfs.timezone: 84 stop(s) carry stop_timezone; 31 differ from every agency zone (display only)