import numpy as np
import epics
import time

from collections import deque
import numpy as np
import time
import pickle
import os
import epics
import matplotlib.pyplot as plt
from slic.gui.widgets import LabeledMathEntry 
from slic.gui.widgets import make_filled_hbox, make_filled_vbox, EXPANDING
from slic.gui.widgets.plotting import PlotPanel




from bstrd import BS, bsstream

from epics import caget, caput, cainfo





channel_TOF = "SATES21-GES1:A1_VALUES"
ch_TOF = BS(channel_TOF)
temp_TOF = ch_TOF.get()


monitor_ch = 'SATES20-CVME-EVR0:Pul11_NEW_DELAY'


current_ch = caget(monitor_ch)

#### Set scan values:

start = 4200#uS
end = 5000#uS

wait_time = 1#seconds

step_size = 10
num_shots = 500

scan_axis = np.arange(start, end+step_size, step_size)

TOF_map     = np.zeros((len(scan_axis), len(temp_TOF)), dtype=float)

temp_TOF = np.zeros_like(temp_TOF,dtype=float)
next(bsstream)

plt.figure()
plt.subplot(1,2,1)

for s,i in enumerate(scan_axis):
    temp_TOF = np.zeros_like(temp_TOF)

    next(bsstream)
    caput(monitor_ch, i)
    time.sleep(wait_time)
    print("Moved %s to %0.2f"%(monitor_ch, i) )
    shot_counter = 0
    while shot_counter <num_shots:
        next(bsstream)
        temp_TOF += ch_TOF.get()/(1.0*num_shots)
        shot_counter += 1
    if s ==0:
        bkg = temp_TOF.copy()
    TOF_map[s,:] = temp_TOF
#    TOF_map[s,:] = temp_TOF - bkg
    plt.plot(temp_TOF, label = i)
plt.legend()




plt.subplot(1,2,2)
plt.pcolormesh(np.arange(0, len(temp_TOF)),scan_axis,  TOF_map)

plt.show()

dic = {
    'map': TOF_map,
    'scan axis': scan_axis,
}


with open(r"test.pickle", "wb") as output_file:
    pickle.dump(dic, output_file)