"""See also the corresponding [section in my thesis](https://stefanheyder.github.io/dissertation/thesis.pdf#nameddest=section.3.3)"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/20_pgssm.ipynb.

# %% auto 0
__all__ = ['simulate_pgssm', 'nb_pgssm_running_example', 'log_probs_y', 'log_prob']

# %% ../../nbs/20_pgssm.ipynb 5
import jax.numpy as jnp
import jax.random as jrn
from jax import lax, vmap
from jaxtyping import Array, Float, PRNGKeyArray
from tensorflow_probability.substrates.jax.distributions import (
    NegativeBinomial as NBinom,
)
from tensorflow_probability.substrates.jax.distributions import Poisson

from .glssm import log_probs_x, simulate_states
from .typing import PGSSM, GLSSMObservationModel, GLSSMState
from .util import mm_time, mm_time_sim

# %% ../../nbs/20_pgssm.ipynb 6
def simulate_pgssm(
    pgssm: PGSSM,
    N: int,  # number of samples
    key: PRNGKeyArray,  # random key
) -> tuple[
    Float[Array, "N n+1 m"], Float[Array, "N n+1 p"]
]:  # simulated states and observations
    u, A, D, Sigma0, Sigma, v, B, dist, xi = pgssm
    key, subkey = jrn.split(key)
    X = simulate_states(GLSSMState(u, A, D, Sigma0, Sigma), N, subkey)
    S = v + mm_time_sim(B, X)

    key, subkey = jrn.split(key)
    Y = dist(S, xi).sample(seed=subkey)

    return X, Y

# %% ../../nbs/20_pgssm.ipynb 8
import jax.scipy.linalg as jsla

from .models.pgssm import nb_pgssm
from .models.stsm import stsm


def nb_pgssm_running_example(
    x0_trend: Float[Array, "m"] = jnp.zeros(2),
    r: Float = 20.0,
    s2_trend: Float = 0.01,
    s2_speed: Float = 0.1,
    alpha: Float = 0.1,
    omega2: Float = 0.01,
    n: int = 100,
    x0_seasonal: Float[Array, "s"] = jnp.zeros(4),
    s2_seasonal: Float = 0.1,
    Sigma0_seasonal: Float[Array, "s s"] = 0.1 * jnp.eye(4),
    s_order: int = 5,
) -> PGSSM:  # the running example for this package
    """a structural time series model with NBinom observations"""
    Sigma0 = jsla.block_diag(0.01 * jnp.eye(2), Sigma0_seasonal)
    x0 = jnp.concatenate((x0_trend, x0_seasonal))

    model = nb_pgssm(
        stsm(
            x0,
            s2_trend,
            s2_speed,
            s2_seasonal,
            n,
            Sigma0,
            omega2,
            s_order,
            alpha,
        ),
        r,
    )
    return model

# %% ../../nbs/20_pgssm.ipynb 12
from .typing import to_states


def log_probs_y(
    x: Float[Array, "n+1 m"],  # states
    y: Float[Array, "n+1 p"],  # observations
    v: Float[Array, "n+1 p"],  # signal biases
    B: Float[Array, "n+1 p m"],  # signal matrices
    dist,  # observation distribution
    xi,  # observation parameters
):
    s = v + (B @ x[:, :, None])[:, :, 0]
    return dist(s, xi).log_prob(y).sum(axis=1)


def log_prob(
    x: Float[Array, "n+1 m"],  # states
    y: Float[Array, "n+1 p"],  # observations
    model: PGSSM,
):
    u, A, D, Sigma0, Sigma, v, B, dist, xi = model
    px = log_probs_x(x, to_states(model)).sum()
    py = log_probs_y(x, y, v, B, dist, xi).sum()
    return px + py
