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

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/50_modified_efficient_importance_sampling.ipynb.

# %% auto 0
__all__ = ['optimal_parameters', 'modified_efficient_importance_sampling']

# %% ../../nbs/50_modified_efficient_importance_sampling.ipynb 5
import jax.numpy as jnp
import jax.random as jrn
from jaxtyping import Float, Array, PRNGKeyArray
from jax import vmap, jit
from .util import converged
from .importance_sampling import log_weights_t, normalize_weights
from functools import partial
from jax.lax import while_loop
from .kalman import kalman, simulation_smoother
from jax.lax import scan
from .util import MVN_degenerate as MVN, mm_sim

from .glssm import mm_sim
from .typing import GLSSM, PGSSM, GLSSMProposal, ConvergenceInformation


@jit
def optimal_parameters(
    signal: Float[Array, "N p"], weights: Float[Array, "N"], log_p: Float[Array, "N"]
):
    ones = jnp.ones_like(weights)[:, None]

    sqrt_W = jnp.sqrt(weights)
    X = jnp.concatenate([ones, signal, -0.5 * signal**2], axis=-1)

    beta, *_ = jnp.linalg.lstsq(
        jnp.einsum("i, ij -> ij", sqrt_W, X), jnp.einsum("i, i -> i", sqrt_W, log_p)
    )

    return beta


def modified_efficient_importance_sampling(
    y: Float[Array, "n+1 p"],  # observations
    model: PGSSM,  # model
    z_init: Float[Array, "n+1 p"],  # initial z estimate
    Omega_init: Float[Array, "n+1 p p"],  # initial Omega estimate
    n_iter: int,  # number of iterations
    N: int,  # number of samples
    key: PRNGKeyArray,  # random key
    eps: Float = 1e-5,  # convergence threshold
) -> tuple[GLSSMProposal, ConvergenceInformation]:
    z, Omega = z_init, Omega_init

    np1, p, m = model.B.shape

    key, crn_key = jrn.split(key)

    v_norm_w = vmap(normalize_weights)
    dist = model.dist
    lw_t = vmap(
        vmap(lambda s, y, xi, z, Omega: log_weights_t(s, y, xi, dist, z, Omega)),
        (0, None, None, None, None),
    )

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

        # in first iteration we don't have old values, converged is True for NaNs
        z_converged = jnp.logical_and(converged(z, z_old, eps), i > 0)
        Omega_converged = jnp.logical_and(converged(Omega, Omega_old, eps), i > 0)
        iteration_limit_reached = i >= n_iter

        return jnp.logical_or(
            jnp.logical_and(z_converged, Omega_converged), iteration_limit_reached
        )

    def _iteration(val):
        i, z, Omega, _, _ = val
        glssm_approx = GLSSM(
            model.u,
            model.A,
            model.D,
            model.Sigma0,
            model.Sigma,
            model.v,
            model.B,
            Omega,
        )
        sim_signal = simulation_smoother(glssm_approx, z, N, crn_key)

        log_weights = lw_t(sim_signal, y, model.xi, z, Omega)
        log_p = dist(sim_signal, model.xi).log_prob(y).sum(axis=-1)
        wls_estimate = vmap(optimal_parameters, (1, 1, 1), 0)(
            sim_signal, v_norm_w(log_weights), log_p
        )

        a = wls_estimate[:, 0]
        b = wls_estimate[:, 1 : (p + 1)]
        c = wls_estimate[:, (p + 1) :]

        z_new = b / c
        Omega_new = vmap(jnp.diag)(1 / c)

        return i + 1, z_new, Omega_new, z, Omega

    _keep_going = lambda *args: jnp.logical_not(_break(*args))

    n_iters, z, Omega, z_old, Omega_old = while_loop(
        _keep_going,
        _iteration,
        (0, z_init, Omega_init, jnp.empty_like(z_init), jnp.empty_like(Omega_init)),
    )

    proposal = GLSSMProposal(
        u=model.u,
        A=model.A,
        D=model.D,
        Sigma0=model.Sigma0,
        Sigma=model.Sigma,
        v=model.v,
        B=model.B,
        Omega=Omega,
        z=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 proposal, information
