# -*- coding: utf-8 -*-
"""
Created on Fri Mar  3 18:15:07 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 RV_function;reload(RV_function);RVfunc=RV_function

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

import sys
f=open("Test_ES_4spots.txt","w")
sys.stdout = f

pl.close('all')

##### spots data pre-processing #####
rdb_data=Table.Table.read('4spots.rdb',format='ascii')
RV_spot=rdb_data['RV_tot'][1:101]
RV_spot=np.array(RV_spot)
RV_spot=RV_spot.astype('Float64')
RV_SPOT=np.concatenate((RV_spot,RV_spot,RV_spot,RV_spot),axis=0)

spots_yy=[]
for i in np.arange(4,401,4):
    a=(RV_SPOT[i-4]+RV_SPOT[i-3]+RV_SPOT[i-2]+RV_SPOT[i-1])*1000/4.
    spots_yy.append(a)

spots_data=[]    
for j in np.arange(1,100,3.3):
       spots_data.append(spots_yy[int(round(j))])
    
##### data and plot #####
# Period(P) ~ 20 e 50 days
# Observations(space) ~ every 4 days
# Error(yerr) ~ 0.20 a 0.50 m
# K=17.353 => planet with 1/4 mass of Jupiter

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([x1+x2 for x1,x2 in zip(y0,spots_data)])
total=np.array([x1+x2 for x1,x2 in zip(y,yerr)])

Xfinal=t
Yfinal=total

##### Lets try GP to fit #####
#kl.ExpSineSquared(theta,l,P) + kl.WhiteNoise(theta)

def sub_tests(trials=20,variation=-0.1):
    theta=17.0;l=1.0;P=25.0
    for i in range(1,trials+1):
        print 'EXAMPLE %i' %i
        kernel1=kl.ExpSquared(theta, l)
        print 'initial kernel ->', kernel1
        likelihood1=lk.likelihood(kernel1,Xfinal,Xfinal,Yfinal,yerr)
        print 'initial likelihood ->', likelihood1
    
        Xcalc=np.linspace(0,100,1000)  
#        [mu,std]=lk.compute_kernel(kernel1,Xfinal,Xcalc,Yfinal,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(Xfinal, Yfinal, yerr=yerr, fmt=".k", capsize=0)
#        pl.xlabel("$x$")
#        pl.ylabel("$y$")
        
        optimization1=opt.committed_optimization(kernel1,Xfinal,Xfinal,Yfinal,yerr,max_opt=10)
        print 'final kernel ->',optimization1[1]
        print 'final likelihood ->', optimization1[0]
        
#        [mu,std]=lk.compute_kernel(optimization1[1],Xfinal,Xcalc,Yfinal,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(Xfinal, Yfinal, yerr=yerr, fmt=".k", capsize=0)
#        pl.xlabel("$x$")
#        pl.ylabel("$y$")
    
        theta=theta+variation;l=l;P=P+variation


def subNoise_tests(trials=20,variation=-0.1):
    theta=17.0;l=1.0;P=24.0;noise=0.5
    for i in range(1,trials+1):
        print 'EXAMPLE %i' %i
        kernel1=kl.ExpSquared(theta, l)+kl.WhiteNoise(noise)
        print 'initial kernel ->', kernel1
        likelihood1=lk.likelihood(kernel1,Xfinal,Xfinal,Yfinal,yerr)
        print 'initial likelihood ->', likelihood1
    
        Xcalc=np.linspace(0,100,1000)  
#        [mu,std]=lk.compute_kernel(kernel1,Xfinal,Xcalc,Yfinal,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(Xfinal, Yfinal, yerr=yerr, fmt=".k", capsize=0)
#        pl.xlabel("$x$")
#        pl.ylabel("$y$")
        
        optimization1=opt.committed_optimization(kernel1,Xfinal,Xfinal,Yfinal,yerr,max_opt=10)
        print 'final kernel ->',optimization1[1]
        print 'final likelihood ->', optimization1[0]
        
#        [mu,std]=lk.compute_kernel(optimization1[1],Xfinal,Xcalc,Yfinal,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(Xfinal, Yfinal, yerr=yerr, fmt=".k", capsize=0)
#        pl.xlabel("$x$")
#        pl.ylabel("$y$")
    
        theta=theta+variation;l=l;P=P+variation;noise=noise+(variation/2.)

def add_tests(trials=20,variation=0.1):
    theta=17.0;l=1.0;P=25.0
    for i in range(1,trials+1):
        print 'EXAMPLE %i' %i
        kernel1=kl.ExpSquared(theta, l)
        print 'initial kernel ->', kernel1
        likelihood1=lk.likelihood(kernel1,Xfinal,Xfinal,Yfinal,yerr)
        print 'initial likelihood ->', likelihood1
    
        Xcalc=np.linspace(0,100,1000)  
#        [mu,std]=lk.compute_kernel(kernel1,Xfinal,Xcalc,Yfinal,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(Xfinal, Yfinal, yerr=yerr, fmt=".k", capsize=0)
#        pl.xlabel("$x$")
#        pl.ylabel("$y$")
        
        optimization1=opt.committed_optimization(kernel1,Xfinal,Xfinal,Yfinal,yerr,max_opt=10)
        print 'final kernel ->',optimization1[1]
        print 'final likelihood ->', optimization1[0]
        
#        [mu,std]=lk.compute_kernel(optimization1[1],Xfinal,Xcalc,Yfinal,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(Xfinal, Yfinal, yerr=yerr, fmt=".k", capsize=0)
#        pl.xlabel("$x$")
#        pl.ylabel("$y$")
    
        theta=theta+variation;l=l;P=P+variation

def addNoise_tests(trials=20,variation=0.1):
    theta=17.0;l=1.0;P=24.0;noise=0.1
    for i in range(1,trials+1):
        print 'EXAMPLE %i' %i
        kernel1=kl.ExpSquared(theta, l)+kl.WhiteNoise(noise)
        print 'initial kernel ->', kernel1
        likelihood1=lk.likelihood(kernel1,Xfinal,Xfinal,Yfinal,yerr)
        print 'initial likelihood ->', likelihood1
    
        Xcalc=np.linspace(0,100,1000)  
#        [mu,std]=lk.compute_kernel(kernel1,Xfinal,Xcalc,Yfinal,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(Xfinal, Yfinal, yerr=yerr, fmt=".k", capsize=0)
#        pl.xlabel("$x$")
#        pl.ylabel("$y$")
        
        optimization1=opt.committed_optimization(kernel1,Xfinal,Xfinal,Yfinal,yerr,max_opt=10)
        print 'final kernel ->',optimization1[1]
        print 'final likelihood ->', optimization1[0]
        
#        [mu,std]=lk.compute_kernel(optimization1[1],Xfinal,Xcalc,Yfinal,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(Xfinal, Yfinal, yerr=yerr, fmt=".k", capsize=0)
#        pl.xlabel("$x$")
#        pl.ylabel("$y$")
    
        theta=theta;l=l;noise=noise+(variation/2.)

sub_tests()