from mantid.simpleapi import *
import jsonArray
from Queue import Queue
import time
import matplotlib
import numpy as np
from canvas import *
import redis
import sys
import datetime
import cStringIO
import pickle
#import Image
import matplotlib.pyplot as plt
import setDataToNeon
#import configure
import json
import random
import threading
from threading import Thread

#class moduleDataGet(QtCore.QThread):
class moduleDataGet(threading.Thread):
    def __init__(self, neonRedis, refreshtime):
        super(moduleDataGet, self).__init__()
        #QtCore.QThread.__init__(self)

        self.refreshtime=refreshtime
        self.neonRedis=neonRedis
        self.module="module01"

        self.reduceCall=True
        self.neonpathDetPid="/MR/workspace/sim/detector/"+self.module+"/pid"
        self.neonpathDetTof="/MR/workspace/sim/detector/"+self.module+"/tof"
        self.neonpathDetValue="/MR/workspace/sim/detector/"+self.module+"/value"
        self.instname="/home/dur/work/online/mr/sim_online/instrument/detector.xml"
        self.instnameMon="/home/dur/work/online/mr/sim_online/instrument/monitor.xml"
        #self.neonpathBeamCenter="/MR/workspace/MantidData/BeamCenter"
        if self.reduceCall:
            self.neonpathIQset="/MR/workspace/MantidData/IQ_reduction"
        else:
            self.neonpathIQset="/MR/workspace/MantidData/IQ_raw"
        
        # set data path to neon
        self.neonpathDetSet="/MR/workspace/MantidData/"+self.module
        print "begin to get module pid data ",self.neonpathDetPid
        self.detPid = self.get1Ddata(self.neonpathDetPid)
        print "begin to get module tof data ",self.neonpathDetTof
        self.tof = self.get1Ddata(self.neonpathDetTof)
        self.tof.append(60000)

    def get1Ddata(self, path):
        be=time.time()
        data=[]
        _json=self.neonRedis.get(path)
        while _json is None:
            print "no 1D data in module01!"
            time.sleep(5)
            _json=self.neonRedis.get(path)
        _array=jsonArray.jsonDecoder(_json)
        for k in _array:
            data.append(int(k))
        #print (time.time()-be), "success in get 1D data in "
        return data
    
    def get2Ddata(self, path):
        be=time.time()
        data=[]
        print "start getting data......"
        _json=self.neonRedis.get(path)
        #print "json: ", type(_json)
        while _json is None:
            print "no 2D data in module01!"
            _json=self.neonRedis.get(path)
            time.sleep(10)
        
        _array=jsonArray.jsonDecoder(_json)
        print "======================================"
        #print sum(_array), "  counts"
        print np.array(_array).shape
        print "======================================"
        for ix in _array:
            for iy in ix:
                data.append(int(iy))
        print (time.time()-be), "success in get 2D data in "
        return data
        #return _array

    def getDataFromNeon(self):
        be=time.time()
        self.detValue=self.get2Ddata(self.neonpathDetValue)
        #print time.time()-be,"finish getDataFrom Neon: " 

    def loadDirectData(self):
        be=time.time()
        #deal with direct data
        LoadISISNexus(Filename='/home/dur/work/online/mr/sim_online/data/Direct_D_0.3.nxs', OutputWorkspace='direct')
        LoadInstrument(Workspace='direct', Filename='/home/dur/work/online/mr/sim_online/data/CSNS_REFL_Definition.xml', MonitorList='1', RewriteSpectraMap='False')
        CropWorkspace(InputWorkspace='direct',OutputWorkspace='inc_detector_ws_orig',StartWorkspaceIndex='3',EndWorkspaceIndex='3')
        CloneWorkspace(InputWorkspace='inc_detector_ws_orig',OutputWorkspace='inc_detector_ws_orig_corr')
        MoveInstrumentComponent(Workspace='inc_detector_ws_orig_corr',ComponentName='point-detector',X='0.0',Z='21.475',RelativePosition='0')
        ConvertUnits(InputWorkspace='inc_detector_ws_orig_corr',OutputWorkspace='inc_det_wave',Target='Wavelength')
        CropWorkspace(InputWorkspace='direct',OutputWorkspace='inc_monitor',StartWorkspaceIndex='0',EndWorkspaceIndex='0')
        ConvertUnits(InputWorkspace='inc_monitor',OutputWorkspace='inc_mon_wave',Target='Wavelength')
        print "  Finish loading Direct Data! ", time.time()-be

    def createWorkspace(self):
        be=time.time()
        self.x_max=7.3
        self.x_min=2.1
        print "length: pid", len(self.detPid)
        print "length: tof", len(self.tof)
        print "length: value", len(self.detValue)
        Mon=self.get1Ddata("/MR/workspace/MantidData/tmp/sample/monitor01")
        #CreateWorkspace(OutputWorkspace='sample_1',DataX=self.tof,DataY=self.detValue, NSpec=6400,UnitX='TOF')
        #CreateWorkspace(OutputWorkspace='sample_2',DataX=self.tof,DataY=Mon, NSpec=1,UnitX='TOF')

        LoadCSNSRaw(OutputWorkspace="sample", PixelID_bank=self.detPid, TimeOfFlight_bank=self.tof, Counts_bank=self.detValue, PixelID_monitor=[1], TimeOfFlight_monitor=self.tof, Counts_monitor=Mon)
        LoadInstrument(Workspace="sample_1", Filename=self.instname, RewriteSpectraMap='True')
        LoadInstrument(Workspace='sample_2', Filename=self.instnameMon, RewriteSpectraMap='True')

        # detector data for reflect
        CloneWorkspace(InputWorkspace='sample_1',OutputWorkspace='ref_det_org')
        Rebin(InputWorkspace='sample_1', OutputWorkspace='sample_tof_rebin', Params='0,16,60000')
        ConvertUnits(InputWorkspace='ref_det_org',OutputWorkspace='ref_det_q',Target='MomentumTransfer', AlignBins=True)

        GroupDetectors(InputWorkspace='sample_tof_rebin', OutputWorkspace='sample_tof', DetectorList='10001-16400')
        ConvertUnits(InputWorkspace='sample_tof',OutputWorkspace='sample_raw_Q',Target='MomentumTransfer',AlignBins=True)

        #MoveInstrumentComponent(Workspace='sample_detector_org',ComponentName='Bank01',X='0.0', Y='0.0',Z='21.475',RelativePosition=False)
        ConvertUnits(InputWorkspace='ref_det_org',OutputWorkspace='ref_det_wave',Target='Wavelength', AlignBins=True)
        #GroupDetectors(InputWorkspace='ref_det_wave', OutputWorkspace='ref_det_wave_group', DetectorList='10001-16400')        
        #CropWorkspace(InputWorkspace='ref_det_wave_group',OutputWorkspace='ref_det_ws',XMin=self.x_min,XMax=self.x_max)
        CropWorkspace(InputWorkspace='ref_det_wave',OutputWorkspace='ref_det_ws',XMin=self.x_min,XMax=self.x_max)
        # monitor data for reflect
        CloneWorkspace(InputWorkspace='sample_2',OutputWorkspace='ref_mon_org')
        ConvertUnits(InputWorkspace='ref_mon_org',OutputWorkspace='ref_mon_wave',Target='Wavelength', AlignBins=True)
        RebinToWorkspace(WorkspaceToRebin='ref_mon_wave', WorkspaceToMatch='ref_det_ws', OutputWorkspace='ref_mon_ws')
        
        print "   INFO: Mantid Workspace updated with module!!",time.time()-be
        #print time.time()-be 

    def noReduceData(self):
        pass


    def dataReduction(self):
        be=time.time()
        #correct sample
        Divide(LHSWorkspace='ref_det_ws',RHSWorkspace='ref_mon_ws',OutputWorkspace='ref_result') 
        #deal with direct
        RebinToWorkspace(WorkspaceToRebin='inc_mon_wave', WorkspaceToMatch='ref_det_ws', OutputWorkspace='inc_mon_ws')
        RebinToWorkspace(WorkspaceToRebin='inc_det_wave', WorkspaceToMatch='ref_det_ws', OutputWorkspace='inc_det_ws')
        Divide(LHSWorkspace='inc_det_ws',RHSWorkspace='inc_mon_ws',OutputWorkspace='inc_result') 
        #correction with direct
        Divide(LHSWorkspace='ref_result',RHSWorkspace='inc_result',OutputWorkspace='IvsLam')
        #convert to Q
        ConvertUnits(InputWorkspace='IvsLam', OutputWorkspace='IvsQ_reduce', Target='MomentumTransfer', AlignBins=True)
         
        GroupDetectors(InputWorkspace='IvsQ_reduce', OutputWorkspace='sample_Q_byM', DetectorList='10001-16400')
        GroupDetectors(InputWorkspace='IvsQ_reduce', OutputWorkspace='sample_Q_byP', DetectorList='10001-16400')

        print " INFO: Data Reduction Finished!"
        print time.time()-be

    def dataSet(self):
        be=time.time()
        setDataToNeon.setXYBankData(self.neonRedis, self.neonpathDetSet) 
        setDataToNeon.setBankRawI(self.neonRedis)

        setDataToNeon.setBankReduction(self.neonRedis)
        #setDataToNeon.setIQxy(self.neonRedis)
        #setDataToNeon.setmonitorI_W(self.neonRedis)
        print "finish module dataSet! ",time.time()-be

    def process(self):
        _be=time.time()
        print "   INFO: getDataFromNeon", datetime.datetime.now()
        self.getDataFromNeon()
        print "   INFO: createWorkspace", datetime.datetime.now()
        self.createWorkspace()
        if self.reduceCall:
            print "   INFO: dataReduction", datetime.datetime.now()
            self.dataReduction()
        else:
            print "   INFO: noReduction", datetime.datetime.now()
            self.noReduceData()
        #print "   INFO: dataSet", datetime.datetime.now()
        self.dataSet()
        print "===================================="
        print " module01 finish!!! ",time.time()-_be, "seconds"
        print "===================================="

    def run(self):
        #time.sleep(self.refreshtime)
        self.loadDirectData()
        print "finish load direct"
        #self.loadCellData()
        #print "finish load  cell!!! "
        while True:
            self.process()
            time.sleep(self.refreshtime)

