"""KS-DFT on a Cyclic Cluster Model reference (RKS-CCM / UKS-CCM).

Kohn-Sham CCM with the **four-center** WSSC-weighted Coulomb (and exact exchange
for hybrids) plus a semi-local exchange-correlation potential. The Coulomb /
exact-exchange operators carry the cyclic boundary conditions through the CCM
four-center (``method="union12"`` or ``"aiccm2026dev-a"``); the XC functional is a
local functional of the cluster density, evaluated on an ordinary molecular Becke
grid over the supercell (XC has no four-center structure). Concretely the KS
matrix is

    F = h^CCM + J^CCM + V_xc[r]  (- a K^CCM for a hybrid),

assembled by the production C++ ``run_rks_scf_with_jk`` driver fed the CCM
four-center JK builder and a supercell XC grid -- the DFT analogue of
:func:`vibeqc.periodic.ccm.scf.run_ccm_rhf_scalable`.

Validation: the isolated-cell limit reproduces vibe-qc's molecular ``run_rks`` /
``run_uks`` (the CCM Coulomb is the molecular Coulomb and the XC grid is the
molecular grid there). Reference: Peintinger & Bredow, J. Comput. Chem. 35, 839
(2014), doi:10.1002/jcc.23550.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from .integrals import ccm_overlap
from .padded import ccm_hcore, ccm_nuclear_repulsion
from .scf import _ccm_scalable_cxx_method, _make_ccm_jk_builder

__all__ = ["CCMKSResult", "run_ccm_rks", "run_ccm_uks"]


@dataclass
class CCMKSResult:
    converged: bool
    n_iter: int
    energy: float                  # total KS-CCM energy per reference cell (Ha)
    energy_per_atom: float
    e_xc: float
    e_coulomb: float
    e_hf_exchange: float           # exact-exchange energy (hybrids; 0 for pure)
    functional: str
    mo_energies: np.ndarray
    mo_coeffs: np.ndarray
    density: np.ndarray
    fock: np.ndarray
    overlap: np.ndarray
    open_shell: bool
    # Per-spin open-shell data (None for RKS). ``mo_energies``/``mo_coeffs`` hold the
    # alpha channel and ``density`` the *total* Pa+Pb; these expose the beta channel
    # and the per-spin densities that the open-shell nuclear gradient needs.
    density_alpha: np.ndarray | None = None
    density_beta: np.ndarray | None = None
    mo_energies_beta: np.ndarray | None = None
    mo_coeffs_beta: np.ndarray | None = None
    exchange_q0: str | None = None  # exchange-q=0 convention label (direct route)


def _ccm_jk_builder(ccm, method, four_center="direct", schwarz_threshold=0.0):
    cxx_method = _ccm_scalable_cxx_method(method, four_center)
    return _make_ccm_jk_builder(ccm, cxx_method, schwarz_threshold)


def _overlap_guard(S, lindep_tol):
    w = np.linalg.eigvalsh(S)
    if w[0] < lindep_tol:
        raise ValueError(
            f"CCM overlap matrix near-singular (min eig {w[0]:.2e} < {lindep_tol:.1e}); "
            "enlarge the cluster or screen diffuse functions (Peintinger & Bredow 2014, Fig. 6)."
        )


def _xc_grid(ccm, grid_options):
    from vibeqc import GridOptions, build_grid
    return build_grid(ccm.supercell, grid_options or GridOptions())


def run_ccm_rks(ccm, functional="pbe", *, method="union12", four_center="direct",
                schwarz_threshold=0.0, max_iter=128, conv_tol=1e-8,
                lindep_tol=1e-7, grid_options=None) -> CCMKSResult:
    """Closed-shell RKS-CCM on ``ccm`` (a :class:`CCMSystem`).

    ``functional`` is any libxc name vibe-qc accepts (``"pbe"``, ``"pbe0"``,
    ``"b3lyp"``, ...). ``method`` selects the CCM four-center (``"union12"`` /
    ``"aiccm2026dev-a"``). ``four_center`` selects the J/K build:
    ``"direct"`` (default, integral-direct, O(nbf**2) memory -- scales to real 3-D
    cells) or ``"full"``/``"dense"`` (the dense O(nbf**4) effective-tensor
    comparison reference); see :func:`~vibeqc.periodic.ccm.scf.run_ccm_rhf_scalable`.
    Returns a :class:`CCMKSResult`.

    Validated: isolated-cell limit reproduces vibe-qc's molecular ``run_rks``.
    """
    from vibeqc import RKSOptions, run_rks_scf_with_jk

    S = ccm_overlap(ccm)
    h, _, _ = ccm_hcore(ccm)
    e_nn = ccm_nuclear_repulsion(ccm)
    n_elec = ccm.supercell.n_electrons()
    if n_elec % 2 != 0:
        raise ValueError(
            f"run_ccm_rks is closed-shell but the cluster has {n_elec} electrons; "
            "use run_ccm_uks for open shells.")
    _overlap_guard(S, lindep_tol)

    opts = RKSOptions()
    opts.functional = functional
    opts.max_iter = max_iter
    opts.conv_tol_energy = conv_tol
    res = run_rks_scf_with_jk(
        ccm.basis, n_elec, np.asarray(S, float), np.asarray(h, float),
        float(e_nn), _ccm_jk_builder(ccm, method, four_center, schwarz_threshold),
        _xc_grid(ccm, grid_options), opts, np.empty((0, 0)))

    e_tot = float(res.energy)
    return CCMKSResult(
        converged=bool(res.converged), n_iter=int(res.n_iter), energy=e_tot,
        energy_per_atom=e_tot / ccm.n_atoms, e_xc=float(res.e_xc),
        e_coulomb=float(res.e_coulomb), e_hf_exchange=float(res.e_hf_exchange),
        functional=functional, mo_energies=np.asarray(res.mo_energies),
        mo_coeffs=np.asarray(res.mo_coeffs), density=np.asarray(res.density),
        fock=np.asarray(res.fock), overlap=np.asarray(S), open_shell=False)


def run_ccm_uks(ccm, functional="pbe", *, method="union12", four_center="direct",
                schwarz_threshold=0.0, max_iter=128, conv_tol=1e-8,
                lindep_tol=1e-7, grid_options=None) -> CCMKSResult:
    """Open-shell UKS-CCM on ``ccm`` (a :class:`CCMSystem`).

    Spin counts come from the cluster supercell's charge + multiplicity (as in
    :func:`vibeqc.periodic.ccm.uhf.run_ccm_uhf`). ``four_center`` selects the J/K
    build (``"direct"`` default, integral-direct / O(nbf**2); ``"full"``/``"dense"``
    the O(nbf**4) reference) -- see :func:`run_ccm_rks`. Returns a
    :class:`CCMKSResult` (``density`` is the total density Pa+Pb). Validated:
    isolated-cell limit reproduces vibe-qc's molecular ``run_uks``; reduces to RKS
    for a closed shell.
    """
    from vibeqc import UKSOptions, run_uks_scf_with_jk

    S = ccm_overlap(ccm)
    h, _, _ = ccm_hcore(ccm)
    e_nn = ccm_nuclear_repulsion(ccm)
    n_elec = ccm.supercell.n_electrons()
    mult = int(ccm.supercell.multiplicity)
    n_unpaired = mult - 1
    if (n_elec - n_unpaired) % 2 != 0 or n_unpaired > n_elec:
        raise ValueError(
            f"cluster electron count {n_elec} and multiplicity {mult} are incompatible.")
    n_alpha = (n_elec + n_unpaired) // 2
    n_beta = n_elec - n_alpha
    _overlap_guard(S, lindep_tol)

    opts = UKSOptions()
    opts.functional = functional
    opts.max_iter = max_iter
    opts.conv_tol_energy = conv_tol
    res = run_uks_scf_with_jk(
        ccm.basis, n_alpha, n_beta, np.asarray(S, float), np.asarray(h, float),
        float(e_nn), _ccm_jk_builder(ccm, method, four_center, schwarz_threshold),
        _xc_grid(ccm, grid_options), opts, np.empty((0, 0)), np.empty((0, 0)))

    e_tot = float(res.energy)
    # UKSResult exposes per-spin MOs/densities; report a MOs + total density.
    Da = np.asarray(res.density_alpha)
    Db = np.asarray(res.density_beta)
    return CCMKSResult(
        converged=bool(res.converged), n_iter=int(res.n_iter), energy=e_tot,
        energy_per_atom=e_tot / ccm.n_atoms, e_xc=float(res.e_xc),
        e_coulomb=float(res.e_coulomb), e_hf_exchange=float(res.e_hf_exchange),
        functional=functional, mo_energies=np.asarray(res.mo_energies_alpha),
        mo_coeffs=np.asarray(res.mo_coeffs_alpha), density=Da + Db,
        fock=np.asarray(res.fock_alpha), overlap=np.asarray(S), open_shell=True,
        density_alpha=Da, density_beta=Db,
        mo_energies_beta=np.asarray(res.mo_energies_beta),
        mo_coeffs_beta=np.asarray(res.mo_coeffs_beta))