from mantid.simpleapi import *
import time
import jsonArray
import numpy as np
import redis
import sys
import math
from math import sqrt
#import matplotlib.pyplot as plt

def setBankReduction(neonRedis):
    be=time.time()
    # I-Q
    ws=["sample_Q_byM", "sample_Q_byP"]
    path=["MR/workspace/MantidData/reduction/monitor/IQ/","MR/workspace/MantidData/reduction/PC/IQ/"]
    for i in range(2):
        name=mtd[ws[i]]
        value=name.readY(0)
        _tmp=name.readX(0)
        print len(value), len(_tmp)
        q=[]
        counts=[]
        for j in range(len(_tmp)-1):
            q.append(_tmp[j]+(_tmp[j+1]-_tmp[j])/2)
        for k in value:
            counts.append(k)
        #np.delete(_tmp)
        #np.delete(value)
        _json_q=jsonArray.jsonEncoder(q)
        _json_counts=jsonArray.jsonEncoder(counts)
        p1=path[i]+"q"
        p2=path[i]+"counts"
        neonRedis.set(p1,_json_q)
        neonRedis.set(p2,_json_counts)
    print "finish set IQ raw"

def setIQxyz(neonRedis):
    be=time.time()
    wname=mtd["ref_det_wave"]
    inst=wname.getInstrument()
    source=inst.getSource()
    sourcePos=source.getPos()
    sam=inst.getSample()
    samPos=sam.getPos()
    beamline=samPos-sourcePos
    numHist=wname.getNumberHistograms()
    _tmp=wname.readX(0)
    wave=[]
    for j in range(len(_tmp)-1):
        wave.append(_tmp[j]+(_tmp[j+1]-_tmp[j])/2)

    for j in range(numHist):

        det=wname.getDetector(j)
        twoTheta=det.getTwoTheta(samPos, beamline)
        sinTheta=matn.sin(math.pi/180*twoTheta*0.5)
        phi=180/math.pi*math.atan(y,x)

        pos=det.getPos()
        x=pos.X()
        y=pos.Y()
        z=pos.Z()
        ans=(z*z)/(sqrt(z*z)*sqrt(x*x+y*y+z*z))
    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)
    print len(qxy)
    print qxy
    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)
    print type(intensity)
    print np.array(intensity).shape
    print "===========I-Qxy============================"
    p1="/MR/workspace/MantidData/IQxyz/x"
    p2="/MR/workspace/MantidData/IQxyz/y"
    p3="/MR/workspace/MantidData/IQxyz/z"
    p4="/MR/workspace/MantidData/IQxyz/counts"
    _json_q=jsonArray.jsonEncoder(qxy)
    _json_counts=jsonArray.jsonEncoder(intensity)
    neonRedis.set(p1,_json_q)
    neonRedis.set(p2,_json_q)
    neonRedis.set(p3,_json_counts)
    print "finish set IQxy raw!!!"

def setBankRawI(neonRedis):
    be=time.time()
    # I-tof
    wname=mtd["sample_tof"]
    _tmp=wname.readX(0)
    tof=[]
    for j in range(len(_tmp)-1):
        tof.append(_tmp[j]+(_tmp[j+1]-_tmp[j])/2)
    value=wname.readY(0)
    _json_tof=jsonArray.jsonEncoder(tof)
    _json_v=jsonArray.jsonEncoder(value)
    p1="/MR/workspace/MantidData/raw/Itof/tof"
    p2="/MR/workspace/MantidData/raw/Itof/counts"
    neonRedis.set(p1,_json_tof)
    neonRedis.set(p2,_json_v)

    # I-q
    wname=mtd["sample_raw_Q"]
    _tmp1=wname.readX(0)
    q=[]
    for j in range(len(_tmp1)-1):
        q.append(_tmp1[j]+(_tmp1[j+1]-_tmp1[j])/2)
    value1=wname.readY(0)
    _json_q=jsonArray.jsonEncoder(q)
    _json_v1=jsonArray.jsonEncoder(value1)
    p1="/MR/workspace/MantidData/raw/IQ/q"
    p2="/MR/workspace/MantidData/raw/IQ/counts"
    neonRedis.set(p1,_json_q)
    neonRedis.set(p2,_json_v1)

    # I-theta
    wname=mtd["sample_1"]
    counts=[]
    scaAngle=[]
    for i in range(500):
        scaAngle.append(5+0.1*i)
        counts.append(0)
    inst=wname.getInstrument()
    source=inst.getSource()
    sourcePos=source.getPos()
    sam=inst.getSample()
    samPos=sam.getPos()
    beamline=samPos-sourcePos
    numHist=wname.getNumberHistograms()
    for j in range(numHist):
        n2=sum(wname.readY(j))
        #counts.append(n2)
        det=wname.getDetector(j)
        twoTheta=180/math.pi*det.getTwoTheta(samPos, beamline)
        for i in range(len(scaAngle)-1):
            if twoTheta>=scaAngle[i] and twoTheta<scaAngle[i+1]:
                counts[i]=counts[i]+n2
            else:
                continue
    _json_the=jsonArray.jsonEncoder(scaAngle)
    _json_v=jsonArray.jsonEncoder(counts)
    p1="/MR/workspace/MantidData/raw/Itheta/theta"
    p2="/MR/workspace/MantidData/raw/Itheta/counts"
    neonRedis.set(p1,_json_q)
    neonRedis.set(p2,_json_v1)
    print "finish I theta"
    #print scaAngle
    #print counts
