src/comparisons.py

"""Derived comparison tables written into the store's `comparisons/` folder.

These are precomputed so the UI can render the Comparison tab instantly instead
of running inference or a sweep of backtests per page load. They are
regenerated after any coverage extension, so they never describe a stale store.

Four artefacts:

* `model_performance.parquet` — per (model, asset, timeframe): OOS Sharpe,
  return and drawdown under the *default* Forecast Follower rule.
* `calibration.parquet` — empirical coverage of the q10-q90 band, plus pinball
  losses. Answers "when this model says it is 80% sure, is it?".
* `directional_accuracy.parquet` — the model against three naive baselines.
* `strategy_timeframe_heatmap.parquet` — the strategy x timeframe OOS-Sharpe
  matrix behind the Comparison tab's time-scale grid.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np
import pandas as pd

from . import config, strategies
from .engine import BacktestConfig, Costs, Validation, run_backtest
from .metrics import (
    calibration_coverage,
    calibration_error,
    directional_accuracy,
    pinball_loss,
)
from .store import SignalStore

log = logging.getLogger("bit.comparisons")

MODEL_PERF = "comparisons/model_performance.parquet"
CALIBRATION = "comparisons/calibration.parquet"
DIRECTIONAL = "comparisons/directional_accuracy.parquet"
HEATMAP = "comparisons/strategy_timeframe_heatmap.parquet"
INDEX_JSON = "comparisons/index.json"

# Strategies shown in the time-scale matrix.
HEATMAP_STRATEGIES = (
    "Buy & Hold (benchmark)",
    "SMA Crossover",
    "RSI Mean Reversion",
    "Bollinger Breakout",
    "MACD Momentum",
    "Sentiment-Gated Momentum",
    "Chronos Forecast Follower",
)


def _default_config(asset: str, timeframe: str, strategy: str) -> BacktestConfig:
    """The single canonical config every comparison number is computed under.

    Costs on, walk-forward validation, six-month locked holdout. Changing this
    changes every published comparison, so it lives in one place.
    """
    return BacktestConfig(
        asset=asset, timeframe=timeframe, strategy=strategy,
        params=strategies.defaults_for(strategy),
        costs=Costs(enabled=True),
        validation=Validation(mode="walk_forward", train_months=12,
                              test_months=3, roll_months=3, holdout_months=6),
    )


@dataclass
class ComparisonReport:
    model_performance: pd.DataFrame
    calibration: pd.DataFrame
    directional: pd.DataFrame
    heatmap: pd.DataFrame

    def is_empty(self) -> bool:
        return all(df.empty for df in
                   (self.model_performance, self.calibration, self.directional, self.heatmap))


# --------------------------------------------------------------------------
# Calibration & directional accuracy
# --------------------------------------------------------------------------


def calibration_row(prices: pd.DataFrame, signals: pd.DataFrame,
                    model_slug: str, asset: str, timeframe: str) -> dict | None:
    """Compare each forecast against the outcome it was predicting.

    The forecast stored at `t` is a one-step-ahead prediction, so it is scored
    against the close at `t+1`, never against the close at `t`.
    """
    if signals.empty or prices.empty:
        return None
    close = prices["close"]
    aligned = signals.reindex(close.index).dropna(subset=["q50"])
    if aligned.empty:
        return None

    actual_next = close.shift(-1).reindex(aligned.index)
    valid = actual_next.notna()
    if valid.sum() < 20:
        return None

    actual_next = actual_next[valid]
    q10 = aligned.loc[valid.index[valid], "q10"]
    q50 = aligned.loc[valid.index[valid], "q50"]
    q90 = aligned.loc[valid.index[valid], "q90"]

    coverage = calibration_coverage(actual_next, q10, q90)
    return {
        "model_slug": model_slug, "asset": asset, "timeframe": timeframe,
        "n": int(valid.sum()),
        "coverage_q10_q90": coverage,
        "nominal_coverage": 0.80,
        "calibration_error": calibration_error(coverage, 0.80),
        "pinball_q10": pinball_loss(actual_next, q10, 0.10),
        "pinball_q50": pinball_loss(actual_next, q50, 0.50),
        "pinball_q90": pinball_loss(actual_next, q90, 0.90),
        "is_placeholder": bool(
            (aligned["inference_version"] == config.PLACEHOLDER_VERSION).any()
            if "inference_version" in aligned.columns else False
        ),
    }


def directional_row(prices: pd.DataFrame, signals: pd.DataFrame,
                    model_slug: str, asset: str, timeframe: str,
                    seed: int = 0) -> dict | None:
    """Model directional hit-rate against three naive baselines.

    Baselines: coin-flip (seeded, so the table is reproducible), momentum
    (continue the last move), and yesterday's-move repeated. All are computed
    causally from data available at the decision bar.
    """
    if signals.empty or prices.empty:
        return None
    close = prices["close"]
    aligned = signals.reindex(close.index).dropna(subset=["q50"])
    if aligned.empty:
        return None

    ref = close.reindex(aligned.index)
    actual_next = close.shift(-1).reindex(aligned.index)
    mask = actual_next.notna() & ref.notna()
    if mask.sum() < 20:
        return None

    ref, actual_next = ref[mask], actual_next[mask]
    q50 = aligned.loc[mask.index[mask], "q50"]

    prev_move = close.diff().reindex(ref.index).fillna(0.0)
    rng = np.random.default_rng(seed)
    coin = pd.Series(rng.choice([-1.0, 1.0], size=len(ref)), index=ref.index)

    return {
        "model_slug": model_slug, "asset": asset, "timeframe": timeframe,
        "n": int(mask.sum()),
        "model_accuracy": directional_accuracy(actual_next, q50, ref),
        "baseline_random": directional_accuracy(actual_next, ref + coin, ref),
        "baseline_momentum": directional_accuracy(actual_next, ref + prev_move, ref),
        "baseline_yesterday_move": directional_accuracy(
            actual_next, ref + prev_move.shift(1).fillna(0.0), ref
        ),
        "is_placeholder": bool(
            (aligned["inference_version"] == config.PLACEHOLDER_VERSION).any()
            if "inference_version" in aligned.columns else False
        ),
    }


# --------------------------------------------------------------------------
# Backtest-derived tables
# --------------------------------------------------------------------------


def _run(strategy: str, prices: pd.DataFrame, signals: pd.DataFrame | None,
         asset: str, timeframe: str):
    cfg = _default_config(asset, timeframe, strategy)
    out = strategies.build(strategy, prices, cfg.params, signals)
    return run_backtest(prices, out, cfg,
                        bars_per_year=config.bars_per_year(asset, timeframe))


def model_performance_row(prices, signals, model_slug, asset, timeframe) -> dict | None:
    """OOS performance of the default Forecast Follower rule over this model."""
    if signals is None or signals.empty:
        return None
    try:
        res = _run("Chronos Forecast Follower", prices, signals, asset, timeframe)
    except Exception as e:
        log.warning("forecast-follower run failed for %s/%s/%s: %s",
                    model_slug, asset, timeframe, e)
        return None

    m = res.metrics_oos
    return {
        "model_slug": model_slug, "asset": asset, "timeframe": timeframe,
        "strategy": "Chronos Forecast Follower",
        "oos_sharpe": m.sharpe, "oos_return": m.total_return,
        "oos_max_drawdown": m.max_drawdown, "oos_sortino": m.sortino,
        "trades": m.trade_count, "win_rate": m.win_rate,
        "costs_paid": res.costs_paid,
        "holdout_sharpe": res.metrics_holdout.sharpe if res.metrics_holdout else float("nan"),
        "is_sharpe": res.metrics_is.sharpe,
        "oos_is_ratio": (m.sharpe / res.metrics_is.sharpe)
                        if res.metrics_is.sharpe not in (0.0, None) else float("nan"),
        "is_placeholder": bool(
            (signals["inference_version"] == config.PLACEHOLDER_VERSION).any()
            if "inference_version" in signals.columns else False
        ),
    }


def heatmap_rows(store: SignalStore, asset: str, model_slug: str | None = None) -> list[dict]:
    """OOS Sharpe for every (strategy, timeframe) pair that has price coverage."""
    rows = []
    for tf in config.TIMEFRAMES:
        prices = store.get_prices(asset, tf)
        if len(prices) < 120:
            continue
        signals = (store.get_signals(model_slug, asset, tf)
                   if model_slug else pd.DataFrame())
        for strategy in HEATMAP_STRATEGIES:
            preset = strategies.PRESETS.get(strategy)
            if preset is None or not preset.available:
                continue
            if preset.needs_signals and (signals is None or signals.empty):
                rows.append({"asset": asset, "strategy": strategy, "timeframe": tf,
                             "oos_sharpe": float("nan"), "trades": 0,
                             "status": "no signal coverage"})
                continue
            try:
                res = _run(strategy, prices, signals, asset, tf)
                rows.append({
                    "asset": asset, "strategy": strategy, "timeframe": tf,
                    "oos_sharpe": res.metrics_oos.sharpe,
                    "oos_return": res.metrics_oos.total_return,
                    "trades": res.metrics_oos.trade_count,
                    "status": "ok",
                })
            except Exception as e:
                log.warning("heatmap cell failed %s/%s/%s: %s", asset, strategy, tf, e)
                rows.append({"asset": asset, "strategy": strategy, "timeframe": tf,
                             "oos_sharpe": float("nan"), "trades": 0,
                             "status": f"error: {type(e).__name__}"})
    return rows


# --------------------------------------------------------------------------
# Regeneration
# --------------------------------------------------------------------------


def regenerate(store: SignalStore, *, assets: list[str] | None = None,
               write: bool = True) -> ComparisonReport:
    """Rebuild every comparison table from what the store currently holds."""
    manifest = store.load_manifest()
    assets = assets or sorted({e.asset for e in manifest.signals.values()}
                              | {p.asset for p in manifest.prices.values()})

    perf, calib, direc, heat = [], [], [], []

    for entry in manifest.signals.values():
        if assets and entry.asset not in assets:
            continue
        prices = store.get_prices(entry.asset, entry.timeframe)
        signals = store.get_signals(entry.model_slug, entry.asset, entry.timeframe)
        if prices.empty or signals.empty:
            continue

        r = model_performance_row(prices, signals, entry.model_slug,
                                  entry.asset, entry.timeframe)
        if r:
            perf.append(r)
        r = calibration_row(prices, signals, entry.model_slug, entry.asset, entry.timeframe)
        if r:
            calib.append(r)
        r = directional_row(prices, signals, entry.model_slug, entry.asset, entry.timeframe)
        if r:
            direc.append(r)

    for asset in assets:
        best_model = None
        candidates = manifest.find_signals(asset=asset)
        if candidates:
            best_model = candidates[0].model_slug
        heat.extend(heatmap_rows(store, asset, best_model))

    report = ComparisonReport(
        model_performance=pd.DataFrame(perf),
        calibration=pd.DataFrame(calib),
        directional=pd.DataFrame(direc),
        heatmap=pd.DataFrame(heat),
    )

    if write:
        if not report.model_performance.empty:
            store.write_table(MODEL_PERF, report.model_performance)
        if not report.calibration.empty:
            store.write_table(CALIBRATION, report.calibration)
        if not report.directional.empty:
            store.write_table(DIRECTIONAL, report.directional)
        if not report.heatmap.empty:
            store.write_table(HEATMAP, report.heatmap)
        store.write_json(INDEX_JSON, {
            "generated_at": pd.Timestamp.now(tz="UTC").isoformat(),
            "assets": list(assets),
            "tables": {
                "model_performance": {"path": MODEL_PERF, "rows": len(report.model_performance)},
                "calibration": {"path": CALIBRATION, "rows": len(report.calibration)},
                "directional_accuracy": {"path": DIRECTIONAL, "rows": len(report.directional)},
                "strategy_timeframe_heatmap": {"path": HEATMAP, "rows": len(report.heatmap)},
            },
            "default_rule": "Chronos Forecast Follower, costs on, walk-forward 12/3/3, 6mo holdout",
        })
    return report


def load_table(store: SignalStore, path: str) -> pd.DataFrame:
    df = store.read_parquet(path)
    return df if df is not None else pd.DataFrame()