import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from scm.plams import Units

file_homo = "C_diamond_3x3x3_q_HOMO.xy"
xy_homo = np.loadtxt(file_homo, usecols=(0, 1), skiprows=1)

to_eV = Units.convert(1, "hartree", "eV")
to_ang = Units.convert(1, "Bohr", "Angstrom")

# HOMO**2
x = xy_homo[:, 0]
y = xy_homo[:, 1]
yy = np.multiply(y, y) * to_eV  # e/bohr^2
yy = yy / to_ang / to_ang  # e/Ang^2


def gaus(x, a, c_x, sigma):
    return a * np.exp(-((x - c_x) ** 2) / 2.0 / sigma**2) / sigma**3 / (2 * np.pi) ** 1.5


a, mean, sigma = 1e-3, 15.0, 2.0
popt, pcov = curve_fit(gaus, x, yy, p0=[a, mean, sigma])
print(popt)

fig, ax = plt.subplots()
ax.plot(x, yy, color="b", lw=2, label=r"DFT defect $\rho$ = $|\varphi (HOMO)|^2$")
ax.plot(x, gaus(x, *popt), color="r", lw=1, ls="--", label=f"Gaussian fit σ={popt[-1]:4.1f} Å")
ax.legend()
ax.set_xlabel("x (Å)")
ax.set_ylabel(r"Charge density (eV/$Å^2$)")
plt.tight_layout()
plt.savefig("charged_defect_C_HOMO_fit.png")
plt.show()