#!/usr/bin/env amspython

from scm.plams import *
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import MaxNLocator


def init_qe(mag, kp):

    # Settings for QE
    s = Settings()
    s.runscript.preamble_lines = ["export SCM_DISABLE_MPI=1"]

    # PP
    pp_dict = {
        "Mg": "GBRV/mg_lda_v1.4.uspp.F.UPF",
        "N": "GBRV/n_lda_v1.2.uspp.F.UPF",
        "O": "GBRV/o_lda_v1.2.uspp.F.UPF",
    }

    pp_string = "\n     ".join(f"{k} {v}" for k, v in pp_dict.items())
    s.input.QuantumEspresso.Pseudopotentials._1 = pp_string

    # k-grid
    s.input.QuantumEspresso.K_Points._h = "automatic"
    s.input.QuantumEspresso.K_Points._1 = f"{kp} {kp} {kp} 0 0 0"

    # SCF
    s.input.QuantumEspresso.Electrons.conv_thr = 1e-8

    # System
    s.input.QuantumEspresso.System = Settings(ecutwfc=60.0, tot_magnetization=mag, nspin=2)

    return s


init()

# Create conventional cell of MgO
mol = Molecule()
mol.add_atom(Atom(symbol="Mg", coords=(0.0, 0.0, 0.0)))
mol.add_atom(Atom(symbol="O", coords=(2.105, 2.105, 2.105)))
mol.lattice = [[0.000, 2.105, 2.105], [2.105, 0.000, 2.105], [2.105, 2.105, 0.000]]
mol_conv = mol.supercell([[-1, 1, 1], [1, -1, 1], [1, 1, -1]])

# Create 2x2x2 supercell
mol_p = mol_conv.supercell(2, 2, 2)
latt_ref = mol.lattice
a2 = latt_ref[0][0] / 2.0

# Defect
mol_0 = mol_p.copy()
mol_0.atoms[1].symbol = "N"

# References
mol_O2 = Molecule()
mol_O2.add_atom(Atom(symbol="O", coords=(a2, a2, a2)))
mol_O2.add_atom(Atom(symbol="O", coords=(a2 + 1.21, a2, a2)))
mol_O2.lattice = latt_ref

mol_N2 = Molecule()
mol_N2.add_atom(Atom(symbol="N", coords=(a2, a2, a2)))
mol_N2.add_atom(Atom(symbol="N", coords=(a2 + 1.1, a2, a2)))
mol_N2.lattice = latt_ref

ef_list = []
kp_list = [1, 2, 3, 4, 5]
mag_dict = {"neutral_perfect": 0, "neutral_defect": 1, "O2": 2, "N2": 0}
for kp in kp_list:

    # Perfect cell
    name = f"neutral_perfect"  # p
    s = init_qe(mag_dict[name], kp)
    s.input.ams.task = "SinglePoint"
    job = AMSJob(settings=s, molecule=mol_p, name=name)
    res = job.run()
    energy_p = job.results.get_energy(unit="eV")

    atoms = mol.atoms
    n_p = len(atoms)

    # Neutral defect
    name = f"neutral_defect"  # 0
    s = init_qe(mag_dict[name], kp)
    s.input.ams.task = "SinglePoint"
    job = AMSJob(settings=s, molecule=mol_0, name=name)
    res = job.run()
    energy_0 = job.results.get_energy(unit="eV")

    # Chemical potentials
    name = f"O2"
    s = init_qe(mag_dict[name], kp)
    s.input.ams.task = "GeometryOptimization"
    job = AMSJob(settings=s, molecule=mol_O2, name=name)
    res = job.run()
    mu_O = job.results.get_energy(unit="eV") / len(mol_O2.atoms)

    name = f"N2"
    s = init_qe(mag_dict[name], kp)
    s.input.ams.task = "GeometryOptimization"
    job = AMSJob(settings=s, molecule=mol_N2, name=name)
    res = job.run()
    mu_N = job.results.get_energy(unit="eV") / len(mol_N2.atoms)

    # Compute neutral vacancy formation energy
    mu = mu_O - mu_N
    E0_f = energy_0 - energy_p + mu
    ef_list.append(E0_f)

    # Print energies
    print(f">>> K-grid {kp}x{kp}x{kp}")
    print(f">> Perfect supercell E_p = {energy_p:6.3f} eV")
    print(f">> Neutral defective supercell E_0 = {energy_0:6.3f} eV")
    print(f">> Chemical potential mu = {mu:6.3f} eV")
    print(f">>> Formation energy E0_f = {E0_f:6.3f} eV")

fig, ax = plt.subplots()
ax.scatter(kp_list, ef_list)
xfit = np.linspace(0, max(kp_list) * 1.1, 100)
ax.plot(xfit, [ef_list[-1]] * len(xfit), color="r", ls="--")
ax.xaxis.set_major_locator(MaxNLocator(integer=True))
ax.ticklabel_format(useOffset=False, style="plain")
ax.set_xlabel("Number of kpoints")
ax.set_ylabel("Defect formation energy (eV)")
plt.tight_layout()
plt.savefig("E0_form.png")

finish()