"""Non-neural probabilistic baselines on the same origins and information set as the neural models.
All produce a Student-t (or Gaussian) predictive distribution for the cumulative H-minute log return and a variance
forecast for the realised variance over the horizon. Volatility estimators use only completed bars at the origin."""
from __future__ import annotations

import numpy as np

from ..data.windows import SymbolGrid, N_DAY


def _hist_returns(g: SymbolGrid, s: int, m: int, n_days: int) -> tuple[np.ndarray, np.ndarray]:
    """Per-minute returns (same-session, present pairs) for the last n_days sessions up to (s, m), with slot ids."""
    rs, slots = [], []
    for d in range(max(0, s - n_days + 1), s + 1):
        last = m if d == s else N_DAY - 1
        lc = g.flat_lc[d * N_DAY: d * N_DAY + last + 1]
        ok = ~np.isnan(lc)
        idx = np.flatnonzero(ok)
        if len(idx) < 2:
            continue
        r = np.diff(lc[idx]); rs.append(r); slots.append(idx[1:])
    if not rs:
        return np.zeros(0), np.zeros(0, dtype=int)
    return np.concatenate(rs), np.concatenate(slots)


class NaiveRandomWalk:
    """Zero mean; scale from the trailing RMS of 1-minute returns in the last `window` minutes; Gaussian."""
    name = "naive_rw"

    def __init__(self, window: int = 390, nu: float | None = None):
        self.window, self.nu = window, nu

    def predict(self, g: SymbolGrid, s: int, m: int, H: int, sigma: float) -> dict:
        var = sigma ** 2 * H
        return {"mu": 0.0, "scale": np.sqrt(var), "nu": self.nu or np.inf, "var": var}


class SeasonalVol:
    """Zero mean; variance = sum over future slots of the past-only time-of-day mean squared return
    (last `n_days` sessions), scaled by today's level ratio (trailing 30-min RV vs its seasonal counterpart)."""
    name = "seasonal_vol"

    def __init__(self, n_days: int = 20):
        self.n_days = n_days

    def predict(self, g: SymbolGrid, s: int, m: int, H: int, sigma: float) -> dict:
        r, sl = _hist_returns(g, s - 1, N_DAY - 1, self.n_days)
        if len(r) < 100:
            return NaiveRandomWalk().predict(g, s, m, H, sigma)
        prof = np.bincount(sl, weights=r * r, minlength=N_DAY) / np.maximum(1, np.bincount(sl, minlength=N_DAY))
        prof = np.where(prof > 0, prof, np.mean(r * r))
        fut = prof[m + 1: m + 1 + H].sum()
        rt, st = _hist_returns(g, s, m, 1)
        recent = rt[-30:]
        level = (np.mean(recent ** 2) / max(prof[max(0, m - 29): m + 1].mean(), 1e-12)) if len(recent) >= 10 else 1.0
        level = float(np.clip(level, 0.25, 4.0))
        var = fut * level
        return {"mu": 0.0, "scale": np.sqrt(var), "nu": np.inf, "var": var}


class EWMAVol:
    """RiskMetrics-style EWMA of squared 1-minute returns (lambda on the minute scale), times H, Student-t(nu)."""
    name = "ewma"

    def __init__(self, lam: float = 0.97, nu: float = 5.0, n_days: int = 3):
        self.lam, self.nu, self.n_days = lam, nu, n_days

    def predict(self, g: SymbolGrid, s: int, m: int, H: int, sigma: float) -> dict:
        r, _ = _hist_returns(g, s, m, self.n_days)
        if len(r) < 30:
            return NaiveRandomWalk().predict(g, s, m, H, sigma)
        w = self.lam ** np.arange(len(r))[::-1]
        v1 = float((w * r * r).sum() / w.sum())
        var = v1 * H
        scale = np.sqrt(var * (self.nu - 2) / self.nu)
        return {"mu": 0.0, "scale": scale, "nu": self.nu, "var": var}