"""Forward backward pass for a GLM-HMM."""

from functools import partial
from typing import Callable, Optional, Tuple

import equinox as eqx
import jax
import jax.numpy as jnp

from ..glm.params import GLMParams
from ..typing import Aux, SolverState
from .m_step_analytical_updates import (
    _analytical_m_step_log_initial_prob,
    _analytical_m_step_log_transition_prob,
)
from .params import GLMHMMParams, GLMScale, HMMParams
from .utils import Array, compute_rate_per_state, initialize_new_session


class GLMHMMState(eqx.Module):
    """State class for the GLMHHM EM-algorithm."""

    data_log_likelihood: float | Array
    previous_data_log_likelihood: float | Array
    log_likelihood_history: Array
    iterations: int


EMCarry = Tuple[GLMHMMParams, GLMHMMState]


def compute_xi_log(
    log_alphas,
    log_betas,
    log_conditional_prob,
    log_normalization,
    is_new_session,
    log_transition_prob,
):
    """
    Compute the sum of the joint posterior (xi) over consecutive latent states in log-space.

    Compute the sum of the joint posterior (xis, eqn. 13.14 of [1]_) over samples, implementing
    the summation required in the eqn. 13.19 of [1]_.

    Parameters
    ----------
    log_alphas :
        Log forward messages, shape ``(n_time_bins, n_states)``
    log_betas :
        Log backward messages, shape ``(n_time_bins, n_states)``.
    log_conditional_prob :
        Log observation likelihoods log p(y_t | z_t), shape ``(n_time_bins, n_states)``.
    log_normalization :
        Log normalization constants from forward pass, shape ``(n_time_bins,)``.
    is_new_session :
        Boolean array, True at start of new sessions, shape ``(n_time_bins,)``.
    log_transition_prob :
        Log transition probability matrix, shape ``(n_states, n_states)``.

    Returns
    -------
    :
        Log expected joint posteriors between time steps, shape ``(n_states, n_states)``.

    References
    ----------
    .. [1] Bishop, C. M. (2006). Pattern recognition and machine learning. Springer.

    """
    # shift alpha so that alpha[t-1] aligns with beta[t]
    norm_log_alpha = log_alphas[:-1] - log_normalization[1:, jnp.newaxis]

    # mask out steps where t is a new session
    norm_log_alpha = jnp.where(
        is_new_session[1:, jnp.newaxis], -jnp.inf, norm_log_alpha
    )

    # Compute xi sum in one matmul
    log_xi_sum = jax.scipy.special.logsumexp(
        norm_log_alpha.T[..., jnp.newaxis]
        + (log_conditional_prob[1:] + log_betas[1:])[jnp.newaxis],
        axis=1,
    )

    return log_xi_sum + log_transition_prob


