import numpy
import theano
from theano import tensor

from pylearn.sampling.hmc import HMC_sampler


def _sampler_on_2d_gaussian(sampler_cls, burnin, n_samples):
    batchsize = 3

    rng = numpy.random.RandomState(234)

    #
    # Define a covariance and mu for a gaussian
    #
    tmp = rng.randn(2, 2).astype(theano.config.floatX)
    tmp[0] += tmp[1]  # induce some covariance
    cov = numpy.dot(tmp, tmp.T)
    cov_inv = numpy.linalg.inv(cov).astype(theano.config.floatX)
    mu = numpy.asarray([5, 9.5], dtype=theano.config.floatX)

    def gaussian_energy(xlist):
        x = xlist
        return 0.5 * (tensor.dot((x - mu), cov_inv) * (x - mu)).sum(axis=1)

    position = theano.shared(rng.randn(batchsize,
                                       2).astype(theano.config.floatX))
    sampler = sampler_cls(position, gaussian_energy)

    print 'initial position', position.get_value(borrow=True)
    print 'initial stepsize', sampler.stepsize.get_value(borrow=True)

    # DRAW SAMPLES
    samples = [sampler.draw() for r in xrange(burnin)]  # burn-in
    samples = numpy.asarray([sampler.draw() for r in xrange(n_samples)])

    assert sampler.avg_acceptance_rate.get_value() > 0
    assert sampler.avg_acceptance_rate.get_value() < 1

    # TEST THAT THEY ARE FROM THE RIGHT DISTRIBUTION

    # samples.shape == (1000, 3, 2)

    print 'target mean:', mu
    print 'empirical mean: ', samples.mean(axis=0)
    #assert numpy.all(abs(mu - samples.mean(axis=0)) < 1)

    print 'final stepsize', sampler.stepsize.get_value()
    print 'final acceptance_rate', sampler.avg_acceptance_rate.get_value()

    print 'target cov', cov
    s = samples[:, 0, :]
    empirical_cov = numpy.cov(samples[:, 0, :].T)
    print ''
    print 'cov/empirical_cov', cov / empirical_cov
    empirical_cov = numpy.cov(samples[:, 1, :].T)
    print 'cov/empirical_cov', cov / empirical_cov
    empirical_cov = numpy.cov(samples[:, 2, :].T)
    print 'cov/empirical_cov', cov / empirical_cov
    return sampler


def test_hmc():
    print ('HMC')
    try:
        orig_mode = theano.config.mode
        if theano.config.mode in ["DebugMode", "DEBUG_MODE"]:
            theano.config.mode = "FAST_RUN"
            sampler = _sampler_on_2d_gaussian(
                HMC_sampler.new_from_shared_positions,
                burnin=3000 / 20, n_samples=90000 / 20)
    finally:
        theano.config.mode = orig_mode


    assert abs(sampler.avg_acceptance_rate.get_value() -
               sampler.target_acceptance_rate) < .1
    assert sampler.stepsize.get_value() >= sampler.stepsize_min
    assert sampler.stepsize.get_value() <= sampler.stepsize_max