"""API surface tests: registry integrity and the pylibxc-compatible shim."""

import numpy as np
import pytest

import funxc
from funxc.libxc_compat import LibXCFunctional


def test_registry_specs_wellformed():
    for fid, spec in funxc.REGISTRY.items():
        assert spec.id == fid
        assert spec.family in ("lda", "gga", "mgga")
        assert spec.kind in ("x", "c", "xc", "k")
        assert spec.libxc_id > 0
        assert spec.kernel_fn is not None
        assert spec.doi


def test_get_spec_is_case_insensitive():
    assert funxc.get_spec("gga_x_pbe") is funxc.get_spec("GGA_X_PBE")
    assert funxc.get_spec("XC_GGA_X_PBE") is funxc.get_spec("GGA_X_PBE")
    with pytest.raises(KeyError):
        funxc.get_spec("GGA_X_DOES_NOT_EXIST")


def test_shim_polarized_shapes():
    f = LibXCFunctional("gga_x_pbe", "polarized")
    n = 4
    rng = np.random.default_rng(0)
    inp = {
        "rho": rng.uniform(0.1, 1.0, (n, 2)),
        "sigma": rng.uniform(0.0, 0.5, (n, 3)),
    }
    out = f.compute(inp)
    assert out["zk"].shape == (n, 1)
    assert out["vrho"].shape == (n, 2)
    assert out["vsigma"].shape == (n, 3)


def test_shim_unpolarized_shapes_and_flat_input():
    f = LibXCFunctional("lda_x", 1)
    rho = np.linspace(0.1, 2.0, 5)
    out = f.compute({"rho": rho})
    assert out["zk"].shape == (5, 1)
    assert out["vrho"].shape == (5, 1)
    out_exc_only = f.compute({"rho": rho}, do_vxc=False)
    assert set(out_exc_only) == {"zk"}


def test_shim_mgga_shapes():
    f = LibXCFunctional("mgga_x_lta", "polarized")
    n = 4
    rng = np.random.default_rng(1)
    inp = {
        "rho": rng.uniform(0.1, 1.0, (n, 2)),
        "sigma": rng.uniform(0.0, 0.5, (n, 3)),
        "lapl": rng.uniform(-1.0, 1.0, (n, 2)),
        "tau": rng.uniform(0.1, 1.0, (n, 2)),
    }
    out = f.compute(inp)
    assert out["zk"].shape == (n, 1)
    assert out["vrho"].shape == (n, 2)
    assert out["vsigma"].shape == (n, 3)
    assert out["vlapl"].shape == (n, 2)
    assert out["vtau"].shape == (n, 2)
    # LTA has no lapl dependence: vlapl must be exactly zero
    assert np.all(out["vlapl"] == 0.0)


def test_shim_metadata():
    f = LibXCFunctional("gga_c_lyp", "polarized")
    assert f.get_number() == 131
    assert f.get_family() == "gga"


def test_composition_blyp():
    """X and C functionals compose additively (BLYP = B88 + LYP)."""
    rho = np.array([[0.4, 0.3]])
    sigma = np.array([[0.02, 0.01, 0.015]])
    b88 = funxc.functional("GGA_X_B88").exc(rho, sigma)
    lyp = funxc.functional("GGA_C_LYP").exc(rho, sigma)
    total = np.asarray(b88) + np.asarray(lyp)
    assert np.all(np.isfinite(total)) and total[0] < 0