"""Volatility modeling for prediction market prices."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pandas as pd
from arch import arch_model


@dataclass
class GARCHResult:
    model_type: str
    params: pd.Series
    std_errors: pd.Series
    log_likelihood: float
    aic: float
    bic: float
    conditional_volatility: pd.Series


@dataclass
class HARResult:
    coefficients: pd.Series
    std_errors: pd.Series
    t_stats: pd.Series
    r_squared: float
    adj_r_squared: float


def compute_realized_volatility_measures(
    returns: pd.Series,
    windows: list[int] | None = None,
) -> pd.DataFrame:
    """Compute realized volatility, bipower variation, and jump component."""
    if windows is None:
        windows = [1, 5, 22]
    result = pd.DataFrame(index=returns.index)

    for w in windows:
        result[f"rv_{w}"] = (returns ** 2).rolling(w).sum().apply(np.sqrt)

    abs_ret = returns.abs()
    result["bv_1"] = (abs_ret * abs_ret.shift(1)).rolling(5).mean() * (np.pi / 2)
    result["jump_1"] = np.maximum(result["rv_5"] ** 2 - result["bv_1"], 0)

    return result


def estimate_har_rv(
    rv_series: pd.DataFrame,
    exog: pd.DataFrame | None = None,
) -> HARResult:
    """HAR-RV model (Corsi 2009).

    RV_{t+1} = c + b_d * RV_t + b_w * RV_t^(w) + b_m * RV_t^(m) + e
    """
    import statsmodels.api as sm

    if "rv_1" not in rv_series.columns:
        raise ValueError("rv_series must contain 'rv_1' column")

    rv_d = rv_series["rv_1"]
    rv_w = rv_series.get("rv_5", rv_d.rolling(5).mean())
    rv_m = rv_series.get("rv_22", rv_d.rolling(22).mean())

    y = rv_d.shift(-1)
    X = pd.DataFrame({"rv_d": rv_d, "rv_w": rv_w, "rv_m": rv_m})
    if exog is not None:
        for col in exog.columns:
            X[col] = exog[col]

    X = sm.add_constant(X)
    valid = pd.concat([y, X], axis=1).dropna()
    if len(valid) < 10:
        return HARResult(
            coefficients=pd.Series(), std_errors=pd.Series(),
            t_stats=pd.Series(), r_squared=np.nan, adj_r_squared=np.nan,
        )

    y_clean = valid.iloc[:, 0]
    X_clean = valid.iloc[:, 1:]

    model = sm.OLS(y_clean, X_clean).fit(cov_type="HAC", cov_kwds={"maxlags": 5})

    return HARResult(
        coefficients=model.params,
        std_errors=model.bse,
        t_stats=model.tvalues,
        r_squared=model.rsquared,
        adj_r_squared=model.rsquared_adj,
    )


def estimate_garch(
    returns: pd.Series,
    model_type: str = "GARCH",
    p: int = 1,
    q: int = 1,
    dist: str = "normal",
) -> GARCHResult:
    """Estimate GARCH-family model using the arch package."""
    clean = returns.dropna()
    scaled = clean * 100

    vol_spec = model_type
    if model_type == "GJR-GARCH":
        vol_spec = "GARCH"
        o = 1
    else:
        o = 0

    am = arch_model(scaled, vol=vol_spec, p=p, o=o, q=q, dist=dist, mean="Constant")
    result = am.fit(disp="off")

    cond_vol = result.conditional_volatility / 100
    cond_vol.index = clean.index

    return GARCHResult(
        model_type=model_type,
        params=result.params,
        std_errors=result.std_err,
        log_likelihood=result.loglikelihood,
        aic=result.aic,
        bic=result.bic,
        conditional_volatility=cond_vol,
    )