src/core.py

"""Data loading and TiRex-2 inference helpers, kept UI-framework-agnostic.
"""

from __future__ import annotations

import yaml
import numpy as np
import pandas as pd

from dataclasses import dataclass
from pathlib import Path
from typing import Any

# Hard model limits (read from the loaded model at runtime; these are fallbacks).
DEFAULT_CONTEXT_LEN = 2048
DEFAULT_FUTURE_LEN = 320

DATA_DIR = Path("data")
DATASET_INFO_PATH = DATA_DIR / "dataset_info.yaml"

SEASONAL_PATTERNS: dict[str, int] = {
    "Weekly cycle": 7,
    "Monthly cycle": 30,
    "Yearly cycle": 365,
}

WEEKEND_PATTERN = "Weekend flag"
HOLIDAY_PATTERN = "Holiday flag"

# Country used for holiday covariates when the caller does not specify one.
DEFAULT_HOLIDAY_COUNTRY = "US"


def infer_default_horizon(df: pd.DataFrame, time_column: str | None = None) -> int:
    """Pick a practical default forecast horizon from a dataset time column."""
    if not time_column or time_column not in df.columns:
        return 64

    parsed = pd.to_datetime(df[time_column], errors="coerce")
    if parsed.isna().any() or len(parsed) < 2:
        return 64

    delta = parsed.diff().dropna().median()
    if pd.isna(delta) or delta <= pd.Timedelta(0):
        return 64
    if delta <= pd.Timedelta(hours=1):
        return 168
    if delta <= pd.Timedelta(days=1):
        return 30
    if delta <= pd.Timedelta(days=8):
        return 12
    return 24


def load_dataset_catalog(path: str | Path = DATASET_INFO_PATH) -> dict[str, dict]:
    """Load example dataset metadata from ``dataset_info.yaml``."""
    info_path = Path(path)
    if not info_path.exists():
        return {}

    with info_path.open("r", encoding="utf-8") as f:
        raw = yaml.safe_load(f) or {}

    catalog: dict[str, dict] = {}
    for item in raw.get("datasets", []):
        name = str(item["name"])
        item_path = Path(item["path"])
        if not item_path.is_absolute():
            item_path = info_path.parent / item_path

        meta = dict(item)
        meta["path"] = str(item_path)
        try:
            preview = load_table(meta["path"], first_row_header=True)
            meta["horizon"] = int(meta.get("horizon") or infer_default_horizon(preview, meta.get("time_column")))
        except Exception:
            meta["horizon"] = int(meta.get("horizon") or 64)
        catalog[name] = meta
    return catalog


# ---------------------------------------------------------------------------
# Table loading
# ---------------------------------------------------------------------------

def load_table(file_or_path, *, filename: str | None = None, first_row_header: bool = True) -> pd.DataFrame:
    """Read a CSV / Excel / Parquet table from a path or an uploaded file object."""
    name = filename or getattr(file_or_path, "name", str(file_or_path))
    ext = str(name).split(".")[-1].lower()
    header = 0 if first_row_header else None
    if ext == "csv":
        return pd.read_csv(file_or_path, header=header)
    if ext in ("xls", "xlsx"):
        return pd.read_excel(file_or_path, header=header)
    if ext == "parquet":
        return pd.read_parquet(file_or_path)
    raise ValueError("Unsupported format. Use CSV, XLS, XLSX, or PARQUET.")


@dataclass
class SeriesTable:
    """A tidy view of a user table: one row per series, plus names."""

    names: list[str]
    values: np.ndarray  # shape [n_series, length], float32, NaNs allowed-but-discouraged

    @property
    def n_series(self) -> int:
        return len(self.names)

    @property
    def length(self) -> int:
        return self.values.shape[1] if self.values.ndim == 2 else 0


def to_series_table(df: pd.DataFrame) -> SeriesTable:
    """Turn a raw dataframe into rows-of-series, auto-detecting optional column names.

    Convention: each numeric column is one series. If the dataframe has non-numeric column
    names, those names are used; otherwise columns are auto-named ``Series 0, 1, ...``.
    """
    if not isinstance(df.columns, pd.RangeIndex) and not pd.api.types.is_numeric_dtype(df.columns):
        # pandas usually absorbs the header into the columns index.
        names = [str(x) for x in df.columns.tolist()]
        data = df
    else:
        names = [f"Series {i}" for i in range(df.shape[1])]
        data = df

    # SeriesTable expects an array where each row is a series.
    # Since our dataframe holds series as columns, we must transpose the extracted values.
    values = data.apply(pd.to_numeric, errors="coerce").to_numpy(dtype=np.float32).T
    
    return SeriesTable(names=names, values=values)


def clip_context(values: np.ndarray, context_len: int) -> np.ndarray:
    """Keep only the last ``context_len`` steps of each series."""
    if values.shape[1] > context_len:
        return values[:, -context_len:]
    return values


# ---------------------------------------------------------------------------
# Generated covariates
# ---------------------------------------------------------------------------

