#!/usr/bin/env amspython

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

################ INPUT ###
# Initialize cell in Angstrom
cell = np.array([[3.518275, 0.0, 0.0], [0.0, 3.518275, 0.0], [0.0, 0.0, 3.518275]])
ecut = 60.0  # in Ha
epsilon = 5.76
nrep_list = [2, 3, 4, 5]

# Gaussian profile
sigma = 1.0  # Angstrom!
tot_charge = -2
# Defect center
c_a1 = 0.5
c_a2 = 0.5
c_a3 = 0.5
##########################

to_eV = Units.conversion_ratio("a.u.", "eV")
to_Ang = Units.conversion_ratio("a.u.", "A")
sigma /= to_Ang


def pl_avg(vol, ni, nj, nk, step):
    x = []
    y = np.zeros((2 * nj + 1))
    for j in range(2 * nj + 1):
        Sum = 0.0
        for i in range(2 * ni + 1):
            for k in range(2 * nk + 1):
                Sum = Sum + vol[i][j][k]
        y[j] = Sum / i / k
        x.append(j * step)
    return np.array(x), y


def get_rho(A2, lx, ly, lz, imax, jmax, kmax, c_a1, c_a2, c_a3, sigma, tot_charge):

    # Define the grid
    a1_list = np.linspace(0, 2 * imax / (2 * imax + 1), 2 * imax + 1)
    a2_list = np.linspace(0, 2 * jmax / (2 * jmax + 1), 2 * jmax + 1)
    a3_list = np.linspace(0, 2 * kmax / (2 * kmax + 1), 2 * kmax + 1)

    # Original center of the Gaussian in cartesian, in the real crystal
    c_crys = np.array([c_a1, c_a2, c_a3])

    # Map the center on the grid
    c_a1_g = np.searchsorted(a1_list, c_a1) / float(2 * imax + 1)
    c_a2_g = np.searchsorted(a2_list, c_a2) / float(2 * jmax + 1)
    c_a3_g = np.searchsorted(a3_list, c_a3) / float(2 * kmax + 1)
    c_crys_n = np.array([c_a1_g, c_a2_g, c_a3_g])
    # In Cartesian
    c_x = np.dot(c_crys_n, A2[:, 0])
    c_y = np.dot(c_crys_n, A2[:, 1])
    c_z = np.dot(c_crys_n, A2[:, 2])

    # Grid in cartesian
    a1_list = a1_list * lx
    a2_list = a2_list * ly
    a3_list = a3_list * lz

    # Checking Gaussian is not spilling
    of_a1_l = True if c_x - 4 * sigma < 0.0 else False
    of_a1_r = True if c_x + 4 * sigma > lx else False
    if of_a1_l and of_a1_r:
        print("Warning: Model charge Sigma too large!")

    of_a2_l = True if c_y - 4 * sigma < 0.0 else False
    of_a2_r = True if c_y + 4 * sigma > ly else False
    if of_a2_l and of_a2_r:
        print("Warning: Model charge Sigma too large!")

    of_a3_l = True if c_z - 4 * sigma < 0.0 else False
    of_a3_r = True if c_z + 4 * sigma > lz else False
    if of_a3_l and of_a3_r:
        print("Warning: Model charge Sigma too large!")

    # Compute the charge density
    norm = (sigma**3) * (2 * np.pi) ** 1.5
    rho_r = []
    for x in a1_list:
        temp1 = []
        for y in a2_list:
            temp2 = []
            for z in a3_list:
                xyz2 = (x - c_x) ** 2 + (y - c_y) ** 2 + (z - c_z) ** 2
                temp2.append(np.exp(-xyz2 / 2.0 / sigma**2) / norm)
            temp1.append(temp2)
        rho_r.append(temp1)
    rho_r = np.array(rho_r) * (tot_charge)

    return rho_r


def FFT(lmax, mmax, nmax, rho_r):

    F_G = np.fft.fftn(rho_r)
    F_G_shift = np.fft.fftshift(F_G)
    rho_G = F_G_shift / (2.0 * lmax + 1) / (2.0 * mmax + 1) / (2.0 * nmax + 1)

    return rho_G


def IFFT(lmax, mmax, nmax, F_G):

    F_G_shift = np.fft.ifftshift(F_G)
    F_r_tr = np.fft.ifftn(F_G_shift) * (2 * lmax + 1) * (2 * mmax + 1) * (2 * nmax + 1)

    return F_r_tr


def poisson_3D(A2, B, lx, ly, lz, imax, jmax, kmax, c_a1, c_a2, c_a3, sigma, tot_charge):

    rho_r = get_rho(A2, lx, ly, lz, imax, jmax, kmax, c_a1, c_a2, c_a3, sigma, tot_charge)
    rho_G = FFT(imax, jmax, kmax, rho_r)

    # Get potential in reciprocal space
    V_G = np.ndarray(shape=(2 * imax + 1, 2 * jmax + 1, 2 * kmax + 1), dtype=complex)
    for i in range(2 * imax + 1):
        for j in range(2 * jmax + 1):
            for k in range(2 * kmax + 1):
                if i == imax and j == jmax and k == kmax:
                    V_G[i][j][k] = 0.0
                else:
                    I, J, K = i - imax, j - jmax, k - kmax
                    G1 = I * B[0][0] + J * B[1][0] + K * B[2][0]
                    G2 = I * B[0][1] + J * B[1][1] + K * B[2][1]
                    G3 = I * B[0][2] + J * B[1][2] + K * B[2][2]
                    modG = G1**2 + G2**2 + G3**2
                    V_G[i][j][k] = 4 * np.pi * rho_G[i][j][k] / (modG * epsilon)

    # Get potential in real space
    V_r = IFFT(imax, jmax, kmax, V_G)
    return V_r, rho_r


