import Manyo
import Manyo.SAS as ms
import Manyo.Utsusemi as mu
import os
from utsusemi.SAS.ana.Reduction.SASUtils import SASBankDic

"""
Base Functions for SAS
Corr.py
- NormByMonCounts
- CorrectDetEffiMulti
- CorrectDetEffi
- I0LambdaCorrection
- P0LambdaCorrection
"""
########################################
def NormByMonCounts( dat, monitor, key="Lamb", XRange="0.6 7.6", factor=1.0 ):
    """
    Normalize data by monitor counts
    @param  dat      (ElementContainerMatrix, ElementContainer) target data
    @param  monitor  (ElementContainer) Monitor data from GetMonHistSAS
    @param  Xrange   (string)  "min-value max-value" for the range to sumalize intensities.
    @param  factor   (float)   factor for normalization I=(Iorg/mon)*factor
    @return None
    """
    if not (isinstance( dat, Manyo.ElementContainerMatrix ) or isinstance( dat, Manyo.ElementContainer)):
        raise UserWarning("dat must be ElementContainerMatrix or ElementContainer")
    if not isinstance( monitor, Manyo.ElementContainer ):
        raise UserWarning("monitor must be ElementContainer")

    xrange_st = XRange.split(" ")
    if len(xrange_st)!=2:
        raise UserWarning("XRange must have 2 values")

    xmin=float(xrange_st[0])
    xmax=float(xrange_st[1])
    monCount = 0.0
    lam_v = monitor("Lamb")
    intens = monitor("Intensity")

    for i in range(len(lam_v)):
        if xmin<=lam_v[i] and lam_v[i]<=xmax:
            monCount += intens[i]

    if (monCount!=0.0):
        dat.MulMySelf((1.0/monCount)*factor)

########################################
def CorrectDetEffi( dat, filename="detectorEffiTBL.dat" ):
    """
    Does the correction of detector efficiency for lambda
    CorrectDetEffi( dat, filename )

    @param  dat       (ElementContainerMatrix)
    @param  filename  (string)
    @retval None
    """
    #fpath = os.path.join( os.environ["UTSUSEMI_USR_DIR"],"ana","xml",filename )
    fpath = mu.FindParamFilePath( filename, "" )
    if not os.path.exists( fpath ):
        raise UserWarning("Could not find file="+filename)

    DEC = ms.DetectorEffiCorrection( dat, fpath );
    DEC.Execute()

    del DEC

########################################
def CorrectDetEffiRPMT( dat, alpha, beta ):
    """
    Does the correction of detector efficiency for lambda
    CorrectDetEffi( dat, filename )

    @param  dat   (ElementContainerMatrix)
    @param  alpha (float) parameter of lambda dependency for detect efficiency of RPMT
    @param  beta  (float) parameter of lambda dependency for detect efficiency of RPMT
    @retval None
    """

    DEC = ms.DetectorEffiCorrectionRPMT( dat )
    DEC.SetParamsFormula1( alpha, beta )
    DEC.Execute()

    del DEC

