from __future__ import division

import numpy as np


def auto_correlation_time(x, s, mu, var):
    b, t, d = x.shape
    act_ = np.zeros([d])
    for i in range(0, b):
        y = x[i] - mu
        p, n = y[:-s], y[s:]
        act_ += np.mean(p * n, axis=0) / var
    act_ = act_ / b
    return act_


def effective_sample_size(x, mu, var):
    """
    Calculate the effective sample size of sequence generated by MCMC.
    Args:
        x: chains with shape batch size, step, dimension
        mu: mean of the variable
        var: variance of the variable

    Returns:
        Effective sample size of the chain
    """
    b, t, d = x.shape
    ess_ = np.ones([d])
    for s in range(1, t):
        p = auto_correlation_time(x, s, mu, var)
        if np.sum(p > 0.05) == 0:
            break
        else:
            for j in range(0, d):
                if p[j] > 0.05:
                    ess_[j] += 2.0 * p[j] * (1.0 - float(s) / t)

    return t / ess_


def acceptance_rate(x):
    """

    Args:
        x: chains with shape batch size, step, dimension

    Returns:

    """
    cnt = x.shape[0] * (x.shape[1] - 1)
    for i in range(0, x.shape[0]):
        for j in range(1, x.shape[1]):
            if np.min(np.equal(x[i, j - 1], x[i, j])):
                cnt -= 1
    return cnt / float(x.shape[0] * (x.shape[1] - 1))


def mean_jump_distance(x):
    """

    Args:
        x: chains with shape batch size, step, dimension

    Returns:

    """
    d = 0
    cnt = x.shape[0] * (x.shape[1] - 1)
    for i in range(0, x.shape[0]):
        for j in range(1, x.shape[1]):
            d += np.linalg.norm(x[i, j - 1] - x[i, j])
    return d / cnt


def gelman_rubin_diagnostic(x, mu=None):
    """

    Args:
        x: chains with shape batch size, step, dimension

    Returns:
    """
    m, n = x.shape[0], x.shape[1]
    theta = np.mean(x, axis=1)
    sigma = np.var(x, axis=1)
    theta_m = mu if mu is not None else np.mean(theta, axis=0)
    b = float(n) / float(m-1) * np.sum((theta - theta_m) ** 2)
    w = 1. / (float(m) * np.sum(sigma, axis=0) + 1e-5)
    v = float(n-1) / float(n) * w + float(m+1) / float(m * n) * b
    r_hat = np.sqrt(v / w)
    return r_hat
