import sys
import os
import numpy as np
from numpy.fft import *
pretty=True #False
if pretty==True:
    %matplotlib notebook 
from ipywidgets import interact, interactive, fixed, interact_manual
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
from matplotlib.colors import LogNorm
from scipy.signal import *
import scipy.io
import scipy
from astropy.utils import iers
iers.conf.auto_download = False
import warnings

import gwosc
import gwpy
import pycbc

print(gwosc.__version__)
print(gwpy.__version__)
print(pycbc.__version__)
print(sys.executable)

from gwosc.datasets import find_datasets, event_gps, run_segment
from gwosc.locate import get_event_urls
from gwosc import datasets

#find_datasets??

V1_events = find_datasets(detector="V1")           # try the other 2 Advanced detectors
print(V1_events)

print("List of available catalogs: \n \n", find_datasets(type="catalog"))    # try other type: "event"

gwtc = datasets.find_datasets(type='event', catalog='O3_Discovery_Papers')
#gwtc = datasets.find_datasets(type='event', catalog='GWTC-2')
#gwtc = datasets.find_datasets(type='catalog', catalog='O3_Discovery_Papers')
print('O3 events:', gwtc)

runs = find_datasets(type='run')
print('Full datasets from runs:', runs)

#run_segment??

print(run_segment('O2_4KHZ_R1'))

gps_ev = event_gps('GW190521-v2')
print(gps_ev)

urls = get_event_urls('GW190521-v2')
print(urls)

from gwpy.timeseries import TimeSeries

#TimeSeries.fetch_open_data??

seg_hlen = 8
l1_data = TimeSeries.fetch_open_data('L1', gps_ev-seg_hlen, gps_ev+seg_hlen, verbose=True)  # cache=True stores file on your computer.
print(l1_data)

#type(gwpy.frequencyseries.frequencyseries.FrequencySeries) 

#gwpy.timeseries.timeseries.TimeSeries.plot??

plot = l1_data.plot(figsize=(8, 3), title='Strain for event from L1');

#TimeSeries.fft??

fft = TimeSeries.fft(l1_data)
plot = fft.abs().plot(xscale="log", yscale="log",figsize=(4,4), title='ASD no window')

from scipy.signal import get_window
window = get_window('hann', l1_data.size)   # Grab window using scipy library
lwin = l1_data * window

#plt.figure(figsize=(3,2))         # matplotlib.pyplot plot!!
#plt.plot(window, 'r'), plt.title('Hann window')
#TimeSeries(window).plot()

fftamp = lwin.fft().abs()
plot_new = fftamp.plot(xscale="log", yscale="log", figsize=(4,4), title='ASD with Hann window')

#gwpy.timeseries.TimeSeries.asd??

asd_l1 = l1_data.asd(fftlength=2, window='hann', method="welch")
plot = asd_l1.plot(title='Welch avg. ASD from 10 to 1500 Hz', figsize=(4,4), color='gwpy:ligo-livingston')
ax = plot.gca()           # get current axis to modify plot axes' intervals
ax.set_xlim(10, 1500)
ax.set_ylim(1e-24, 1e-20);

##gwpy.timeseries.TimeSeries.spectrogram??

##gwpy.timeseries.TimeSeries.spectrogram2??

specgram = l1_data.spectrogram2(fftlength=1, overlap=0.8, window='hann') ** (1/2.)
plot = specgram.plot(title='Not so good.', figsize=(7,5))

adjust=False
if adjust==True:
    ax = plot.gca()
    ax.set_yscale('log')
    ax.set_ylim(10, 1500)
    ax.colorbar(
        clim=(1e-24, 1e-20),
        norm="log",
        label=r"Strain noise [$1/\sqrt{\mathrm{Hz}}$]",
    )
    ax.set_title('Better!')
    #plot  # refresh

#gwpy.timeseries.TimeSeries.q_transform??

l1_q = l1_data.q_transform(frange=(30, 200), outseg=(gps_ev-0.25,gps_ev+0.25))   # outseg to zoom around merger
plot = l1_q.plot(title='GW190521 at L1', figsize=(7,5))
plot.colorbar(label="Normalised energy");

