#! /usr/bin/python

from numpy import *
#import nds2
import matplotlib
matplotlib.use("Agg", warn=False)
import pylab
from scipy import signal
import cdsutils
from scipy.optimize import leastsq
from scipy.optimize import curve_fit
import matplotlib.gridspec as gridspec


## Parametric function: 'v' is the vector of parameters, 'x' the independent variable
def decay(t,A0,Q):
    f0 = 9.848
    tau = Q / (pi*f0)
    return A0*exp(-t/tau)


def decimate(x, q, n=None, ftype='iir', axis=-1):
    if not isinstance(q, int):
        raise TypeError("q must be an integer")
    if n is None:
        if ftype == 'fir':
            n = 30
        elif q>32:
            n = 2
        else:
            n = 3
    if ftype == 'fir':
        b = signal.firwin(n + 1, 1. / q, window='hamming')
        a = 1.
    else:
        b, a = signal.cheby1(n, 0.05, 0.8 / q)
        #b, a = signal.butter(4, 0.8 / q, 'low', analog=False)

    #print b, a
    y = signal.filtfilt(b, a, x)
    sl = [slice(None)] * y.ndim
    sl[axis] = slice(None, None, q)
    return y[sl]


def bandpass(x,f1,f2,fs):
    nyq = fs/2
    b, a = signal.butter(4, [f1/nyq, f2/nyq], 'bandpass')
    y = signal.filtfilt(b, a, x)
    return y


def lowpass(x,fc,fs):
    nyq = fs/2
    b, a = signal.butter(4, fc/nyq, 'lowpass')
    y = signal.filtfilt(b, a, x)
    return y



# First grab the RMS monitor data

start_gps = 1117361776  # 10:16:00 Jun 3 2014
meas_duration = int(floor(3600*3.8))

chans = ['H1:OAF-BOUNCE_ITMX_RMS_OUTMON']

x = cdsutils.getdata(chans,meas_duration,start_gps)
Fs = x[0].sample_rate

bounce_RMS = lowpass(x[0].data,1,Fs)
t_RMS = arange(0,meas_duration,1.0/Fs)



# Now grab DARM_ERR in 10-second snippets

chans = ['H1:CAL-DELTAL_EXTERNAL_DQ']

# grab 10 seconds of data every 50 seconds, starting at 10:10 UTC and going until ???
step_duration = 200

duration = 20

t_DARM = arange(0,meas_duration,step_duration)
datapoints = zeros(len(t_DARM))

demod_I = zeros(len(t_DARM))
demod_Q = zeros(len(t_DARM))

for i in range(len(t_DARM)):

    gps_start = start_gps + i*step_duration

    #x = conn.fetch(gps_start,gps_start+gps_duration,chans)
    x = cdsutils.getdata(chans,duration,gps_start)
    Fs = x[0].sample_rate

    DARM = x[0].data
    #t_step = arange(0,duration,1.0/Fs)

    # frequency of ITMX bounce mode, +/- 0.001 Hz
    f0 = 9.848

    # decimation frequency
    Fs0 = 256

    # decimate data, multiply by large factor to bring near unity
    w = decimate(DARM, int(Fs/Fs0))*1e8

    t_dec = arange(0,duration,1.0/Fs0)

    # first bandpass the data around f0
    y = bandpass(w,9.0,13.0,Fs0)

    # demodulate and low-pass
    sinewave = sin(2*pi*f0*t_dec)
    cosinewave = cos(2*pi*f0*t_dec)

    Iphase = y*sinewave
    Qphase = y*cosinewave

    # only use 3 seconds to 9 seconds to avoid filter transients
    i1 = Fs0*3
    i2 = Fs0*(duration-1)

    Ipass = lowpass(Iphase,2.,Fs0)
    Qpass = lowpass(Qphase,2.,Fs0)

    data_demod = sqrt(Ipass**2 + Qpass**2)

    #datapoints[i] = max(abs(data_demod[i1:i2]))
    datapoints[i] = median(data_demod[i1:i2])




fignum=0

# Plot RMS and demod data together

fignum+=1
pylab.figure(fignum)

ax1 = pylab.subplot(1,1,1)

#plot with a 1-second pad to avoid filter transients
ax1.plot(t_RMS[16:-16],bounce_RMS[16:-16])

ax1.set_xlabel('time (s)')
ax1.set_ylabel('RMS [counts]', color='b')
for tl in ax1.get_yticklabels():
    tl.set_color('b')

ax2 = ax1.twinx()
ax2.plot(t_DARM, datapoints, 'ro')
ax2.set_ylabel('DEMOD [arb]', color='r')
for tl in ax2.get_yticklabels():
    tl.set_color('r')

pylab.grid(True)
pylab.title('ITMX Bounce mode decay, f0 = 9.848Hz')

pylab.savefig('bounceRMS.png')



# plot the last sample of demod data for debugging


fignum+=1
pylab.figure(fignum)

pylab.plot(t_dec[i1:i2],Ipass[i1:i2])
pylab.plot(t_dec[i1:i2],Qpass[i1:i2])

pylab.plot(t_dec[i1:i2],data_demod[i1:i2])
pylab.grid(True)

pylab.savefig('bounce_demod.png')




gs = gridspec.GridSpec(2,1,height_ratios=[3,1])

fignum+=1
pylab.figure(fignum)

# fit RMS data

t = t_DARM
data = datapoints
param_start = [0.8,0.57e6]
pfit, fitCov = curve_fit(decay,t,data,param_start)

print pfit

pylab.subplot(gs[0])

#plot with a 1-second pad to avoid filter transients
pylab.plot(t,data,'bo',markersize=6)
pylab.plot(t,decay(t,pfit[0],pfit[1]),'k--',linewidth=1.5)
pylab.ylabel('Demodulated bounce [arb]')
pylab.grid(True)
pylab.legend(('Demod data','Fit'))
pylab.title('Fit to Demod data')

pylab.subplot(gs[1])

pylab.plot(t,data-decay(t,pfit[0],pfit[1]),'bo',markersize=6)

pylab.ylabel('residuals')
pylab.xlabel('time [sec]')
pylab.grid(True)

pylab.savefig('bounceDemod_fit.png')





fignum+=1
pylab.figure(fignum)

# fit RMS data

t = t_RMS[16:-16]
data = bounce_RMS[16:-16]
param_start = [2800,0.57e6]
pfit, fitCov = curve_fit(decay,t,data,param_start)

print pfit

pylab.subplot(gs[0])

#plot with a 1-second pad to avoid filter transients
pylab.plot(t,data,'r.',markersize=3)
pylab.plot(t,decay(t,pfit[0],pfit[1]),'k--',linewidth=1.5)
pylab.ylabel('RMS Monitor [counts]')
pylab.grid(True)
pylab.legend(('RMS data','Fit'))
pylab.title('Fit to RMS data')

pylab.subplot(gs[1])

pylab.plot(t,data-decay(t,pfit[0],pfit[1]),'k.',markersize=3)

pylab.ylabel('residuals')
pylab.xlabel('time [sec]')
pylab.grid(True)

pylab.savefig('bounceRMS_fit.png')