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

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/45_cross_entropy_method.ipynb.

# %% auto 0
__all__ = ['proposal_from_moments', 'simulate_cem', 'log_pdf', 'log_weight_cem', 'posterior_markov_proposal',
           'cross_entropy_method']

# %% ../../nbs/45_cross_entropy_method.ipynb 1
from typing import Tuple

import jax.numpy as jnp
import jax.random as jrn
import jax.scipy as jsp
import tensorflow_probability.substrates.jax.distributions as tfd
from jax import jit, vmap
from jax.lax import fori_loop, scan, while_loop
from jaxtyping import Array, Float, PRNGKeyArray
from tensorflow_probability.substrates.jax.distributions import (
    MultivariateNormalFullCovariance as MVN,
)

from .importance_sampling import ess_pct, normalize_weights
from .laplace_approximation import laplace_approximation
from .pgssm import log_prob as log_prob_joint
from .util import converged

# %% ../../nbs/45_cross_entropy_method.ipynb 6
from functools import partial
from typing import NamedTuple

import jax.scipy.linalg as jsla

from .typing import PGSSM, MarkovProposal, Observations
from isssm.util import (
    degenerate_cholesky,
    location_antithetic,
    mm_sim,
    mm_time_sim,
    scale_antithethic,
)


def proposal_from_moments(
    mean: Float[Array, "n+1 m"],  # mean $v$
    consecutive_covs: Float[
        Array, "n 2*m 2*m"
    ],  # covariances $\operatorname{Cov}((U_t, U_{t + 1}))
) -> MarkovProposal:  # corresponding proposal
    """Find the unique Gaussian Markov proposal that matches means and consecutive covariances."""
    _, m = mean.shape

    # prepend Cov((0, U_0))
    initial_cov = jsla.block_diag(jnp.zeros((m, m)), consecutive_covs[0][:m, :m])
    all_covs = jnp.concatenate(
        (
            initial_cov[None],
            consecutive_covs,
        ),
        axis=0,
    )

    chols = degenerate_cholesky(all_covs)

    # J_tt and J_tp1t dont matter for the first entry
    J_tt = chols[1:, :m, :m]
    J_tp1t = chols[1:, m:, :m]
    Rs = chols[:, m:, m:]

    return MarkovProposal(mean=mean, R=Rs, J_tt=J_tt, J_tp1t=J_tp1t)

# %% ../../nbs/45_cross_entropy_method.ipynb 10
def simulate_cem(
    proposal: MarkovProposal,  # proposal
    N: int,  # number of samples
    key: PRNGKeyArray,  # random number seed
) -> Float[Array, "N n+1 m"]:
    np1, m = proposal.mean.shape
    key, subkey = jrn.split(key)
    u = MVN(loc=jnp.zeros(m), covariance_matrix=jnp.eye(m)).sample((N, np1), subkey)

    # transpose to have time in first dimension
    eps = mm_time_sim(proposal.R, u).transpose((1, 0, 2))
    lower_tri_solve = partial(jsla.solve_triangular, lower=True)

    def _iteration(carry, inputs):
        (x_prev,) = carry
        eps, J_tt, J_tp1t = inputs

        x_next = mm_sim(J_tp1t, vmap(lower_tri_solve, (None, 0))(J_tt, x_prev)) + eps

        return (x_next,), x_next

    J_tt_ext = jnp.concatenate([jnp.eye(m)[None], proposal.J_tt], axis=0)
    J_tp1t_ext = jnp.concatenate([jnp.eye(m)[None], proposal.J_tp1t], axis=0)
    _, x = scan(_iteration, (jnp.zeros((N, m)),), (eps, J_tt_ext, J_tp1t_ext))

    mean = proposal.mean
    samples = x.transpose((1, 0, 2)) + proposal.mean

    u = u.reshape((N, -1))
    l_samples = location_antithetic(samples, mean)
    s_samples = scale_antithethic(u, samples, mean)
    ls_samples = scale_antithethic(u, l_samples, mean)

    return jnp.concatenate((samples, l_samples, s_samples, ls_samples), axis=0)

# %% ../../nbs/45_cross_entropy_method.ipynb 12
from .util import mm_time


