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

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/60_maximum_likelihood_estimation.ipynb.

# %% auto 0
__all__ = ['vmm', 'gnll', 'gnll_full', 'mle_glssm', 'mle_glssm_ad', 'pgnll', 'initial_theta', 'mle_pgssm']

# %% ../../nbs/60_maximum_likelihood_estimation.ipynb 3
import jax.numpy as jnp
import jax.random as jrn
from jax import vmap
from jaxtyping import Float, Array, PRNGKeyArray
from .kalman import kalman
from jax import jit
from scipy.optimize import minimize as minimize_scipy
from .laplace_approximation import laplace_approximation
from isssm.modified_efficient_importance_sampling import (
    modified_efficient_importance_sampling,
)
from .importance_sampling import normalize_weights
from .typing import GLSSM, PGSSM

# %% ../../nbs/60_maximum_likelihood_estimation.ipynb 6
from .util import MVN_degenerate as MVN

vmm = vmap(jnp.matmul, (0, 0))


@jit
def gnll(
    y: Float[Array, "n+1 p"],  # observations $y_t$
    x_pred: Float[Array, "n+1 m"],  # predicted states $\hat X_{t+1\bar t}$
    Xi_pred: Float[Array, "n+1 m m"],  # predicted state covariances $\Xi_{t+1\bar t}$
    B: Float[Array, "n+1 p m"],  # state observation matrices $B_{t}$
    Omega: Float[Array, "n+1 p p"],  # observation covariances $\Omega_{t}$
) -> Float:  # gaussian negative log-likelihood
    """Gaussian negative log-likelihood"""
    y_pred = vmm(B, x_pred)
    Psi_pred = vmm(vmm(B, Xi_pred), jnp.transpose(B, (0, 2, 1))) + Omega

    return -MVN(y_pred, Psi_pred).log_prob(y).sum()


@jit
def gnll_full(y: Float[Array, "n+1 p"], model: GLSSM):  # observations $y_t$
    filtered = kalman(y, model)
    return gnll(y, filtered.x_pred, filtered.Xi_pred, model.B, model.Omega)

# %% ../../nbs/60_maximum_likelihood_estimation.ipynb 9
from scipy.optimize import minimize as minimize_scipy
from jax.scipy.optimize import minimize as minimize_jax
from scipy.optimize import OptimizeResult
from jax.scipy.optimize import OptimizeResults


def mle_glssm(
    y: Float[Array, "n+1 p"],  # observations $y_t$
    model_fn,  # parameterize GLSSM
    theta0: Float[Array, "k"],  # initial parameter guess
    aux,  # auxiliary data for the model
    options=None,  # options for the optimizer
) -> OptimizeResult:  # result of MLE optimization
    """Maximum likelihood estimation for GLSSM"""

    @jit
    def f(theta: Float[Array, "k"]) -> Float:
        model = model_fn(theta, aux)
        # improve numerical stability by dividing by number of observations
        n_obs = y.size
        return gnll_full(y, model) / n_obs

    return minimize_scipy(f, theta0, method="BFGS", options=options)


def mle_glssm_ad(
    y: Float[Array, "n+1 p"],  # observations $y_t$
    model_fn,  # parameterize GLSSM
    theta0: Float[Array, "k"],  # initial parameter guess
    aux,  # auxiliary data for the model
    options=None,  # options for the optimizer
) -> OptimizeResults:  # result of MLE optimization
    """Maximum likelihood estimation for GLSSM using automatic differentiation"""

    def f(theta: Float[Array, "k"]) -> Float:
        model = model_fn(theta, aux)
        # improve numerical stability by dividing by number of observations
        n_obs = y.size
        return gnll_full(y, model) / n_obs

    return minimize_jax(f, theta0, method="BFGS", options=options)

# %% ../../nbs/60_maximum_likelihood_estimation.ipynb 16
from jax.scipy.special import logsumexp
from .importance_sampling import pgssm_importance_sampling
from .kalman import kalman
from .typing import GLSSM, PGSSM


