#!/usr/bin/env python

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

from zoetrope import aniplot as plt
from bstrd import BS, bsstream
plt.blit = False
plt.style.use('ggplot')

def unpumped(events, *arrays):
    laser = events[:, 18]
    darkShot = events[:, 21]
    background_shots = np.logical_and.reduce((laser, darkShot))
    return [a[background_shots] for a in arrays]

def pumped(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]

# config
chname_diode  = "SLAAR11-LSCP1-FNS:CH0:VAL_GET"
#chname_diode = "SARES11-GES1:CH1_VAL_GET"
chname_events = "SAR-CVME-TIFALL4:EvtSet"
chname_i0 = "SAROP11-PBPS110:INTENSITY"
length = 100

# create channel
ch_diode  = BS(chname_diode)
ch_events = BS(chname_events)
ch_i0 = BS(chname_i0)

n = 50 
sigs = np.empty(n)
evts = np.empty((n, 256))
i0s = np.empty(n)

# create a buffer for the plotting
pp_sigs = deque(maxlen=length)

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

# some plot settings
plt.suptitle(chname_diode)
plt.fig.set_figheight(5)
plt.fig.set_figwidth(15)
plt.tight_layout()

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

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

    sigs_p = pumped(evts, sigs)
    sigs_u = unpumped(evts, sigs)
    i0s_p = pumped(evts, i0s)
    i0s_u = unpumped(evts, i0s)

    sig_p = np.mean(np.asarray(sigs_p))#/np.asarray(i0s_p))
    sig_u = np.mean(np.asarray(sigs_u))#/np.asarray(i0s_u))

    sig = -np.log10(sig_p/sig_u)

    pp_sigs.append(sig)

    xs = np.arange(len(pp_sigs))
    pd.set(xs, pp_sigs)

    # this, I need to move into the library
    pd.ax.relim()
    pd.ax.autoscale_view()

bsstream.close()