ensemble.py

"""Ensemble blending + calibration nudges for the FutureQuery forecaster.

This module is deliberately dependency-light (numpy only) so the blending
rules can be unit-tested without loading TimesFM / Chronos. ``main.py`` runs
the two models and hands their raw forecast arrays to :func:`full_response`;
the short-series / no-history guards live in :func:`short_series_response`.

Rules (per the Nov-2025 research notes baked into the task):

* < 10 points                     -> use_forecast = False, warning set.
* question_type == "event"        -> ensemble = 0.5*TimesFM + 0.5*Chronos-2
  (equal weight: the models showed negative R^2 on binary event markets, so
  we do not let either dominate).
* question_type == "numeric"      -> Chronos-2 0.6 / TimesFM 0.4 (Chronos-2
  beats TimesFM on fev-bench for numeric trajectories).
* base-rate prior                 -> if the latest price is within 0.1 of 0.5,
  pull the ensemble 15% back toward 0.5 (regression to the mean).
* extremise (events only)         -> if confidence > 0.70, sharpen with the
  standard superforecasting transform 0.5 + 1.3*(p - 0.5).
"""

from __future__ import annotations

from typing import Dict, List, Optional, Sequence

import numpy as np

# --- weights & thresholds (single source of truth) -----------------------
MIN_POINTS = 10
EVENT_WEIGHTS = (0.5, 0.5)        # (timesfm, chronos2)
NUMERIC_WEIGHTS = (0.4, 0.6)      # (timesfm, chronos2) -> Chronos-2 favoured
BASE_RATE_BAND = 0.1              # |last - 0.5| <= 0.1 triggers the prior
BASE_RATE_PULL = 0.15            # move 15% toward 0.5
EXTREMISE_GATE = 0.70            # only extremise confident forecasts
EXTREMISE_K = 1.3               # superforecasting extremisation factor
TREND_EPS = 0.02                 # 2 percentage points = flat band for events


def _clamp01(arr: np.ndarray) -> np.ndarray:
    return np.clip(arr, 0.0, 1.0)


def _trend(forecast: np.ndarray, reference: float, numeric: bool) -> str:
    """up / down / flat relative to the last observed value."""
    if forecast.size == 0:
        return "flat"
    delta = float(forecast[-1]) - float(reference)
    eps = max(TREND_EPS * abs(reference), 1e-9) if numeric else TREND_EPS
    if delta > eps:
        return "up"
    if delta < -eps:
        return "down"
    return "flat"


def _persistence(prices: np.ndarray, horizon: int, event: bool) -> np.ndarray:
    """Flat continuation used when we cannot really forecast."""
    if prices.size:
        return np.full(int(horizon), float(prices[-1]), dtype=float)
    return np.full(int(horizon), 0.5 if event else 0.0, dtype=float)


def _block(forecast: np.ndarray, reference: float, numeric: bool) -> Dict:
    return {
        "forecast": [round(float(x), 4) for x in forecast],
        "trend": _trend(forecast, reference, numeric),
    }


def blend(
    prices: Sequence[float],
    timesfm_fc: Sequence[float],
    chronos_fc: Sequence[float],
    question_type: str,
    horizon: int,
) -> Dict:
    """Blend two model forecasts into the ensemble (the heart of Phase 1)."""
    price_arr = np.asarray(prices, dtype=float)
    last = float(price_arr[-1]) if price_arr.size else 0.5
    event = question_type == "event"
    numeric = not event

    tf = np.asarray(timesfm_fc, dtype=float)
    ch = np.asarray(chronos_fc, dtype=float)
    n = min(tf.size, ch.size, int(horizon))
    if n == 0:
        n = int(horizon)
    tf = tf[:n]
    ch = ch[:n]

    if event:
        tf = _clamp01(tf)
        ch = _clamp01(ch)
        w_tf, w_ch = EVENT_WEIGHTS
    else:
        w_tf, w_ch = NUMERIC_WEIGHTS

    ensemble = w_tf * tf + w_ch * ch

    if event:
        # 1) base-rate prior: regress near-coin-flip forecasts toward 0.5
        if abs(last - 0.5) <= BASE_RATE_BAND:
            ensemble = ensemble + BASE_RATE_PULL * (0.5 - ensemble)
        # 2) extremise only confident forecasts
        headline = float(np.mean(ensemble))
        confidence = max(headline, 1.0 - headline)
        if confidence > EXTREMISE_GATE:
            ensemble = 0.5 + EXTREMISE_K * (ensemble - 0.5)
        ensemble = _clamp01(ensemble)

    brier = _brier_estimate(ensemble) if event else None

    return {
        "timesfm": _block(tf, last, numeric),
        "chronos2": _block(ch, last, numeric),
        "ensemble": {
            **_block(ensemble, last, numeric),
            "brier_estimate": brier,
        },
    }


def _brier_estimate(ensemble: np.ndarray) -> Optional[float]:
    """Expected Brier score if the outcome were Bernoulli(p): p*(1-p).

    This is a self-consistency proxy reported at forecast time (we don't yet
    know the resolution). It peaks at 0.25 for p=0.5 and shrinks as the
    forecast gets confident. The realised Brier score is tracked separately in
    calibration.py once markets resolve.
    """
    if ensemble.size == 0:
        return None
    p = float(np.clip(ensemble[-1], 0.0, 1.0))
    return round(p * (1.0 - p), 2)


def short_series_response(
    prices: Sequence[float],
    horizon: int,
    question_type: str,
    warning: str,
) -> Dict:
    """Schema-valid response when we refuse to forecast (too short / no data)."""
    event = question_type == "event"
    price_arr = np.asarray(prices, dtype=float)
    flat = _persistence(price_arr, horizon, event)
    last = float(price_arr[-1]) if price_arr.size else (0.5 if event else 0.0)
    block = _block(flat, last, not event)
    return {
        "timesfm": dict(block),
        "chronos2": dict(block),
        "ensemble": {**block, "brier_estimate": None},
        "use_forecast": False,
        "warning": warning,
    }


def full_response(
    prices: Sequence[float],
    timesfm_fc: Sequence[float],
    chronos_fc: Sequence[float],
    question_type: str,
    horizon: int,
) -> Dict:
    payload = blend(prices, timesfm_fc, chronos_fc, question_type, horizon)
    payload["use_forecast"] = True
    payload["warning"] = None
    return payload


def headline_probability(payload: Dict) -> Optional[float]:
    """Representative probability of an ensemble payload (last horizon point)."""
    forecast: List[float] = payload.get("ensemble", {}).get("forecast", [])
    if not forecast:
        return None
    return float(np.clip(forecast[-1], 0.0, 1.0))