from ase.io import read
import numpy
import pylab
atoms=read("LOCPOT",format="vasp")

locpot_file=open("LOCPOT","r")
for i in range(9+len(atoms)):
    locpot_file.readline()
line=locpot_file.readline()
tmp=line.split()
nxgrid=int(tmp[0])
nygrid=int(tmp[1])
nzgrid=int(tmp[2])
ngrid=nxgrid*nygrid*nzgrid

print(nxgrid,nygrid,nzgrid,ngrid)

locpot_data=[]
while line!="":
    line=locpot_file.readline()
    tmp=line.split()
    for datum in tmp:
        locpot_data.append(float(datum))
locpot_file.close()
print(len(locpot_data))
locpot_data=numpy.array(locpot_data)
#locpot_data=locpot_data*atoms.get_volume()
locpot_data=numpy.reshape(locpot_data,[nzgrid,nygrid,nxgrid])
locpot_avg=numpy.mean(locpot_data,axis=-1)
locpot_avg=numpy.mean(locpot_avg,axis=-1)
print(numpy.shape(locpot_data))
print(numpy.shape(locpot_avg))
pylab.figure()
zvals=numpy.linspace(0,atoms.cell[2,2],1260,endpoint=False)
print(zvals)
pylab.plot(zvals,locpot_avg,'k-')
pylab.xlabel(r"Height, z / $\AA$")
pylab.ylabel(r"Electrostatic Potental / eV")
pylab.savefig("LOCPOT_Average.png")

