"""Simulation and components, see also [the corresponding chapter in my thesis](https://stefanheyder.github.io/dissertation/thesis.pdf#nameddest=section.3.2)."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/00_glssm.ipynb.

# %% auto 0
__all__ = ['simulate_states', 'simulate_glssm', 'log_probs_x', 'log_probs_y', 'log_prob']

# %% ../../nbs/00_glssm.ipynb 1
import jax.numpy as jnp
import jax.random as jrn
from jax import vmap
from jax.lax import scan
from jaxtyping import Array, Float, PRNGKeyArray

from .kalman import kalman
from .typing import GLSSM, GLSSMObservationModel, GLSSMState, Observations, States
from .util import MVN_degenerate as MVN
from .util import mm_sim

# %% ../../nbs/00_glssm.ipynb 6
def simulate_states(
    state: GLSSMState,
    N: int,  # number of samples to draw
    key: PRNGKeyArray,  # the random state
) -> Float[Array, "N n+1 m"]:  # array of N samples from the state distribution
    """Simulate states of a GLSSM"""
    u, A, D, Sigma0, Sigma = state

    _, _, l = D.shape

    def sim_next_states(carry, inputs):
        x_prev, key = carry
        u, A, D, Sigma = inputs

        key, subkey = jrn.split(key)
        eps = MVN(jnp.zeros(l), Sigma).sample(N, seed=subkey)

        samples = u + mm_sim(A, x_prev) + mm_sim(D, eps)

        return (samples, key), samples

    key, subkey = jrn.split(key)
    x0 = MVN(u[0], Sigma0).sample(N, subkey)
    key, subkey = jrn.split(key)
    init = (x0, subkey)

    _, X = scan(sim_next_states, init, (u[1:], A, D, Sigma))

    X = jnp.concatenate([x0[None], X], axis=0)

    return X.transpose((1, 0, 2))

# %% ../../nbs/00_glssm.ipynb 7
from .typing import to_observation_model, to_states


def simulate_glssm(
    glssm: GLSSM,
    N: int,  # number of sample paths
    key: PRNGKeyArray,  # the random state
) -> (
    Float[Array, "N n+1 m"],
    Float[Array, "N n+1 p"],
):  # tuple of two arrays each with of N samples from the state/observation distribution
    """Simulate states and observations of a GLSSM"""

    v, B, Omega = to_observation_model(glssm)
    key, subkey = jrn.split(key)
    X = simulate_states(to_states(glssm), N, subkey).transpose((1, 0, 2))

    S = v[:, None, :] + vmap(mm_sim, (0, 0))(B, X)

    # samples x time x space
    X = X.transpose((1, 0, 2))

    S = S.transpose((1, 0, 2))

    key, subkey = jrn.split(key)
    Y = MVN(S, Omega).sample(seed=subkey)

    return X, Y

# %% ../../nbs/00_glssm.ipynb 15
from .util import append_to_front, mm_time


def log_probs_x(
    x: States, state: GLSSMState  # the states  # the state model
) -> Float[Array, "n+1"]:  # log probabilities $\log p(x_t \vert x_{t-1})$
    """log probabilities $\\log p(x_t | x_{t-1})$"""
    u, A, D, Sigma0, Sigma = state
    _, _, l = D.shape
    log_p0 = MVN(u[0], Sigma0).log_prob(x[0])
    x_prev = x[:-1]

    DT = D.transpose((0, 2, 1))

    eps = mm_time(DT, x[1:] - u[1:] - (A @ x_prev[:, :, None])[:, :, 0])

    log_p = MVN(jnp.zeros(l), Sigma).log_prob(eps)
    return append_to_front(log_p0, log_p)


def log_probs_y(
    y: Observations,  # the observations
    x: States,  # the states
    obs_model: GLSSMObservationModel,  # the observation model
) -> Float[Array, "n+1"]:  # log probabilities $\log p(y_t \vert x_t)$
    """log probabilities $\\log p(y_t | x_t)$"""
    v, B, Omega = obs_model
    y_pred = v + (B @ x[:, :, None])[:, :, 0]
    return MVN(y_pred, Omega).log_prob(y)


def log_prob(x: States, y: Observations, glssm: GLSSM) -> Float:  # $\log p(x,y)$
    """joint log probability of states and observations"""
    log_p_x = jnp.sum(log_probs_x(x, to_states(glssm)))
    log_p_y_given_x = jnp.sum(log_probs_y(y, x, to_observation_model(glssm)))
    return log_p_x + log_p_y_given_x
