# This is a tutorial for plotting waveform using pycbc.
# Not for light hearted, text based code. 
# pycbc is needed check with pycbc instruction to for installation help
# Also look at pycbc tutorial for plotting waveform
kb=1024
sps=4*kb  # Sampling frequency 4K should be fine!
dt=1/sps  # by defination
tmax=16   # sec length of time series if needed
df=1/tmax # by defination
flower=10 # Hz for the Advance LIGO/VIRGO/etc
m1=35     # Solar masses, feel free to change it
m2=30     # Solar masses, feel free to change it
distance=500 # Mpc, feel free to change it 
### For simplisity, we are not using the other source parameters
### You may add them. look at pycbc wavform module documentation
### for full details 
fd_model='IMRPhenomC' # This is pycbc name for fdwaveform
td_model='IMRPhenomC' # This is pycbc name for tdwaveform
## Above two waveform need not be same if want to compare!




###
## basic libaries 
import numpy as np
import matplotlib.pyplot as plt
from pycbc.waveform import get_fd_waveform
from pycbc.waveform import get_td_waveform


# Time to get fD wave form from pycbc
hp_fd, hc_fd= get_fd_waveform(approximant=fd_model,
	mass1=m1, mass2=m2, distance=distance,
	f_lower=flower, delta_f=df)
# Time to get td waveform from pycbc 
hp_td, hc_td= get_td_waveform(approximant=td_model,
	mass1=m1, mass2=m2, distance=distance,
	f_lower=flower, delta_t=dt)
# Let's plot them 
f1=plt.figure(figsize=[10,8])

plt.title('Plot of fd waveform '+fd_model, size=18)
plt.plot(hp_fd.sample_frequencies, hp_fd.real(),color='C0', label=r'$real(h_p)$')
plt.plot(hp_fd.sample_frequencies, hp_fd.imag(),color='C1', label=r'$imag(h_p)$')
plt.plot(hp_fd.sample_frequencies, np.sqrt(hp_fd.squared_norm()),color='k', label=r'$abs(h_p)$' )
plt.plot(hp_fd.sample_frequencies, hc_fd.real(),ls='--', color='C2', label=r'$real(h_x)$')
plt.plot(hp_fd.sample_frequencies, hc_fd.imag(),ls='--',color='C3', label=r'$imag(h_x)$')
plt.plot(hp_fd.sample_frequencies, np.sqrt(hc_fd.squared_norm()),ls='--', color='r', label=r'$abs(h_x)$' )
plt.grid(alpha=0.3)
plt.xlabel('Frequency[Hz]')
plt.legend()
plt.show(block=False)
### Too many overlaping plot plot's. You can cleanup the plots.
### You will know more about initial phase and relation between plus and cross polariation
f2=plt.figure(figsize=[10,8])
plt.title('Plot of td waveform '+td_model, size=18)
plt.plot(hp_td.sample_times, hp_td,color='C0', label=r'$h_p$')
plt.plot(hp_td.sample_times, hc_td,ls='--', color='C2', label=r'$h_x$')
plt.xlabel('time[sec]')
#plt.xlim([-8,0.1])
plt.grid(alpha=0.3)
plt.legend()
plt.show(block=False)
## 
## 
## Play around axis limits and learn more about signal tapering to avoid the artifacts due to Fourier transform
## Now we convert fd waveform to td wavefrom using build in function in pycbc.
## We also try converting td waveform to fd wavefrom.
## We only do this for hp at this stage
tdp=hp_fd.to_timeseries()
f3=plt.figure(figsize=[10,8])
plt.title('Plot of td waveform converted from fd '+td_model, size=18)
plt.plot(tdp.sample_times, tdp, color='C0', label='hp')
plt.xlabel('time[sec]')
plt.grid(alpha=0.3)
plt.legend()
plt.show(block=False)
### 
## Not as good as td waveform from original pycbc! Why?
fdp=hp_td.to_frequencyseries()
f4=plt.figure(figsize=[10,8])
plt.title('Plot of fd waveform converted from td '+fd_model, size=18)
plt.plot(fdp.sample_frequencies, fdp.real(), color='C0', label='hp.real')
plt.xlabel('Frequecny[Hz]')
plt.grid(alpha=0.3)
plt.legend()
plt.show()
## Not too bad ? Why. Check with TaylorF2 waveform## 
##