#!/usr/bin/python
# Copyright (C) 2012  Alan Weinstein
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# A copy of the GNU General Public License may be found at
# http://www.gnu.org/copyleft/gpl.html
# or write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

"""
plot_data.py
This is example python code to read in detector strain data,
and plot it in various ways.
"""

##############################################
# import everything we need, up front:
import numpy as np
import h5py
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
from matplotlib import mlab 

# the directory containing the data files:
dir = './'

#############
# loop over detectors:
for IFO in ['H2','L1']:

   # First, the 256s of data around GRB051103
   # sampling frequency in Hz and duration in seconds
   fs = 4096
   Duration = 256
   GPSstart = 815045078

   # open the file and get the data (strain time series):
   fn = dir+IFO+'-'+'STRAIN_'+str(fs)+'Hz-'+str(GPSstart)+'-'+str(Duration)+'.hdf5'
   f = h5py.File(fn,'r') 
   d = f[IFO+':LSC-STRAIN']
   data = d[:]
   f.close()

   # print out some basic data checks
   print('Contents of '+fn) 
   print('IFO, Npoints = '+IFO+', '+str(data.size))
   print('Duration = '+str(data.size/(float(fs*Duration)))+' * '+str(fs)+' * '+str(Duration))
   print('min,max,mean,std = '+str(min(data))+', '+str(max(data))+', '+str(data.mean())+
         ', '+str(data.std().tolist()))

   ## histogram the raw data:
   hmax = max(data)
   plt.figure()
   hn, hbins, patches = plt.hist(data,bins=100, range=(-hmax, hmax))
   # make a corresponding Gaussian:
   N = hn.max()
   mu = data.mean()
   sigma = data.std().tolist()
   y = N*np.exp(-np.power((hbins-mu)/sigma,2)/2.0)
   # plot it:
   plt.plot(hbins,y, 'r--', linewidth=1)
   plt.xlim((-hmax, hmax))
   plt.xlabel("data value (strain)")
   plt.ylabel("N")
   tit =  'N = {0}, mean = {1:7.2e}, std = {2:7.2e}'.format(len(data),data.mean(),data.std(0).tolist())
   plt.title(tit)
   plt.savefig(IFO+'_hhist.png')

   # plot the first bunch of data points 
   Npoints = 10000
   datn = data[0:Npoints]
   # make a time vector
   timn = np.arange(0,Npoints)/float(fs)
   plt.figure()
   plt.plot(timn,datn,'b')
   plt.xlabel('time (s)')
   plt.ylabel(IFO+'strain data value (strain)')
   tit =  'N = {0}, mean = {1:7.2e}, std = {2:7.2e}'.format(len(data),data.mean(),data.std(0).tolist())
   plt.savefig(IFO+'_hoft.png')

   # make the asd
   psdm = matplotlib.mlab.psd(data,fs,fs,noverlap=fs/2, window=np.blackman(fs)) 
   asdm = np.sqrt(psdm[0]) 
   freqm = psdm[1] 

   # plot the asd
   plt.figure()
   plt.loglog(freqm, asdm,'b')
   plt.xlim((20, 4196)) 
   plt.ylim((1.e-23, 1.e-18)) 
   plt.grid(True) 
   plt.xlabel("f (Hz)")
   plt.ylabel("ASD")
   plt.title(IFO+' noise ASD around GRB051103')
   plt.savefig(IFO+'_asd.png') 

   plt.figure()
   # make and plot the spectrogram
   Pxx, freqs, bins, im = plt.specgram(data, NFFT=1024, Fs=fs)
   plt.xlabel("Time since 815045078 (s)") 
   plt.ylabel("frequency (Hz)") 
   plt.title(IFO+' spectrogram around GRB051103')
   plt.savefig(IFO+'_spectrogram.png')

   ##############
   # Now, the 2190s of data around GRB051103
   # sampling frequency in Hz and duration in seconds
   fs = 16384
   Duration = 2190
   GPSstart = 815043278

   # open the file and get the data (strain time series):
   fn = dir+IFO+'-'+'STRAIN_'+str(fs)+'Hz-'+str(GPSstart)+'-'+str(Duration)+'.hdf5'
   f = h5py.File(fn,'r') 
   d = f[IFO+':LSC-STRAIN']
   data = d[:]
   f.close()

   # print out some basic data checks
   print('Contents of '+fn) 
   print('IFO, Npoints = '+IFO+', '+str(data.size))
   print('Duration = '+str(data.size/(float(fs*Duration)))+' * '+str(fs)+' * '+str(Duration))
   print('min,max,mean,std = '+str(min(data))+', '+str(max(data))+', '+str(data.mean())+
         ', '+str(data.std().tolist()))

   ## histogram the raw data:
   hmax = max(data)
   plt.figure()
   hn, hbins, patches = plt.hist(data,bins=100, range=(-hmax, hmax))
   # make a corresponding Gaussian:
   N = hn.max()
   mu = data.mean()
   sigma = data.std().tolist()
   y = N*np.exp(-np.power((hbins-mu)/sigma,2)/2.0)
   # plot it:
   plt.plot(hbins,y, 'r--', linewidth=1)
   plt.xlim((-hmax, hmax))
   plt.xlabel("data value (strain)")
   plt.ylabel("N")
   tit =  'N = {0}, mean = {1:7.2e}, std = {2:7.2e}'.format(len(data),data.mean(),data.std(0).tolist())
   plt.title(tit)
   plt.savefig(IFO+'_hhistf.png')

   # make the asd
   # pick out the 256s overlap: data = data[(1800*16384):(2056*16384-1)]
   psdf = matplotlib.mlab.psd(data,fs,fs,noverlap=fs/2, window=np.blackman(fs)) 
   asdf = np.sqrt(psdf[0]) 
   freqf = psdf[1] 

   # now read in the spectrum from the distributed file:
   fn = dir+IFO+'-'+'SPEC-'+str(GPSstart)+'-'+str(Duration)+'.hdf5' 
   f = h5py.File(fn,'r') 
   dspec = f[IFO+':SPECTRUM']
   asdd = dspec[:,1]
   freqd = dspec[:,0]
   f.close()

   # plot the asd, this time with both the short and long data stretches
   plt.figure()
   plt.loglog(freqm, asdm,'b', label='256 s @ 4096 Hz')
   plt.loglog(freqf, asdf,'r', label='2190 s @ 16384 Hz')
   plt.loglog(freqd, asdd,'g--', label='SPEC file')
   plt.xlim((20, 4196)) 
   plt.ylim((1.e-23, 1.e-18)) 
   plt.grid(True) 
   plt.xlabel("f (Hz)")
   plt.ylabel("ASD")
   plt.legend()
   plt.title(IFO+' noise ASD around GRB051103')
   plt.savefig(IFO+'_asd.png')