#v1_q = v1_data.q_transform(frange=(30, 200), outseg=(gps_ev-0.25,gps_ev+0.25))
#plot = v1_q.plot(title='GW190521 at V1', figsize=(7,5))
#plot.colorbar(label="Normalised energy");

import pycbc.filter
from pycbc import types
from pycbc.psd import interpolate, inverse_spectrum_truncation
from pycbc.waveform import get_td_waveform
from pycbc.detector import Detector

#get_td_waveform??

f_samp = 4096
hp, hc = get_td_waveform(approximant="SEOBNRv4_opt",       # others, such as "IMRPhenomD"
                         mass1=12,                         # in solar masses
                         mass2=15,
                         delta_t=1.0/f_samp,
                         f_lower=30,
                         distance=100,                     # in Mpc
                         inclination=0)                    # in radians  

type(hp)

plt.figure(figsize=(10,4.5))
plt.plot(np.arange(0,len(hp))/f_samp, hp+hc, 'g')
plt.xlabel('[s]')
plt.title(r'$m1 = 12M_\odot$, $m2 = 15M_\odot$');

plt.figure(figsize=(8,5))
plt.subplot(2,1,1)
plt.plot(hp.sample_times, hp, 'b', label='plus pol.')
plt.legend()
plt.subplot(2,1,2)
plt.plot(hc.sample_times, hc, 'r', label='cross pol.')
plt.legend()
plt.tight_layout(pad=1)

from pycbc.detector import Detector, get_available_detectors

print(get_available_detectors())

for abv, long_name in get_available_detectors():
    d = Detector(abv)
    print("{} Latitude {} Longitude {}".format(abv, d.latitude, d.longitude))    # Latitude and longitude in radians

#Detector.antenna_pattern??

### Right ascension, declination, polarization angle and inclination angle (radians), GPS time
ra = 1.7
dec = 1.7
pol = 0.2
inc = np.pi/4
time = 1242442967.4

### Recall how we produced the above waveform
hp, hc = get_td_waveform(approximant="SEOBNRv4_opt", mass1=12, mass2=15,
                         f_lower=30, delta_t=1.0/f_samp, inclination=0,
                         distance=10)

det = Detector("L1")                                  # Let's chooose the LIGO Livingston detector
fp, fc = d.antenna_pattern(ra, dec, pol, time)        # weight factors for h+ and hx
print("The antenna pattern factors are fp={}, fc={}".format(fp, fc))

### Apply the antenna pattern factors to the two polarizations, then sum and plot
ht = fp * hp + fc * hc

### ALTERNATIVE: the project_wave_function.
hp.start_time = hc.start_time = time
ht = det.project_wave(hp, hc, ra, dec, pol)
plt.figure(figsize=(10,4.5))
plt.plot(ht.sample_times, ht, 'g')
plt.xlabel('[s]')
plt.xlim(min(ht.sample_times), max(ht.sample_times))
plt.title(r'$m1 = 12M_\odot$, $m2 = 15M_\odot$ at ' + det.name);

from pycbc.frame import read_frame
from pycbc.filter import highpass_fir, lowpass_fir
from pycbc.waveform import get_fd_waveform
from pycbc.psd import welch, interpolate
from pycbc.catalog import Merger
import pycbc.catalog

gwtc2_events = pycbc.catalog.Catalog(source='O3_Discovery_Papers')
print(gwtc2_events.names)

plt.figure(figsize=(9,7))
i_c=1
for ifo in ['V1','H1', 'L1']:
    # Read data and remove low frequency content
    h1 = Merger('GW170814',source='gwtc-1').strain(ifo)
    h1 = highpass_fir(h1, 15, 8)
    merger_time = Merger('GW170814',source='gwtc-1').time

    # Calculate the noise spectrum
    psd = interpolate(welch(h1), 1.0 / h1.duration)

    # whiten
    white_strain = (h1.to_frequencyseries() / psd ** 0.5).to_timeseries()

    # remove some of the high and low
    smooth = highpass_fir(white_strain, 20, 8)
    smooth = lowpass_fir(white_strain, 300, 8)

    # time shift and flip L1
    if ifo == 'H1':
        itf_color='gwpy:ligo-hanford'
    elif ifo == 'L1':
        itf_color='gwpy:ligo-livingston'
    elif ifo == 'V1':
        itf_color='gwpy:virgo'
    
    plt.subplot(3,1,i_c)
    plt.plot(smooth.sample_times, smooth, label=ifo, color=itf_color)
    plt.legend()
    plt.xlim(merger_time-0.3, merger_time+0.3)
    plt.ylim(-150, 150)
    plt.ylabel(r'$h_w$')
    plt.grid()
    plt.xlabel('GPS Time (s)')
    i_c+=1

