from PyQt5 import QtCore, QtGui, QtWidgets, uic
#import guidata
#from  guidata.dataset import datatypes
#from guidata.dataset import dataitems
import numpy
import pySAXS
from pySAXS.LS import LSusaxs
from pySAXS.guisaxs.dataset import *
from pySAXS.guisaxs.qt import plugin

classlist=['USAXSTransmission','USAXSDesmearing'] #need to be specified

class USAXSTransmission(plugin.pySAXSplugin):
    menu="Data Treatment"
    subMenu="USAXS"
    subMenuText="Background and Data correction"
    icon="hist.png"
    wavelength=1.54189
        
    def execute(self):
        '''
        The user click on "USAXS - Data correction"
        '''
        if self.selectedData is None:
            QtWidgets.QMessageBox.information(self.parent,"pySAXS", "No data are selected", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.NoButton)
            return
        
        '''if not('rock' in self.data_dict):#not(self.data_dict.has_key('rock')):
            QtWidgets.QMessageBox.information(self.parent, "pySAXS","Load data and rocking curve before correcting by transmission \n\
            (Try to rename the rocking curve datas in 'rock')",buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.NoButton)
            return
        '''
        self.dlg = dlgUsaxs(self.selectedData, self.parent)
        '''
        #generate gui using guidata
        #here we use guidata to generate a dialog box
        rocklist=['rock']
        rocklist+=self.ListOfDatasChecked()
        
        items = {
         "wavelength": dataitems.FloatItem("Wavelength",self.wavelength),
         #"thickness": dataitems.FloatItem("Thickness of sample (cm)",0.1),
         "backgroundData": dataitems.FloatItem("Background data :",0.0,unit='cps/s'),
         "backgroundRC": dataitems.FloatItem("Background Rocking curve :",0.0,unit='cps/s'),
         "shiftdata": dataitems.FloatItem("Shift Data :",0.0,unit="steps"),
         "shiftrock": dataitems.FloatItem("Shift Rocking curve :",0.0,unit="steps"),
         "qmin":dataitems.FloatItem("q min after substraction :",3e-4),
         "keepdatacentered":dataitems.BoolItem("Keep intermediary centered datas",False),
         "rock":dataitems.ChoiceItem("Rocking curve :",rocklist),
         "desmearing":dataitems.BoolItem("Continue with desmearing ?",False),
         }
        clz = type("Transmission and Data correction :", (datatypes.DataSet,), items)
        self.form = clz()
        if self.form.edit():
            #ok
            self.calculate()
        '''
            


