from __future__ import division  # divide integers as if they were floating point
import os, sys, glob, pickle
import optparse
import pandas
import pycbc
from pycbc.waveform import get_td_waveform
import matplotlib
matplotlib.use('Agg')
import pylab
import scipy
import numpy as np
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
from math import *
#import pycbc.noise
import pycbc.psd

plotDir = '../plots/'
if not os.path.isdir(plotDir):
    os.makedirs(plotDir)


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.
fmax=2000.
fref = 30.
deltaF=(fmax-fmin)/50000
#ampOrder = 1
#phOrder = 3
# needed only for detector response:
z = 0.
ra = 0.
dec = 0.
psi = 0.*lal.PI_4
time = 1186302519
approx=lalsim.IMRPhenomPv2_NRTidal

# create BBH signal
lambda1 = 0.
lambda2=0.

hpf1, hcf1 = lalsim.SimInspiralChooseFDWaveform(phiRef,deltaF,m1_SI,m2_SI,s1x,s1y,s1z,
			s2x,s2y,s2z,fmin,fmax,fref,distance,inclination,lambda1,lambda2,None,None,0,0,approx)

# the waveform is generated in the frequency domain as a COMPLEX frequency series
realhpf1=hpf1.data.data
realhcf1=hcf1.data.data
print(len(realhpf1))
# get real signal
realhpf1=realhpf1.real
realhcf1=realhcf1.real
#realhpf1=np.pad(realhpf1,(int(ceil(fmin/deltaF)),0),'constant',constant_values=(0,0))
#realhcf1=np.pad(realhcf1,(int(ceil(fmin/deltaF)),0),'constant',constant_values=(0,0))
f1=fmin+deltaF*np.arange(len(realhpf1))
#realhpt1=np.fft.irfft(np.pad(np.array(realhpf1), (int(ceil(100/deltaF)),int(ceil(len(realhpf1)+100/deltaF))), 'constant', constant_values=(0,0)))
#realhct1=np.fft.irfft(np.pad(np.array(realhcf1), (int(ceil(100/deltaF)),int(ceil(len(realhcf1)+100/deltaF))), 'constant', constant_values=(0,0)))
#realhpf1_flip= np.flip(realhpf1,0)
#realhcf1_flip= np.flip(realhcf1,0)
#realhpf1_new= np.append(realhpf1_flip,realhpf1)
#realhcf1_new= np.append(realhcf1_flip,realhcf1)
#f1=-(deltaF*len(realhpf1_new)/2)+deltaF*np.arange(len(realhpf1_new))
#realhpt1=np.fft.ifft(realhpf1_new)
#realhct1=np.fft.ifft(realhcf1_new)

# get signal in time domain
realhpt1=np.fft.irfft(np.pad(np.array(realhpf1), (int(ceil(100/deltaF)),0), 'constant', constant_values=(0, 0)))
realhct1=np.fft.irfft(np.pad(np.array(realhcf1), (int(ceil(100/deltaF)),0), 'constant', constant_values=(0, 0)))
#tp1,realhpt1=scipy.signal.istft([np.pad(np.array(realhpf1), (int(ceil(100/deltaF)),int(ceil(len(realhcf1)+100/deltaF))), 'constant', constant_values=(0, 0)),deltaF*np.arange(int(2*ceil(len(realhcf1)+100/deltaF)))],fs=fs,input_onesided=True)
#tc1,realhct1=scipy.signal.istft([np.pad(np.array(realhcf1), (int(ceil(100/deltaF)),int(ceil(len(realhcf1)+100/deltaF))), 'constant', constant_values=(0, 0)),deltaF*np.arange(int(2*ceil(len(realhcf1)+100/deltaF)))],fs=fs,input_onesided=True)
#tp1,realhpt1=scipy.signal.istft(np.array([realhpf1,f1]), fs=1/deltaF,  input_onesided=True)
#tc1,realhct1=scipy.signal.istft(np.array([realhcf1,f1]), fs=1/deltaF,  input_onesided=True)
deltaT=1/(f1[len(f1)-1])
realhpt1_new=realhpt1[int(ceil(len(realhpt1)/2)):]
realhct1_new=realhct1[int(ceil(len(realhct1)/2)):]
t1=-(deltaT*len(realhpt1_new))+deltaT*np.arange(len(realhpt1_new))

# create first signal with tidal effects
lambda1 =100.
lambda2=100.
hpf2, hcf2 = lalsim.SimInspiralChooseFDWaveform(phiRef,deltaF,m1_SI,m2_SI,s1x,s1y,s1z,
                        s2x,s2y,s2z,fmin,fmax,fref,distance,inclination,lambda1,lambda2,None,None,0,0,approx)
