# -*- coding: utf-8 -*-
"""
Created on Fri Feb  3 11:36:58 2017

@author: camacho
"""
import  numpy as np
import matplotlib.pyplot as pl
pl.close("all")

##### RV FUNCTION 1 - circular orbit
def RV_circular(P=365,K=0.1,T=0,gamma=0,time=100,space=20):
    #parameters
    #P = period in days
    #K =  semi-amplitude of the signal
    #T = velocity at zero phase
    #gamma = average velocity of the star
    #time  = time of the simulation
    #space => I want an observation every time/space days
    
    t=np.linspace(0,time,space)
    RV=[K*np.sin(2*np.pi*x/P - T) + gamma for x in t]
    RV=[x for x in RV] #m/s 
    return [t,RV]

##### RV FUNCTION 2 - keplerian orbit
def RV_kepler(P=365,e=0,K=0.1,T=0,gamma=0,w=np.pi,time=100,space=20):
    #parameters
    #P = period in days
    #e = eccentricity
    #K = RV amplitude
    #gamma = constant system RV
    #T = zero phase
    #w = longitude of the periastron
    #time = time of the simulation
    #space => I want an observation every time/space days
    
    t=np.linspace(0,time,space)

    #mean anomaly
    Mean_anom=[2*np.pi*(x1-T)/P  for x1 in t]
       
    #eccentric anomaly -> E0=M + e*sin(M) + 0.5*(e**2)*sin(2*M)
    E0=[x + e*np.sin(x)  + 0.5*(e**2)*np.sin(2*x) for x in Mean_anom]
    #mean anomaly -> M0=E0 - e*sin(E0)
    M0=[x - e*np.sin(x) for x in E0]

    i=0
    while i<100:
        #[x + y for x, y in zip(first, second)]
        calc_aux=[x2-y for x2,y in zip(Mean_anom,M0)]    
        E1=[x3 + y/(1-e*np.cos(x3)) for x3,y in zip(E0,calc_aux)]
        M1=[x4 - e*np.sin(x4) for x4 in E0]   
        i+=1
        E0=E1
        M0=M1

    nu=[2*np.arctan(np.sqrt((1+e)/(1-e))*np.tan(x5/2)) for x5 in E0]
    RV=[ gamma + K*(e*np.cos(w)+np.cos(w+x6)) for x6 in nu]
    RV=[x for x in RV] #m/s 
    return t,RV

#Examples    
#a=RV_circular()
#pl.figure('RV_circular with P=365')
#pl.plot(a[0],a[1],':',)
#pl.title('planet of 365 days orbit')
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')

#b=RV_circular(P=100)
#pl.figure('RV_circular with P=100')
#pl.title('planet of 100 days orbit')
#pl.plot(b[0],b[1],':',)
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')

#c=RV_kepler(P=100,e=0,w=np.pi,time=100)
#pl.figure()
#pl.plot(c[0],c[1],':',)
#pl.title('P=100, e=0, w=pi, time=100')
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')

#d=RV_kepler(P=100,e=0, w=np.pi,time=500)
#pl.figure()
#pl.title('P=100, e=0, w=pi, time=25')
#pl.plot(d[0],d[1],'-',)
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')