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

CO molecule is used as an example. 

Usage: python -i plotwavefunction.py <bandno>

"""

#Adjust the name of the nc output file here (if you have a different one)
ncfilename = 'CO.nc'

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

atoms = Dacapo.ReadAtoms(filename=ncfilename)
calc = atoms.GetCalculator()


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