"""GARCH(1,1) / EGARCH(1,1) with Student-t innovations on 5-minute training returns per instrument (arch package),
forecasting the variance over the next H minutes by aggregating the term structure of the fitted process on the
5-minute grid (H/5 steps; H < 5 scaled proportionally). Mean fixed at zero. Fitted once per instrument on the
training period; used out of sample with the last `n_days` days of 5-minute returns as filtering history."""
from __future__ import annotations

import warnings

import numpy as np

from ..data.windows import SymbolGrid


def five_minute_returns(g: SymbolGrid, s0: int, s1: int) -> np.ndarray:
    out = []
    for s in range(s0, s1 + 1):
        c = g.flat_b5[s * 78:(s + 1) * 78, 3]; m = g.flat_m5[s * 78:(s + 1) * 78]
        idx = np.flatnonzero(m)
        if len(idx) >= 2:
            out.append(np.diff(c[idx]))
    return np.concatenate(out) if out else np.zeros(0)


class GARCHBaseline:
    def __init__(self, kind: str = "garch", scale: float = 1e4):
        self.kind, self.scale = kind, scale
        self.models: dict[str, object] = {}
        self.name = kind

    def fit(self, symbol: str, r5: np.ndarray) -> bool:
        from arch import arch_model
        if len(r5) < 500:
            return False
        y = r5 * self.scale
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            if self.kind == "garch":
                am = arch_model(y, mean="Zero", vol="GARCH", p=1, q=1, dist="t")
            else:
                am = arch_model(y, mean="Zero", vol="EGARCH", p=1, o=1, q=1, dist="t")
            res = am.fit(disp="off", show_warning=False)
        self.models[symbol] = res
        return True

    def predict(self, symbol: str, r5_hist: np.ndarray, H: int) -> dict | None:
        res = self.models.get(symbol)
        if res is None or len(r5_hist) < 50:
            return None
        from arch import arch_model
        y = r5_hist * self.scale
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            if self.kind == "garch":
                am = arch_model(y, mean="Zero", vol="GARCH", p=1, q=1, dist="t")
            else:
                am = arch_model(y, mean="Zero", vol="EGARCH", p=1, o=1, q=1, dist="t")
            steps = max(1, int(np.ceil(H / 5)))
            fc = am.fix(res.params.values).forecast(horizon=steps, method="simulation" if self.kind != "garch" else "analytic",
                                                     simulations=500 if self.kind != "garch" else None, reindex=False)
        v5 = fc.variance.values[-1, :steps] / self.scale ** 2
        var = v5.sum() * (H / (5 * steps))
        nu = float(res.params.get("nu", 6.0)); nu = float(np.clip(nu, 2.5, 50))
        return {"mu": 0.0, "scale": np.sqrt(var * (nu - 2) / nu), "nu": nu, "var": var}