#!/usr/bin/env python

from time import sleep
from collections import deque
import numpy as np
import scipy.signal
import epics

from zoetrope import aniplot as plt
from scipy.stats.stats import pearsonr
from bstrd import BS, bsstream

def goodshots(events, *arrays):
    #fel = events[:, 13]
    laser = events[:, 18]
    darkShot = events[:, 21]
    pumped_shots = np.logical_and.reduce((laser, np.logical_not(darkShot)))
    return [a[pumped_shots] for a in arrays]

plt.blit = False
plt.style.use('ggplot')

chname_diode = "SARES12-LSCP11-FNS:CH4:VAL_GET"
#chname_diode = "SLAAR-LSCP1-FNS:CH4:VAL_GET"
chname_events = "SAR-CVME-TIFALL4:EvtSet"
chname_laserRate = "SIN-TIMAST-TMA:Evt-23-Freq-SP"

# create channel
ch_diode  = BS(chname_diode)
ch_events = BS(chname_events)
ch_laserRate = epics.PV(chname_laserRate)


n = 1000
sigs = np.empty(n)
evts = np.empty((n, 256))

# create the empty plot
pd = plt.plot([0])

plt.suptitle("{}".format(chname_diode))
plt.fig.set_figheight(4)
plt.fig.set_figwidth(8)
plt.tight_layout()

for counter, data in zip(plt.show(), bsstream):
    print(counter)

    for i in range(n):
        sigs[i] = ch_diode.get()
        evts[i] = ch_events.get()
        next(bsstream) # this gets the next set of data

    laserRate = ch_laserRate.get()
    if not(laserRate):
        laserRate = 100

    goodsigs, _ = goodshots(evts, sigs, sigs)    
    x=np.arange(0,np.ceil(laserRate),np.ceil(laserRate)/len(goodsigs))
    vals = np.hstack((goodsigs-np.mean(goodsigs), np.zeros_like(goodsigs)))
    power = np.abs(np.fft.fft(vals))**2

    print(x[0])
    print(len(x))
    print(x[-1])
    print(len(power))

    plt.clf()
    plt.plot(x[0:200]/2, power[0:200])
    plt.show()

bsstream.close()