def log_pdf(x: Float[Array, "n+1 m"], proposal: MarkovProposal):
    np1, m = x.shape

    lower_tri_solve = partial(jsla.solve_triangular, lower=True)
    centered = x - proposal.mean
    eps = centered[1:] - mm_time(
        proposal.J_tp1t, vmap(lower_tri_solve)(proposal.J_tt, centered[:-1])
    )
    eps = jnp.concatenate((centered[:1], eps), axis=0)

    nu = vmap(lower_tri_solve)(proposal.R, eps)

    return (
        MVN(jnp.zeros(m), jnp.eye(m)).log_prob(nu).sum()
        - jnp.log(vmap(jnp.diag)(proposal.R)).sum()
    )


def log_weight_cem(
    x: Float[Array, "n+1 m"],  # point at which to evaluate the weights
    y: Observations,  # observations
    model: PGSSM,  # modle
    proposal: MarkovProposal,  # proposal
) -> Float:  # log weights
    log_p = log_prob_joint(x, y, model)
    log_g = log_pdf(x, proposal)

    return log_p - log_g

# %% ../../nbs/45_cross_entropy_method.ipynb 15
from .kalman import kalman, smoother
from .typing import GLSSM


def _joint_cov(Xi_smooth_t, Xi_smooth_tp1, Xi_filt_t, Xi_pred_tp1, A_t):
    """Joint covariance of conditional Markov process"""
    off_diag = (
        Xi_filt_t @ A_t.T @ jnp.linalg.pinv(Xi_pred_tp1, hermitian=True) @ Xi_smooth_tp1
    )  # jnp.linalg.solve(Xi_pred_tp1, Xi_smooth_tp1)
    return jnp.block([[Xi_smooth_t, off_diag], [off_diag.T, Xi_smooth_tp1]])


def posterior_markov_proposal(
    y: Observations, model: GLSSM  # observations  # model
) -> MarkovProposal:  # Markov proposal of posterior X|Y
    """calculate the Markov proposal of the smoothing distribution using the Kalman smoother"""
    filtered = kalman(y, model)
    _, Xi_filter, _, Xi_pred = filtered
    x_smooth, Xi_smooth = smoother(filtered, model.A)

    covs = vmap(_joint_cov)(
        Xi_smooth[:-1], Xi_smooth[1:], Xi_filter[:-1], Xi_pred[1:], model.A
    )

    return proposal_from_moments(x_smooth, covs)

# %% ../../nbs/45_cross_entropy_method.ipynb 17
from functools import partial

from .typing import GLSSM


def cross_entropy_method(
    model: PGSSM,  # model
    y: Observations,  # observations
    N: int,  # number of samples to use in the CEM
    key: PRNGKeyArray,  # random number seed
    n_iter: int,  # number of iterations
) -> MarkovProposal:  # the CEM proposal
    """iteratively perform the CEM to find an optimal proposal"""
    key, subkey_crn = jrn.split(key)

    proposal, info = laplace_approximation(y, model, n_iter)
    glssm_la = GLSSM(
        model.u,
        model.A,
        model.D,
        model.Sigma0,
        model.Sigma,
        model.v,
        model.B,
        proposal.Omega,
    )
    initial = posterior_markov_proposal(proposal.z, glssm_la)

    def _iteration(i, vals):
        proposal, _ = vals

        model_log_weights = partial(log_weight_cem, y=y, model=model, proposal=proposal)

        samples = simulate_cem(proposal, N, subkey_crn)

        _N, np1, m = samples.shape

        log_w = vmap(model_log_weights)(samples)
        w = normalize_weights(log_w)

        mean = jnp.sum(w[:, None, None] * samples, axis=0)
        x_consecutive = jnp.concatenate((samples[:, :-1], samples[:, 1:]), axis=-1)

        cov_IS = partial(jnp.cov, aweights=w, rowvar=False)

        # ensure matrices (in 1D case jnp.cov reduces to (np1, m))
        cov_shape = (np1 - 1, 2 * m, 2 * m)
        consecutive_covs = vmap(cov_IS, 1)(x_consecutive).reshape(cov_shape)

        new_proposal = proposal_from_moments(mean, consecutive_covs)

        return new_proposal, log_w

    final_proposal, log_w = fori_loop(
        0, n_iter, _iteration, (initial, jnp.empty(4 * N))
    )

    return final_proposal, log_w
