import numpy as np
import matplotlib.pyplot as plt
import pickle
from lmfit import Model,Parameters
from matplotlib import cm
from PIL import Image as img

def g2exp(x,peak,cen,wid,base):
	return base+peak*np.exp(-2*np.abs(x-cen)/wid)

foldername=''

#pickle data is tuple(filename,lmfit.result)
with open(foldername+'processed_data_multiprocessed_20260519_1943.pickle','rb') as f:
	temp_result=pickle.load(f)
with open(foldername+'processed_data_multiprocessed_20260519_1943_corrected_values.pickle','rb') as f:
	temp_result1=pickle.load(f)

current_start=6
current_end=6
current_step=1
num_trials=1

h_start=0
h_end=14000
h_step=700
v_start=0
v_end=-17000
v_step=-850

num_current_steps=int((current_end-current_start)/current_step+1)
num_h_steps=int((h_end-h_start)/h_step+1)
num_v_steps=int((v_end-v_start)/v_step+1)

current_range=np.arange(current_start,current_end,current_step)
h_range=np.arange(h_start,h_end+h_step,h_step)
v_range=np.arange(v_start,v_end+v_step,v_step)

peak=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
cen=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
wid=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
base=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
peak_sigma=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
cen_sigma=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
wid_sigma=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
base_sigma=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))

peak_near=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
cen_near=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
wid_near=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
base_near=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
peak_sigma_near=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
cen_sigma_near=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
wid_sigma_near=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))
base_sigma_near=np.full((num_h_steps,num_v_steps,num_current_steps), np.nan)#np.zeros((num_h_steps,num_v_steps,num_current_steps))

for item in temp_result:
	filename=item[0][:-3].split('_')
	h_index=int((int(filename[0])-h_start)/h_step)
	v_index=int((int(filename[1])-v_start)/v_step)
	current_index=int((int(filename[2][:-2])-current_start)/current_step)-1
	loc=(h_index,v_index,current_index)
	print(filename)
	print(loc)
	print(item[2])
	print('\n')

	#if loc = (11,4,):
	#	print()

	#print(item[1].summary()['params'])
	peak[loc]=item[1].summary()['params'][0][1]#item[1].summary()['best_values']['peak']
	cen[loc]=item[1].summary()['params'][1][1]#item[1].summary()['best_values']['cen']
	wid[loc]=item[1].summary()['params'][2][1]#item[1].summary()['best_values']['wid']
	base[loc]=item[1].summary()['params'][3][1]#item[1].summary()['best_values']['base']
	peak_sigma[loc]=item[1].summary()['params'][0][7]#item[1].summary()['best_values']['peak']
	cen_sigma[loc]=item[1].summary()['params'][1][7]#item[1].summary()['best_values']['cen']
	wid_sigma[loc]=item[1].summary()['params'][2][7]#item[1].summary()['best_values']['wid']
	base_sigma[loc]=item[1].summary()['params'][3][7]#item[1].summary()['best_values']['base']
	#print(item[1].fit_report())
	#print(item[1].summary()['params'][1][1])
	#print(item[1].summary()['params'][1][7])
	#print(item[1].summary().keys())

print('going through corrected values')

for item in temp_result1:
	filename=item[0][:-3].split('_')
	h_index=int((int(filename[0])-h_start)/h_step)
	v_index=int((int(filename[1])-v_start)/v_step)
	current_index=int((int(filename[2][:-2])-current_start)/current_step)-1
	loc=(h_index,v_index,current_index)
	print(filename)
	print(loc)
	print(item[2])
	print('\n')

	#if loc = (11,4,):
	#	print()

	#print(item[1].summary()['params'])
	peak[loc]=item[1].summary()['params'][0][1]#item[1].summary()['best_values']['peak']
	cen[loc]=item[1].summary()['params'][1][1]#item[1].summary()['best_values']['cen']
	wid[loc]=item[1].summary()['params'][2][1]#item[1].summary()['best_values']['wid']
	base[loc]=item[1].summary()['params'][3][1]#item[1].summary()['best_values']['base']
	peak_sigma[loc]=item[1].summary()['params'][0][7]#item[1].summary()['best_values']['peak']
	cen_sigma[loc]=item[1].summary()['params'][1][7]#item[1].summary()['best_values']['cen']
	wid_sigma[loc]=item[1].summary()['params'][2][7]#item[1].summary()['best_values']['wid']
	base_sigma[loc]=item[1].summary()['params'][3][7]#item[1].summary()['best_values']['base']
	#print(item[1].fit_report())
	#print(item[1].summary()['params'][1][1])
	#print(item[1].summary()['params'][1][7])
	#print(item[1].summary().keys())



"""
#masking based on tc
mask=wid>10
wid[mask]=np.nan
cen[mask]=np.nan
mask=wid<1
wid[mask]=np.nan
cen[mask]=np.nan
"""


#masking based on accidentals
mask=base<225
wid[mask]=np.nan
cen[mask]=np.nan

"""
#masking based on distance
mask=cen*1e-9*3/2*1e8<1200
print(mask)
wid[mask]=np.nan
cen[mask]=np.nan
"""

process_fail=[(19,11),(12,14),(11,18),(10,18),(11,20),(2,16),(2,17)]#[(10,18),(11,20),(11,10),(12,4),(12,10),(12,13),(12,14),(12,15),(12,18)]

for i in range(21): #horizontal
	for j in range(21): #vertical
		if (i,j) in process_fail:
			cen[i,j]=np.nan