# the waveform is generated in the frequency domain as a COMPLEX frequency series
realhpf2=hpf2.data.data
realhcf2=hcf2.data.data
# get real signal
realhpf2=realhpf2.real
realhcf2=realhcf2.real
f2=fmin+deltaF*np.arange(len(realhpf2))
# get signal in time domain
realhpt2=np.fft.irfft(np.pad(np.array(realhpf2), (int(ceil(100/deltaF)),0), 'constant', constant_values=(0, 0)))
realhct2=np.fft.irfft(np.pad(np.array(realhcf2), (int(ceil(100/deltaF)),0), 'constant', constant_values=(0, 0)))
deltaT2=1/(f2[len(f2)-1])
realhpt2_new=realhpt2[int(ceil(len(realhpt2)/2)):]
realhct2_new=realhct2[int(ceil(len(realhct2)/2)):]
t2=-(deltaT2*len(realhpt2_new))+deltaT2*np.arange(len(realhpt2_new))

#create second signal with tidal effects
lambda1 =1000.
lambda2=1000.
hpf3, hcf3 = lalsim.SimInspiralChooseFDWaveform(phiRef,deltaF,m1_SI,m2_SI,s1x,s1y,s1z,
                        s2x,s2y,s2z,fmin,fmax,fref,distance,inclination,lambda1,lambda2,None,None,0,0,approx)
# the waveform is generated in the frequency domain as a COMPLEX frequency series
realhpf3=hpf3.data.data
realhcf3=hcf3.data.data
# get real signal
realhpf3=realhpf3.real
realhcf3=realhcf3.real
f3=fmin+deltaF*np.arange(len(realhpf3))
# get signal in time domain
realhpt3=np.fft.irfft(np.pad(np.array(realhpf3), (int(ceil(100/deltaF)),0), 'constant', constant_values=(0, 0)))
realhct3=np.fft.irfft(np.pad(np.array(realhcf3), (int(ceil(100/deltaF)),0), 'constant', constant_values=(0, 0)))
deltaT3=1/(f3[len(f3)-1])
realhpt3_new=realhpt3[int(ceil(len(realhpt3)/2)):]
realhct3_new=realhct3[int(ceil(len(realhct3)/2)):]
t3=-(deltaT3*len(realhpt3_new))+deltaT3*np.arange(len(realhpt3_new))

realhpt1=realhpt1_new
realhct1=realhct1_new
realhpt2=realhpt2_new
realhct2=realhct2_new
realhpt3=realhpt3_new
realhct3=realhct3_new

#plot BBH signal to check for mistakes
plotName="%s/frequency_signal.pdf"%(plotDir)
plt.figure(figsize=(12,8))
#plt.plot(f1,realhpf1_new,label="plus")
#plt.plot(f1,realhcf1_new,label='cross')
plt.plot(f1,realhpf1,label="plus")
plt.plot(f1,realhcf1,label='cross')
plt.legend(loc='best')
plt.xlabel=('Frequency (Hz)')
plt.ylabel=('Signal')
plt.savefig(plotName,bbox_inches='tight')
plt.close()

plotName="%s/time_signal.pdf"%(plotDir)
plt.figure(figsize=(12,8))
plt.plot(t1,realhpt1,label="plus")
plt.plot(t1,realhct1,label='cross')
plt.legend(loc='best')
plt.xlabel=('Time (s)')
plt.ylabel=('Signal')
plt.xlim(-6,0)
plt.savefig(plotName,bbox_inches='tight')
plt.close()

# plot all signals for comparison
plotName="%s/frequency_allsignals.pdf"%(plotDir)
plt.figure(figsize=(12,8))
plt.subplot(211)
plt.plot(f1,realhpf1,label="BBH")
plt.xlabel=('Frequency (Hz)')
plt.ylabel=('Signal')
plt.legend(loc='best')
plt.subplot(212)
plt.plot(f1,realhpf1,label="BBH")
plt.plot(f2,realhpf2,label="Lambda1,2=100")
plt.plot(f3,realhpf3,label="Lambda1,2=1000")
plt.xlabel=('Frequency (Hz)')
plt.ylabel=('Signal')
plt.legend(loc='best')
plt.savefig(plotName,bbox_inches='tight')
plt.close()

