import numpy as np
from numpy.lib import recfunctions as rfn
import matplotlib.pyplot as plt
import sys

#energyStep=100
energyMax=40000
energyBins=400
clusterSize=3

dtypeCluster = [('frameNr', np.int32),('coord', (np.int16,2)),('data',(np.int32,clusterSize*clusterSize))]

def read_cluster_file(filename):
    fd = open(filename,'rb')
    clusters = np.fromfile(fd,dtype=dtypeCluster)
    fd.close
    #rfn.drop_fields(clusters, 'frameNr')
    uCl = rfn.structured_to_unstructured(clusters)
    return uCl

def getEnergyArray(data):
    off=3
    enArray = np.sum(data[:,off:], axis=-1)
    print(data.shape)
    Q = np.empty((4,data.shape[0]))
    print(Q.shape, data.shape)
    Q[0,:]=data[:,off+0]+data[:,off+1]+data[:,off+3]+data[:,off+4]
    Q[1,:]=data[:,off+1]+data[:,off+2]+data[:,off+4]+data[:,off+5]
    Q[2,:]=data[:,off+3]+data[:,off+4]+data[:,off+6]+data[:,off+7]
    Q[3,:]=data[:,off+4]+data[:,off+5]+data[:,off+7]+data[:,off+8]
    print(Q)
    quadArray = np.max(Q,axis=0)
    return enArray,quadArray

def spectrum_roi(x,y,en,xmin,xmax,ymin,ymax):
    global energyBins
    global energymax
    roi=np.where(np.logical_and(np.logical_and(np.logical_and(x>=xmin,x<=xmax),y>=ymin),y<=ymax))
    spectrum3,xedges2=np.histogram(en[roi],bins=energyBins,range=(0,energyMax))
    return spectrum3,xedges2[1:]

def image_cut(x,y,en,emin,emax):
    energyCut=np.where(np.logical_and(en>=emin,en<=emax))
    print(x[energyCut].shape,y[energyCut].shape)
    image,xedges,yedges=np.histogram2d(x[energyCut],y[energyCut],bins=400)
    return image


fname="/mnt/moench_data/tomcat_fluorescence_24022020/clusters_tomo/cu_fibers_27keV_17.clust"
if len(sys.argv)>1:
    fname=sys.argv[1]

cl=read_cluster_file(fname)
print("file read")

enADC,quadADC=getEnergyArray(cl)
print("energy array done",enADC,quadADC)

en=enADC*1000./150.
quad=quadADC*1000./150.
print("energy conversion done:",en,quad)

x=cl[:,1]
y=cl[:,2]

#him=get_hyperimage(x,y,cl)
#print("hyperimage done")
xmin=0
xmax=400
ymin=00
ymax=400


spectrum2,xedges1= spectrum_roi(x,y,quad,xmin,xmax,ymin,ymax)
fig2, axs2 = plt.subplots()
axs2.plot(xedges1,spectrum2,"b-")
fig2.show()
print("sp plotted")

emin=0
emax=40000
image=image_cut(x,y,quad,emin,emax)
fig, axs = plt.subplots()
v=axs.imshow(image,vmax=np.mean(image)*5*np.sqrt(np.var(image)),origin='upper',cmap=plt.cm.jet)
fig.colorbar(v, ax=axs)
fig.show()

print("To get the spectrum for a certain ROI use: sp,x=spectrum_roi(x,y,en,xmin,xmax,ymin,ymax)")
print("To the get the image with a certain energy cut use im=image_cut(x,y,en,emin,emax)")