# -*- coding: utf-8 -*-
"""
Created on Wed Mar  1 14:55:48 2017

@author: camacho
"""

#import Kernel;reload(Kernel);kl=Kernel
#import Kernel_likelihood;reload(Kernel_likelihood);lk=Kernel_likelihood
#import Kernel_optimization;reload(Kernel_optimization);opt=Kernel_optimization
#import Kernel_MCMCoptimization;reload(Kernel_MCMCoptimization); mcmc=Kernel_MCMCoptimization
import Gedi as gedi

import george
import george.kernels as ge

import RV_function;reload(RV_function);RVfunc=RV_function

import numpy as np;np.random.seed(123)
import matplotlib.pylab as pl
import astropy.table as Table

pl.close()

##### data and plot #####
# Periodo(P) ~ 20 e 50 dias
# Mediçoes(space) ~ a cada 4 dias
# Erro(yerr) ~ 0.20 a 0.50 m
 
test1=RVfunc.RV_circular(P=25,K=17.353,T=0,gamma=0,time=100,space=30)
t=np.linspace(0,100,30) #np.linspace(0,time,space)
y0=np.array(test1[1])
yerr=np.array([np.random.uniform(0.2,0.5) for x in y0])
y=np.array([x + y for x,y in zip(y0,yerr)])

#total=[x1+x2 for x1,x2 in  zip(y,RV_SPOT)]
#pl.errorbar(t,y0,yerr,fmt='*')
#pl.xlabel('time (days)')
#pl.ylabel('RV (m/s)')


#print '########## Calculations from gedi ##########'
#kernel1=gedi.kernel.ExpSineGeorge(2.0/1.7**2,5.0) \
#        *gedi.kernel.ExpSineGeorge(2.0/1.0**2,10) \
#        *gedi.kernel.ExpSineGeorge(2.0/3.0**2,11) \
#        *gedi.kernel.ExpSineGeorge(2.0/1.1**2,9) \
#        +gedi.kernel.ExpSineGeorge(2.0/1.0**2,10) \
#        *gedi.kernel.ExpSineGeorge(2.0/3.0**2,11) \
#        *gedi.kernel.ExpSineGeorge(2.0/1.1**2,9)
#print 'kernel ->',kernel1;print''
#
#like1=gedi.kernel_likelihood.likelihood(kernel1,t,y,yerr)
#print 'likelihood ->',like1
#
#grad1=gedi.kernel_likelihood.gradient_likelihood(kernel1,t,y,yerr)
#print 'gradients ->',grad1
#
#
#print '########## Calculations from george ##########'
#kernel=ge.ExpSine2Kernel(2.0/1.7**2,5.0) \
#        *ge.ExpSine2Kernel(2.0/1.0**2,10.0) \
#        *ge.ExpSine2Kernel(2.0/3.0**2,11.0) \
#        *ge.ExpSine2Kernel(2.0/1.1**2,9) \
#        +ge.ExpSine2Kernel(2.0/1.0**2,10.0) \
#        *ge.ExpSine2Kernel(2.0/3.0**2,11.0) \
#        *ge.ExpSine2Kernel(2.0/1.1**2,9)
#gp = george.GP(kernel)
#gp.compute(t,yerr)
#
#print 'initial kernel_george ->', kernel
#print 'initial likelihood_george ->', gp.lnlikelihood(y)
#print 'initial gradient_george ->',gp.grad_lnlikelihood(y)
#### OPTIMIZE HYPERPARAMETERS
#import scipy.optimize as op
## 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)
#
## You need to compute the GP once before starting the optimization.
#gp.compute(x,yerr)
#
#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 likelihood_george ->', gp.lnlikelihood(y)
#print 'final kernel_george ->', kernel #kernel final
#print '#####################################'
#print ''


#kernel1=gedi.kernel.Matern_52(10,10)+gedi.kernel.WhiteNoise(1.0)
#kernel1=gedi.kernel.ExpSquared(15.0,10.0)+gedi.kernel.WhiteNoise(1.0)
#parameters1=[[10.0,20.0],[1.0,10.0],[10.0,20.0],[1.0,10.0],[0.1,5.0]]

#print 'Let it run!'
#start=0;stop=10
#while start<stop:
#    kernel1=gedi.kernel.ExpSineSquared(10,1,10) + gedi.kernel.WhiteNoise(10.0)
#    parameters1=[[15.0,20.0],[1.0,4.0],[20.0,30.0],[0.1,1]]
#    runs=5000
#    burn_in=1000
#    gedi.kernel_mcmc.MCMC(kernel1,t,y,yerr,parameters1,runs)
#    start=+1    
    
print '##### EXEMPLO MCMC -> ESS+WN, 1st run'
kernel1=gedi.kernel.ExpSineSquared(10,1,10) + gedi.kernel.WhiteNoise(10.0)
parameters1=[[16.0,18.0],[1.0,4.0],[20.0,30.0],[0.1,1]]

runs=50000
burns=round(runs/3)
trial=gedi.kernel_mcmc.MCMC(kernel1,t,y,yerr,parameters1,runs,burns)
print trial[0], trial[1]

xt=np.linspace(0,runs,runs-burns)
pl.figure()
pl.title('log marginal likelihood')
pl.plot(xt,trial[2],'k-')

