""" plot wavefunction, together with atoms using VTK. 

CO molecule is used as an example. 

Usage: python -i plotwavefunction.py <bandno>

"""
from ASE import Atom,ListOfAtoms
from ASE.Visualization.VTK import VTKPlotWaveFunction
from ASE.Visualization.VTK import VTKPlotAtoms
from Dacapo import Dacapo
import os.path,sys

args = sys.argv[1:]
if len(args) != 1:
   band = 0
else: 
   band = int(args[0])

print 'plotting wavefunction for band number ',band

# check if we have already calculated the nc file
if os.path.isfile('CO_in_a_box.nc'): 

     # start calculator from file 
     atoms = Dacapo.ReadAtoms(filename='CO_in_a_box.nc')
     calc = atoms.GetCalculator()

else: 

     # start from scratch
     atoms = ListOfAtoms([Atom('C', (2, 2, 2)),
                          Atom('O', (3.1, 2, 2))],
                          cell=(4, 4, 4), periodic=1)

     calc = Dacapo(planewavecutoff=300,    # in eV
                   nbands=6,               # 1 extra empty bands
                   out='CO_in_a_box.nc',txtout='CO_in_a_box.txt')

     atoms.SetCalculator(calc)

# make a combined plot of the wavefunction and the atoms
atomplot = VTKPlotAtoms(atoms)
wfplot = VTKPlotWaveFunction(atoms,band=band,parent=atomplot)
atomplot.Update()

# The appearence of the individual atomic elements can also be changed.
# Finding aluminium atom avatar (found in the dictionary of species avatars)
p1_O=atomplot.GetDictOfSpecies()['O']
p1_C=atomplot.GetDictOfSpecies()['C']

# Change the radius to 2.0 AA and set the color to blue
p1_C.SetRadius(0.5)
p1_O.SetRadius(0.7)

#The color is given in an rgb (red,green,blue) scale
p1_C.SetColor((1,0,0))
p1_O.SetColor((0,1,0))
atomplot.Render()

# The unitcell can be removed
unitcell=atomplot.unitcell
atomplot.RemoveAvatar(unitcell)
atomplot.Render()