def forward_pass(
    log_initial_prob: Array,
    log_transition_prob: Array,
    log_conditional_prob: Array,
    is_new_session: Array,
) -> Tuple[Array, Array]:
    """
    Forward pass of an HMM in log-space.

    This function performs the recursive forward pass over time using ``jax.lax.scan``,
    computing the filtered log-probabilities (``log alpha``, eqn. 13.34 and 13.36 of [1]_) at each time step.
    At the start of a new session, the recursion is reset using the initial state distribution.
    All computations are performed in log-space for numerical stability.

    Parameters
    ----------
    log_initial_prob :
        Initial state log-probability distribution, array of shape ``(n_states,)``.

    log_transition_prob :
        Log-transition matrix of shape ``(n_states, n_states)``, where entry ``log T[i, j]`` is the
        log-probability of transitioning from state ``i`` to state ``j``.

    log_conditional_prob :
        Array of shape ``(n_time_bins, n_states)``, representing the observation log-likelihood
        ``log p(y_t | z_t)`` at each time step for each state.

    is_new_session :
        Boolean array of shape ``(n_time_bins,)`` indicating the start of new sessions. When
        ``is_new_session[t]`` is True, the recursion at time ``t`` is reset using ``log_initial_prob``.

    Returns
    -------
    log_alphas :
        Array of shape ``(n_time_bins, n_states)``, containing the filtered log-probabilities
        at each time step. ``log_alphas[t]`` corresponds to the log forward message at time ``t``.

    log_normalizers :
        Array of shape ``(n_time_bins,)`` containing the log-normalization constants at each
        time step. The sum of these values gives the log-likelihood of the sequence.

    Notes
    -----
    All operations are performed in log-space to avoid numerical underflow/overflow.
    Normalization is performed at each time step using ``logsumexp`` for numerical stability.

    Equivalent pseudocode in standard Python (log-space version):

    .. code-block:: python

        n_time_bins, n_states = log_py_z.shape
        log_alphas = np.full((n_time_bins, n_states), -np.inf)
        log_c = np.full(n_time_bins, -np.inf)

        for t in range(n_time_bins):
            if new_sess[t]:
                log_alphas[t] = log_initial_prob + log_py_z[t]
            else:
                log_alphas[t] = log_py_z[t] + logsumexp(
                    log_transition_prob + log_alphas[t - 1][None, :],
                    axis=1
                )

            log_c[t] = logsumexp(log_alphas[t])
            log_alphas[t] -= log_c[t]

    References
    ----------
    .. [1] Bishop, C. M. (2006). Pattern recognition and machine learning. Springer.

    """

    def initial_compute(log_posterior, _):
        return log_posterior + log_initial_prob

    def transition_compute(log_posterior, log_alpha_previous):
        log_exp_transition = jax.scipy.special.logsumexp(
            log_transition_prob + log_alpha_previous[None, :],  # Broadcasting
            axis=1,  # Sum over second axis (after transpose)
        )
        return log_posterior + log_exp_transition

    def body_fn(carry, xs):
        log_alpha_previous = carry
        log_posterior, is_new_session = xs

        log_alpha = jax.lax.cond(
            is_new_session,
            initial_compute,
            transition_compute,
            log_posterior,
            log_alpha_previous,
        )
        log_const = jax.scipy.special.logsumexp(log_alpha)
        log_alpha = log_alpha - log_const
        return log_alpha, (log_alpha, log_const)

    init = jnp.full_like(log_conditional_prob[0], -jnp.inf)  # log(0)
    log_transition_prob = log_transition_prob.T
    _, (log_alphas, log_normalizers) = jax.lax.scan(
        body_fn, init, (log_conditional_prob, is_new_session)
    )
    return log_alphas, log_normalizers


