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

# file_p = 'neutral_perfect_3x3x3_HOMO.xy'
file_0_coul = "C_diamond_3x3x3_0_coul.xy"
file_q_coul = "C_diamond_3x3x3_q_coul.xy"
file_m_coul = "V_3_r.txt"
file_q = "C_diamond_3x3x3_q_HOMO.xy"
xy_0c = np.loadtxt(file_0_coul, usecols=(0, 1), skiprows=1)
xy_qc = np.loadtxt(file_q_coul, usecols=(0, 1), skiprows=1)
xy_mc = np.loadtxt(file_m_coul, usecols=(0, 1))
xy_q = np.loadtxt(file_q, usecols=(0, 1), skiprows=1)

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

# Coul
x_c = xy_0c[:, 0]
y_0c = xy_0c[:, 1]
y_qc = xy_qc[:, 1]
y_q0_c = (y_qc - y_0c) * to_eV  # eV

x_m = xy_mc[:, 0]
y_mc = xy_mc[:, 1] * (-1)

# HOMO**2
x = xy_q[:, 0]
y_q = xy_q[:, 1]
y_pq = np.multiply(y_q, y_q) * to_eV  # e/bohr^2
y_pq = y_pq / 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, y_pq, p0=[a, mean, sigma])
print(popt)

fig, ax = plt.subplots(2, sharex=True)
ax[0].plot(x, y_pq, color="b", lw=2, label=r"DFT defect $\rho$ = $|\Psi (HOMO)|^2$")
ax[0].plot(x, gaus(x, *popt), color="r", lw=1, ls="--", label=f"Gaussian fit σ={popt[-1]:4.1f} Å")
ax[0].legend()

y_shift = min(y_q0_c) - min(y_mc)
x_shift = x_c[np.argmax(y_q0_c)] - x_m[np.argmax(y_mc)]
ax[1].plot(x_c, y_q0_c, color="g", lw=2, label="DFT defect V(r) [q-0]")
ax[1].plot(x_m + x_shift, y_mc + y_shift, color="r", lw=1, ls="--", label="Model V(r)")
# ax[1].plot(xx_c,(yy_q0_c-yy_mc+y_shift),color='r',lw=1,ls='--',label='Model V_r')
ax[1].legend()

ax[1].set_xlabel("x (Å)")
ax[0].set_ylabel("Charge density (eV/Å^2)")
ax[1].set_ylabel("Potential (eV)")
plt.tight_layout()
plt.savefig("charged_defect_C_potential.png")
plt.show()