#!/usr/local/bin/python2.7
from __future__ import division  # divide integers as if they were floating point
import sys
import pycbc
from pycbc.waveform import get_td_waveform
import pylab
import scipy
from math import *
import numpy as np
#matplotlib inline
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import h5py
import json
# LIGO-specific stuff:
import lal as lal
import lalsimulation as lalsim
import readligo as rl
import scipy.signal as sig
import pycbc.noise.reproduceable
import pycbc.psd
from pycbc.filter import resample_to_delta_t, highpass

#print lalsim.EccentricTD

H1=lalsim.DetectorPrefixToLALDetector('H1')
L1=lalsim.DetectorPrefixToLALDetector('L1')
V1=lalsim.DetectorPrefixToLALDetector('V1')

# Visualize different tidal distortions

# Set sampling rate
fs = 16384 
# get the parameters needed to generate the waveform
deltaT = 1. / fs  # requires from __future__ import division
m1_SI = 1.4 * lal.MSUN_SI
m2_SI = 1.4 * lal.MSUN_SI
s1x = s1y = s1z = 0.
s2x = s2y = s2z = 0.
dist = 40. 
distance = dist * 1.e6 * lal.PC_SI
inclination = lal.PI_4
phiRef = 0
eccentricity = 0.
fmin = 100.
fref = 30.
#ampOrder = 1
#phOrder = 3
# needed only for detector response:
z = 0.
ra = 0.
dec = 0.
psi = 0.*lal.PI_4
time = 1186302519
approx = lalsim.TaylorF2
# approx=lalsim.IMRPhenomPv2_NRTidal
WFdict = lal.CreateDict()

Lambda = 0.
lalsim.SimInspiralWaveformParamsInsertTidalLambda1(WFdict, Lambda)
lalsim.SimInspiralWaveformParamsInsertTidalLambda2(WFdict, Lambda)
hpt1, hct1 = lalsim.SimInspiralTD(m1_SI,m2_SI, s1x, s1y, s1z, s2x, s2y,s2z,
           distance, inclination, phiRef, 0., eccentricity, 0., deltaT, fmin, fref, WFdict , approx)
hd1=lalsim.SimDetectorStrainREAL8TimeSeries(hpt1, hct1, ra, dec, psi, H1)
h1 = hd1.data.data
t1 = float(hd1.epoch)+hd1.deltaT*np.arange(hd1.data.length)
t1f=float(hpt1.epoch)+hpt1.deltaT*np.arange(hpt1.data.length)
deltaF1=1/(t1f[len(t1f)-1]-t1f[0])
print(deltaF1)
print(len(t1f))

Lambda = 100.
lalsim.SimInspiralWaveformParamsInsertTidalLambda1(WFdict, Lambda)
lalsim.SimInspiralWaveformParamsInsertTidalLambda2(WFdict, Lambda)
hpt2, hct2 = lalsim.SimInspiralTD(m1_SI,m2_SI, s1x, s1y, s1z, s2x, s2y,s2z,
           distance, inclination, phiRef, 0., eccentricity, 0., deltaT, fmin, fref, WFdict , approx)
hd2=lalsim.SimDetectorStrainREAL8TimeSeries(hpt2, hct2, ra, dec, psi, H1)
h2 = hd2.data.data
t2 = float(hd2.epoch)+hd2.deltaT*np.arange(hd2.data.length)
t2f=float(hpt2.epoch)+hpt2.deltaT*np.arange(hpt2.data.length)
deltaF2=1/(t2f[len(t2f)-1]-t2f[0])


Lambda = 1000.
lalsim.SimInspiralWaveformParamsInsertTidalLambda1(WFdict, Lambda)
lalsim.SimInspiralWaveformParamsInsertTidalLambda2(WFdict, Lambda)
hpt3, hct3 = lalsim.SimInspiralTD(m1_SI,m2_SI, s1x, s1y, s1z, s2x, s2y,s2z,
           distance, inclination, phiRef, 0., eccentricity, 0., deltaT, fmin, fref, WFdict , approx)
hd3=lalsim.SimDetectorStrainREAL8TimeSeries(hpt3, hct3, ra, dec, psi, H1)
h3 = hd3.data.data
t3 = float(hd3.epoch)+hd3.deltaT*np.arange(hd3.data.length)
t3f=float(hpt3.epoch)+hpt3.deltaT*np.arange(hpt3.data.length)
deltaF3=1/(t3f[len(t3f)-1]-t3f[0])



# plot signal and compare to BBH
plt.figure(figsize=(18,8))
plt.subplot(211)
plt.plot(t1, h1, label='BBH')
plt.xlabel('Time (s)')
plt.ylabel('Signal (m)')
plt.xlim(-0.06,-0.015)
plt.legend()
plt.subplot(212)
plt.plot(t1, h1, label='BBH')
plt.plot(t2, h2, label='h2:Lambda1,2=100')
plt.plot(t3, h3, label='h3:Lambda1,2=1000')
plt.xlabel('Time (s)')
plt.ylabel('Signal (m)')
plt.xlim(-0.06,-0.015)
plt.legend()
plt.show()

# create complex signal and get phase in time domain
tcomplexh1=[0]*len(hpt1.data.data)
phaset1=[0]*len(hpt1.data.data)
for i in range(len(tcomplexh1)):
	tcomplexh1[i]=complex(hpt1.data.data[i],hct1.data.data[i])
	phaset1[i]=np.angle(tcomplexh1[i])
