from __future__ import division, print_function
import numpy as np


def sample_HMC(f, num_samples, Lmax, epsilon, x0, verbose=False,
               thin=1, burn=0, RNG=np.random.RandomState(0)):
    """
    A straight-forward HMC implementation. The mass matrix is assumed to be the
    identity.

    f is a python function that returns the energy and its gradient

      f(x) = E(x), dE(x)/dx

    we then generate samples from the distribution

      pi(x) = exp(-E(x))/Z

    - num_samples is the number of samples to generate.
    - Lmax, epsilon are parameters of the hmc procedure to be tuned.
    - x0 is the starting position for the procedure.
    - verbose is a flag which turns on the display of the running accept ratio.
    - thin is an integer which specifies the thinning interval
    - burn is an integer which specifies how many initial samples to discard.
    - RNG is a random number generator

    The total number of iterations is given by

      burn + thin * num_samples

    the return shape is always num_samples x D
    """

    # burn some samples if needed.
    if burn > 0:
        if verbose:
            print("burn-in sampling started")
        samples = sample_HMC(f, num_samples=burn, Lmax=Lmax,
                             epsilon=epsilon, x0=x0,
                             verbose=verbose, thin=1, burn=0)
        if verbose:
            print("burn-in sampling ended")
        x0 = samples[-1]

    D = x0.size
    samples = np.zeros((num_samples, D))
    samples[0] = x0.copy()
    x = x0.copy()
    logprob, grad = f(x0)
    logprob, grad = -logprob, -grad

    accept_count_batch = 0

    for t in range(1, num_samples * thin):

        # Output acceptance rate every 100 iterations
        if(((t+1) % 100) == 0):
            if verbose:
                print("Iteration: ", t+1,
                      "\t Acc Rate: ", 1. * accept_count_batch, "%")
            accept_count_batch = 0

        # make a copy of the old state.
        x_old, logprob_old, grad_old = x.copy(), logprob, grad.copy()
        p_old = RNG.randn(D)

        # Standard HMC - begin leapfrogging
        premature_reject = False
        p = p_old + 0.5 * epsilon * grad
        for l in range(RNG.randint(1, Lmax + 1)):
            x += epsilon * p
            logprob, grad = f(x)
            logprob, grad = -logprob, -grad
            if np.any(np.isnan(grad)):  # pragma: no cover
                premature_reject = True
                break
            p += epsilon * grad
        p -= 0.5*epsilon * grad
        # leapfrogging done

        # reject the proposal if there are numerical errors.
        if premature_reject:  # pragma: no cover
            print("warning: numerical instability.\
                  Rejecting this proposal prematurely")
            x, logprob, grad = x_old, logprob_old, grad_old
            if t % thin == 0:
                samples[t / thin] = x_old
            continue

        # work out whether to accpet the proposal
        log_accept_ratio = logprob - 0.5 * p.dot(p) -\
            logprob_old + 0.5 * p_old.dot(p_old)
        logu = np.log(RNG.rand())

        if logu < log_accept_ratio:  # accept
            if t % thin == 0:
                samples[t/thin] = x
            accept_count_batch += 1
        else:  # reject
            if t % thin == 0:
                samples[t/thin] = x_old
            x, logprob, grad = x_old, logprob_old, grad_old
    return samples