from pylearn.sampling.mcmc import *

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

    rng = np.random.RandomState(234)

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

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


    position = 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

    # DRAW SAMPLES

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

    assert sampler.avg_acceptance_rate > 0
    assert sampler.avg_acceptance_rate < 1

    # TEST THAT THEY ARE FROM THE RIGHT DISTRIBUTION

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

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


    print 'final stepsize', sampler.stepsize
    print 'final acceptance_rate', sampler.avg_acceptance_rate

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

def test_mcmc():
    print ('MCMC')
    sampler = _sampler_on_2d_gaussian(MCMC_sampler, burnin=3000, n_samples=90000)