import sys
import pysdds
from PyQt5 import QtWidgets,QtGui


from matplotlib.figure import Figure
import matplotlib.patches as patches

from matplotlib.backends.backend_qt5agg import (
    FigureCanvasQTAgg as FigureCanvas,
    NavigationToolbar2QT as NavigationToolbar)

from ui.ElegantPlotGui import Ui_ElegantPlotGUI
import numpy as np
import h5py



class ElegantPlot(QtWidgets.QMainWindow, Ui_ElegantPlotGUI):
    def __init__(self,parent=None):
        super(ElegantPlot, self).__init__()
        self.setupUi(self)
        self.parent=parent
        self.version = '1.0.1'
        self.data = None
        self.setWindowIcon(QtGui.QIcon("rsc/Audrey_Icon.png"))
        self.setWindowTitle("Elegant Results Plotting Window")

        self.PBetax.toggled.connect(self.doplot)
        self.PBetay.toggled.connect(self.doplot)
        self.PSizex.toggled.connect(self.doplot)
        self.PSizey.toggled.connect(self.doplot)
        self.PEmitx.toggled.connect(self.doplot)
        self.PEmity.toggled.connect(self.doplot)
        self.PDuration.toggled.connect(self.doplot)
        self.PEnergy.toggled.connect(self.doplot)
        self.PSpread.toggled.connect(self.doplot)
        self.PCEnergy.toggled.connect(self.doplot)
        self.PCx.toggled.connect(self.doplot)
        self.PCy.toggled.connect(self.doplot)

        self.PDistTE.clicked.connect(self.doPSplot)
        self.PDistXPx.clicked.connect(self.doPSplot)
        self.PDistYPy.clicked.connect(self.doPSplot)
        self.PDistTX.clicked.connect(self.doPSplot)
        self.PDistTY.clicked.connect(self.doPSplot)
        self.PDistXY.clicked.connect(self.doPSplot)
        self.currentDist=None
        self.dist={}
        self.Q=0
        self.initmpl(self.mplvl, self.mplwindow)



    def initmpl(self,mplvl,mplwindow):
        self.fig=Figure()
        self.axes=self.fig.add_subplot(111)
        self.axes2 = self.axes.twinx()
        self.canvas = FigureCanvas(self.fig)
        mplvl.addWidget(self.canvas)
        self.canvas.draw()
        self.toolbar=NavigationToolbar(self.canvas,mplwindow, coordinates=True)
        mplvl.addWidget(self.toolbar)

    def newData(self,data,dists=[]):
        if not self.isVisible():
            self.show()
        print('Updating Plotting Data')
        self.data=data
        # scale some data
        self.data['ecnx']=[ele*1e9 for ele in  self.data['ecnx']] # nm
        self.data['ecny'] = [ele * 1e9 for ele in self.data['ecny']]  # nm
        self.data['Sx'] = [ele * 1e6 for ele in self.data['Sx']]  # micrometer
        self.data['Sy'] = [ele * 1e6 for ele in self.data['Sy']]  # micrometer
        self.data['Cx'] = [ele * 1e6 for ele in self.data['Cx']]  # micrometer
        self.data['Cy'] = [ele * 1e6 for ele in self.data['Cy']]  # micrometer
        self.data['pCentral'] = [ele * 0.511 for ele in self.data['pCentral']]  # MeV
        self.data['St'] = [ele * 1e15 for ele in self.data['St']]  # fs
        self.data['Sdelta'] = [ele * 1e2 for ele in self.data['Sdelta']]  # %
        self.data['Cdelta'] = [ele * 1e2 for ele in self.data['Cdelta']]  # %

        self.PDistList.clear()
        for dist in dists:
            self.PDistList.addItem(dist)
        self.currentDist=None
        self.updateOpticsTable()
        self.doplot()

    def updateOpticsTable(self):
        self.UITwissValues.clear()
        if self.data is None:
            return

        nrow = len(self.data['s'])
        ncol = 14
        self.UITwissValues.setColumnCount(ncol)
        self.UITwissValues.setRowCount(nrow)
        self.UITwissValues.setHorizontalHeaderLabels(['Element','s','Size in x','Size in y','Size in t','Energy','Energy Spread','Emit x','Emit y','Betax','Betay','<x>','<y>','<E>'])


        for i in range(nrow):
            self.UITwissValues.setItem(i, 0, QtWidgets.QTableWidgetItem(self.data['ElementName'][i]))
            self.UITwissValues.setItem(i, 1, QtWidgets.QTableWidgetItem('%10.6f' % self.data['s'][i]))
            self.UITwissValues.setItem(i, 2, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Sx'][i]))
            self.UITwissValues.setItem(i, 3, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Sy'][i]))
            self.UITwissValues.setItem(i, 4, QtWidgets.QTableWidgetItem('%10.6f' % self.data['St'][i]))
            self.UITwissValues.setItem(i, 5, QtWidgets.QTableWidgetItem('%10.6f' % self.data['pCentral'][i]))
            self.UITwissValues.setItem(i, 6, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Sdelta'][i]))
            self.UITwissValues.setItem(i, 7, QtWidgets.QTableWidgetItem('%10.6f' % self.data['ecnx'][i]))
            self.UITwissValues.setItem(i, 8, QtWidgets.QTableWidgetItem('%10.6f' % self.data['ecny'][i]))
            self.UITwissValues.setItem(i, 9, QtWidgets.QTableWidgetItem('%10.6f' % self.data['betaxBeam'][i]))
            self.UITwissValues.setItem(i, 10, QtWidgets.QTableWidgetItem('%10.6f' % self.data['betayBeam'][i]))
            self.UITwissValues.setItem(i, 11, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Cx'][i]))
            self.UITwissValues.setItem(i, 12, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Cy'][i]))
            self.UITwissValues.setItem(i, 13, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Cdelta'][i]))
        self.UITwissValues.resizeColumnsToContents()
        self.UITwissValues.verticalHeader().hide()
        self.UITwissValues.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)

    def loadDist(self,filename):
        dist = pysdds.read(filename)
        cols=['t','p','x','xp','y','yp']
        self.dist.clear()
        for col in cols:
            self.dist[col] = dist.col(col)[0]
        self.Q = dist.par('Charge').data[0]
        return filename



    def doPSplot(self):
        dist = self.PDistList.currentText()
        if not dist == self.currentDist:
            self.currentDist = self.loadDist(dist)
        but = self.sender().objectName()

        if 'DistTE' in but:
            self.LPS(self.dist['t'],self.dist['p'],self.Q,dist)
        elif 'DistXPx' in but:
            self.TPS(self.dist['x'],self.dist['xp'],True,dist)
        elif 'DistYPy' in but:
            self.TPS(self.dist['y'], self.dist['yp'], False,dist)
        elif 'DistTX' in but:
            self.TimePS(self.dist['t'], self.dist['x'], self.Q,True,dist)
        elif 'DistTY' in but:
            self.TimePS(self.dist['t'], self.dist['y'], self.Q,False,dist)
        elif 'DistXY' in but:
            self.TS(self.dist['x'], self.dist['y'],dist)

    def LPS(self,t,p,Q,title=None):
        if title is None:
            title = self.parent.UIDistList.currentText()
        x =  np.array(t).ravel()*1e15  # to fs
        y =  np.array(p).ravel()*0.511 # to MEV
        x=x-np.mean(x)
        self.PS(x,y,Q=Q,title = title, labx=r'$t$ (fs)',laby=r'$E$ (MeV)',laby2=r'$I$ (A)',legend1='Energy Distribution',legend2='Current')


    def TPS(self,xin,pin,isX=True,title=None):
        if title is None:
            title = self.parent.UIDistList.currentText()
        x =  np.array(xin).ravel()*1e3  # to mm
        y =  np.array(pin).ravel()*1e3 # to mrad
        labx=r'$y$ (mm)'
        laby=r'$y\prime$ (mrad)'
        if isX:
            labx = r'$x$ (mm)'
            laby = r'$x\prime$ (mrad)'
        self.PS(x, y, Q=None, title=title, labx=labx, laby=laby, laby2='',legend1='Distribution in Momentum', legend2='Distribution in Space')

    def TimePS(self,t,xin,Q,isX=True,title=None):
        if title is None:
            title = self.parent.UIDistList.currentText()
        x =  np.array(t).ravel()*1e15  # to fs
        y =  np.array(xin).ravel()*1e3 # to mm
        x=x-np.mean(x)
        labx=r'$t$ (fs)'
        laby=r'$y$ (mm)'
        if isX:
            laby = r'$x$ (m)'
        self.PS(x, y, Q=Q, title=title, labx=labx, laby=laby, laby2=r'$I$ (A)',legend1='Distribution in Space', legend2='Current')

    def TS(self,xin,yin,title=None):
        if title is None:
            title = self.parent.UIDistList.currentText()
        x =  np.array(xin).ravel()*1e3  # to mm
        y =  np.array(yin).ravel()*1e3 # to mm
        labx=r'$x$ (mm)'
        laby=r'$y$ (mm)'
        self.PS(x, y, Q=None, title=title, labx=labx, laby=laby, laby2='',legend1='Distribution in Y', legend2='Distribution in X')


    def PS(self,x,y,Q = None,title='',labx='',laby='',laby2='',legend1='',legend2=''):

        xmin = x.min()
        xmax = x.max()
        ymin = y.min()
        ymax = y.max()
        dx=(xmax-xmin)*0.1
        dy=(ymax-ymin)*0.1
        xmax+=dx
        xmin-=dx
        ymax+=dy
        ymin-=dy
        range=[[xmin,xmax],[ymin,ymax]]
        N = 300
        img,xed,yed = np.histogram2d(x,y,N,range=range)
        xdist, xval =  np.histogram(x,N,range=range[0])
        if Q:
            cal = np.sum(xdist)*(xval[1]-xval[0])*1e-12/Q
            xdist = xdist/cal*1e3

        ydist, yval = np.histogram(y, N,range=range[1])
        ydist = ydist / np.max(ydist) * 0.3 * (xmax - xmin) + xmin

        self.axes.clear()
        self.axes2.clear()
        self.axes.imshow(np.transpose(np.fliplr(img)), aspect='auto', cmap='gnuplot2',interpolation='bicubic',  # was viridis andf
                         extent=[xmin, xmax, ymin, ymax])
        self.axes2.plot(xval[:-1],xdist,'y')
        self.axes.plot(ydist,yval[:-1], 'c', label=legend1)
        self.axes.plot([], [], 'y', label=legend2)
        self.axes.legend()

        self.axes.set_xlabel(labx)
        self.axes.set_ylabel(laby)
        if len(laby2)> 1:
            self.axes2.yaxis.set_visible(True)
        else:
            self.axes2.yaxis.set_visible(False)
        self.axes2.set_ylim([0,3*np.max(xdist)])
        self.axes2.set_xlim([xmin,xmax])
        self.axes2.set_ylabel(laby2)
        self.axes.set_title(title)
        self.canvas.draw()


    def doplot(self):
        if self.data is None:
            return

        z0=float(str(self.PStart.text()))
        z1=float(str(self.PEnd.text()))
        if z0 > z1:
            tmp = z1
            z1 = z0
            z0 = tmp

        filt={}
        filt['Sx']=self.PSizex.isChecked()
        filt['Sy'] = self.PSizey.isChecked()
        filt['ecnx']=self.PEmitx.isChecked()
        filt['ecny']=self.PEmity.isChecked()
        filt['betaxBeam']=self.PBetax.isChecked()
        filt['betayBeam']=self.PBetay.isChecked()
        filt['Cx']=self.PCx.isChecked()
        filt['Cy']=self.PCy.isChecked()
        filt['St']=self.PDuration.isChecked()
        filt['pCentral']=self.PEnergy.isChecked()
        filt['Sdelta'] = self.PSpread.isChecked()
        filt['Cdelta'] = self.PCEnergy.isChecked()
        color={}
        color['Sx']={'color':(0,0,1,1),'label':r'$\sigma_x$ ($\mu$m)','legend':r'$\sigma_x$','dash':False}
        color['Sy'] = {'color':(1, 0, 0, 1),'label':r'$\sigma_y$ ($\mu$m)','legend':r'$\sigma_y$','dash':False}
        color['ecnx'] ={'color': (0, 0.4, 1, 1),'label':r'$\epsilon_{nx}$ (nm)','legend':r'$\epsilon_{nx}$','dash':False}
        color['ecny'] = {'color':(1, 0.4, 0, 1),'label':r'$\epsilon_{ny}$ (nm)','legend':r'$\epsilon_{ny}$','dash':False}
        color['betaxBeam'] = {'color': (0, 0, 1, 1), 'label': r'$\beta_x$ (m)', 'legend': r'$\beta_x$','dash':False}
        color['betayBeam'] = {'color': (1, 0, 0, 1), 'label': r'$\beta_y$ (m)', 'legend': r'$\beta_y$','dash':False}
        color['Cx'] = {'color': (0, 0, 1, 1), 'label': r'$<x>$ ($\mu$m)', 'legend': r'$<x>$','dash':True}
        color['Cy'] = {'color': (1, 0, 0, 1), 'label': r'$<y>$ ($\mu$m)', 'legend': r'$<y>$','dash':True}
        color['pCentral'] = {'color':(0,0.5,0,1),'label':r'$E_0$ (MeV)', 'legend': r'$E_0$','dash':False}
        color['Sdelta'] = {'color': (0, 0.5, 0, 1), 'label': r'$\sigma_\delta$ (%)', 'legend': r'$\sigma_\delta$', 'dash': False}
        color['Cdelta'] = {'color': (0, 0.5, 0, 1), 'label': r'$<\delta>$ (%)', 'legend': r'$<\delta>$', 'dash': True}
        color['St'] = {'color': (0.5, 0, 0.5, 1), 'label': r'$\sigma_t$ (fs)', 'legend': r'$\sigma_t$', 'dash': False}

        s = np.array(self.data['s'])
        i1 = np.argmin(np.abs(s - z0))
        i2 = np.argmin(np.abs(s - z1))

        self.axes.clear()
        self.axes2.clear()

        ylabel = r''
        for key in filt.keys():
            if filt[key]:
                if key in color.keys():
                    self.plotSingle(s[i1:i2],self.data[key][i1:i2], color[key]['color'],color[key]['legend'],dashed=color[key]['dash'])
                    ylabel += color[key]['label']+', '

        if len(ylabel) < 3:
            self.canvas.draw()
            return

        self.axes.legend(bbox_to_anchor=(0.15, 0.85))
        self.axes.set_xlabel('s (m)')
        self.axes.set_ylabel(ylabel[0:-2])

        self.plotLayout(s,self.data['ElementName'])
        self.axes.set_xlim([s[i1], s[i2]])
        ylim = self.axes.get_ylim()
        dl = np.abs(ylim[1] - ylim[0])
        yl = [ylim[0], ylim[1] + 0.2 * dl]
        self.axes.set_ylim(yl)
        self.axes2.set_xlim([s[i1], s[i2]])
        self.canvas.draw()

        return

    def plotLayout(self, s, elements):
        splitquads = False
        sstart = 0

        s1 = [np.min(s), np.max(s)]
        s2 = [0.9, 0.9]
        self.axes2.plot(s1, s2, 'k')
        for i, name in enumerate(elements):
            if '.MBND' in name and '.B' in name:
                s1 = s[i - 1]
                s2 = s[i]
                self.axes2.add_patch(patches.Rectangle((s1, 0.9), (s2 - s1), 0.03, facecolor='blue', edgecolor="none"))
            if ',MSEX' in name and '.SEX':
                s1 = s[i - 1]
                s2 = s[i]
                self.axes2.add_patch(
                    patches.Rectangle((s1, 0.87), (s2 - s1), 0.06, facecolor='green', edgecolor="none"))

            if '.UIND' in name and '.WIG' in name:
                s1 = s[i - 1]
                s2 = s[i]
                self.axes2.add_patch(
                    patches.Rectangle((s1, 0.88), (s2 - s1), 0.04, facecolor='purple', edgecolor="none"))

            if '.RACC' in name or '.RTDS' in name:
                s1 = s[i - 1]
                s2 = s[i]
                self.axes2.add_patch(patches.Rectangle((s1, 0.89), (s2 - s1), 0.02, facecolor='cyan', edgecolor="none"))

            if '.MQUA' in name and '.Q':
                if splitquads == True:
                    if '.Q2' in name:
                        s1 = sstart
                        s2 = s[i]
                        self.axes2.add_patch(
                            patches.Rectangle((s1, 0.85), (s2 - s1), 0.1, facecolor='red', edgecolor="none"))
                        splitquads = False
                else:
                    if '.Q1' in name:
                        splitquads = True
                        sstart = s[i]
                    else:
                        s1 = s[i - 1]
                        s2 = s[i]
                        self.axes2.add_patch(
                            patches.Rectangle((s1, 0.85), (s2 - s1), 0.1, facecolor='red', edgecolor="none"))

        self.axes2.set_ylim([0, 1])
        self.axes2.yaxis.set_visible(False)
        return


    def plotSingle(self, x, y, color, legend, dashed=False):

        if dashed:
            self.axes.plot(x, y, '--', color=color, label=legend)
        else:
            self.axes.plot(x, y, color=color, label=legend)




    # --------------------------------
    # Main routine


if __name__ == '__main__':
    QtWidgets.QApplication.setStyle(QtWidgets.QStyleFactory.create("plastique"))
    app = QtWidgets.QApplication(sys.argv)
    if len(sys.argv) > 1:
        arg=int(sys.argv[1])
    else:
        arg=0
    plot = ElegantPlot()
    plot.show()
    sys.exit(app.exec_())