from Manyo import HeaderBase, ElementContainer, ElementContainerArray, ElementContainerMatrix, Int4Vector,\
                StringVector, DoubleVector ##, DoubleVectorContainer

from Manyo import SearchInHeader

from Manyo import ListToDoubleVector as l2d
import os, sys, string
import numpy

from sub1 import simplehistogram, ecsub, eccopy


def ChopperMaskAuto(ecm, lowerlimit = 1.0e-4, upperlimit = 1.0e4, wrinteonfile = True, filename = None):
    """
    This function make masking file automatically.
    At first, it calculate the sum of each ElementContainer and find their medium(not average) value.
    Then if the sum of intensity is less than lowerlimit x medium or more than upperlimit x medium, 
    the index of the ElementContainer in ElementContainerMatrix is stored.
    Finally, the indexes can be stored in file if @writeonfile == True. If masking file is not appointed as function arguement,
    the default filname is "YYYYMMDD_HHMMSS.msk". (year, month, day, hour, minute, second)
    If @writeonfile is not True, it returns its masking table.
    """
    
    
    flux={}
    tableindex=[]
    maskedindex=[]
    ecanumber = ecm.PutTableSize()
    for ecaindex in range(ecanumber):
        ecnumber = ecm(ecaindex).PutTableSize() 
        for ecindex in range(ecnumber):
            tempec=ecm(ecindex, ecaindex)
            #tableindex.append((ecaindex, ecindex))
            flux[(ecaindex, ecindex)]=tempec.Sum(tempec.PutYKey())
    flux0 = numpy.array(flux.values(), dtype='float')
    flux0.sort()
    flux_medium = flux0[len(flux0)/2-1]
    #tableindex = tuple(tableindex)
    for index0 in flux.keys():
        if flux[index0] <= lowerlimit*flux_medium or flux[index0] >= upperlimit*flux_medium :
            maskedindex.append(index0)
    
    maskedindex.sort()    
    if writeonfile == True :
        if filename == None :
            import time
            tm = time.localtime()
            fn = str("%04d%02d%02d_%02d%02d%02d.msk"%(tm[0],tm[1],tm[2],tm[3],tm[4],tm[5]))
        else :
            fn=filename
        mfile = open(fn, "w")
        for index0 in maskedindex :
            mfile.write("%s.%s\n"%(index0[0],index0[1]))
        mfile.close()
        #return
        return None # added by inamura (2008.09.03)
    else :
        return maskedindex
    
def ChopperMask(ecm, masked=None, maskingfile=None):
    """
    This function at first read masked index as list of index or file.
    Then, it write "MASKED" key as 1 in HeaderBase in each ElementContainer
    
    """
    
    maskindex = []
    if maskingfile != None :
        mkfile = open(maskingfile, "r")
        lines = mkfile.readlines()
        for line in lines :
            ll = line.split()
            if len(ll) == 0:
                pass
            elif ll[0] is "#" :
                pass
            else:
                for ls in ll:
                    qs = ls.split(".")
                    if len(qs) == 1 or (len(qs)==2 and len(qs[1])==0) :
                        maskindex.append((int(qs[0]),))
                    elif len(qs) == 2:
                        maskindex.append((int(qs[0]), int(qs[1])))
        maskindex = tuple(maskindex)
    elif masked != None :
        maskindex = masked    

    #if type(ecm) != type(ElementContainerMatrix()):
    if not isinstance( ecm, ElementContainerMatrix ):
        raise AttributeError, "Input must be Manyo.ElementContainerMatrix"
    if type(maskindex) not in [type([1,1]), type((1,))]:
        raise AttributeError, "Parameters appointing detector to be masked must be Python.List or Python.Tuple"
    ## [Inamura 091222,130702]-->
    HIS = SearchInHeader(ecm)
    ## <--[Inamura 091222,130702]
    for mm in maskindex :
        ## [Inamura 091222,130702]-->
        #print mm
        HIS.SearchArray("PSDID",mm[0])
        res = HIS.PutResultIndex(0)
        if res.size()==0:
            print "PSD-ID:",mm[0]," is not found."
            continue
        if len(mm) == 1:
            mm = (res[0],)
        elif len(mm) == 2:
            mm = (res[0],mm[1])
        else:
            continue
        ##<--[Inamura 091222,130702]
        if len(mm) == 1:
            ecm(mm[0]).PutHeaderPointer().OverWrite("MASKED",1)
            for jj in range(ecm(mm[0]).PutTableSize()):
                tempheader = ecm(mm[0], jj).PutHeader()
                tempheader.OverWrite("MASKED",1)
                ecm(mm[0],jj).InputHeader(tempheader)
        elif len(mm) == 2:
            tempheader = ecm(mm[0], mm[1]).PutHeader()
            tempheader.OverWrite("MASKED",1)
            ecm(mm[0],mm[1]).InputHeader(tempheader)
            echeader = ecm(mm[0]).PutHeaderPointer()
            vind = echeader.PutInt4Vector("MASKEDElementContainer")
            vind.append(mm[1])
            echeader.OverWrite("MASKEDElementContainer", vind)
    ecmhd = ecm.PutHeaderPointer()
    ecmhdprc = ecmhd.PutStringVector("DATAPROCESSED")
    ecmhdprc.append("MASK")
    ecmhd.OverWrite("DATAPROCESSED", ecmhdprc)
    ## [Inamura 091222]-->
    del HIS
    ## <--[Inamura 091222]
    return ecm
    