def setXYBankData(neonRedis, neonpathDetSet):
    be=time.time()
    name1=mtd["sample_1"]
    numHist=name1.getNumberHistograms()
    _tmpv=[]
    x=[]
    y=[]
    xaxis=[]
    yaxis=[]
    for j in range(numHist):
        _det=name1.getDetector(j)
        _pos=_det.getPos()
        _n2=sum(name1.readY(j))
        x.append(_pos.X())
        y.append(_pos.Y())
        _tmpv.append(_n2)
    xaxis=list(set(x))
    yaxis=list(set(y))
    xaxis.sort()
    yaxis.sort()
    #print "get position"
    print len(xaxis), len(yaxis), len(_tmpv)
    vaxis=[[] for r in range(len(yaxis))]
    num=0
    for p in range(len(yaxis)):
        for q in range(len(xaxis)):
            vaxis[p].append(_tmpv[num])
            num=num+1
    print np.array(xaxis).shape, np.array(yaxis).shape, np.array(vaxis).shape
    #print np.array(xaxis)
    #print np.array(yaxis)
    #print np.array(vaxis)
    _json_x=jsonArray.jsonEncoder(xaxis)
    _json_y=jsonArray.jsonEncoder(yaxis)
    _json_v=jsonArray.jsonEncoder(vaxis)
    _p1=neonpathDetSet+"/xy_image/x"
    _p2=neonpathDetSet+"/xy_image/y"
    _p4=neonpathDetSet+"/xy_image/value"
    neonRedis.set(_p1, _json_x)
    neonRedis.set(_p2, _json_y)
    neonRedis.set(_p4, _json_v)
    print "finish XY Bank!", _p1, _p2, _p4

def setMonTofCounts(neonRedis, monitor):
    rawname=str(monitor)+"_rebin"
    name=mtd[rawname]
    counts=name.readY(0)
    tof=[]
    _tmp=name.readX(0)
    for j in range(len(_tmp)-1):
        tof.append(_tmp[j]+(_tmp[j+1]-_tmp[j])/2)
    print np.array(tof)
    print np.array(counts)
    _p1="/SANS/workspace/MantidData/"+monitor+"/tof"
    _p2="/SANS/workspace/MantidData/"+monitor+"/counts"
    _json_x=jsonArray.jsonEncoder(tof)
    _json_y=jsonArray.jsonEncoder(counts)
    print _p1, _p2
    neonRedis.set(_p1, _json_x)
    neonRedis.set(_p2, _json_y)
    co=sum(name.readY(0))
    print monitor, "  total neutron counts: ",co

def setXYMonitor(neonRedis, monitor):
    x=[]
    y=[]
    z=[]
    xaxis=[]
    yaxis=[]
    vtmp=[]
    rawName=monitor+"_2"   
    name=mtd[rawName]
    numHist=name.getNumberHistograms()
    print "hist of monitor: ", numHist
    for i in range(numHist):
        det=name.getDetector(i)
        pos=det.getRelativePos()
        #pos=det.getPos()
        n2=sum(name.readY(i))
        x.append(pos.X())
        y.append(pos.Y())
        #z.append(pos.Z())
        vtmp.append(n2)
    xaxis=list(set(x))
    yaxis=list(set(y))
    xaxis.sort()
    yaxis.sort()
    vaxis=[[] for m in range(len(yaxis))]
    n=0
    for j in range(len(yaxis)):
        for k in range(len(xaxis)):
            vaxis[j].append(vtmp[n])
            n=n+1
    _p1="/SANS/workspace/MantidData/"+monitor+"/axis_x"
    _p2="/SANS/workspace/MantidData/"+monitor+"/axis_y"
    _p3="/SANS/workspace/MantidData/"+monitor+"/value"
    print _p1, _p2, _p3
    _json_x=jsonArray.jsonEncoder(xaxis)
    _json_y=jsonArray.jsonEncoder(yaxis)
    _json_v=jsonArray.jsonEncoder(vaxis)
    neonRedis.set(_p1, _json_x)
    neonRedis.set(_p2, _json_y)
    neonRedis.set(_p3, _json_v)
    