plotName="%s/time_allsignals.pdf"%(plotDir)
plt.figure(figsize=(12,8))
plt.subplot(211)
plt.plot(t1,realhpt1,label="BBH")
plt.xlabel=('Time (s)')
plt.ylabel=('Signal')
plt.xlim(-0.05,0)
plt.legend(loc='best')
plt.subplot(212)
plt.plot(t1,realhpt1,label="BBH")
plt.plot(t2,realhpt2,label="Lambda1,2=100")
plt.plot(t3,realhpt3,label="Lambda1,2=1000")
plt.xlabel=('Time (s)')
plt.ylabel=('Signal')
plt.xlim(-0.05,0)
plt.legend(loc='best')
plt.savefig(plotName,bbox_inches='tight')
plt.close()


# create complex signal hplus+i*hcross and get phase in frequency domain
complexhf1=np.zeros(len(realhpf1),dtype=complex)
phasef1=np.zeros(len(realhpf1))
for i in range(len(complexhf1)):
        complexhf1[i]=complex(realhpf1[i],realhcf1[i])
        phasef1[i]=np.angle(complexhf1[i])
complexhf2=np.zeros(len(realhpf2),dtype=complex)
phasef2=np.zeros(len(realhpf2))
for i in range(len(complexhf2)):
        complexhf2[i]=complex(realhpf2[i],realhcf2[i])
        phasef2[i]=np.angle(complexhf2[i])
complexhf3=np.zeros(len(realhpf3),dtype=complex)
phasef3=np.zeros(len(realhpf3))
for i in range(len(complexhf3)):
        complexhf3[i]=complex(realhpf3[i],realhcf3[i])
        phasef3[i]=np.angle(complexhf3[i])
      
# create complex signal hplus+i*hcross and get phase in time domain
complexht1=np.zeros(len(realhpt1),dtype=complex)
phaset1=np.zeros(len(realhpt1))
for i in range(len(complexht1)):
        complexht1[i]=complex(realhpt1[i],realhct1[i])
        phaset1[i]=np.angle(complexht1[i])
complexht2=np.zeros(len(realhpt2),dtype=complex)
phaset2=np.zeros(len(realhpt2))
for i in range(len(complexht2)): 
        complexht2[i]=complex(realhpt2[i],realhct2[i])
        phaset2[i]=np.angle(complexht2[i])
complexht3=np.zeros(len(realhpt3),dtype=complex)
phaset3=np.zeros(len(realhpt3)) 
for i in range(len(complexht3)):
        complexht3[i]=complex(realhpt3[i],realhct3[i])
        phaset3[i]=np.angle(complexht3[i])

# compare phase of signal with tidal distortion to phase of BBH in frequency domain
deltaphf2=np.zeros(len(phasef1))
deltaphf3=np.zeros(len(phasef1))
for i in range(len(phasef1)):
	deltaphf2[i]=phasef2[i]-phasef1[i]
        deltaphf3[i]=phasef3[i]-phasef1[i]
deltaphf2=np.unwrap(deltaphf2)
deltaphf3=np.unwrap(deltaphf3)
plotName="%s/frequency_phase.pdf"%(plotDir)
plt.figure(figsize=(12,8))
plt.plot(f1,deltaphf2,label="Lambda1,2=100")
plt.plot(f1,deltaphf3,label="Lambda1,2=1000")
#plt.xlabel('Frequency (Hz)')
#plt.ylabel('Phase')
plt.legend(loc='best')
plt.savefig(plotName,bbox_inches="tight")
plt.close()

# compare phase of signal with tidal distortion to phase of BBH in time domain
deltapht2=np.zeros(len(phaset1))
deltapht3=np.zeros(len(phaset1))
for i in range(len(phaset1)):
        deltapht2[i]=phaset2[i]-phaset1[i]
        deltapht3[i]=phaset3[i]-phaset1[i]
deltapht2=np.unwrap(deltapht2)
deltapht3=np.unwrap(deltapht3)
plotName="%s/time_phase.pdf"%(plotDir)
plt.figure(figsize=(12,8))
plt.plot(t1,deltapht2,label="Lambda1,2=100")
plt.plot(t1,deltapht3,label="Lambda1,2=1000")
#plt.xlabel('Time (s)')
#plt.ylabel('Phase')
plt.legend(loc='best')
plt.xlim(-10,0.5)
plt.ylim(-1,20)
plt.savefig(plotName,bbox_inches="tight")
plt.close()

# #parameter estimation:
# print(deltaF)
# print(deltaT2)
# #generate noise for lambda=100 signal
# flow = 30.0
# flen = int(2048 / deltaF) + 1
# psd = pycbc.psd.aLIGOZeroDetHighPower(flen, deltaF, flow)
# tsamples = len(realhpt2)-1
# ts = pycbc.noise.noise_from_psd(tsamples, deltaT2, psd, seed=127)

#add noise to signal
#for i in range(len(realhpt2)):
#	realhpt2[i]=realhpt2[i]+ts.data.data[i]