import numpy as np 
import matplotlib
import matplotlib.pyplot as plt

rho_steps = 81
z_steps = 401
f1 = 'field_out_2d_u=exp.dat'
f2 = 'field_out_2d_u=3.dat'

d1 = np.genfromtxt(f1)
d2 = np.genfromtxt(f2)

z = d1[:,0] # in microns
rho = d1[:,1] # in microns

lambda1=22.4
lambda2=70.5

#this is I=|E|^2 and not electric field strength
efield = (d1[:,3]-d1[:,4].min())/d1[:,4].max() * lambda1
efield_sum = (d1[:,4]-d1[:,4].min())/d1[:,4].max() * lambda1

efield_2 = (d2[:,3]-d2[:,4].min())/d2[:,4].max() * lambda2
efield_sum_2 = (d2[:,4]-d2[:,4].min())/d2[:,4].max() * lambda2

efield = np.reshape(efield, (rho_steps*2, z_steps))
efield_sum = np.reshape(efield_sum, (rho_steps*2, z_steps))
efield_2 = np.reshape(efield_2, (rho_steps*2, z_steps))
efield_sum_2 = np.reshape(efield_sum_2, (rho_steps*2, z_steps))

#plotting
extent = [  z[0], z[-1],-rho[-1], rho[-1] ]
aspect_ratio = 1
inter = 'none'
cmap_set = 'jet' #'jet' #'hot'
#inter = 'spline36'

matplotlib.rcParams.update({'font.size': 14})
matplotlib.rcParams['xtick.direction'] = 'out'
matplotlib.rcParams['ytick.direction'] = 'out'
matplotlib.rc('font',**{'family':'sans-serif'})

fig, axes = plt.subplots(nrows=2, ncols=2, sharex=False, sharey=False)
fig.subplots_adjust(hspace=-0.49,wspace=0.05)#,left=0.3,right=0.7)

for ax in axes.ravel():
	#ax.set_adjustable('box-forced')
	ax.set_xticks([])
	ax.set_yticks([])
	ax.set_xlim([-2.8,2.8])
	ax.set_ylim([-1.7,1.7])

axes[0,0].set_title('one side')
axes[0,1].set_title('4PI')

axes[1,0].set_xticks([-2,0,2])
axes[1,1].set_xticks([-2,0,2])
axes[1,0].set_xlabel('z ($\mu$m)')
axes[1,1].set_xlabel('z ($\mu$m)')
axes[1,0].xaxis.set_ticks_position('bottom')
axes[1,1].xaxis.set_ticks_position('bottom')

axes[0,0].set_yticks([-1,0,1])
axes[1,0].set_yticks([-1,0,1])
axes[0,0].set_ylabel('x ($\mu$m)')
axes[1,0].set_ylabel('x ($\mu$m)')
axes[0,0].yaxis.set_ticks_position('left')
axes[1,0].yaxis.set_ticks_position('left')

im1 = axes[0,0].imshow(efield,extent=extent,aspect=aspect_ratio,vmin=0,vmax=lambda1,interpolation=inter,cmap=cmap_set)
im2 = axes[0,1].imshow(efield_sum,extent=extent,aspect=aspect_ratio,vmin=0,vmax=lambda1,interpolation=inter,cmap=cmap_set)
im3 = axes[1,0].imshow(efield_2,extent=extent,aspect=aspect_ratio,vmin=0,vmax=lambda2,interpolation=inter,cmap=cmap_set)
im4 = axes[1,1].imshow(efield_sum_2,extent=extent,aspect=aspect_ratio,vmin=0,vmax=lambda2,interpolation=inter,cmap=cmap_set)

#format colorbar
pad = 0.03
shrink = 0.35
aspect = 12
cb1=fig.colorbar(im1, ax=[axes[0,0],axes[0,1]], shrink=shrink, aspect=aspect, pad=pad)
cb2=fig.colorbar(im3, ax=[axes[1,0],axes[1,1]], shrink=shrink, aspect=aspect, pad=pad)
cb1.set_ticks([0+lambda1*0.04,lambda1-lambda1*0.1])
cb2.set_ticks([0+lambda2*0.04,lambda2-lambda2*0.1])
cb1.ax.set_yticklabels(['0',str(round(lambda1/100,2))])
cb2.ax.set_yticklabels(['0',str(round(lambda2/100,2))])
cb1.ax.yaxis.set_ticks_position('none') 
cb2.ax.yaxis.set_ticks_position('none')
cb1.ax.set_title('$\Lambda$')
cb2.ax.set_title('$\Lambda$')

#adding scalebar
#axes[1,1].text(1.5, -1.3,'1$\mu$m', color='w')
#axes[1,1].plot((1.5,2.5),(-1.5,-1.5),'w',linewidth=3)

#adding label
#label=['a','b','c','d']
#for i in range(0,4):
	#axes.ravel()[i].text(-2.6, 1.2 ,label[i], color='w')

plt.savefig('field_hori.pdf', bbox_inches='tight')
#plt.show()