from __future__ import print_function
import Manyo as mm
import Cmm
import scipy
import scipy.optimize
import numpy
#-------------------------------------------------------------------------------------------------
def GetHistogram():
    return 0
#-------------------------------------------------------------------------------------------------
def MovingAverage(ECM,sample):
    ecm=mm.ElementContainerMatrix()
    ecaSize = ECM.PutTableSize()
    for i in range(ecaSize):
        eca=MovingAverageECA(ECM.PutPointer(i),sample)
        ecm.Add(eca)
    return ecm
#-------------------------------------------------------------------------------------------------
def CalcTOFOffset(ECM, energyMin, energyMax,
                       energyLimitMin=-1, energyLimitMax=-1,
                       peakLimitMin=-1, peakLimitMax=-1,
                       outputFilename="tofOffset.txt", failedFilename="failedData.txt"):
    offsetArray,failedData = CalcOriginOfEnergy(ECM,energyMin,energyMax,
                                   energyLimitMin,energyLimitMax,peakLimitMin,peakLimitMax)

    WriteFitFailedData(failedFilename, failedData)
    WriteTOFOffset(outputFilename, offsetArray)
#-------------------------------------------------------------------------------------------------
def MovingAverageECA(ECA,sample):
    if sample >1:

        eca    = mm.ElementContainerArray(ECA.PutHeaderPointer())
        ecSize = ECA.PutTableSize()
        ec     = ECA.PutPointer(0)
        xkey,ykey,ekey    = [ec.PutXKey(),ec.PutYKey(),ec.PutEKey()]
        xunit,yunit,eunit = [ec.PutUnit(key) for key in [xkey,ykey,ekey]]


        totalCount=ec.PutHeaderPointer().PutDouble("TotalCounts")
        for k in range(1,sample-1):
            ecTmp = ECA.PutPointer(k)
            ec = ec + ecTmp
            totalCount += ecTmp.PutHeaderPointer().PutDouble("TotalCounts")


        for j in range(ecSize-sample):
            ecj         = ECA.PutPointer(j)
            ech         = ecj.PutHeaderPointer()
            ecjCount    = ech.PutDouble("TotalCounts")

            ecTmp       = ECA.PutPointer(sample-1+j)
            ec          = ec + ecTmp
            totalCount += ecTmp.PutHeaderPointer().PutDouble("TotalCounts")

            ech.OverWrite("TotalCounts",totalCount)
            ec.InputHeader(ech)
            ec.SetUnit(xkey,xunit)
            ec.SetUnit(ykey,yunit)
            ec.SetUnit(ekey,eunit)
            eca.Add(ec)

            ec          = ec - ecj
            totalCount -= ecjCount

        ec          = ec + ecj
        ec.SetUnit(xkey,xunit)
        ec.SetUnit(ykey,yunit)
        ec.SetUnit(ekey,eunit)
        totalCount += ecjCount

        for l in range(ecSize-sample,ecSize):
            ech= ECA.PutPointer(l).PutHeaderPointer()
            ech.OverWrite("TotalCounts",totalCount)
            #ec.InputHeader(ech)
            eca.Add(ec)
    else:
        eca=ECA
    return eca

#-------------------------------------------------------------------------------------------------
def PutValArray(ec,energyMin,energyMax):
    ecXn = numpy.array(ec.PutXList())
    ecX  = (ecXn[1:]+ecXn[:-1])*0.5
    ecY  = numpy.array(ec.PutYList())

    minReginArray = ecX>energyMin
    maxReginArray = ecX<energyMax
    regionArray   = numpy.logical_and(minReginArray,maxReginArray)

    ecX           = ecX[regionArray]
    ecY           = ecY[regionArray]
    return ecX,ecY
#-------------------------------------------------------------------------------------------------
def CalcOriginOfEnergy(ECM,energyMin,energyMax,
                       energyLimitMin=-1,energyLimitMax=-1,peakLimitMin=-1,peakLimitMax=-1):
    ecaSize = ECM.PutTableSize()
    failedData=[]
    offsetPsds=[]
    xarray,yarray = PutValArray(ECM.Put(0).Put(0),energyMin,energyMax)
    for i in range(ecaSize):
        ECA=ECM.PutPointer(i)
        ecSize=ECA.PutTableSize()
        offsetPixels=[]
        for j in range(ecSize):
            EC=ECA.PutPointer(j)

            mask = EC.PutHeaderPointer().PutInt4("MASKED")
            if mask==1:
                offset   =  0.0
            else:
                xarray,yarray = PutValArray(EC,energyMin,energyMax)
                guess = EstimateFittingInitial(xarray,yarray)
                try:
                    param = ExecGaussFit(xarray,yarray,guess)
                    flag,msg = CheckLimitCondition(param,energyLimitMin,energyLimitMax,
                                                         peakLimitMin,  peakLimitMax)
                    if flag:
                        offset = CalcTOFOffsetVal(EC,param[1])
                    else:
                        print("   Out of Range [",msg,"] (psd=%d,pixel=%d)"%(i,j))
                        failedData.append([i,j,xarray,yarray])
                        offset   = 0.0

                except Exception as e:
                    if "Number of calls to function has reached maxfev" in str(e):
                        print("   Number of calls to function has reached maxfev (psd=%d,pixel=%d)"%(i,j))
                    else:
                        print("   "+str(e)+"(psd=%d,pixel=%d)"%(i,j))
                    failedData.append([i,j,xarray,yarray])
                    offset   =  0.0
            offsetPixels.append(offset)
        offsetPsds.append(offsetPixels)
    return offsetPsds,failedData
