#!/usr/bin/env python

from time import sleep
from collections import deque
import numpy as np
import scipy.signal
import epics
from epics import caput
from zoetrope import aniplot as plt
from bstrd import BS, bsstream

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

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

# config
STAGEpv = epics.PV('SLAAR11-LMOT-M452:MOTOR_1.VAL') # global globi
chname_amplitude  = "SARES11-SPEC125-M1.edge_amplitude"
chname_jitter = "SARES11-SPEC125-M1.edge_position"
chname_events = "SAR-CVME-TIFALL4:EvtSet"
chname_iZero = "SAROP11-PBPS110:INTENSITY"
length = 500
mm2fs = 6671.2 # lightspeed and delay stages

# create channel
ch_amp  = BS(chname_amplitude)
ch_jitter = BS(chname_jitter)
ch_events = BS(chname_events)
ch_iZero = BS(chname_iZero)

n = 100
amp = np.empty(n)
jitter = np.empty(n)
iZero = np.empty(n)
evts = np.empty((n, 256))

# create a buffer for the plotting
jitter_vals = deque(maxlen=length)
jitter_stds = deque(maxlen=length)
amplitude_vals = deque(maxlen=length)
iZero_vals = deque(maxlen=length)

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

# some plot settings
plt.suptitle("arrival time, fs")
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()
        jitter[i] = ch_jitter.get()
        amp[i] = ch_amp.get()
        iZero[i] = ch_iZero.get()
        next(bsstream) # this gets the next set of data

    edge_amp, edge_pos, i_zero = goodshots(evts, amp, jitter, iZero)

    jitter_avg = np.mean(edge_pos[edge_amp > 10])
    jitter_std = np.std(edge_pos[edge_amp > 10])
    jitter_vals.append(jitter_avg)
    jitter_stds.append(jitter_std)
    
    edge_amp_avg = np.mean(edge_amp)
    amplitude_vals.append(edge_amp_avg)

    xs = np.arange(len(jitter_vals))
    pd.set(xs, amplitude_vals)

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

#    if (np.abs(jitter_avg) > 100) & (np.abs(jitter_avg) < 500):
#        moveTo = STAGEpv.get()*mm2fs - jitter_avg
#        STAGEpv.put(moveTo/mm2fs)
#        print('Stage moved.')

bsstream.close()