import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as st
import statsmodels.api as sm

np.random.seed(42)

def f(y):
    return np.exp(-((y**2 - 0.5) ** 2))


def get_K(auto_covars):
    K = None
    for k in range(int(len(auto_covars) / 2)):
        if (auto_covars[2 * k] + auto_covars[2 * k + 1]) <= 0:
            K = k
            break
    if K is None:
        K = int(len(auto_covars) / 2 - 1)
    return K


def sigma_ipse(auto_covars):
    K = get_K(auto_covars)
    return -auto_covars[0] + 2 * np.sum(
        [auto_covars[2 * k] + auto_covars[2 * k + 1] for k in range(1, K + 1)]
    )


def sigma_imse(auto_covars):
    K = get_K(auto_covars)
    return -auto_covars[0] + 2 * np.sum(
        [
            np.min(
                [auto_covars[2 * j] + auto_covars[2 * j + 1] for j in range(1, k + 1)]
            )
            for k in range(1, K + 1)
        ]
    )


def sigma_bm(Xchain, a=0.5):
    n = len(Xchain)
    Nb = int(np.floor(n**a))
    Nl = int(np.floor(n / Nb))
    batch_means = []
    mcmc_mean = np.mean(Xchain)
    for i in range(1, Nb + 1):
        batch_means.append(np.mean([Xchain[j] for j in range((i - 1) * Nl, i * Nl)]))

    batch_means = np.array(batch_means)
    return np.sum((batch_means - mcmc_mean) ** 2) * n / (Nb - 1) / Nb


def compute_auto_covars(Xchain):
    Xmean = np.mean(Xchain)
    N = len(Xchain)
    autocov = np.zeros(N)
    for k in range(N):
        for i in range(N - k):
            autocov[k] += ((Xchain[i + k]) - Xmean) * (Xchain[i] - Xmean)
        autocov[k] = (1 / (N - 1)) * autocov[k]
    return autocov


# Simulation parameters
alpha = 0.05
coeff = st.norm.ppf(1 - alpha / 2)
sigma = 4.0
p = 2

# Start the chain
Xchain = np.array([0])
n = 0
error1s = []
error2s = []
error3s = []

iteration = 0
iter_min = 20  # force a minimum number of iteration before checking stopping criterion
iter_stop = 1000  # max iteration limit

ns = []
while True:
    # Compute MH step
    Xn = Xchain[n]
    Xnew = st.norm.rvs() * sigma
    uniform = st.uniform.rvs()
    ratio = np.min([1, (f(Xnew) * st.norm.pdf(Xn)) / (f(Xn) * st.norm.pdf(Xnew))])
    if uniform <= ratio:
        Xchain = np.append(Xchain, Xnew)
    else:
        Xchain = np.append(Xchain, Xn)
    n = n + 1

    if n >= iter_min:
        # Compute phi(X)
        Xchain_p = Xchain**p

        # Compute autocovariances
        # Autocovariances computed from scratch, not incrementally.
        # Can be reimplemented for speed
        auto_covars = compute_auto_covars(Xchain_p)

        # Compute variances
        sigma1 = sigma_ipse(auto_covars)
        sigma2 = sigma_imse(auto_covars)
        sigma3 = sigma_bm(Xchain_p)

        error1 = coeff * np.sqrt(sigma1) / np.sqrt(n)
        error2 = coeff * np.sqrt(sigma2) / np.sqrt(n)
        error3 = coeff * np.sqrt(sigma3) / np.sqrt(n)

        print(n, error1, error2, error3)
        error1s.append(error1)
        error2s.append(error2)
        error3s.append(error3)

        ns.append(n)
    if n >= iter_stop:
        break

plt.loglog(ns, error1s, label="ipse")
plt.loglog(ns, error2s, label="imse")
plt.loglog(ns, error3s, label="bm")
plt.loglog(ns, np.array(ns) ** -0.5, "--", color="red", label=r"$n^{-1/2}$")
plt.grid(which="both")
plt.xlabel("Chain length")
plt.ylabel("95% confidence interval")
plt.legend()
plt.show()