# -*- coding: utf-8 -*-

import Manyo as mm
import os
################################
#      VanadiumCorrection      #
################################
class VanadiumCorrection(object):
    """
    Vanadium Correction
    Data Correction using White Vanadium or Monochromatic Vanadium
    """
    ################################
    def __init__(self, SampleData=None, WhiteVan=None, MonoVan=None, WhiteVanM=None):
        """
        Constructor
        
        @SampleData (ElementContainerMatrix) Sample data
        @WhiteVan (UInt4) White Vanadium run number
        @MonoVan (UInt4) Monochromatic Vanadium run number
        @WhiteVanM (UInt4) White Vanadium run number for Monochromatic Vanadium run
        """
        self.inst_code = 'SIK'
        self.SampleData=None  
        self.WV_run = None
        self.WV_EC = None
        self.MV_run = None
        self.MV_EC = None
        self.WVM_run = None
        self.WVM_EC = None

        self.PathToVanData = "./"
        self.EventDataPath = "../RawData/SIK10_081120000000/"
        self.DetectorInfoFile = "./DetectorInfo_16PSD2.xml"

        self.CommentHead = "Vanadium >> "
        self.DoFlag = False
        self.SetFlag = [0,0,0,0]

        self.PeakPosi_m1 = 0.0
        self.Lms = 0.0

        self.Set(SampleData, WhiteVan, MonoVan, WhiteVanM)


    ################################
    def Set(self, SampleData=None, WhiteVan=None, MonoVan=None, WhiteVanM=None):
        """
        Set Parameters
        Set Parameters for the process of Vanadium Correction
        @SampleData (ElementContainerMatrix) Sample data
        @WhiteVan (UInt4) White Vanadium run number
        @MonoVan (UInt4) Monochromatic Vanadium run number
        @WhiteVanM (UInt4) White Vanadium run number for Monochromatic Vanadium run
        """
        if SampleData!=None: # sample data ElementContainerMatrix import
            #if type(SampleData)==type(mm.ElementContainerMatrix()):
            if isinstance( SampleData, mm.ElementContainerMatrix ):
                self.SampleData=SampleData
                self.SetFlag[0]=2

        if WhiteVan!=None:   
            if type(WhiteVan)==type(0):
                fn = self._VanDataExist(WhiteVan,1)
                if fn!=None:
                    self.WV_run = int(WhiteVan)
                    self.WV_EC = self.LoadFromFile(fn)
                    self.SetFlag[1]=2
                else:
                    print self.CommentHead+" Warnning : There is no WhiteVan data of run number you inputted."
                    self.WV_run = int(WhiteVan)
                    self.WV_EC = None
                    self.SetFlag[1]=1
                                   
        if MonoVan!=None:
            if type(MonoVan)==type(0):
                fn = self._VanDataExist(MonoVan,2)
                if fn!=None:
                    self.MV_run = int(MonoVan)
                    self.MV_EC = self.LoadFromFile(fn)
                    self.SetFlag[2]=2
                else:
                    print self.CommentHead+" Warnning : There is no MonoVan data of run number you inputted."
                    self.MV_run = int(MonoVan)
                    self.MV_EC = None
                    self.SetFlag[2]=1

        if WhiteVanM!=None:
            if type(WhiteVanM)==type(0):
                fn = self._VanDataExist(WhiteVanM,3)
                if fn!=None:
                    self.WVM_run = int(WhiteVanM)
                    self.WVM_EC = self.LoadFromFile(fn)
                    self.SetFlag[3]=2
                else:
                    print self.CommentHead+" Warnning : There is no WhiteVan for MonoVan data of run number you inputted."
                    self.WVM_run = int(WhiteVanM)
                    self.WVM_EC = None
                    self.SetFlag[3]=1

    ################################
    def DoVanCorrection(self):
        """
        Do Vanadium Correction
        Do Vanadium Correction using White Vanadium or Monochromatic Vanadium
        """
        Flag = True
        
        if self.SampleData==None:
            print self.CommentHead+"Error : Sample data have never been set."
            Flag = False
        if (self.WV_EC==None)and(self.MV_EC==None):
            print self.CommentHead+"Error : Both Vanadium run number have never been set."
            Flag = False

        if Flag:                  # If there are Sample data and Vanadium 
            process_str = ""
                                  # if we have W_s, MV, W_mv
            if (self.WV_EC!=None)and(self.MV_EC!=None)and(self.WVM_EC!=None):
                index_WV = self.WV_EC.PutX()
                count_WV = self.WV_EC.PutY()
                index_MV = self.MV_EC.PutX()
                count_MV = self.MV_EC.PutY()
                index_WVM = self.WVM_EC.PutX()
                count_WVM = self.WVM_EC.PutY()
                if index_WV.size()!=index_MV.size(): 
                    print self.CommentHead+" Warning : Index of WhiteVan and one of MonoVan are different."
                index_v = []
                count_v = []
                for i in range(index_WV.size()): 
                    index_v.append( index_WV[i] )
                    count_v.append( count_WV[i]/count_WVM[i]*count_MV[i] )
                process_str = "WV+MV+WVM"

            elif (self.WV_EC!=None)and(self.MV_EC==None): # if we have only W_s
                index_v = self.WV_EC.PutX()
                count_v = self.WV_EC.PutY()
                process_str = "WV"
                
            else:                                        # if we have only MV
                index_v = self.MV_EC.PutX()
                count_v = self.MV_EC.PutY()
                process_str = "MV"

            print "[DoCorrection] len(count_v)=",len(count_v)
            for i in range(1,self.SampleData.PutTableSize()):  # Do correction
                for j in range(self.SampleData(i).PutTableSize()):
                    pixel_id = self.SampleData(i,j).PutHeaderPointer().PutInt4("PIXELID")
                    ccc = count_v[pixel_id]
                    if ccc==0.0:
                        ccc=1.0e-10
                    dammy = self.SampleData(i,j).MulMySelf(1.0/ccc)
            del dammy
            hd = self.SampleData.PutHeaderPointer()
            process = hd.PutStringVector("DATAPROCESSED")
            process.append(":VanCorr(" + process_str + ")")
            hd.OverWrite("DATAPROCESSED", process)
            self.DoFlag = True

        else :
            print self.CommentHead+" ALART!! : We CANNOT correct data by Vanadium data."

    ################################
    def PutCorrectedData(self):
        """
        Put corrected data
        Put data corrected by Vanadium
        @retval (ElementContainerMatrix) 
        """
        if self.DoFlag:
            return self.SampleData
        else :
            print self.CommentHead+" ALART!! : We CANNOT Do vanadium corrections."
            return self.SampleData

    ################################
    def _VanDataExist(self,runnum,which):
        """
        Check file path to Vanadium data
        Check file path to vanadium data and return its path
        @param runnum(Int4) Run Number for Vanadium
        @param which(Int4) 1:WhiteVan for sample, 2:MonoVan, 3:WhiteVan for MonoVan
        @retval (None or String)
        """
        DataName = self._ReturnVanDataFile(runnum,which)
        if DataName!="None":
            fn = self.PathToVanData + DataName
            if os.path.exists(fn):
                return fn
            else:
                print "There is no file name =",fn
                return None
        else:
            print "ReturnVanDataFile=None"
            return None

    ################################
    def SetEventDataToHistXMLFile(self, FileName):
        """
        Set XML file 
        Set XML file into this class for MakeVanadiumData
        @param (String) file name of XML information.
        """
        self.DetectorInfoFile=FileName

    ################################
    def PutWhiteVanEC(self):
        """
        Put ElementContainer for White Vanadium 
        Put ElementContainer for White Vanadium calculated for the vanadium corrections
        @retval (ElementContainer)  
        """
        if self.WV_EC==None:
            print "There is no data for White Vanadium."
            return -1
        else:
            return self.WV_EC

    ################################
    def PutMonoVanEC(self):
        """
        Put ElementContainer for Monochromatic Vanadium 
        Put ElementContainer for Monochromatic Vanadium calculated for the vanadium corrections
        @retval (ElementContainer)  
        """
        if self.MV_EC==None:
            print "There is no data for Monochro Vanadium."
            return -1
        else:
            return self.MV_EC

    ################################
    def _ReturnVanDataFile(self,runnum,which):
        """
        Return data file name of analysed vanadium
        Get file name of analysed vanadium data from DB and return it in future
        @param runnum (Int) Runnumber
        @param which(Int) 1:WhiteVan, 2:MonoVan, 3:WhiteVan for MonoVan
        @retval string
        """
        stamp = ""
        st_runnum="%6d" % (runnum)
        if (which == 1)|(which == 3):
            stamp = "wvd"
        elif which == 2:
            stamp = "mvd"
        
        print "[ReturnVanDataFile] stamp=",stamp
        if stamp!="":
            ret = self.inst_code + st_runnum.replace(' ','0') + "." + stamp
            print ret
            return ret
        else:
            return "None"
        
    ################################
    def SetMonitorInfo(self,Lms):
        """
        Set the length between monitor 1 to sample
        
        @param Lms(Double) Length between monitor 1 to sample
        
        """
        self.Lms = Lms

    ################################
    def MakeVanData(self):
        """
        Make and save data for Vanadium correction
        Make and save data for Vanadium correction for White and Monochromatic
        """
        for which in range(1,3):
            if (self.SetFlag[which]==1):
                self._MakeVanadiumEC(which)
                self.SaveToFile(which)


    ################################
    def _MakeVanadiumEC(self,Which):
        """
        Make data for vanadium correction
        Make data for vanadium correction 
        @param Which (UInt4) 1: White Van, 2: Mono Van, 3: White Van for Mono van
        """
        import EventDataToHist as EVT
        import MonitorAnalysis as MON
        from DataReduct_SJY import Tof2Energy, Tof2EnergyWhite
        from sub1 import simplehistogram, ecsub, eccopy

        if Which==1:
            runnum = self.WV_run
        elif Which==2:
            runnum = self.MV_run
        elif Which==3:
            runnum = self.WVM_run
        else:
            return

        e2h = EVT.EventDataToHist(self.inst_code,runnum,debug=0)
        e2h.ImportDammyMonitor()
        ret = e2h.SetParamFile(self.DetectorInfoFile,0)
        if ret==-1:
            print "XML file is no correct."
            return
        e2h.SetPathToData(self.EventDataPath)
        e2h.MakeEventDataFileName()
        VanDat=e2h.PutElementContainers()

        MaxPixelId = e2h.EtoH.PutMaxOfPixelId()
        
        if Which==2:
            M_mv = MON.MonitorAnalysis(dat)
            M_mv.SearchPeakPosition(0)
        
            Ei, Ei_unit = M_mv.PutIncidentEnergy()
            delta_Ei, delta_Ei_unit = M_mv.PutIncidentEnergyWidth()
            Ei_ini = Ei-delta_Ei/2.0
            Ei_fin = Ei+delta_Ei/2.0
            self.PeakPosi_m1 = mon.PutPeakPosi(0)
            Tof2Energy(VanDat, self.Lms, self.PeakPosi_m1)
 
        else:
            Ei, Ei_unit = (0.0,"meV")
            delta_Ei, delta_Ei_unit = (0.0,"meV")
            Ei_ini = 20.0
            Ei_fin = 50.0
            print "Tof2Energy White--"
            Tof2EnergyWhite(VanDat, 0.0 )
        
        V_EC = mm.ElementContainer()
        list_pixelId=[]
        list_integ=[]
        list_dam_err=[]
        for i in range(MaxPixelId+1):
            list_pixelId.append(i)
            list_integ.append(0.0)
            list_dam_err.append(1.0)
        list_pixelId.append(MaxPixelId)
        
        if VanDat(0).PutHeaderPointer().PutString("TYPE")=="MONITOR":
            start_i = 1
        else:
            start_i = 0
        VanDat(1,0).Dump()
        print "Max of PixelID = ",MaxPixelId
        print "length of List_pixelId =",len(list_pixelId)
        for psd in range(start_i, VanDat.PutTableSize()):
            for pixel in range(VanDat(psd).PutTableSize()):
                pixel_id=VanDat(psd)(pixel).PutHeaderPointer().PutInt4("PIXELID")
                sh = simplehistogram(VanDat(psd)(pixel))
                try:
                    (d1,d2)=sh.integrate(ini=(Ei_ini),fin=(Ei_fin))
                    list_integ[pixel_id] = float(d1)
                    list_dam_err[pixel_id] = float(d2)
                except:
                    print "Errored Id =",pixel_id
                    return
                               
        for i in range(len(list_integ)):
            print list_integ[i]

        V_EC.Add("PixelId",list_pixelId)
        V_EC.Add("IntegMonVan",list_integ)
        V_EC.Add("DammyErr", list_dam_err)
        V_EC.SetKeys("PixelId","IntegMonVan","DammyErr")

        if Which==1:
            self.WV_EC = V_EC
        elif Which==2:
            self.MV_EC = V_EC
        else:
            self.WVM_EC = V_EC
        
        del VanDat, V_EC, list_pixelId, list_integ, list_dam_err
        
        

    ################################
    def SaveToFile(self, Which=0):
        """
        Save data for Vanadium corrections 
        Save data for Vanadium corrections with NeXus format
        @param Which (UInt4) 1: White Van, 2: Mono Van, 3: White Van for Mono van
        """
        runnum = None
        if Which==1:
            runnum = self.WV_run
        elif Which==2:
            runnum = self.MV_run
        elif Which==3:
            runnum = self.WVM_run

        if runnum==None:
            print "Vanadium infor not setting."
            return
        
        if self.DetectorInfoFile==None:
            print "DetectorInfo file is not set."
            return

        filename = self._ReturnVanDataFile(runnum,Which)
        
        fp = open(self.DetectorInfoFile,"r")
        XmlContents = fp.read()
        fp.close()

        w = mm.WriteNeXusFile( filename, "VanadiumCorrect", 1 )
        w.MakeOpenGroup( "Entry1", "NXentry" )
        w.MakeOpenGroup( "Vanadium", "NXdata" )
        w.WriteData( self.DetectorInfoFile, XmlContents )
        if Which==1:
            w.WriteData( "WhiteVan", self.WV_EC )
        elif Which==2:
            w.WriteData( "MonoVan", self.MV_EC )
        else:
            w.WriteData( "WhiteVan", self.WVM_EC )

        w.CloseGroup()
        w.CloseGroup()
        
        del w


    ################################
    def LoadFromFile(self, filename):
        """
        Load data from NeXus file
        Load data for Vanadium corrections from NeXus file
        @param filename (String) filename
        """
        if filename[-3:]=="wvd":
            tag = "WhiteVan"
        elif filename[-3:]=="mvd":
            tag = "MonoVan"
        else:
            print "Filename is not good for van data"
            return None

        rf = mm.ReadNeXusFile( filename )
        vs1 = rf.GetNextEntry()
        rf.OpenGroup( vs1[0], vs1[1] )
        
        vs2 = rf.GetNextEntry()
        rf.OpenGroup( vs2[0], vs2[1] )
        
        VanEC = rf.ReadElementContainer( tag )
        
        rf.CloseGroup()
        rf.CloseGroup()
        
        VanEC.SetKeys("PixelId","IntegMonVan","DammyErr")
        return VanEC
        
            