def backward_pass(
    log_transition_prob: Array,
    log_conditional_prob: Array,
    log_normalizers: Array,
    is_new_session: Array,
):
    """
    Run the backward pass of the HMM inference algorithm to compute log-beta messages.

    This function performs the backward recursion step of the forward–backward algorithm,
    using ``jax.lax.scan`` in reverse to compute log-beta messages at each time step, computing
    the ``log beta`` parameters, see eqn. 13.35 and 13.38 of [1]_.
    It handles session boundaries by resetting the beta messages when a new session starts.
    All computations are performed in log-space for numerical stability.

    Parameters
    ----------
    log_transition_prob :
        Log-transition matrix of shape ``(n_states, n_states)``, where entry ``log T[i, j]`` is the
        log-probability of transitioning from state ``i`` to state ``j``.

    log_conditional_prob :
        Array of shape ``(n_time_bins, n_states)``, representing the observation log-likelihoods
        ``log p(y_t | z_t)`` at each time step for each state.

    log_normalizers :
        Array of shape ``(n_time_bins,)`` containing the log-normalization constants from the forward
        pass. These are used to normalize the backward recursion in log-space.

    is_new_session :
        Boolean array of shape ``(n_time_bins,)`` indicating the start of new sessions. When
        ``is_new_session[t]`` is True, the backward message at time ``t`` is reset to a vector of zeros
        (corresponding to log(1) for each state).

    Returns
    -------
    log_betas :
        Array of shape ``(n_time_bins, n_states)``, containing the log-beta messages at each time step.
        The indexing is aligned with the forward pass, such that ``log_betas[t]`` corresponds to the
        log backward message at time ``t``.

    Notes
    -----
    This implementation follows the standard HMM backward equations (Bishop, 2006, Eq. 13.38–13.39),
    adapted to log-space, including reinitialization for segmented sequences.

    Equivalent pseudocode in standard Python (log-space version):

    .. code-block:: python

        n_time_bins, n_states = log_py_z.shape
        log_betas = np.full((n_time_bins, n_states), -np.inf)
        log_betas[-1] = np.zeros(n_states)  # log(1) = 0

        for t in range(n_time_bins - 2, -1, -1):
            if new_sess[t + 1]:
                log_betas[t] = np.zeros(n_states)
            else:
                log_betas[t] = logsumexp(
                    log_transition_prob + (log_betas[t + 1] + log_py_z[t + 1])[None, :],
                    axis=1
                ) - log_c[t + 1]

    References
    ----------
    .. [1] Bishop, C. M. (2006). Pattern recognition and machine learning. Springer.

    """
    init = jnp.zeros_like(log_conditional_prob[0])

    def initial_compute(log_posterior, *_):
        # Initialize with log(ones) = zeros
        return jnp.zeros_like(log_posterior)

    def backward_step(log_posterior, log_betas, log_normalization):
        # Normalize (log of Equation 13.62)
        return (
            jax.scipy.special.logsumexp(
                log_transition_prob + (log_posterior + log_betas)[None, :], axis=1
            )
            - log_normalization
        )

    def body_fn(carry, xs):
        log_posterior, log_norm, is_new_sess = xs
        log_beta = jax.lax.cond(
            is_new_sess,
            initial_compute,
            backward_step,
            log_posterior,
            carry,
            log_norm,
        )
        return log_beta, carry

    # Keeping the output betas because I am interested in
    # all outputs, including the last one.
    _, log_betas = jax.lax.scan(
        body_fn,
        init,
        (log_conditional_prob, log_normalizers, is_new_session),
        reverse=True,
    )
    return log_betas