def extend_time_index(time_values, length: int) -> pd.DatetimeIndex:
    """Parse and extend a time column to ``length`` timestamps."""
    parsed = pd.to_datetime(pd.Series(time_values), errors="coerce")
    if parsed.isna().any():
        raise ValueError("Selected time column contains values that could not be parsed as dates/times.")
    if len(parsed) < 2:
        raise ValueError("Selected time column needs at least two timestamps to infer future steps.")

    freq = pd.infer_freq(parsed) if len(parsed) >= 3 else None
    if freq is not None:
        return pd.date_range(parsed.iloc[0], periods=length, freq=freq)

    deltas = parsed.diff().dropna()
    step = deltas.median()
    if pd.isna(step) or step <= pd.Timedelta(0):
        raise ValueError("Selected time column must be sorted with a positive regular interval.")
    return pd.DatetimeIndex([parsed.iloc[0] + i * step for i in range(length)])


def supported_holiday_countries() -> list[str]:
    """ISO country codes for which the ``holidays`` package can build a calendar."""
    import holidays

    return sorted(holidays.list_supported_countries())


def holiday_flag(time_index: pd.DatetimeIndex, country: str) -> np.ndarray:
    """A future-known 0/1 flag marking public holidays for ``country``."""
    import holidays

    years = range(int(time_index.year.min()), int(time_index.year.max()) + 1)
    try:
        calendar = holidays.country_holidays(country, years=years)
    except NotImplementedError as exc:
        raise ValueError(f"'{country}' is not a supported holiday calendar.") from exc
    dates = time_index.normalize().date
    return np.fromiter((d in calendar for d in dates), dtype=np.float32, count=len(dates))


def _cycle_fraction(period: int, steps: np.ndarray, time_index: pd.DatetimeIndex | None) -> np.ndarray:
    """Position within a cycle in ``[0, 1)``, calendar-aware when a time index exists."""
    if time_index is None:
        return (steps % period) / period
    if period == 7:  # weekly -> day of week
        return time_index.dayofweek.to_numpy(dtype=np.float32) / 7.0
    if period == 30:  # monthly -> fractional position in the month
        day = time_index.day.to_numpy(dtype=np.float32) - 1
        return day / time_index.days_in_month.to_numpy(dtype=np.float32)
    if period == 365:  # yearly -> day of year
        return (time_index.dayofyear.to_numpy(dtype=np.float32) - 1) / 365.0
    return (steps % period) / period


def seasonal_covariates(
    labels: list[str],
    length: int,
    *,
    time_values=None,
    country: str = DEFAULT_HOLIDAY_COUNTRY,
) -> tuple[list[str], np.ndarray]:
    """Create future-known calendar covariates that extend to any requested length.

    Cyclic patterns are encoded as a ``sin``/``cos`` Fourier pair (period-aligned to the
    calendar when a time column is present), which represents the full cycle unambiguously
    - unlike a single half-wave, ``sin(pi * t / period)``, where e.g. Monday and Sunday map
    to the same value. Weekend and holiday patterns are 0/1 flags and require a time column.
    """
    if not labels:
        return [], np.empty((0, length), dtype=np.float32)

    time_index = extend_time_index(time_values, length) if time_values is not None else None
    steps = np.arange(length, dtype=np.float32)
    names: list[str] = []
    values: list[np.ndarray] = []
    for label in labels:
        if label == WEEKEND_PATTERN:
            if time_index is None:
                continue
            names.append("weekend_flag")
            values.append((time_index.dayofweek >= 5).astype(np.float32))
        elif label == HOLIDAY_PATTERN:
            if time_index is None:
                continue
            names.append(f"holiday_{country.lower()}")
            values.append(holiday_flag(time_index, country))
        elif label in SEASONAL_PATTERNS:
            period = SEASONAL_PATTERNS[label]
            short = label.split(" cycle")[0].strip().lower().replace("-", "_").replace(" ", "_")
            angle = 2 * np.pi * _cycle_fraction(period, steps, time_index)
            names.append(f"{short}_sin")
            values.append(np.sin(angle).astype(np.float32))
            names.append(f"{short}_cos")
            values.append(np.cos(angle).astype(np.float32))

    if not values:
        return [], np.empty((0, length), dtype=np.float32)
    return names, np.asarray(values, dtype=np.float32)


# ---------------------------------------------------------------------------
# Forecasting
# ---------------------------------------------------------------------------

@dataclass
class ForecastResult:
    names: list[str]
    context: list[np.ndarray]        # per series: observed history used as input [Tc]
    quantiles: np.ndarray            # forecast [n_series, Q, H]
    quantile_levels: list[float]
    inference_s: float
    multivariate: bool
    truth: list[np.ndarray] | None = None  # per series: held-out actuals [H] (future-cov holdout)
    cov_names: list[str] | None = None     # covariate series used, if any
    cov_mode: str | None = None            # "past" | "future" when covariates are used
    timeseries: Any | None = None           # TiRex-2 TimeseriesType used for inference
    x_values: np.ndarray | None = None      # absolute x-axis values for context + future
    prediction_start: int | None = None     # absolute index where the forecast begins

    @property
    def horizon(self) -> int:
        return self.quantiles.shape[-1]

    def median_idx(self) -> int:
        levels = np.asarray(self.quantile_levels)
        return int(np.abs(levels - 0.5).argmin())

    def q_idx(self, q: float) -> int:
        levels = np.asarray(self.quantile_levels)
        return int(np.abs(levels - q).argmin())


