src/store.py

"""Read and write access to the signal store.

The store is a Hugging Face dataset repo laid out as parquet under stable
paths. `local_root` is a working mirror: reads fall back to the Hub when a file
is absent locally, writes stage into the mirror, and `flush()` pushes them as
one commit.

The rule that shapes this whole module is **frozen at issue**. An archived
forecast row is written once and never rewritten. Resolution appends to a
different tree; it does not revise the archive. So `append_forecast` is
first-write-wins on its dedup key rather than last-write-wins, which is the
opposite of what a cache would do and is the entire point: a track record whose
history can be edited after the fact is not evidence of anything.
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import threading
from pathlib import Path

import numpy as np
import pandas as pd

from . import config

log = logging.getLogger("arena.store")


class StoreError(RuntimeError):
    pass


# --------------------------------------------------------------------------
# Schemas
# --------------------------------------------------------------------------

# One row per (forecast, step). Flat rather than array-valued because every
# consumer -- the resolver, the standings, the chart -- wants steps as rows.
FORECAST_COLUMNS = [
    "forecast_id", "model_slug", "asset", "timeframe",
    "issued_ts", "target_ts", "step", "horizon_bars",
    "seed", "n_samples", "context_len", "inference_version",
    "dispersion", "path_ref", "backfilled",
] + [f"q{int(round(q * 100)):02d}" for q in config.QUANTILE_LEVELS]

QUANTILE_COLUMNS = [f"q{int(round(q * 100)):02d}" for q in config.QUANTILE_LEVELS]

# Resolved outcomes. Append-only, never revised.
TRACKRECORD_COLUMNS = [
    "forecast_id", "model_slug", "asset", "timeframe",
    "issued_ts", "target_ts", "step", "horizon_bars",
    "realized_close", "q10", "q50", "q90", "q20", "q80",
    "inside_80", "inside_60",
    "abs_error", "pct_error", "resolved_ts", "backfilled",
]

STANDINGS_COLUMNS = [
    "model_slug", "asset_class", "resolved_count", "forecast_count",
    "coverage_80", "coverage_60", "median_abs_pct_error",
    "calibration_gap", "grade", "updated_ts",
]


def forecast_id(model_slug: str, asset: str, timeframe: str,
                issued_ts, seed: int) -> str:
    """The dedup key, as a short stable hash.

    Deliberately *excludes* `inference_version`. Re-running the same request
    after a model upgrade must not append a second row: the first answer is the
    one that was issued, and the archive records what was issued. The version
    that produced it is kept as a column so the row is still traceable.
    """
    payload = "|".join([
        str(model_slug), str(asset), str(timeframe),
        _iso(issued_ts), str(int(seed)),
    ])
    return hashlib.sha256(payload.encode()).hexdigest()[:16]


def _utc(ts) -> pd.Timestamp:
    t = pd.Timestamp(ts)
    return t.tz_localize("UTC") if t.tzinfo is None else t.tz_convert("UTC")


def _iso(ts) -> str:
    return _utc(ts).isoformat()


def now_utc() -> pd.Timestamp:
    """A tz-aware "now".

    `pd.Timestamp.utcnow()` returns tz-aware on current pandas and tz-naive on
    older ones, so localising it unconditionally raises. Every "now" in the app
    comes through here so that difference is handled once.
    """
    return _utc(pd.Timestamp.utcnow())


# --------------------------------------------------------------------------
# Store
# --------------------------------------------------------------------------


class ArenaStore:
    """Parquet access to the Arena's trees in the signal store.

    Set `offline=True` (or leave `repo_id=None`) for a purely local store --
    what the tests and the backfill script use before they have a token.
    """

    def __init__(self, repo_id: str | None = config.STORE_REPO,
                 local_root: str | os.PathLike | None = None,
                 token: str | None = None, offline: bool = False,
                 revision: str = "main") -> None:
        self.repo_id = repo_id
        self.revision = revision
        self.token = token or os.environ.get("HF_WRITE_TOKEN") or os.environ.get("HF_TOKEN")
        self.offline = offline or repo_id is None
        self.local_root = Path(local_root
                               or os.environ.get("BIT_STORE_CACHE", ".cache/store"))
        self.local_root.mkdir(parents=True, exist_ok=True)
        self._pending: set[str] = set()
        self._price_cache: dict[tuple[str, str], pd.DataFrame] = {}
        self._flush_timer: threading.Timer | None = None
        self._lock = threading.RLock()

    # -- low-level --------------------------------------------------------

    def _local(self, repo_path: str) -> Path:
        return self.local_root / repo_path

    def _fetch(self, repo_path: str) -> Path | None:
        local = self._local(repo_path)
        if local.exists():
            return local
        if self.offline:
            return None
        try:
            from huggingface_hub import hf_hub_download

            return Path(hf_hub_download(
                repo_id=self.repo_id, repo_type=config.STORE_REPO_TYPE,
                filename=repo_path, revision=self.revision, token=self.token,
            ))
        except Exception:
            # Absent from the Hub is a normal "no data" answer, not a fault.
            return None

    def read_parquet(self, repo_path: str) -> pd.DataFrame | None:
        p = self._fetch(repo_path)
        if p is None:
            return None
        try:
            return pd.read_parquet(p)
        except Exception as e:  # pragma: no cover
            log.warning("unreadable parquet %s: %s", repo_path, e)
            return None

    def read_json(self, repo_path: str) -> dict | None:
        p = self._fetch(repo_path)
        if p is None:
            return None
        try:
            return json.loads(Path(p).read_text())
        except Exception as e:  # pragma: no cover
            log.warning("unreadable json %s: %s", repo_path, e)
            return None

    def _stage(self, repo_path: str, write) -> None:
        local = self._local(repo_path)
        local.parent.mkdir(parents=True, exist_ok=True)
        write(local)
        with self._lock:
            self._pending.add(repo_path)

    def write_parquet(self, repo_path: str, df: pd.DataFrame) -> None:
        self._stage(repo_path, lambda p: df.to_parquet(p, index=False))
        if repo_path.startswith("prices/"):
            self.invalidate_prices()

    def write_json(self, repo_path: str, payload: dict) -> None:
        self._stage(repo_path,
                    lambda p: Path(p).write_text(json.dumps(payload, indent=2, sort_keys=True)))

    @property
    def pending(self) -> list[str]:
        with self._lock:
            return sorted(self._pending)

    def schedule_flush(self, message: str = "arena: archive forecasts",
                       every: float = 30.0) -> None:
        """Commit staged files in the background, batched.

        Flushing inline put an HTTP commit on the end of every forecast --
        seconds of latency on a request whose actual work took milliseconds,
        and one commit per click in the store's history. This coalesces: the
        first call starts a timer, later calls within the window join the same
        commit.

        Deliberately not `huggingface_hub.CommitScheduler`, which commits an
        entire folder. This mirror also holds price files pulled *down* from
        the Hub, and a whole-folder scheduler would keep re-committing those.
        Only what was explicitly staged is ever pushed.
        """
        if self.offline:
            return
        with self._lock:
            if not self._pending or self._flush_timer is not None:
                return
            timer = threading.Timer(every, self._flush_now, args=(message,))
            timer.daemon = True
            self._flush_timer = timer
        timer.start()

    def _flush_now(self, message: str) -> None:
        with self._lock:
            self._flush_timer = None
        try:
            pushed = self.flush(message)
            if pushed:
                log.info("archived %d file(s) to the store", pushed)
        except Exception as e:      # pragma: no cover - network
            # The rows are already in the local mirror; the next flush retries.
            log.warning("could not push the archive: %s", e)

    def flush(self, message: str = "arena: update") -> int:
        """Push everything staged as one commit. Returns the file count."""
        with self._lock:
            paths = sorted(self._pending)
            if not paths or self.offline:
                self._pending.clear()
                return 0
        from huggingface_hub import CommitOperationAdd, HfApi

        ops = [CommitOperationAdd(path_in_repo=p, path_or_fileobj=str(self._local(p)))
               for p in paths]
        HfApi(token=self.token).create_commit(
            repo_id=self.repo_id, repo_type=config.STORE_REPO_TYPE,
            operations=ops, commit_message=message, revision=self.revision,
        )
        with self._lock:
            self._pending.clear()
        return len(ops)

    # -- prices -----------------------------------------------------------

    def get_prices(self, asset: str, timeframe: str,
                   start=None, end=None) -> pd.DataFrame:
        """OHLCV slice with a `ts` column, sorted ascending.

        Reads the existing price cache written by the shared pipeline; the
        Arena never fetches a price in the request path.

        The full series is memoised per (asset, timeframe). Without it the
        backfill re-read and re-concatenated ~9,600 rows of parquet for every
        one of its thousands of as-of forecasts, which dominated the run;
        `invalidate_prices` drops the memo after a refresh writes new bars.
        """
        key = (asset, timeframe)
        cached = self._price_cache.get(key)
        if cached is not None:
            return self._slice(cached, start, end)
        full = self._read_prices(asset, timeframe)
        self._price_cache[key] = full
        return self._slice(full, start, end)

    @staticmethod
    def _slice(frame: pd.DataFrame, start, end) -> pd.DataFrame:
        if start is None and end is None:
            return frame
        out = frame
        if start is not None:
            out = out[out["ts"] >= _utc(start)]
        if end is not None:
            out = out[out["ts"] <= _utc(end)]
        return out.reset_index(drop=True)

    def invalidate_prices(self, asset: str | None = None,
                          timeframe: str | None = None) -> None:
        if asset is None:
            self._price_cache.clear()
            return
        self._price_cache.pop((asset, timeframe), None)

    def _read_prices(self, asset: str, timeframe: str) -> pd.DataFrame:
        start = end = None
        years = _years_between(start, end)
        frames = []
        for y in years:
            df = self.read_parquet(config.prices_path(asset, timeframe, y))
            if df is not None and len(df):
                frames.append(df)
        if not frames:
            return pd.DataFrame(columns=["ts", "open", "high", "low", "close", "volume"])
        out = pd.concat(frames, ignore_index=True)
        out["ts"] = pd.to_datetime(out["ts"], utc=True)
        return out.drop_duplicates(subset=["ts"]).sort_values("ts").reset_index(drop=True)

    # -- forecasts (the archive) ------------------------------------------

    def get_forecasts(self, model_slug: str, asset: str, timeframe: str,
                      years: list[int] | None = None) -> pd.DataFrame:
        years = years or _years_between(None, None)
        frames = []
        for y in years:
            df = self.read_parquet(config.signals_path(model_slug, asset, timeframe, y))
            if df is not None and len(df):
                frames.append(df)
        if not frames:
            return pd.DataFrame(columns=FORECAST_COLUMNS)
        out = pd.concat(frames, ignore_index=True)
        out["issued_ts"] = pd.to_datetime(out["issued_ts"], utc=True)
        out["target_ts"] = pd.to_datetime(out["target_ts"], utc=True)
        return out

    def append_forecast(self, rows: pd.DataFrame) -> int:
        """Archive one forecast. Returns rows actually written (0 if a dup).

        First write wins. A forecast already present under this id is left
        exactly as it was -- see the module docstring.
        """
        if rows is None or not len(rows):
            return 0
        missing = [c for c in FORECAST_COLUMNS if c not in rows.columns]
        if missing:
            raise StoreError(f"forecast rows missing columns {missing}")

        fid = str(rows["forecast_id"].iloc[0])
        if rows["forecast_id"].nunique() != 1:
            raise StoreError("append_forecast takes one forecast at a time")

        model_slug = str(rows["model_slug"].iloc[0])
        asset = str(rows["asset"].iloc[0])
        tf = str(rows["timeframe"].iloc[0])
        year = int(_utc(rows["issued_ts"].iloc[0]).year)
        path = config.signals_path(model_slug, asset, tf, year)

        existing = self.read_parquet(path)
        if existing is not None and len(existing):
            if fid in set(existing["forecast_id"].astype(str)):
                return 0
            merged = pd.concat([existing, rows[FORECAST_COLUMNS]], ignore_index=True)
        else:
            merged = rows[FORECAST_COLUMNS].copy()

        merged = merged.sort_values(["issued_ts", "step"]).reset_index(drop=True)
        self.write_parquet(path, merged)
        return len(rows)

    # -- track record -----------------------------------------------------

    def get_trackrecord(self, model_slug: str, asset: str,
                        timeframe: str) -> pd.DataFrame:
        df = self.read_parquet(config.trackrecord_path(model_slug, asset, timeframe))
        if df is None or not len(df):
            return pd.DataFrame(columns=TRACKRECORD_COLUMNS)
        df["issued_ts"] = pd.to_datetime(df["issued_ts"], utc=True)
        df["target_ts"] = pd.to_datetime(df["target_ts"], utc=True)
        return df

    def append_trackrecord(self, model_slug: str, asset: str, timeframe: str,
                           rows: pd.DataFrame) -> int:
        """Append resolved outcomes. Idempotent on (forecast_id, step)."""
        if rows is None or not len(rows):
            return 0
        path = config.trackrecord_path(model_slug, asset, timeframe)
        existing = self.get_trackrecord(model_slug, asset, timeframe)
        if len(existing):
            seen = set(zip(existing["forecast_id"].astype(str),
                           existing["step"].astype(int)))
            mask = [(str(f), int(s)) not in seen
                    for f, s in zip(rows["forecast_id"], rows["step"])]
            rows = rows[np.asarray(mask)]
            if not len(rows):
                return 0
            merged = pd.concat([existing, rows[TRACKRECORD_COLUMNS]], ignore_index=True)
        else:
            merged = rows[TRACKRECORD_COLUMNS].copy()
        merged = merged.sort_values(["issued_ts", "step"]).reset_index(drop=True)
        self.write_parquet(path, merged)
        return len(rows)

    def all_trackrecords(self) -> pd.DataFrame:
        """Every resolved row in the store, for the standings rebuild.

        Walks the local mirror. The scheduled job syncs the tree first, so this
        sees the whole store; in the Space it sees whatever has been fetched,
        which is why standings are *read* there and never rebuilt.
        """
        root = self.local_root / "arena" / "trackrecord"
        if not root.exists():
            return pd.DataFrame(columns=TRACKRECORD_COLUMNS)
        frames = []
        for p in sorted(root.rglob("*.parquet")):
            try:
                frames.append(pd.read_parquet(p))
            except Exception as e:  # pragma: no cover
                log.warning("skipping unreadable trackrecord %s: %s", p, e)
        if not frames:
            return pd.DataFrame(columns=TRACKRECORD_COLUMNS)
        out = pd.concat(frames, ignore_index=True)
        out["issued_ts"] = pd.to_datetime(out["issued_ts"], utc=True)
        out["target_ts"] = pd.to_datetime(out["target_ts"], utc=True)
        return out

    # -- standings and registry -------------------------------------------

    def get_standings(self) -> pd.DataFrame:
        df = self.read_parquet(config.STANDINGS_PATH)
        if df is None:
            return pd.DataFrame(columns=STANDINGS_COLUMNS)
        return df

    def put_standings(self, df: pd.DataFrame) -> None:
        self.write_parquet(config.STANDINGS_PATH, df[STANDINGS_COLUMNS])

    def get_panels(self) -> dict:
        """Precomputed Track Record panels. See `trackrecord.build_panels`."""
        doc = self.read_json(config.PANELS_PATH)
        if not doc:
            return {"version": 1, "panels": {}, "totals": {}}
        doc.setdefault("panels", {})
        doc.setdefault("totals", {})
        return doc

    def put_panels(self, doc: dict) -> None:
        self.write_json(config.PANELS_PATH, doc)

    def get_registry(self) -> dict:
        reg = self.read_json(config.REGISTRY_PATH)
        if not reg:
            return {"version": 1, "models": {}}
        reg.setdefault("models", {})
        return reg

    def put_registry(self, registry: dict) -> None:
        self.write_json(config.REGISTRY_PATH, registry)


def _years_between(start, end) -> list[int]:
    """Year partitions to look in. Defaults to a window around now."""
    now = now_utc()
    s = _utc(start).year if start is not None else now.year - 3
    e = _utc(end).year if end is not None else now.year + 1
    return list(range(int(s), int(e) + 1))