"""types used throughout this package"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/99_typings.ipynb.

# %% auto 0
__all__ = ['InitialState', 'Observations', 'States', 'GLSSMState', 'GLSSMObservationModel', 'GLSSM', 'FilterResult',
           'SmootherResult', 'PGSSM', 'to_states', 'to_observation_model', 'GLSSMProposal', 'to_glssm',
           'ConvergenceInformation', 'MarkovProposal']

# %% ../../nbs/99_typings.ipynb 2
from typing import NamedTuple
import tensorflow_probability.substrates.jax.distributions as tfd
from jaxtyping import Float, Array, Bool

# %% ../../nbs/99_typings.ipynb 3
InitialState = Float[Array, "m"]
Observations = Float[Array, "n+1 p"]
States = Float[Array, "n+1 m"]

# %% ../../nbs/99_typings.ipynb 5
class GLSSMState(NamedTuple):
    u: Float[Array, "n+1 m"]  # state bias
    A: Float[Array, "n m m"]  # state transition matrix
    D: Float[Array, "n m l"]  # noise embedding matrix
    Sigma0: Float[Array, "m m"]  # initial state covariance matrix
    Sigma: Float[Array, "n l l"]  # state covariance matrix


class GLSSMObservationModel(NamedTuple):
    v: Float[Array, "n+1 p"]  # observation bias
    B: Float[Array, "n+1 p m"]  # observation matrix
    Omega: Float[Array, "n+1 p p"]  # observation covariance matrix


class GLSSM(NamedTuple):
    u: Float[Array, "n+1 m"]  # state bias
    A: Float[Array, "n m m"]  # state transition matrix
    D: Float[Array, "n m l"]  # noise embedding matrix
    Sigma0: Float[Array, "m m"]  # initial state covariance matrix
    Sigma: Float[Array, "n l l"]  # state covariance matrix
    v: Float[Array, "n+1 p"]  # observation bias
    B: Float[Array, "n+1 p m"]  # observation matrix
    Omega: Float[Array, "n+1 p p"]  # observation covariance matrix

# %% ../../nbs/99_typings.ipynb 7
class FilterResult(NamedTuple):
    x_filt: Float[Array, "n+1 m"]
    Xi_filt: Float[Array, "n+1 m m"]
    x_pred: Float[Array, "n+1 m"]
    Xi_pred: Float[Array, "n+1 m m"]


class SmootherResult(NamedTuple):
    x_smooth: Float[Array, "n+1 m"]
    Xi_smooth: Float[Array, "n+1 m m"]

# %% ../../nbs/99_typings.ipynb 9
class PGSSM(NamedTuple):
    u: Float[Array, "n+1 m"]
    A: Float[Array, "n m m"]
    D: Float[Array, "n m l"]
    Sigma0: Float[Array, "m m"]
    Sigma: Float[Array, "n l l"]
    v: Float[Array, "n+1 p"]
    B: Float[Array, "n+1 p m"]
    dist: tfd.Distribution
    xi: Float[Array, "n+1 p"]


def to_states(model: GLSSM | PGSSM) -> GLSSMState:
    return GLSSMState(
        u=model.u, A=model.A, D=model.D, Sigma0=model.Sigma0, Sigma=model.Sigma
    )


def to_observation_model(model: GLSSM) -> GLSSMObservationModel:
    return GLSSMObservationModel(v=model.v, B=model.B, Omega=model.Omega)

# %% ../../nbs/99_typings.ipynb 11
class GLSSMProposal(NamedTuple):
    u: Float[Array, "n+1 m"]
    A: Float[Array, "n m m"]
    D: Float[Array, "n m l"]
    Sigma0: Float[Array, "m m"]
    Sigma: Float[Array, "n l l"]
    v: Float[Array, "n+1 p"]
    B: Float[Array, "n+1 p m"]
    Omega: Float[Array, "n+1 p p"]
    z: Float[Array, "n+1 p"]


def to_glssm(proposal: GLSSMProposal) -> GLSSM:
    return GLSSM(
        proposal.u,
        proposal.A,
        proposal.D,
        proposal.Sigma0,
        proposal.Sigma,
        proposal.v,
        proposal.B,
        proposal.Omega,
    )


class ConvergenceInformation(NamedTuple):
    converged: Bool
    n_iter: int
    delta: Float

# %% ../../nbs/99_typings.ipynb 13
class MarkovProposal(NamedTuple):
    mean: Float[Array, "n+1 m"]
    R: Float[Array, "n+1 m m"]
    J_tt: Float[Array, "n m m"]  # lower triangular
    J_tp1t: Float[Array, "n m m"]
