"""
Volatility Clustering Model — model how volatility clusters over time.

Improves risk management by predicting volatility regime changes.
"""
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional

import numpy as np
from numba import njit


@njit(cache=True)
def compute_volatility_regime(
    current_vol: float,
    long_term_vol: float,
    vol_of_vol: float,
    recent_returns: np.ndarray,
) -> float:
    """
    Compute volatility regime score (0-1).

    Model:
      - High current vol relative to long-term → regime = 1
      - Low current vol relative to long-term → regime = 0
      - vol_of_vol adjusts sensitivity

    Returns regime score (0=low vol, 1=high vol).
    """
    if long_term_vol <= 0:
        return 0.5

    vol_ratio = current_vol / long_term_vol
    # Sigmoid mapping: vol_ratio=1 → 0.5, vol_ratio>1 → >0.5, vol_ratio<1 → <0.5
    regime = 1.0 / (1.0 + math.exp(-2.0 * (vol_ratio - 1.0)))
    return regime


@njit(cache=True)
def predict_volatility(
    current_vol: float,
    long_term_vol: float,
    vol_of_vol: float,
    time_horizon_s: float,
    mean_reversion_rate: float = 0.05,
) -> float:
    """
    Predict volatility after time_horizon_s.

    Model: GARCH-like mean reversion toward long-term volatility.
    """
    if long_term_vol <= 0:
        return current_vol

    # Mean reversion toward long-term
    predicted = current_vol + (long_term_vol - current_vol) * (1 - math.exp(-mean_reversion_rate * time_horizon_s))

    # Add vol-of-vol noise
    noise = vol_of_vol * math.sqrt(time_horizon_s / 86400.0)  # annualized
    predicted += noise * (2.0 * ((hash(str(current_vol)) % 1000) / 1000.0) - 1.0)

    return max(0.001, predicted)


class VolatilityClusteringModel:
    """
    Volatility clustering model for the CWM.

    Tracks volatility regime and predicts future volatility.
    """

    def __init__(self) -> None:
        self._vol_history: list[float] = []
        self._long_term_vol: float = 15.0  # default
        self._vol_of_vol: float = 5.0  # default

    def update(self, volatility: float) -> None:
        """Update with current volatility."""
        self._vol_history.append(volatility)
        if len(self._vol_history) > 1000:
            self._vol_history = self._vol_history[-500:]
        # Update long-term estimate
        if len(self._vol_history) > 50:
            self._long_term_vol = sum(self._vol_history[-200:]) / len(self._vol_history[-200:])

    def regime(self) -> float:
        """Get current volatility regime (0=low, 1=high)."""
        if not self._vol_history:
            return 0.5
        current = self._vol_history[-1]
        return compute_volatility_regime(current, self._long_term_vol, self._vol_of_vol, np.array([]))

    def predict(self, time_horizon_s: float = 60.0) -> float:
        """Predict volatility after time_horizon_s."""
        if not self._vol_history:
            return self._long_term_vol
        current = self._vol_history[-1]
        return predict_volatility(current, self._long_term_vol, self._vol_of_vol, time_horizon_s)

    @property
    def current_volatility(self) -> float:
        return self._vol_history[-1] if self._vol_history else 0.0

    @property
    def long_term_volatility(self) -> float:
        return self._long_term_vol

    @property
    def vol_of_vol(self) -> float:
        if len(self._vol_history) < 20:
            return 0.0
        recent = self._vol_history[-50:]
        mean = sum(recent) / len(recent)
        variance = sum((x - mean) ** 2 for x in recent) / len(recent)
        return math.sqrt(variance)