#-------------------------------------------------------------------------------------------------
def WriteTOFOffset(filename,offset_array):
    f=open(filename,"w")
    for array in offset_array:
        text = ""
        for val in array:
            text += "%.4f "%val
        text = text.strip()+"\n"
        f.write(text)
    f.close()
#-------------------------------------------------------------------------------------------------
def WriteFitFailedData(filename,failedData):

    f=open(filename,"w")
    for psd,pixel,dataX,dataY in failedData:
        text="# psd:%d, pixel:%d"%(psd,pixel)+"\n"
        f.write(text)

        for x,y in zip(dataX,dataY):
            text ="%f %f\n"%(x,y)
            f.write(text)

        text ="\n"
        f.write(text)
    f.close()
#-------------------------------------------------------------------------------------------------
def CheckLimitCondition(param,energyLimitMin=-1,energyLimitMax=-1,peakLimitMin=-1,peakLimitMax=-1):
    if energyLimitMin!=energyLimitMax:
        if param[1]<energyLimitMin:
            return False,"Lower Energy"
        if param[1]>energyLimitMax:
            return False,"Higher Energy"
    if peakLimitMin!=peakLimitMax:
        if param[0]<peakLimitMin:
            return False,"Lower Peak"
        if param[0]>peakLimitMax:
            return False,"Higher Peak"

    return True,""
#-------------------------------------------------------------------------------------------------
def GaussFunction(x,height,center,sigma,constant):

    return height*numpy.exp( -0.5*((x-center)/sigma)**2 ) + constant
#-------------------------------------------------------------------------------------------------
def ExecGaussFit(xarray,yarray,guess):
    #popt,pconv = scipy.optimize.curve_fit(GaussFunction,xarray,yarray,guess,maxfev=1000)
    popt,pconv = scipy.optimize.curve_fit(GaussFunction,xarray,yarray,guess)

    return popt
#-------------------------------------------------------------------------------------------------
def ExecGaussFitEC(ec,energyMin,energyMax,height,center,fwhm,back):
    ecX,ecY=PutValArray(ec,energyMin,energyMax)
    if height<0 and center<0 and fwhm<0 and back<0:
        height,center,fwhm,back = EstimateFittingInitial(ecX,ecY)

    popt = ExecGaussFit(ecX,ecY,[height,center,fwhm,back])
    newX = numpy.arange(ecX[0],ecX[-1],(ecX[-1]-ecX[0])/1000.0)
    newY = GaussFunction(newX,popt[0],popt[1],popt[2],popt[3])
    ecOut= mm.ElementContainer()

    newX -= (newX[1] - newX[0])*0.5
    newX = numpy.hstack([newX,numpy.array([newX[-1]+(newX[1] - newX[0])])])
    ecOut.Add("x",newX.tolist())
    ecOut.Add("y",newY.tolist())
    ecOut.Add("e",numpy.zeros(len(newY)).tolist())
    ecOut.SetKeys("x","y","e")
    return ecOut
#-------------------------------------------------------------------------------------------------
def CalcTOFOffsetVal(EC,dHW):
    ech=EC.PutHeaderPointer()

    ef = ech.PutDouble("Ef")
    l1 = ech.PutDouble("L1")
    #l2 = ech.PutDouble("L2")
    #l3 = ech.PutDouble("L3")
    TOF_offset = 2286.2869 * l1 * ( 1.0/numpy.sqrt(ef) - 1.0/numpy.sqrt(ef+dHW) )

    return TOF_offset
#-------------------------------------------------------------------------------------------------
# Estimate Gaussian Function
#-------------------------------------------------------------------------------------------------
def EstimateFittingInitial(x,y):
    sumVal     = GetIntegral(x,y)
    maxY,maxX  = GetMaximum(x,y)
    binwidth   = GetBinWidth(x,y)
    fwhm       = GetFWHM(sumVal,maxY,binwidth)
    background = 0.0
    return maxY,maxX,fwhm,background
#-------------------------------------------------------------------------------------------------
def GetBinWidth(x,y):
    #binwidth=x[1]-x[0]
    binwidth = numpy.float(x.max()-x.min())/x.size
    return binwidth
#-------------------------------------------------------------------------------------------------
def GetIntegral(x,y):
    return float(y.sum())
#-------------------------------------------------------------------------------------------------
def GetMaximum(x,y):
    #z=zip(y,x)
    #return max(z)
    sumY=y.sum()
    if sumY==0:
        return 0.0,0.0
    else:
        return numpy.float(y.max()),numpy.float((x*y).sum())/y.sum()
#-------------------------------------------------------------------------------------------------
def GetFWHM(area,height,binwidth):
    if height==0:
        return 1
    else:
        return area/height/numpy.sqrt(2*numpy.pi)*numpy.sqrt(2*numpy.log(2))*(binwidth)
#-------------------------------------------------------------------------------------------------
