import pylab
import sys
import matplotlib.pyplot as plt
import matplotlib.patches as ptc
from matplotlib import colors, ticker, cm


#-------------
# Introduction
#-------------
'''
This script needs to be run with the .nc file containing the Hartree 
difference potential from the HarteeDifferencePotential analysis as input file.

Example of run:
atkpython hdp.py device.nc
'''

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

# z-distance between atomic layers in the left and right electrode material in Ang (Determines the width
# of the Gaussian kernels used to calculate the average potential)
gauss_left = 2.275
gauss_right = 5.406

# z-coordinate of the interface position in Ang calculated as the midpoint between 
# the last atom of the left material and the first atom in the right material
zcoord_int = 27.2734

# Output filename
fname = 'hdp'

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

fname2 = sys.argv[1]

# Load the Hartee difference potential
pot = nlread(fname2, HartreeDifferencePotential)[0]

#----------------------------------------------------------
# Calculate the averaged electrostatic difference potential
#----------------------------------------------------------

# Average the potential over the xy-plane
c, pot_z = pot.axisProjection("average","c")
pot_z = pot_z.inUnitsOf(eV)

# c is the fractional coordinates of all grid points, these are converted to Cartesian coordinates for plotting
c = numpy.array(c)
z = c*pot.primitiveVectors()[2][2].inUnitsOf(Ang)
dz = pot.volumeElement()[2][2].inUnitsOf(Ang) # The grid spacing in the z-direction

# For each grid point, average over the z-distance between atoms (in a homogenerous material, 
# this would be the same distance for all points, say the lattice constant)
av_z = numpy.array([0.0]*len(pot_z))
ix = numpy.where(z < zcoord_int)[0] # Left material
av_z[ix] = gauss_left
ix = numpy.where(z >= zcoord_int)[0] # Right material
av_z[ix] = gauss_right

# Number of grid points used to sample the Gaussians, based on the average averaging distance 
a = numpy.average(av_z)
n = 2*int(a/dz+1)

# Gaussian kernels for performing average, one for each grid point
index = numpy.arange(-n,n+1)
kernels = [numpy.exp(-(index*dz/a)**2) for a in av_z]
weights = [sum(kernel) for kernel in kernels] # Create weights for nomalization

# Perform the averaging (n points in the grid are excluded at both ends of the grid)
av_pot = numpy.array([ sum(pot_z[i-n:i+n+1]*kernels[i])/weights[i] for i in range(n, len(pot_z)-n)])
r_pot = av_pot[len(pot_z)-2*n-1]
zmax = z[n:len(pot_z)-n][-1]
zmin = z[n:len(pot_z)-n][0]

#-------------------------
# Print the data to a file
#-------------------------

out = open(fname+".dat","w")
nz = len(z[n:len(pot_z)-n])
for i in range(nz):
    out.write("%f %f\n"%(z[i+n], av_pot[i]-r_pot))

#------------
# Plot figure
#------------

plt.figure(figsize=(5,2))
ax = pylab.subplot(111)

emin = min(av_pot-r_pot)-0.5
emax = max(av_pot-r_pot)+0.5

ax.plot(z[n:len(pot_z)-n], av_pot-r_pot,color='blue',linewidth=0.75)
#ax.plot([zcoord_int,zcoord_int],[emin,emax],color='black',linestyle='-',linewidth=0.5)

ax.set_ylabel(r'$\langle\Delta$ $V_H\rangle$ (eV)', fontsize=10)
ax.set_xlabel(r'Cell Length Z ($\AA$)', fontsize=10)
ax.set_ylim(emin,emax)
ax.set_xlim(zmin,zmax)

ax.spines['top'].set_linewidth(0.5)
ax.spines['bottom'].set_linewidth(0.5)
ax.spines['right'].set_linewidth(0.5)
ax.spines['left'].set_linewidth(0.5)
ax.xaxis.set_tick_params(width=0.25,zorder=20)
ax.yaxis.set_tick_params(width=0.25,zorder=20)
ax.tick_params(direction='in', pad=7.5)

plt.tight_layout()
plt.savefig(fname+".png",dpi=300)