class dlgUsaxs(QtWidgets.QDialog):
    def __init__(self, selectedData,parent):
        self.selectedData=selectedData
        self.parent=parent
        if parent is not None:
            self.printTXT = parent.printTXT
            self.workingdirectory = parent.workingdirectory
        '''
        datas=self.parent.data_dict[self.selectedData]
        q=numpy.array(datas.q)#*10'''
        self.data_dict=self.parent.data_dict
        QtWidgets.QDialog.__init__(self, parent)
        self.ui = uic.loadUi(pySAXS.UI_PATH + "dlgUsaxs.ui", self)
        self.icon = QtGui.QIcon(pySAXS.ICON_PATH + "hist.png")
        self.setWindowIcon(self.icon)
        self.ui.buttonBox.clicked.connect(self.click)
        #fill the list of rock
        #self.ui.comboRock.addItem(str('rock'))

        #checkedElements = dataset.getListOfChecked(self.parentwindow.data_dict)
        for el in self.parent.ListOfDatasChecked():
            self.ui.comboRock.addItem(str(el))
        self.ui.show()

    def click(self,obj=None):
        name=obj.text()
        #print(str(name))
        if name=="OK":
            self.close()
        elif name=="Apply":
            self.calculate()
        elif name=="Close":
            self.close()

    def calculate(self):
        '''
        do the transmission correction after dialog box
        '''
        # --- do the correction
        self.printTXT("-----USAXS Data correction --------")
        self.printTXT("Thickness is fixed at 1")
        self.thick = 1
        self.steprock = 0#float(self.ui.shiftrock.text())
        self.stepexp = 0#self.form.shiftdata
        self.backgrdData = float(self.ui.edtBackgroundData.text())#self.form.backgroundData
        self.backgrdRC = float(self.ui.edtBackgroundRC.text())#self.form.backgroundRC
        self.wavelength = float(self.ui.edtWavelength.text())#self.form.wavelength
        keepcentered = False
        # print dir(self.form)
        # if hasattr(self.form,'keepdatacentered'):
        keepcentered = self.ui.chkKeepdatacentered.isChecked()#self.form.keepdatacentered

        # if hasattr(self.form,'desmearing'):
        desmearing = self.ui.chkDesmearing.isChecked()#self.form.desmearing
        # print("desmearing ?",desmearing)

        qminimum = float(self.ui.edtQmin.text())#self.form.qmin
        n = 15
        a1 = 10000.
        a2 = 1e-10
        it = 2000
        tol = 1e-15

        # --- setting the datas
        Iexp = self.data_dict[self.selectedData].i - self.backgrdData
        qexp = self.data_dict[self.selectedData].q
        rockname=self.ui.comboRock.currentText()
        Irock = self.data_dict[rockname].i - self.backgrdRC
        qrock = self.data_dict[rockname].q

        a0exp = Iexp[numpy.argmax(Iexp)]
        a0rock = Irock[numpy.argmax(Irock)]
        Thetaexp = qexp
        Thetarock = qrock

        # --- "Fitting to the Gaussian"
        Fitexp, Thetaexp_sel, Iexp_sel, FitParamexp = LSusaxs.FitGauss(Thetaexp, Iexp, n, a0exp, a1, a2, tol, it)
        Fitrock, Thetarock_sel, Irock_sel, FitParamrock = LSusaxs.FitGauss(Thetarock, Irock, n, a0rock, a1, a2, tol, it)
        self.printTXT("Data center found at ", FitParamexp[2])
        self.printTXT("RC center found at ", FitParamrock[2])

        DeltaThetaexp = -FitParamexp[2]  # center
        DeltaThetarock = -FitParamrock[2]
        # ---"-----Fitting procedure end-----"
        # ---"--Steps in Theta-----------------------------------------------"
        shiftexp = abs(Thetaexp_sel[0] - Thetaexp_sel[1])
        shiftrock = abs(Thetarock_sel[0] - Thetarock_sel[1])
        # ---------------------------------------------------------------"

        NewThetaexp_sel = LSusaxs.Qscalemod(DeltaThetaexp, Thetaexp_sel, self.stepexp * shiftexp)
        NewThetarock_sel = LSusaxs.Qscalemod(DeltaThetarock, Thetarock_sel, self.steprock * shiftrock)

        CorrThetaexp_positive = numpy.repeat(NewThetaexp_sel, NewThetaexp_sel > 0.0)
        Iexp_sel_positive = numpy.repeat(Iexp_sel, NewThetaexp_sel > 0.0)
        ####self.somme=LSusaxs.somme(self.CorrThetaexp_positive,self.Iexp_sel_positive)

        # ---"New Theta scales after zero shifting and manual shifting (if any)------"
        qnewexp = LSusaxs.Qscalemod(DeltaThetaexp, Thetaexp, self.stepexp * shiftexp)
        qnewrock = LSusaxs.Qscalemod(DeltaThetarock, Thetarock, self.steprock * shiftexp)
        if keepcentered:
            self.data_dict[self.selectedData + ' centered'] = dataset(self.selectedData + ' centered', qnewexp, Iexp,
                                                                      'centered datas', type='calculated',
                                                                      parent=[self.selectedData])
            self.data_dict['rock' + ' centered'] = dataset('rock' + ' centered', qnewrock, Irock, 'rock datas',
                                                           type='calculated', parent=['rock'])
        # ------------------------------------------------"
        self.printTXT("max value for experimental datas : ", FitParamexp[0])
        self.printTXT("max value for rocking curve datas : ", FitParamrock[0])
        self.TransmissionValue = FitParamexp[0] / FitParamrock[0]
        self.printTXT('Transmission of the sample (%)= ', self.TransmissionValue)
        self.printTXT('Sample thickness (cm.)=', self.thick)
        qnewpos = numpy.repeat(qnewrock, qnewrock >= 0.0)
        Inewpos = numpy.repeat(Irock, qnewrock >= 0.0)
        self.printTXT('minimum Theta taken for central beam calculation= ', qnewpos[0])
        # print 'Intensity at minimum Theta value taken for central beam calculation= ', Inewpos[0]
        # --- 'Area in the central beam'
        thetanewpos = LSusaxs.QtoTheta(qnewpos, self.wavelength)
        sommetheta = 2. * LSusaxs.somme(thetanewpos, Inewpos)
        # print sommet
        # somme=2.*LSusaxs.somme(qnewpos,Inewpos)
        self.printTXT('Central beam area(counts.s^-1.rad^-2)= ', sommetheta)
        qnew, ITcorr = LSusaxs.TrCorrectedProf(qnewexp, Iexp, qnewrock, Irock, self.thick, sommetheta,
                                               self.TransmissionValue)
        # qnew,ITcorr=LSusaxs.TrCorrectedProf(qexp,Iexp,qrock,Irock,self.thick,somme,self.TransmissionValue)
        newname = self.selectedData + ' substracted'
        # ITcorr=LSusaxs.TrCorrectedProf(qnewexp,Iexp,qnewrock,Irock,self.thick,somme,self.TransmissionValue)[1]
        self.data_dict[newname] = dataset(newname, \
                                          numpy.repeat(qnew, qnew > qminimum), \
                                          numpy.repeat(ITcorr, qnew > qminimum), \
                                          'substracted datas', type='scaled', \
                                          parent=[self.selectedData])

        self.parent.redrawTheList()
        self.parent.Replot()
        if desmearing:
            # continue with desmearing
            self.printTXT('No desmearing')
            pass
            desmear = USAXSDesmearing(self.parent, self.selectedData + ' substracted', noGUI=True)
            desmear.setParameters({'ExtrapolationType': LSusaxs.POWERLAW})
            desmear.execute()
        else:
            self.printTXT('No desmearing')
    