tcomplexh2=[0]*len(hpt2.data.data)
phaset2=[0]*len(hpt1.data.data)
for i in range(len(tcomplexh2)):
	tcomplexh2[i]=complex(hpt2.data.data[i],hct2.data.data[i])
	phaset2[i]=np.angle(tcomplexh2[i])
tcomplexh3=[0]*len(hpt3.data.data)
phaset3=[0]*len(hpt1.data.data)
for i in range(len(tcomplexh3)):
	tcomplexh3[i]=complex(hpt3.data.data[i],hct3.data.data[i])
	phaset3[i]=np.angle(tcomplexh3[i])

# plot phase difference from BBH in time domain
deltapht2=[0]*len(hpt1.data.data)
deltapht3=[0]*len(hpt1.data.data)
for i in range(len(hpt1.data.data)):
	deltapht2[i]=phaset2[i]-phaset1[i]
	deltapht3[i]=phaset3[i]-phaset1[i]
deltapht2=np.unwrap(deltapht2)
deltapht3=np.unwrap(deltapht3)
plt.figure(figsize=(18,8))
plt.plot(t2f, deltapht2, label='h2:Lambda1,2=100')
plt.plot(t3f, deltapht3, label='h3:Lambda1,2=1000')
plt.xlabel('Time (s)')
plt.ylabel('Phase')
plt.xlim(-0.5,0.02)
plt.ylim(ymax=40)
plt.legend()
plt.show()


# get fft of signal and phase in frequency domain
fcomplexh1=np.fft.fft(np.array(tcomplexh1))
phasef1=[0]*len(fcomplexh1)
for i in range(len(fcomplexh1)):
	phasef1[i]=np.angle(fcomplexh1[i])
fcomplexh2=np.fft.fft(np.array(tcomplexh2))
phasef2=[0]*len(fcomplexh2)
for i in range(len(fcomplexh2)):
	phasef2[i]=np.angle(fcomplexh2[i])
fcomplexh3=np.fft.fft(np.array(tcomplexh3))
phasef3=[0]*len(fcomplexh3)
for i in range(len(fcomplexh3)):
	phasef3[i]=np.angle(fcomplexh3[i])

# plot phase difference from BBH in frequency domain
deltaphf2=[0]*len(hpt1.data.data)
deltaphf3=[0]*len(hpt1.data.data)
for i in range(len(hpt1.data.data)):
	deltaphf2[i]=phasef2[i]-phasef1[i]
	deltaphf3[i]=phasef3[i]-phasef1[i]
deltaphf2=np.unwrap(deltaphf2)
deltaphf3=np.unwrap(deltaphf3)
plt.figure(figsize=(18,8))
plt.plot(float(fmin)+deltaF2*np.arange(hpt2.data.length), deltaphf2, label='h2:Lambda1,2=100')
plt.plot(float(fmin)+deltaF3*np.arange(hpt3.data.length), deltaphf3, label='h3:Lambda1,2=1000')
plt.xlim(0,3000)
plt.ylim(0,45)
plt.xlabel('Frequency (Hz)')
plt.ylabel('Phase')
plt.legend()
plt.show()

# # PSD
# plt.figure(figsize=(18,8))
# PSD1,freqs1=mlab.psd(h1, Fs=fs, NFFT=fs)
# PSD2,freqs2=mlab.psd(h2, Fs=fs, NFFT=fs)
# PSD3,freqs3=mlab.psd(h3, Fs=fs, NFFT=fs)
# plt.subplot(211)
# plt.loglog(freqs1, PSD1, label='BBH')
# # plt.axis([10, 2000, 1e-46, 1e-36])
# plt.grid('on')
# plt.ylabel('PSD')
# plt.xlabel('Freq (Hz)')
# plt.legend()
# plt.subplot(212)
# plt.loglog(freqs1, PSD1, label='BBH')
# plt.loglog(freqs2, PSD2, label='h2:Lambda1,2=100')
# plt.loglog(freqs3, PSD3, label='h3:Lambda1,2=1000')
# # plt.axis([10, 2000, 1e-46, 1e-36])
# plt.grid('on')
# plt.ylabel('PSD')
# plt.xlabel('Freq (Hz)')
# plt.legend()

# plt.show()

# # plot spectrogram
# plt.figure(figsize=(18,8))
# NFFT = 16384
# window = np.blackman(NFFT)
# plt.subplot(131)
# spec_power1, freqst1, bins1, im1 = plt.specgram(h1, NFFT=NFFT, Fs=fs, window=window)
# plt.ylim(ymax=1000)
# plt.xlabel('Time (s)')
# plt.ylabel('Freq (Hz)')
# plt.subplot(132)
# spec_power2, freqst2, bins2, im2 = plt.specgram(h2, NFFT=NFFT, Fs=fs, window=window)
# plt.ylim(ymax=1000)
# plt.xlabel('Time (s)')
# plt.ylabel('Freq (Hz)')
# plt.subplot(133)
# spec_power3, freqst3, bins3, im3 = plt.specgram(h3, NFFT=NFFT, Fs=fs, window=window)
# plt.ylim(ymax=1000)
# plt.xlabel('Time (s)')
# plt.ylabel('Freq (Hz)')
# plt.show()