""" workfunction and dipole correction for a Al slab 
"""

import os
import Numeric as num
from ASE.Visualization.gnuplot import gnuplot
from ASE.Visualization.VTK import VTKPlotArray
from Dacapo import Dacapo
from ASE import Atom, ListOfAtoms
from ASE.Utilities.List import List
import os.path,sys
from ASE.IO.Cube import WriteCube

sys.path.remove('/home/camp/lhansen/CamposASE')
sys.path.remove('/usr/lib/python2.2/site-packages/CamposASE')

# from Visualization import *

# check if we have already calculated the nc file
if not os.path.isfile('Al111-O.nc'):

    # setup the Al slab
    alslab =  ListOfAtoms([ Atom('Al',(0, 0, 0)),
                        Atom('Al',(0, 0.5, 0)),
                        Atom('Al',(0.5, 0, 0)),
                        Atom('Al',(0.5, 0.5, 0)),
                        Atom('Al',(0.3333333, 0.166666667, 0.2338/1.25)),
                        Atom('Al',(0.3333333, 0.666666667, 0.2338/1.25)),
                        Atom('Al',(0.8333333, 0.166666667, 0.2338/1.25)),
                        Atom('Al',(0.8333333, 0.666666667, 0.2338/1.25 ))])

    alslab.append(Atom('O',(0.0,0.0, -0.12),tag=1))

    unitcell = [[5.72756492761103, 0, 0],
                [-2.86378246380552, 4.96021672913593, 0],
                [0,0,12.5]]

    alslab.SetUnitCell(unitcell)

    calc = Dacapo(planewavecutoff = 340,
                  densitycutoff = 500,
                  dipole = True,
                  nbands = 30, 
                  out = 'Al111-O.nc', 
                  txtout = 'Al111-O.txt') 


    alslab.SetCalculator(calc)

else: 
    alslab = Dacapo.ReadAtoms('Al111-O.nc') 
    calc = alslab.GetCalculator()


electro = calc.GetElectrostaticPotential()

nx,ny,nz = num.shape(electro) 

# get the x-axis 
unitcell = alslab.GetUnitCell()
xaxis = [float(z)/float(nz)*unitcell[2,2] for z in range(nz)]

# xy average of potential
xyaverage = [num.sum(num.sum(electro[:,:,z]))/(nx*ny) for z in range(nz)]
potential = List(zip(xaxis,xyaverage))
potential.legend = 'Electrostatic potential'
potential.ylabel = 'Electrostatic potential (eV)'
potential.xlabel = 'Distance along z-axis (Angstrom)'

# add fermilevel to the plot
fermilevel = calc.GetFermilevel()
fermilevel = List(zip(xaxis,[fermilevel for z in range(nz)]))
fermilevel.legend = 'Fermi level'

# make the plots
plot = gnuplot(potential) 
fermiplot = gnuplot(fermilevel,parent=plot) 
plot.Update()

electro = calc.GetDensityArray()

# vtkplot = VTKPlotArray(electro,unitcell)
WriteCube(alslab,electro[:,:,:],'test.cube')