#splitting near and far

for i in range(21): #horizontal
	for j in range(21): #vertical
		if i >=14 and j <=15 and j>=10:
			cen_near[i,j]=cen[i,j]
			cen[i,j]=np.nan
		if i <=1 and j >=2 and j <=4:
			cen_near[i,j]=cen[i,j]
			cen[i,j]=np.nan

#plt.imshow(np.flip(np.rot90(cen[:,:,0]),1))
#plt.imshow(np.flip(np.rot90(cen[:,:,0]*1e-9*3/2*1e8),1))
#plt.imshow(np.flip(np.rot90(wid[:,:,0]),1),vmax=10)
#plt.imshow(np.flip(np.rot90(base[:,:,0]),1))
#plt.colorbar()
#plt.show()

x=np.arange(0,21)
y=np.arange(0,21)
x,y=np.meshgrid(y,x)

plotstyle='image'

font_size=17

if plotstyle=='surface':
	fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
	xx,yy=x.ravel(),y.ravel()
	bottom=0#(np.ones_like(cen[:,:,0])*np.nanmin(cen[:,:,0]*1e-9*3/2*1e8)).ravel()
	temp_cen=cen
	temp_cen[np.isnan(temp_cen)]=np.nanmin(cen[:,:,0])
	top=np.flip(np.rot90(temp_cen[:,:,0]*1e-9*3/2*1e8),1).ravel()-np.nanmin(cen[:,:,0]*1e-9*3/2*1e8)#np.ones_like(cen[:,:,0])*np.max(cen[:,:,0])
	print(top)
	print(bottom)
	width = depth = 1
	ax.bar3d(xx, yy, bottom, width, depth, top, shade=True)
	#surf = ax.plot_surface(x,y,np.flip(np.rot90(-1*cen[:,:,0]*1e-9*3/2*1e8),1),cmap=cm.coolwarm)
	
elif plotstyle=='image':
	fig, axs = plt.subplots(nrows=2,figsize=(6,8))#(figsize=(8.27,11.69))
	fig.subplots_adjust(left=0,right=0.9,top=0.95,bottom=0.05,hspace=0)
	#ax.set_title('distance (m) with masking, accidentals>225',fontsize=font_size)
	im2=axs[1].imshow(np.flip(np.rot90(cen_near[:,:,0]*1e-9*3/2*1e8),1),aspect=850/700,cmap='vanimo')#'managua'
	cbar2=fig.colorbar(im2,ticks=np.arange(214.5,217,0.5))
	cbar2.ax.tick_params(labelsize=font_size)
	im1=axs[1].imshow(np.flip(np.rot90(cen[:,:,0]*1e-9*3/2*1e8),1),aspect=850/700,cmap='managua')#'managua'
	cbar1=fig.colorbar(im1)
	cbar1.ax.tick_params(labelsize=font_size)
	axs[1].set_xticks(np.arange(0,22,5),np.arange(0,22,5),fontsize=font_size)
	axs[1].set_yticks(np.arange(0,22,5),np.flip(np.arange(0,22,5)),fontsize=font_size)
	for i in range(20):
		axs[1].axvline(x=i+0.5, color='0.8', linestyle='-')
		axs[1].axhline(y=i+0.5, color='0.8', linestyle='-')
	poi=[(2,20-12),(9,20-7),(18,20-15)]
	for point in poi:
		line_color='r'
		axs[1].plot((point[0]-0.5,point[0]+0.5),(point[1]-0.5,point[1]-0.5),color=line_color)
		axs[1].plot((point[0]-0.5,point[0]+0.5),(point[1]+0.5,point[1]+0.5),color=line_color)
		axs[1].plot((point[0]-0.5,point[0]-0.5),(point[1]-0.5,point[1]+0.5),color=line_color)
		axs[1].plot((point[0]+0.5,point[0]+0.5),(point[1]-0.5,point[1]+0.5),color=line_color)

	pos1=axs[1].get_position()
	#image=img.open('20260506_121808_cropped_again_again_again.jpg')
	#axs[0].imshow(image,aspect=0.9)
	image=img.open('20260506_121808_cropped_again_again.jpg')
	axs[0].imshow(image,aspect=pos1.height/pos1.width)
	axs[0].set_axis_off()
	pos0=axs[0].get_position()
	#pos1=axs[1].get_position()
	axs[0].set_position([pos1.x0, pos0.y0, pos0.width, pos0.height])
	
	fig.text(0.975, 0.25, 'Distance (m)', horizontalalignment='center',verticalalignment='center',rotation='vertical',fontsize=font_size)

	#image=img.open('20260506_121808_cropped_again.png')
	#image=image.resize((150,int(150*2269/1937)))
	#fig.figimage(image,465,75)

#print(temp_result)
plt.show()

fig, ax = plt.subplots(figsize=(11.69,8.27))
ax.set_title('accidentals (counts/s)',fontsize=font_size)
im=ax.imshow(np.flip(np.rot90(base[:,:,0]),1),cmap='managua')
ax.set_xticks(np.arange(0,22,5),np.arange(0,22,5),fontsize=font_size)
ax.set_yticks(np.arange(0,22,5),np.flip(np.arange(0,22,5)),fontsize=font_size)
cbar=fig.colorbar(im)
cbar.ax.tick_params(labelsize=font_size)
for i in range(20):
	ax.axvline(x=i+0.5, color='0', linestyle='-')
	ax.axhline(y=i+0.5, color='0', linestyle='-')
#plt.show()