"""Markov regime detection for hz.run() pipelines.

Wraps the Rust HMM (Hidden Markov Model) for real-time regime classification.
Each tick, calls filter_step() on the latest return -- O(N^2) where N is
typically 2-3, so effectively zero overhead.

Usage:
    # Train offline on historical returns, then use in live pipeline
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(historical_returns, max_iters=100)

    hz.run(
        pipeline=[
            markov_regime(model, feed="btc"),
            my_strategy,  # receives regime via ctx.params["regime"]
        ],
        ...
    )

    # Or train inline during backtest warmup
    hz.run(
        pipeline=[
            markov_regime(n_states=2, warmup=100, feed="btc"),
            my_strategy,
        ],
        ...
    )
"""

from __future__ import annotations

from collections import deque
from typing import Callable

from horizon._horizon import MarkovRegimeModel, prices_to_returns
from horizon.context import Context


def markov_regime(
    model: MarkovRegimeModel | None = None,
    n_states: int = 2,
    warmup: int = 100,
    feed: str | None = None,
    param_name: str = "regime",
) -> Callable[[Context], None]:
    """Create a pipeline function that classifies the current market regime.

    On each tick, computes the log return from the latest price, feeds it
    through the HMM forward filter, and injects regime info into ctx.params.

    Injected params:
        ctx.params[param_name] -> int: most likely state (0=calm, 1=volatile, ...)
        ctx.params[param_name + "_probs"] -> list[float]: state probabilities
        ctx.params[param_name + "_vol_state"] -> float: P(highest-vol state)

    Args:
        model: Pre-trained MarkovRegimeModel. If None, trains online after warmup.
        n_states: Number of states (only used if model is None).
        warmup: Ticks to collect before online training (only if model is None).
        feed: Feed name to read price from. None = first available feed.
        param_name: Key to inject into ctx.params.

    Returns:
        Pipeline function compatible with hz.run() and hz.backtest().
    """
    state = _MarkovState(
        model=model,
        n_states=n_states,
        warmup=warmup,
        feed=feed,
        param_name=param_name,
    )

    def _regime_fn(ctx: Context) -> None:
        state.update(ctx)

    return _regime_fn


class _MarkovState:
    """Internal state for the markov_regime pipeline function."""

    __slots__ = (
        "model", "n_states", "warmup", "feed", "param_name",
        "_prices", "_prev_price", "_trained",
    )

    def __init__(
        self,
        model: MarkovRegimeModel | None,
        n_states: int,
        warmup: int,
        feed: str | None,
        param_name: str,
    ):
        self.model = model
        self.n_states = n_states
        self.warmup = warmup
        self.feed = feed
        self.param_name = param_name
        self._prices: deque[float] = deque(maxlen=max(warmup + 10, 200))
        self._prev_price: float = 0.0
        self._trained = model is not None and model.is_trained()

    def update(self, ctx: Context) -> None:
        # Extract price from feed
        price = self._get_price(ctx)
        if price <= 0.0:
            return

        self._prices.append(price)

        # Need at least 2 prices for a return
        if self._prev_price <= 0.0:
            self._prev_price = price
            return

        # Compute log return
        log_return = 0.0
        if self._prev_price > 0.0:
            log_return = _safe_log(price / self._prev_price)
        self._prev_price = price

        # Auto-train if no pre-trained model and we have enough data
        if not self._trained and len(self._prices) >= self.warmup:
            self._auto_train()

        if not self._trained:
            return

        # Online filter step -- the hot path, O(N^2)
        probs = self.model.filter_step(log_return)
        regime = max(range(len(probs)), key=lambda i: probs[i])

        # Inject into context
        ctx.params[self.param_name] = regime
        ctx.params[self.param_name + "_probs"] = probs
        ctx.params[self.param_name + "_vol_state"] = probs[-1]  # highest-vol state

    def _get_price(self, ctx: Context) -> float:
        if self.feed is not None:
            fdata = ctx.feeds.get(self.feed)
            if fdata is not None and fdata.price > 0:
                return fdata.price
            return 0.0

        # First available feed
        for fdata in ctx.feeds.values():
            if fdata.price > 0:
                return fdata.price
        return 0.0

    def _auto_train(self) -> None:
        price_list = list(self._prices)
        if len(price_list) < 10:
            return

        returns = []
        for i in range(1, len(price_list)):
            if price_list[i - 1] > 0.0:
                returns.append(_safe_log(price_list[i] / price_list[i - 1]))
            else:
                returns.append(0.0)

        if len(returns) < 10:
            return

        self.model = MarkovRegimeModel(n_states=self.n_states)
        self.model.fit(returns, max_iters=50, tol=1e-6)
        self._trained = True


def _safe_log(x: float) -> float:
    """Safe natural log that handles edge cases."""
    if x <= 0.0:
        return 0.0
    import math
    return math.log(x)