@partial(jax.jit, static_argnames=["inverse_link_function", "log_likelihood_func"])
def forward_backward(
    X: Array,
    y: Array,
    log_initial_prob: Array,
    log_transition_prob: Array,
    glm_params: GLMParams,
    glm_scale: GLMScale,
    inverse_link_function: Callable[[Array], Array],
    log_likelihood_func: Callable[[Array, Array, Array], Array],
    is_new_session: Array | None = None,
):
    """
    Run the forward-backward Baum-Welch algorithm.

    Run the forward-backward Baum-Welch algorithm [1]_ that compute a posterior distribution over latent
    states. It handles session boundaries by resetting the ``alpha`` and ``beta`` messages when a new
    session starts.

    Parameters
    ----------
    X :
        Design matrix, pytree with leaves of shape ``(n_time_bins, n_features)``.

    y :
        Observations, pytree with leaves of shape ``(n_time_bins,)``.

    log_initial_prob :
        Log of the initial latent state probability, pytree with leaves of shape ``(n_states, 1)``.

    log_transition_prob :
        Latent state log-transition matrix, pytree with leaves of shape ``(n_states, n_states)``.
        ``transition_prob[i, j]`` is the probability of transitioning from state ``i`` to state ``j``.

    glm_params :
        GLM coefficients of shape ``(n_features, n_states)``
        and intercept of shape ``(n_states,)`` as GLMParams.

    glm_scale :
        The scale parameter of the likelihood. Shape (n_states,) or (n_neurons, n_states).

    inverse_link_function :
        Function mapping linear predictors to the mean of the observation distribution
        (e.g., exp for Poisson, sigmoid for Bernoulli).

    log_likelihood_func :
        Function computing the elementwise log-likelihood of observations given predicted mean values.
        Must return an array of shape ``(n_time_bins, n_states)``.

    is_new_session :
        Boolean array marking the start of a new session.
        If unspecified or empty, treats the full set of trials as a single session.

    Returns
    -------
    log_posteriors :
        Marginal log-posterior distribution over latent states, shape ``(n_time_bins, n_states)``.

    log_joint_posterior :
        Joint log-posterior distribution between consecutive time steps summed
        over samples, shape ``(n_states, n_states)``.

    log_likelihood :
        Total log-likelihood of the observation sequence under the model.

    log_likelihood_norm :
        The normalized total likelihood.

    log_alphas :
        Log forward messages (log alpha values), shape ``(n_time_bins, n_states)``.

    log_betas :
        Log backward messages (log beta values), shape ``(n_time_bins, n_states)``.

    References
    ----------
    .. [1] Bishop, C. M. (2006). *Pattern recognition and machine learning*. Springer.
    """
    # Initialize variables
    n_time_bins = y.shape[0]
    is_new_session = initialize_new_session(y.shape[0], is_new_session)
    predicted_rate_given_state = compute_rate_per_state(
        X, glm_params, inverse_link_function
    )

    # Compute likelihood given the fixed weights
    # Data likelihood p(y|z) from emissions model
    # NOTE:
    # For N neurons and S samples, we want the total likelihood
    # across neurons for each sample and latent state.

    # Example helper:
    # def combined_likelihood(log_likelihood_func, y, rate):
    #     # log_likelihood_func takes (y, rate) and returns log-likelihood per neuron:
    #     #   y:    shape (S, N)
    #     #   rate: shape (S, N, K)  # K = number of latent states
    #
    #     assert y.ndim == 2      # (samples, neurons)
    #     assert rate.ndim == 3   # (samples, neurons, states)
    #
    #     # vmap over the state axis: apply log_likelihood_func for each state
    #     # Result: shape (S, N, K)
    #     log_like = jax.vmap(log_likelihood_func, in_axes=(None, 2), out_axes=2)(y, rate)
    #
    #     # Combine neurons assuming conditional independence:
    #     # sum log-likelihoods over neurons (axis=1), then exponentiate for stability
    #     # Final shape: (S, K)
    #     return jnp.exp(jnp.sum(log_like, axis=1))
    #
    # Here, log_likelihood_func is the ``log_likelihood`` method from
    # nemos.observation_models.Observations with ``aggregate_sample_scores = lambda x:x``

    log_conditionals = log_likelihood_func(
        y, predicted_rate_given_state, jnp.exp(glm_scale.log_scale)
    )

    # Compute forward pass
    log_alphas, log_normalization = forward_pass(
        log_initial_prob, log_transition_prob, log_conditionals, is_new_session
    )  # these are equivalent to the forward pass with python loop

    # Compute backward pass
    log_betas = backward_pass(
        log_transition_prob, log_conditionals, log_normalization, is_new_session
    )

    log_likelihood = jnp.sum(
        log_normalization
    )  # Store log-likelihood, log of Equation 13.63

    likelihood_norm = jnp.exp(log_likelihood / n_time_bins)  # Normalize

    # Posteriors
    # ----------
    # Compute posterior distributions
    # Gamma - Equations 13.32, 13.64 from [1]
    log_posteriors = log_alphas + log_betas

    # xis Equations 13.43 and 13.65 from [1]
    # Posterior over consecutive states summed across time steps
    log_joint_posterior = compute_xi_log(
        log_alphas,
        log_betas,
        log_conditionals,
        log_normalization,
        is_new_session,
        log_transition_prob,
    )
    return (
        log_posteriors,
        log_joint_posterior,
        log_likelihood,
        likelihood_norm,
        log_alphas,
        log_betas,
    )


