#!/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('ptwire.nc'):
    lat = 2.41
    atoms = ListOfAtoms([
          Atom('Pt', ([0*lat,2*lat,2*lat]),tag=0)],
          cell=([lat,4*lat,4*lat]), periodic=True)

    # Dacapo calculator:
    calc = Dacapo(planewavecutoff=350, nbands=8,kpts=(15,1,1),xc='PBE',out='ptwire.nc')
    atoms.SetCalculator(calc)

    # Displace kpoints sligthly, so that the symmetry program in dacapo does
    # not use inversion symmetry to reduce kpoints.
    kpoints = calc.GetBZKPoints()
    kpoints[:,0] += 2e-5
    calc.SetBZKPoints(kpoints)

energy = atoms.GetPotentialEnergy()


# Wannier part
from ASE.Utilities.Wannier.Wannier import Wannier

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

# Use 5 d-orbitals centered at the atom and 1 s-orbital centered at the Pt-Pt bond
# The radius of both types of orbitals is set to 0.4 Angstroms  
initialwannier=[
[[0],2,0.4],
[[0.5,0.5,0.5],0,0.4]
]

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

# Perform the localization
wannier.Localize()

# Translate all the WFs to cell (1,0,0)
wannier.TranslateAllWannierFunctionsToCell([1,0,0])

# Band diagram
from ASE.Utilities.Wannier.HoppingParameters import HoppingParameters

hop=HoppingParameters(wannier,couplingradius=9.0)

# Get a band with 50 k-points
k,band=hop.GetBandDiagram(50,[0,0,0],[0.5,0,0])

# Plot the band using Gnuplot
import Gnuplot
plot=Gnuplot.Gnuplot()
plot('set data style lines')
plot.plot(Gnuplot.Data(k,band[0]),Gnuplot.Data(k,band[1]),Gnuplot.Data(k,band[2]),Gnuplot.Data(k,band[3]),Gnuplot.Data(k,band[4]),Gnuplot.Data(k,band[5]))

# Write the band diagram as a NetCDF file.
hop.WriteBandDiagramToNetCDFFile('ptband.nc',50,[0,0,0],[0.5,0,0])

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