"""
Pure NumPy Kalman Filter for Pairs Trading.

Implements a linear Kalman filter with time-varying observation matrix
suited for estimating the evolving hedge ratio βₜ and intercept αₜ
in the cointegrating regression:

    Yₜ = αₜ + βₜ Xₜ + vₜ          (observation)
    [αₜ, βₜ]ᵀ = [αₜ₋₁, βₜ₋₁]ᵀ + wₜ   (state transition, random walk)

Design decisions:
  - Pure NumPy (no scipy, no pykalman) → zero external deps beyond NumPy
  - Time-varying H matrix: Hₜ = [1, Xₜ] — adapts every observation
  - Diagonal process covariance Q controls adaptability:
      High Q → fast adaptation, noisy estimates (overfit risk)
      Low Q  → slow adaptation, smooth estimates (lag risk)
  - Scalar observation noise R controls measurement noise filtering
  - State dimension = 2 (α, β); observation dimension = 1 (Y)
  - Online filtering mode: update() called per observation
  - Offline smoothing mode: smooth() runs RTS smoother over full series

Reference:
  R. E. Kalman (1960). "A New Approach to Linear Filtering
  and Prediction Problems."
"""

from __future__ import annotations

import numpy as np
from dataclasses import dataclass, field
from typing import Optional, Tuple


@dataclass
class KalmanState:
    """Holds the Kalman filter state at a single timestep."""

    alpha: float          # Intercept estimate
    beta: float           # Hedge ratio estimate
    cov: np.ndarray       # 2×2 state covariance matrix
    log_likelihood: float = 0.0  # Contribution to log-likelihood


