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

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/40_importance_sampling.ipynb.

# %% auto 0
__all__ = ['prediction_percentiles', 'log_weights_t', 'log_weights', 'pgssm_importance_sampling', 'normalize_weights', 'ess',
           'ess_lw', 'ess_pct', 'mc_integration', 'future_prediction_interval', 'predict', 'prediction']

# %% ../../nbs/40_importance_sampling.ipynb 4
from functools import partial

import jax.numpy as jnp
from jax import vmap
from jaxtyping import Array, Float

from tensorflow_probability.substrates.jax.distributions import (
    MultivariateNormalFullCovariance as MVN,
)

from .typing import PGSSM


def log_weights_t(
    s_t: Float[Array, "p"],  # signal
    y_t: Float[Array, "p"],  # observation
    xi_t: Float[Array, "p"],  # parameters
    dist,  # observation distribution
    z_t: Float[Array, "p"],  # synthetic observation
    Omega_t: Float[Array, "p p"],  # synthetic observation covariance, assumed diagonal
) -> Float:  # single log weight
    """Log weight for a single time point."""
    p_ys = dist(s_t, xi_t).log_prob(y_t).sum()

    # omega_t = jnp.sqrt(jnp.diag(Omega_t))
    # g_zs = MVN_diag(s_t, omega_t).log_prob(z_t).sum()
    g_zs = MVN(s_t, Omega_t).log_prob(z_t).sum()

    return p_ys - g_zs


def log_weights(
    s: Float[Array, "n+1 p"],  # signals
    y: Float[Array, "n+1 p"],  # observations
    dist,  # observation distribution
    xi: Float[Array, "n+1 p"],  # observation parameters
    z: Float[Array, "n+1 p"],  # synthetic observations
    Omega: Float[Array, "n+1 p p"],  # synthetic observation covariances:
) -> Float:  # log weights
    """Log weights for all time points"""
    p_ys = dist(s, xi).log_prob(y).sum()

    # avoid triangular solve problems
    # omega = jnp.sqrt(vmap(jnp.diag)(Omega))
    # g_zs = MVN_diag(s, omega).log_prob(z).sum()
    g_zs = MVN(s, Omega).log_prob(z).sum()

    return p_ys - g_zs

# %% ../../nbs/40_importance_sampling.ipynb 7
from functools import partial

import jax.random as jrn

from jaxtyping import Array, Float, PRNGKeyArray

from .kalman import FFBS, simulation_smoother
from .typing import GLSSM, PGSSM


def pgssm_importance_sampling(
    y: Float[Array, "n+1 p"],  # observations
    model: PGSSM,  # 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
) -> tuple[
    Float[Array, "N n+1 m"], Float[Array, "N"]
]:  # importance samples and weights
    u, A, D, Sigma0, Sigma, v, B, dist, xi = model

    glssm = GLSSM(u, A, D, Sigma0, Sigma, v, B, Omega)

    key, subkey = jrn.split(key)
    s = simulation_smoother(glssm, z, N, subkey)

    model_log_weights = partial(log_weights, y=y, dist=dist, xi=xi, z=z, Omega=Omega)

    lw = vmap(model_log_weights)(s)

    return s, lw

# %% ../../nbs/40_importance_sampling.ipynb 13
from jaxtyping import Array, Float


def normalize_weights(
    log_weights: Float[Array, "N"],  # log importance sampling weights
) -> Float[Array, "N"]:  # normalized importance sampling weights
    """Normalize importance sampling weights."""
    max_weight = jnp.max(log_weights)

    log_weights_corrected = log_weights - max_weight

    weights = jnp.exp(log_weights_corrected)

    return weights / weights.sum()

# %% ../../nbs/40_importance_sampling.ipynb 16
from jaxtyping import Array, Float


def ess(
    normalized_weights: Float[Array, "N"],  # normalized weights
) -> Float:  # the effective sample size
    """Compute the effective sample size of a set of normalized weights"""
    return 1 / (normalized_weights**2).sum()


def ess_lw(
    log_weights: Float[Array, "N"],  # the log weights
) -> Float:  # the effective sample size
    """Compute the effective sample size of a set of log weights"""
    return ess(normalize_weights(log_weights))


def ess_pct(
    log_weights: Float[Array, "N"],  # log weights
) -> Float:  # the effective sample size in percent, also called efficiency factor
    (N,) = log_weights.shape
    return ess_lw(log_weights) / N * 100

# %% ../../nbs/40_importance_sampling.ipynb 19
from jax import jit


@jit
def mc_integration(samples: Float[Array, "N ..."], log_weights: Float[Array, "N"]):
    return jnp.einsum("i...,i->...", samples, normalize_weights(log_weights))

# %% ../../nbs/40_importance_sampling.ipynb 22
from jax import jit
from scipy.optimize import minimize

from .glssm import simulate_states
from .kalman import kalman
from .typing import GLSSMProposal, GLSSMState
from .util import mm_time_sim


