# The PyMca X-Ray Fluorescence Toolkit
# This file is part of the PyMca X-ray Fluorescence Toolkit developed at
# the ESRF by the Software group.
__author__ = "V.A. Sole - ESRF Data Analysis"
__contact__ = "example@nnthink.com"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
import os
import numpy
import logging
arctan = numpy.arctan
from PyMca5.PyMcaMath.fitting import SpecfitFuns
from PyMca5.PyMcaMath.fitting.Gefit import LeastSquaresFit

_logger = logging.getLogger(__name__)

try:
    HOME=os.getenv('HOME')
except:
    HOME=None
if HOME is not None:
    os.environ['HOME'] = HOME
else:
    os.environ['HOME'] = "."
SPECFITFUNCTIONS_DEFAULTS={'FileAction':0,
                 'infile':os.environ['HOME']+'/.specfitdefaults.py',
                 'outfile':os.environ['HOME']+'/.specfitdefaults.py',
                 'Geometry':"600x400+50+50",
                 'NoConstrainsFlag':0,
                 'BackgroundIndex':1,
                 'PosFwhmFlag':1,
                 'HeightAreaFlag':1,
                 'SameFwhmFlag':1,
                 'PositionFlag':0,
                 'EtaFlag':0,
                 'WeightFlag':0,
                 'Yscaling':1.0,
                 'Xscaling':1.0,
                 'FwhmPoints':8,
                 'AutoFwhm':0,
                 'Sensitivity':2.5,
                 'ForcePeakPresence':0,
                 'McaMode':0,
                 #Hypermet
                 'HypermetTails':15,
                 'QuotedFwhmFlag':0,
                 'MaxFwhm2InputRatio':1.5,
                 'MinFwhm2InputRatio':0.4,
                 #short tail parameters
                 'MinGaussArea4ShortTail':50000.,
                 'InitialShortTailAreaRatio':0.050,
                 'MaxShortTailAreaRatio':0.100,
                 'MinShortTailAreaRatio':0.0010,
                 'InitialShortTailSlopeRatio':0.70,
                 'MaxShortTailSlopeRatio':2.00,
                 'MinShortTailSlopeRatio':0.50,
                 #long tail parameters
                 'MinGaussArea4LongTail':1000.0,
                 'InitialLongTailAreaRatio':0.050,
                 'MaxLongTailAreaRatio':0.300,
                 'MinLongTailAreaRatio':0.010,
                 'InitialLongTailSlopeRatio':20.0,
                 'MaxLongTailSlopeRatio':50.0,
                 'MinLongTailSlopeRatio':5.0,
                 #step tail
                 'MinGaussHeight4StepTail':5000.,
                 'InitialStepTailHeightRatio':0.002,
                 'MaxStepTailHeightRatio':0.0100,
                 'MinStepTailHeightRatio':0.0001,
                 #Hypermet constraints
                 'QuotedPositionFlag':1,
                 'DeltaPositionFwhmUnits':0.5,
                 'SameSlopeRatioFlag':1,
                 'SameAreaRatioFlag':1}
CFREE       = 0
CPOSITIVE   = 1
CQUOTED     = 2
CFIXED      = 3
CFACTOR     = 4
CDELTA      = 5
CSUM        = 6
CIGNORED    = 7

