Aligning Seasonal Travel Patterns Across Multiple Years

import pandas as pd
import numpy as np
from typing import Literal


def align_seasonal_mobility(
    df: pd.DataFrame,
    date_col: str = "timestamp",
    volume_col: str = "trips",
    reference_year: int = 2023,
    leap_handling: Literal["shift", "drop"] = "shift",
) -> pd.DataFrame:
    """
    Align multi-year mobility data to a fixed 365-day seasonal reference frame.

    Parameters
    ----------
    df : pd.DataFrame
        Must contain `date_col` (datetime-like) and `volume_col` (numeric).
    date_col : str
        Column holding observation timestamps. UTC-aware or naive both accepted.
    volume_col : str
        Numeric column to carry through alignment (e.g. trip counts, km driven).
    reference_year : int
        Calendar year onto which all aligned dates are projected (e.g. 2023).
    leap_handling : {"shift", "drop"}
        "shift" — days after Feb 28 in a leap year are shifted back by 1 to keep
                  365-day continuity and preserve all observations.
        "drop"  — Feb 29 observations are removed; simpler but loses ~0.27% of
                  leap-year data and can introduce gaps in rolling windows.

    Returns
    -------
    pd.DataFrame with columns:
        seasonal_date  — datetime mapped to reference_year (for cross-year groupby)
        doy_aligned    — integer 1–365
        doy_sin        — sin(2π · doy_aligned / 365)
        doy_cos        — cos(2π · doy_aligned / 365)
        <volume_col>   — original volume, unchanged

    Raises
    ------
    ValueError  — missing columns, empty DataFrame, or unknown leap_handling value.
    """
    if df.empty:
        raise ValueError("Input DataFrame is empty.")
    for col in (date_col, volume_col):
        if col not in df.columns:
            raise ValueError(f"Missing required column: '{col}'")
    if leap_handling not in ("shift", "drop"):
        raise ValueError(f"leap_handling must be 'shift' or 'drop', got '{leap_handling}'")

    df = df.copy()
    df[date_col] = pd.to_datetime(df[date_col], utc=False)

    # Strip timezone info if present so Timestamp arithmetic stays naive.
    if df[date_col].dt.tz is not None:
        df[date_col] = df[date_col].dt.tz_convert("UTC").dt.tz_localize(None)

    df["_doy_raw"] = df[date_col].dt.dayofyear
    df["_is_leap"] = df[date_col].dt.is_leap_year

    if leap_handling == "shift":
        # Feb 29 is doy 60 in a leap year.  All days > 60 get shifted back by 1
        # so June 1 always maps to doy 152, regardless of whether the source year
        # is a leap year.
        leap_post = df["_is_leap"] & (df["_doy_raw"] > 60)
        df["doy_aligned"] = df["_doy_raw"].where(~leap_post, df["_doy_raw"] - 1)
    else:
        # Drop Feb 29 entirely (doy == 60 in a leap year).
        drop_mask = df["_is_leap"] & (df["_doy_raw"] == 60)
        if drop_mask.any():
            dropped = drop_mask.sum()
            import warnings
            warnings.warn(
                f"Dropping {dropped} Feb-29 observations (leap_handling='drop').",
                stacklevel=2,
            )
        df = df[~drop_mask].copy()
        df["doy_aligned"] = df["_doy_raw"]

    df["doy_aligned"] = df["doy_aligned"].astype("int32")

    # Project onto reference year so cross-year groupby works on real dates.
    ref_start = pd.Timestamp(f"{reference_year}-01-01")
    df["seasonal_date"] = ref_start + pd.to_timedelta(df["doy_aligned"] - 1, unit="D")

    # Cyclical encoding: map doy onto a unit circle.
    # Using 365 (not 365.25) because doy_aligned is already bounded to [1, 365].
    angle = 2 * np.pi * df["doy_aligned"] / 365
    df["doy_sin"] = np.sin(angle)
    df["doy_cos"] = np.cos(angle)

    # Clean up temporaries.
    df.drop(columns=["_doy_raw", "_is_leap"], inplace=True)

    return df[["seasonal_date", "doy_aligned", "doy_sin", "doy_cos", volume_col]]

from datetime import date, timedelta
import calendar


