#!/opt/gfa/python-3.5/latest/bin/python 

import sys
import datetime
import os.path
import time
from os import walk
import numpy as np
import sys
import re


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

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


 

class OpticsPlotting:

    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 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)




    def plot(self,data,filt,z0,z1):

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

        s=np.array(data['S'])
        if z0>z1:
            tmp=z1
            z1=z0
            z0=tmp

        i1=np.argmin(np.abs(s-z0))
        i2=np.argmin(np.abs(s-z1))

        ylabel=r''
        if filt['BETX']:
            self.plotSingle(s[i1:i2],data['BETX'][i1:i2],(0,0,1,1),r'$\beta_{x}$')
            ylabel=ylabel+r'$\beta_x$ (m), '
        if filt['BETY']:
            self.plotSingle(s[i1:i2],data['BETY'][i1:i2],(1,0,0,1),r'$\beta_{y}$')
            ylabel=ylabel+r'$\beta_y$ (m), '
        if filt['ALFX']:
            self.plotSingle(s[i1:i2],data['ALFX'][i1:i2],(0,0,1,1),r'$\alpha_{x}$')
            ylabel=ylabel+r'$\alpha_x$ (rad), '
        if filt['ALFY']:
            self.plotSingle(s[i1:i2],data['ALFY'][i1:i2],(1,0,0,1),r'$\alpha_{y}$')
            ylabel=ylabel+r'$\alpha_y$ (rad), '
        if filt['DX']:
            self.plotSingle(s[i1:i2],data['DX'][i1:i2],(0,0,1,1),r'$\eta_{x}$')
            ylabel=ylabel+r'$\eta_x$ (m), '
        if filt['DY']:
            self.plotSingle(s[i1:i2],data['DY'][i1:i2],(1,0,0,1),r'$\eta_{y}$')
            ylabel=ylabel+r'$\eta_y$ (m), '
        if filt['RE56']:
            self.plotSingle(s[i1:i2],data['RE56'][i1:i2],(0,0,0,1),r'$R_{56}$')
            ylabel=ylabel+r'$R_{56}$ (m), '
        if filt['Energy']:
            self.plotSingle(s[i1:i2],data['Energy'][i1:i2],(0,1,0,1),r'$E$')
            ylabel=ylabel+r'$E$ (MeV), '


        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,data['NAME'])
        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:
                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:
                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:
                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 'ACC' in name or 'TDS' 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:
                if splitquads == True:
                    if 'END' 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 'START' 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




#