########################################
def CorrectDetEffiMulti( dat, BankDefault="All", Bank20atm="SCC,BW", DefaultData="detectorEffiTBL.dat", Atm20Data="detectorEffiTBL.dat" ):
    """
    Does the correction of detector efficiency for lambda
    CorrectDetEffi( dat, filename )

    @param  dat       (ElementContainerMatrix)
    @param  BankDefault (string)
    @param  Bank20atm   (string)
    @param  DefaultData (string)
    @param  Atm20Data   (string)
    @retval None
    """
    if type(dat)!=list:
        dat = [dat]

    # Set Bank Info
    default_banks_list = []
    if BankDefault.upper()!="ALL":
        default_banks_tmp = BankDefault.replace(" ",",")
        default_banks_tmp = default_banks_tmp.split(",")
        for a_bank in default_banks_tmp:
            if a_bank!="":
                if a_bank=="SM":
                    default_banks_list.append( "SMBL16" )
                    default_banks_list.append( "SMBL14" )
                else:
                    default_banks_list.append( a_bank )
    else:
        for a_bank in SASBankDic:
            if a_bank=="SM":
                default_banks_list.append( "SMBL16" )
                default_banks_list.append( "SMBL14" )
            else:
                default_banks_list.append( a_bank )
    if len(default_banks_list)==0:
        return

    atm20_banks_list = []
    atm20_banks_tmp = Bank20atm.replace(" ",",")
    atm20_banks_tmp = atm20_banks_tmp.split(",")
    for a_bank in atm20_banks_tmp:
        if a_bank!="":
            atm20_banks_list.append( a_bank )
            if a_bank in default_banks_list:
                default_banks_list.remove( a_bank )

    # Set data for detector efficiency
    default_dpath = mu.FindParamFilePath( DefaultData )
    atm20_dpath = mu.FindParamFilePath( Atm20Data )
    if not os.path.exists( default_dpath ):
        raise UserWarning("Could not find file="+DefaultData)
    if not os.path.exists( atm20_dpath ):
        raise UserWarning("Could not find file="+Atm20Data)

    # Execute correction for each banks
    for a_dat in dat:
        bank_vect = a_dat.PutHeader().PutStringVector("BANKNAMELIST")
        for i in range(bank_vect.size()):
            if bank_vect[i] in default_banks_list:
                print( "Bank , i = %s, %i"%(bank_vect[i],SASBankDic[ bank_vect[i] ]) )
                DEC = ms.SASDetectorEfficiencyCorrection();
                DEC.LoadDataFile( default_dpath )
                DEC.Execute( a_dat,SASBankDic[ bank_vect[i] ] )
                del DEC
            else:
                print( "Bank(20atm) , i = %s, %i"%(bank_vect[i],SASBankDic[ bank_vect[i] ]) )
                DEC = ms.SASDetectorEfficiencyCorrection();
                DEC.LoadDataFile( atm20_dpath )
                DEC.Execute( a_dat, SASBankDic[ bank_vect[i] ] )
                del DEC

#########################################
def I0LambdaCorrection(dat,I0Lamb,PolData="None"):
    """
    I0(lambda) correction
    @param dat      (ElementContainerMatrix)
    @param I0Lamb   (ElementContainer)
    @param PolData  (str) file path to text data for the polarization vs lambda
    """
    i0_corr = ms.I0LambdaCorrection(dat)
    if PolData=="None":
        if i0_corr.Execute(I0Lamb):
            pass
        else:
            raise UserWarning("I0LambdaCorrection >> failed to execute")
    else:
        if i0_corr.Execute(I0Lamb,PolData):
            pass
        else:
            raise UserWarning("I0LambdaCorrection >> failed to execute with PolData"+PolData)
    del i0_corr

#########################################
def P0LambdaCorrection(dat1,dat2,PolData):
    """
    P0(lambda dependency of poralization) correction
    @param dat1     (ElementContainerMatrix) I^{+}_obs
    @param dat2     (ElementContainerMatrix) I^{-}_obs
    @param PolData  (str) file path to text data for the polarization vs lambda
    @retval RET list of ElementContainerMatrix [I^{+}, I^{-}]
    """
    if not (isinstance( dat1, Manyo.ElementContainerMatrix ) or isinstance( dat2, Manyo.ElementContainerMatrix)):
        raise UserWarning("dat1 and dat2 must be ElementContainerMatrix")
    P0LC = ms.P0LambdaCorrection( dat1, dat2 )
    if P0LC.SetPolarizationTable( PolData ):
        Ip = Manyo.ElementContainerMatrix()
        Im = Manyo.ElementContainerMatrix()
        if P0LC.Execute(Ip,Im):
            del P0LC
            return [Ip,Im]
    else:
        raise UserWarning("Poralization data is not found.")

#########################################
def NormByKickers(dat):
    """
    Normalize Intensity by the number of kickers
    @param dat   (ElementContainerMatrix)
    @retval None
    """
    num_kickers = dat.PutHeaderPointer().PutInt4("KICKERCOUNT")
    if num_kickers!=0:
        dat.MulMySelf( 1.0/float(num_kickers) )