def nth_weekday_of_month(year: int, month: int, weekday: int, n: int) -> date:
    """Return the nth occurrence of `weekday` (0=Mon … 6=Sun) in the given month."""
    first = date(year, month, 1)
    first_weekday = first.weekday()
    delta = (weekday - first_weekday) % 7
    target = first + timedelta(days=delta + (n - 1) * 7)
    if target.month != month:
        raise ValueError(f"No {n}th weekday {weekday} in {year}-{month:02d}")
    return target


def thanksgiving_offset_days(year: int, reference_year: int = 2023) -> int:
    """
    Return how many days the Thanksgiving window in `year` is offset from
    the reference-year Thanksgiving.  Negative means earlier that year.
    """
    ref_tg = nth_weekday_of_month(reference_year, 11, 3, 4)  # 4th Thursday Nov
    yr_tg  = nth_weekday_of_month(year, 11, 3, 4)
    ref_doy = ref_tg.timetuple().tm_yday
    yr_doy  = yr_tg.timetuple().tm_yday
    return yr_doy - ref_doy  # positive → later in the year than reference


def apply_holiday_window_offset(
    df: pd.DataFrame,
    date_col: str,
    volume_col: str,
    reference_year: int = 2023,
) -> pd.DataFrame:
    """
    Shift each year's rows by the Thanksgiving offset so peak demand aligns
    before the 365-day normalization step.  Returns a modified copy.
    """
    df = df.copy()
    df[date_col] = pd.to_datetime(df[date_col])
    df["_year"] = df[date_col].dt.year
    df["_offset"] = df["_year"].map(
        lambda y: thanksgiving_offset_days(y, reference_year)
    )
    df[date_col] = df[date_col] - pd.to_timedelta(df["_offset"], unit="D")
    df.drop(columns=["_year", "_offset"], inplace=True)
    return df

import logging

logger = logging.getLogger(__name__)


def validate_alignment(aligned: pd.DataFrame, volume_col: str = "trips") -> None:
    """Assert structural invariants on the output of align_seasonal_mobility."""
    # 1. Shape and non-emptiness.
    assert not aligned.empty, "Aligned DataFrame is empty."
    expected_cols = {"seasonal_date", "doy_aligned", "doy_sin", "doy_cos", volume_col}
    missing = expected_cols - set(aligned.columns)
    assert not missing, f"Missing output columns: {missing}"

    # 2. doy_aligned must be within [1, 365] — no leap-day bleed-through.
    assert aligned["doy_aligned"].between(1, 365).all(), (
        "doy_aligned contains values outside [1, 365]; leap-year handling failed."
    )

    # 3. Cyclical features must lie on the unit circle (allow floating-point epsilon).
    unit_check = (aligned["doy_sin"] ** 2 + aligned["doy_cos"] ** 2)
    assert np.allclose(unit_check, 1.0, atol=1e-9), (
        "sin² + cos² ≠ 1; cyclical encoding is corrupted."
    )

    # 4. Year-boundary continuity: Dec 31 (doy 365) and Jan 1 (doy 1) must be
    #    close in cyclical feature space.  Angular distance should be near 2π/365.
    day1   = aligned[aligned["doy_aligned"] == 1][["doy_sin", "doy_cos"]].values
    day365 = aligned[aligned["doy_aligned"] == 365][["doy_sin", "doy_cos"]].values
    if day1.size > 0 and day365.size > 0:
        dot = np.dot(day1[0], day365[0])   # cosine similarity on unit vectors
        angle_gap = np.degrees(np.arccos(np.clip(dot, -1, 1)))
        assert angle_gap < 3.0, (
            f"Year-boundary angular gap is {angle_gap:.2f}°; expected < 3° (≈1 day)."
        )

    # 5. seasonal_date must all map to reference year (single-year range).
    years = aligned["seasonal_date"].dt.year.unique()
    assert len(years) == 1, f"seasonal_date spans multiple years: {years}"

    logger.info(
        "Alignment validated: %d rows, doy range [%d, %d], volume sum %.0f",
        len(aligned),
        aligned["doy_aligned"].min(),
        aligned["doy_aligned"].max(),
        aligned[volume_col].sum(),
    )