# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/99_util.ipynb.

# %% auto 0
__all__ = ['LOFM', 'LOLT', 'mm_sim', 'mm_time', 'mm_time_sim', 'degenerate_cholesky', 'MVN_degenerate', 'converged',
           'append_to_front', 'location_antithetic', 'scale_antithethic']

# %% ../../nbs/99_util.ipynb 2
import jax.numpy as jnp
from jax import vmap
from jaxtyping import Array, Float, Bool
import tensorflow_probability.substrates.jax as tfp
from tensorflow_probability.substrates.jax.distributions import (
    MultivariateNormalLinearOperator as MVNLO,
)

# %% ../../nbs/99_util.ipynb 6
LOFM = tfp.tf2jax.linalg.LinearOperatorFullMatrix
LOLT = tfp.tf2jax.linalg.LinearOperatorLowerTriangular


def degenerate_cholesky(Sigma):
    evals, evecs = jnp.linalg.eigh(Sigma)
    # transpose for QR
    # ensure positive eigenvalues
    sqrt_cov = jnp.einsum("...ij,...j->...ji", evecs, jnp.sqrt(jnp.abs(evals)))
    Q, R = jnp.linalg.qr(sqrt_cov, mode="complete")
    # ensure positive diagonal
    R = R * jnp.sign(jnp.einsum("...ii->...i", R)[..., None])
    L = R.swapaxes(-1, -2)
    return L


def MVN_degenerate(
    loc: Array, cov: Array
) -> tfp.distributions.MultivariateNormalLinearOperator:
    L = degenerate_cholesky(cov)
    return MVNLO(loc=loc, scale=LOLT(L))

# %% ../../nbs/99_util.ipynb 10
def converged(
    new: Float[Array, "..."],  # the new array
    old: Float[Array, "..."],  # the old array
    eps: Float,  # tolerance
) -> Bool:  # whether the arrays are close enough
    """check that sup-norm of relative change is smaller than tolerance"""
    # Avoid division by zero by using absolute difference where old is zero or close to zero
    # Use a small threshold to determine when to use absolute vs relative difference
    threshold = 1e-10
    is_small = jnp.abs(old) < threshold
    # For small values, use absolute difference
    # For larger values, use relative difference
    diff = jnp.where(is_small, jnp.abs(new - old), jnp.abs((new - old) / old))
    is_close = jnp.max(diff) < eps
    any_nans = jnp.isnan(new).sum() > 0
    return jnp.logical_or(is_close, any_nans)

# %% ../../nbs/99_util.ipynb 13
# multiply $B_t$ and $X^i_t$
mm_sim = vmap(jnp.matmul, (None, 0))
# matmul with $(B_t)_{t}$ and $(X_t)_{t}$
mm_time = vmap(jnp.matmul, (0, 0))
# matmul with $(B_t)_{t}$ and $(X^i_t)_{i,t}$
mm_time_sim = vmap(mm_time, (None, 0))

# %% ../../nbs/99_util.ipynb 16
def append_to_front(a0: Float[Array, "..."], a: Float[Array, "n ..."]):
    return jnp.concatenate([a0[None], a], axis=0)

# %% ../../nbs/99_util.ipynb 19
from tensorflow_probability.substrates.jax.distributions import Chi2


def location_antithetic(samples: Float[Array, "N ..."], mean: Float[Array, "N ..."]):
    return 2 * mean[None] - samples


def scale_antithethic(
    u: Float[Array, "N n+1 k"],
    samples: Float[Array, "N n+1 p"],
    mean: Float[Array, "n+1 p"],
):
    N, l = u.shape
    # ensure dtype is Float64
    chi_dist = Chi2(l * jnp.ones(1))

    c = jnp.linalg.norm(u, axis=1) ** 2
    c_prime = chi_dist.quantile(1.0 - chi_dist.cdf(c))

    return mean[None] + jnp.sqrt(c_prime / c)[:, None, None] * (samples - mean[None])
