#!/usr/bin/env plams
# postprocessing script for calculating surface energies
# part of the "Crystals and Surfaces" tutorial in the Amsterdam Modeling Suite
# first run the surface_energy_jobs.py script
# set the variable jobdir to the directory with calculation results
# USAGE: $AMSBIN/plams surface_energy_postprocessing.py

from scipy import stats
import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

config.erase_workdir = True

###################
# set jobdir to the plams working directory where the jobs were run
###################
jobdir = "slab_1010_normal"

jobs = load_all(jobdir)

bulk_e_good = bulk_e_normal = 0
data = []
for path, job in jobs.items():
    if job.name == "lattice_optimization":
        bulk_e_good = job.results.get_energy()
        n_bulk_atoms = len(job.results.get_main_molecule())
    elif job.name == "normalkspace_bulk":
        bulk_e_normal = job.results.get_energy()
    else:
        mol = job.results.get_main_molecule()
        n_slab_atoms = len(mol)
        energy = job.results.get_energy()
        surfacearea = np.linalg.det(toASE(mol).get_cell()[:2, :2])  # square angstrom
        data.append([n_slab_atoms, energy, surfacearea])

n_atoms_index = 0
e_index = 1
area_index = 2

data = np.array(data)
data = data[data[:, n_atoms_index].argsort()]  # sort w.r.t first column, i.e. n_atoms

print("#atoms, energy, surfacearea")
print(data)

use_last_n_datapoints = 3
x_data = data[-use_last_n_datapoints:, n_atoms_index]  # n_atomss
y_data = data[-use_last_n_datapoints:, e_index]  # energies
slope, intercept, r_value, p_value, std_err = stats.linregress(x_data, y_data)
# the intercept = 2*A*ESurf, the area is the same for all jobs
Esurf = intercept / (2 * surfacearea)
conversion = Units.convert(1.0, "Hartree", "J") / Units.convert(1.0, "angstrom", "m") ** 2


# create a plot

max_n_atoms = np.max(data[:, n_atoms_index])
linear_fit_data = np.array([[0, intercept], [max_n_atoms, slope * max_n_atoms + intercept]])
plt.plot(data[:, n_atoms_index], data[:, e_index], "ro", linear_fit_data[:, 0], linear_fit_data[:, 1], "b-")
plt.style.use("seaborn")
plt.ylabel("Energy (Ha)")
plt.xlabel(r"# atoms in ZnO(10$\bar{1}$0) slab")
min_x = 0
max_x = int(max_n_atoms * 1.1)
min_y = int(data[-1, e_index] * 1.1)
max_y = 1
plt.axis([min_x, max_x, min_y, max_y])
plt.annotate(
    "intercept = {:.6f} Ha".format(intercept),
    xy=(0.0, 0.0),
    xytext=(0.1 * (max_x - min_x) + min_x, 0.95 * (max_y - min_y) + min_y),
    arrowprops={"arrowstyle": "-|>"},
)
plt.title(r"Slab energies for ZnO(10$\bar{1}$0) (k-space quality normal)")
plt.savefig("{}.png".format(jobdir))


print("Surface energy from slope (method 1): {}".format(Esurf * conversion))
print("Effective bulk energy: {}".format(slope))
print("Good bulk energy: {}".format(bulk_e_good / n_bulk_atoms))
print("Normal bulk energy: {}".format(bulk_e_normal / n_bulk_atoms))

naive_Esurf = (data[-1, e_index] - (data[-1, n_atoms_index] / n_bulk_atoms) * bulk_e_good) / (2 * surfacearea)
naive_Esurf *= conversion
print("Naive surface energy w.r.t. good bulk (method 2): {}".format(naive_Esurf))
naive_Esurf = (data[-1, e_index] - (data[-1, n_atoms_index] / n_bulk_atoms) * bulk_e_normal) / (2 * surfacearea)
naive_Esurf *= conversion
print("Naive surface energy w.r.t. normal bulk (method 2): {}".format(naive_Esurf))