#!/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 scipy.stats.stats import pearsonr
from bstrd import BS, bsstream
plt.blit = False
plt.style.use('ggplot')

def pumpedshots(events, *arrays):
    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]


chname_tt1 = "SARES11-SPEC125-M1.edge_amplitude"
chname_tt2 = "SARES11-SPEC125-M1.edge_amplitude2"
chname_i0 = "SAROP11-PBPS110:INTENSITY"
chname_events = "SAR-CVME-TIFALL4:EvtSet"
length = 500

# create channel
ch_tt1  = BS(chname_tt1)
ch_tt2 = BS(chname_tt2)
ch_i0 = BS(chname_i0)
ch_events = BS(chname_events)

n = 200
tt1 = np.empty(n)
tt2 = np.empty(n)
i0s = np.empty(n)
evts = np.empty((n, 256))


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

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

# some plot settings
plt.suptitle("{}, {}".format(chname_tt1, chname_tt2))
plt.fig.set_figheight(5)
plt.fig.set_figwidth(5)
plt.tight_layout()

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

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


    pumptt1, pumptt2 = pumpedshots(evts, tt1, tt2)
    pearscoeff1, _ = pearsonr(pumptt1, pumptt2)

    #xs = np.arange(len(pp_sigs))
    plt.clf()
    plt.scatter(pumptt1, pumptt2)
    plt.title("{}".format(round(pearscoeff1, 4)))
    plt.show()

bsstream.close()