from pycbc.filter import matched_filter
from pycbc.waveform import get_fd_waveform

# Generate a template to filter with
hp, hc = get_fd_waveform(approximant="IMRPhenomD", mass1=30.5, mass2=25.3,
                         f_lower=20, delta_f=1.0/h1.duration)
hp.resize(len(h1) // 2 + 1)

# Calculate the complex (two-phase SNR)
snr = matched_filter(hp, h1, psd=psd, low_frequency_cutoff=20.0)

# Remove regions corrupted by filter wraparound
snr = snr[len(snr) // 3: len(snr) * 2 // 3]


plt.figure()
plt.plot(snr.sample_times, abs(snr))
plt.ylabel('signal-to-noise')
plt.xlabel('GPS Time (s)');

gps_ev

# Hanford H1
h1_data = TimeSeries.fetch_open_data('H1', gps_ev-seg_hlen, gps_ev+seg_hlen)
asd_h1 = h1_data.asd(fftlength=4, window='hann', method="welch")

# Virgo V1
v1_data = TimeSeries.fetch_open_data('V1', gps_ev-seg_hlen, gps_ev+seg_hlen)
asd_v1 = v1_data.asd(fftlength=4, window='hann', method="welch")

plot_all = asd.plot(title='Welch avg. ASD from 10 to 1500 Hz', figsize=(4,4), color='gwpy:ligo-livingston')
ax2 = plot_all.gca()           # get current axis 
ax2.plot(asd_h1, label='H1', color='gwpy:ligo-hanford')
ax2.plot(asd_v1, label='V1', color='gwpy:virgo')
ax2.plot(asd_l1, label='L1', color='gwpy:ligo-livingston')

ax2.set_xlim(10,1500)
ax2.set_ylim(1e-24, 1e-20)
ax2.set_ylabel(r'Strain noise [$1/\sqrt{\mathrm{Hz}}$]')
ax2.legend()
plot;

### 1.
event_list = datasets.find_datasets(type='events', catalog='GWTC-1-confident')
print(event_list)

dt=16
gps_GW = event_gps('GW170817-v3')
h1_GW = TimeSeries.fetch_open_data('H1', gps_GW-dt, gps_GW+dt)
l1_GW = TimeSeries.fetch_open_data('L1', gps_GW-dt, gps_GW+dt)
v1_GW = TimeSeries.fetch_open_data('V1', gps_GW-dt, gps_GW+dt)

### 2.
h1_qGW = h1_GW.q_transform(frange=(30, 500), qrange=(100,110), outseg=(gps_GW-3,gps_GW+3))   # recall outseg to zoom
l1_qGW = l1_GW.q_transform(frange=(30, 500), qrange=(100,110), outseg=(gps_GW-3,gps_GW+3))   
V1_q = v1_GW.q_transform(frange=(30, 500), qrange=(100,110), outseg=(gps_GW-3,gps_GW+3))   

### 3. The signal is clearly visible at LIGO Hanford
plotGW_h1 = h1_qGW.plot(title='GW190521 at H1', figsize=(7,5))
plotGW_l1 = l1_qGW.plot(title='GW190521 at L1', figsize=(7,5))
plotGW_v1 = V1_q.plot(title='GW190521 at V1', figsize=(7,5))
ax_GW = plot.gca()
ax_GW.set_yscale('log')
ax_GW.colorbar(label="Normalised energy");

### 4. An instrumental glitch is saturating the LIGO Livingston data!

itfs=['L1','H1','V1']
ii=0
for itf1 in itfs:
    ii+=1
    for itf2 in itfs[ii:]:
        tt =Detector.light_travel_time_to_detector(Detector(itf1), Detector(itf2))
        print("Light travel time from {} to {} is {} seconds".format(itf1, itf2, tt))