from __future__ import annotations

import jax
import jax.numpy as jnp


def IntegratedRTSSmoother(kernel, t_states, obsid, instid, stateid, kalman_results):
    """
    Wrapper for jitted integrated_rts_smoother function

    Parameters:
        kernel  : IntegratedStateSpaceModel kernel
        t_states: Array of size K, sorted time coordinate of all states (exposure starts and ends)
        obsid   : Array of size N, which observation (0,...,N-1) is being made at each state k
        instid  : Array of size N, which instrument (0,...,Ninst-1) recorded observation n
        stateid : Array of size K, 0 for exposure-start, 1 for exposure-end
        kalman_results: output from Kalman filter (m_filtered, P_filtered, m_predicted, P_predicted)

    Returns:
        m_filtered: filtered means
        P_filtered: filtered covariances
        m_predicted: predicted means
        P_predicted: predicted covariances
    """

    # Model components
    A_aug = kernel.transition_matrix
    RESET = kernel.reset_matrix

    return integrated_rts_smoother(A_aug, RESET, t_states, obsid, instid, stateid, *kalman_results)


@jax.jit
def integrated_rts_smoother(
    A_aug,
    RESET,
    t_states,
    obsid,
    instid,
    stateid,
    m_filtered,
    P_filtered,
    m_predicted,
    P_predicted,
):
    """
    Jax implementation of the integrated RTS smoothing algorithm

    See Section 3.2.2 in Rubenzahl & Hattori et al. (in prep)
    for detailed description of the algorithm and notation.
    """

    def step(carry, k):
        # Outputs from Kalman filter, unpacked for notational consistency
        m_k = m_filtered[k]
        P_k = P_filtered[k]
        m_pred_next = m_predicted[k + 1]  # has superscript minus
        P_pred_next = P_predicted[k + 1]  # has superscript minus

        # Unpack state and covariance from last iteration
        m_hat_next, P_hat_next = carry

        # Compute smoothing gain
        Delta = t_states[k + 1] - t_states[k]
        A_k = A_aug(0, Delta)

        def smooth_start():
            """RTS smooth an exposure-start state"""

            m_k_pre = m_predicted[k]  # pre-reset start state
            P_k_pre = P_predicted[k]  # pre-reset start covariance

            Reset = RESET(instid[obsid[k]])
            AR = A_k @ Reset
            G_k = jnp.linalg.solve(P_pred_next.T, (P_k_pre @ AR.T).T).T
            m_hat_k = m_k_pre + G_k @ (m_hat_next - m_pred_next)
            P_hat_k = P_k_pre + G_k @ (P_hat_next - P_pred_next) @ G_k.T
            return m_hat_k, P_hat_k

        def smooth_end():
            """RTS smooth an exposure-end state"""
            G_k = jnp.linalg.solve(P_pred_next.T, (P_k @ A_k.T).T).T
            m_hat_k = m_k + G_k @ (m_hat_next - m_pred_next)
            P_hat_k = P_k + G_k @ (P_hat_next - P_pred_next) @ G_k.T
            return m_hat_k, P_hat_k

        m_hat_k, P_hat_k = jax.lax.cond(
            stateid[k] == 0,
            lambda _: smooth_start(),
            lambda _: smooth_end(),
            operand=None,
        )

        return (m_hat_k, P_hat_k), (m_hat_k, P_hat_k)

    # Start smoothing from final filtered state
    init_carry = (m_filtered[-1], P_filtered[-1])

    # Run backward from N-2 down to 0
    K = len(t_states)  # number of iterations
    _, outputs = jax.lax.scan(step, init_carry, jnp.arange(K - 2, -1, -1))
    m_smooth_reversed, P_smooth_reversed = outputs

    # Reverse outputs (with final filtered=smoothed state) to match time order
    m_smooth = jnp.vstack([m_smooth_reversed[::-1], m_filtered[-1][None, :]])
    P_smooth = jnp.vstack([P_smooth_reversed[::-1], P_filtered[-1][None, :, :]])
    return m_smooth, P_smooth