@partial(
    jax.jit,
    static_argnames=[
        "m_step_fn_glm_params",
        "m_step_fn_glm_scale",
        "inverse_link_function",
    ],
)
def run_m_step(
    params: GLMHMMParams,
    X: Array,
    y: Array,
    log_posteriors: Array,
    log_joint_posterior: Array,
    is_new_session: Array,
    m_step_fn_glm_params: Callable[
        [GLMParams, Array, Array, Array], Tuple[GLMParams, SolverState, Aux]
    ],
    m_step_fn_glm_scale: (
        Callable[[GLMScale, Array, Array, Array], Tuple[GLMScale, SolverState]] | None
    ),
    inverse_link_function: Callable[[Array], Array],
    dirichlet_prior_alphas_init_prob: Array | None = None,
    dirichlet_prior_alphas_transition: Array | None = None,
) -> Tuple[GLMHMMParams, SolverState]:
    r"""
    Perform the M-step of the EM algorithm for GLM-HMM.

    Parameters
    ----------
    params:
        The current model parameters.
    X:
        Design matrix of observations, shape (n_samples, n_features).
    y:
        Target responses, shape ``(n_samples,)`` or ``(n_samples, n_neurons)``.
    log_posteriors:
        Log-posterior probabilities over states, shape ``(n_samples, n_states)``.
    log_joint_posterior:
        Log joint posterior probabilities over pairs of states summed over samples. Shape ``(n_states, n_states)``.
        :math:`\sum_t P(z_{t-1}, z_t \mid X, y, \theta_{\text{old}})`.
    is_new_session:
        Boolean mask marking the first observation of each session. Shape ``(n_samples,)``.
    m_step_fn_glm_params:
        Callable that performs the M-step update for GLM parameters (coefficients and intercepts).
        Should have signature: ``f(glm_params, X, y, posteriors) -> (updated_params, state, aux)``.
        The regularizer/prior for the GLM parameters should be configured within this callable.
    m_step_fn_glm_scale:
        Callable that performs the M-step update for GLM scales.
        Should have signature: ``f(scale, glm_params, X, y, posteriors) -> (updated_scale, state)``.
    inverse_link_function:
        Inverse link function for computing the predicted rates per each state.
    dirichlet_prior_alphas_init_prob:
        Prior for the initial states, shape ``(n_states,)``.
    dirichlet_prior_alphas_transition:
        Prior for the transition probabilities, shape ``(n_states, n_states)``.

    Returns
    -------
    params:
        The updated model parameters.
    state:
        State returned by the solver.

    Notes
    -----
    The current implementation requires all Dirichlet prior parameters alpha >= 1.
    Support for sparse priors (0 < alpha < 1) may be added in a future version
    using alternative optimization methods such as proximal gradient descent.
    """
    posteriors = jnp.exp(log_posteriors)

    # Update Initial state probability Eq. 13.18
    log_initial_prob = _analytical_m_step_log_initial_prob(
        log_posteriors,
        is_new_session=is_new_session,
        dirichlet_prior_alphas=dirichlet_prior_alphas_init_prob,
    )
    log_transition_prob = _analytical_m_step_log_transition_prob(
        log_joint_posterior, dirichlet_prior_alphas=dirichlet_prior_alphas_transition
    )

    with jax.disable_jit(False):
        # Minimize negative log-likelihood to update GLM weights
        optimized_projection_weights, state, _ = m_step_fn_glm_params(
            params.glm_params, X, y, posteriors
        )
        predicted_rate = compute_rate_per_state(
            X, optimized_projection_weights, inverse_link_function=inverse_link_function
        )

    if m_step_fn_glm_scale is not None:
        # Gaussian, Gamma, and other have a scale.
        glm_scale, state_scale, _ = m_step_fn_glm_scale(
            params.glm_scale, y, predicted_rate, posteriors
        )
    else:
        # Poisson, Bernoulli etc. do not have a scale
        # just keep carrying the scale
        glm_scale = params.glm_scale

    params = GLMHMMParams(
        hmm_params=HMMParams(log_initial_prob, log_transition_prob),
        glm_params=optimized_projection_weights,
        glm_scale=glm_scale,
    )
    return (
        params,
        state,
    )


