src/tsfm_bench/eval/prequential.py

from __future__ import annotations

import hashlib
import json
import math
import re
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Iterable

import numpy as np
import pandas as pd
from gluonts.dataset.common import ListDataset

from tsfm_bench.data.ts_bench import TsBenchDataSource
from tsfm_bench.eval.canonical_metrics import (
    CANONICAL_METRIC_VERSION,
    CANONICAL_QUANTILES,
    canonical_metrics,
)


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _slug(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", "_", value.lower()).strip("_")


def _atomic_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(_json_safe(payload), indent=2, allow_nan=False) + "\n",
        encoding="utf-8",
    )
    temporary.replace(path)


def _json_safe(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(key): _json_safe(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_json_safe(item) for item in value]
    if isinstance(value, (float, np.floating)):
        return float(value) if math.isfinite(float(value)) else None
    if isinstance(value, np.integer):
        return int(value)
    if isinstance(value, np.bool_):
        return bool(value)
    return value


def _parse_time(value: object) -> datetime:
    parsed = pd.Timestamp(str(value))
    if parsed.tzinfo is None:
        parsed = parsed.tz_localize("UTC")
    else:
        parsed = parsed.tz_convert("UTC")
    return parsed.to_pydatetime()


def forecast_deadline(task: dict[str, Any]) -> datetime:
    """A forecast must be persisted before its first target period starts."""
    first_target = pd.Period(task["context_end"], freq=task["frequency"]) + 1
    return first_target.start_time.tz_localize("UTC").to_pydatetime()


def model_is_admitted(model_spec: dict[str, Any], issued_at: str) -> bool:
    admitted_at = model_spec.get("admitted_at")
    if not admitted_at:
        return False
    return _parse_time(admitted_at) <= _parse_time(issued_at)


def stable_dataset_name(source: Any, ds_name: str) -> str:
    if isinstance(source, TsBenchDataSource):
        return source.leaderboard_dataset_name(ds_name)
    meta = source.get_metadata(ds_name)
    return f"{ds_name.split('/')[0].lower()}/{meta.frequency}/short"


def _context_length(source: Any, ds_name: str, target_size: int, prediction_length: int) -> int:
    if isinstance(source, TsBenchDataSource):
        configured = source._tasks[ds_name].context_length
        if configured is not None:
            return max(1, min(int(configured), target_size))
    return max(1, target_size - prediction_length)


def observed_window(source: Any, ds_name: str) -> dict[str, Any]:
    """Return only observations already available at task issue time."""

    meta = source.get_metadata(ds_name)
    record = next(iter(source.stream(ds_name)))
    target = np.asarray(record.target, dtype=np.float64)
    if target.ndim != 1:
        raise ValueError(
            f"Prequential protocol currently requires univariate input; got {target.shape}"
        )
    finite = np.isfinite(target)
    target = target[finite]
    if target.size < 2:
        raise ValueError(f"Insufficient observed context for {ds_name}")

    prediction_length = (
        source.get_prediction_length(ds_name)
        if isinstance(source, TsBenchDataSource)
        else int(getattr(getattr(source, "_settings", None), "prediction_length", 24))
    )
    periods = pd.period_range(record.start, periods=len(finite), freq=record.freq)
    periods = periods[finite]
    issue_cutoff = pd.Timestamp.now(tz="UTC")
    not_future = np.asarray(
        [period.start_time.tz_localize("UTC") <= issue_cutoff for period in periods],
        dtype=bool,
    )
    target = target[not_future]
    periods = periods[not_future]
    if target.size < 2:
        raise ValueError(f"No causal observed context is available for {ds_name}")
    length = _context_length(source, ds_name, target.size, prediction_length)
    context = target[-length:]
    context_periods = periods[-length:]
    dataset = stable_dataset_name(source, ds_name)
    fetched_at = str(getattr(source, "fetched_at", lambda: _utc_now())())
    return {
        "source_task_name": ds_name,
        "dataset": dataset,
        "domain": meta.domain,
        "frequency": meta.frequency,
        "prediction_length": int(prediction_length),
        "context": context.tolist(),
        "context_timestamps": [period.start_time.isoformat() for period in context_periods],
        "context_start": str(context_periods[0]),
        "context_end": str(context_periods[-1]),
        "data_fetched_at": fetched_at,
    }


def forecast_context(predictor: Any, task: dict[str, Any]) -> dict[str, Any]:
    context = np.asarray(task["context"], dtype=np.float32)
    entry = {
        "item_id": task["dataset"].split("/")[0],
        "start": pd.Period(task["context_start"], freq=task["frequency"]),
        "target": context,
    }
    dataset = ListDataset([entry], freq=task["frequency"])
    forecasts = list(predictor.predict(dataset))
    if len(forecasts) != 1:
        raise ValueError(f"Expected one forecast, received {len(forecasts)}")
    forecast = forecasts[0]
    horizon = int(task["prediction_length"])
    mean = np.asarray(forecast.mean, dtype=np.float64).reshape(-1)
    if mean.size < horizon:
        raise ValueError(f"Mean forecast length {mean.size} < {horizon}")
    mean = mean[:horizon]
    quantiles: dict[str, list[float]] = {}
    for level in CANONICAL_QUANTILES:
        values = np.asarray(forecast.quantile(level), dtype=np.float64).reshape(-1)
        if values.size < horizon:
            raise ValueError(f"Quantile {level:g} length {values.size} < {horizon}")
        quantiles[f"{level:g}"] = values[:horizon].tolist()
    if not np.all(np.isfinite(mean)) or any(
        not np.all(np.isfinite(values)) for values in quantiles.values()
    ):
        raise ValueError("Forecast contains non-finite values")
    return {"mean": mean.tolist(), "quantiles": quantiles}


@dataclass
class PrequentialCycle:
    resolved_rows: dict[str, list[list[Any]]] = field(default_factory=dict)
    resolved_meta: dict[str, list[dict[str, Any]]] = field(default_factory=dict)
    issued_tasks: int = 0
    issued_forecasts: int = 0
    pending_tasks: int = 0
    failed_forecasts: list[dict[str, str]] = field(default_factory=list)


class PrequentialStore:
    """Persistent issue-now/score-later state for the live evaluator."""

    def __init__(self, output_root: Path):
        self.output_root = Path(output_root)
        self.root = self.output_root / "prequential"
        self.tasks_dir = self.root / "tasks"
        self.forecasts_dir = self.root / "forecasts"
        self.tasks_dir.mkdir(parents=True, exist_ok=True)
        self.forecasts_dir.mkdir(parents=True, exist_ok=True)

    def _task_path(self, task_id: str) -> Path:
        return self.tasks_dir / f"{task_id}.json"

    def _forecast_path(self, task_id: str, model_slug: str) -> Path:
        return self.forecasts_dir / task_id / f"{model_slug}.json"

    def forecast_exists(self, task_id: str, model_slug: str) -> bool:
        return self._forecast_path(task_id, model_slug).exists()

    def tasks(self, *, status: str | None = None) -> list[dict[str, Any]]:
        tasks: list[dict[str, Any]] = []
        for path in sorted(self.tasks_dir.glob("*.json")):
            try:
                task = json.loads(path.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError):
                continue
            if status is None or task.get("status") == status:
                tasks.append(task)
        return tasks

    def open_task(self, dataset: str) -> dict[str, Any] | None:
        for task in self.tasks(status="open"):
            if task.get("dataset") == dataset:
                return task
        return None

    def ensure_task(self, window: dict[str, Any]) -> tuple[dict[str, Any], bool]:
        existing = self.open_task(window["dataset"])
        if existing is not None:
            return existing, False
        identity = (
            f"{window['dataset']}|{window['context_end']}|"
            f"h:{window['prediction_length']}"
        )
        digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()[:20]
        task_id = f"future-{digest}"
        if self._task_path(task_id).exists():
            # Never overwrite a resolved/expired release when the source is stale.
            return json.loads(self._task_path(task_id).read_text()), False
        task = {
            **window,
            "task_id": task_id,
            "issued_at": _utc_now(),
            "status": "open",
            "scored_models": [],
            "metric_version": CANONICAL_METRIC_VERSION,
        }
        _atomic_json(self._task_path(task_id), task)
        return task, True

    def issue_forecast(
        self,
        task: dict[str, Any],
        model_spec: dict[str, Any],
        *,
        model_name: str,
        model_slug: str,
        predictor: Any,
    ) -> bool:
        path = self._forecast_path(task["task_id"], model_slug)
        if path.exists():
            return False
        if not model_is_admitted(model_spec, task["issued_at"]):
            return False
        if task["status"] != "open" or _parse_time(_utc_now()) >= forecast_deadline(task):
            return False
        predictions = forecast_context(predictor, task)
        completed_at = _utc_now()
        if _parse_time(completed_at) >= forecast_deadline(task):
            return False
        payload = {
            "task_id": task["task_id"],
            "dataset": task["dataset"],
            "model": model_name,
            "model_id": model_spec["model_id"],
            "admitted_at": model_spec.get("admitted_at"),
            "issued_at": task["issued_at"],
            "forecast_created_at": completed_at,
            "protocol": "prequential-v1",
            "metric_version": CANONICAL_METRIC_VERSION,
            **predictions,
        }
        _atomic_json(path, payload)
        return True

    def _current_windows(self, source: Any, dataset_names: Iterable[str]) -> dict[str, dict[str, Any]]:
        windows: dict[str, dict[str, Any]] = {}
        for ds_name in dataset_names:
            try:
                window = observed_window(source, ds_name)
            except Exception:
                continue
            windows[window["dataset"]] = window
        return windows

    def resolve_ready(
        self,
        source: Any,
        dataset_names: Iterable[str],
    ) -> PrequentialCycle:
        cycle = PrequentialCycle()
        current = self._current_windows(source, dataset_names)
        for task in self.tasks(status="open"):
            forecast_paths = sorted((self.forecasts_dir / task["task_id"]).glob("*.json"))
            if not forecast_paths and _parse_time(_utc_now()) >= forecast_deadline(task):
                task.update(status="expired", expired_at=_utc_now(),
                            expiration_reason="No frozen forecasts before target start")
                _atomic_json(self._task_path(task["task_id"]), task)
                continue
            window = current.get(task["dataset"])
            if window is None:
                continue
            current_periods = [
                pd.Period(timestamp, freq=task["frequency"])
                for timestamp in window["context_timestamps"]
            ]
            cutoff = pd.Period(task["context_end"], freq=task["frequency"])
            horizon = int(task["prediction_length"])
            expected = [cutoff + step for step in range(1, horizon + 1)]
            # A provider can return the currently forming candle/aggregate.
            # Score only after the whole target window has closed.
            if expected[-1].end_time.tz_localize("UTC") >= pd.Timestamp(_utc_now()):
                continue
            observations = dict(zip(current_periods, window["context"]))
            if current_periods and min(current_periods) > expected[0]:
                task.update(status="expired", expired_at=_utc_now(),
                            expiration_reason="Exact target window no longer available")
                _atomic_json(self._task_path(task["task_id"]), task)
                continue
            if not all(period in observations for period in expected):
                continue
            actual_periods = expected
            actual = [float(observations[period]) for period in expected]
            scored = set(task.get("scored_models", []))
            for forecast_path in forecast_paths:
                try:
                    forecast = json.loads(forecast_path.read_text(encoding="utf-8"))
                except (OSError, json.JSONDecodeError):
                    continue
                model_name = str(forecast.get("model", "")).strip()
                if not model_name or model_name in scored:
                    continue
                created_at = forecast.get("forecast_created_at", forecast.get("issued_at"))
                if _parse_time(created_at) >= forecast_deadline(task):
                    continue
                metrics = canonical_metrics(
                    context=task["context"],
                    actual=actual,
                    mean=forecast["mean"],
                    quantiles=forecast["quantiles"],
                )
                evaluated_at = _utc_now()
                release_time = actual_periods[-1].start_time.tz_localize("UTC").isoformat()
                self._append_release(task, forecast, metrics, release_time, evaluated_at)
                self._save_snapshot(task, forecast, metrics, actual, actual_periods, evaluated_at)
                cycle.resolved_rows.setdefault(model_name, []).append(
                    self._compatibility_row(task, forecast, metrics)
                )
                cycle.resolved_meta.setdefault(model_name, []).append(
                    {
                        "dataset": task["dataset"],
                        "task_id": task["task_id"],
                        "issued_at": task["issued_at"],
                        "resolved_at": evaluated_at,
                        "target_end": release_time,
                        "prediction_length": horizon,
                        "metric_version": CANONICAL_METRIC_VERSION,
                    }
                )
                scored.add(model_name)
            task["scored_models"] = sorted(scored)
            if forecast_paths and len(scored) >= len(forecast_paths):
                task["status"] = "resolved"
                task["resolved_at"] = _utc_now()
                task["target_start"] = actual_periods[0].start_time.isoformat()
                task["target_end"] = actual_periods[-1].start_time.isoformat()
            _atomic_json(self._task_path(task["task_id"]), task)
        cycle.pending_tasks = len(self.tasks(status="open"))
        return cycle

    def _append_release(
        self,
        task: dict[str, Any],
        forecast: dict[str, Any],
        metrics: dict[str, Any],
        release_time: str,
        evaluated_at: str,
    ) -> None:
        payload = {
            "release_id": task["task_id"],
            "release_time": release_time,
            "evaluated_at": evaluated_at,
            "issued_at": task["issued_at"],
            "protocol": "prequential-v1",
            "dataset": task["dataset"],
            "domain": task["domain"],
            "frequency": task["frequency"],
            "prediction_length": task["prediction_length"],
            "model": forecast["model"],
            "model_id": forecast["model_id"],
            "metric_version": metrics["metric_version"],
            "MSE": metrics["MSE"],
            "RMSE": metrics["RMSE"],
            "MAPE": metrics["MAPE"] if math.isfinite(float(metrics["MAPE"])) else None,
            "CRPS": metrics["CRPS"] if math.isfinite(float(metrics["CRPS"])) else None,
            "MAPE_eligible": metrics["MAPE_eligible"],
        }
        # Keep canonical rows separate from the historical mixed-scale cache.
        # This makes it impossible for a new live cycle to silently combine old
        # raw-scale scores with context-normalised scores.
        log_path = self.output_root / "evaluation_metrics_canonical.jsonl"
        log_path.parent.mkdir(parents=True, exist_ok=True)
        with log_path.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(payload, allow_nan=False) + "\n")

    def refresh_public_snapshots(self, source: Any, dataset_names: Iterable[str]) -> int:
        """Expose frozen predictions and closed targets without publishing task caches.

        Keep one snapshot per model/dataset. Once a complete evaluation exists it
        remains visible until the next complete evaluation replaces it. Before
        that, the first live window is shown with missing targets explicitly null.
        This operation never issues forecasts or writes leaderboard scores.
        """
        current = self._current_windows(source, dataset_names)
        written = 0
        now = pd.Timestamp(_utc_now())
        for task in sorted(self.tasks(), key=lambda t: t["issued_at"]):
            if task.get("status") not in {"open", "expired"}:
                continue
            window = current.get(task["dataset"])
            if window is None:
                continue
            freq = task["frequency"]
            cutoff = pd.Period(task["context_end"], freq=freq)
            expected = [cutoff + step for step in range(1, int(task["prediction_length"]) + 1)]
            observations = dict(zip(
                [pd.Period(t, freq=freq) for t in window["context_timestamps"]],
                window["context"],
            ))
            actuals = [observations.get(p) if p.end_time.tz_localize("UTC") < now else None
                       for p in expected]
            for path in sorted((self.forecasts_dir / task["task_id"]).glob("*.json")):
                try:
                    forecast = json.loads(path.read_text())
                    created_at = forecast["forecast_created_at"]
                    if (forecast["task_id"] != task["task_id"]
                            or _parse_time(created_at) >= forecast_deadline(task)):
                        continue
                    predictions = {"mean": forecast["mean"],
                                   "p10": forecast["quantiles"]["0.1"],
                                   "p50": forecast["quantiles"]["0.5"],
                                   "p90": forecast["quantiles"]["0.9"]}
                    if any(len(v) != len(expected) or not np.all(np.isfinite(v))
                           for v in predictions.values()):
                        continue
                    destination = self._public_snapshot_path(task, forecast)
                    previous = json.loads(destination.read_text()) if destination.exists() else {}
                    if previous.get("status") == "resolved":
                        continue
                    values = list(actuals)
                    if previous.get("task_id") == task["task_id"]:
                        # A partial collection must not erase previously closed targets.
                        old_values = previous.get("actuals", [])
                        if len(old_values) == len(values):
                            values = [v if v is not None else old for v, old in zip(values, old_values)]
                    snapshot = {
                        "schema_version": 1,
                        "status": "expired" if task["status"] == "expired" else "pending",
                        "protocol": "prequential-v1",
                        "dataset": task["dataset"], "domain": task["domain"],
                        "model": forecast["model"], "model_id": forecast["model_id"],
                        "task_id": task["task_id"], "issued_at": task["issued_at"],
                        "forecast_created_at": created_at, "updated_at": _utc_now(),
                        "freq": freq, "prediction_length": len(expected),
                        "context": task["context"], "actuals": values, **predictions,
                        "timestamps": [*task["context_timestamps"],
                                       *(p.start_time.isoformat() for p in expected)],
                    }
                    _atomic_json(destination, snapshot)
                    written += 1
                except (OSError, ValueError, KeyError, TypeError):
                    continue
        return written

    def _public_snapshot_path(self, task: dict, forecast: dict) -> Path:
        dataset_key = task["dataset"].replace("/", "__").replace(" ", "_")
        return self.output_root / "public_snapshots" / _slug(forecast["model"]) / f"{dataset_key}.json"

    def _save_snapshot(
        self,
        task: dict[str, Any],
        forecast: dict[str, Any],
        metrics: dict[str, Any],
        actual: list[float],
        actual_periods: list[pd.Period],
        evaluated_at: str,
    ) -> None:
        quantiles = forecast["quantiles"]
        snapshot = {
            "schema_version": 1,
            "status": "resolved",
            "dataset": task["dataset"],
            "domain": task["domain"],
            "model": forecast["model"],
            "model_id": forecast["model_id"],
            "task_id": task["task_id"],
            "issued_at": task["issued_at"],
            "forecast_created_at": forecast["forecast_created_at"],
            "protocol": "prequential-v1",
            "evaluated_at": evaluated_at,
            "freq": task["frequency"],
            "context_length": len(task["context"]),
            "prediction_length": task["prediction_length"],
            "context": task["context"],
            "actuals": actual,
            "actual": actual,
            "mean": forecast["mean"],
            "p10": quantiles.get("0.1", []),
            "p50": quantiles.get("0.5", forecast["mean"]),
            "p90": quantiles.get("0.9", []),
            "quantiles": quantiles,
            "predictions": {
                "mean": forecast["mean"],
                "p10": quantiles.get("0.1", []),
                "p50": quantiles.get("0.5", forecast["mean"]),
                "p90": quantiles.get("0.9", []),
                "quantiles": quantiles,
            },
            "metrics": metrics,
            "timestamps": [
                *task["context_timestamps"],
                *(period.start_time.isoformat() for period in actual_periods),
            ],
        }
        model_slug = _slug(forecast["model"])
        _atomic_json(self._public_snapshot_path(task, forecast), snapshot)
        dataset_key = task["dataset"].replace("/", "__").replace(" ", "_")
        _atomic_json(
            self.output_root / model_slug / "forecasts" / f"{dataset_key}.json",
            snapshot,
        )
        domain_slug = _slug(task["domain"])
        dataset_slug = _slug(task["dataset"])
        stamp = pd.Timestamp(evaluated_at).strftime("%Y%m%d_%H%M%S")
        _atomic_json(
            self.output_root
            / "saved_evals"
            / domain_slug
            / dataset_slug
            / model_slug
            / f"{stamp}.json",
            snapshot,
        )

    @staticmethod
    def _compatibility_row(
        task: dict[str, Any],
        forecast: dict[str, Any],
        metrics: dict[str, Any],
    ) -> list[Any]:
        return [
            task["dataset"],
            forecast["model"],
            metrics["MSE"],
            metrics["MSE[0.5]"],
            float("nan"),
            float("nan"),
            metrics["MAPE"],
            float("nan"),
            float("nan"),
            metrics["RMSE"],
            float("nan"),
            float("nan"),
            metrics["CRPS"],
            task["domain"],
            1,
            task["prediction_length"],
        ]