def BackgroundSubtraction(ecm, method, info): # change ec->ecm by inamura (2008.09.03)
    
    """
    Background Subtraction Method
    
    @  method  
        0 : Constant Subtraction for All histogram 
        1 : Average of Both End Points
        2 : Average of Both End range
        3 : Background is given by List or Array
    @ info 
        if method is 0 : Constant 
                     1 : (Number of Point at starting, Number of point at ending) 
                     2 : (range of left end point, right end point)
                     3 : BG array
    Notice >
        if method 1 or 2 and length of info is 1, then it is used for both.
        
        if BG array is 1-dimensional and the size is same as ElementContainerArray number
            in ElementContainerMatrix, one element in BG array is considered as BG for 
            ElementContainerArray  
        
    """
    BGSubMethod = { "By Constant" : 0,              
                    "By Average of End Point" : 1,
                    "By Average of End Range": 2,
                    "By Array" : 3 }
    
    #ecm = eccopy(ec)  removed by inamura (2008.09.03)

    if type(ecm) == type(None):
        raise AttributeError, "Input must be Manyo.ElementContainerMatrix"
    BGS_Method = 0
    
    
    if type(method) == type(1) and method in BGSubMethod.values():
        BGS_Method = method
    elif type(method) == type("a") and method in BGSubMethod.keys():
        BGS_Method = BGSubMethod(method)
    else :
        raise AttributeError, "Background Subtraction Method is not specified"
    
    if BGS_Method == 0:
        if type(info) not in (type(1), type(1.0)):
            raise ValueError, "Must be value" 
        #if type(ecm) is type(ElementContainer()):
        if isinstance( ecm, ElementContainer ):
            ecsub(ecm, info)
        #elif type(ecm) is type(ElementContainerArray()):
        elif isinstance( ecm, ElementContainerArray ):
            for ind in range(ecm.PutTableSize()):
                ec=ecm.PutPointer(ind)
                ecsub(ec, info)
        #elif type(ecm) is type(ElementContainerMatrix()):
        elif isinstance( ecm, ElementContainerMatrix ):
            for ind1 in range(ecm.PutTableSize()):
                eca = ecm.PutPointer(ind1)
                print "---[eca table size=",ind1,eca.PutTableSize()
                for ind2 in range(eca.PutTableSize()):
                    ec = eca.PutPointer(ind2)
                    ecsub(ec, info)
    
    elif BGS_Method == 1 :
        if type(info) in [type(1), type(1.0)]:
            nn = numpy.array([info])
        else:
            nn = numpy.array(info)
        
        if len(info) == 1:
            nl = nr = nn[0]
        elif len(info) == 2:
            nl = nn[0]
            nr = nn[1]
        else :
            raise ValueError, "Must be lentgh 2 integer sequence - list, tuple, array" 
                
        #if type(ecm) is type(ElementContainer()):
        if isinstance( ecm, ElementContainer ):
            yd = numpy.array(ecm.PutYList())
            bg = (yd[:nl].sum()+yd[-nr:].sum())/float(nl+nr)
            ecsub(ecm, bg)
        #elif type(ecm) is type(ElementContainerArray()):
        elif not isinstance( ecm, ElementContainerArray ):
            for ind in range(ecm.PutTableSize()):
                ec=ecm.PutPointer(ind)
                yd = numpy.array(ec.PutYList())
                bg = (yd[:nl].sum()+yd[-nr:].sum())/float(nl+nr)
                ecsub(ec, bg)
        #elif type(ecm) is type(ElementContainerMatrix()):
        elif isinstance( ecm, ElementContainerMatrix ):   
            for ind1 in range(ecm.PutTableSize()):
                eca = ecm.PutPointer(ind1)
                for ind2 in range(eca.PutTableSize()):
                    ec = eca.PutPointer(ind2)
                    yd = numpy.array(ec.PutYList())
                    bg = (yd[:nl].sum()+yd[-nr:].sum())/float(nl+nr)
                    ecsub(ec, bg)
        else :
            raise ValueError, "Wrong Background Number"
        
        
    elif BGS_Method == 2:
        if type(info) in [type(1), type(1.0)]:
            nn = numpy.array([info])
        else:
            nn = numpy.array(info)
        
        if len(nn) == 1:
            nl = nr = float(nn[0])
        elif len(info) == 2:
            nl = float(nn[0])
            nr = float(nn[1])
        else :
            raise ValueError, "Must be lentgh 2 integer sequence - list, tuple, array" 
                
        #if type(ecm) is type(ElementContainer()):
        if isinstance( ecm, ElementContainer ):
            xd0 = numpy.array(ecm.PutXList())
            xd = (xd0[1:]+xd0[:-1])*0.5
            yd = numpy.array(ecm.PutYList())
            bgr = []
            for ind in range(len(xd)):
                if abs(xd[ind]-xd0[0]) <= nl or abs(xd0[-1]-xd[ind]) <= nr : 
                    bgr.append(yd[ind])
            bgr = numpy.array(bgr)
            bg = bgr.sum()/len(bgr)
            ecsub(ecm, bg)
        #elif type(ecm) is type(ElementContainerArray()):
        elif not isinstance( ecm, ElementContainerArray ):
            for ind in range(ecm.PutTableSize()):
                ec=ecm.PutPointer(ind)
                xd0 = numpy.array(ec.PutXList())
                xd = (xd0[1:]+xd0[:-1])*0.5
                yd = numpy.array(ec.PutYList())
                bgr = []
                for ind in range(len(xd)):
                    if abs(xd[ind]-xd0[0]) <= nl or abs(xd0[-1]-xd[ind]) <= nr : 
                        bgr.append(yd[ind])
                bgr = numpy.array(bgr)
                bg = bgr.sum()/len(bgr)
                ecsub(ec, bg)
        #elif type(ecm) is type(ElementContainerMatrix()):
        elif isinstance( ecm, ElementContainerMatrix ):
            for ind1 in range(ecm.PutTableSize()):
                eca = ecm.PutPointer(ind1)
                for ind2 in range(eca.PutTableSize()):
                    ec = eca.PutPointer(ind2)
                    xd0 = numpy.array(ec.PutX())
                    xd = (xd0[1:]+xd0[:-1])*0.5
                    yd = numpy.array(ec.PutY())
                    bgr = []
                    for ind in range(len(xd)):
                        if abs(xd[ind]-xd0[0]) <= nl or abs(xd0[-1]-xd[ind]) <= nr : 
                            bgr.append(yd[ind])
                    bgr = numpy.array(bgr)
                    bg = bgr.sum()/len(bgr)
                    ecsub(ec, bg)
        
        else :
            raise ValueError , "Wrong Background Range"
        
    elif BGS_Method == 3:
        BGarr = info
            
        #if type(ecm) == type(ElementContainerArray()):
        if isinstance( ecm, ElementContainerArray ):
            if len(BGarr) != ecm.PutTableSize():
                raise ValueError, "Wrong Backgroud Array size"
            for ind in range(ecm.PutTableSize()):
                ec=ecm.PutPointer(ind)
                ecsub(ec, BGarr[ind])
        #elif type(ecm) == type(ElementContainerMatrix()) :
        elif isinstance( ecm, ElementContainerMatrix ):
            if len(BGarr) != ecm.PutTableSize():
                raise ValueError, "Wrong Backgroud Array size"
            for ind1 in range(ecm.PutTableSize()):
                eca = ecm.PutPointer(ind1)
                for ind2 in range(eca.PutTableSize()):
                    ec = eca.PutPointer(ind2)
                    if type(BGarr[ind1]) in (type(1.0), type(1)):
                        ecsub(ec, BGarr[ind1])
                    elif type(BGarr[ind1]) in (type([1,]), type((1,)), type(numpy.array([1,]))):
                        if len(BGarr[ind1]) == eca.PutTableSize():
                            ecsub(ec, BGarr[ind1, ind2])
                        else :
                            raise ValueError, "Wrong Backgroud Array size"
                    else :
                        raise ValueError, "Wrong Backgroud Array size"
    else :
        raise ValueError, "Wrong Backgroud Method or Information"
    return ecm
    

