"""Golden-data tests: funxc vs libxc 7.0.0 vendored regression files.

Every registered functional is compared against libxc/testsuite/regression
for all four systems, both spin modes, orders 0 (zk) and 1 (vrho, vsigma).

Tolerance policy: |got - ref| <= ATOL + RTOL*|ref|.
Typical agreement is machine precision (zk ~1e-15 relative). The absolute
floor covers two known noise sources that libxc itself carries:
  * roundoff residue of analytically-cancelling terms at zeta = +-1 (the
    reference files store values like -1.1e-16 where the exact result is 0);
  * expression-grouping noise ~(A*t^2)^2 * eps in the saturated large-t
    limit of the PBE H(t) correlation term, visible only at density-tail
    points (rho ~ 1e-11) where |vsigma| ~ 1e-6 and the noise is ~4e-10.
"""

import numpy as np
import pytest

import funxc

from libxc_reference import (
    SYSTEMS,
    has_regression,
    load_input,
    load_regression,
    unpolarize,
)

ATOL = 2e-9
RTOL = 1e-7

# Points where the vendored regression *file* disagrees with the libxc 7.0.0
# *library* itself (module libxc/7.0.0): the testsuite files ship with the
# source tree and some predate behavior changes. Each entry was verified by
# comparing the installed .so against the file — funxc matches the library
# (e.g. LTA Li vrho(b) point 4: library -2.38273e-8, funxc -2.38273e-8,
# stale file -1.05e-14, a leftover of an older density screen).
STALE_GOLDEN = {
    ("MGGA_X_LTA", "Li", "pol", "vrho(b)", 4),
}


def family_dir(fid: str) -> str:
    return "_".join(fid.lower().split("_")[:2])


FUNCTIONALS = [
    fid for fid in sorted(funxc.REGISTRY) if has_regression(family_dir(fid), fid.lower())
]


def check(got, ref, label, stale=()):
    got = np.asarray(got).reshape(-1)
    err = np.abs(got - ref)
    tol = ATOL + RTOL * np.abs(ref)
    bad = err > tol
    if len(stale):
        bad[np.asarray(stale, dtype=int)] = False
    assert not bad.any(), (
        f"{label}: {bad.sum()}/{bad.size} points exceed tolerance; "
        f"worst at i={err.argmax()}: got={got[err.argmax()]:.12e} "
        f"ref={ref[err.argmax()]:.12e}"
    )


@pytest.mark.parametrize("spin", ["pol", "unpol"])
@pytest.mark.parametrize("system", SYSTEMS)
@pytest.mark.parametrize("fid", FUNCTIONALS)
def test_against_libxc(fid, system, spin):
    fdir = family_dir(fid)
    f = funxc.functional(fid, polarized=(spin == "pol"))

    inp = load_input(system)
    if spin == "unpol":
        inp = unpolarize(inp)
    rho, sigma = inp["rho"], inp["sigma"]

    if f.spec.family == "lda":
        out = f.exc_vxc(rho)
    elif f.spec.family == "gga":
        out = f.exc_vxc(rho, sigma)
    else:
        out = f.exc_vxc(rho, sigma, inp["lapl"], inp["tau"])

    try:
        ref0 = load_regression(fdir, fid.lower(), system, spin, 0)
    except FileNotFoundError:
        # libxc keeps some combinations in regression/<family>/disabled
        # (e.g. gga_c_lyp on the fully polarized H atom); follow suit.
        pytest.skip(f"libxc testsuite has no {fid} {system} {spin} data")
    assert ref0["_func_id"] == f.spec.libxc_id
    check(out["zk"], ref0["zk"], f"{fid} {system} {spin} zk")

    try:
        ref1 = load_regression(fdir, fid.lower(), system, spin, 1)
    except FileNotFoundError:
        # Some combinations disable only the derivative data (e.g.
        # gga_xc_b97_3c on polarized H); zk was still checked above.
        pytest.skip(f"libxc testsuite has no {fid} {system} {spin} order-1 data")
    if spin == "pol":
        columns = {
            "vrho": ("(a)", "(b)"),
            "vsigma": ("(aa)", "(ab)", "(bb)"),
            "vlapl": ("(a)", "(b)"),
            "vtau": ("(a)", "(b)"),
        }
        for key, comps in columns.items():
            if key not in out or f"{key}{comps[0]}" not in ref1:
                continue
            got = np.asarray(out[key])
            for i, comp in enumerate(comps):
                stale = [
                    idx
                    for (sf, ss, sm, sc, idx) in STALE_GOLDEN
                    if (sf, ss, sm, sc) == (fid, system, spin, f"{key}{comp}")
                ]
                check(
                    got[:, i], ref1[f"{key}{comp}"],
                    f"{fid} {system} {key}{comp}", stale=stale,
                )
    else:
        for key in ("vrho", "vsigma", "vlapl", "vtau"):
            if key in out and key in ref1:
                stale = [
                    idx
                    for (sf, ss, sm, sc, idx) in STALE_GOLDEN
                    if (sf, ss, sm, sc) == (fid, system, spin, key)
                ]
                check(out[key], ref1[key], f"{fid} {system} {key}", stale=stale)


def test_all_registered_functionals_have_golden_data():
    missing = [fid for fid in sorted(funxc.REGISTRY) if fid not in FUNCTIONALS]
    assert not missing, f"no regression data for: {missing}"