""" Wannier orbitals for the ethylene molecule 
"""

from ASE import Atom, ListOfAtoms
from Dacapo import Dacapo
from ASE.Utilities.Wannier.Wannier import Wannier
import os


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

    # ethylene molecule
    a = 6.0  # Size of unit cell (Angstrom)
    atoms = ListOfAtoms([
                        Atom('H', (-1.235, 0.936 , 0 ),tag=0),
                        Atom('H', ( 1.235,-0.936 , 0 ),tag=1),
                        Atom('H', ( 1.235, 0.936 , 0 ),tag=1),
                        Atom('H', (-1.235,-0.936 , 0 ),tag=1),
                        Atom('C', ( 0.660, 0.000 , 0 ),tag=1),
                        Atom('C', (-0.660, 0.000 , 0 ),tag=1)],
                        cell=(a,a,a), periodic=True)


    pos = atoms.GetCartesianPositions() + 3
    atoms.SetCartesianPositions(pos)

    # Dacapo calculator:
    calc = Dacapo(planewavecutoff=400, nbands=8, xc='PBE',out='ethylene.nc')
    atoms.SetCalculator(calc)
    tot = atoms.GetPotentialEnergy()

atoms = Dacapo.ReadAtoms('ethylene.nc')
calc = atoms.GetCalculator()

# Initialize the Wannier class
wannier = Wannier(numberofwannier=6,calculator=calc) 

# Perform the localization with specified convergence criterion
wannier.Localize(tolerance=1.0e-8)

# Read the centers and radii of the WFs
for center in  wannier.GetCenters(): 
    print center

centers = wannier.GetCentersAsAtoms()

# plot centers together with atoms
try:
    # first try with VMD
    from ASE.Visualization.VMD import VMD
    VMD(atoms,centers)
except:
    from ASE.Visualization.RasMol import RasMol
    centers.extend(atoms)
    rasmol = RasMol(centers)
    

# make a 3D isosurface plot using VTK
from ASE.Visualization.VTK import VTKPlotElectronicState,VTKPlotAtoms
state = wannier.GetElectronicState(0)

plot = VTKPlotElectronicState(state)
plot.SetRepresentationToIsoSurface2([0.1])
atomsplot = VTKPlotAtoms(atoms,parent=plot) 
plot.Update()
atomsplot.RemoveAvatar(atomsplot.GetAvatars()[0])
plot.Render()