import Gedi as gedi
import george
import numpy as np;
import matplotlib.pylab as pl; pl.close('all')
from time import time,sleep
import scipy.optimize as op
import sys
#####  INITIAL DATA ###########################################################

nrep = 1
pontos=[]       
temposQP=[]
temposmulti=[]
georgeQP=[]
sleeptime=10
lista=[10,20,50,100,200,500]
#for i in np.arange(100,650,200):
#for i in np.arange(100,1400,350):

### Functions george
# Define the objective function (negative log-likelihood in this case).
def nll(p):
    # Update the kernel parameters and compute the likelihood.
    gp.kernel[:] = p
    ll = gp.lnlikelihood(y, quiet=True)

    # The scipy optimizer doesn't play well with infinities.
    return -ll if np.isfinite(ll) else 1e25

# And the gradient of the objective function.
def grad_nll(p):
    # Update the kernel parameters and compute the likelihood.
    gp.kernel[:] = p
    return -gp.grad_lnlikelihood(y, quiet=True)

### Functions gedi                
def nll_gedi(p):
    global kernel
    # Update the kernel parameters and compute the likelihood.
    kernel= gedi.kernel_optimization.new_kernel(kernel,np.exp(p))
    ll = gedi.kernel_likelihood.likelihood(kernel,x,y,yerr)
    # The scipy optimizer doesn't play well with infinities.
    return -ll if np.isfinite(ll) else 1e25

# And the gradient of the objective function.
def grad_nll_gedi(p):
    global kernel
    # Update the kernel parameters and compute the likelihood.
    kernel= gedi.kernel_optimization.new_kernel(kernel,np.exp(p))
    return -np.array(gedi.kernel_likelihood.gradient_likelihood(kernel,x,y,yerr))
###############################################################################


### Things to run   
for i0, i in enumerate(lista):
    f=open("{0}.txt".format(i),"w")
    sys.stdout = f
    
    print i
    pontos.append(i)
    print 'pontos', pontos
    x = 10 * np.sort(np.random.rand(2*i))
    yerr = 0.2 * np.ones_like(x)
    y = np.sin(x) + yerr * np.random.randn(len(x))
       
    av = []
    for _ in range(nrep):
        start= time()
        kernel= gedi.kernel.QuasiPeriodic(15.0,2.0,1.0,10.0)
        print 'Initial gedi kernel =',kernel
        print 'Initial gedi likelihood =',gedi.kernel_likelihood.likelihood(kernel,x,y,yerr)
        
        # Run the optimization routine.
        p0_gedi = np.log(kernel.pars)
        results_gedi = op.minimize(nll_gedi, p0_gedi, jac=grad_nll_gedi)
        
        kernel= gedi.kernel_optimization.new_kernel(kernel,np.exp(results_gedi.x))
        print 'Final gedi kernel =',kernel
        print 'Final gedi likelihood =',gedi.kernel_likelihood.likelihood(kernel,x,y,yerr)
        print 
        tempo1= time() - start

        av.append(tempo1)

    temposQP.append(sum(av) / float(nrep))
    print 'temposQP', temposQP
    sleep(sleeptime*i0)
###############################################################################
    av = []
    for _ in range(nrep):
        start= time()
        kernel= gedi.kernel.ExpSineSquared(15.0, 2.0, 10.0)* \
                    gedi.kernel.ExpSquared(1.0,1.0)
        print 'Initial gedi kernel =',kernel
        print 'Initial gedi likelihood =',gedi.kernel_likelihood.likelihood(kernel,x,y,yerr)
        
        # Run the optimization routine.
        p0_gedi = np.log(kernel.pars)
        results_gedi = op.minimize(nll_gedi, p0_gedi, jac=grad_nll_gedi)
        
        kernel= gedi.kernel_optimization.new_kernel(kernel,np.exp(results_gedi.x))
        print 'Final gedi kernel =',kernel
        print 'Final gedi likelihood =',gedi.kernel_likelihood.likelihood(kernel,x,y,yerr)
        print 
        tempo1= time() - start

        av.append(tempo1)

    temposmulti.append(sum(av) / float(nrep))
    print 'temposmult', temposmulti
      
    sleep(sleeptime*i0)    
    av = []
    for _ in range(nrep):

        start = time() # Calculation using george 
        kernelg1 = 15.0**2*george.kernels.ExpSine2Kernel(2/2.0**2,10.0)* \
                    george.kernels.ExpSquaredKernel(1.0**2)
        
        # You need to compute the GP once before starting the optimization.
        gp = george.GP(kernelg1, mean=np.mean(y))
        gp.compute(x,yerr)
        
        # Print the initial ln-likelihood.
        print 'Initial george kernel', kernelg1
        print 'Initial george likelihood', gp.lnlikelihood(y)
        
        # Run the optimization routine.
        p0 = gp.kernel.vector
        results = op.minimize(nll, p0, jac=grad_nll)
        
        # Update the kernel and print the final log-likelihood.
        gp.kernel[:] = results.x
        print 'Final george kernel =',gp.kernel
        print 'Final george likelihood= ', gp.lnlikelihood(y)        
        print  
        tempog1= time() - start
        av.append(tempog1)
        
    georgeQP.append(sum(av) / float(nrep))
    print 'georgeQP', georgeQP
    ###########################################################################
    sys.stdout = sys.__stdout__
    f.close()
    sleep(sleeptime*i0)

N = pontos   
pl.figure()
pl.loglog(N, temposQP, 'r-')
pl.loglog(N, temposmulti, 'b-o')
pl.loglog(N, georgeQP, 'b--')

pl.xlim(0.9*N[0], 1.1*N[-1])
pl.xlabel('Number of points')
pl.ylabel('Time')
#pl.title('Covariance matrix calculations')
pl.legend(['gedi QP', 'gedi ESS*ES','george ESS*ES'],loc='upper left')
pl.xticks(fontsize = 18);pl.yticks(fontsize=18)
pl.savefig('speedtest_6.png')
#pl.close('all')