#!/usr/bin/env python

# Localized Wannier orbitals for a Pt wire

from Dacapo import Dacapo
from ASE import Atom,ListOfAtoms
from ASE.Visualization.RasMol import RasMol
import os

# check if we have already calculated the nc file
if not os.path.isfile('pt4.nc'):
    
    # Pt wire
    lat = 2.41
    atoms = ListOfAtoms([
                        Atom('Pt', ([0*lat,2*lat,2*lat]),tag=0),
                        Atom('Pt', ([1*lat,2*lat,2*lat]),tag=0),
                        Atom('Pt', ([2*lat,2*lat,2*lat]),tag=0),
                        Atom('Pt', ([3*lat,2*lat,2*lat]),tag=0)],
                        cell=([4*lat,4*lat,4*lat]), periodic=True)


    # Dacapo calculator:
    calc = Dacapo(planewavecutoff=350, nbands=30, xc='PBE',out='pt4.nc')

    atoms.SetCalculator(calc)

    plot = RasMol(atoms.Repeat((2,2,2)))
    tot = atoms.GetPotentialEnergy()

# Wannier part
from Dacapo import Dacapo
from ASE.Utilities.Wannier.Wannier import Wannier
import Numeric as num

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

#Initialize Wannier class
wannier = Wannier(numberofwannier=4*(5+1),calculator=calc) 

# Perform localization 
wannier.Localize()

# Print the WF centers
centers = wannier.GetCenters()
for center in centers:
    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(10)

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