#def Tof2Energy(ecm, L2p, Lms, t_offset):
def Tof2Energy(ecm, Lms, t_offset):        #changed by inamura (2008.08.04)
    """
    Function to convert Time-of-Flight to Energy 
    
    Arguements:
        ecm : Elememt Container Matrix
        Ei (meV) : Incident Neutron Energy
        L2p (m) : list of sample and detector pixel distance
        Lms (m) : monitor to sample distance
        t_offset (microsec) : time-of-flight shift. 
    """
    #Ei = ecm.PutHeader().PutDouble("IncidentNeutronEnergy")
    Ei = ecm.PutHeader().PutDouble("Ei") # changed by inamura (2008.09.04)
    import numpy
    if ecm(0).PutHeaderPointer().PutString("TYPE")=="MONITOR":
        start_ind0 = 1
    else:
        start_ind0 = 0
    for ind0 in range(start_ind0,ecm.PutTableSize()):
        print "PSD=",ind0
        for ind1 in range(ecm(ind0).PutTableSize()):
            xkey = ecm(ind0, ind1).PutXKey()
            ykey = ecm(ind0, ind1).PutYKey()
            ekey = ecm(ind0, ind1).PutEKey()
            tof = numpy.array(ecm(ind0, ind1).PutX()) - t_offset
            vi = 437.4 * numpy.sqrt(Ei)
            mn = 1.6749e-27  # Neutron mass
            J2meV = 6.24151e21 # Joule to meV
            '''
            if type(Lms) == type([]):
                energy = Ei - 0.5 * mn * (L2p*L2p)/(tof*1.0e-6 - Lms[ind0][ind1]/vi)**2
            else :
                energy = Ei - 0.5 * mn * (L2p*L2p)/(tof*1.0e-6 - Lms/vi)**2 * J2meV
            '''
            PosiVect = ecm(ind0, ind1).PutHeaderPointer().PutDoubleVector("PixelPosition")
            L2p = (numpy.sqrt(PosiVect[0]**2 + PosiVect[1]**2 + PosiVect[2]**2))/1000.0 # [m]
            energy = Ei - 0.5 * mn * (L2p*L2p)/(tof*1.0e-6 - Lms/vi)**2 * J2meV
                #### removed and added by inamura (2008.09.04)

            ectemp = ecm.PutPointer(ind0).PutPointer(ind1)
            ectemp.Add("energy", l2d(energy), "meV")
            ectemp.SetKeys("energy", ykey, ekey)
            ectemp.Remove(xkey)
    hd = ecm.PutHeaderPointer()
    strvector = hd.PutStringVector("DATAPROCESSED")
    strvector.append("TOF to ENERGY CONVERSION")
    hd.OverWrite("DATAPROCESSED", strvector)
    return ecm

