#!/usr/bin/env amspython
# coding: utf-8

# ## Initial imports

import matplotlib.pyplot as plt
from ase.build import bulk as ase_build_bulk
from scm.plams import *

init()


# ## Metal band structure relative to Fermi energy

Cu = fromASE(ase_build_bulk("Cu", "fcc", a=3.6))  # primitive cell
plot_molecule(Cu)


s = Settings()
s.input.ams.Task = "SinglePoint"
s.input.DFTB.Periodic.BandStructure.Enabled = "Yes"
s.input.DFTB.Model = "SCC-DFTB"
s.input.DFTB.ResourcesDir = "DFTB.org/matsci-0-3"
s.runscript.nproc = 1


job = AMSJob(settings=s, name="Cu", molecule=Cu)
job.run()


x, y_spin_up, y_spin_down, labels, fermi_energy = job.results.get_band_structure(unit="eV")
plot_band_structure(x, y_spin_up, None, labels, fermi_energy, zero="fermi")
plt.ylim(-10, 10)
plt.ylabel("$E - E_{Fermi}$ (eV)")
plt.xlabel("Path")
plt.title("Cu with DFTB.org/matsci-0-3")
plt.show()


# ## Semiconductor band structure relative to VBM
#
# For a semiconductor like ZnO you can also choose to put the zero at the VBM ('vbm') or CBM ('cbm')

ZnO = fromASE(ase_build_bulk("ZnO", "wurtzite", a=3.2, c=5.3, u=0.375))
plot_molecule(ZnO, rotation=("60x,60y,80z"))


s = Settings()
s.input.ams.Task = "SinglePoint"
s.input.DFTB.Periodic.BandStructure.Enabled = "Yes"
s.input.DFTB.Model = "GFN1-xTB"
s.runscript.nproc = 1
job = AMSJob(settings=s, molecule=ZnO, name="ZnO")
job.run()


# The below call to ``plot_band_structure`` plots both the spin up and spin down. The spin-down bands are plotted as dashed lines. Note that in this case there is no spin polarization so the spin-down bands perfectly overlap the spin-up bands.

plot_band_structure(*job.results.get_band_structure(unit="eV"), zero="vbmax")
plt.ylim(-10, 10)
plt.ylabel("$E - E_{VBM}$ (eV)")
plt.xlabel("Path")
plt.title("ZnO with GFN1-xTB")
plt.show()


# ## Spin-up and spin-down band structures
# If you perform a spin-polarized calculation you get both spin-up and spin-down bands. Below a spin-polarized DFT+U calculation on NiO is performed together with the BAND engine.

d = 2.085
mol = Molecule()
mol.add_atom(Atom(symbol="Ni", coords=(0, 0, 0)))
mol.add_atom(Atom(symbol="O", coords=(d, d, d)))
mol.lattice = [[0.0, d, d], [d, 0.0, d], [d, d, 0.0]]
plot_molecule(mol)


s = Settings()
s.input.ams.task = "SinglePoint"
s.input.band.Unrestricted = "yes"
s.input.band.XC.GGA = "BP86"
s.input.band.Basis.Type = "DZ"
s.input.band.NumericalQuality = "Basic"
s.input.band.HubbardU.Enabled = "Yes"
s.input.band.HubbardU.UValue = "0.6 0.0"
s.input.band.HubbardU.LValue = "2 -1"
s.input.band.BandStructure.Enabled = "Yes"

job = AMSJob(settings=s, molecule=mol, name="NiO")
job.run()


plot_band_structure(*job.results.get_band_structure(unit="eV"), zero="vbmax")
plt.ylim(-10, 10)
plt.ylabel("$E - E_{VBM}$ (eV)")
plt.xlabel("Path")
plt.title("NiO with DFT+U")
plt.show()