"""Loaders for libxc's vendored testsuite data (golden reference).

Input files (libxc/testsuite/input/<system>): first line npoints, then
9 columns: rhoa rhob sigmaaa sigmaab sigmabb lapla laplb taua taub.

Regression files (libxc/testsuite/regression/<family>/<name>.<system>.<spin>.<order>.bz2):
first line "func_id npoints order", then a column-header line, then values.
"""

import bz2
from pathlib import Path

import numpy as np

LIBXC_ROOT = Path(__file__).resolve().parent.parent / "libxc"
INPUT_DIR = LIBXC_ROOT / "testsuite" / "input"
REGRESSION_DIR = LIBXC_ROOT / "testsuite" / "regression"

SYSTEMS = ("BrOH", "BrOH+", "H", "Li")


def load_input(system: str) -> dict[str, np.ndarray]:
    """Return the 9-column grid data for one test system."""
    raw = np.loadtxt(INPUT_DIR / system, skiprows=1)
    return {
        "rho": raw[:, 0:2],  # (N, 2): rhoa, rhob
        "sigma": raw[:, 2:5],  # (N, 3): sigmaaa, sigmaab, sigmabb
        "lapl": raw[:, 5:7],
        "tau": raw[:, 7:9],
    }


def unpolarize(inp: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    """Collapse spin channels the way xc-regression.c does for nspin=1."""
    return {
        "rho": inp["rho"].sum(axis=1),
        "sigma": inp["sigma"][:, 0] + 2.0 * inp["sigma"][:, 1] + inp["sigma"][:, 2],
        "lapl": inp["lapl"].sum(axis=1),
        "tau": inp["tau"].sum(axis=1),
    }


def load_regression(
    family_dir: str, functional: str, system: str, spin: str, order: int
) -> dict[str, np.ndarray]:
    """Return reference outputs keyed by column name (zk, vrho(a), ...)."""
    path = (
        REGRESSION_DIR
        / family_dir
        / f"{functional}.{system}.{spin}.{order}.bz2"
    )
    with bz2.open(path, "rt") as fh:
        header = fh.readline().split()
        func_id, npoints, forder = int(header[0]), int(header[1]), int(header[2])
        assert forder == order
        names = fh.readline().split()
        data = np.loadtxt(fh)
    data = np.asarray(data).reshape(npoints, len(names))
    out = {name: data[:, i] for i, name in enumerate(names)}
    out["_func_id"] = func_id
    return out


def has_regression(family_dir: str, functional: str) -> bool:
    return any(
        (REGRESSION_DIR / family_dir).glob(f"{functional}.*.bz2")
    )