#!/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
from datetime import datetime

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
threshold = 7 


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

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

while True:
#    print(counter)

    timeofdata = datetime.now()

    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 > threshold])
    jitter_std = np.std(edge_pos[edge_amp > threshold])

    if len(edge_pos[edge_amp > threshold] > 125):

        print("{} statistics on {} shots: offset of {} fs, jitter of {} fs rms".format(timeofdata.strftime('%Y-%m-%d-%H:%M:%S'), len(edge_pos[edge_amp > threshold]), round(jitter_avg, 2), round(jitter_std, 2)))

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

    else:
        print ("{} Where are my X-rays?!".format(timeofdata.strftime('%Y-%m-%d-%H:%M:%S')))

bsstream.close()