#########################################
def MakeDataListForSIB(dataList="A,B",A=None,B=None,C=None,D=None,E=None,F=None,G=None,H=None,I=None,J=None,K=None,L=None,M=None):
    """
    Make Data list for ScaleIntensityBanks
    @param dataList (string) list contents using A,B,... and M
    @param A [Def:EMPTYDAT] Data for a bank
    @param B [Def:EMPTYDAT] Data for a bank
    @param C [Def:EMPTYDAT] Data for a bank
    @param D [Def:EMPTYDAT] Data for a bank
    @param E [Def:EMPTYDAT] Data for a bank
    @param F [Def:EMPTYDAT] Data for a bank
    @param G [Def:EMPTYDAT] Data for a bank
    @param H [Def:EMPTYDAT] Data for a bank
    @param I [Def:EMPTYDAT] Data for a bank
    @param J [Def:EMPTYDAT] Data for a bank
    @param K [Def:EMPTYDAT] Data for a bank
    @param L [Def:EMPTYDAT] Data for a bank
    @param M [Def:EMPTYDAT] Data for a bank
    @retval listCBB
    """
    exec_code = "listCBB=["+dataList+"]"
    print("MakeDataListForCBB : exec_code = "+exec_code)
    exec(exec_code)
    return listCBB

#########################################
def ScaleIntensityBanks(datalist, SCC=1.0, SC=1.0, SM=1.0, SU=1.0, SD=1.0, MU=1.0, MD=1.0, MR=1.0, ML=1.0, HU=1.0, HD=1.0, HL=1.0, BW=1.0):
    """
    Scaling the intensitiy for each bank by given factor
    @param datalist [Def:listCBB]  (ElementContainerMatrix of list of ElementContainerMatrix)
    @param SCC (float) Scale factor for SCC bank
    @param SC  (float) Scale factor for SC bank
    @param SM  (float) Scale factor for SM bank
    @param SU  (float) Scale factor for SU bank
    @param SD  (float) Scale factor for SD bank
    @param MU  (float) Scale factor for MU bank
    @param MD  (float) Scale factor for MD bank
    @param MR  (float) Scale factor for MR bank
    @param ML  (float) Scale factor for ML bank
    @param HU  (float) Scale factor for HU bank
    @param HD  (float) Scale factor for HD bank
    @param HL  (float) Scale factor for HL bank
    @param BW  (float) Scale factor for BW bank
    @retval None
    """
    if type(datalist)!=list:
        datalist = [datalist]

    for a_data in datalist:
        bankName_list=a_data.PutHeader().PutStringVector("BANKNAMELIST")
        for a_bankName in bankName_list:
            bankId=SASBankDic[a_bankName]
            for i in range(a_data.PutSize()):
                if bankId == a_data(i).PutHeaderPointer().PutInt4("BANKID"):
                    if a_bankName=="SMBL16" or a_bankName=="SMBL14":
                        a_bankName="SM"
                    exec_code = 'a_data(i).MulMySelf( %s )'%(a_bankName)
                    #print( "exec_code= ",exec_code )
                    exec(exec_code)
#########################################
def LambdaDependCorrect(data, P1="ALL:lam_dep.txt", P2="None", P3="None", P4="None", P5="None", P6="None", P7="None", P8="None", P9="None"):
    """
    Correction for lambda dependency on intensity by given text file.
    Params format : <Bank(s)>:<FactorFile path>
    <Bank(s)> can include several banks with comma (,) like "SM,SC"

    @param data [Def:DAT] (ElementContainerMatrix or list of it)
    @param P1 (string)
    @param P2 (string)
    @param P3 (string)
    @param P4 (string)
    @param P5 (string)
    @param P6 (string)
    @param P7 (string)
    @param P8 (string)
    @param P9 (string)
    @retval None
    """
    if type(data)!=list:
        data=[data]

    LD = ms.LambdaDependCorrection()
    params_list = [P1,P2,P3,P4,P5,P6,P7,P8,P9]
    for i,a_params in enumerate(params_list):
        a_params = a_params.strip().replace(" ","")
        if a_params.upper()=="NONE":
            continue
        params_v=a_params.split(":")
        if len(params_v)<2:
            raise UserWarning("LambdaDependCorrect >>> P{} params infalid".format(i))

        bank_list_org = params_v[0].split(",")
        bank_list = []
        for a_bank in bank_list_org:
            if LD.CheckBankName(a_bank):
                bank_list.append( a_bank )
            else:
                print( "Invalid Bank Name={} is ignored.".format(a_bank) )

        for a_bank in bank_list:
            if LD.SetCurveTextDep(a_bank, params_v[1]):
                pass
            else:
                raise UserWarning("LambdaDependCorrect >>> Failed executed by params={}".format(params_v[1]))

    for a_data in data:
        LD.SetTarget(a_data)
        if LD.Execute():
            pass
        else:
            raise UserWarning("LambdaDependCorrect >>> Failed on executing")
    return