def future_prediction_interval(dist, signal_samples, xi, log_weights, p):
    def integer_ecdf(y):
        return (
            dist(signal_samples, xi).cdf(y).squeeze(axis=-1)
            * normalize_weights(log_weights)
        ).sum()

    def ecdf(y):
        y_floor = jnp.floor(y)
        y_ceil = jnp.ceil(y)
        y_gauss = y - y_floor

        return integer_ecdf(y_floor) * (1 - y_gauss) + integer_ecdf(y_ceil) * y_gauss

    def pinball_loss(y, p):
        return (jnp.abs(ecdf(y) - p).sum()) ** 2

    mean = mc_integration(dist(signal_samples, xi).mean(), log_weights)
    result = minimize(pinball_loss, mean, args=(p,), method="Nelder-Mead")
    return result.x


def _prediction_percentiles(Y, weights, probs):
    Y_sorted = jnp.sort(Y)
    weights_sorted = weights[jnp.argsort(Y)]
    cumsum = jnp.cumsum(weights_sorted)

    # find indices of cumulative sum closest to probs
    # take corresponding Y_sorted values
    # with linear interpolation if necessary

    indices = jnp.searchsorted(cumsum, probs)
    indices = jnp.clip(indices, 1, len(Y_sorted) - 1)
    left_indices = indices - 1
    right_indices = indices
    left_cumsum = cumsum[left_indices]
    right_cumsum = cumsum[right_indices]
    left_Y = Y_sorted[left_indices]
    right_Y = Y_sorted[right_indices]
    # linear interpolation
    quantiles = left_Y + (probs - left_cumsum) / (right_cumsum - left_cumsum) * (
        right_Y - left_Y
    )
    return quantiles


prediction_percentiles = vmap(
    vmap(_prediction_percentiles, (1, None, None), 1), (2, None, None), 2
)


def predict(
    model: PGSSM,
    y: Float[Array, "n+1 p"],
    proposal: GLSSMProposal,
    future_model: PGSSM,
    N: int,
    key: PRNGKeyArray,
):
    key, subkey = jrn.split(key)
    signal_samples, log_weights = pgssm_importance_sampling(
        y, model, proposal.z, proposal.Omega, N, subkey
    )
    (N,) = log_weights.shape

    signal_model = GLSSM(
        proposal.u,
        proposal.A,
        proposal.D,
        proposal.Sigma0,
        proposal.Sigma,
        proposal.v,
        proposal.B,
        proposal.Omega,
    )

    @jit
    def future_sample(signal_sample, key):
        x_filt, Xi_filt, _, _ = kalman(signal_sample, signal_model)
        state = GLSSMState(
            future_model.u.at[0].set(x_filt[-1]),
            future_model.A,
            future_model.D,
            Xi_filt[-1],
            future_model.Sigma,
        )

        (x,) = simulate_states(state, 1, key)
        return x

    key, *subkeys = jrn.split(key, N + 1)
    subkeys = jnp.array(subkeys)

    future_x = vmap(future_sample)(signal_samples, subkeys)
    future_s = mm_time_sim(future_model.B, future_x)
    future_y = future_model.dist(future_s, future_model.xi).mean()

    return (future_x, future_s, future_y), log_weights

# %% ../../nbs/40_importance_sampling.ipynb 24
from .kalman import to_signal_model


def prediction(
    f: callable,
    y,
    proposal: GLSSMProposal,
    model: PGSSM,
    N: int,
    key: PRNGKeyArray,
    probs: Float[Array, "k"],
    prediction_model=None,
):
    if prediction_model is None:
        prediction_model = model

    key, subkey = jrn.split(key)
    signal_samples, log_weights = pgssm_importance_sampling(
        y, model, proposal.z, proposal.Omega, N, subkey
    )
    N = signal_samples.shape[0]

    signal_model = to_signal_model(proposal)

    def state_sample(signal_sample, key):
        (x_sample,) = FFBS(signal_sample, signal_model, 1, key)
        return x_sample

    key, *subkeys = jrn.split(key, N + 1)
    subkeys = jnp.array(subkeys)
    x_samples = vmap(state_sample)(signal_samples, subkeys)
    s_samples = mm_time_sim(prediction_model.B, x_samples)

    key, subkey = jrn.split(key)
    y_prime_samples = prediction_model.dist(
        s_samples, prediction_model.xi[None]
    ).sample(seed=subkey)

    f_samples = vmap(f)(x_samples, s_samples, y_prime_samples)

    mean_f = mc_integration(f_samples, log_weights)
    sd_f = jnp.sqrt(mc_integration(f_samples**2, log_weights) - mean_f**2)

    if f_samples.ndim == 3:
        percentiles = prediction_percentiles(
            f_samples, normalize_weights(log_weights), probs
        )
    elif f_samples.ndim == 2:
        percentiles = vmap(_prediction_percentiles, (1, None, None), 1)(
            f_samples, normalize_weights(log_weights), probs
        )
    elif f_samples.ndim == 1:
        percentiles = _prediction_percentiles(
            f_samples, normalize_weights(log_weights), probs
        )

    return mean_f, sd_f, percentiles
