#!/usr/bin/env python
# this function gets some data (from the 40m) and saves it as
# a .mat file for the matlabs
# Ex. python -O getData.py


from scipy.io import savemat,loadmat
import scipy.signal as sig
from astropy.time import Time
import nds2
import os, sys, time
import argparse


parser = argparse.ArgumentParser()
parser.add_argument("-ifo", "--ifo",
                    help="interferometer: L1, H1, C1, A1", default="C1")
parser.add_argument("-dur", "--duration",
                    help="duration in seconds", default=8)
parser.add_argument("-fs", "--fs",
                    help="sample rate to decimate to [Hz]; Ex. 128", default=128)


args = parser.parse_args()

ifo  = str(args.ifo)
fs   = int(args.fs)
dur  = int(args.duration)

# channel names
chan_head = ifo + ':'
#fname = 'ChanList_S2L.txt'
fname = 'chanlist.txt'

with open(fname, 'r') as f:
    chanlines = f.read().split()
channels = [chan_head + line for line in chanlines]

# 3 hour stretch with moderate high uSeism
portNumber = 31200
scale = 'utc'
if ifo == 'L1':
    ndsServer  = 'nds.ligo.caltech.edu'
    times      = '2019-11-15 03:15:00'     #
elif ifo == 'H1':
    ndsServer  = 'nds.ligo-wa.caltech.edu'
    times      = '2017-03-14 10:05:00'     # 3 hour stretch with moderate high uSeism
elif ifo == 'C1':
    ndsServer  = '131.215.115.200'
    times      = '2019-11-20 05:00:00'     # 
    portNumber = 31200
else:
    sys.exit("unknown IFO specified")

# Setup connection to the NDS

conn       = nds2.connection(ndsServer, portNumber)

# Setup start and stop times
# good double coinc time = '2017-03-11 14:00:00'
#times   = '2017-03-11 14:00:00'
t       = Time(times, format='iso', scale=scale)
t_start = int(t.gps)
if __debug__:
    print(t_start)



# Data will be downsampled to `fsup` Hz
fsup = fs
if __debug__:
    print("Output sample rate: {} Hz".format(fsup))


if __debug__:
    print("List of Channels:...")
    print("\n".join(channels))

print("Getting data from " + ndsServer + "...")
tic = time.time()
data = conn.fetch(t_start, t_start + dur, channels)

# get the data and stack it into a single matrix where the data are the columns
vdata = []
for k in range(len(channels)):
    fsdown = data[k].channel.sample_rate
    if fsup > fsdown:
        print('WARNING: bad downsample freq!!')
    down_factor = int(fsdown // fsup)

    fir_aa = sig.firwin(20 * down_factor + 1, 0.8 / down_factor,
        window='blackmanharris')

    # Using fir_aa[1:-1] cuts off a leading and trailing zero
    downdata = sig.decimate(data[k].data, down_factor,
                            ftype = sig.dlti(fir_aa[1:-1], 1.0),
                            zero_phase = True)
    vdata.append(downdata)

# save to a hdf5 format that matlab can read (why is compression off by default?)
save_dir = 'Data/'
funame = save_dir + ifo + '_data_array.mat'

# make the save_dir if it DNE
if os.path.exists(save_dir):
    abcdef = 42  # dummy line
else:
    os.mkdir(save_dir)

savemat(funame,
            mdict={'data': vdata, 'fsample': fsup, 'chans': channels},
            do_compression=True)
print("Data saved as " + funame)

if __debug__:
    print(" ")
    print("Channel name        = " + data[0].channel.name)
    print("Sample rate         = " + str(fsup) + " Hz")
    print("Number of samples   = " + str(vdata[0].size))
    print("GPS Start time      = " + str(data[0].gps_seconds))
    print("Data retrieval time = " + str(round(time.time() - tic,3)) + " seconds.")

# uncomment this stuff to get info on what fields are in the data
#dir(data[0])
#dir(data[0].channel)