from collections import deque
import random
import wx
import numpy as np
import epics

from datetime import datetime
from scipy.stats.stats import pearsonr

#from epics.wx import MotorPanel

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

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

# config
nshots = 100
qlength = 10000
histlength = 1
i0thres = 0.2

# channels
chname_tt1    = "SARES11-SPEC125-M1.edge_position"
chname_tt2    = "SARES11-SPEC125-M1.edge_position2"
chname_i0     = "SAROP11-PBPS122:INTENSITY"
chname_events = "SAR-CVME-TIFALL4:EvtSet"

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

iso_format = "%Y-%m-%d %H:%M:%S"

class MainPanel(wx.Panel):

    def __init__(self, parent):
        super().__init__(parent)

        self.evts = np.empty((nshots,256))
        self.tt1  = np.empty(nshots)
        self.tt2  = np.empty(nshots)
        self.i0  = np.empty(nshots)

        self.tt1hist  = deque(maxlen=histlength)
        self.tt2hist  = deque(maxlen=histlength)
        self.tt1s = deque(maxlen=qlength)
        self.tt2s = deque(maxlen=qlength)
        self.ttdiff   = deque(maxlen=qlength)
        
        self.plot_hist_tt1  = plot_hist_tt1  = PlotPanel(self, figsize=(2,3))
        self.plot_hist_tt2  = plot_hist_tt2  = PlotPanel(self, figsize=(2,3))	
        self.plot_corr = plot_corr = PlotPanel(self, figsize=(2,3))
        self.plot_diff = plot_diff = PlotPanel(self, figsize=(2,3))

        plots1 = (plot_hist_tt1,  plot_hist_tt2)
        plots2 = (plot_corr,)
        plots3 = (plot_diff,)
        
        hb_plot1 = make_filled_hbox(plots1)
        hb_plot2 = make_filled_hbox(plots2)
        hb_plot3 = make_filled_hbox(plots3)

        btn_clearQ = wx.Button(self, label="Clear plots")
        btns = (btn_clearQ,)
        hb_btns = make_filled_hbox(btns)

        widgets = (hb_plot1, hb_plot2, hb_plot3, hb_btns)
        box = make_filled_vbox(widgets, border=1)
        self.SetSizerAndFit(box)

        btn_clearQ.Bind(wx.EVT_BUTTON, self.on_click_clearQ)

        self.timer = wx.Timer(self)
        self.Bind(wx.EVT_TIMER, self.on_update, self.timer)
        self.timer.Start(100)


    def on_update(self, _event):

        self.time = datetime.now().strftime(iso_format)

        self.evts = np.empty((nshots,256))
        self.tt1  = np.empty(nshots)
        self.tt2  = np.empty(nshots)
        self.int  = np.empty(nshots)

        for i in range(nshots):
            tempa = ch_events.get()
            self.evts[i] = tempa
            tempb = ch_i0.get()
            self.i0[i] = tempb
            tempc = ch_tt1.get()
            self.tt1[i] = tempc
            tempd = ch_tt2.get()
            self.tt2[i] = tempd
            next(bsstream)
            wx.GetApp().Yield()

        self.tt1p = np.ravel(pumped(self.evts, self.i0, self.tt1))
        self.tt2p = np.ravel(pumped(self.evts, self.i0, self.tt2))

        if (len(self.tt1p) > 1):
            print ('{} -- {} = {:.2f} a.u.'.format(self.time, chname_i0, np.nanmean(self.i0)))

            self.tt1hist.append(self.tt1p)
            self.tt2hist.append(self.tt2p)
            self.tt1s.append(np.mean(self.tt1p))
            self.tt2s.append(np.mean(self.tt2p))
            #self.ttdiff.append(np.array(self.tt1s) - np.array(self.tt2s))
            self.ttdiff.append(np.nanmean(np.array(self.tt1p) - np.array(self.tt2p)))

        else:    
            print ('{} -- {} = {:.2f} a.u. --> No Xrays'.format(self.time, chname_i0, np.nanmean(self.int)))    
            self.tt1hist.append(np.nan)
            self.tt2hist.append(np.nan)
            self.tt1s.append(np.nan)
            self.tt2s.append(np.nan)
            self.ttdiff.append(np.nan)

        self.draw_plot()

    def draw_plot(self):
        self.plot_hist_tt1.clear()
        self.plot_hist_tt2.clear()
        self.plot_corr.clear()
        self.plot_diff.clear()

        self.plot_hist_tt1.set_xlabel('Arr time tt (fs)')
        self.plot_hist_tt2.set_xlabel('Arr time tt2 (fs)')
        self.plot_corr.set_xlabel('Arr time tt1 (fs)')
        self.plot_corr.set_ylabel('Arr time tt2 (fs)')
        self.plot_diff.set_xlabel('time ago, a.u.')
        self.plot_diff.set_ylabel('Difference (fs)')
        
        if (len(self.tt1p) > 1):        
            self.plot_hist_tt1.set_title('tt1: {:.2f} fs, jitter: {:.2f}'.format(np.nanmean(self.tt1p), np.nanstd(self.tt1p)))
            self.plot_hist_tt2.set_title('tt2: {:.2f} fs, jitter: {:.2f}'.format(np.nanmean(self.tt2p), np.nanstd(self.tt2p)))

            self.plot_hist_tt1.hist(np.ravel(self.tt1hist), facecolor='dodgerblue', edgecolor='black')
            self.plot_hist_tt2.hist(np.ravel(self.tt2hist), facecolor='green', edgecolor='black')

            self.plot_corr.set_title('Corr: {:.4f}'.format(pearsonr(np.ravel(self.tt1p), np.ravel(self.tt2p))[0]))
            self.plot_corr.plot(np.ravel(self.tt1p), np.ravel(self.tt2p), '.', color='purple')
            self.plot_corr.grid()

            #self.plot_diff.plot(np.arange(0, len(self.tt1s)), np.array(self.tt1p) - np.array(self.tt2p), '.', color='orangered')
            #self.plot_diff.plot(np.array(self.tt1p) - np.array(self.tt2p), '.', color='orangered')
        
            self.plot_diff.set_title('Diff: {:.2f} fs'.format(self.ttdiff[-1]))
            self.plot_diff.plot(np.arange(0, len(self.ttdiff)), np.array(self.ttdiff), '.', color='orangered')
        
        self.plot_hist_tt1.draw()
        self.plot_hist_tt2.draw()
        self.plot_corr.draw()
        self.plot_diff.draw()


    def on_click_clearQ(self, _event):
        
        self.tt1s.clear()
        self.tt2s.clear()
        self.ttdiff.clear()