#!/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):

    s = Settings()
    s.input.ams.Task = 'GeometryOptimization'
    s.input.QuantumESPRESSO = Settings()
    s.input.QuantumESPRESSO.K_Points._h = 'automatic'
    s.input.QuantumESPRESSO.K_Points._1 = f'{kp} {kp} {kp} 0 0 0'
    s.input.QuantumESPRESSO.System.ecutwfc = '60.0'
    if mag != 0:
        s.input.QuantumESPRESSO.System.nspin = 'Collinear'
        s.input.QuantumESPRESSO.System.tot_magnetization = f'{mag}'
    s.runscript.preamble_lines = ['export SCM_DISABLE_MPI=1']
    s.runscript.postamble_lines = []

    return s

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