def _em_step(
    carry: EMCarry,
    X: Array,
    y: Array,
    inverse_link_function: Callable[[Array], Array],
    log_likelihood_func: Callable[[Array, Array, Array], Array],
    m_step_fn_glm_params: Callable[
        [GLMParams, Array, Array, Array], Tuple[GLMParams, SolverState]
    ],
    m_step_fn_glm_scale: (
        Callable[[Array, Array, Array, Array], Tuple[Array, SolverState]] | None
    ),
    is_new_session: Array,
) -> EMCarry:
    """
    Execute a single EM iteration combining E-step and M-step.

    Performs one complete EM cycle: computes posterior distributions over
    latent states (E-step), then updates all model parameters to maximize
    the expected complete-data log-likelihood (M-step).

    Parameters
    ----------
    carry :
        Tuple of current parameters and state:
        ``((log_init_prob, log_trans_matrix, glm_params, log_scale), previous_state)``
    X :
        Design matrix of observations.
    y :
        Target responses.
    inverse_link_function :
        Function mapping linear predictors to predicted rates.
    log_likelihood_func :
        Log-likelihood function for the E-step.
    m_step_fn_glm_params :
        M-step update function for GLM coefficients and intercepts.
    m_step_fn_glm_scale :
        M-step update function for scale parameters.
    is_new_session :
        Boolean array marking session boundaries.

    Returns
    -------
    carry :
        Updated tuple of parameters and state with new log-likelihood values:
        ``((log_init_prob, log_trans_matrix, glm_params, glm_scale), new_state)``

    Notes
    -----
    This function is designed to be called repeatedly by ``eqx.internal.while_loop``
    until convergence criteria are met. The carry structure allows JAX to efficiently
    compile and execute the EM loop.
    """

    params, previous_state = carry

    log_posteriors, log_joint_posterior, _, new_log_like, _, _ = forward_backward(
        X,
        y,
        params.hmm_params.log_initial_prob,
        params.hmm_params.log_transition_prob,
        params.glm_params,
        params.glm_scale,
        inverse_link_function,
        log_likelihood_func,
        is_new_session,
    )

    new_params, _ = run_m_step(
        params,
        X,
        y,
        log_posteriors=log_posteriors,
        log_joint_posterior=log_joint_posterior,
        is_new_session=is_new_session,
        m_step_fn_glm_params=m_step_fn_glm_params,
        m_step_fn_glm_scale=m_step_fn_glm_scale,
        inverse_link_function=inverse_link_function,
    )

    new_state = GLMHMMState(
        iterations=previous_state.iterations + 1,
        data_log_likelihood=new_log_like,
        previous_data_log_likelihood=previous_state.data_log_likelihood,
        log_likelihood_history=previous_state.log_likelihood_history.at[
            previous_state.iterations
        ].set(new_log_like),
    )

    return new_params, new_state


def em_step(
    params: GLMHMMParams,
    state: GLMHMMState,
    X: Array,
    y: Array,
    inverse_link_function: Callable,
    log_likelihood_func: Callable,
    m_step_fn_glm_params: Callable,
    m_step_fn_glm_scale: Callable,
    is_new_session: Array,
) -> Tuple[GLMHMMParams, GLMHMMState]:
    """
    Perform a single EM iteration step for GLM-HMM.

    This function provides a clean public API for running a single EM iteration,
    compatible with the optimization API pattern used by solvers. It wraps the
    internal `_em_step` function which operates on EMCarry tuples.

    Parameters
    ----------
    params : GLMHMMParams
        Current GLM-HMM parameters containing GLM coefficients/intercepts and
        HMM initial/transition probabilities.
    state : GLMHMMState
        Current EM algorithm state containing iteration count and log-likelihood history.
    X : Array
        Design matrix of observations.
    y : Array
        Target responses.
    inverse_link_function : Callable
        Elementwise function mapping linear predictors to rates.
    log_likelihood_func : Callable
        Function computing the log-likelihood.
    m_step_fn_glm_params : Callable
        Callable that performs the M-step update for GLM parameters.
    m_step_fn_glm_scale : Callable
        Callable that performs the M-step update for GLM scale.
    is_new_session : Array
        Boolean mask for the first observation of each session.

    Returns
    -------
    updated_params : GLMHMMParams
        Updated parameters after one EM iteration.
    updated_state : GLMHMMState
        Updated state after one EM iteration.
    """
    # Pack params and state into EMCarry format (log-space for HMM params)
    carry = params, state

    # Run the internal EM step
    params, state = _em_step(
        carry,
        X=X,
        y=y,
        inverse_link_function=inverse_link_function,
        log_likelihood_func=log_likelihood_func,
        m_step_fn_glm_params=m_step_fn_glm_params,
        m_step_fn_glm_scale=m_step_fn_glm_scale,
        is_new_session=is_new_session,
    )

    return params, state


