"""Regime detection engine: HMM-based state classification with trend/mean-reversion overlay."""

from __future__ import annotations

import logging
import math
import threading
import time
from collections import deque
from dataclasses import dataclass

from horizon._horizon import MarkovRegimeModel

logger = logging.getLogger("horizon.fund.regime")

_MIN_OBS_FOR_FIT = 30
_TREND_WINDOW = 20
_TREND_THRESHOLD = 0.3
_MEAN_REV_THRESHOLD = -0.3


@dataclass(frozen=True)
class RegimeSnapshot:
    """Point-in-time regime classification."""

    regime: str
    confidence: float
    volatility: float
    trend_strength: float
    mean_reversion_score: float
    timestamp: float


class RegimeDetector:
    """Wraps Rust MarkovRegimeModel (3-state HMM) with trend/mean-reversion overlay.

    HMM states are mapped to calm/volatile/crisis by sorting emission
    parameters by variance. On top, OLS trend strength and lag-1
    autocorrelation produce combined labels: trending_up, trending_down,
    mean_reverting, volatile, quiet.
    """

    def __init__(self, lookback: int = 500) -> None:
        self._model = MarkovRegimeModel(n_states=3)
        self._returns: deque[float] = deque(maxlen=lookback)
        self._history: deque[RegimeSnapshot] = deque(maxlen=10_000)
        self._lock = threading.Lock()
        self._current: RegimeSnapshot | None = None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def update(self, returns: list[float]) -> RegimeSnapshot:
        """Feed a batch of returns, refit HMM if enough data, classify regime.

        Args:
            returns: New return observations.

        Returns:
            Current regime snapshot.
        """
        with self._lock:
            for r in returns:
                if not math.isfinite(r):
                    continue
                self._returns.append(r)

            obs = list(self._returns)
            n = len(obs)

            # Defaults before we have enough data
            base_state = "quiet"
            confidence = 0.0
            volatility = 0.0
            trend_strength = 0.0
            mr_score = 0.0

            if n >= _MIN_OBS_FOR_FIT:
                try:
                    self._model.fit(obs, max_iters=100, tol=1e-6)
                except Exception:
                    logger.debug("HMM fit failed, using fallback regime")

                if self._model.is_trained():
                    # Filter on latest observation to get state probabilities
                    probs = self._model.filter_step(obs[-1])
                    current_state = int(self._model.current_regime())
                    confidence = probs[current_state] if current_state < len(probs) else 0.0

                    # Map HMM states to semantic labels by sorting emission variance
                    emissions = self._model.emission_params()  # list[(mean, std)]
                    indexed = list(enumerate(emissions))
                    indexed.sort(key=lambda x: x[1][1])  # sort by std ascending
                    # indexed[0] = calm, indexed[1] = moderate, indexed[2] = crisis
                    state_map = {}
                    for rank, (idx, _) in enumerate(indexed):
                        if rank == 0:
                            state_map[idx] = "calm"
                        elif rank == 1:
                            state_map[idx] = "moderate"
                        else:
                            state_map[idx] = "crisis"

                    base_state = state_map.get(current_state, "quiet")

                    # Current volatility from emission std of active state
                    volatility = emissions[current_state][1] if current_state < len(emissions) else 0.0

                    # Median volatility across states
                    vols = sorted(e[1] for e in emissions)
                    median_vol = vols[len(vols) // 2] if vols else 0.0
                else:
                    volatility = _std(obs)

            # Trend strength and mean-reversion score from recent returns
            if n >= _TREND_WINDOW:
                tail = obs[-_TREND_WINDOW:]
                trend_strength = _ols_slope(tail)
                mr_score = _lag1_autocorrelation(tail)
            elif n >= 5:
                tail = obs[-n:]
                trend_strength = _ols_slope(tail)
                mr_score = _lag1_autocorrelation(tail)

            # Combined classification
            regime = self._classify(
                base_state, volatility,
                _median_vol(self._model) if (n >= _MIN_OBS_FOR_FIT and self._model.is_trained()) else 0.0,
                trend_strength, mr_score,
            )

            snap = RegimeSnapshot(
                regime=regime,
                confidence=confidence,
                volatility=volatility,
                trend_strength=trend_strength,
                mean_reversion_score=mr_score,
                timestamp=time.time(),
            )
            self._current = snap
            self._history.append(snap)
            return snap

    def current_regime(self) -> RegimeSnapshot | None:
        """Return latest regime snapshot, or None if never updated."""
        with self._lock:
            return self._current

    def regime_history(self, limit: int = 50) -> list[RegimeSnapshot]:
        """Return recent regime snapshots."""
        with self._lock:
            return list(self._history)[-limit:]

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    @staticmethod
    def _classify(
        base_state: str,
        volatility: float,
        median_vol: float,
        trend_strength: float,
        mr_score: float,
    ) -> str:
        if base_state == "crisis" or (median_vol > 0 and volatility > 2.0 * median_vol):
            return "volatile"
        if abs(trend_strength) > _TREND_THRESHOLD:
            return "trending_up" if trend_strength > 0 else "trending_down"
        if mr_score < _MEAN_REV_THRESHOLD:
            return "mean_reverting"
        return "quiet"


# ------------------------------------------------------------------
# Pure helpers (no state)
# ------------------------------------------------------------------

def _ols_slope(returns: list[float]) -> float:
    """OLS slope of returns normalized by std. Closed-form linear regression."""
    n = len(returns)
    if n < 2:
        return 0.0
    s = _std(returns)
    if s < 1e-12:
        return 0.0
    mean_i = (n - 1) / 2.0
    mean_r = sum(returns) / n
    num = sum((i - mean_i) * (r - mean_r) for i, r in enumerate(returns))
    den = sum((i - mean_i) ** 2 for i in range(n))
    if den < 1e-12:
        return 0.0
    return (num / den) / s


def _lag1_autocorrelation(returns: list[float]) -> float:
    """Lag-1 autocorrelation of a return series."""
    n = len(returns)
    if n < 3:
        return 0.0
    mean = sum(returns) / n
    var = sum((r - mean) ** 2 for r in returns)
    if var < 1e-12:
        return 0.0
    cov = sum((returns[i] - mean) * (returns[i + 1] - mean) for i in range(n - 1))
    return cov / var


def _std(xs: list[float]) -> float:
    """Population standard deviation."""
    n = len(xs)
    if n < 2:
        return 0.0
    m = sum(xs) / n
    return math.sqrt(sum((x - m) ** 2 for x in xs) / n)


def _median_vol(model: MarkovRegimeModel) -> float:
    """Median emission std across HMM states."""
    if not model.is_trained():
        return 0.0
    vols = sorted(e[1] for e in model.emission_params())
    return vols[len(vols) // 2] if vols else 0.0
