import pylab
import sys
import matplotlib.pyplot as plt

plt.figure(figsize=(8,6))
ax = pylab.subplot(111)

fname = sys.argv[1]

pot = nlread(fname, GridValues)[-1]

# First average over XY
c, pot_z = pot.axisProjection("average","c")
pot_z = pot_z.inUnitsOf(Ang**-3)

# c is the fractional coordinates of all grid points - convert to Cartesian for plotting
c = numpy.array(c)
z= c*pot.primitiveVectors()[2][2].inUnitsOf(Ang)

# This is the grid spacing in z
dz = pot.volumeElement()[2][2].inUnitsOf(Ang)

# For each grid point, define a distance to average over which the potential is averaged
# For a homogenerous material, this would be the same distance for all points, say the lattice constant
av_z = [3.5]*len(pot_z)  

# Number of cells used for the averaging, based on the average averaging distance 
a = numpy.average(av_z)
# "n" can also be set to a pure number
n = 2*int(a/dz+1)

# Make 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]
# Normalize
weights = [sum(kernel) for kernel in kernels]

# 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)])

plt.plot(z[n:len(pot_z)-n], av_pot,color='red',linewidth=3)

ax.set_ylabel(r'EDD ($\AA^{-3}$)', fontsize=20)
ax.set_xlabel(r'Cell Length ($\AA$)', fontsize=20)

ax.grid(color='k',linewidth=1)

plt.tight_layout()

plt.savefig('edp_macro_avg.png',dpi=300)