class USAXSDesmearing(plugin.pySAXSplugin):
    '''
    class for desmearing USAXS
    '''
    menu="Data Treatment"
    subMenu="USAXS"
    subMenuText="Desmearing"
    icon="find.png"
    
    def execute(self):
        '''
        The user click on "USAXS - Data correction"
        '''
        if self.selectedData is None:
            QtGui.QMessageBox.information(self.parent,"pySAXS", "No data are selected", buttons=QtGui.QMessageBox.Ok, defaultButton=QtGui.QMessageBox.NoButton)
            return
        
        if not(self.data_dict.has_key('resfunc')):
            QtGui.QMessageBox.information(self.parent, "pySAXS","Load resolution function before deconvolution \n\
            (or rename  in 'resfunc')",buttons=QtGui.QMessageBox.Ok, defaultButton=QtGui.QMessageBox.NoButton)
            return
        
        self.printTXT("-----USAXS desmearing --------")
        #read the resolution function
        resx=self.data_dict['resfunc'].q
        resy=self.data_dict['resfunc'].i
        
        if not self.noGUI:
            #call the transmission dialog box with feedback to OnUSAXSTransmission()
            self.typeList=(LSusaxs.POWERLAW,LSusaxs.CONSTANTBACKGROUND)
            
             #generate gui using guidata
            #here we use guidata to generate a dialog box
            
            items = {
             "ExtrapolationType": dataitems.ChoiceItem("Type of extrapolation for high q",self.typeList),
             }
            clz = type("USAXS Desmearing :", (datatypes.DataSet,), items)
            self.form = clz()
            if not self.form.edit():
                #not ok
                return
            print(self.form.ExtrapolationType)
            self.ExtrapolationType=self.form.ExtrapolationType#self.typeList[self.form.ExtrapolationType]        
        #ExtrapolationType should be defined
        self.printTXT('Extrapolation type : ',self.ExtrapolationType)
        it=3
        ns=1
        self.wavelength=1.54189
        Isous=self.data_dict[self.selectedData].i
        qsous=self.data_dict[self.selectedData].q
        #----------- desmearing
        #try:
        theta=LSusaxs.QtoTheta(qsous,self.wavelength)
        Idec=LSusaxs.lake(theta,Isous,it,self.ExtrapolationType,ns,resx=resx,resy=resy) #lake
        #qdec=LSusaxs.ThetatoQ(qsous,self.wavelength)
        qdec=qsous
        #except:
        #QtGui.QMessageBox.information(self.parent, "pySAXS","Error occured when desmearing...",buttons=QtGui.QMessageBox.Ok, defaultButton=QtGui.QMessageBox.NoButton)
        #-------------------------------------------
        #qdec=qsous   
        self.data_dict[self.selectedData+' deconvoluted']=dataset(self.selectedData+' deconvoluted',qdec,Idec,self.selectedData+'deconvoluted',True,type='calculated',parent=[self.selectedData])
        self.redrawTheList()
        self.Replot()
    
    