#!/usr/bin/env python

from Dacapo import Dacapo
from ASE.Visualization.VTK import VTKPlotArray, VTKPlotAtoms

atoms_dimer = Dacapo.ReadAtoms("2H2O.nc")
atoms_molecule1 = Dacapo.ReadAtoms("H2Oa.nc")
atoms_molecule2 = Dacapo.ReadAtoms("H2Ob.nc")

calc_dimer = atoms_dimer.GetCalculator()
calc_molecule1 = atoms_molecule1.GetCalculator()
calc_molecule2 = atoms_molecule2.GetCalculator()


dens_dimer = calc_dimer.GetDensityArray()
dens_molecule1 = calc_molecule1.GetDensityArray()
dens_molecule2 = calc_molecule2.GetDensityArray()

dens_diff = dens_dimer - (dens_molecule1 + dens_molecule2)

plot = VTKPlotArray(dens_diff, atoms_dimer.GetUnitCell())
plot.SetTranslation((0,0,0))

atomplot = VTKPlotAtoms(atoms_dimer, parent = plot)
atomplot.RemoveAvatar(atomplot.unitcell)
plot.Render()


#If one wants the atoms white
#atomavatar.SetAtomColors((1,1,1))

#Get a handle on the isosurfaces and only look at two
#at +-15
isosurface1=plot.GetIsoSurface1()
isosurface1.SetContourRange((-3,3))
isosurface1.SetColorRange((-3,3))
isosurface1.SetNumberOfContours(2)
plot.Render()

#The isosurface at -15 is colored blue, the one at +15 yellow
plot.GetColorTable().SetRGBColors(['blue','yellow'])
plot.GetAvatars()[0].GetActorProperty().SetOpacity(0.35)
plot.Render()