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

efield = (d1[:,3]-d1[:,4].min())/d1[:,4].max()
efield_sum = (d1[:,4]-d1[:,4].min())/d1[:,4].max()

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

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 = [ -rho[-1], rho[-1], z[0], z[-1] ]
extent = [ -rho[-1], rho[-1], z[0], z[-1] ]
aspect_ratio = 1
inter = 'none'
cmap_set = 'hot'
#inter = 'spline36'

matplotlib.rcParams.update({'font.size': 16})
fig, axes = plt.subplots(nrows=2, ncols=2, sharex=False, sharey=False)
fig.subplots_adjust(hspace=0,wspace=0.2,left=0.3,right=0.7)

limits=[]
for ax in axes.ravel():
	#ax.set_adjustable('box-forced')
	ax.set_xticks([])
	ax.set_yticks([])
	ax.set_ylim([-3,3])

axes[1,0].set_xticks([-1,0,1])
axes[1,1].set_xticks([-1,0,1])
axes[1,0].set_xlabel('x (um)')
axes[1,1].set_xlabel('x (um)')

axes[0,0].set_yticks([-2,0,2])
axes[1,0].set_yticks([-2,0,2])
axes[0,0].set_ylabel('z (um)')
axes[1,0].set_ylabel('z (um)')

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

cb1=fig.colorbar(im1, ax=[axes[0,0],axes[0,1]] ,fraction=0.0312)
cb2=fig.colorbar(im3, ax=[axes[1,0],axes[1,1]] ,fraction=0.0312)
cb1.set_ticks([0,1])
cb2.set_ticks([0,1])
#fig.subplots_adjust(right=0.8)
#cbar_ax = fig.add_axes([0.95, 0.15, 0.02, 0.7])


plt.savefig('test.eps', bbox_inches='tight')
plt.show()