#########################################
def Tof2EnergyWhite(ecm, t_offset):        #Added by inamura (2008.11.20)
    """
    Function to convert Time-of-Flight to Energy on White beam
    
    Arguements:
        ecm : Elememt Container Matrix
        t_offset (microsec) : time-of-flight shift. 
    """
    import numpy
    if ecm(0).PutHeaderPointer().PutString("TYPE")=="MONITOR":
        start_ind0 = 1
    else:
        start_ind0 = 0
    L1 = ecm.PutHeaderPointer().PutDouble("L1") # [mm]
    L1 = L1/1000.0 #[m]
    for ind0 in range(start_ind0,ecm.PutTableSize()):
        print "psd = ",ind0
        for ind1 in range(ecm(ind0).PutTableSize()):
            xkey = ecm(ind0, ind1).PutXKey()
            ykey = ecm(ind0, ind1).PutYKey()
            ekey = ecm(ind0, ind1).PutEKey()
            tof_o = ecm(ind0,ind1).PutX()
            for i in range(len(tof_o)):
                if tof_o[i]==0:
                    tof_o[i]=1.0e-6
            tof = numpy.array(tof_o) - t_offset
            mn = 1.6749e-27  # Neutron mass
            J2meV = 6.24151e21 # Joule to meV
            
            PosiVect = ecm(ind0, ind1).PutHeaderPointer().PutDoubleVector("PixelPosition")
            L2p = (numpy.sqrt(PosiVect[0]**2 + PosiVect[1]**2 + PosiVect[2]**2))/1000.0 # [m]
            energy = 0.5 * mn * ((L1+L2p)*(L1+L2p))/(tof*1.0e-6)**2 * J2meV
            
            ectemp = ecm.PutPointer(ind0).PutPointer(ind1)
            ectemp.Add("energy", l2d(energy), "meV")
            ectemp.SetKeys("energy", ykey, ekey)
            ectemp.Remove(xkey)
    hd = ecm.PutHeaderPointer()
    strvector = hd.PutStringVector("DATAPROCESSED")
    strvector.append("TOF to ENERGY CONVERSION")
    hd.OverWrite("DATAPROCESSED", strvector)
    return ecm


