import pylab
import sys
import matplotlib.pyplot as plt

#------------------------
# User defined parameters
#------------------------

# Name of the file containing the IVCurve analysis
fname = 'device_120.hdf5'

# Name of .dat files containing the Hartree difference potential
hdp_files = ['hdp0.05','hdp0.1','hdp0.15','hdp0.2','hdp0.25','hdp0.3']

# The Schottky barrier calculated with pldos_hdp.py at zero bias in eV
phipot = 0.606

# The conduction band minima of the bulk right electrode w.r.t the chemical potential of the left electrode  calculated with pldos_hdp.py in eV
CB_min = 0.00832

# Name of output file containing the barriers associated with the thermionic emission process in eV
# for each bias point
fname2 = 'phi_F'

# Name of output file containing the energy of maximum spectral current in eV for each bias point
fname3 = 'Emax'

#-----------
# Load files
#-----------

iv = nlread(fname, IVCurve)[0]
transmission_spectra = iv.transmissionSpectra()

# Consider positive bias only
pos_spectra = []
bias = []
for t in transmission_spectra:
  if float(t.bias()) > 0:
    pos_spectra.append(t)
    bias.append(t.bias().inUnitsOf(Volt))

I = []
# Loop through biases
for spectrum in pos_spectra:
	I.append(spectrum.spectralCurrent())
	E = spectrum.energies()

I = numpy.array(I)
E = numpy.array(E)
bias = numpy.array(bias)

#----------------------------------------------------------------------
# Calculate the barrier associated with the thermionic emission process
#----------------------------------------------------------------------

fig = plt.figure(figsize=(7,3))
ax1 = fig.add_subplot(121)

m = float(len(bias))

phimin = []
phimax = []
for n in range(len(hdp_files)):
	V = numpy.loadtxt(hdp_files[n]+'.dat')
	phimin.append(V[-1,1]) 
	phimax.append(max(V[:,1]))
	
	emax = max(V[:,1])+bias[n]+CB_min
	if n == 0:
		emin = V[-1,1]+bias[n]+CB_min

	ax1.plot(V[:,0],V[:,1]+bias[n]+CB_min,color=pylab.cm.jet(n/m),linewidth=1)

phi_F = numpy.array(phimax) - numpy.array(phimin)

zmin = min(V[:,0])
zmax = max(V[:,0])

#-------------------------------------------------
# Calculate the energy of maximum spectral current
#-------------------------------------------------

Emax = []
for n in range(len(bias)): 
    Imax = numpy.where(abs(I[n,:]) == numpy.amax(abs(I[n,:])))[0][0]
    Emax.append(E[Imax])

#--------------
# Write to file
#--------------

out = open(fname2+".dat","w")
for n in range(len(bias)):
	out.write("%f %f \n"%(bias[n], phi_F[n]))

out = open(fname3+".dat","w")
for n in range(len(bias)):
	out.write("%f %f \n"%(bias[n], Emax[n]))

for n in range(len(bias)):
	print()
	print("Bias:", bias[n], "V")
	print("The barrier of thermionic emission =", phi_F[n], 'eV')
	print("The energy of maximum spectral current =", Emax[n]+bias[n]/2, 'eV')
print()
print("--------------------------------------------------------------------------------")

#--------------------------------------------
# Plot the Hartree difference potential
#--------------------------------------------

ax1.set_xlabel(r'Cell Length Z ($\AA$)', fontsize=15)
ax1.set_ylabel(r'$E$-$\chi$-$\mu_{L}$ (eV)', fontsize=15)
ax1.set_xlim(zmin,zmax)
ax1.set_ylim(emin-0.05,emax+0.05)
for label in ax1.xaxis.get_ticklabels()[::2]:
    label.set_visible(False)

#--------------------------
# Plot the Spectral Current
#--------------------------

ax2 = fig.add_subplot(122)

for n in range(len(bias)): 
    ax2.plot(abs(I[n,:]),E+bias[n]/2,linewidth=1,color=plt.cm.jet((n)/m),label=str(bias[n])+" V")

ax2.plot([5*10**-10,1],[phipot,phipot],color='black',linewidth=1,zorder=10,linestyle='--')

ax2.text(I[-1,Imax],phipot+0.02,'$\Phi^{pot}$')

ax2.set_xscale('log')
ax2.set_ylabel('$E$-$\mu_{L}$ (eV)', fontsize=15)
ax2.set_xlabel('$I(E)$ (A eV$^{-1}$)', fontsize=15)
ax2.set_ylim(emin-0.05,emax+0.05)
ax2.set_xlim(1*10**-9,I[-1,Imax]+10*I[-1,Imax])

ax2.set_position([0.55,0.1,0.5,0.8])
ax2.legend(bbox_to_anchor=(1.05, 1), loc=2,fontsize=10,frameon=False,ncol=1, borderaxespad=0.)
plt.tight_layout()
plt.savefig('spectral_current.png',dpi=300)