from kernel_exp_family.estimators.finite.gaussian import KernelExpFiniteGaussian
from kernel_exp_family.estimators.lite.gaussian import KernelExpLiteGaussian
from kernel_exp_family.examples.tools import visualise_array, pdf_grid
from kernel_hmc.densities.gaussian import IsotropicZeroMeanGaussian
from kernel_hmc.proposals.kmc import KMCStatic
from kernel_hmc.tools.log import Log
import matplotlib.pyplot as plt
import numpy as np


# depends on theano, which is an optional dependency
try:
    from kernel_hmc.densities.bananas import Banana, sample_banana
except ImportError:
    print("Skipping example, depends on theano which is unavailable")
    exit(0)

Log.set_loglevel(20)

def visualise_trajectory(Qs, acc_probs, log_pdf_q, target_log_pdf=None):
    assert Qs.ndim == 2
    
    if not target_log_pdf is None:
        Xs = np.linspace(-30, 30, 75)
        Ys = np.linspace(-10, 20, len(Xs))
        D, G = pdf_grid(Xs, Ys, target_log_pdf)
    
    plt.figure(figsize=(10, 12))
    plt.subplot(411)
    visualise_array(Xs, Ys, D)
    plt.plot(Qs[:, 0], Qs[:, 1])
    plt.plot(Qs[0, 0], Qs[0, 1], 'r*', markersize=15)
    plt.title("Log-pdf surrogate")
    
    plt.subplot(412)
    visualise_array(Xs, Ys, G)
    plt.plot(Qs[:, 0], Qs[:, 1])
    plt.plot(Qs[0, 0], Qs[0, 1], 'r*', markersize=15)
    plt.title("Gradient norm surrogate")
    
    plt.subplot(413)
    plt.title("Acceptance probability")
    plt.xlabel("Leap frog iteration")
    plt.plot(acc_probs)
    plt.plot([0, len(acc_probs)], [np.mean(acc_probs) for _ in range(2)], 'r--')
    plt.xlim([0, len(acc_probs)])
    
    plt.subplot(414)
    plt.title("Target log-pdf")
    plt.xlabel("Leap frog iteration")
    plt.plot(log_pdf_q)
    plt.xlim([0, len(log_pdf_q)])
    
if __name__ == '__main__':
    D = 2
    target = Banana()
    N = 1000
    X = sample_banana(N, D)
    
    for surrogate in [
                        KernelExpFiniteGaussian(gamma=.5, lmbda=0.001, m=1000, D=D),
                        KernelExpLiteGaussian(sigma=25., lmbda=0.001, D=D)
                      ]:
        
        surrogate.fit(X)
        
        momentum = IsotropicZeroMeanGaussian(D=D, sigma=.1)
        
        # fixed HMC parmaeters
        num_steps = 1000
        step_size = .01
        
        kmc = KMCStatic(surrogate, target, momentum, num_steps, num_steps, step_size, step_size)
        
        current = sample_banana(N=1, D=D)[0]
        current = np.array([0, -3])
        current_log_pdf = target.log_pdf(current)
        Qs, acc_probs, log_pdf_q = kmc._proposal_trajectory(current, current_log_pdf)
        
        visualise_trajectory(Qs, acc_probs, log_pdf_q, surrogate)
        plt.suptitle("%s" % surrogate.__class__.__name__)
    plt.show()

@@ -281,32 +281,55 @@
 ZeroMeanGaussian
+, %5C%0A    sample_gaussian
 %0Afrom kernel_hmc
@@ -451,16 +451,32 @@
 s np%0A%0A%0A#
+ banana gradient
  depends
@@ -560,17 +560,16 @@
 s.banana
-s
  import 
@@ -594,109 +594,84 @@
 ana%0A
-except ImportError:%0A    print(%22Skipping example, depends on theano which is unavailable%22)%0A    exit(0)
+    banana_available = True%0Aexcept ImportError:%0A    banana_available = False
 %0A%0ALo
@@ -1788,24 +1788,67 @@
 :%0A    D = 2%0A
+    N = 1000%0A%0A    if banana_available:%0A    
     target =
@@ -1861,25 +1861,16 @@
 a()%0A    
-N = 1000%0A
     X = 
@@ -1885,24 +1885,110 @@
 anana(N, D)%0A
+    else:%0A        target = IsotropicZeroMeanGaussian()%0A        X = sample_gaussian(N)%0A
     %0A    for
@@ -2141,17 +2141,17 @@
 (sigma=2
-5
+0
 ., lmbda
@@ -2499,53 +2499,8 @@
    %0A
-        current = sample_banana(N=1, D=D)%5B0%5D%0A
     
@@ -2800,16 +2800,25 @@
 name__)%0A
+        %0A
     plt.

change target to gaussian if banana is not available