"""Kalman Filter pipeline functions.

Pipeline functions:
    kalman_tracker - Track latent state with Kalman filter
    kalman_hedge_ratio - Dynamic hedge ratio estimation
    kalman_spread - Dynamic spread tracking
"""

from __future__ import annotations

from typing import Any, Callable

from horizon._horizon import KalmanFilter
from horizon.context import Context


def kalman_tracker(
    feed: str,
    state_dim: int = 1,
    obs_dim: int = 1,
) -> Callable[[Context], dict[str, Any]]:
    """Create a Kalman filter tracking pipeline function.

    Tracks a latent state variable from noisy feed observations.

    Args:
        feed: Feed name to read price data from.
        state_dim: State vector dimension.
        obs_dim: Observation vector dimension.

    Returns:
        Pipeline function: (Context) -> dict with keys:
            state, predicted_state, innovation, log_likelihood
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    filters: dict[str, KalmanFilter] = {}

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

        if market_id not in filters:
            filters[market_id] = KalmanFilter(state_dim, obs_dim)

        kf = filters[market_id]

        result = {
            "state": kf.state(),
            "predicted_state": [],
            "innovation": 0.0,
            "log_likelihood": 0.0,
        }

        if fd is not None and fd.price > 0:
            obs = [fd.price] + [0.0] * (obs_dim - 1)
            obs = obs[:obs_dim]

            predicted = kf.predict()
            result["predicted_state"] = predicted

            state = kf.update(obs)
            result["state"] = state

            if predicted:
                result["innovation"] = fd.price - predicted[0]

            result["log_likelihood"] = kf.log_likelihood(obs)

        return result

    _kalman_tracker.__name__ = "kalman_tracker"
    return _kalman_tracker


def kalman_hedge_ratio(
    feed_a: str,
    feed_b: str,
) -> Callable[[Context], dict[str, Any]]:
    """Create a dynamic hedge ratio estimator using Kalman filter.

    Models: price_b = beta * price_a + alpha + noise
    Uses a 2-state Kalman filter to track (alpha, beta) dynamically.

    Args:
        feed_a: Feed name for the first asset.
        feed_b: Feed name for the second asset.

    Returns:
        Pipeline function: (Context) -> dict with keys:
            hedge_ratio, intercept, spread, spread_zscore
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    filters: dict[str, KalmanFilter] = {}
    spread_history: dict[str, list[float]] = {}

    def _kalman_hedge_ratio(ctx: Context) -> dict[str, Any]:
        market_id = ctx.market.id if ctx.market else "__default__"
        fd_a = ctx.feeds.get(feed_a)
        fd_b = ctx.feeds.get(feed_b)

        if market_id not in filters:
            kf = KalmanFilter(2, 1)
            kf.set_process_noise([[1e-5, 0.0], [0.0, 1e-5]])
            kf.set_measurement_noise([[1e-3]])
            filters[market_id] = kf
            spread_history[market_id] = []

        kf = filters[market_id]
        spreads = spread_history[market_id]

        result = {
            "hedge_ratio": 0.0,
            "intercept": 0.0,
            "spread": 0.0,
            "spread_zscore": 0.0,
        }

        if fd_a is not None and fd_b is not None and fd_a.price > 0 and fd_b.price > 0:
            kf.set_observation([[1.0, fd_a.price]])
            kf.predict()
            state = kf.update([fd_b.price])

            if len(state) >= 2:
                alpha = state[0]
                beta = state[1]
                spread = fd_b.price - beta * fd_a.price - alpha

                result["hedge_ratio"] = beta
                result["intercept"] = alpha
                result["spread"] = spread

                spreads.append(spread)
                if len(spreads) > 100:
                    spreads[:] = spreads[-100:]

                if len(spreads) >= 10:
                    import math
                    mean_s = sum(spreads) / len(spreads)
                    var_s = sum((s - mean_s) ** 2 for s in spreads) / len(spreads)
                    std_s = math.sqrt(max(var_s, 1e-15))
                    result["spread_zscore"] = (spread - mean_s) / std_s if std_s > 1e-10 else 0.0

        return result

    _kalman_hedge_ratio.__name__ = "kalman_hedge_ratio"
    return _kalman_hedge_ratio


def kalman_spread(
    feed_a: str,
    feed_b: str,
) -> Callable[[Context], dict[str, Any]]:
    """Create a dynamic spread tracker using Kalman filter.

    Tracks the spread between two assets with time-varying mean.

    Args:
        feed_a: Feed name for the first asset.
        feed_b: Feed name for the second asset.

    Returns:
        Pipeline function: (Context) -> dict with keys:
            spread, filtered_spread, spread_vol
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    filters: dict[str, KalmanFilter] = {}

    def _kalman_spread(ctx: Context) -> dict[str, Any]:
        market_id = ctx.market.id if ctx.market else "__default__"
        fd_a = ctx.feeds.get(feed_a)
        fd_b = ctx.feeds.get(feed_b)

        if market_id not in filters:
            kf = KalmanFilter(1, 1)
            kf.set_process_noise([[1e-4]])
            kf.set_measurement_noise([[1e-3]])
            filters[market_id] = kf

        kf = filters[market_id]

        result = {
            "spread": 0.0,
            "filtered_spread": 0.0,
            "spread_vol": 0.0,
        }

        if fd_a is not None and fd_b is not None and fd_a.price > 0 and fd_b.price > 0:
            spread = fd_b.price - fd_a.price
            result["spread"] = spread

            kf.predict()
            state = kf.update([spread])
            result["filtered_spread"] = state[0] if state else spread

            cov = kf.covariance()
            if cov and cov[0]:
                import math
                result["spread_vol"] = math.sqrt(max(cov[0][0], 0.0))

        return result

    _kalman_spread.__name__ = "kalman_spread"
    return _kalman_spread