class SpecfitFunctions(object):
    def __init__(self,config=None):
        if config is None:
            self.config=SPECFITFUNCTIONS_DEFAULTS
        else:
            self.config=config

    def gauss(self,pars,x):
       """
       A fit function.
       """
       return SpecfitFuns.gauss(pars,x)

    def agauss(self,pars,x):
       """
       A fit function.
       """
       return SpecfitFuns.agauss(pars,x)

    def hypermet(self,pars,x):
       """
       A fit function.
       """
       return SpecfitFuns.ahypermet(pars,x,self.config['HypermetTails'],0)

    def lorentz(self,pars,x):
       """
       Fit function.
       """
       return SpecfitFuns.lorentz(pars,x)

    def alorentz(self,pars,x):
       """
       Fit function.
       """
       return SpecfitFuns.alorentz(pars,x)

    def pvoigt(self,pars,x):
       """
       Fit function.
       """
       return SpecfitFuns.pvoigt(pars,x)

    def apvoigt(self,pars,x):
       """
       Fit function.
       """
       return SpecfitFuns.apvoigt(pars,x)

    def splitgauss(self,pars,x):
       """
       Asymmetric gaussian.
       """
       return SpecfitFuns.splitgauss(pars,x)

    def splitlorentz(self,pars,x):
       """
       Asymmetric lorentz.
       """
       return SpecfitFuns.splitlorentz(pars,x)

    def splitpvoigt(self,pars,x):
       """
       Asymmetric pseudovoigt.
       """
       return SpecfitFuns.splitpvoigt(pars,x)

    def stepdown(self,pars,x):
        """
        A fit function.
        """
        return SpecfitFuns.downstep(pars,x)

    def stepup(self,pars,x):
        """
        A fit function.
        """
        return SpecfitFuns.upstep(pars,x)

    def slit(self,pars,x):
        """
        A fit function.
        """
        return 0.5*SpecfitFuns.slit(pars,x)

    def atan(self, pars, x):
        return pars[0] * (0.5 + (arctan((1.0*x-pars[1])/pars[2]) / numpy.pi))

    def periodic_gauss(self, pars, x):
        """
        Fit function periodic_gauss(pars, x)
        pars = [npeaks, delta, height, position, fwhm]
        """
        newpars = numpy.zeros((int(pars[0]), 3), numpy.float64)
        for i in range(int(pars[0])):
            newpars[i, 0] = pars[2]
            newpars[i, 1] = pars[3] + i * pars[1]
            newpars[:, 2] = pars[4]
        return SpecfitFuns.gauss(newpars,x)

    def gauss2(self,param0,t0):
        param=numpy.array(param0)
        t=numpy.array(t0)
        dummy=2.3548200450309493*(t-param[1])/param[2]
        return param[0] * self.myexp(-0.5 * dummy * dummy)

    def myexp(self,x):
        # put a (bad) filter to avoid over/underflows
        # with no python looping
        return numpy.exp(x * numpy.less(abs(x),250))-1.0*numpy.greater_equal(abs(x),250)

    def indexx(x):
        #adapted from runningReport (Mike Fletcher, Python newsgroup)
        set = map(None,x,range(len(x)))
        set.sort() #sorts by values and then by index
        return map(lambda x:x[1],set)

    def bkg_constant(self, pars, x):
        """
        Constant background
        """
        return pars[0] * numpy.ones_like(x)

    def bkg_linear(self, pars, x):
        """
        Linear background
        """
        return pars[0] + pars[1] * x

    def bkg_internal(self,pars,x):
        """
        Internal Background
        """
        #fast
        #slow: recalculate the background as function of the parameters
        yy=SpecfitFuns.subac(self.ydata*1.0,
                             pars[0],pars[1])
        nrx=shape(x)[0]
        nry=shape(yy)[0]
        if nrx == nry:
            return SpecfitFuns.subac(yy,pars[0],pars[1])
        else:
            return SpecfitFuns.subac(numpy.take(yy,numpy.arange(0,nry,2)),pars[0],pars[1])

    def bkg_none(self,pars,x):
        """
        Internal Background
        """

    def fun(self,param, t):
        gterm = param[2] * numpy.exp(-0.5 * ((t - param[3]) * (t - param[3]))/param[4])
        bterm = param[1] * t + param[0]
        return gterm + bterm

    def estimate(self,x,y,z,xscaling=1.0,yscaling=1.0):
        ngauss = input(' Number of Gaussians : ')
        ngauss=int(ngauss)
        if ngauss < 1:
            ngauss=1
        newpar=[]
        for i in range(ngauss):
            _logger.info("Defining Gaussian numer %d ", i+1)
            newpar.append(input('Height   = '))
            newpar.append(input('Position = '))
            newpar.append(input('FWHM     = '))
        return newpar,numpy.zeros((3,len(newpar)),numpy.float64)

    def seek(self,y,x=None,yscaling=None,
                           fwhm=None,
                           sensitivity=None,
                           mca=None):
        """
        SpecfitFunctions.seek(self,y,
                                x=None,
                                yscaling=None,fwhm=None,sensitivity=None,
                                mca=None)
        It searches for peaks in the y array. If x it is given, it gives back
        the closest x(s) to the position of the peak(s). Otherways it gives back
        the index of the closest point to the peak.
        """
        if yscaling is None:
            yscaling=self.config['Yscaling']

        if fwhm is None:
            if self.config['AutoFwhm']:
                fwhm=self.guess_fwhm(x=x,y=y)
            else:
                fwhm=self.config['FwhmPoints']
        if sensitivity is None:
            sensitivity=self.config['Sensitivity']
        if mca is None:
            mca=self.config['McaMode']

        search_fwhm=int(max(fwhm,3))
        search_sensitivity=max(1.0,sensitivity)
        mca=1.0
        if mca:
            ysearch = numpy.array(y) * yscaling
        else:
            ysearch = numpy.ones([len(y) + 2 * search_fwhm,], numpy.float64)
            if y[0]:
                ysearch[0:(search_fwhm+1)]=ysearch[0:(search_fwhm+1)]*y[0]*yscaling
            else:
                ysearch[0:(search_fwhm+1)]=ysearch[0:(search_fwhm+1)]*yscaling*sum(y[0:3])/3.0
            ysearch[-1:-(search_fwhm+1):-1]=ysearch[-1:-(search_fwhm+1):-1]*y[len(y)-1]*yscaling
            ysearch[search_fwhm:(search_fwhm+len(y))]=y*yscaling
        npoints=len(ysearch)
        if npoints > (1.5)*search_fwhm:
            peaks=SpecfitFuns.seek(ysearch,0,npoints,
                                    search_fwhm,
                                    search_sensitivity)
        else:
            peaks=[]

        if len(peaks) > 0:
            if mca == 0:
                for i in range(len(peaks)):
                    peaks[i]=int(peaks[i]-search_fwhm)
            for i in peaks:
                if (i < 1) | (i > (npoints-1)):
                    peaks.remove(i)
            if x is not None:
                if len(x) == len(y):
                    for i in range(len(peaks)):
                        peaks[i]=x[int(max(peaks[i],0))]
        return peaks

    def guess_fwhm(self, **kw):
        if 'y' in kw:
            y=kw['y']
        else:
            return self.config['FwhmPoints']
        if 'x' in kw:
            x=kw['x']
        else:
            x=numpy.arange(len(y))*1.0

        #set at least a default value for the fwhm
        fwhm=4

        zz=SpecfitFuns.subac(y,1.000,1000)
        yfit=y-zz

        #now I should do some sort of peak search ...
        maximum = max(yfit)
        idx = numpy.nonzero(yfit == maximum)[0]
        pos = numpy.take(x,idx)[-1]
        posindex = idx[-1]
        height = yfit[posindex]
        imin = posindex
        while ((yfit[imin] > 0.5*height) & (imin >0)):
            imin=imin - 1
        imax=posindex
        while ((yfit[imax] > 0.5*height) & (imax <(len(yfit)-1))):
            imax=imax + 1
        fwhm=max(imax-imin-1,fwhm)

        return fwhm

    def estimate_constant(self, xx, yy, zzz, xscaling=1.0, yscaling=None):
        estimated_par = [min(yy)]
        constraints = [[0], [0], [0]]
        return estimated_par, constraints

    def estimate_linear(self, xx, yy, zzz, xscaling=1.0, yscaling=None):
        # compute strip bg and use it to estimate the linear bg parameters
        zz = SpecfitFuns.subac(yy, 1.000, 10000)
        n = float(len(zz))
        Sy = numpy.sum(zz)
        Sx = float(numpy.sum(xx))
        Sxx = float(numpy.sum(xx * xx))
        Sxy = float(numpy.sum(xx * zz))

        deno = n * Sxx - (Sx * Sx)
        if deno != 0:
            bg = (Sxx * Sy - Sx * Sxy) / deno
            slope = (n * Sxy - Sx * Sy) / deno
        else:
            bg = 0.0
            slope = 0.0
        estimated_par = [bg, slope]
        constraints = [[0, 0], [0, 0], [0, 0]]
        return estimated_par, constraints

    def estimate_gauss(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
       if yscaling == None:
            try:
                yscaling=self.config['Yscaling']
            except:
                yscaling=1.0
       if yscaling == 0:
            yscaling=1.0
       fittedpar=[]
       zz=SpecfitFuns.subac(yy,1.000,10000)

       npoints = len(zz)
       if self.config['AutoFwhm']:
            search_fwhm=self.guess_fwhm(x=xx,y=yy)
       else:
            search_fwhm=int(float(self.config['FwhmPoints']))
       search_sens=float(self.config['Sensitivity'])
       search_mca=int(float(self.config['McaMode']))

       if search_fwhm < 3:
            search_fwhm = 3
            self.config['FwhmPoints']=3

       if search_sens < 1:
            search_sens = 1
            self.config['Sensitivity']=1

       if npoints > 1.5*search_fwhm:
            peaks=self.seek(yy,fwhm=search_fwhm,
                               sensitivity=search_sens,
                               yscaling=yscaling,
                               mca=search_mca)
       else:
            peaks = []
       if not len(peaks):
            mca = int(float(self.config.get('McaMode', 0)))
            forcePeak =  int(float(self.config.get('ForcePeakPresence', 0)))
            if (not mca) and forcePeak:
                delta = yy - zz
                peaks  = [int(numpy.nonzero(delta == delta.max())[0])]

       largest_index = 0
       if len(peaks) > 0:
            j = 0
            for i in peaks:
                if j == 0:
                    sig=5*abs(xx[npoints-1]-xx[0])/npoints
                    peakpos = xx[int(i)]
                    if abs(peakpos) < 1.0e-16:
                        peakpos = 0.0
                    param = numpy.array([yy[int(i)] - zz[int(i)], peakpos ,sig])
                    largest = param
                else:
                    param2 = numpy.array([yy[int(i)] - zz [int(i)], xx[int(i)] ,sig])
                    param = numpy.concatenate((param,param2))
                    if (param2[0] > largest[0]):
                        largest = param2
                        largest_index = j
                j = j + 1
            xw = numpy.resize(xx,(npoints,1))
            if (0):
                sy = numpy.sqrt(abs(yy))
                yw = numpy.resize(yy-zz,(npoints,1))
                sy = numpy.resize(sy,(npoints,1))
                datawork = numpy.concatenate((xw,yw,sy),1)
            else:
                yw = numpy.resize(yy-zz,(npoints,1))
                datawork = numpy.concatenate((xw,yw),1)
            cons = numpy.zeros((3,len(param)),numpy.float64)
            cons [0] [0:len(param):3] = CPOSITIVE
            #force peaks to stay around their position
            if (1):
                cons [0] [1:len(param):3] = CQUOTED
                #FWHM should be given in terms of X not of points!
                if len(xw) > search_fwhm:
                    fwhmx = numpy.fabs(xw[int(search_fwhm)]-xw[0])
                    cons [1] [1:len(param):3] = param[1:len(param):3]-0.5*fwhmx
                    cons [2] [1:len(param):3] = param[1:len(param):3]+0.5*fwhmx
                else:
                    cons [1] [1:len(param):3] = numpy.ones(shape(param[1:len(param):3]),numpy.float64)*min(xw)
                    cons [2] [1:len(param):3] = numpy.ones(shape(param[1:len(param):3]),numpy.float64)*max(xw)

            if 0:
                cons [0] [2:len(param):3] = CFACTOR
                cons [1] [2:len(param):3] = 2
                cons [2] [2:len(param):3] = 1.0
                cons [0] [2] = CPOSITIVE
                cons [1] [2] = 0
                cons [2] [2] = 0
            else:
                cons [0] [2:len(param):3] = CPOSITIVE
            fittedpar, chisq, sigmapar = LeastSquaresFit(SpecfitFuns.gauss,param,
                                            datawork,
                                            weightflag=self.config['WeightFlag'],
                                            maxiter=4,constrains=cons.tolist())
            #I already have the estimation
            #print self.seek(yw,\
            #                   fwhm=search_fwhm,\
            #                   sensitivity=search_sens,\
       #     #Use SPEC estimate ...
       #         #forget SPEC estimation of fwhm
       cons = numpy.zeros((3,len(fittedpar)),numpy.float64)
       j=0
       for i in range(len(peaks)):
                #Setup height area constrains
                if self.config['NoConstrainsFlag'] == 0:
                    if self.config['HeightAreaFlag']:
                        cons[0] [j] = 1
                        cons[1] [j] = 0
                        cons[2] [j] = 0
                j=j+1

                #Setup position constrains
                if self.config['NoConstrainsFlag'] == 0:
                    if self.config['PositionFlag']:
                                cons[0][j]=2
                                cons[1][j]=min(xx)
                                cons[2][j]=max(xx)
                j=j+1

                #Setup positive FWHM constrains
                if self.config['NoConstrainsFlag'] == 0:
                    if self.config['PosFwhmFlag']:
                                cons[0][j]=1
                                cons[1][j]=0
                                cons[2][j]=0
                    if self.config['SameFwhmFlag']:
                        if (i != largest_index):
                                cons[0][j]=4
                                cons[1][j]=3*largest_index+2
                                cons[2][j]=1.0
                j=j+1
       return fittedpar,cons

    def estimate_lorentz(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
       fittedpar,cons = self.estimate_gauss(xx,yy,zzz,xscaling,yscaling)
       return  fittedpar,cons

    def estimate_agauss(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
        fittedpar,cons = self.estimate_gauss(xx,yy,zzz,xscaling,yscaling)
        #get the number of found peaks
        npeaks=int(len(cons[0])/3)
        for i in range(npeaks):
            height = fittedpar[3*i]
            fwhm = fittedpar[3*i+2]
            fittedpar[3*i] = (height * fwhm / (2.0*numpy.sqrt(2*numpy.log(2))))*numpy.sqrt(2 * numpy.pi)
        return fittedpar,cons

    def estimate_alorentz(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
        fittedpar,cons = self.estimate_gauss(xx,yy,zzz,xscaling,yscaling)
        #get the number of found peaks
        npeaks=int(len(cons[0])/3)
        for i in range(npeaks):
            height = fittedpar[3*i]
            fwhm = fittedpar[3*i+2]
            fittedpar[3*i] = (height * fwhm * 0.5 * numpy.pi)
        return fittedpar,cons

    def estimate_splitgauss(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
        fittedpar, cons = self.estimate_gauss(xx,yy,zzz,xscaling,yscaling)
        #get the number of found peaks
        npeaks=int(len(cons[0])/3)
        estimated_parameters = []
        estimated_constraints = numpy.zeros((3, 4*npeaks),numpy.float64)
        for i in range(npeaks):
            for j in range(3):
                estimated_parameters.append(fittedpar[3*i+j])
            estimated_parameters.append(fittedpar[3*i+2])
            estimated_constraints[0][4*i]   = cons [0][3*i]
            estimated_constraints[0][4*i+1] = cons [0][3*i+1]
            estimated_constraints[0][4*i+2] = cons [0][3*i+2]
            estimated_constraints[0][4*i+3] = cons [0][3*i+2]
            estimated_constraints[1][4*i]   = cons [1][3*i]
            estimated_constraints[1][4*i+1] = cons [1][3*i+1]
            estimated_constraints[1][4*i+2] = cons [1][3*i+2]
            estimated_constraints[1][4*i+3] = cons [1][3*i+2]
            estimated_constraints[2][4*i]   = cons [2][3*i]
            estimated_constraints[2][4*i+1] = cons [2][3*i+1]
            estimated_constraints[2][4*i+2] = cons [2][3*i+2]
            estimated_constraints[2][4*i+3] = cons [2][3*i+2]
            if cons[0][3*i+2] == 4:
                #same FWHM case
                estimated_constraints[1][4*i+2] = int(cons[1][3*i+2]/3) * 4 + 2
                estimated_constraints[1][4*i+3] = int(cons[1][3*i+2]/3) * 4 + 3
        return estimated_parameters, estimated_constraints

    def estimate_splitlorentz(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
        fittedpar, cons = self.estimate_gauss(xx,yy,zzz,xscaling,yscaling)
        #get the number of found peaks
        npeaks=int(len(cons[0])/3)
        estimated_parameters = []
        estimated_constraints = numpy.zeros((3, 4*npeaks),numpy.float64)
        for i in range(npeaks):
            for j in range(3):
                estimated_parameters.append(fittedpar[3*i+j])
            estimated_parameters.append(fittedpar[3*i+2])
            estimated_constraints[0][4*i]   = cons [0][3*i]
            estimated_constraints[0][4*i+1] = cons [0][3*i+1]
            estimated_constraints[0][4*i+2] = cons [0][3*i+2]
            estimated_constraints[0][4*i+3] = cons [0][3*i+2]
            estimated_constraints[1][4*i]   = cons [1][3*i]
            estimated_constraints[1][4*i+1] = cons [1][3*i+1]
            estimated_constraints[1][4*i+2] = cons [1][3*i+2]
            estimated_constraints[1][4*i+3] = cons [1][3*i+2]
            estimated_constraints[2][4*i]   = cons [2][3*i]
            estimated_constraints[2][4*i+1] = cons [2][3*i+1]
            estimated_constraints[2][4*i+2] = cons [2][3*i+2]
            estimated_constraints[2][4*i+3] = cons [2][3*i+2]
            if cons[0][3*i+2] == 4:
                #same FWHM case
                estimated_constraints[1][4*i+2] = int(cons[1][3*i+2]/3) * 4 + 2
                estimated_constraints[1][4*i+3] = int(cons[1][3*i+2]/3) * 4 + 3
        return estimated_parameters, estimated_constraints

    def estimate_pvoigt(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
       fittedpar,cons = self.estimate_gauss(xx,yy,zzz,xscaling,yscaling)
       npeaks=int(len(cons[0])/3)
       newpar=[]
       newcons=numpy.zeros((3,4*npeaks),numpy.float64)
       #find out related parameters proper index
       if self.config['NoConstrainsFlag'] == 0:
            if self.config['SameFwhmFlag']:
                j=0
                #get the index of the free FWHM
                for i in range(npeaks):
                    if cons[0][3*i+2] != 4:
                        j=i
                for i in range(npeaks):
                    if i != j:
                       cons[1][3*i+2] = 4*j+2
       for i in range(npeaks):
            newpar.append(fittedpar[3*i])
            newpar.append(fittedpar[3*i+1])
            newpar.append(fittedpar[3*i+2])
            newpar.append(0.5)
            newcons[0][4*i]=cons[0][3*i]
            newcons[0][4*i+1]=cons[0][3*i+1]
            newcons[0][4*i+2]=cons[0][3*i+2]
            newcons[1][4*i]=cons[1][3*i]
            newcons[1][4*i+1]=cons[1][3*i+1]
            newcons[1][4*i+2]=cons[1][3*i+2]
            newcons[2][4*i]=cons[2][3*i]
            newcons[2][4*i+1]=cons[2][3*i+1]
            newcons[2][4*i+2]=cons[2][3*i+2]
            #Eta constrains
            newcons[0][4*i+3]=0
            newcons[1][4*i+3]=0
            newcons[2][4*i+3]=0
            if self.config['NoConstrainsFlag'] == 0:
                if self.config['EtaFlag']:
                      newcons[0][4*i+3]=2
                      newcons[1][4*i+3]=0.0
                      newcons[2][4*i+3]=1.0
       return  newpar,newcons

    def estimate_splitpvoigt(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
       fittedpar,cons = self.estimate_gauss(xx,yy,zzz,xscaling,yscaling)
       npeaks=int(len(cons[0])/3)
       newpar=[]
       newcons=numpy.zeros((3,5*npeaks),numpy.float64)
       #find out related parameters proper index
       if self.config['NoConstrainsFlag'] == 0:
            if self.config['SameFwhmFlag']:
                j=0
                #get the index of the free FWHM
                for i in range(npeaks):
                    if cons[0][3*i+2] != 4:
                        j=i
                for i in range(npeaks):
                    if i != j:
                       cons[1][3*i+2] = 4*j+2
       for i in range(npeaks):
            #height
            newpar.append(fittedpar[3*i])
            #position
            newpar.append(fittedpar[3*i+1])
            #fwhm1
            newpar.append(fittedpar[3*i+2])
            #fwhm2 equal to the first
            newpar.append(fittedpar[3*i+2])
            #eta
            newpar.append(0.5)
            newcons[0][5*i]=cons[0][3*i]
            newcons[0][5*i+1]=cons[0][3*i+1]
            newcons[0][5*i+2]=cons[0][3*i+2]
            newcons[0][5*i+3]=cons[0][3*i+2]
            newcons[1][5*i]=cons[1][3*i]
            newcons[1][5*i+1]=cons[1][3*i+1]
            newcons[1][5*i+2]=cons[1][3*i+2]
            newcons[1][5*i+3]=cons[1][3*i+2]
            newcons[2][5*i]=cons[2][3*i]
            newcons[2][5*i+1]=cons[2][3*i+1]
            newcons[2][5*i+2]=cons[2][3*i+2]
            newcons[2][5*i+3]=cons[2][3*i+2]
            if cons[0][3*i+2] == 4:
                newcons[1][5*i+3] = newcons[1][5*i+2]+1
            #Eta constrains
            newcons[0][5*i+4]=0
            newcons[1][5*i+4]=0
            newcons[2][5*i+4]=0
            if self.config['NoConstrainsFlag'] == 0:
                if self.config['EtaFlag']:
                      newcons[0][5*i+4]=2
                      newcons[1][5*i+4]=0.0
                      newcons[2][5*i+4]=1.0
       return  newpar,newcons

    def estimate_apvoigt(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
       fittedpar,cons = self.estimate_pvoigt(xx,yy,zzz,xscaling,yscaling)
       npeaks=int(len(cons[0])/4)
       for i in range(npeaks):
            height = fittedpar[4*i]
            fwhm = fittedpar[4*i+2]
            fittedpar[4*i] = 0.5*(height * fwhm * 0.5 * numpy.pi)+\
                         0.5*(height * fwhm / (2.0*numpy.sqrt(2*numpy.log(2))))* numpy.sqrt(2 * numpy.pi)
       return fittedpar,cons

    def estimate_hypermet(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
       """
       if yscaling == None:
            try:
                yscaling=self.config['Yscaling']
            except:
                yscaling=1.0
       if yscaling == 0:
            yscaling=1.0
       """
       fittedpar,cons = self.estimate_gauss(xx,yy,zzz,xscaling,yscaling)
       npeaks=int(len(cons[0])/3)
       newpar=[]
       newcons=numpy.zeros((3,8*npeaks),numpy.float64)
       main_peak=0
       #find out related parameters proper index
       if self.config['NoConstrainsFlag'] == 0:
            if self.config['SameFwhmFlag']:
                j=0
                #get the index of the free FWHM
                for i in range(npeaks):
                    if cons[0][3*i+2] != 4:
                        j=i
                for i in range(npeaks):
                    if i != j:
                       cons[1][3*i+2] = 8*j+2
                main_peak=j
       for i in range(npeaks):
                if fittedpar[3*i] > fittedpar[3*main_peak]:
                    main_peak = i

       for i in range(npeaks):
            height = fittedpar[3*i]
            position=fittedpar[3*i+1]
            fwhm = fittedpar[3*i+2]
            area = (height * fwhm / (2.0*numpy.sqrt(2*numpy.log(2))))*numpy.sqrt(2 * numpy.pi)
            #the gaussian parameters
            newpar.append(area)
            newpar.append(position)
            newpar.append(fwhm)
            #Avoid zero derivatives because of not calculating contribution
            g_term    = 1
            st_term   = 1
            lt_term   = 1
            step_term = 1
            if self.config['HypermetTails'] != 0:
                    g_term    = self.config['HypermetTails'] & 1
                    st_term   = (self.config['HypermetTails']>>1) & 1
                    lt_term   = (self.config['HypermetTails']>>2) & 1
                    step_term = (self.config['HypermetTails']>>3) & 1
            if g_term == 0:
                    #fix the gaussian parameters
                    newcons[0][8*i]=CFIXED
                    newcons[0][8*i+1]=CFIXED
                    newcons[0][8*i+2]=CFIXED
            #the short tail parameters
            if ((area*yscaling) < \
               self.config['MinGaussArea4ShortTail']) | \
               (st_term == 0):
               newpar.append(0.0)
               newpar.append(0.0)
               newcons[0][8*i+3]=CFIXED
               newcons[1][8*i+3]=0.0
               newcons[2][8*i+3]=0.0
               newcons[0][8*i+4]=CFIXED
               newcons[1][8*i+4]=0.0
               newcons[2][8*i+4]=0.0
            else:
               newpar.append(self.config['InitialShortTailAreaRatio'])
               newpar.append(self.config['InitialShortTailSlopeRatio'])
               newcons[0][8*i+3]=CQUOTED
               newcons[1][8*i+3]=self.config['MinShortTailAreaRatio']
               newcons[2][8*i+3]=self.config['MaxShortTailAreaRatio']
               newcons[0][8*i+4]=CQUOTED
               newcons[1][8*i+4]=self.config['MinShortTailSlopeRatio']
               newcons[2][8*i+4]=self.config['MaxShortTailSlopeRatio']
            #the long tail parameters
            if ((area*yscaling) < \
               self.config['MinGaussArea4LongTail']) | \
               (lt_term == 0):
               newpar.append(0.0)
               newpar.append(0.0)
               newcons[0][8*i+5]=CFIXED
               newcons[1][8*i+5]=0.0
               newcons[2][8*i+5]=0.0
               newcons[0][8*i+6]=CFIXED
               newcons[1][8*i+6]=0.0
               newcons[2][8*i+6]=0.0
            else:
               newpar.append(self.config['InitialLongTailAreaRatio'])
               newpar.append(self.config['InitialLongTailSlopeRatio'])
               newcons[0][8*i+5]=CQUOTED
               newcons[1][8*i+5]=self.config['MinLongTailAreaRatio']
               newcons[2][8*i+5]=self.config['MaxLongTailAreaRatio']
               newcons[0][8*i+6]=CQUOTED
               newcons[1][8*i+6]=self.config['MinLongTailSlopeRatio']
               newcons[2][8*i+6]=self.config['MaxLongTailSlopeRatio']
            #the step parameters
            if ((height*yscaling) < \
               self.config['MinGaussHeight4StepTail'])| \
               (step_term == 0):
               newpar.append(0.0)
               newcons[0][8*i+7]=CFIXED
               newcons[1][8*i+7]=0.0
               newcons[2][8*i+7]=0.0
            else:
               newpar.append(self.config['InitialStepTailHeightRatio'])
               newcons[0][8*i+7]=CQUOTED
               newcons[1][8*i+7]=self.config['MinStepTailHeightRatio']
               newcons[2][8*i+7]=self.config['MaxStepTailHeightRatio']
       if npeaks > 0:
          if g_term:
            if self.config['HeightAreaFlag']:
                for i in range(npeaks):
                    newcons[0] [8*i] = CPOSITIVE
            if self.config['PosFwhmFlag']:
                for i in range(npeaks):
                    newcons[0] [8*i+2] = CPOSITIVE
            if self.config['SameFwhmFlag']:
                for i in range(npeaks):
                    if i != main_peak:
                        newcons[0] [8*i+2] = CFACTOR
                        newcons[1] [8*i+2] = 8*main_peak+2
                        newcons[2] [8*i+2] = 1.0
            if self.config['QuotedPositionFlag']:
                for i in range(npeaks):
                    delta=self.config['DeltaPositionFwhmUnits'] * \
                      int(float(self.config['FwhmPoints']))
                    #that was delta in points
                    #I need it in terms of FWHM
                    delta=self.config['DeltaPositionFwhmUnits'] * fwhm
                    newcons[0][8*i+1] = CQUOTED
                    newcons[1][8*i+1] = newpar[8*i+1]-delta
                    newcons[2][8*i+1] = newpar[8*i+1]+delta
          if self.config['SameSlopeRatioFlag']:
            for i in range(npeaks):
                if i != main_peak:
                    newcons[0] [8*i+4] = CFACTOR
                    newcons[1] [8*i+4] = 8*main_peak+4
                    newcons[2] [8*i+4] = 1.0
                    newcons[0] [8*i+6] = CFACTOR
                    newcons[1] [8*i+6] = 8*main_peak+6
                    newcons[2] [8*i+6] = 1.0
          if self.config['SameAreaRatioFlag']:
            for i in range(npeaks):
                if i != main_peak:
                    newcons[0] [8*i+3] = CFACTOR
                    newcons[1] [8*i+3] = 8*main_peak+3
                    newcons[2] [8*i+3] = 1.0
                    newcons[0] [8*i+5] = CFACTOR
                    newcons[1] [8*i+5] = 8*main_peak+5
                    newcons[2] [8*i+5] = 1.0
       return  newpar,newcons

    def estimate_stepdown(self,xxx,yyy,zzz,xscaling=1.0,yscaling=1.0):
        crappyfilter=[-0.25,-0.75,0.0,0.75,0.25]
        cutoff = 2
        yy = numpy.convolve(yyy, crappyfilter, mode=1)[2:-2]
        if max(yy) > 0:
            yy = yy * max(yyy)/max(yy)
        xx=xxx[2:-2]
        # use height and not area
        fittedpar,cons=self.estimate_gauss(xx,yy,zzz,xscaling,yscaling)
        npeaks=int(len(fittedpar)/3)
        largest_index=0
        largest=[fittedpar[3*largest_index],
                 fittedpar[3*largest_index+1],
                 fittedpar[3*largest_index+2]]
        newcons=numpy.zeros((3,3),numpy.float64)
        for i in range(npeaks):
            if fittedpar[3*i] > largest[0]:
                largest_index=i
                largest=[fittedpar[3*largest_index],
                         fittedpar[3*largest_index+1],
                         fittedpar[3*largest_index+2]]
        largest[0] = max(yyy)-min(yyy)
        #Setup constrains
        if self.config['NoConstrainsFlag'] == 0:
                #Setup height constrains
                if self.config['HeightAreaFlag']:
                        cons[0] [0] = 1
                        cons[1] [0] = 0
                        cons[2] [0] = 0

                #Setup position constrains
                if self.config['PositionFlag']:
                                cons[0][1]=2
                                cons[1][1]=min(xxx)
                                cons[2][1]=max(xxx)

                #Setup positive FWHM constrains
                if self.config['PosFwhmFlag']:
                                cons[0][2]=1
                                cons[1][2]=0
                                cons[2][2]=0

        return largest,cons

    def estimate_slit(self, xxx, yyy, zzz, xscaling=1.0, yscaling=1.0):
        largestup, cons = self.estimate_stepup(xxx, yyy, zzz, xscaling, yscaling)
        largestdown, cons = self.estimate_stepdown(xxx, yyy, zzz, xscaling, yscaling)
        height = 0.5*(largestup[0] + largestdown[0])
        position = 0.5*(largestup[1] + largestdown[1])
        fwhm = numpy.fabs(largestdown[1] - largestup[1])
        beamfwhm = 0.5 * (largestup[2] + largestdown[1])
        beamfwhm = min(beamfwhm, fwhm / 10.0)
        beamfwhm = max(beamfwhm,(max(xxx)-min(xxx))*3.0/len(xxx))
        #own estimation
        yy=yyy-zzz
        height=max(yyy-zzz)
        i1=numpy.nonzero(yy>=0.5*height)[0]
        xx=numpy.take(xxx,i1)
        position=(xx[0]+xx[-1])/2.0
        fwhm=xx[-1]-xx[0]
        largest=[height,position,fwhm,beamfwhm]
        cons=numpy.zeros((3,4),numpy.float64)
        #Setup constrains
        if self.config['NoConstrainsFlag'] == 0:
                #Setup height constrains
                if self.config['HeightAreaFlag']:
                        cons[0] [0] = 1
                        cons[1] [0] = 0
                        cons[2] [0] = 0

                #Setup position constrains
                if self.config['PositionFlag']:
                                cons[0][1]=2
                                cons[1][1]=min(xxx)
                                cons[2][1]=max(xxx)

                #Setup positive FWHM constrains
                if self.config['PosFwhmFlag']:
                                cons[0][2]=1
                                cons[1][2]=0
                                cons[2][2]=0

                #Setup positive FWHM constrains
                if self.config['PosFwhmFlag']:
                                cons[0][3]=1
                                cons[1][3]=0
                                cons[2][3]=0
        return largest,cons

    def estimate_stepup(self, xxx, yyy, zzz, xscaling=1.0, yscaling=1.0):
        crappyfilter = [0.25, 0.75, 0.0, -0.75, -0.25]
        cutoff = 2
        yy = numpy.convolve(yyy, crappyfilter, mode=1)[2:-2]
        if max(yy) > 0:
            yy = yy * max(yyy)/max(yy)
        xx = xxx[2:-2]
        fittedpar, cons = self.estimate_gauss(xx, yy, zzz, xscaling, yscaling)
        npeaks = int(len(fittedpar) / 3)
        largest_index = 0
        largest = [fittedpar[3 * largest_index],
                   fittedpar[3 * largest_index + 1],
                   fittedpar[3 * largest_index + 2]]
        newcons=numpy.zeros((3,3),numpy.float64)
        for i in range(npeaks):
            if fittedpar[3*i] > largest[0]:
                largest_index=i
                largest=[fittedpar[3*largest_index],
                         fittedpar[3*largest_index+1],
                         fittedpar[3*largest_index+2]]
        largest[0] = max(yyy)-min(yyy)
        #Setup constrains
        if self.config['NoConstrainsFlag'] == 0:
                #Setup height constrains
                if self.config['HeightAreaFlag']:
                        cons[0] [0] = 1
                        cons[1] [0] = 0
                        cons[2] [0] = 0

                #Setup position constrains
                if self.config['PositionFlag']:
                                cons[0][1]=2
                                cons[1][1]=min(xxx)
                                cons[2][1]=max(xxx)

                #Setup positive FWHM constrains
                if self.config['PosFwhmFlag']:
                                cons[0][2]=1
                                cons[1][2]=0
                                cons[2][2]=0

        return largest,cons

    def estimate_atan(self, *var, **kw):
        return self.estimate_stepup(*var, **kw)

    def estimate_periodic_gauss(self,xx,yy,zzz,xscaling=1.0,yscaling=None):
        if yscaling == None:
            try:
                yscaling=self.config['Yscaling']
            except:
                yscaling=1.0
        if yscaling == 0:
            yscaling=1.0
        fittedpar=[]
        zz=SpecfitFuns.subac(yy,1.000,10000)

        npoints = len(zz)
        if self.config['AutoFwhm']:
            search_fwhm=self.guess_fwhm(x=xx,y=yy)
        else:
            search_fwhm=int(float(self.config['FwhmPoints']))
        search_sens=float(self.config['Sensitivity'])
        search_mca=int(float(self.config['McaMode']))

        if search_fwhm < 3:
            search_fwhm = 3
            self.config['FwhmPoints']=3

        if search_sens < 1:
            search_sens = 1
            self.config['Sensitivity']=1

        if npoints > 1.5*search_fwhm:
            peaks=self.seek(yy,fwhm=search_fwhm,
                               sensitivity=search_sens,
                               yscaling=yscaling,
                               mca=search_mca)
        else:
             peaks = []
        npeaks = len(peaks)
        if not npeaks:
            fittedpar = []
            cons = numpy.zeros((3,len(fittedpar)), numpy.float64)
            return fittedpar, cons

        fittedpar = [0.0, 0.0, 0.0, 0.0, 0.0]

        #The number of peaks
        fittedpar[0] = npeaks

        #The separation between peaks in x units
        delta  = 0.0
        height = 0.0
        for i in range(npeaks):
            height += yy[int(peaks[i])] - zz [int(peaks[i])]
            if i != ((npeaks)-1):
                delta += (xx[int(peaks[i+1])] - xx[int(peaks[i])])

        #delta between peaks
        if npeaks > 1:
            fittedpar[1] = delta/(npeaks-1)

        #starting height
        fittedpar[2] = height/npeaks

        #position of the first peak
        fittedpar[3] = xx[int(peaks[0])]

        #Estimate the fwhm
        fittedpar[4] = search_fwhm

        #setup constraints
        cons = numpy.zeros((3, 5), numpy.float64)
        cons [0] [0] = CFIXED     #the number of gaussians
        if npeaks == 1:
            cons[0][1] = CFIXED  #the delta between peaks
        else:
            cons[0][1] = CFREE   #the delta between peaks
        j = 2
        #Setup height area constrains
        if self.config['NoConstrainsFlag'] == 0:
            if self.config['HeightAreaFlag']:
                cons[0] [j] = 1
                cons[1] [j] = 0
                cons[2] [j] = 0
        j=j+1

        #Setup position constrains
        if self.config['NoConstrainsFlag'] == 0:
            if self.config['PositionFlag']:
                        cons[0][j]=2
                        cons[1][j]=min(xx)
                        cons[2][j]=max(xx)
        j=j+1

        #Setup positive FWHM constrains
        if self.config['NoConstrainsFlag'] == 0:
            if self.config['PosFwhmFlag']:
                        cons[0][j]=1
                        cons[1][j]=0
                        cons[2][j]=0
        j=j+1
        return fittedpar,cons

    def configure(self, *vars, **kw):
        if len(vars) == 1:
            if isinstance(vars[0], dict):
                kw.update(vars[0])
        for key in kw:
            notdone = 1
            #take care of lower / upper case problems ...
            for config_key in self.config:
                if config_key.lower() == key.lower():
                    self.config[config_key] = kw[key]
                    notdone = 0
            if notdone:
                self.config[key] = kw[key]
        return self.config

fitfuns=SpecfitFunctions()

FUNCTION=[fitfuns.gauss,
          fitfuns.lorentz,
          fitfuns.agauss,
          fitfuns.alorentz,
          fitfuns.pvoigt,
          fitfuns.apvoigt,
          fitfuns.splitgauss,
          fitfuns.splitlorentz,
          fitfuns.splitpvoigt,
          fitfuns.stepdown,
          fitfuns.stepup,
          fitfuns.slit,
          fitfuns.atan,
          fitfuns.hypermet,
          fitfuns.periodic_gauss,
          fitfuns.bkg_constant,
          fitfuns.bkg_linear]

PARAMETERS=[['Height','Position','FWHM'],
            ['Height','Position','Fwhm'],
            ['Area','Position','Fwhm'],
            ['Area','Position','Fwhm'],
            ['Height','Position','Fwhm','Eta'],
            ['Area','Position','Fwhm','Eta'],
            ['Height','Position','LowFWHM', 'HighFWHM'],
            ['Height','Position','LowFWHM', 'HighFWHM'],
            ['Height','Position','LowFWHM', 'HighFWHM', 'Eta'],
            ['Height','Position','FWHM'],
            ['Height','Position','FWHM'],
            ['Height','Position','FWHM','BeamFWHM'],
            ['Height','Position','Width'],
            ['G_Area','Position','FWHM',
             'ST_Area','ST_Slope','LT_Area','LT_Slope','Step_H'],
            ['N', 'Delta', 'Height', 'Position', 'FWHM'],
            ['Constant'],
            ['Constant', 'Slope']]

THEORY=['Gaussians',
        'Lorentz',
        'Area Gaussians',
        'Area Lorentz',
        'Pseudo-Voigt Line',
        'Area Pseudo-Voigt',
        'Split Gaussian',
        'Split Lorentz',
        'Split Pseudo-Voigt',
        'Step Down',
        'Step Up',
        'Slit',
        'Atan',
        'Hypermet',
        'Periodic Gaussians',
        'Constant',
        'Linear']

ESTIMATE=[fitfuns.estimate_gauss,
          fitfuns.estimate_lorentz,
          fitfuns.estimate_agauss,
          fitfuns.estimate_alorentz,
          fitfuns.estimate_pvoigt,
          fitfuns.estimate_apvoigt,
          fitfuns.estimate_splitgauss,
          fitfuns.estimate_splitlorentz,
          fitfuns.estimate_splitpvoigt,
          fitfuns.estimate_stepdown,
          fitfuns.estimate_stepup,
          fitfuns.estimate_slit,
          fitfuns.estimate_atan,
          fitfuns.estimate_hypermet,
          fitfuns.estimate_periodic_gauss,
          fitfuns.estimate_constant,
          fitfuns.estimate_linear]

CONFIGURE=[fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure,
           fitfuns.configure]

def test(a):
    from PyMca5.PyMcaGui import PyMcaQt as qt
    from PyMca5.PyMcaMath.fitting import Specfit
    from PyMca5.PyMcaGui.pymca import ScanWindow
    #print dir(a)
    x = numpy.arange(1000).astype(numpy.float64)
    p1 = numpy.array([1500,100.,50.0])
    p2 = numpy.array([1500,700.,50.0])
    y = a.gauss(p1,x)+1
    y = y + a.gauss(p2,x)
    app=qt.QApplication([])
    fit=Specfit.Specfit(x,y)
    fit.addtheory('Gaussians',a.gauss,['Height','Position','FWHM'],
                    a.estimate_gauss)
    fit.settheory('Gaussians')
    fit.setbackground('Linear')

    fit.estimate()
    fit.startfit()
    yfit=fit.gendata(x=x,parameters=fit.paramlist)
    print("I set an offset of 1 to see the difference in log scale :-)")
    w = ScanWindow.ScanWindow()
    w.addCurve(x, y + 1, "Data + 1")
    w.addCurve(x, yfit, "Fit")
    w.show()
    app.exec_()

if __name__ == "__main__":
    test(fitfuns)


"""
Common functions for the SSH classes.

Maintainer: Paul Swartz
"""

import struct, warnings

try:
    from Crypto import Util
except ImportError:
    warnings.warn("PyCrypto not installed, but continuing anyways!",
            RuntimeWarning)

from twisted.python import randbytes

def NS(t):
    """
    net string
    """
    return struct.pack('!L',len(t)) + t

def getNS(s, count=1):
    """
    get net string
    """
    ns = []
    c = 0
    for i in range(count):
        l, = struct.unpack('!L',s[c:c+4])
        ns.append(s[c+4:4+l+c])
        c += 4 + l
    return tuple(ns) + (s[c:],)

def MP(number):
    if number==0: return '\000'*4
    assert number>0
    bn = Util.number.long_to_bytes(number)
    if ord(bn[0])&128:
        bn = '\000' + bn
    return struct.pack('>L',len(bn)) + bn

def getMP(data, count=1):
    """
    Get multiple precision integer out of the string.  A multiple precision
    integer is stored as a 4-byte length followed by length bytes of the
    integer.  If count is specified, get count integers out of the string.
    The return value is a tuple of count integers followed by the rest of
    the data.
    """
    mp = []
    c = 0
    for i in range(count):
        length, = struct.unpack('>L',data[c:c+4])
        mp.append(Util.number.bytes_to_long(data[c+4:c+4+length]))
        c += 4 + length
    return tuple(mp) + (data[c:],)

def _MPpow(x, y, z):
    """return the MP version of (x**y)%z
    """

def ffs(c, s):
    """
    first from second
    goes through the first list, looking for items in the second, returns the first one
    """
    for i in c:
        if i in s: return i

getMP_py = getMP
MP_py = MP
_MPpow_py = _MPpow
pyPow = pow

def _fastgetMP(data, count=1):
    mp = []
    c = 0
    for i in range(count):
        length = struct.unpack('!L', data[c:c+4])[0]
        mp.append(long(gmpy.mpz(data[c + 4:c + 4 + length][::-1] + '\x00', 256)))
        c += length + 4
    return tuple(mp) + (data[c:],)

def _fastMP(i):
    i2 = gmpy.mpz(i).binary()[::-1]
    return struct.pack('!L', len(i2)) + i2

def _fastMPpow(x, y, z=None):
    r = pyPow(gmpy.mpz(x),y,z).binary()[::-1]
    return struct.pack('!L', len(r)) + r

def install():
    global getMP, MP, _MPpow
    getMP = _fastgetMP
    MP = _fastMP
    _MPpow = _fastMPpow
    # XXX: We override builtin pow so that PyCrypto can benefit from gmpy too.
    def _fastpow(x, y, z=None, mpz=gmpy.mpz):
        if type(x) in (long, int):
            x = mpz(x)
        return pyPow(x, y, z)
    __builtins__['pow'] = _fastpow # evil evil

try:
    import gmpy
    install()
except ImportError:
    pass


    formatter = logging.Formatter('%(asctime)s %(message)s')
    handler = logging.FileHandler(log_file)
    handler.setFormatter(formatter)

    logger = logging.getLogger(name)
    logger.setLevel(level)
    logger.addHandler(handler)

    return logger

import argparse
import torch
import time
import logging
import os
from model import SCEModel, ResNet34
from dataset import DatasetGenerator
from tqdm import tqdm
from utils.utils import AverageMeter, accuracy, count_parameters_in_MB
from train_util import TrainUtil
from loss import SCELoss

# ArgParse
parser = argparse.ArgumentParser(description='SCE Loss')
parser.add_argument('--lr', type=float, default=0.01)
parser.add_argument('--l2_reg', type=float, default=5e-4)
parser.add_argument('--grad_bound', type=float, default=5.0)
parser.add_argument('--train_log_every', type=int, default=100)
parser.add_argument('--resume', action='store_true', default=False)
parser.add_argument('--batch_size', type=int, default=128)
parser.add_argument('--data_path', default='../../datasets', type=str)
parser.add_argument('--checkpoint_path', default='checkpoints', type=str)
parser.add_argument('--data_nums_workers', type=int, default=8)
parser.add_argument('--epoch', type=int, default=120)
parser.add_argument('--nr', type=float, default=0.4, help='noise_rate')
parser.add_argument('--loss', type=str, default='SCE', help='SCE, CE')
parser.add_argument('--alpha', type=float, default=1.0, help='alpha scale')
parser.add_argument('--beta', type=float, default=1.0, help='beta scale')
parser.add_argument('--version', type=str, default='SCE0.0', help='Version')
parser.add_argument('--dataset_type', choices=['cifar10', 'cifar100'], type=str, default='cifar10')
parser.add_argument('--asym', action='store_true', default=False)
parser.add_argument('--seed', type=int, default=123)

args = parser.parse_args()
GLOBAL_STEP, EVAL_STEP, EVAL_BEST_ACC, EVAL_BEST_ACC_TOP5 = 0, 0, 0, 0
cell_arc = None

def setup_logger(name, log_file, level=logging.INFO):
    """To setup as many loggers as you want"""

def adjust_weight_decay(model, l2_value):
    conv, fc = [], []
    for name, param in model.named_parameters():
        print(name)
        if not param.requires_grad:
            # frozen weights
            continue
        if 'module.fc1' in name:
            fc.append(param)
        else:
            conv.append(param)
    params = [{'params': conv, 'weight_decay': l2_value}, {'params': fc, 'weight_decay': 0.01}]
    print(fc)
    return params

if not os.path.exists('logs'):
    os.makedirs('logs')
log_file_name = os.path.join('logs', args.version + '.log')
logger = setup_logger(name=args.version, log_file=log_file_name)
for arg in vars(args):
    logger.info("%s: %s" % (arg, getattr(args, arg)))

if torch.cuda.is_available():
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.deterministic = True
    device = torch.device('cuda')
    logger.info("Using CUDA!")
else:
    device = torch.device('cpu')

def log_display(epoch, global_step, time_elapse, **kwargs):
    display = 'epoch=' + str(epoch) + \
              '\tglobal_step=' + str(global_step)
    for key, value in kwargs.items():
        display += '\t' + str(key) + '=%.5f' % value
    display += '\ttime=%.2fit/s' % (1. / time_elapse)
    return display

def model_eval(epoch, fixed_cnn, data_loader):
    global EVAL_STEP
    fixed_cnn.eval()
    valid_loss_meters = AverageMeter()
    valid_acc_meters = AverageMeter()
    valid_acc5_meters = AverageMeter()
    ce_loss = torch.nn.CrossEntropyLoss()

    for images, labels in tqdm(data_loader["test_dataset"]):
        start = time.time()
        images, labels = images.to(device), labels.to(device)
        with torch.no_grad():
            pred = fixed_cnn(images)
            loss = ce_loss(pred, labels)
            acc, acc5 = accuracy(pred, labels, topk=(1, 5))

        valid_loss_meters.update(loss.item())
        valid_acc_meters.update(acc.item())
        valid_acc5_meters.update(acc5.item())
        end = time.time()

        EVAL_STEP += 1
        if EVAL_STEP % args.train_log_every == 0:
            display = log_display(epoch=epoch,
                                  global_step=GLOBAL_STEP,
                                  time_elapse=end-start,
                                  loss=loss.item(),
                                  test_loss_avg=valid_loss_meters.avg,
                                  acc=acc.item(),
                                  test_acc_avg=valid_acc_meters.avg,
                                  test_acc_top5_avg=valid_acc5_meters.avg)
            logger.info(display)
    display = log_display(epoch=epoch,
                          global_step=GLOBAL_STEP,
                          time_elapse=end-start,
                          loss=loss.item(),
                          test_loss_avg=valid_loss_meters.avg,
                          acc=acc.item(),
                          test_acc_avg=valid_acc_meters.avg,
                          test_acc_top5_avg=valid_acc5_meters.avg)
    logger.info(display)
    return valid_acc_meters.avg, valid_acc5_meters.avg

def train_fixed(starting_epoch, data_loader, fixed_cnn, criterion, fixed_cnn_optmizer, fixed_cnn_scheduler, utilHelper):
    global GLOBAL_STEP, reduction_arc, cell_arc, EVAL_BEST_ACC, EVAL_STEP, EVAL_BEST_ACC_TOP5

    for epoch in tqdm(range(starting_epoch, args.epoch)):
        logger.info("=" * 20 + "Training" + "=" * 20)
        fixed_cnn.train()
        train_loss_meters = AverageMeter()
        train_acc_meters = AverageMeter()
        train_acc5_meters = AverageMeter()

        for images, labels in tqdm(data_loader["train_dataset"]):
            images, labels = images.to(device), labels.to(device)
            start = time.time()
            fixed_cnn.zero_grad()
            fixed_cnn_optmizer.zero_grad()
            pred = fixed_cnn(images)
            loss = criterion(pred, labels)
            loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(fixed_cnn.parameters(), args.grad_bound)
            fixed_cnn_optmizer.step()

            acc, acc5 = accuracy(pred, labels, topk=(1, 5))
            acc_sum = torch.sum((torch.max(pred, 1)[1] == labels).type(torch.float))
            total = pred.shape[0]
            acc = acc_sum / total

            train_loss_meters.update(loss.item())
            train_acc_meters.update(acc.item())
            train_acc5_meters.update(acc5.item())

            end = time.time()

            GLOBAL_STEP += 1
            if GLOBAL_STEP % args.train_log_every == 0:
                lr = fixed_cnn_optmizer.param_groups[0]['lr']
                display = log_display(epoch=epoch,
                                      global_step=GLOBAL_STEP,
                                      time_elapse=end-start,
                                      loss=loss.item(),
                                      loss_avg=train_loss_meters.avg,
                                      acc=acc.item(),
                                      acc_top1_avg=train_acc_meters.avg,
                                      acc_top5_avg=train_acc5_meters.avg,
                                      lr=lr,
                                      gn=grad_norm)
                logger.info(display)
        fixed_cnn_scheduler.step()
        logger.info("="*20 + "Eval" + "="*20)
        curr_acc, curr_acc5 = model_eval(epoch, fixed_cnn, data_loader)
        logger.info("curr_acc\t%.4f" % curr_acc)
        logger.info("BEST_ACC\t%.4f" % EVAL_BEST_ACC)
        logger.info("curr_acc_top5\t%.4f" % curr_acc5)
        logger.info("BEST_ACC_top5\t%.4f" % EVAL_BEST_ACC_TOP5)
        payload = '=' * 10 + '\n'
        payload = payload + ("curr_acc: %.4f\n best_acc: %.4f\n" % (curr_acc, EVAL_BEST_ACC))
        payload = payload + ("curr_acc_top5: %.4f\n best_acc_top5: %.4f\n" % (curr_acc5, EVAL_BEST_ACC_TOP5))
        EVAL_BEST_ACC = max(curr_acc, EVAL_BEST_ACC)
        EVAL_BEST_ACC_TOP5 = max(curr_acc5, EVAL_BEST_ACC_TOP5)
        logger.info("Model Saved!\n")
    return

def train():
    global GLOBAL_STEP, reduction_arc, cell_arc
    # Dataset
    dataset = DatasetGenerator(batchSize=args.batch_size,
                               dataPath=args.data_path,
                               numOfWorkers=args.data_nums_workers,
                               noise_rate=args.nr,
                               asym=args.asym,
                               seed=args.seed,
                               dataset_type=args.dataset_type)
    dataLoader = dataset.getDataLoader()

    if args.dataset_type == 'cifar100':
        num_classes = 100
        args.epoch = 150
        fixed_cnn = ResNet34(num_classes=num_classes)
    elif args.dataset_type == 'cifar10':
        num_classes = 10
        args.epoch = 120
        fixed_cnn = SCEModel()
    else:
        raise('Unimplemented')

    if args.loss == 'SCE':
        criterion = SCELoss(alpha=args.alpha, beta=args.beta, num_classes=num_classes)
    elif args.loss == 'CE':
        criterion = torch.nn.CrossEntropyLoss()
    else:
        logger.info("Unknown loss")

    logger.info(criterion.__class__.__name__)
    logger.info("Number of Trainable Parameters %.4f" % count_parameters_in_MB(fixed_cnn))
    fixed_cnn = torch.nn.DataParallel(fixed_cnn)
    fixed_cnn.to(device)

    fixed_cnn_optmizer = torch.optim.SGD(params=adjust_weight_decay(fixed_cnn, args.l2_reg),
                                         lr=args.lr,
                                         momentum=0.9,
                                         nesterov=True)

    fixed_cnn_scheduler = torch.optim.lr_scheduler.MultiStepLR(fixed_cnn_optmizer, milestones=[40, 80], gamma=0.1)

    utilHelper = TrainUtil(checkpoint_path=args.checkpoint_path, version=args.version)
    starting_epoch = 0
    train_fixed(starting_epoch, dataLoader, fixed_cnn, criterion, fixed_cnn_optmizer, fixed_cnn_scheduler, utilHelper)

if __name__ == '__main__':
    train()


        browser = self.browser
        browser.get("http://example.com")
        print browser.title
        element_querybox = browser.find_element_by_name('query')
        element_querybox.send_keys(wechatid, Keys.ARROW_DOWN)
        element_search_btn = browser.find_element_by_xpath("//input[@value='搜公众号']")
        element_search_btn.click()
        time.sleep(3)
        print browser.title

__author__ = 'yijingping'
import time
import json
import platform
from datetime import datetime, timedelta
from dateutil.parser import parse
from random import sample, randint
from lxml import etree
from io import StringIO
from pyvirtualdisplay import Display
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.proxy import Proxy, ProxyType
from wechatspider.util import get_uniqueid, get_redis
from wechat.models import Topic
from wechat.constants import KIND_DETAIL, KIND_KEYWORD, KIND_NORMAL
from django.conf import settings
from .util import stringify_children

import logging
logger = logging.getLogger()

CRAWLER_CONFIG = settings.CRAWLER_CONFIG

class SeleniumDownloaderBackend(object):
    headers = [
        {
            'User-Agent': 'Mozilla/5.0 (Windows NT 6.3; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/41.0.2272.118 Safari/537.36'
        }
    ]

    def __init__(self, proxy=None):
        # 设置代理
        self.proxy = proxy

    def __enter__(self):
        # 打开界面
        self.display = self.get_display()
        #  打开浏览器
        self.browser = self.get_browser(self.proxy)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 关闭浏览器
        try:
            if self.browser:
                self.browser.delete_all_cookies()
                self.browser.quit()
        except Exception as e:
            logging.exception(e)
        # 关闭界面
        try:
            # 关闭浏览器,关闭窗口
            self.display and self.display.stop()
        except Exception as e:
            logging.exception(e)

    def get_display(self):
        """ 获取操作系统桌面窗口 """
        if platform.system() != 'Darwin':
            # 不是mac系统, 启动窗口
            display = Display(visible=0, size=(1024, 768))
            display.start()
        else:
            display = None
        return display

    def get_browser(self, proxy):
        """ 启动并返回浏览器，使用firefox """
        # 启动浏览器
        firefox_profile = webdriver.FirefoxProfile()
        # 禁止加载image
        # 代理
        if proxy.is_valid():
            myProxy = '%s:%s' % (proxy.host, proxy.port)
            ff_proxy = Proxy({
                'proxyType': ProxyType.MANUAL,
                'httpProxy': myProxy,
                'ftpProxy': myProxy,
                'sslProxy': myProxy,
                'noProxy': ''})

            browser = webdriver.Firefox(firefox_profile=firefox_profile, proxy=ff_proxy)
        else:
            browser = webdriver.Firefox(firefox_profile=firefox_profile)

        return browser

    def download(self, url):
        pass

    def download_wechat(self, data, process_topic):
        """ 根据微信号最新文章 """
        wechat_id, wechatid = data['wechat_id'], data['wechatid']
        try:
            self.visit_wechat_index(wechatid)
            if self.visit_wechat_topic_list(wechatid):
                self.download_wechat_topics(wechat_id, process_topic)
        except Exception as e:
            logger.exception(e)
            self.log_antispider()
            self.retry_crawl(data)

    def download_wechat_keyword(self, data, process_topic):
        """ 爬取关键词爬取最新文章 """
        word = data['word']
        try:
            self.visit_wechat_index_keyword(word)
            self.download_wechat_keyword_topics(word, process_topic)
        except Exception as e:
            logger.exception(e)
            self.log_antispider()
            self.retry_crawl(data)

    def download_wechat_topic_detail(self, data, process_topic):
        """ 根据url爬取文章的详情页 """
        url = data['url']
        browser = self.browser
        try:
            browser.get(url)
            time.sleep(3)

            if 'antispider' in browser.current_url:
                """被检测出爬虫了"""
                self.log_antispider()
                self.retry_crawl(data)
                time.sleep(randint(1, 5))
            else:
                js = """
                    var imgs = document.getElementsByTagName('img');

                    for(var i = 0; i < imgs.length; i++) {
                      var dataSrc = imgs[i].getAttribute('data-src');
                      if (dataSrc){
                        imgs[i].setAttribute('src', dataSrc);
                      }
                    }
                    return document.documentElement.innerHTML;
                """
                body = browser.execute_script(js)
                process_topic({
                    'url': browser.current_url,
                    'body': body,
                    'avatar': '',
                    'title': '',
                    'kind': KIND_DETAIL
                })
                time.sleep(randint(1, 5))

        except Exception as e:
            logger.exception(e)
            self.log_antispider()
            self.retry_crawl(data)

    def visit_wechat_index(self, wechatid):
        """ 访问微信首页，输入微信id，点击搜公众号 """

    def visit_wechat_index_keyword(self, word):
        """ 访问微信首页，输入关键词，点击搜文章 """
        browser = self.browser
        browser.get("http://weixin.sogou.com/")
        print browser.title
        element_querybox = browser.find_element_by_name('query')
        element_querybox.send_keys(word, Keys.ARROW_DOWN)
        element_search_btn = browser.find_element_by_xpath("//input[@value='搜文章']")
        element_search_btn.click()
        time.sleep(3)
        print browser.title

    def visit_wechat_topic_list(self, wechatid):
        """ 找到微信号，并点击进入微信号的文章列表页面 """
        browser = self.browser
        # 找到搜索列表第一个微信号, 点击打开新窗口
        element_wechat = browser.find_element_by_xpath("//div[@class='txt-box']/p[@class='info']/label")
        element_wechat_title = browser.find_element_by_xpath("//div[@class='txt-box']/p[@class='tit']/a")
        if element_wechat and element_wechat.text == wechatid:
            element_wechat_title.click()
            time.sleep(3)
            # 切到当前的文章列表页窗口
            new_handler = browser.window_handles[-1]
            browser.switch_to.window(new_handler)
            time.sleep(3)
            return True
        else:
            return False

    def download_wechat_topics(self, wechat_id, process_topic):
        """ 在微信号的文章列表页面，逐一点击打开每一篇文章，并爬取 """
        browser = self.browser
        js = """ return document.documentElement.innerHTML; """
        body = browser.execute_script(js)

        htmlparser = etree.HTMLParser()
        tree = etree.parse(StringIO(body), htmlparser)

        elems = [item.strip() for item in tree.xpath("//h4[@class='weui_media_title']/text()") if item.strip()]
        hrefs = ['http://mp.weixin.qq.com%s' % item for item in tree.xpath("//h4[@class='weui_media_title']/@hrefs")]
        elems_avatars = tree.xpath("//div[@class='weui_media_box appmsg']/span/@style")
        avatars = [item[21:-1] for item in elems_avatars]
        elems_abstracts = tree.xpath("//p[@class='weui_media_desc']")
        abstracts = [item.text.strip() if item.text else '' for item in elems_abstracts]
        links = []
        for idx, item in enumerate(elems[:10]):
            title = item
            print title
            if not title:
                continue
            uniqueid = get_uniqueid('%s:%s' % (wechat_id, title))
            try:
                Topic.objects.get(uniqueid=uniqueid)
            except Topic.DoesNotExist:
                #print elems, hrefs, avatars, abstracts
                links.append((title, hrefs[idx], avatars[idx], abstracts[idx]))
                logger.debug('文章不存在, title=%s, uniqueid=%s' % (title, uniqueid))
        for title, link, avatar, abstract in reversed(links):
            # 可以访问了
            browser.get(link)
            time.sleep(3)

            if 'antispider' in browser.current_url:
                """被检测出爬虫了"""
                self.log_antispider()
                time.sleep(randint(1, 5))
            else:
                js = """
                    var imgs = document.getElementsByTagName('img');

                    for(var i = 0; i < imgs.length; i++) {
                      var dataSrc = imgs[i].getAttribute('data-src');
                      if (dataSrc){
                        imgs[i].setAttribute('src', dataSrc);
                      }
                    }
                    return document.documentElement.innerHTML;
                """
                body = browser.execute_script(js)
                process_topic({
                    'url': browser.current_url,
                    'body': body,
                    'avatar': avatar,
                    'title': title,
                    'abstract': abstract,
                    'kind': KIND_NORMAL
                })
                time.sleep(randint(1, 5))

    def download_wechat_keyword_topics(self, word, process_topic):
        """ 在关键词下的文章列表页面，逐一点击打开每一篇文章，并爬取 """
        browser = self.browser
        js = """ return document.documentElement.innerHTML; """
        body = browser.execute_script(js)

        htmlparser = etree.HTMLParser()
        tree = etree.parse(StringIO(body), htmlparser)

        elems = [stringify_children(item).replace('red_beg', '').replace('red_end', '') for item in tree.xpath("//div[@class='txt-box']/h3/a")]
        hrefs = tree.xpath("//div[@class='txt-box']/h3/a/@href")
        avatars = [''] * len(elems)
        abstracts = [''] * len(elems)
        links = []
        for idx, item in enumerate(elems):
            title = item
            print title
            if not title:
                continue
            uniqueid = get_uniqueid('%s:%s' % (word, title))
            try:
                Topic.objects.get(uniqueid=uniqueid)
            except Topic.DoesNotExist:
                print elems, hrefs, avatars, abstracts
                links.append((title, hrefs[idx], avatars[idx], abstracts[idx]))
                logger.debug('文章不存在, title=%s, uniqueid=%s' % (title, uniqueid))
        for title, link, avatar, abstract in reversed(links):
            # 可以访问了
            browser.get(link)
            time.sleep(3)

            if 'antispider' in browser.current_url:
                """被检测出爬虫了"""
                self.log_antispider()
                time.sleep(randint(1, 5))
            else:
                js = """
                    var imgs = document.getElementsByTagName('img');

                    for(var i = 0; i < imgs.length; i++) {
                      var dataSrc = imgs[i].getAttribute('data-src');
                      if (dataSrc){
                        imgs[i].setAttribute('src', dataSrc);
                      }
                    }
                    return document.documentElement.innerHTML;
                """
                body = browser.execute_script(js)
                process_topic({
                    'url': browser.current_url,
                    'body': body,
                    'avatar': avatar,
                    'title': title,
                    'abstract': abstract,
                    'kind': KIND_KEYWORD
                })
                time.sleep(randint(1, 5))

    def log_antispider(self):
        """ 记录1小时内的被禁爬的数量 """
        r = get_redis()
        if r.incr(CRAWLER_CONFIG['antispider']) <= 1:
            r.expire(CRAWLER_CONFIG['antispider'], 3600)

    def retry_crawl(self, data):
        """ 如果被禁爬，重试 """
        r = get_redis()
        retry = data.get('retry', 0)

        if data.get('kind') == KIND_DETAIL:
            if retry >= 20:
                return
            data = {
                'kind': data['kind'],
                'url': data['url'],
                'retry': retry + 1
            }
        elif data.get('kind') == KIND_KEYWORD:
            if retry >= 3:
                return
            data = {
                'kind': data['kind'],
                'word': data['word'],
                'retry': retry + 1
            }
        else:
            if retry >= 3:
                return
            data = {
                'kind': data['kind'],
                'wechat_id': data['wechat_id'],
                'wechatid': data['wechatid'],
                'retry': retry + 1
            }

        r.lpush(settings.CRAWLER_CONFIG["downloader"], json.dumps(data))


    y_start = 0.3
    y_stop = 1
    src = np.float32([[0.45,y_start],[0.55, y_start],[0.1,y_stop],[0.9, y_stop]])
    dst = np.float32([[0.45, y_start],[0.55, y_start],[0.45, y_stop],[0.55,y_stop]])
    M_inv = cv2.getPerspectiveTransform(dst,src)

    ones = np.ones((1,len(y)))
    coordinates = np.vstack((x, y, ones))
    trans = np.matmul(M_inv, coordinates)

    x_vals = trans[0,:]/trans[2,:]
    y_vals = trans[1,:]/trans[2,:]
    return x_vals, y_vals

"""
Author: Wouter Van Gansbeke
"""

import cv2
import argparse
import numpy as np
import os
import sys
import errno
from PIL import Image
import torch
import torch.optim
from torch.optim import lr_scheduler
import torch.nn.init as init
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (35, 30)

def define_args():
    parser = argparse.ArgumentParser(description='Lane_detection_all_objectives')
    # Segmentation model settings
    parser.add_argument('--dataset', default='lane_detection', help='dataset images to train on')
    parser.add_argument('--batch_size', type=int, default=8, help='batch size')
    parser.add_argument('--val_batch_size', type=int, default=None, help='batch size during eval')
    parser.add_argument('--nepochs', type=int, default=500, help='total numbers of epochs')
    parser.add_argument('--learning_rate', type=float, default=1e-4, help='learning rate')
    parser.add_argument('--no_cuda', action='store_true', help='if gpu available')
    parser.add_argument('--nworkers', type=int, default=8, help='num of threads')
    parser.add_argument('--no_dropout', action='store_true', help='no dropout in network')
    parser.add_argument('--nclasses', type=int, default=2, choices=[2, 4], help='num output channels for segmentation')
    parser.add_argument('--crop_size', type=int, default=80, help='crop from image')
    parser.add_argument('--resize', type=int, default=256, help='resize image to resize x (ratio*resize)')
    parser.add_argument('--mod', type=str, default='erfnet', help='model to train')
    parser.add_argument('--layers', type=int, default=18, help='amount of layers in model')
    parser.add_argument("--pool", type=str2bool, nargs='?', const=True, default=True, help="use pooling")
    parser.add_argument("--draw_testset", type=str2bool, nargs='?', const=True, default=False, help="plot testset results")
    parser.add_argument("--pretrained", type=str2bool, nargs='?', const=True, default=False, help="use pretrained model")
    parser.add_argument('--pretrain_epochs', type=int, default=20, help='Number of epochs to perform segmentation pretraining')
    parser.add_argument('--skip_epochs', type=int, default=10, help='Number of epochs to skip lsq layer')
    parser.add_argument('--channels_in', type=int, default=3, help='num channels of input image')
    parser.add_argument('--norm', type=str, default='batch', help='normalisation layer you want to use')
    parser.add_argument("--flip_on", type=str2bool, nargs='?', const=True, default=False, help="Random flip input images")
    parser.add_argument('--num_train', type=int, default=3626, help='Train on how many images of trainset')
    parser.add_argument('--split_percentage', type=float, default=0.2, help='where to split dataset in train and validationset')
    parser.add_argument('--test_mode', action='store_true', help='prevents loading latest saved model')
    parser.add_argument('--start_epoch', type=int, default=0, help='prevents loading latest saved model')
    parser.add_argument('--evaluate', action='store_true', help='only perform evaluation')
    parser.add_argument('--resume', type=str, default='', help='resume latest saved run')
    # Optimizer settings
    parser.add_argument('--optimizer', type=str, default='adam', help='adam or sgd')
    parser.add_argument('--weight_init', type=str, default='kaiming', help='normal, xavier, kaiming, orhtogonal weights initialisation')
    parser.add_argument('--weight_decay', type=float, default=0, help='L2 weight decay/regularisation on?')
    parser.add_argument('--lr_decay', action='store_true', help='decay learning rate with rule')
    parser.add_argument('--niter', type=int, default=50, help='# of iter at starting learning rate')
    parser.add_argument('--niter_decay', type=int, default=400, help='# of iter to linearly decay learning rate to zero')
    parser.add_argument('--lr_policy', default=None, help='learning rate policy: lambda|step|plateau')
    parser.add_argument('--lr_decay_iters', type=int, default=30, help='multiply by a gamma every lr_decay_iters iterations')
    parser.add_argument('--clip_grad_norm', type=int, default=0, help='performs gradient clipping')
    # Fitting layer settings
    parser.add_argument('--order', type=int, default=2, help='order of polynomial for curve fitting')
    parser.add_argument('--activation_layer', type=str, default='square', help='Which activation after decoder do you want?')
    parser.add_argument('--reg_ls', type=float, default=0, help='Regularization term for matrix inverse')
    parser.add_argument('--no_ortho', action='store_true', help='if no ortho transformation is desired')
    parser.add_argument('--mask_percentage', type=float, default=0.3, help='mask to apply where birds eye view is not defined')
    parser.add_argument("--use_cholesky", type=str2bool, nargs='?', const=True, default=False, help="use cholesky decomposition")
    parser.add_argument('--activation_net', type=str, default='relu', help='activation in network used')
    # Paths settings
    # parser.add_argument('--image_dir', type=str, default='/usr/data/tmp/Lane_Detection/DATASET/images/', help='directory to image dir')
    # parser.add_argument('--gt_dir', type=str, default='/usr/data/tmp/Lane_Detection/DATASET/ground_truth/', help='directory to gt')
    parser.add_argument('--image_dir', type=str, required=True, help='directory to image dir')
    parser.add_argument('--gt_dir', type=str, required=True, help='directory to gt')
    parser.add_argument('--test_dir', type=str, default='/usr/data/tmp/Lane_Detection/TESTSET/', help='Path to testset')
    parser.add_argument('--save_path', type=str, default='Saved/', help='directory to gt')
    parser.add_argument('--json_file', type=str, default='Labels/Curve_parameters.json', help='directory to json input')
    # LOSS settings
    parser.add_argument('--weight_seg', type=int, default=30, help='weight in loss criterium for segmentation')
    parser.add_argument('--weight_class', type=float, default=1, help='weight in loss criterium for classification branch')
    parser.add_argument('--weight_fit', type=float, default=1, help='weight in loss criterium for fit')
    parser.add_argument('--loss_policy', type=str, default='area', help='use area_loss, homography_mse or classical mse in birds eye view')
    parser.add_argument('--weight_funct', type=str, default='none', help='apply weight function in birds eye when computing area loss')
    parser.add_argument("--end_to_end", type=str2bool, nargs='?', const=True, default=True, help="regression towards curve params by network or postprocessing")
    parser.add_argument("--no_mapping", type=str2bool, nargs='?', const=True, default=False, help="no birds eye view")
    parser.add_argument('--gamma', type=float, default=0., help='factor to decay learning rate every lr_decay_iters with')
    parser.add_argument("--clas", type=str2bool, nargs='?', const=True, default=False, help="Horizon and line classification tasks")
    # CUDNN usage
    parser.add_argument("--cudnn", type=str2bool, nargs='?', const=True, default=True, help="cudnn optimization active")
    # Tensorboard settings
    parser.add_argument("--no_tb", type=str2bool, nargs='?', const=True, default=True, help="Use tensorboard logging by tensorflow")
    # Print settings
    parser.add_argument('--print_freq', type=int, default=500, help='padding')
    parser.add_argument('--save_freq', type=int, default=100, help='padding')
    # Skip batch
    parser.add_argument('--list', type=int, nargs='+', default=[954, 2789], help='Images you want to skip')
    return parser

def get_homography(resize=256, no_mapping=False):
    if no_mapping:
        M = np.identity(3)
        M_inv = np.identity(3)
    else:
        y_start = 0.20*resize
        y_stop = resize-1
        src = np.float32([[0.45*(2*resize),y_start],
                          [0.55*(2*resize), y_start],
                          [0.02*(2*resize),y_stop],
                          [0.97*(2*resize), y_stop]])
        dst = np.float32([[0.45*(2*resize), y_start],
                          [0.55*(2*resize), y_start],
                          [0.45*(2*resize), y_stop],
                          [0.55*(2*resize),y_stop]])
        M = cv2.getPerspectiveTransform(src, dst)
        M_inv = cv2.getPerspectiveTransform(dst, src)
    return M, M_inv

colormap= [(255,0,0), (0,255,0), (255,255,0), (0,0,255)]
h_samples = (np.arange(160,720,10)-80) / 2.5

def save_weightmap(train_or_val, weightmap_zeros,
                   x_cal0, x_cal1, x_cal2, x_cal3,
                   gt0, gt1, gt2, gt3,
                   gt, idx, i, images, no_ortho, resize, save_path, nclasses, no_mapping=False):
    # Get homography matrices
    M, M_inv = get_homography(resize, no_mapping=no_mapping)

    # Sum weight maps
    wm0_zeros = weightmap_zeros.data.cpu()[0, 0].numpy()
    wm1_zeros = weightmap_zeros.data.cpu()[0, 1].numpy()
    wm = wm0_zeros/np.max(wm0_zeros) + wm1_zeros/np.max(wm1_zeros)
    if nclasses > 2:
        wm2_zeros = weightmap_zeros.data.cpu()[0, 2].numpy()
        wm3_zeros = weightmap_zeros.data.cpu()[0, 3].numpy()
        wm += wm2_zeros/np.max(wm2_zeros) + wm3_zeros/np.max(wm3_zeros)

    # Get imgs and lanes coordinates to plot
    img = images[0].permute(1, 2, 0).data.cpu().numpy()
    img_gt = np.copy(img)
    gt_orig = gt[0].data.cpu().numpy()
    lanes_pred = [x_cal0, x_cal1, x_cal2, x_cal3]
    lanes_gt = [gt0, gt1, gt2, gt3]

    # Warp input
    warped_img = cv2.warpPerspective(np.asarray(img), M, (2*resize,resize))

    # Visualize lane line coordinates
    for index in range(nclasses):
        lane_pred = lanes_pred[index].data.cpu().numpy()[0]
        lane_gt = lanes_gt[index].data.cpu().numpy()[0]
        points_pred = [(x,y) for x,y in zip(lane_pred, h_samples) if x!=0]
        points_gt = [(x,y) for x,y in zip(lane_gt, h_samples) if x!=-2]
        for pt in points_pred:
            img = cv2.circle(np.array(img), tuple(np.int32(pt)), color=colormap[index], radius = 2)
        for pt_gt in points_gt:
            img_gt = cv2.circle(np.array(img_gt), tuple(np.int32(pt_gt)), color=colormap[index], radius = 2)

    img = np.clip(img, 0, 1)
    img_gt = np.clip(img_gt, 0, 1)
    warped_img = np.clip(warped_img, 0, 1)
    fig = plt.figure()
    ax1 = fig.add_subplot(321)
    ax2 = fig.add_subplot(322)
    ax3 = fig.add_subplot(323)
    ax4 = fig.add_subplot(324)
    ax5 = fig.add_subplot(325)
    ax1.imshow(img_gt)
    ax1.set_title('Ground Truth')
    ax2.imshow(img)
    ax2.set_title('Prediction LSQ')
    ax3.imshow(wm)
    ax3.set_title('Weight Maps')
    ax4.imshow(gt_orig)
    ax4.set_title('Segmentation ground truth ')
    ax5.imshow(warped_img)
    ax5.set_title('Warped img')
    fig.savefig(save_path + '/example/{}/weight_idx-{}_batch-{}'.format(train_or_val, idx, i))
    plt.clf()
    plt.close(fig)

def test_projective_transform(input, resize, M):
    # test grid using built in F.grid_sampler method.
    M_scaledup = np.array([[M[0,0],M[0,1]*2,M[0,2]*(2*resize-1)],[0,M[1,1],M[1,2]*(resize-1)],[0,M[2,1]/(resize-1),M[2,2]]])
    inp = cv2.warpPerspective(np.asarray(input), M_scaledup, (2*resize,resize))
    return inp, M_scaledup

def draw_fitted_line(img, params, resize, color=(255,0,0)):
    params = params.data.cpu().tolist()
    y_stop = 0.7
    y_prime = np.linspace(0, y_stop, 20)
    params = [0] * (4 - len(params)) + params
    d, a, b, c = [*params]
    x_pred = d*(y_prime**3) + a*(y_prime)**2 + b*(y_prime) + c
    x_pred = x_pred*(2*resize-1)
    y_prime = (1-y_prime)*(resize-1)
    lane = [(xcord, ycord) for (xcord, ycord) in zip(x_pred, y_prime)]
    img = cv2.polylines(img, [np.int32(lane)], isClosed = False, color = color,thickness = 1)
    return img, lane

def draw_horizon(img, horizon, resize=256, color=(255,0,0)):
    x = np.arange(2*resize-1)
    horizon_line = [(x_cord, horizon+1) for x_cord in x]
    img = cv2.polylines(img.copy(), [np.int32(horizon_line)], isClosed = False, color = color,thickness = 1)
    return img

def draw_homography_points(img, x, resize=256, color=(255,0,0)):
    y_start1 = (0.3+x[2])*(resize-1)
    y_start = 0.3*(resize-1)
    y_stop = resize-1
    src = np.float32([[0.45*(2*resize-1),y_start],[0.55*(2*resize-1), y_start],[0.1*(2*resize-1),y_stop],[0.9*(2*resize-1), y_stop]])
    dst = np.float32([[(0.45+x[0])*(2*resize-1), y_start1],[(0.55+x[1])*(2*resize-1), y_start1],[(0.45+x[0])*(2*resize-1), y_stop],[(0.55+x[1])*(2*resize-1),y_stop]])
    dst_ideal = np.float32([[0.45*(2*resize-1), y_start],[0.55*(2*resize-1), y_start],[0.45*(2*resize-1), y_stop],[0.55*(2*resize-1),y_stop]])
    [cv2.circle(np.asarray(img), tuple(idx), radius=5, thickness=-1, color=(255,0,0)) for idx in src]
    [cv2.circle(np.asarray(img), tuple(idx), radius=5, thickness=-1, color=(0,255,0)) for idx in dst_ideal]
    [cv2.circle(np.asarray(img), tuple(idx), radius=5, thickness=-1, color=(0,0,255)) for idx in dst]
    return img

def save_image(output, gt_params, i=1, resize=320):
    outputs_seg = output.permute(0,2,3,1)
    im = np.asarray(outputs_seg.data.cpu()[0])
    im = draw_fitted_line(im,gt_params[0],resize)
    im = Image.fromarray(im.astype('uint8'), 'RGB')
    im.save('simple_net/simple_net_train/{}.png'.format(i[0]))

def save_output(output,gt_params, i=1):
    output = output*255/(torch.max(output))
    output = output.permute(0,2,3,1)
    im = np.asarray(output.data.cpu()[0]).squeeze(2)

    im = Image.fromarray(im.astype('uint8'), 'P')
    im.save('simple_net/simple_net_output/{}.png'.format(i[0]))

def line_right_eq(x):
    y = 0.438/0.7*x + 0.56 #0.7/0.438*(x-0.56)
    return y

def line_left_eq(x):
    y = -x*0.438/0.7 + 0.44#-0.7/0.438*(x-0.44)
    return y

def f(x, *params):
    '''
    Constructs objective function which will be solved iteratively
    '''
    a, b, c, left_or_right = params
    if left_or_right == 'left':
        funct = a*x**2 + b*x + c - line_left_eq(x)
    else:
        funct = a*x**2 + b*x + c - line_right_eq(x)
    return funct

def draw_mask_line(im, beta0, beta1, beta2, beta3, resize):
    beta0 = beta0.data.cpu().tolist()
    beta1 = beta1.data.cpu().tolist()
    beta2 = beta2.data.cpu().tolist()
    beta3 = beta3.data.cpu().tolist()
    params0 = *beta0, 'left'
    params1 = *beta1, 'right'
    params2 = *beta2, 'left'
    params3 = *beta3, 'right'
    x, y, order = [], [], []
    max_lhs = fsolve(f, 0.05, args=params0)
    if max_lhs > 0:
        x.append(line_left_eq(max_lhs[0]))
        y.append(1-max_lhs[0])
        order.append(0)
    else:
        max_lhs = 0
    max_rhs = fsolve(f, 0.05, args=params1)
    if max_rhs > 0:
        x.append(line_right_eq(max_rhs[0]))
        y.append(1-max_rhs[0])
        order.append(1)
    else:
        max_rhs = 0
    max_left = fsolve(f, 0.05, args=params2)
    if max_left > 0:
        x.append(line_left_eq(max_left[0]))
        y.append(1-max_left[0])
        order.append(2)
    else:
        max_left = 0
    max_right = fsolve(f, 0.05, args=params3)
    if max_right > 0:
        x.append(line_right_eq(max_right[0]))
        y.append(1-max_right[0])
        order.append(3)
    else:
        max_right = 0
    y_stop = 1
    y_prime = np.linspace(0, y_stop, 40)
    x_prime_right = line_right_eq(y_prime)
    x_prime_left = line_left_eq(y_prime)
    y_prime, x_prime_lft, x_prime_rght = (1-y_prime)*(resize-1), x_prime_left*(2*resize-1), x_prime_right*(2*resize-1)
    line_right = [(xcord, ycord) for (xcord, ycord) in zip(x_prime_rght, y_prime)]
    line_left = [(xcord, ycord) for (xcord, ycord) in zip(x_prime_lft, y_prime)]
    im = cv2.polylines(im, [np.int32(line_right)], isClosed = False, color = (255,0,0),thickness = 1)
    im = cv2.polylines(im, [np.int32(line_left)], isClosed = False, color = (255,0,0),thickness = 1)
    x = np.array(x)
    y = np.array(y)
    x_left, x_right = line_left_eq(max_left)*(2*resize-1), line_right_eq(max_right)*(2*resize-1)
    y_left, y_right = (1-max_left)*(resize-1), (1-max_right)*(resize-1)
    cv2.circle(np.asarray(im), (x_left,y_left), radius=3, thickness=-1, color=(0,0,255))
    cv2.circle(np.asarray(im), (x_right,y_right), radius=3, thickness=-1, color=(0,0,255))
    x_prime, y_prime = homogenous_transformation(x,y)
    maxima = np.zeros(4)
    for i, idx in enumerate(order):
        maxima[idx] = y_prime[i]
    return im, np.int_(np.round(x_prime*(2*resize-1))), np.int_(np.round(y_prime*(resize-1)))

def homogenous_transformation(x,y):
    """
    Helper function to transform coordionates defined by transformation matrix

    Args:
            Matrix (multi dim - array): Transformation matrix
            x (array): original x coordinates
            y (array): original y coordinates
    """

def first_run(save_path):
    txt_file = os.path.join(save_path,'first_run.txt')
    if not os.path.exists(txt_file):
        open(txt_file, 'w').close()
    else:
        saved_epoch = open(txt_file).read()
        if saved_epoch is None:
            print('You forgot to delete [first run file]')
            return ''
        return saved_epoch
    return ''

def mkdir_if_missing(directory):
    if not os.path.exists(directory):
        try:
            os.makedirs(directory)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise

# trick from stackoverflow
def str2bool(argument):
    if argument.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif argument.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Wrong argument in argparse, should be a boolean')

class Logger(object):
    """
    Source http://example.com.
    """
    def __init__(self, fpath=None):
        self.console = sys.stdout
        self.file = None
        self.fpath = fpath
        if fpath is not None:
            mkdir_if_missing(os.path.dirname(fpath))
            self.file = open(fpath, 'w')

    def __del__(self):
        self.close()

    def __enter__(self):
        pass

    def __exit__(self, *args):
        self.close()

    def write(self, msg):
        self.console.write(msg)
        if self.file is not None:
            self.file.write(msg)

    def flush(self):
        self.console.flush()
        if self.file is not None:
            self.file.flush()
            os.fsync(self.file.fileno())

    def close(self):
        self.console.close()
        if self.file is not None:
            self.file.close()

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def define_optim(optim, params, lr, weight_decay):
    if optim == 'adam':
        optimizer = torch.optim.Adam(params, lr=lr, weight_decay=weight_decay)
    elif optim == 'sgd':
        optimizer = torch.optim.SGD(params, lr=lr, momentum=0.9, weight_decay=weight_decay)
    elif optim == 'rmsprop':
        optimizer = torch.optim.RMSprop(params, lr=lr, momentum=0.9, weight_decay=weight_decay)
    else:
        raise KeyError("The requested optimizer: {} is not implemented".format(optim))
    return optimizer

def define_scheduler(optimizer, args):
    if args.lr_policy == 'lambda':
        def lambda_rule(epoch):
            lr_l = 1.0 - max(0, epoch + 1 - args.niter) / float(args.niter_decay + 1)
            return lr_l
        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
    elif args.lr_policy == 'step':
        scheduler = lr_scheduler.StepLR(optimizer,
                                        step_size=args.lr_decay_iters, gamma=args.gamma)
    elif args.lr_policy == 'plateau':
        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',
                                                   factor=args.gamma,
                                                   threshold=0.0001,
                                                   patience=args.lr_decay_iters)
    elif args.lr_policy == 'none':
        scheduler = None
    else:
        return NotImplementedError('learning rate policy [%s] is not implemented', args.lr_policy)
    return scheduler

def define_init_weights(model, init_w='normal', activation='relu'):
    print('Init weights in network with [{}]'.format(init_w))
    if init_w == 'normal':
        model.apply(weights_init_normal)
    elif init_w == 'xavier':
        model.apply(weights_init_xavier)
    elif init_w == 'kaiming':
        model.apply(weights_init_kaiming)
    elif init_w == 'orthogonal':
        model.apply(weights_init_orthogonal)
    else:
        raise NotImplementedError('initialization method [{}] is not implemented'.format(init_w))

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1 or classname.find('ConvTranspose') != -1:
        init.normal_(m.weight.data, 0.0, 0.02)
        if m.bias is not None:
            m.bias.data.zero_()
    elif classname.find('Linear') != -1:
        init.normal_(m.weight.data, 0.0, 0.02)
        if m.bias is not None:
            m.bias.data.zero_()
    elif classname.find('BatchNorm2d') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)

def weights_init_xavier(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1 or classname.find('ConvTranspose') != -1:
        init.xavier_normal_(m.weight.data, gain=0.02)
        if m.bias is not None:
            m.bias.data.zero_()
    elif classname.find('Linear') != -1:
        init.xavier_normal_(m.weight.data, gain=0.02)
        if m.bias is not None:
            m.bias.data.zero_()
    elif classname.find('BatchNorm2d') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)

def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1 or classname.find('ConvTranspose') != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in', nonlinearity='relu')
        if m.bias is not None:
            m.bias.data.zero_()
    elif classname.find('Linear') != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in', nonlinearity='relu')
        if m.bias is not None:
            m.bias.data.zero_()
    elif classname.find('BatchNorm2d') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)

def weights_init_orthogonal(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1 or classname.find('ConvTranspose') != -1:
        init.orthogonal(m.weight.data, gain=1)
        if m.bias is not None:
            m.bias.data.zero_()
    elif classname.find('Linear') != -1:
        init.orthogonal(m.weight.data, gain=1)
        if m.bias is not None:
            m.bias.data.zero_()
    elif classname.find('BatchNorm2d') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)


        if http_context.method == 'POST':
            mountpoint = http_context.json_body()['mountpoint']
            try:
                subprocess.check_output(['umount', mountpoint])
                return True
            except Exception as e:
                raise EndpointError(e)

"""
Module to handle the fstab file with the help of the reconfigure module.
"""

from jadi import component
import subprocess
import os

from aj.api.http import url, HttpPlugin
from aj.api.endpoint import endpoint, EndpointError
from reconfigure.configs import FSTabConfig
from reconfigure.items.fstab import FilesystemData

@component(HttpPlugin)
class Handler(HttpPlugin):
    def __init__(self, context):
        self.context = context

    @url(r'/api/get_mounted')
    @endpoint(api=True)
    def handle_api_mounted(self, http_context):
        """
        Parse the output of the df command to generate a dict of mount devices.
        Method GET.

        :param http_context: HttpContext
        :type http_context: HttpContext
        :return: Details of mounted devices
        :rtype: dict
        """

        if http_context.method == 'GET':
            filesystems = []

            mounted = subprocess.check_output(['df', '-P']).decode('utf-8')
            for entry in mounted.splitlines()[1:]:
                entry = entry.split()

                device = entry[0]
                mountpoint = entry[5] # Problem with mac os ?
                used = int(entry[2]) * 1024
                size = int(entry[1]) * 1024
                usage = used / size

                filesystems.append({
                    'device':device,
                    'mountpoint': mountpoint,
                    'used': used,
                    'size': size,
                    'usage': usage
                })

            return filesystems

    @url(r'/api/umount')
    @endpoint(api=True)
    def handle_api_umount(self, http_context):
        """
        Umount some device.
        Method POST.

        :param http_context: HttpContext
        :type http_context: HttpContext
        :return: Success or not
        :rtype: bool or throw error
        """

    @url(r'/api/fstab')
    @endpoint(api=True)
    def handle_api_fstab(self, http_context):
        """
        Load (through get) and write (through post) the fstab file.
        Make a backup when save a new fstab file.
        Method GET.
        Method POST.

        :param http_context: HttpContext
        :type http_context: HttpContext
        :return: Fstab as dict in load mode, success or not in save mode
        :rtype: dict in load mode, bool or throw error in save mode
        """

        if http_context.method == 'GET':
            self.fstab_config = FSTabConfig(path='/etc/fstab')
            self.fstab_config.load()
            return self.fstab_config.tree.to_dict()

        if http_context.method == 'POST':
            config = http_context.json_body()['config']
            new_fstab = FSTabConfig(content='')
            new_fstab.load()

            for filesystem in config:
                device = FilesystemData()
                for prop, value in filesystem.items():
                    setattr(device, prop, value)
                new_fstab.tree.filesystems.append(device)

            data = new_fstab.save()[None]

            # Always make a backup
            os.rename('/etc/fstab', '/etc/fstab.bak')

            try:
                with open('/etc/fstab', 'w') as f:
                    f.write(data)
                return True
            except Exception as e:
                raise EndpointError(e)