def _pgnll(
    gnll: Float,  # surrogate gaussian negative log-likelihood
    unnormalized_log_weights: Float[
        Array, "N"
    ],  # unnormalized log-weights $\log w(X^i)$
) -> Float:  # the approximate negative log-likelihood
    """Internal Log-Concave Negative Log-Likelihood"""
    (N,) = unnormalized_log_weights.shape
    weights = normalize_weights(unnormalized_log_weights)
    return (
        gnll - logsumexp(unnormalized_log_weights) + jnp.log(N)
    )  # - (jnp.var(weights) / (2 * N * jnp.mean(weights) ** 2))


def pgnll(
    y: Float[Array, "n+1 p"],  # observations
    model: PGSSM,  # the model
    z: Float[Array, "n+1 p"],  # synthetic observations
    Omega: Float[Array, "n+1 p p"],  # covariance of synthetic observations
    N: int,  # number of samples
    key: PRNGKeyArray,  # random key
) -> Float:  # the approximate negative log-likelihood
    """Log-Concave Negative Log-Likelihood"""

    key, subkey = jrn.split(key)

    _, log_weights = pgssm_importance_sampling(y, model, z, Omega, N, subkey)

    return _pgnll(
        gnll_full(
            z,
            GLSSM(
                model.u,
                model.A,
                model.D,
                model.Sigma0,
                model.Sigma,
                model.v,
                model.B,
                Omega,
            ),
        ),
        log_weights,
    )

# %% ../../nbs/60_maximum_likelihood_estimation.ipynb 19
from .importance_sampling import log_weights
from .laplace_approximation import posterior_mode


def initial_theta(
    y: Float[Array, "n+1 p"],  # observations $y_t$
    model_fn,  # parameterized PGSSM
    theta0: Float[Array, "k"],  # initial parameter guess
    aux,  # auxiliary data for the model
    n_iter_la: int,  # number of LA iterations
    options=None,  # options for the optimizer
    jit_target=True,  # whether to jit the function
):
    """Initial value for Maximum Likelihood Estimation for PGSSMs"""

    def f(theta):
        model = model_fn(theta, aux)

        proposal, info = laplace_approximation(y, model, n_iter_la)

        u, A, D, Sigma0, Sigma, v, B, Omega, z = proposal
        glssm_la = GLSSM(u, A, D, Sigma0, Sigma, v, B, Omega)

        signal = posterior_mode(proposal)
        _, _, x_pred, Xi_pred = kalman(z, glssm_la)

        negloglik = gnll(z, x_pred, Xi_pred, B, Omega) - log_weights(
            signal, y, model.dist, model.xi, z, Omega
        )
        # improve numerical stability by dividing by number of observations
        n_obs = y.size
        return negloglik / n_obs

    if jit_target:
        f = jit(f)

    result = minimize_scipy(f, theta0, method="BFGS", jac="3-point", options=options)
    return result

# %% ../../nbs/60_maximum_likelihood_estimation.ipynb 24
def mle_pgssm(
    y: Float[Array, "n+1 p"],  # observations $y_t$
    model_fn,  # parameterized LCSSM
    theta0: Float[Array, "k"],  # initial parameter guess
    aux,  # auxiliary data for the model
    n_iter_la: int,  # number of LA iterations
    N: int,  # number of importance samples
    key: Array,  # random key
    options=None,  # options for the optimizer
) -> Float[Array, "k"]:  # MLE
    """Maximum Likelihood Estimation for PGSSMs"""

    @jit
    def f(theta, key):
        model = model_fn(theta, aux)

        proposal_la, _ = laplace_approximation(y, model, n_iter_la)

        key, subkey = jrn.split(key)
        proposal_meis, _ = modified_efficient_importance_sampling(
            y, model, proposal_la.z, proposal_la.Omega, n_iter_la, N, subkey
        )

        key, subkey = jrn.split(key)
        # improve numerical stability by dividing by number of observations
        n_obs = y.size
        return pgnll(y, model, proposal_meis.z, proposal_meis.Omega, N, subkey) / n_obs

    key, subkey = jrn.split(key)
    result = minimize_scipy(
        f, theta0, method="BFGS", jac="3-point", options=options, args=(subkey,)
    )
    return result