def get_energy(A2, V_r, rho_r, imax, jmax, kmax):

    Vol = np.dot(A2[0], np.cross(A2[1], A2[2]))
    dV = (1.0 / float(2 * imax + 1) / float(2 * jmax + 1) / float(2 * kmax + 1)) * Vol
    ss = 0.0
    for l in range(2 * imax + 1):
        for m in range(2 * jmax + 1):
            for n in range(2 * kmax + 1):
                ss += 0.5 * V_r[l][m][n] * rho_r[l][m][n] * dV
    e_model = np.real(ss) * to_eV
    print(f">>> Volume {Vol*to_Ang*to_Ang*to_Ang:6.3f} Å^3 Model energy: {e_model:6.3f} eV")

    return e_model, Vol * to_Ang * to_Ang * to_Ang


def get_fit(e_list, ilatt_list):

    L2, L3, L4 = ilatt_list[0], ilatt_list[1], ilatt_list[2]
    A = np.array([[L2, L2**3, 1.0], [L3, L3**3, 1.0], [L4, L4**3, 1.0]])
    A_inv = np.linalg.inv(A)
    X_u = np.dot(A_inv, e_list)

    return X_u


def func(a, f1, f2, f3):

    return f1 + f2 / (a) + f3 / (a**3)


def fit_model(e_list, a_list, sigma):

    # Fit
    popt, _ = optimize.curve_fit(func, a_list, e_list, p0=(1, 1, 1))

    # Correction
    Em_iso = round(popt[0], 3)

    # Plot
    fig, ax = plt.subplots()
    ax.scatter(1 / a_list, e_list, color="k", marker="o", s=250, label="Uncorrected")

    xmax = max(1 / a_list) * 1.1
    xfit = np.linspace(0, xmax, 100)
    ax.plot(xfit, func(1 / xfit, *popt), color="r", lw=4.0, label="Fit")

    ax.set_ylabel(r"$E_m$ (eV)", fontsize=32)
    ax.set_xlabel(r"$1/a$ ($\AA^{-1}$)", fontsize=32)

    ax.set_xlim(0, xmax)
    #    ax.set_title(f"Sigma = {sigma:6.3f}, Em_iso = {Em_iso:6.3f} eV",fontsize=20)
    ax.set_title(f"Em(iso) = {Em_iso:6.3f} eV", fontsize=20)

    plt.tight_layout()
    plt.savefig("E_corr.png")
    plt.show()

    return Em_iso


def main():

    a_list = []
    e_list = []
    for nrep in nrep_list:

        # Here putting things in a.u. (??)
        A2 = cell * nrep / to_Ang

        # Lattice vectors in real space
        a_1 = A2[0]
        a_2 = A2[1]
        a_3 = A2[2]
        lx = np.linalg.norm(a_1)
        ly = np.linalg.norm(a_2)
        lz = np.linalg.norm(a_3)

        # Reciprocal lattice vectors
        B = 2 * np.pi * np.linalg.inv(np.transpose(A2))
        b_1 = B[0]
        b_2 = B[1]
        b_3 = B[2]

        # Grid in reciprocal space
        Gmax = np.sqrt(2 * ecut)
        imax = int(Gmax / np.sqrt(np.dot(b_1, b_1))) + 1
        jmax = int(Gmax / np.sqrt(np.dot(b_2, b_2))) + 1
        kmax = int(Gmax / np.sqrt(np.dot(b_3, b_3))) + 1

        # Compute charge density and potential
        V_r, rho_r = poisson_3D(A2, B, lx, ly, lz, imax, jmax, kmax, c_a1, c_a2, c_a3, sigma, tot_charge)

        # Planar avg
        x, y = pl_avg(V_r, imax, jmax, kmax, lx / (2 * imax + 1))
        np.savetxt(f"V_{nrep}_r.txt", np.c_[x * to_Ang, y * to_eV])
        #        plt.plot(x*to_Ang,y*to_eV)
        #        plt.show()

        # Compute energy
        e_model, omega = get_energy(A2, V_r, rho_r, imax, jmax, kmax)

        # This is E_periodic
        e_list.append(e_model)
        a_list.append(lx)

    # Fit and get Eisolated
    Eiso = fit_model(np.array(e_list), np.array(a_list), sigma)

    for i in range(len(nrep_list)):

        nrep = nrep_list[i]
        Eper = e_list[i]
        Ecorr = Eiso - Eper
        print(f">>> {nrep}x{nrep}x{nrep} E_latt = E_iso - E_per = {Eiso:6.3f} - {Eper:6.3f} = {Ecorr:6.3f} eV")


if __name__ == "__main__":
    init()
    main()
    finish()