# speclines.py
# read in, plot, and find lines in S5,S6 hi-resolution spectra
# AJW, 6/14

###################################################
import numpy
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
#from matplotlib import mlab 
import scipy.io

###################################################
# specify the run and the detectors:
srun = 'S5'
det = ['S5H1','S5H2','S5L1']
detc = {'S5H1':'r','S5H2':'b','S5L1':'g'}
detf = {'S5H1':1.e-22,'S5H2':2.e-22,'S5L1':1.e-22}

srun = 'S6'
det = ['S6H1','S6L1']
detc = {'S6H1':'r','S6L1':'g'}
detf = {'S6H1':1.e-22,'S6L1':1.e-22}

###################################################
# fit the continuum part of the spectrum
from scipy import optimize
clipl = 0.1
fitfunc = lambda p, f: p[0]*f**(-3.4)+p[1]*f**(-0.)+p[2]*f**1. # Target function
errfunc = lambda p, f, y: numpy.clip((1.-y/fitfunc(p, f)),-clipl,clipl) # Distance to the target function
p0 = numpy.array([2.0*60.**(3.4), 0.05, 1.4e-3 ]) # Initial guess for the parameters

###################################################
# read in the hi-res spectra from
# https://ldas-jobs.ligo.caltech.edu/~keithr/spectra/spectra_withS5H2/
# https://ldas-jobs.ligo.caltech.edu/~keithr/spectra/spectra_withS5H2/S5spectradata.mat
# https://ldas-jobs.ligo.caltech.edu/~keithr/spectra/S6spectradata.mat

fndir = './'
fnmat = srun+'spectradata.mat'
specdata = scipy.io.loadmat(fndir+fnmat,struct_as_record=True)
# need to use "ravel" to compress the data into a properly-formed numpy array.
fr   = numpy.ravel(specdata['frequency'])
asds = {}
for deti in det:
  asds[deti] = numpy.ravel(specdata[deti+'amppsdwt'])
  p1, success = optimize.leastsq(errfunc, detf[deti]*p0, args=(fr, asds[deti]))
  asds[deti+'fit'] = fitfunc(p1,fr)

print('Spectra with {} frequency bins of size {} from {} to {} Hz'\
    .format(fr.size, fr[1]-fr[0], fr[0], fr[-1]))

###################################################
# plot the spectrum with the fits
plt.figure()
for deti in det:
  plt.loglog(fr,asds[deti],detc[deti], label=deti)
for deti in det:
  plt.loglog(fr,asds[deti+'fit'],detc[deti]+'--', label=deti+'fit')
plt.xlim([40, 2000])
plt.grid(b=True, which='both')
plt.xlabel('frequency (Hz)')
plt.ylabel(srun+' noise spectra (strain/rtHz)')
plt.legend()
plt.savefig(srun+'spec.png')
#plt.show()

##############################################
# identify peaks in each spectrum:
for deti in det:

    # simple peak finder:
    # first downsample
    fa = fr[9:-9].reshape(-1, 18).mean(axis=1)
    aa = asds[deti][9:-9].reshape(-1, 18).mean(axis=1)
    af = asds[deti+'fit'][9:-9].reshape(-1, 18).mean(axis=1)
    #aa = aa/af
    # exclude the first and last 5 bins:
    fae = fa[5:-5]
    aae = aa[5:-5]
    afe = af[5:-5]
    # get "background" from +-1 bin:
    aav1 = (aa[4:-6]+aa[6:-4])/2.
    # get "background" from +-5 bins;
    aav2 = (aa[:-10]+aa[10:])/2.
    rm1 = aae/aav1
    rm2 = aae/aav2
    bi = (rm1 > 1.003)&(rm2 > 1.04)
    bi = (rm1 > 1.01)&(rm2 > 1.1)
    #ba = (aa[5:-5]>aa[4:-6])&(aa[5:-5]>aa[6:-4])&(aa[5:-5]>aa[3:-7])&(aa[5:-5]>aa[7:-3])& \
    #     (aa[5:-5]>aa[2:-8])&(aa[5:-5]>aa[8:-2])&(aa[5:-5]>aa[1:-9])&(aa[5:-5]>aa[9:-1])
    alines = numpy.array([fae[bi],aae[bi],aae[bi]/afe[bi]]).transpose()
    
    # output table of lines:
    numpy.savetxt(deti+'lines.txt',alines,'%10.4f %12.4e %10.4f')
    
    # plot the spectrum with the identified peaks
    plt.figure()
    #plt.loglog(fr,asd,detc[deti],label=deti)
    plt.loglog(fa,aa,detc[deti],label=deti)
    plt.loglog(fae[bi],aae[bi],'*')
    plt.xlim([40, 2000])
    plt.grid(b=True, which='both')
    plt.xlabel('frequency (Hz)')
    plt.ylabel(srun+' noise spectra (strain/rtHz)')
    plt.legend(loc='lower right')
    plt.savefig(deti+'peaks.png')
    #plt.show()

    # plot the SNR with the identified peaks
    plt.figure()
    plt.loglog(fa,aa/af,detc[deti],label=deti)
    plt.loglog(fae[bi],aae[bi]/afe[bi],'*')
    plt.xlim([40, 2000])
    plt.grid(b=True, which='both')
    plt.xlabel('frequency (Hz)')
    plt.ylabel(srun+' SNR')
    plt.legend(loc='lower right')
    plt.savefig(deti+'peaksnr.png')
    #plt.show()
    
# Calibration lines in S5, from JZ, 2/7/13:
# H1: 46.7, 393.1, 1144.3Hz
# L1: 54.7, 396.7, 1151.5Hz