class KalmanFilter:
    """
    Pure-NumPy linear Kalman filter for the state-space model:

      State:       xₜ = F xₜ₋₁ + wₜ,    wₜ ~ N(0, Q)
      Observation: yₜ = Hₜ xₜ + vₜ,     vₜ ~ N(0, R)

    where:
      - xₜ = [αₜ, βₜ]ᵀ  (2×1 state vector)
      - F = I₂           (random walk transition)
      - Q = diag(q_α, q_β) or scalar × I₂
      - Hₜ = [1, Xₜ]    (1×2, time-varying)
      - R = scalar       (observation noise variance)

    Usage:
        kf = KalmanFilter(transition_covariance=1e-4, observation_covariance=1e-2)
        for x, y in zip(X_series, Y_series):
            state = kf.update(x, y)
            print(state.alpha, state.beta)
    """

    def __init__(
        self,
        transition_covariance: float = 1e-4,
        observation_covariance: float = 1e-2,
        initial_state_covariance: float = 1.0,
        initial_alpha: float = 0.0,
        initial_beta: float = 1.0,
    ) -> None:
        """
        Args:
            transition_covariance:
                Diagonal value(s) for process noise Q.
                Higher = faster adaptation, more noise.
                Can be float (both states) or (q_alpha, q_beta) tuple.
            observation_covariance:
                Scalar measurement noise R.
                Higher = smoother estimates (trust model more than data).
            initial_state_covariance:
                Initial uncertainty (diagonal of P₀).
            initial_alpha, initial_beta:
                Initial state estimates.
        """
        # State dimension
        self.n_states = 2

        # Transition matrix: identity (random walk)
        self.F = np.eye(self.n_states, dtype=np.float64)

        # Process noise covariance Q
        if isinstance(transition_covariance, (int, float)):
            self.Q = np.eye(self.n_states) * transition_covariance
        else:
            self.Q = np.diag(transition_covariance)

        # Observation noise (scalar)
        self.R = np.atleast_2d(observation_covariance).astype(np.float64)

        # Initial state
        self.x = np.array([[initial_alpha], [initial_beta]], dtype=np.float64)

        # Initial state covariance
        self.P = np.eye(self.n_states) * initial_state_covariance

        # Bookkeeping
        self.n_obs = 0
        self.history: list[KalmanState] = []

    # ── Properties ──────────────────────────────────────────

    @property
    def alpha(self) -> float:
        """Current intercept estimate."""
        return float(self.x[0, 0])

    @property
    def beta(self) -> float:
        """Current hedge ratio estimate."""
        return float(self.x[1, 0])

    # ── Core Filtering ──────────────────────────────────────

    def update(self, X_t: float, Y_t: float) -> KalmanState:
        """
        Single Kalman filter update step.

        Args:
            X_t: Independent variable observation (e.g., X asset price)
            Y_t: Dependent variable observation (e.g., Y asset price)

        Returns:
            KalmanState with current α, β, covariance, and log-likelihood.
        """
        self.n_obs += 1

        # ── Prediction ──
        x_pred = self.F @ self.x                     # (2×1)
        P_pred = self.F @ self.P @ self.F.T + self.Q  # (2×2)

        # ── Observation matrix (time-varying!) ──
        H = np.array([[1.0, X_t]], dtype=np.float64)  # (1×2)

        # ── Innovation ──
        y_pred = (H @ x_pred)[0, 0]                      # predicted Y
        innovation = Y_t - y_pred                     # scalar

        S = H @ P_pred @ H.T + self.R                 # innovation covariance (1×1)
        S_inv = 1.0 / S[0, 0] if S[0, 0] > 0 else 1e10

        # ── Kalman gain ──
        K = P_pred @ H.T * S_inv                      # (2×1)

        # ── Update ──
        self.x = x_pred + K * innovation              # (2×1)
        self.P = P_pred - K @ H @ P_pred              # (2×2)
        # Ensure symmetry
        self.P = (self.P + self.P.T) / 2.0

        # ── Log-likelihood contribution ──
        ll = -0.5 * (
            np.log(2 * np.pi * S[0, 0]) +
            innovation * innovation * S_inv
        )

        state = KalmanState(
            alpha=float(self.x[0, 0]),
            beta=float(self.x[1, 0]),
            cov=self.P.copy(),
            log_likelihood=float(ll),
        )
        self.history.append(state)
        return state

    def update_batch(self, X: np.ndarray, Y: np.ndarray) -> list[KalmanState]:
        """Filter a full series of observations. Online (forward pass only)."""
        results = []
        for i in range(len(X)):
            state = self.update(float(X[i]), float(Y[i]))
            results.append(state)
        return results

    def compute_spread(self, X_t: float, Y_t: float) -> float:
        """
        Compute the Kalman-estimated spread at a given observation.

        spreadₜ = Yₜ - (αₜ + βₜ Xₜ)

        Positive spread → Y is overpriced relative to X → short Y, long X.
        Negative spread → Y is underpriced relative to X → long Y, short X.
        """
        return Y_t - (self.alpha + self.beta * X_t)

    # ── Smoothing (RTS) ────────────────────────────────────

    def smooth(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Rauch-Tung-Striebel (RTS) smoother.

        Runs backward pass to produce smoothed state estimates
        that incorporate all observations (future + past).

        Returns:
            (smoothed_alpha, smoothed_beta) as 1-D arrays.
        """
        n = len(self.history)
        if n == 0:
            return np.array([]), np.array([])

        # Forward states and covariances
        x_fwd = np.array([[s.alpha, s.beta] for s in self.history]).T  # (2×n)
        P_fwd = np.array([s.cov for s in self.history])                # (n×2×2)

        # Initialize smoothed
        x_smooth = np.zeros_like(x_fwd)
        x_smooth[:, -1] = x_fwd[:, -1]

        # Backward pass
        for t in range(n - 2, -1, -1):
            P_next = P_fwd[t + 1]           # (2×2)
            P_curr = P_fwd[t]               # (2×2)

            # Smoothing gain
            P_pred = self.F @ P_curr @ self.F.T + self.Q
            try:
                C = P_curr @ self.F.T @ np.linalg.inv(P_pred)
            except np.linalg.LinAlgError:
                C = np.zeros((2, 2))

            x_smooth[:, t] = x_fwd[:, t] + C @ (x_smooth[:, t + 1] - self.F @ x_fwd[:, t])

        return x_smooth[0, :], x_smooth[1, :]

    # ── Utility ─────────────────────────────────────────────

    def likelihood(self) -> float:
        """Total log-likelihood of the filtered series."""
        return sum(s.log_likelihood for s in self.history)

    def reset(self) -> None:
        """Reset filter to initial state (for warm-start / retune)."""
        self.x = np.array([[0.0], [1.0]], dtype=np.float64)
        self.P = np.eye(self.n_states) * 1.0
        self.n_obs = 0
        self.history.clear()


class KalmanPairsTrader:
    """
    Production-grade Kalman-filter-based pairs trading engine.

    Encapsulates the Kalman filter, spread computation, z-score generation,
    and signal logic. Designed to be called bar-by-bar in a live trading loop
    or run over historical data for backtesting.

    Architecture:
        ┌─────────────┐
        │  Price Feed  │──Xₜ, Yₜ──→ KalmanFilter.update()
        └─────────────┘                    │
                              ┌────────────▼────────────┐
                              │  αₜ, βₜ, spreadₜ       │
                              │  zₜ = (spreadₜ - μ) / σ │
                              │  signal = f(zₜ, θ)      │
                              └─────────────────────────┘

    Signal logic:
        z > +z_entry  → Y overpriced → SHORT Y, LONG X
        z < -z_entry  → Y underpriced → LONG Y, SHORT X
        |z| < z_exit  → close position (mean reversion complete)

    Usage:
        trader = KalmanPairsTrader(
            transition_covariance=1e-4,
            z_entry=2.0,
            z_exit=0.5,
        )
        for x, y in zip(prices_X, prices_Y):
            signal = trader.step(x, y)
            if signal != 0:
                execute(signal)
    """

    def __init__(
        self,
        transition_covariance: float = 1e-4,
        observation_covariance: float = 1e-2,
        z_entry: float = 2.0,
        z_exit: float = 0.5,
        z_stop: float = 4.0,
        warmup_bars: int = 50,
        z_score_lookback: int = 100,
    ) -> None:
        """
        Args:
            transition_covariance: Q diagonal — controls β adaptation speed.
            observation_covariance: R scalar — measurement noise filter.
            z_entry: Z-score threshold for opening positions.
            z_exit: Z-score threshold for closing positions.
            z_stop: Stop-loss threshold (close immediately if |z| exceeds this).
            warmup_bars: Minimum observations before trading.
            z_score_lookback: Rolling window for z-score μ and σ estimation.
        """
        self.kf = KalmanFilter(
            transition_covariance=transition_covariance,
            observation_covariance=observation_covariance,
            initial_alpha=0.0,
            initial_beta=1.0,
        )
        self.z_entry = z_entry
        self.z_exit = z_exit
        self.z_stop = z_stop
        self.warmup_bars = warmup_bars
        self.z_score_lookback = z_score_lookback

        # Rolling spread history for z-score normalization
        self._spreads: list[float] = []

        # Current position state
        self.position: int = 0   # +1 = long Y/short X, -1 = short Y/long X
        self.entry_spread: float = 0.0

    # ── Properties ──────────────────────────────────────────

    @property
    def alpha(self) -> float:
        return self.kf.alpha

    @property
    def beta(self) -> float:
        return self.kf.beta

    @property
    def spread(self) -> float:
        return self._spreads[-1] if self._spreads else 0.0

    # ── Core Step ───────────────────────────────────────────

    def step(self, X_t: float, Y_t: float) -> dict:
        """
        Process one observation and return a signal.

        Args:
            X_t: Independent variable price (denominator asset)
            Y_t: Dependent variable price (numerator asset)

        Returns:
            Dict with keys: signal (int), spread (float), z_score (float),
            alpha (float), beta (float), position (int)
        """
        # Update Kalman filter
        self.kf.update(X_t, Y_t)

        # Compute spread
        spread = self.kf.compute_spread(X_t, Y_t)
        self._spreads.append(spread)

        # Trim spread history to lookback
        lookback = min(self.z_score_lookback, len(self._spreads))
        recent = self._spreads[-lookback:]

        # Z-score computation
        mu = np.mean(recent)
        sigma = np.std(recent, ddof=1)
        z = (spread - mu) / sigma if sigma > 1e-12 else 0.0

        # Signal generation
        signal = 0  # 0 = hold / no action

        if self.kf.n_obs < self.warmup_bars:
            signal = 0
        elif self.position == 0:
            # No position — look for entry
            if z > self.z_entry:
                signal = -1  # Y overpriced → SHORT Y, LONG X
            elif z < -self.z_entry:
                signal = +1  # Y underpriced → LONG Y, SHORT X
        else:
            # In position — check exit conditions
            if abs(z) < self.z_exit:
                signal = -self.position  # close
            elif abs(z) > self.z_stop:
                signal = -self.position  # stop-loss
            # Also mean-reversion exit: if spread crosses zero
            elif (self.position > 0 and spread > 0) or (self.position < 0 and spread < 0):
                signal = -self.position  # profit-taking on mean cross

        # Update position
        if signal != 0 and self.position == 0:
            self.position = signal
            self.entry_spread = spread
        elif signal != 0 and self.position != 0:
            self.position = 0
            self.entry_spread = 0.0

        return {
            "signal": signal,
            "spread": spread,
            "z_score": z,
            "alpha": self.alpha,
            "beta": self.beta,
            "position": self.position,
        }

    def reset(self) -> None:
        """Reset trader state (for backtest runs)."""
        self.kf.reset()
        self._spreads.clear()
        self.position = 0
        self.entry_spread = 0.0