def run_forecast(
    model,
    values: np.ndarray,
    names: list[str],
    *,
    horizon: int,
    multivariate: bool,
    context_len: int,
    tta_diff: bool | None = None,
    tta_sign_flip: bool | None = None,
    cov_values: np.ndarray | None = None,
    cov_names: list[str] | None = None,
    cov_mode: str = "future",
    prediction_start: int | None = None,
) -> ForecastResult:
    """Forecast a stack of series with TiRex-2.

    The dashboard forecast path uses one target series and optional covariates. The target
    and covariates are sliced into a TiRex-2 ``TimeseriesType`` so ``prediction_start``
    controls where the forecast begins, rather than always forecasting after the table end.

    Covariates (optional, ``cov_values`` is ``[n_cov, T]`` aligned to the targets):
      * ``cov_mode="past"``   -> passed as ``past_covariates`` (history only); targets are
        forecast from ``prediction_start`` using covariate history only.
      * ``cov_mode="future"`` -> covariates from the context window through the forecast
        horizon are passed as ``future_covariates``. This requires covariate values through
        ``prediction_start + horizon``.
    """
    import time

    import torch
    from tirex2 import TimeseriesType

    quantile_levels = [round(float(q), 6) for q in model.quantiles]

    predict_kwargs = {}
    if tta_diff is not None:
        predict_kwargs["tta_diff"] = tta_diff
    if tta_sign_flip is not None:
        predict_kwargs["tta_sign_flip"] = tta_sign_flip

    target_values = np.asarray(values, dtype=np.float32)
    if target_values.ndim == 1:
        target_values = target_values[None, :]
    if target_values.shape[0] != 1:
        raise ValueError("Select exactly one target series to forecast.")

    target = target_values[0]
    n_time = target.shape[0]
    if n_time < 2:
        raise ValueError("Target series must contain at least two time steps.")

    forecast_start = n_time if prediction_start is None else int(prediction_start)
    if forecast_start < 1 or forecast_start > n_time:
        raise ValueError(f"Forecast start must be between 1 and {n_time}.")

    context_start = max(0, forecast_start - context_len)
    context = np.ascontiguousarray(target[context_start:forecast_start], dtype=np.float32)
    if len(context) < 1:
        raise ValueError("Forecast start leaves no target history for the model.")
    if np.isnan(context).any():
        raise ValueError("Target context contains NaN values. Please clean or impute the selected series.")

    truth_values = target[forecast_start:min(forecast_start + horizon, n_time)]
    truth = [np.asarray(truth_values, dtype=np.float32)] if len(truth_values) else None

    has_cov = cov_values is not None and len(cov_values) > 0
    past_covariates = None
    future_covariates = None

    if has_cov:
        cov = np.asarray(cov_values, dtype=np.float32)
        if cov.ndim == 1:
            cov = cov[None, :]
        if cov.shape[1] < n_time:
            raise ValueError("Covariates must be aligned to the target and at least as long as the target.")
        if cov_mode == "future":
            cov_end = forecast_start + horizon
            if cov.shape[1] < cov_end:
                raise ValueError(
                    "Future-known covariates need values through the full forecast horizon "
                    f"(need index {cov_end - 1}, have {cov.shape[1] - 1})."
                )
            cov_slice = np.ascontiguousarray(cov[:, context_start:cov_end], dtype=np.float32)
            if np.isnan(cov_slice).any():
                raise ValueError("Future covariates contain NaN values in the context or forecast window.")
            future_covariates = torch.from_numpy(cov_slice)
        else:
            cov_slice = np.ascontiguousarray(cov[:, context_start:forecast_start], dtype=np.float32)
            if np.isnan(cov_slice).any():
                raise ValueError("Past covariates contain NaN values in the context window.")
            past_covariates = torch.from_numpy(cov_slice)

    start = time.monotonic()
    ts = TimeseriesType(
        target=torch.from_numpy(context[None, :]),
        past_covariates=past_covariates,
        future_covariates=future_covariates,
    )
    out = model.forecast([ts], prediction_length=horizon, output_type="numpy", **predict_kwargs)[0]
    quantiles = np.asarray(out, dtype=np.float32)  # [V, Q, H]
    if quantiles.ndim == 2:
        quantiles = quantiles[None, :, :]
    inference_s = time.monotonic() - start
    x_len = len(context) + max(quantiles.shape[-1], len(truth_values), ts.future_length)
    x_values = np.arange(context_start, context_start + x_len)

    return ForecastResult(
        names=names[:1],
        context=[context],
        quantiles=quantiles,
        quantile_levels=quantile_levels,
        inference_s=inference_s,
        multivariate=has_cov,
        truth=truth,
        cov_names=list(cov_names) if has_cov else None,
        cov_mode=cov_mode if has_cov else None,
        timeseries=ts,
        x_values=x_values,
        prediction_start=forecast_start,
    )