"""Particle Filter pipeline functions.

Pipeline functions:
    particle_tracker - Track latent state with particle filter
"""

from __future__ import annotations

from typing import Any, Callable

from horizon._horizon import ParticleFilter
from horizon.context import Context


def particle_tracker(
    feed: str,
    n_particles: int = 1000,
    process_noise: float = 0.01,
    measurement_noise: float = 0.05,
) -> Callable[[Context], dict[str, Any]]:
    """Create a particle filter tracking pipeline function.

    Tracks a latent state variable from noisy feed observations
    using Sequential Monte Carlo methods.

    Args:
        feed: Feed name to read price data from.
        n_particles: Number of particles.
        process_noise: State transition noise std dev.
        measurement_noise: Observation noise std dev.

    Returns:
        Pipeline function: (Context) -> dict with keys:
            estimate, variance, ess, resampled
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    filters: dict[str, ParticleFilter] = {}

    def _particle_tracker(ctx: Context) -> dict[str, Any]:
        market_id = ctx.market.id if ctx.market else "__default__"
        fd = ctx.feeds.get(feed)

        result = {
            "estimate": 0.0,
            "variance": 0.0,
            "ess": 0.0,
            "resampled": False,
        }

        if fd is not None and fd.price > 0:
            if market_id not in filters:
                filters[market_id] = ParticleFilter(
                    n_particles, fd.price, process_noise, measurement_noise
                )

            pf = filters[market_id]
            state = pf.update(fd.price)
            result["estimate"] = state.estimate
            result["variance"] = state.variance
            result["ess"] = state.ess
            result["resampled"] = state.resampled

        return result

    _particle_tracker.__name__ = "particle_tracker"
    return _particle_tracker