def run_prequential_cycle(
    *,
    source: Any,
    dataset_names: list[str],
    model_specs: list[dict[str, Any]],
    output_root: Path,
    predictor_factory: Callable[[dict[str, Any], int], Any],
    model_name: Callable[[dict[str, Any]], str],
    model_slug: Callable[[dict[str, Any]], str],
) -> PrequentialCycle:
    store = PrequentialStore(output_root)
    cycle = store.resolve_ready(source, dataset_names)

    for ds_name in dataset_names:
        window = observed_window(source, ds_name)
        if store.open_task(window["dataset"]) is None and _parse_time(_utc_now()) >= forecast_deadline(window):
            cycle.failed_forecasts.append({"model": "*", "dataset": window["dataset"],
                                           "error": "Source is stale: first target period already started"})
            continue
        task, created = store.ensure_task(window)
        if task["status"] != "open" or _parse_time(_utc_now()) >= forecast_deadline(task):
            continue
        if created:
            cycle.issued_tasks += 1
        for spec in model_specs:
            name = model_name(spec)
            slug = model_slug(spec)
            if not model_is_admitted(spec, task["issued_at"]):
                continue
            if store.forecast_exists(task["task_id"], slug):
                continue
            try:
                predictor = predictor_factory(spec, int(task["prediction_length"]))
                if store.issue_forecast(
                    task,
                    spec,
                    model_name=name,
                    model_slug=slug,
                    predictor=predictor,
                ):
                    cycle.issued_forecasts += 1
            except Exception as exc:
                cycle.failed_forecasts.append(
                    {"model": name, "dataset": task["dataset"], "error": str(exc)}
                )
    cycle.pending_tasks = len(store.tasks(status="open"))
    store.refresh_public_snapshots(source, dataset_names)
    return cycle