# from https://github.com/jwalton3141/jwalton3141.github.io/blob/master/assets/posts/ESS/rwmh.py
# ref Gelman, Andrew, J. B. Carlin, Hal S. Stern, David B. Dunson, Aki Vehtari, and Donald B. Rubin. 2013. Bayesian Data Analysis. Third Edition. London: Chapman & Hall / CRC Press.

import jax.numpy as jnp


def ESS(x):
    """
    Compute the effective sample size of estimand of interest.
    Vectorised implementation.
    x : m_chaines, n_iter, dim

    Output:
    list of ints of length dim: ESS for each dims
    """
    if x.shape < (2,):
        raise ValueError(
            "Calculation of effective sample size"
            "requires multiple chains of the same length."
        )
    try:
        m_chains, n_iter = x.shape  # 1d variable
    except ValueError:
        return [ESS(y.T) for y in x.T]  # apply sequentially to each dimension

    def variogram(t):
        return ((x[:, t:] - x[:, : (n_iter - t)]) ** 2).sum() / (
            m_chains * (n_iter - t)
        )

    post_var = gelman_rubin(x)
    assert post_var > 0

    t = 1
    rho = jnp.ones(n_iter)
    negative_autocorr = False

    # Iterate until the sum of consecutive estimates of autocorrelation is negative
    while not negative_autocorr and (t < n_iter):
        rho = rho.at[t].set(1 - variogram(t) / (2 * post_var))

        if not t % 2:
            negative_autocorr = sum(rho[t - 1 : t + 1]) < 0

        t += 1

    return int(m_chains * n_iter / (1 + 2 * rho[1:t].sum()))


def gelman_rubin(x):
    """
    Estimate the marginal posterior variance. Vectorised implementation.
    x : m_chaines, n_iter
    """
    m_chains, n_iter = x.shape

    # Calculate between-chain variance
    B_over_n = ((jnp.mean(x, axis=1) - jnp.mean(x)) ** 2).sum() / (m_chains - 1)

    # Calculate within-chain variances
    W = ((x - x.mean(axis=1, keepdims=True)) ** 2).sum() / (m_chains * (n_iter - 1))

    # (over) estimate of variance
    s2 = W * (n_iter - 1) / n_iter + B_over_n

    return s2