def check_log_likelihood_increment(state: GLMHMMState, tol: float) -> Array:
    """
    Check EM convergence using absolute tolerance on log-likelihood.

    Parameters
    ----------
    state :
        Current EM state containing likelihood history.
    tol :
        Absolute tolerance threshold.

    Returns
    -------
    : Array
        Boolean indicating convergence.
    """
    delta = jnp.abs(state.data_log_likelihood - state.previous_data_log_likelihood)
    return delta < tol


@partial(
    jax.jit,
    static_argnames=[
        "inverse_link_function",
        "log_likelihood_func",
        "m_step_fn_glm_params",
        "m_step_fn_glm_scale",
        "maxiter",
        "check_convergence",
        "tol",
    ],
)
def em_glm_hmm(
    params: GLMHMMParams,
    X: Array,
    y: Array,
    inverse_link_function: Callable,
    log_likelihood_func: Callable,
    m_step_fn_glm_params: Callable,
    m_step_fn_glm_scale: Callable | None,
    is_new_session: Optional[Array] = None,
    maxiter: int = 10**3,
    tol: float = 1e-8,
    check_convergence: Callable = check_log_likelihood_increment,
) -> Tuple[GLMHMMParams, GLMHMMState]:
    """
    Perform EM optimization for a GLM-HMM.

    Uses equinox while_loop for efficient early stopping when convergence
    criteria are met.

    Parameters
    ----------
    params:
        Initial GLM-HMM parameters. This includes:
        - the GLM coef, shape ``(n_features, n_states)`` or ``(n_features, n_neurons, n_states)`` .
        - the GLM intercept, shape  ``(n_states, )`` or ``(n_neurons, n_states)`` .
        - the HMM initial probabilities, shape ``(n_states,)``.
        - the HMM transition probabilities, shape ``(n_states, n_states)``.
    X:
        Design matrix of observations.
    y:
        Target responses.
    inverse_link_function:
        Elementwise function mapping linear predictors to rates.
    log_likelihood_func:
        Function computing the log-likelihood.
    m_step_fn_glm_params:
        Callable that performs the M-step update for GLM parameters (coefficients and intercepts).
        Should have signature: ``f(glm_params, X, y, posteriors) -> (updated_params, state)``.
        Typically created by configuring a solver with the appropriate regularizer/prior.
    m_step_fn_glm_scale:
        If provided, callable that performs the M-step updated for the scale parameter of
        the distribution.
        Must have signature ``f(scale, glm_params, X, y, posteriors) -> (updated_scale, state)``
    is_new_session:
        Boolean mask for the first observation of each session.
    maxiter:
        Maximum number of EM iterations.
    tol:
        The tolerance for the convergence criterion.
    check_convergence:
        Callable receiving the state and computing the convergence.

    Returns
    -------
    params:
        The fitted GLM-HMM parameters.
    state:
        Final GLMHMMState containing all parameters and diagnostics.
    """
    is_new_session = initialize_new_session(y.shape[0], is_new_session)

    state = GLMHMMState(
        data_log_likelihood=-jnp.array(jnp.inf),
        previous_data_log_likelihood=-jnp.array(jnp.inf),
        log_likelihood_history=jnp.full(maxiter, jnp.nan),
        iterations=0,
    )

    em_step_fn_while = eqx.Partial(
        lambda *args, **kwargs: _em_step(*args, **kwargs),
        X=X,
        y=y,
        inverse_link_function=inverse_link_function,
        log_likelihood_func=log_likelihood_func,
        m_step_fn_glm_params=m_step_fn_glm_params,
        m_step_fn_glm_scale=m_step_fn_glm_scale,
        is_new_session=is_new_session,
    )

    def stopping_condition_while(carry):
        _, new_state = carry
        return ~check_convergence(new_state, tol)

    init_carry = params, state
    params, state = eqx.internal.while_loop(
        stopping_condition_while,
        em_step_fn_while,
        init_carry,
        max_steps=maxiter,
        kind="lax",
    )

    return params, state