class monitorDataGet(threading.Thread):
#class monitorDataGet(QtCore.QThread):
    def __init__(self, neonRedis, refreshtime, monitorName):
        super(monitorDataGet, self).__init__()
	    #QtCore.QThread.__init__(self)

        self.refreshtime=refreshtime
        self.neonRedis=neonRedis
        self.monitor=monitorName
        #self.path=pathName
        self.path="/MR/workspace/sim"
        self.neonpathMonPid=self.path+"/"+self.monitor+"/pid"
        self.neonpathMonTof=self.path+"/"+self.monitor+"/tof"
        self.neonpathMonValue=self.path+"/"+self.monitor+"/value"
        # set data path to neon
        self.neonpathMonSet="/MR/workspace/MantidData/"+self.monitor
        #get basic data from neon
        print "begin to get monitor pid data ",self.neonpathMonPid
        self.monPid = self.get1Ddata(self.neonpathMonPid)
        print "begin to get monitor tof data ",self.neonpathMonTof
        self.monTof = self.get1Ddata(self.neonpathMonTof)
        self.monTof.append(60000)

        
    def setRefMonitor(self):
        monValue1=self.get2Ddata('/MR/workspace/sim/cell/transmission/'+self.monitor+'/value')
        print "get cell monitor ok !"
        monValue2=self.get2Ddata('/MR/workspace/sim/direct/transmission/'+self.monitor+'/value')
        print "get direct monitor ok !"
        LoadCSNSRaw(OutputWorkspace="Mon1", PixelID_bank=[0,1], TimeOfFlight_bank=[0,16], Counts_bank=[0,1], PixelID_monitor=self.monPid, TimeOfFlight_monitor=self.monTof, Counts_monitor=monValue1)
        LoadCSNSRaw(OutputWorkspace="Mon2", PixelID_bank=[0,1], TimeOfFlight_bank=[0,16], Counts_bank=[0,1], PixelID_monitor=self.monPid, TimeOfFlight_monitor=self.monTof, Counts_monitor=monValue2)

        SumSpectra(InputWorkspace="Mon1_2", OutputWorkspace="tmp_1_"+self.monitor, IncludeMonitors=False)
        SumSpectra(InputWorkspace="Mon2_2", OutputWorkspace="tmp_2_"+self.monitor, IncludeMonitors=False)
        name1=mtd["tmp_1_"+self.monitor] 
        name2=mtd["tmp_2_"+self.monitor] 
        counts1=name1.readY(0) 
        counts2=name2.readY(0) 
        p1="/MR/workspace/MantidData/tmp/cell/"+self.monitor 
        p2="/MR/workspace/MantidData/tmp/direct/"+self.monitor 
        _json1=jsonArray.jsonEncoder(counts1)
        _json2=jsonArray.jsonEncoder(counts2)
        self.neonRedis.set(p1, _json1)
        self.neonRedis.set(p2, _json2)
        print "successful in set direct and cell ",self.monitor
        print "==============================================================="

    def get1Ddata(self, path):
        be=time.time()
        data=[]
        _json=self.neonRedis.get(path)
        while _json is None:
            print "no 1D data in module!",self.monitor
            time.sleep(5)
            _json=self.neonRedis.get(path)
        _array=jsonArray.jsonDecoder(_json)
        for k in _array:
            data.append(int(k))
        #print (time.time()-be), "success in get 1D data in ",self.monitor
        return data
    
    def get2Ddata(self, path):
        be=time.time()
        data=[]
        _json=self.neonRedis.get(path)
        #print "json: ", type(_json)
        while _json is None:
            print "no 2D data in module!",self.monitor
            _json=self.neonRedis.get(path)
            time.sleep(10)
        
        _array=jsonArray.jsonDecoder(_json)
        for ix in _array:
            for iy in ix:
                data.append(int(iy))
        #print (time.time()-be), "success in get 2D data in ",self.monitor
        return data

    def getDataFromNeon(self):
        be=time.time()
        self.monValue=self.get2Ddata(self.neonpathMonValue)
        #print time.time()-be,"finish getDataFrom Neon: ", self.monitor

    def createWorkspace(self):
        be=time.time()
        print "length: monpid", len(self.monPid)
        print "length: montof", len(self.monTof)
        print "length: monvalue", len(self.monValue)
        LoadCSNSRaw(OutputWorkspace=self.monitor, PixelID_bank=[0,1], TimeOfFlight_bank=[0,16], Counts_bank=[0,1], PixelID_monitor=self.monPid, TimeOfFlight_monitor=self.monTof, Counts_monitor=self.monValue)
        #CloneWorkspace(InputWorkspace=self.monitor+"_2", OutputWorkspace=str(self.monitor))
        LoadInstrument(Workspace=self.monitor+"_2", Filename='./instrument/'+self.monitor+'.xml', RewriteSpectraMap='True')
        SumSpectra(InputWorkspace=self.monitor+"_2", OutputWorkspace=self.monitor+"_sum", IncludeMonitors=False)
        Rebin(InputWorkspace=self.monitor+"_sum", OutputWorkspace=self.monitor+'_rebin', Params='5000,32,40000')
        print "   INFO: Mantid Workspace updated. ",time.time()-be
        #print time.time()-be 
        rawname=self.monitor+"_sum"
        name1=mtd[rawname] 
        counts1=name1.readY(0)
        print len(counts1) 
        p1="/MR/workspace/MantidData/tmp/sample/"+self.monitor 
        _json1=jsonArray.jsonEncoder(counts1)
        self.neonRedis.set(p1, _json1)
        print "successful in set sample ",self.monitor
        print "==============================================================="
    def dataSet(self):
        #self.get3DXYZMonitor()
        be=time.time()
        setDataToNeon.setMonTofCounts(self.neonRedis, self.monitor)
        setDataToNeon.setXYMonitor(self.neonRedis, self.monitor)
        print "finish dataSet! ",time.time()-be

    def process(self):
        _be=time.time()
        print "   INFO: getDataFromNeon", datetime.datetime.now()
        self.getDataFromNeon()
        print "   INFO: createWorkspace", datetime.datetime.now()
        self.createWorkspace()
        print "   INFO: dataSet", datetime.datetime.now()
        self.dataSet()
        print "===================================="
        print self.monitor, "finish!!! ",time.time()-_be, "seconds"
        print "===================================="

    def run(self):
        #get cell and direct monitor data
        #self.setRefMonitor()
        #print "load cell and direct successfully!!!"
        while True:
            self.process()
            time.sleep(self.refreshtime)

    
class connectNeon():

    def __init__(
        self,
        ip,
        port,
        timeout,
        ):

        self.status=False
        client=None
        self.server=None
        begin=time.time()

        while True:
            if not client:
                self.server = redis.Redis(host=ip, port=port, db=0)
                self.status=True
                print "   INFO: NEON Started"
                break
            else:
                if time.time()-begin>timeout:
                    print "   ERROR: Connect NEON timeout"  
                    sys.exit()
                    #break
                try:
                    print "   INFO: Attempt to connect NEON"  
                    client=redis.client_list()
                except:
                    pass

    def getServer(self):
        return self.server