# Four axes, returned as a 2-d array
f, axarr = pl.subplots(2, 2)
axarr[0, 0].plot(xt, trial[3][0])
axarr[0, 0].set_title('amplitude')
axarr[0, 1].plot(xt, trial[3][1])
axarr[0, 1].set_title('lenght scale')
axarr[1, 0].plot(xt, trial[3][2])
axarr[1, 0].set_title('period')
axarr[1, 1].plot(xt, trial[3][3])
axarr[1, 1].set_title('white noise')
pl.setp([a.get_xticklabels() for a in axarr[0, :]], visible=False)

##### computation  of the graphic
Xcalc=np.linspace(0,100,500) 
[mu,std]=gedi.kernel_likelihood.compute_kernel(trial[0],t,Xcalc,y,yerr)
pl.figure() #Graphics
pl.fill_between(Xcalc, mu+std, mu-std, color="k", alpha=0.1)
pl.plot(Xcalc, mu+std, color="k", alpha=1, lw=0.25)
pl.plot(Xcalc, mu-std, color="k", alpha=1, lw=0.25)
pl.plot(Xcalc, mu, color="k", alpha=1, lw=0.5)
pl.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0)
pl.title('Pos-optimization')
pl.xlabel("$time$")
pl.ylabel("$y$")


#print '##### EXEMPLO MCMC -> ESS+WN, 2nd run'
#kernel2=gedi.kernel.ExpSineSquared(10,1,10) + gedi.kernel.WhiteNoise(10.0)
#parameters2=[[15.0,17.0],[2.25,4.0],[24.9,25.0],[0,0.2]]
#
#runs=100000
#trial=gedi.kernel_mcmc.MCMC(kernel2,t,y,yerr,parameters2,runs)
#xt=np.linspace(0,runs,runs)
#
#pl.figure()
#pl.title('log marginal likelihood')
#pl.plot(xt,trial[2],'k-')
## Four axes, returned as a 2-d array
#f, axarr = pl.subplots(2, 2)
#axarr[0, 0].plot(xt, trial[3][0])
#axarr[0, 0].set_title('amplitude')
#axarr[0, 1].plot(xt, trial[3][1])
#axarr[0, 1].set_title('lenght scale')
#axarr[1, 0].plot(xt, trial[3][2])
#axarr[1, 0].set_title('period')
#axarr[1, 1].plot(xt, trial[3][3])
#axarr[1, 1].set_title('white noise')
#pl.setp([a.get_xticklabels() for a in axarr[0, :]], visible=False)
##pl.setp([a.get_yticklabels() for a in axarr[:, 1]], visible=False)
#
###### computation  of the graphic
#Xcalc=np.linspace(0,100,500) 
#[mu,std]=gedi.kernel_likelihood.compute_kernel(trial[0],t,Xcalc,y,yerr)
#pl.figure() #Graphics
#pl.fill_between(Xcalc, mu+std, mu-std, color="k", alpha=0.1)
#pl.plot(Xcalc, mu+std, color="k", alpha=1, lw=0.25)
#pl.plot(Xcalc, mu-std, color="k", alpha=1, lw=0.25)
#pl.plot(Xcalc, mu, color="k", alpha=1, lw=0.5)
#pl.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0)
#pl.title('Pos-optimization')
#pl.xlabel("$time$")
#pl.ylabel("$y$")
#
###############################################################################
#print '##### EXEMPLO MCMC -> RQ+WN'
#kernel1=gedi.kernel.RatQuadratic(10,1,1) + gedi.kernel.WhiteNoise(10)
#parameters=[[15.0,20.0],[1.0,4.0],[1.0,4.0],[0.1,1]]
#
#runs=100000
#trial=gedi.kernel_mcmc.MCMC(kernel1,t,y,yerr,parameters,runs)
#xt=np.linspace(0,runs,runs)
#
#pl.figure()
#pl.title('log marginal likelihood')
#pl.plot(xt,trial[2],'k-')
#
## Four axes, returned as a 2-d array
#f, axarr = pl.subplots(2, 2)
#axarr[0, 0].plot(xt, trial[3][0])
#axarr[0, 0].set_title('amplitude')
#axarr[0, 1].plot(xt, trial[3][1])
#axarr[0, 1].set_title('alpha')
#axarr[1, 0].plot(xt, trial[3][2])
#axarr[1, 0].set_title('lenght scale')
#axarr[1, 1].plot(xt, trial[3][3])
#axarr[1, 1].set_title('white noise')
#pl.setp([a.get_xticklabels() for a in axarr[0, :]], visible=False)
#
###### computation  of the graphic
#Xcalc=np.linspace(0,100,500) 
#[mu,std]=gedi.kernel_likelihood.compute_kernel(trial[0],t,Xcalc,y,yerr)
#pl.figure() #Graphics
#pl.fill_between(Xcalc, mu+std, mu-std, color="k", alpha=0.1)
#pl.plot(Xcalc, mu+std, color="k", alpha=1, lw=0.25)
#pl.plot(Xcalc, mu-std, color="k", alpha=1, lw=0.25)
#pl.plot(Xcalc, mu, color="k", alpha=1, lw=0.5)
#pl.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0)
#pl.title('Pos-optimization')
#pl.xlabel("$time$")
#pl.ylabel("$y$")