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

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/30_laplace_approximation.ipynb.

# %% auto 0
__all__ = ['vmm', 'vdiag', 'vvmap', 'SmoothState', 'PseudoObs', 'PseudoObsCov', 'default_link', 'laplace_approximation',
           'posterior_mode']

# %% ../../nbs/30_laplace_approximation.ipynb 5
from functools import partial

import jax.numpy as jnp
import jax.random as jrn
from jax import grad, jacfwd, jacrev, jit, vmap
from jax.lax import scan, while_loop
from jaxtyping import Array, Float

from .kalman import kalman, smoother
from .pgssm import simulate_pgssm
from .typing import GLSSM, PGSSM, InitialState
from .util import converged

# %% ../../nbs/30_laplace_approximation.ipynb 7
from jax.scipy.optimize import minimize

from .kalman import smoothed_signals
from .typing import ConvergenceInformation, GLSSMProposal, to_glssm

vmm = jit(vmap(jnp.matmul))
vdiag = jit(vmap(jnp.diag))
vvmap = lambda fun: vmap(vmap(fun))

SmoothState = Float[Array, "n+1 m"]
PseudoObs = Float[Array, "n+1 p"]
PseudoObsCov = Float[Array, "n+1 p p"]

default_link = lambda y: jnp.log(y + 1.0)


def _initial_guess(xi_ti, y_ti, dist, link=default_link):
    result = minimize(
        lambda s_ti: -dist(s_ti, xi_ti).log_prob(y_ti).sum(),
        jnp.atleast_1d(default_link(y_ti)),
        method="BFGS",
    )
    return jnp.squeeze(result.x)


def laplace_approximation(
    y: Float[Array, "n+1 p"],  # observation
    model: PGSSM,
    n_iter: int,  # number of iterations
    log_lik=None,  # log likelihood function
    d_log_lik=None,  # derivative of log likelihood function
    dd_log_lik=None,  # second derivative of log likelihood function
    eps: Float = 1e-5,  # precision of iterations
    link=default_link,  # default link to use in initial guess
) -> tuple[GLSSMProposal, ConvergenceInformation]:
    u, A, D, Sigma0, Sigma, v, B, dist, xi = model
    np1, p, m = B.shape

    s_init = vvmap(partial(_initial_guess, dist=dist, link=link))(xi, y)

    def default_log_lik(s_ti, xi_ti, y_ti):
        return dist(s_ti, xi_ti).log_prob(y_ti).sum()

    if log_lik is None:
        log_lik = default_log_lik

    if d_log_lik is None:
        d_log_lik = jacfwd(log_lik, argnums=0)
    if dd_log_lik is None:
        dd_log_lik = jacrev(d_log_lik, argnums=0)

    vd_log_lik = jit(vvmap(d_log_lik))
    vdd_log_lik = jit(vvmap(dd_log_lik))

    def _break(val):
        _, i, z, Omega, z_old, Omega_old = val

        z_converged = jnp.logical_and(converged(z, z_old, eps), i > 0)
        Omega_converged = jnp.logical_and(converged(Omega, Omega_old, eps), i > 0)
        all_converged = jnp.logical_and(z_converged, Omega_converged)

        iteration_limit_reached = i >= n_iter

        return jnp.logical_or(all_converged, iteration_limit_reached)

    def _iteration(val):
        s, i, z_old, Omega_old, _, _ = val

        grad = vd_log_lik(s, xi, y)
        Gamma = -vdd_log_lik(s, xi, y)
        # assume hessian is diagonal
        Omega = vdiag(1.0 / Gamma)

        z = s + grad / Gamma
        approx_glssm = GLSSM(u, A, D, Sigma0, Sigma, v, B, Omega)

        filtered = kalman(z, approx_glssm)
        s_new = smoothed_signals(filtered, z, approx_glssm)

        return s_new, i + 1, z, Omega, z_old, Omega_old

    empty_z = jnp.empty_like(s_init)
    empty_Omega = jnp.empty((np1, p, p))
    init = (s_init, 0, empty_z, empty_Omega, empty_z, empty_Omega)

    _keep_going = lambda *args: jnp.logical_not(_break(*args))
    _, n_iters, z, Omega, z_old, Omega_old = while_loop(_keep_going, _iteration, init)

    final_proposal = GLSSMProposal(u, A, D, Sigma0, Sigma, v, B, Omega, z)
    delta_z = jnp.max(jnp.abs(z - z_old))
    delta_Omega = jnp.max(jnp.abs(Omega - Omega_old))
    information = ConvergenceInformation(
        converged=jnp.logical_and(
            converged(z, z_old, eps), converged(Omega, Omega_old, eps)
        ),
        n_iter=n_iters,
        delta=jnp.max(jnp.array([delta_z, delta_Omega])),
    )
    return final_proposal, information


def posterior_mode(proposal: GLSSMProposal) -> Float[Array, "n+1 p"]:
    glssm = to_glssm(proposal)
    return smoothed_signals(kalman(proposal.z, glssm), proposal.z, glssm)