#########################################
def NormalizedByMonitor(ecm_sample, monitorindex = 0, ini = None, fin = None): 
                        # changed by inamrua (2008.09.08)
    """
    Normalization by Monitor
    
    @param  dat             Params 
    @param  monitorindex    int     index of monitor for normalization
    @param  ini             float   inital value for integrating monitor data
    @param  fin             float   final value for integrating monitor data
    @retval ECM             elementContainerMatrix
    """
    t_i = ini     # added by inamura(2008.09.03)
    t_f = fin     # added by inamura(2008.09.03)
    monitorhistogram = simplehistogram(ecm_sample(0, monitorindex)) # changed by inamura (2008.09.03)
    m_i, m_e = monitorhistogram.sum(t_i, t_f)
    for ecaindex in range(1, ecm_sample.PutTableSize()):
        for ecindex in range(ecm_sample(ecaindex).PutTableSize()):
            '''
            ectemp = ecm_sample.PutPointer(ecaindex).PutPointer(ecindex)
            xk, yk, ek = ectemp.PutXKey(), ectemp.PutYKey(), ectemp.PutEKey()
            y0 = numpy.array(ectemp.PutY(), dtype='float')
            e0 = numpy.array(ectemp.PutE(), dtype='float')
            y1 = y0/m_i
            e1 = numpy.sqrt((e0/m_i)**2 + (m_e*y0/m_i**2)**2)
            newy = l2d(y1)
            newe = l2d(e1)
            ectemp.Replace(yk, newy)
            ectemp.Replace(ek, newe)
            ectemp.SetKeys(xk, yk, ek)
            '''
            ecm_sample(ecaindex,ecindex).MulMySelf(1/m_i) # changed by inamura (2008.09.03)
    hd = ecm_sample.PutHeaderPointer()
    strvector = hd.PutStringVector("DATAPROCESSED")
    strvector.append("MonitorCountNormalization")
    hd.OverWrite("DATAPROCESSED", strvector)
    return ecm_sample
    
    
    
    
def Rebin(ecm, Ei, Ef, dE):
    earray0 = numpy.linspace(Ei, Ef, dE)
    earray = l2d(earray0)
    for ind0 in range(ecm.PutTableSize()):
        for ind1 in range(ecm(ind0).PutTableSize()):
            tempec=ecm.PutPointer(ind0).PutPointer(ind1)
            tempec.Rebin(earray)
    return ecm
        
