from mantid.simpleapi import *
import jsonArray
import time
import datetime
import json
import multiprocessing
from multiprocessing import Process, Manager
import shutil
import numpy as np
import path

class moduleIO(multiprocessing.Process):
    def __init__(self, neonRedis, refreshtime, moduleList, myDict, logger):
        super(moduleIO, self).__init__()
        self.refreshtime=refreshtime
        self.neonRedis_write=neonRedis
        self.neonRedis_read=neonRedis
        self.logger=logger
        self.moduleList=moduleList
        self.instname="./instrument/detector.xml"
        self.myDict=myDict
        self.confJudge = True

        self.neonpathDetValue= []
        self.neonpathDetPid= []
        self.neonpathDetTof= []
        for num in range(len(self.moduleList)):
            self.neonpathDetValue.append("/SANS/workspace/detector/module"+str(num+1)+"/value")
            self.neonpathDetPid.append("/SANS/workspace/detector/module"+str(num+1)+"/pid")
            self.neonpathDetTof.append("/SANS/workspace/detector/module"+str(num+1)+"/tof")

    def getProtonCharge(self):
        data=[]
        path="/SANS/drone/proton_charge"
        _json=self.neonRedis_read.get(path)
        data=jsonArray.jsonDecoder(_json)
        val=data[1][-1]
        if(type(val)==int):
	    if val == 0:
		return 1
	    else:
            	return val
        else:
            return 1

    def get1Ddata(self, path):
        data=[]
        _json=self.neonRedis_read.get(path)
        data=jsonArray.jsonDecoder(_json)
        return data

    def getValue(self, path):
        data = []
        _json=self.neonRedis_read.get(path)
        _array=jsonArray.jsonDecoder(_json)
        for ix in _array:
            for iy in ix:
                data.append(iy)
        return data

    def createWorkspace(self):
        LoadCSNSRaw(OutputWorkspace="sample", PixelID_bank=self.detPid, TimeOfFlight_bank=self.detTof, Counts_bank=self.detValue, PixelID_monitor=[1,2,3], TimeOfFlight_monitor=[0,8], Counts_monitor=[2,4,6])
        LoadInstrument(Workspace="sample_1", Filename=self.instname, RewriteSpectraMap='True')
        SaveNexus(InputWorkspace='sample_1', Filename='./detector.nxs')

    def cropworkspace(self):
	for i in range(12*len(self.moduleList)):
	    CropWorkspace(InputWorkspace="sample_1", OutputWorkspace=str(i), StartWorkspaceIndex=i*125, EndWorkspaceIndex=i*125+124)
	    SumSpectra(InputWorkspace=str(i), OutputWorkspace="tofCounts"+str(i), IncludeMonitors=False)
	    name=mtd["tofCounts"+str(i)]
	    counts=[]
	    counts=name.readY(0)
	    _tmp=name.readX(0)
	    tof=[]
	    path1="/SANS/workspace/MantidData/tube"+str(i+1)+"/tof"
	    path2="/SANS/workspace/MantidData/tube"+str(i+1)+"/counts"
            for j in range(len(_tmp)-1):
                tof.append(_tmp[j]+(_tmp[j+1]-_tmp[j])/2)
	    self.neonRedis_write.set(path1,jsonArray.jsonEncoder(tof))
	    self.neonRedis_write.set(path2,jsonArray.jsonEncoder(counts))
	    self.logger.debug("set I-tof for tube"+str(i+1))
 
    def noReduceData(self):
        #find beam center
        #tof-counts
        SumSpectra(InputWorkspace='sample_1', OutputWorkspace="tofCounts", IncludeMonitors=False)
	proton=self.getProtonCharge()
        #print proton
        #print type(proton)
        CreateSingleValuedWorkspace(OutputWorkspace="__pro", DataValue=proton)
	Divide(LHSWorkspace="tofCounts", RHSWorkspace="__pro", OutputWorkspace="tofCounts_PC")	

        ConvertToHistogram(InputWorkspace='sample_1', OutputWorkspace='sample_1')
        ConvertUnits(InputWorkspace='sample_1', OutputWorkspace='sample_1_wave', Target='Wavelength', AlignBins=True)
        SumSpectra(InputWorkspace='sample_1_wave', OutputWorkspace="waveCounts", IncludeMonitors=False)
	Divide(LHSWorkspace="waveCounts", RHSWorkspace="__pro", OutputWorkspace="waveCounts_PC")	
        FindCenterOfMassPosition(InputWorkspace="sample_1_wave", Output="beam_center")
	name=mtd["sample_1_wave"]
        w=name.readX(0)
	max_w=w[len(w)-1]
        min_w=w[0]
        step=(max_w-min_w)/100
	param_a='0.001,'+str(-step)+','+str(max_w)
	self.logger.warning("wave: "+param_a)
        Rebin(InputWorkspace='sample_1_wave', OutputWorkspace='sample_wave_rebin', Params=param_a)
	ConvertUnits(InputWorkspace='sample_wave_rebin', OutputWorkspace='tmp_q', Target='MomentumTransfer', AlignBins=True)
	SumSpectra(InputWorkspace='tmp_q', OutputWorkspace="tmp_q", IncludeMonitors=False)
        
        #Q1D(DetBankWorkspace='sample_wave_rebin', OutputWorkspace='sample_Q', OutputBinning='0.008,-0.01,2.0', SolidAngleWeighting=False)
	name=mtd["tmp_q"]
        #q=name.readX(0)
	#max_q=q[len(q)-1]
	#max_q=5.0
        #min_q=q[0]
        #step=(max_q-min_q)/100
	#param_b='0.001,'+str(step)+','+str(max_q)
	#self.logger.warning("q: "+param_b)
        #Rebin(InputWorkspace='tmp_q', OutputWorkspace='sample_Q_rebin', Params=param_b)
        Rebin(InputWorkspace='tmp_q', OutputWorkspace='sample_Q_rebin', Params='0.001,0.004,5')
	Divide(LHSWorkspace="sample_Q_rebin", RHSWorkspace="__pro", OutputWorkspace="sample_Q_rebin_PC")	
	#ConvertUnits(InputWorkspace='sample_Q_rebin', OutputWorkspace='sample_wave', Target='Wavelength')

        Qxy(InputWorkspace='sample_wave_rebin', OutputWorkspace='sample_Qxy', MaxQxy=0.25, DeltaQ=0.01, SolidAngleWeighting=False)

    def setBeamCenter(self, path):
        value=[]
        wname=mtd["beam_center"]
        x=(wname.row(0))['Value']
        y=(wname.row(1))['Value']
        value.append(x)
        value.append(y)
        _json_v=jsonArray.jsonEncoder(value)
        self.neonRedis_write.set(path, _json_v)

    def set1DData(self,wsname, path1,path2):
        name=mtd[wsname]
        value=name.readY(0)
        _tmp=name.readX(0)
        x=[]
        y=[]
        for j in range(len(_tmp)-1):
            x.append(_tmp[j]+(_tmp[j+1]-_tmp[j])/2)
        for i in value:
            y.append(i)
        self.neonRedis_write.set(path1,jsonArray.jsonEncoder(x))
        self.neonRedis_write.set(path2,jsonArray.jsonEncoder(y))


    def setIQxy(self, path1,path2,path3):
        wname=mtd["sample_Qxy"]
        _tmp=wname.readX(0)
        qxy=[]
        for j in range(len(_tmp)-1):
            qxy.append(_tmp[j]+(_tmp[j+1]-_tmp[j])/2)
        numHist=wname.getNumberHistograms()
        intensity=[[] for i in range(numHist)]
        for j in range(numHist):
            _n2=wname.readY(j)
            for p in _n2:
                intensity[j].append(p)
        self.neonRedis_write.set(path1,jsonArray.jsonEncoder(qxy))
        self.neonRedis_write.set(path2,jsonArray.jsonEncoder(qxy))
        self.neonRedis_write.set(path3,jsonArray.jsonEncoder(intensity))

    def setXY(self, path1, path2, path3):
        name=mtd["sample_1"]
        counts=[]
        xaxis=[]
        yaxis=[]
        for j in range(name.getNumberHistograms()):
            _det=name.getDetector(j)
            _pos=_det.getRelativePos()
            xaxis.append(_pos.X()-0.004)
            yaxis.append(_pos.Y()-0.004)
            counts.append(sum(name.readY(j)))
        xaxis=list(set(xaxis))
        yaxis=list(set(yaxis))
        xaxis.sort()
        yaxis.sort()
        xaxis.append(xaxis[-1]+0.008)
        yaxis.append(yaxis[-1]+0.008)
        for i in range(len(xaxis)):
            xaxis[i]=xaxis[i]*1000-500
        for i in range(len(yaxis)):
            yaxis[i]=yaxis[i]*1000-500
        vaxis=[[] for r in range(len(yaxis)-1)]
        num=0
        for p in range(len(yaxis)-1):
            for q in range(len(xaxis)-1):
                vaxis[p].append(counts[num])
                num=num+1
        self.neonRedis_write.set(path1, jsonArray.jsonEncoder(xaxis))
        self.neonRedis_write.set(path2, jsonArray.jsonEncoder(yaxis))
        self.neonRedis_write.set(path3, jsonArray.jsonEncoder(vaxis))
        self.logger.debug("set XY bank!"+path1+path2+path3)
        # set 128 tubes
        #nx=126
        #ny=2
        for num in range(120):
            Y=[]
            Y.append(yaxis[num])
            Y.append(yaxis[num+1])
            Z=[[] for r in range(1)]
            for p in range(1):
                Z[p].extend(vaxis[num])
            #X=[[] for k in range(ny)]
            #Y=[[] for k in range(ny)]
            #Z=[[] for k in range(ny)]
            #for i in range(nx):
            #    Y[0].append(yaxis[num])
            #    Y[1].append(yaxis[num+1])
            #for j in range(ny):
            #    for i in range(nx):
            #        X[j].append(xaxis[i])
            #        Z[j].append(vaxis[num][i])
            _p1="/SANS/workspace/MantidData/tube"+str(num+1)+"/xy_image/x"
            _p2="/SANS/workspace/MantidData/tube"+str(num+1)+"/xy_image/y"
            _p3="/SANS/workspace/MantidData/tube"+str(num+1)+"/xy_image/value"
            self.neonRedis_write.set(_p1,jsonArray.jsonEncoder(xaxis))
            self.neonRedis_write.set(_p2,jsonArray.jsonEncoder(Y))
            self.neonRedis_write.set(_p3,jsonArray.jsonEncoder(Z))
        self.logger.debug("finish XY for 120 tubes!")

    def dataSet(self):
        self.setBeamCenter("/SANS/workspace/MantidData/BeamCenter")
        self.setXY("/SANS/workspace/MantidData/xy_image/x","/SANS/workspace/MantidData/xy_image/y","/SANS/workspace/MantidData/xy_image/value")
        self.set1DData("tofCounts", "/SANS/workspace/MantidData/raw/tof","/SANS/workspace/MantidData/raw/counts")
        self.set1DData("tofCounts_PC", "/SANS/workspace/MantidData/reduction/PC/tof","/SANS/workspace/MantidData/reduction/PC/counts")
        self.set1DData("sample_Q_rebin", "/SANS/workspace/MantidData/raw/q","/SANS/workspace/MantidData/raw/intensity")
        self.set1DData("sample_Q_rebin_PC", "/SANS/workspace/MantidData/reduction/PC/q","/SANS/workspace/MantidData/reduction/PC/intensity")
        self.set1DData("waveCounts", "/SANS/workspace/MantidData/raw/wavelength","/SANS/workspace/MantidData/raw/wcounts")
        self.set1DData("waveCounts_PC", "/SANS/workspace/MantidData/reduction/PC/wavelength","/SANS/workspace/MantidData/reduction/PC/wcounts")
        self.setIQxy("/SANS/workspace/MantidData/IQxy/x","/SANS/workspace/MantidData/IQxy/y","/SANS/workspace/MantidData/IQxy/counts")        

    def process(self):
        if self.myDict['configure']:
            if self.confJudge:
                self.detPid=[]
                self.detTof=[]
                try:
                    for i in range(10):
                        self.detPid.extend(self.get1Ddata(self.neonpathDetPid[i]))
                    self.detTof=self.get1Ddata(self.neonpathDetTof[0])
                except:
		    self.logger.error("no data in Pid and Tof in module")
                    self.myDict['error']=True
                    self.myDict['configure']=False

		path.clearRedis(self.neonRedis_write)
                self.logger.debug("finish configure")
                self.confJudge=False
            else:
                pass
        else:
            self.confJudge=True            
        
        if self.myDict['running']:
	    be=time.time()
            self.logger.debug("start running program for module")
            self.detValue=[]
            try:
                for i in range(10):
                    self.detValue.extend(self.getValue(self.neonpathDetValue[i]))
            except:
                self.logger.warning("no data for module")
                npid=len(self.detPid)
                ntof=len(self.detTof)
                for i in range(npid*(ntof-1)):
                    self.detValue.append(0)
            self.neonRedis_write.set("/SANS/workspace/MantidData/neutroncounts",json.dumps(sum(self.detValue)))        
             
            if (len(self.detValue)!=(len(self.detTof)-1)*len(self.detPid)):
                self.logger.warning("data size doesn't match!")
                
            else:
                if self.myDict['running']:
                    self.logger.debug("   INFO: createWorkspace")
                    #try:
                    self.createWorkspace()
                    self.cropworkspace()
                    self.noReduceData()
		    '''
                    except:
                        self.myDict['error']=True
                        self.myDict['running']=False
                        self.logger.error("error in get process data in mantid for module")
		    '''

                    if self.myDict['running']:
                        self.logger.debug("   INFO: dataSet")
                        #try:  
                        self.dataSet()
	                shutil.move("./detector.nxs","./tmp/detector.nxs" )
			self.logger.debug(" finish move nexus file!!!")
			self.logger.warning(" finish time is "+str(time.time()-be))
                        #except:
                        #    self.myDict['error']=True
                        #    self.myDict['running']=False
                        #    self.logger.error("error in get set data to neon for module")
                    else:
                        pass
                else:
                    pass
        else:
            pass

    def run(self):            
        while True:
            self.process()
            time.sleep(self.refreshtime)

