import numpy as np
from ..hmc import hmc
import matplotlib.pyplot as pp


def lnprob_gaussian(x, icov):
    logp = -np.dot(x, np.dot(icov, x)) / 2.0
    grad = -np.dot(x, icov)
    return logp, grad


def test_1():
    dim = 2
    x0 = np.zeros(dim)
    cov = np.array([[1, 1.98], [1.98, 4]])
    # cov = np.array([[1, 0], [0, 1]])
    icov = np.linalg.inv(cov)

    samples, diag = hmc(lnprob_gaussian, x0, args=(icov,),
                  n_samples=10**3, n_omit=0,
                  steps=10, epsilon=0.2, return_diagnostics=True,
                  display=True)

    pp.plot(samples[:, 0], samples[:, 1])
    pp.show()


# def test_hmc():
#     import f_energy as en
#     dim = 5
#     x0 = np.zeros(dim)
#     M = np.eye(dim)
#     samples = hmc(en.f_multivariate_normal, x0, en.g_multivariate_normal, args=(M,),
#                   nsamples=10**3,nomit=10**3,steps=100,stepadj=.05)
#     C = np.cov(samples,rowvar=0,bias=1)
#     np.testing.assert_array_almost_equal(M,np.linalg.inv(C),1)

@@ -290,47 +290,8 @@
 %5D%5D)%0A
-    # cov = np.array(%5B%5B1, 0%5D, %5B0, 1%5D%5D)%0A
     
@@ -332,14 +332,8 @@
 ples
-, diag
  = h
@@ -402,17 +402,17 @@
 les=10**
-3
+4
 , n_omit
@@ -412,17 +412,21 @@
  n_omit=
-0
+10**3
 ,%0A      
@@ -444,17 +444,17 @@
  steps=1
-0
+5
 , epsilo
@@ -458,16 +458,17 @@
 ilon=0.2
+5
 , return
@@ -484,371 +484,49 @@
 ics=
-True,%0A                  display=True)%0A%0A    pp.plot(samples%5B:, 0%5D, samples%5B:, 1%5D)%0A    pp.show()%0A%0A%0A# def test_hmc():%0A#     import f_energy as en%0A#     dim = 5%0A#     x0 = np.zeros(dim)%0A#     M = np.eye(dim)%0A#     samples = hmc(en.f_multivariate_normal, x0, en.g_multivariate_normal, args=(M,),%0A#                   nsamples=10**3,nomit=10**3,steps=100,stepadj=.05)%0A# 
+False,%0A                  display=False)%0A%0A
     
@@ -544,16 +544,17 @@
 samples,
+ 
 rowvar=0
@@ -558,16 +558,17 @@
 r=0,
+ 
 bias=1)%0A
 #   
@@ -567,10 +567,8 @@
 =1)%0A
-# 
     
@@ -608,26 +608,15 @@
 ual(
-M,np.linalg.inv(C),
+cov, C, 
 1)%0A

clean up test