"""
Goedecker-Teter-Hutter (GTH) local pseudopotentials.

The local pseudopotential in real space is:
  V_loc(r) = -Z_ion/r · erf(r/√2 r_loc)
             + exp(-r²/2r_loc²) · (C1 + C2(r/r_loc)² + C3(r/r_loc)⁴ + C4(r/r_loc)⁶)

Its 3D Fourier transform (infinite-space, Ha·Bohr³) for G ≠ 0:
  V̂_loc(G) = -4π Z_ion/G² · exp(-G²r_loc²/2)
              + (2π)^(3/2) r_loc³ · exp(-G²r_loc²/2)
              · [C1 + C2(3 - x) + C3(15 - 10x + x²) + C4(105 - 105x + 21x² - x³)]
  where x = G² r_loc².

The total ionic potential in the cell is:
  V̂_ion(G) = Σ_I  exp(-i G·R_I) · V̂_loc(G)          [Ha·Bohr³]
  V_ion(r)  = (1/V) Σ_G  V̂_ion(G) · exp(i G·r)        [Ha]

G = 0 is set to zero; the divergence cancels with the Hartree G=0 term
for a charge-neutral cell.

References:
  Goedecker, Teter, Hutter, Phys. Rev. B 54, 1703 (1996).
  Krack, Theor. Chem. Acc. 114, 145 (2005).
"""

import numpy as np

# GTH-LDA local pseudopotential parameters
# (Z_ion, r_loc, C1, C2, C3, C4)
_GTH_LDA = {
    # Hard GTH-LDA: r_loc=0.2 Bohr → needs ecut ≥ 100 Ha (production use)
    'H':  (1,  0.200000, -4.063326,  0.677832,  0.000000,  0.000000),
    # Soft GTH-LDA: r_loc=0.5 Bohr → converges at ecut ≥ 10 Ha (demo/testing)
    # C1=0 chosen so ε_1s ≈ −0.24 Ha (≈ LDA value of −0.228 Ha)
    'H_soft': (1, 0.500000,  0.000000,  0.000000,  0.000000,  0.000000),
    'He': (2,  0.200000, -9.111202,  1.698768,  0.000000,  0.000000),
    'Li': (3,  0.400000, -8.236204,  8.360586, -3.028050,  0.000000),
    'C':  (4,  0.338471, -8.513760,  1.228451,  0.000000,  0.000000),
    'N':  (5,  0.289157,-12.234831,  1.767960,  0.000000,  0.000000),
    'O':  (6,  0.247631,-16.582292,  2.395364,  0.000000,  0.000000),
    'F':  (7,  0.218525,-21.307361,  3.072807,  0.000000,  0.000000),
}


def valence_electrons(symbol: str) -> int:
    return _GTH_LDA[symbol][0]


def use_soft_pseudopotentials(mol) -> "Molecule":
    """
    Return a copy of mol where all atoms use soft pseudopotentials
    (e.g., 'H' → 'H_soft') where available.  Useful for low-ecut tests.
    """
    from copy import deepcopy
    from .molecule import Atom
    m = deepcopy(mol)
    for atom in m.atoms:
        soft_key = atom.symbol + '_soft'
        if soft_key in _GTH_LDA:
            atom.symbol = soft_key
    return m


def n_valence_electrons(mol) -> int:
    return sum(valence_electrons(a.symbol) for a in mol.atoms)


def _pseudo_G_values(symbol: str, G2: np.ndarray) -> np.ndarray:
    """
    GTH local pseudopotential V̂_loc(G) in Ha·Bohr³ (per atom, no structure factor).

    Parameters
    ----------
    symbol : str      Element.
    G2     : (N,) array  |G|² in Bohr⁻².

    Returns
    -------
    V_G : (N,) complex  Values in Ha·Bohr³.
    """
    Z_ion, r_loc, C1, C2, C3, C4 = _GTH_LDA[symbol]
    rl2 = r_loc * r_loc
    V_G = np.zeros(len(G2), dtype=complex)

    nz = G2 > 1e-14
    g2 = G2[nz]
    x = rl2 * g2                              # G² r_loc²
    ef = np.exp(-x / 2)

    erf_term   = -4 * np.pi * Z_ion / g2 * ef
    poly       = C1 + C2*(3 - x) + C3*(15 - 10*x + x**2) + C4*(105 - 105*x + 21*x**2 - x**3)
    gauss_term = (2 * np.pi)**1.5 * r_loc**3 * ef * poly

    V_G[nz] = erf_term + gauss_term
    # G = 0: set to zero; divergence cancelled by Hartree neutrality condition
    V_G[~nz] = 0.0
    return V_G


def build_ionic_potential(mol, cell, offset=None) -> np.ndarray:
    """
    Build the total GTH ionic potential on the real-space grid.

    Uses the full FFT mesh (density cutoff) for accuracy.

    Parameters
    ----------
    mol    : Molecule
    cell   : Cell
    offset : (3,) array, optional
        Shift added to atomic coordinates so the molecule sits in the cell.

    Returns
    -------
    V_ion_r : (N_r,) float array  Ionic potential in Hartree.
    """
    G2_flat = cell._G2_full.ravel()       # |G|² on full mesh, (N_r,)
    G_flat  = cell.G_full                 # G-vectors on full mesh, (N_r, 3)

    V_ion_G = np.zeros(cell.N_r, dtype=complex)

    for atom in mol.atoms:
        pos = atom.coords.copy()
        if offset is not None:
            pos = pos + offset

        S_G      = np.exp(-1j * (G_flat @ pos))        # structure factor (N_r,)
        V_G_atom = _pseudo_G_values(atom.symbol, G2_flat)
        V_ion_G += S_G * V_G_atom

    # G=0: enforce zero (charge neutrality cancellation)
    V_ion_G[G2_flat < 1e-14] = 0.0

    # f̂(G) = V̂_ion → V_ion(r) = (1/V) Σ_G f̂(G) exp(iGr) = ifftn(f̂/dV)
    return cell.potential_G_to_r(V_ion_G)