"""Belief-volatility surface tracker pipeline for hz.run().

Maintains per-market BeliefVolSurface instances that estimate "belief vol"
 - a tradable parameter analogous to implied vol in options.

Usage:
    hz.run(
        pipeline=[
            hz.belief_vol_tracker("polymarket"),
            hz.logit_market_maker(),
        ],
        ...
    )
"""

from __future__ import annotations

import time
from typing import Callable

from horizon._horizon import BeliefVolSurface
from horizon.context import Context


def belief_vol_tracker(
    feed_name: str | None = None,
    decay: float = 0.99,
    min_observations: int = 20,
) -> Callable[[Context], None]:
    """Create a pipeline function that tracks belief volatility per market.

    Injects into ctx.params:
       - "belief_vol": current belief volatility estimate
       - "jump_prob": estimated jump probability

    Args:
        feed_name: Feed to read price/bid/ask from.
        decay: EWMA decay factor (0.5-0.9999).
        min_observations: Minimum observations before EM estimation.

    Returns:
        Pipeline function: (Context) -> None
    """
    surfaces: dict[str, BeliefVolSurface] = {}

    def _tracker(ctx: Context) -> None:
        market_id = ctx.market.condition_id if ctx.market else ""
        if not market_id:
            return

        # Get or create surface for this market
        if market_id not in surfaces:
            surfaces[market_id] = BeliefVolSurface(decay=decay, min_observations=min_observations)
        surface = surfaces[market_id]

        # Extract feed data
        price = 0.0
        bid = 0.0
        ask = 0.0
        volume = 0.0

        if feed_name and feed_name in ctx.feeds:
            fd = ctx.feeds[feed_name]
            price = fd.price
            bid = fd.bid
            ask = fd.ask
            volume = getattr(fd, "volume", 0.0)
        elif ctx.feeds:
            fd = next(iter(ctx.feeds.values()))
            price = fd.price
            bid = fd.bid
            ask = fd.ask
            volume = getattr(fd, "volume", 0.0)

        if price <= 0.0 or price >= 1.0:
            return

        timestamp = time.time()
        vol = surface.update(price, timestamp, bid, ask, volume)

        ctx.params["belief_vol"] = vol
        ctx.params["jump_prob"] = surface.jump_probability()

    _tracker.__name__ = "belief_vol_tracker"
    return _tracker
