#!/usr/bin/python
# -*- coding: utf-8 -*-

import os,sys,time
ver = float(sys.version_info[0])+float(sys.version_info[1])*0.1
if ver<2.5:
    from elementtree.ElementTree import ElementTree,fromstring,XML,tostring
else:
    from xml.etree.ElementTree import ElementTree,fromstring,XML,tostring

from numpy import *
import time

import Manyo as mm
import Manyo.Utsusemi as mu
import vis.VisContMParams as VCP
      
################################
#      CuiCtrlVisContM
################################
class CuiCtrlVisContM(object):
    """
    Class to control VisualContM from CUI
    
    ver 0.1 [110726] first release
    Change History
    [inamura 160619] Bugfix of treating SampleRotateSteps
    """
    Def_type = ["Visual","OutToD4Mat"]
    #####################################
    def __init__(self, filepath="", type="OutToD4Mat", ecm=None):
        """
        Constructor
        
        @param self
        @param filepath (String) path to xml file of parameters
        @param ecm (ElementContainerMatrix)
        @retval None
        """
        if not type in self.Def_type:
            print "type=",type," is invalid."
            raise
        
        self.VCP = VCP.VisContMParams()

        self.type     = type
        self.viewAxis = None
        self.param    = None
        self.plotid   = None
        self.ecm      = ecm
        self.map      = None
        self.folding  = [0,0,0,0]
        self.VisualCalcSqe = None
        self.MessageTag = "CuiCtrlVisContM >>>"
        
        if filepath!="":
            try:
                self.VCP.SetParamFile( filepath )
            except:
                print "File not found.",filepath
                raise
            try:
                self.VCP.LoadVisualContParams()
            except:
                print "XML format may be invalid."
                raise

            self.MakeSliceParams()

        if self.ecm!=None:
            self.SetECM()

        
    #####################################
    def MakeSliceParams(self,type=""):
        """
        Make parameters for slicing from imported XML.
        @param type (String) the name of visualization software to be used 
        """
        
        if (self.VCP.proj_axes == None) or (self.VCP.paxis_info == None):
            return

        if (type!="") and (type in self.Def_type):
            self.type = type

        """
        if self.VCP.EnableVisType!=self.type:
            ss="The type of parameter XML is not correct for "+self.type
            raise "IncorrectXMLFile",ss
        """
            
        # make parameter for slicing funcions on VisualCont(2)
        self.viewAxis = zeros([4,4], float64)
        self.param = zeros([12],float64)

        self.plotid = [ [], [], [] ]
        
        if self.type=="Visual":
            print self.MessageTag,"Type=Visual"

            shift = 0
            index = 0
            flag_t = 0

            for i in range(4):
                (type_str,range_axis,folding) = self.VCP.paxis_info[ str(i) ]

                va = self.VCP.proj_axes[ str(i) ]
                for j in range(4):
                    self.viewAxis[index][j]=va[j]
                if type_str=="x":
                    self.plotid[0].append(index)
                    for j in range(len(range_axis)):
                        self.param[j]=range_axis[j]
                   
                elif type_str=="y":
                    self.plotid[1].append(index)
                    for j in range(len(range_axis)):
                        self.param[j+3]=range_axis[j]
                    
                elif (type_str=="t"):
                    self.plotid[2].append(index)
                    if flag_t==0:
                        for j in range(len(range_axis)):
                            self.param[j+6]=range_axis[j]
                        flag_t=1
                    else:
                        for j in range(len(range_axis)):
                            self.param[j+8]=range_axis[j]

                else:
                    print self.MessageTag,"Unknown type = ",type_str
                    raise

                self.folding[index]=folding

                index+=1


        elif self.type=="OutToD4Mat":
            print self.MessageTag,"Type=OutToD4Mat"

            for i in range(4):
                (type_str,range_axis,folding) = self.VCP.paxis_info[ str(i) ]

                va = self.VCP.proj_axes[ str(i) ]
                for j in range(4):
                    self.viewAxis[i][j]=va[j]

                if (type_str=="t"):
                    if (len(range_axis)==2):
                        range_axis.append(range_axis[1]-range_axis[0])
                    elif (len(range_axis)==3):
                        range_axis[2] = range_axis[1]-range_axis[0]
                    else:
                        print self.MessageTag,"index = %d : range_axis number(%d) is invalid" % (i,len(range_axis))
                        raise

                for j in range(len(range_axis)):
                    self.param[i*3 + j]=range_axis[j]

                self.folding[i]=folding
                
        else:
            print self.MessageTag,"Type is invalid.",type

    #####################################            
    def SetPhi(self, phi=0.0):
        """
        Set Phi ( fai in KCS products )
        @param phi (Double) phi angle from given orientation of sample
        """
        if self.VCP.Header=={}:
            print self.MessageTag,"Original XML file is not loaded."
            raise
        
        #self.VCP.smpInfo.fai = float(fai)
        self.VCP.SetOrientationPhi(phi)

    #####################################
    def SetECM(self, ecm=None):
        """
        Set ElementContainerMatrix for slicing.
        @param ecm (ElementContainerMatrix)
        """
        
        #if ecm!=None and type(ecm)==type(mm.ElementContainerMatrix()):
        if ecm!=None and isinstance( ecm, mm.ElementContainerMatrix ):
            self.ecm = ecm

    #####################################
    def DoProjection(self):
        """
        """
        if self.VCP.Header=={}:
            print self.MessageTag,"Did not make parameters."
            raise

        print self.MessageTag,"Projection..."

        ## Get parameters from XML
        LC  = self.VCP.Header['SimpleLatticeConsts']
        UV  = self.VCP.Header['SampleUvec']
        VV  = self.VCP.Header['SampleVvec']
        if 'SampleRotatePhi' in self.VCP.Header:
            PHI = self.VCP.Header['SampleRotatePhi']
        else:
            PHI = 0
        ##[inamura 160619]-->
        if  self.VCP.Header.has_key('SampleRotateSteps'):
            RL= self.VCP.Header['SampleRotateSteps']
        else:
            RL=None
        ## Set parameters to Header of ElementContianerMatrix
        ecm_h = self.ecm.PutHeaderPointer()
        lcv = mm.MakeDoubleVector()
        for i in range(6):
            lcv.append( LC[i] )
        if ecm_h.CheckKey("SimpleLatticeConsts")==1:
            ecm_h.OverWrite("SimpleLatticeConsts",lcv)
        else:
            ecm_h.Add("SimpleLatticeConsts",lcv)
            
        uv = mm.MakeDoubleVector()
        vv = mm.MakeDoubleVector()
        for i in range(3):
            uv.append( UV[i] )
            vv.append( VV[i] )
        if ecm_h.CheckKey("SampleUvec")==1:
            ecm_h.OverWrite("SampleUvec",uv)
        else:
            ecm_h.Add("SampleUvec",uv)
        if ecm_h.CheckKey("SampleVvec")==1:
            ecm_h.OverWrite("SampleVvec",vv)
        else:
            ecm_h.Add("SampleVvec",vv)
        
        if ecm_h.CheckKey("SampleRotatePhi")==1:
            ecm_h.OverWrite("SampleRotatePhi", PHI )
        else:
            ecm_h.Add("SampleRotatePhi", PHI )
        
        ##[inamura 160619]-->
        ##[inamura 150907]-->
        sample_rot_steps = mm.MakeDoubleVector()
        if RL!=None and len(RL)!=0:
            for a_rl in RL:
                if a_rl[0]=="X":
                    sample_rot_steps.append( 0.0 )
                elif a_rl[0]=="Y":
                    sample_rot_steps.append( 1.0 )
                elif a_rl[0]=="Z":
                    sample_rot_steps.append( 2.0 )
                else:
                    continue
                sample_rot_steps.append( a_rl[1] )
        else:
            sample_rot_steps.append( 1.0 ) ## Axis-Y is 1
            sample_rot_steps.append( PHI ) ## rotate angle

        ##<--[inamura 160619]
        if ecm_h.CheckKey("SampleRotateSteps")==1:
            ecm_h.OverWrite("SampleRotateSteps", sample_rot_steps )
        else:
            ecm_h.Add("SampleRotateSteps", sample_rot_steps )
        ##<--[inamura 150907]
        
        ## Calc projection
        self.VisualCalcSqe = mu.UtsusemiSqeCalc()
        
        self.VisualCalcSqe.SetTarget( self.ecm )
        viewAxis_in = mm.MakeDoubleVector()
        for i in range(4):
            for j in range(4):
                viewAxis_in.append( self.viewAxis[i][j] )
        all_u = self.VisualCalcSqe.Projection(viewAxis_in)
        
        
    #####################################
    def DoSlice(self):
        """
        Make sliced data from given ElementContainerMatrix and slicing parameters.
        
        """
        
        self.DoProjection()
            
        ## Slicing
        Ranges = []
        AxType = mm.MakeStringVector()
        Foldings = mm.MakeDoubleVector()
        for i in range(4):
            (type_str, range_axis, folding) = self.VCP.paxis_info[ str(i) ]
            print "type_str=",type_str
            AxType.append( type_str )
            Foldings.append( folding )
            Ranges.append( mm.MakeDoubleVector() )
            for j in range( len(range_axis) ):
                Ranges[i].append( range_axis[j] )

        ret = self.VisualCalcSqe.Slice( Ranges[0],Ranges[1],Ranges[2],Ranges[3],
                                        AxType,Foldings )
        ## Make map format                
        xv = self.VisualCalcSqe.PutXArray()
        yv = self.VisualCalcSqe.PutYArray()
        II = self.VisualCalcSqe.PutDArray()
        EE = self.VisualCalcSqe.PutEArray()
        print "size=",xv.size(),yv.size(),II.size(),EE.size()
        h = []
        v = []
        for i in range(xv.size()):
            h.append(xv[i])
        for i in range(yv.size()):
            v.append(yv[i])
        
        h = array(h)
        v = array(v)

        H, V = meshgrid(h, v)

        numx = xv.size()-1
        numy = yv.size()-1
        print "numx,numy=",numx,numy
        data =zeros((numy,numx), float64)
        err =zeros((numy,numx), float64)
        
        for j in range(numx):
            for k in range(numy):
                data[k,j] = II[ j*numy + k ]
                err[k,j] = EE[ j*numy + k ]

        self.map = ( data, err, H, V, 1 )

    #####################################
    def DoSliceToEca(self):
        """
        Make sliced data from given ElementContainerMatrix and slicing parameters.
        
        """
        
        self.DoProjection()
            
        ## Slicing
        Ranges = []
        AxType = mm.MakeStringVector()
        Foldings = mm.MakeDoubleVector()
        for i in range(4):
            (type_str, range_axis, folding) = self.VCP.paxis_info[ str(i) ]
            print "type_str=",type_str
            AxType.append( type_str )
            Foldings.append( float(folding) )
            Ranges.append( mm.MakeDoubleVector() )
            for j in range( len(range_axis) ):
                Ranges[i].append( range_axis[j] )

        ret = self.VisualCalcSqe.Slice( Ranges[0],Ranges[1],Ranges[2],Ranges[3],AxType,Foldings,True )
        ECA = mm.ElementContainerArray()
        self.VisualCalcSqe.SetSlicedElementContainerArray( ECA )
        return ECA
                
    #####################################    
    def SendToM2Plot(self):
        """
        Send sliced data to M2Plot.
        """
        if self.map == None:
            print self.MessageTag,"Did not do slicing."
            raise
        
        import vis.M2Plot as M2P
        return M2P.M2Plot(self.map)

    #####################################
    def SaveFigureOfM2Plot(self,filepath="test.png",zmin=-1,zmax=-1,isLogPlot=False):
        """
        Save a figure for sliced data plotted by CuiD2Chart
        [inamura 120426]
        @param filepath (string) Output file (extension must be .png)
        @param zmin (double) min z axis range
        @param zmax (double) max z axis range
        @param isLogPlot (bool) use log scaling on Intensity or not
        @retval None
        """
        if self.map == None:
            print self.MessageTag,"Did not do slicing."
            raise


        #make Main Title
        Titles=["","","",""]
        Ei = self.ecm.PutHeader().PutDouble("Ei")
        sample = self.VCP.Header['SampleName']
        Titles[0] = "%s ( Ei= %.1f (meV) )" % ( sample, Ei )
        
        # make Sub Title
        UV = self.VCP.Header['SampleUvec']
        VV = self.VCP.Header['SampleVvec']
        str1 = "U=[%4.2f, %4.2f, %4.2f]\n" % (UV[0],UV[1],UV[2])
        str2 = "V=[%4.2f, %4.2f, %4.2f]\n" % (VV[0],VV[1],VV[2])
        if 'SampleRotatePhi' in self.VCP.Header:
            str3 = "Psi=%5.2f\n" % self.VCP.Header['SampleRotatePhi']
        else:
            rotstep_list = self.VCP.Header['SampleRotateSteps']
            tmp_rot_str="RotStep="
            for a_rotstep in rotstep_list:
                tmp_rot_str+=("%s:%g,"%(a_rotstep[0],a_rotstep[1]))
            str3 = "%s\n"%(tmp_rot_str[:-1])

        AX1=self.VCP.proj_axes["0"]
        AX2=self.VCP.proj_axes["1"]
        AX3=self.VCP.proj_axes["2"]
        AX4=self.VCP.proj_axes["3"]
        str4 = "ax1=[%4.2f,%4.2f,%4.2f,%4.2f]\n" % (AX1[0],AX1[1],AX1[2],AX1[3])
        str5 = "ax2=[%4.2f,%4.2f,%4.2f,%4.2f]\n" % (AX2[0],AX2[1],AX2[2],AX2[3])
        str6 = "ax3=[%4.2f,%4.2f,%4.2f,%4.2f]\n" % (AX3[0],AX3[1],AX3[2],AX3[3])
        str7 = "ax4=[%4.2f,%4.2f,%4.2f,%4.2f]\n" % (AX4[0],AX4[1],AX4[2],AX4[3])

        nn = []
        str8 = ""
        for i in range(len(self.VCP.paxis_info)):
            PX=self.VCP.paxis_info[str(i)]
            if PX[0]=="t":
                str8 += "Thick=[ax%d:%s - %s]\n" % ( i,PX[1][0],PX[1][1] )
            if PX[0]=="x":
                Titles[2]=self.VCP.proj_axes[str(i)][4]
            if PX[0]=="y":
                Titles[3]=self.VCP.proj_axes[str(i)][4]
                
        str8 = str8[:-1]
        
        Titles[1] = str1+str2+str3+";"+str4+str5+str6+str7+";"+str8

        # make figure
        import vis.CuiD2Chart as CDC

        cdc = CDC.CuiD2Chart( self.map, zmin, zmax, Titles )
        cdc.MakePlot()
        cdc.SaveFile( filepath )

        del cdc
        
    #####################################
    def OutToD4Mat(self, filename):
        """
        """
        filepath = str(filename)
        
        self.DoProjection()
        Ranges = []
        #Foldings = mm.MakeInt4Vector()
        Foldings = mm.MakeDoubleVector()
        for i in range(4):
            (type_str, range_axis, folding) = self.VCP.paxis_info[ str(i) ]
            Foldings.append( folding )
            Ranges.append( mm.MakeDoubleVector() )
            for j in range( len(range_axis) ):
                Ranges[i].append( range_axis[j] )
        
        ret = self.VisualCalcSqe.ConvertToD4Mat( filepath,
                                                 Ranges[0],Ranges[1],Ranges[2],Ranges[3],
                                                 Foldings )



###########################    
def Test(filepath,ecm,type="VisualCont"):
    cui = CuiCtrlVisCont(filepath,type,ecm)
    cui.DoSlice()
    cui.SendToM2Plot()
    
    return cui
