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

z_2d = np.reshape(d1[:,0],(rho_steps*2, z_steps))
rho_2d = np.reshape(d1[:,1],(rho_steps*2, z_steps))
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' #'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.05,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('unidirectional')
axes[0,1].set_title('bidirectional')

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)

levels1 = np.linspace(0,lambda1,11)
levels2 = np.linspace(0,lambda1,6)
levels3 = np.linspace(0,lambda2,11)
levels4 = np.linspace(0,lambda2,6)

im1 = axes[0,0].contour(z_2d,rho_2d,efield,levels1,cmap=cmap_set)
im2 = axes[0,1].contour(z_2d,rho_2d,efield_sum,levels2,cmap=cmap_set)
#im2 = axes[0,1].contour(z_2d,rho_2d,efield_sum,levels1,colors='k')
#axes[0,1].clabel(im2, inline=1)
im3 = axes[1,0].contour(z_2d,rho_2d,efield_2,levels3,cmap=cmap_set)
im4 = axes[1,1].contour(z_2d,rho_2d,efield_sum_2,levels4,cmap=cmap_set)


#format colorbar
pad = 0.03
shrink = 0.8
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(np.round(levels2,1))
cb2.set_ticks(levels4)
#cb1.ax.set_yticklabels(['0',str(lambda1)+'%'])
#cb2.ax.set_yticklabels(['0',str(lambda2)+'%'])
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('contour.pdf', bbox_inches='tight')
plt.show()