@partial(
    jax.jit,
    static_argnames=["inverse_link_function", "log_likelihood_func", "return_index"],
)
def max_sum(
    X: Array,
    y: Array,
    initial_prob: Array,
    transition_prob: Array,
    glm_params: GLMParams,
    inverse_link_function: Callable,
    log_likelihood_func: Callable[[Array, Array], Array],
    is_new_session: Array | None = None,
    return_index: bool = False,
):
    """
    Find maximum a posteriori (MAP) state path via the max-sum algorithm.

    This function implements the max-sum algorithm for a GLM-HMM, also known as Viterbi algorithm.

    Parameters
    ----------
    X :
        Design matrix, pytree with leaves of shape ``(n_time_bins, n_features)``.

    y :
        Observations, pytree with leaves of shape ``(n_time_bins,)``.

    initial_prob :
        Initial latent state probability, pytree with leaves of shape ``(n_states, 1)``.

    transition_prob :
        Latent state transition matrix, pytree with leaves of shape ``(n_states, n_states)``.
        ``transition_prob[i, j]`` is the probability of transitioning from state ``i`` to state ``j``.

    glm_params :
        Length two tuple with the GLM coefficients of shape ``(n_features, n_states)``
        and intercept of shape ``(n_states,)``.

    inverse_link_function :
        Function mapping linear predictors to the mean of the observation distribution
        (e.g., exp for Poisson, sigmoid for Bernoulli).

    is_new_session :
        Boolean array marking the start of a new session.
        If unspecified or empty, treats the full set of trials as a single session.

    return_index:
        If False, return 1-hot encoded map states, if True, return map state indices.

    Returns
    -------
    map_path:
        The MAP state path.

    """
    is_new_session = initialize_new_session(y.shape[0], is_new_session)
    predicted_rate_given_state = compute_rate_per_state(
        X, glm_params, inverse_link_function
    )
    log_emission = log_likelihood_func(y, predicted_rate_given_state)

    log_transition = jnp.log(transition_prob)
    log_init = jnp.log(initial_prob)
    n_states = initial_prob.shape[0]

    def forward_max_sum(omega_prev, xs):
        log_em, is_new_sess = xs

        def reset_chain(omega_prev, log_em):
            # New session: reset to initial distribution
            omega = log_init + log_em
            max_prob_state = jnp.full(n_states, -1)  # Boundary marker
            return omega, max_prob_state

        def continue_chain(omega_prev, log_em):
            # Continue existing session: Viterbi step
            step = log_em[None, :] + log_transition + omega_prev[:, None]
            max_prob_state = jnp.argmax(step, axis=0)
            omega = step[max_prob_state, jnp.arange(n_states)]
            return omega, max_prob_state

        omega, max_prob_state = jax.lax.cond(
            is_new_sess,
            reset_chain,
            continue_chain,
            omega_prev,
            log_em,
        )

        return omega, (omega, max_prob_state)

    init_omega = log_init + log_emission[0]
    _, (omegas, max_prob_states) = jax.lax.scan(
        forward_max_sum, init_omega, (log_emission[1:], is_new_session[1:])
    )

    # Backward pass
    best_final_state = jnp.argmax(omegas[-1])
    # Prepend initial omega and exclude last one, which is already considered.
    omegas = jnp.concatenate([init_omega[None, :], omegas[:-1]], axis=0)

    def backward_max_sum(current_state_idx, xs):
        max_prob_st, omega_t = xs

        def session_boundary(state_idx, max_prob, omega):
            # Hit a session start, pick best state at this boundary
            return jnp.argmax(omega)

        def continue_backward(state_idx, max_prob, omega):
            # Normal backtracking
            return max_prob[state_idx]

        is_boundary = max_prob_st[current_state_idx] == -1

        prev_state_idx = jax.lax.cond(
            is_boundary,
            session_boundary,
            continue_backward,
            current_state_idx,
            max_prob_st,
            omega_t,
        )

        return prev_state_idx, prev_state_idx

    _, map_path = jax.lax.scan(
        backward_max_sum, best_final_state, (max_prob_states, omegas), reverse=True
    )

    # Append the final state
    map_path = jnp.concatenate([map_path, jnp.array([best_final_state])])

    if not return_index:
        map_path = jax.nn.one_hot(map_path, n_states, dtype=jnp.int32)

    return map_path
