"""
The one output format every model produces, and how it is scored.

A forecast is a set of quantiles of the natural-log price at each horizon.
Quantiles work for any model (closed-form, simulated, bootstrapped, mixtures)
and make scoring simple. Working in log price makes errors relative: a CRPS of
0.1 is roughly "typically 10% off", in 2013 or in 2026.
"""

from dataclasses import dataclass

import numpy as np
import pandas as pd
from scipy.stats import norm

N_LEVELS = 100
# Midpoints of 100 equal-probability bins: 0.005, 0.015, ..., 0.995.
LEVELS = (np.arange(N_LEVELS) + 0.5) / N_LEVELS


@dataclass(frozen=True)
class Forecast:
    origin: pd.Timestamp
    horizons: np.ndarray  # days after origin, shape (H,)
    log_quantiles: np.ndarray  # log price at LEVELS, shape (H, N_LEVELS), rows nondecreasing

    @classmethod
    def normal(cls, origin, horizons, mean, sd) -> "Forecast":
        """Normal distribution in log price (i.e. lognormal price) at each horizon."""
        mean = np.broadcast_to(np.asarray(mean, dtype=float), np.shape(horizons))
        sd = np.broadcast_to(np.asarray(sd, dtype=float), np.shape(horizons))
        q = mean[:, None] + sd[:, None] * norm.ppf(LEVELS)[None, :]
        return cls(pd.Timestamp(origin), np.asarray(horizons), q)

    @classmethod
    def from_samples(cls, origin, horizons, samples) -> "Forecast":
        """Empirical quantiles of simulated log prices, shape (n_samples, H)."""
        q = np.quantile(np.asarray(samples), LEVELS, axis=0).T
        return cls(pd.Timestamp(origin), np.asarray(horizons), q)

    @property
    def dates(self) -> pd.DatetimeIndex:
        return self.origin + pd.to_timedelta(self.horizons, unit="D")

    def quantile(self, level: float) -> np.ndarray:
        """Log price at an arbitrary level, interpolated between grid levels."""
        return np.array([np.interp(level, LEVELS, row) for row in self.log_quantiles])

    def interval(self, coverage: float) -> tuple[np.ndarray, np.ndarray]:
        """Central interval in log price holding `coverage` probability."""
        tail = (1 - coverage) / 2
        return self.quantile(tail), self.quantile(1 - tail)


def crps(log_quantiles: np.ndarray, outcome: np.ndarray) -> np.ndarray:
    """
    Continuous ranked probability score, from quantiles, in log-price units.

    CRPS is twice the pinball loss integrated over all quantile levels; the
    quantile grid gives the integral directly. It rewards sharpness and
    calibration together and has no free parameters to game. Lower is better.
    """
    outcome = np.asarray(outcome, dtype=float)
    u = outcome[..., None] - log_quantiles
    pinball = u * (LEVELS - (u < 0))
    return 2 * pinball.mean(axis=-1)


def pit(log_quantiles: np.ndarray, outcome: np.ndarray) -> np.ndarray:
    """Probability integral transform: forecast CDF evaluated at the outcome."""
    return np.array(
        [
            np.interp(y, q, LEVELS, left=0.0, right=1.0)
            for q, y in zip(np.atleast_2d(log_quantiles), np.atleast_1d(outcome), strict=True)
        ]
    )