"""
Self-Consistent Field (SCF) driver for plane-wave LDA-DFT.

Uses FFT-based Davidson iterative diagonalization — H is never built explicitly.
Cost per SCF step: O(n_occ × N_G × log N_G) for H applications.

KS Hamiltonian: H = T + V_KS(r)
  T|ψ⟩_G  = (|k+G|²/2) ψ(G)                       [diagonal in G-space]
  V|ψ⟩(G) = FFT[ V_KS(r) · IFFT[ψ](r) ]            [FFT convolution]
"""

import numpy as np
from .cell import Cell
from .xc import xc_energy, xc_potential
from .pseudo import build_ionic_potential, n_valence_electrons
from .diag import davidson


class SCFResult:
    def __init__(self):
        self.converged = False
        self.e_total = 0.0
        self.e_kinetic = 0.0
        self.e_hartree = 0.0
        self.e_xc = 0.0
        self.e_ion_elec = 0.0
        self.e_ion_ion = 0.0
        self.eigenvalues = None
        self.rho_r = None
        self.n_iter = 0

    def __repr__(self):
        lines = [
            f"SCF {'converged' if self.converged else 'NOT converged'} "
            f"in {self.n_iter} iterations",
            f"  E_total    = {self.e_total:+.8f} Ha",
            f"  E_kinetic  = {self.e_kinetic:+.8f} Ha",
            f"  E_hartree  = {self.e_hartree:+.8f} Ha",
            f"  E_xc       = {self.e_xc:+.8f} Ha",
            f"  E_ion-elec = {self.e_ion_elec:+.8f} Ha",
            f"  E_ion-ion  = {self.e_ion_ion:+.8f} Ha",
        ]
        if self.eigenvalues is not None:
            eig_str = "  ".join(f"{e:.6f}" for e in self.eigenvalues)
            lines.append(f"  Eigenvalues (Ha): {eig_str}")
        return "\n".join(lines)


def run_scf(
    cell: Cell,
    mol,
    offset=None,
    n_iter_max: int = 100,
    tol: float = 1e-6,
    mix: float = 0.4,
    diis_start: int = 4,
    diis_size: int = 8,
    diag_tol: float = 1e-7,
    verbose: bool = True,
) -> SCFResult:
    """
    Plane-wave LDA-DFT SCF (closed shell, Γ-point).

    Parameters
    ----------
    cell       : Cell
    mol        : Molecule
    offset     : (3,) shift added to atomic coords (from Cell.from_molecule).
    n_iter_max : Maximum SCF iterations.
    tol        : Density convergence: √(∫|Δρ|² dV) < tol.
    mix        : Linear mixing fraction.
    diis_start : Iteration to begin Pulay DIIS.
    diis_size  : DIIS history depth.
    diag_tol   : Residual tolerance for the Davidson diagonaliser.
    verbose    : Print iteration table.
    """
    n_elec = n_valence_electrons(mol)
    # Restricted calculation: pair electrons, last orbital singly occupied if odd
    n_occ = (n_elec + 1) // 2
    occ   = np.full(n_occ, 2.0)
    if n_elec % 2 == 1:
        occ[-1] = 1.0

    V_ion  = build_ionic_potential(mol, cell, offset)   # (N_r,) fixed

    # Initial density: uniform
    rho = np.full(cell.N_r, n_elec / cell.volume)

    result       = SCFResult()
    result.e_ion_ion = mol.nuclear_repulsion

    diis_rho: list = []
    diis_err: list = []
    guess = None          # warm-start Davidson between SCF iterations

    if verbose:
        print(f"{'Iter':>4}  {'E_total':>16}  {'ΔE':>12}  {'|Δρ|':>12}")
        print("-" * 56)

    e_prev = 0.0

    for iteration in range(1, n_iter_max + 1):
        # ----- KS potential -----
        V_H  = cell.hartree_potential(rho)
        v_xc = xc_potential(rho)
        V_KS = V_ion + V_H + v_xc        # (N_r,)

        # ----- Diagonalise -----
        evals, evecs = davidson(
            cell, V_KS, n_occ,
            n_extra=max(4, n_occ),
            tol=diag_tol,
            guess=guess,
        )
        guess = evecs.copy()              # warm start next iteration

        # ----- New density -----
        orbs_G = evecs                   # (n_occ, N_G)
        rho_new = cell.density_from_orbitals(orbs_G, occ)

        # ----- Energies -----
        # E_kin = Σ_n f_n Σ_G KE_G |ψ̂_n(G)|² / V
        e_kin      = float(np.einsum('n,nG,G->', occ, np.abs(orbs_G)**2, cell.KE_G) / cell.volume)
        e_hartree  = 0.5 * float(np.dot(V_H, rho_new) * cell.dV)
        e_xc       = xc_energy(rho_new, cell.dV)
        e_ion_elec = float(np.dot(V_ion, rho_new) * cell.dV)
        e_total    = e_kin + e_hartree + e_xc + e_ion_elec + result.e_ion_ion

        delta_E   = e_total - e_prev
        e_prev    = e_total

        # ----- Convergence -----
        delta_rho = rho_new - rho
        conv      = float(np.sqrt(np.dot(delta_rho, delta_rho) * cell.dV))

        if verbose:
            print(f"{iteration:>4}  {e_total:>16.8f}  {delta_E:>+12.2e}  {conv:>12.2e}")

        # ----- Density mixing -----
        diis_rho.append(rho_new.copy())
        diis_err.append(delta_rho.copy())
        if len(diis_rho) > diis_size:
            diis_rho.pop(0)
            diis_err.pop(0)

        if iteration >= diis_start and len(diis_err) >= 2:
            rho = _pulay_mix(diis_rho, diis_err)
        else:
            rho = (1 - mix) * rho + mix * rho_new

        # Enforce ρ ≥ 0 and particle conservation
        rho = np.maximum(rho, 0.0)
        rho *= n_elec / (np.sum(rho) * cell.dV)

        if conv < tol and iteration > 2:
            result.converged = True
            break

    result.n_iter      = iteration
    result.e_total     = e_total
    result.e_kinetic   = e_kin
    result.e_hartree   = e_hartree
    result.e_xc        = e_xc
    result.e_ion_elec  = e_ion_elec
    result.eigenvalues = evals
    result.rho_r       = rho_new

    if verbose:
        print()
        print(result)

    return result


def _pulay_mix(rho_hist: list, err_hist: list) -> np.ndarray:
    """Pulay DIIS: optimal linear combination of stored densities."""
    m   = len(err_hist)
    B   = np.zeros((m + 1, m + 1))
    for i in range(m):
        for j in range(m):
            B[i, j] = np.dot(err_hist[i], err_hist[j])
    B[m, :] = -1.0
    B[:, m] = -1.0
    B[m, m] = 0.0
    rhs      = np.zeros(m + 1)
    rhs[m]   = -1.0
    try:
        c = np.linalg.solve(B, rhs)[:m]
    except np.linalg.LinAlgError:
        c = np.zeros(m); c[-1] = 1.0
    return sum(ci * ri for ci, ri in zip(c, rho_hist))