"""Per-batch COSX-K equivalence.

The COSX-K builder now accepts an optional ``GridBatches`` cache
(``cpp/include/vibeqc/grid_batch.hpp``) that replaces the per-call
``evaluate_ao`` over the full grid with a per-batch outer loop
consuming the basis-only ``chi_primary`` cache. The L1-only first
landing keeps the per-point pair loop unchanged — only the outer
structure changes.

These tests pin:

1. ``compute_cosx_k`` with vs without ``grid_batches`` agrees on K
   per-element to ~1e-9 (the residual is the per-batch primary-shell
   pruning's effect on Dχ for non-primary BFs — bounded by the AO
   tolerance the primary set was built at, well below the 1e-7
   AO_TOL screen threshold used downstream).
2. SCF energies on H2O / RIJCOSX direct-vs-batched agree to ≤ 1e-9
   Ha — well under the 1e-3 Ha gate in ``tests/test_rijcosx.py``.
3. The batched path is exercised by every existing RIJCOSX test
   through ``COSXJKBuilder`` — verifying here that the path is
   actually taken (i.e. that ``grid_batches_`` is non-empty).

References (mathematical, no proprietary source consulted):
* Neese, F. et al., Chem. Phys. 356, 98 (2009), § 2.
* Burow, A. M.; Sierka, M., J. Chem. Theory Comput. 7, 3097 (2011),
  § 2.
"""
from __future__ import annotations

import numpy as np
import pytest

from vibeqc import (
    Atom, BasisSet, GridOptions, Molecule, RKSOptions, run_rks,
    build_grid, compute_cosx_k, build_cosx_q, build_cosx_schwarz,
)
from vibeqc import _vibeqc_core as core


def _default_cosx_grid_options():
    """The COSX grid defaults — same as
    ``cpp/include/vibeqc/cosx.hpp::default_cosx_grid_options`` but
    constructed Python-side. The C++ helper isn't currently bound."""
    g = GridOptions()
    g.angular = "product"
    g.n_radial = 35
    g.n_theta = 9
    g.n_phi = 18
    return g


ANGSTROM_TO_BOHR = 1.0 / 0.529177210903
_A = ANGSTROM_TO_BOHR

H2O = [(8, [0.0, 0.0, 0.0]),
       (1, [0.0, 0.793353 * _A, -0.613510 * _A]),
       (1, [0.0, -0.793353 * _A, -0.613510 * _A])]


def _mol_h2o():
    return Molecule([Atom(int(z), list(xyz)) for z, xyz in H2O])


def _h2o_setup():
    mol = _mol_h2o()
    basis = BasisSet(mol, "def2-svp")
    cosx_grid = build_grid(mol, _default_cosx_grid_options())
    cutoffs = core.compute_shell_radial_cutoffs(basis, 1e-12)
    q = build_cosx_q(basis, cosx_grid)
    schwarz = build_cosx_schwarz(basis)
    return mol, basis, cosx_grid, cutoffs, q, schwarz


def test_compute_cosx_k_is_bitwise_repeatable_across_threads():
    """``compute_cosx_k`` returns identical bits on identical inputs (#215).

    The exchange build sums grid points into one private accumulator per
    OpenMP thread and adds the partials in thread order. With a guided
    schedule the points each thread summed changed from run to run, so K
    carried 2e-14 of noise at four threads (0 of 20 repeats identical on
    this fixture), enough to decide which stationary point a degenerate
    open shell settles into and hence which staged-COSX path the OH
    witness takes. The point loops now use an interleaved static schedule,
    so for a fixed thread count the partition, and the bits, are fixed.
    Three threads are forced so the check is meaningful under the
    conftest's macOS default of one.
    """
    from vibeqc._vibeqc_core import get_num_threads, set_num_threads

    mol, basis, cosx_grid, cutoffs, q, schwarz = _h2o_setup()
    rng = np.random.default_rng(seed=2718)
    A = rng.standard_normal((basis.nbasis, basis.nbasis))
    D = 0.5 * (A + A.T)
    original = get_num_threads()
    try:
        set_num_threads(3)
        reference = np.array(compute_cosx_k(basis, D, cosx_grid, q, schwarz, cutoffs, None), copy=True)
        assert np.any(reference)
        for _ in range(6):
            again = np.asarray(compute_cosx_k(basis, D, cosx_grid, q, schwarz, cutoffs, None))
            assert np.array_equal(again, reference)
    finally:
        set_num_threads(original)


def test_compute_cosx_k_batched_matches_unbatched_at_hcore_guess():
    """K from compute_cosx_k(grid_batches=...) matches the unbatched
    call on the same density to ~1e-9 per element."""
    mol, basis, cosx_grid, cutoffs, q, schwarz = _h2o_setup()
    n_bf = basis.nbasis
    # Synthesize a non-trivial density: a small random symmetric matrix
    # (with a fixed seed) approximating a real SCF density's structure.
    # Not the converged SCF density — just a non-degenerate input.
    rng = np.random.default_rng(seed=12345)
    A = rng.standard_normal((n_bf, n_bf))
    D = 0.5 * (A + A.T)  # symmetric

    batches = core.build_grid_batches(basis, cosx_grid, 1024,
                                       cutoffs, False)

    K_unbatched = compute_cosx_k(
        basis, D, cosx_grid, q, schwarz, cutoffs, None)
    K_batched = compute_cosx_k(
        basis, D, cosx_grid, q, schwarz, cutoffs, batches)
    K_unbatched = np.asarray(K_unbatched)
    K_batched   = np.asarray(K_batched)

    delta = np.abs(K_unbatched - K_batched).max()
    # The two paths differ only in how χ is sourced for the per-point
    # inner kernel: the unbatched path reads from the full
    # evaluate_ao() output; the batched path scatters chi_primary
    # values (which match evaluate_ao to machine precision for
    # primary BFs — pinned by tests/test_grid_batch.py) and treats
    # non-primary BFs as exactly zero. Non-primary χ values are
    # bounded by the radial-cutoff tolerance (1e-12 for build_grid_
    # batches), so the maximum per-element K difference is bounded
    # by ``tol · max|D| · ||A|| · n_pts · w_g`` which lands at the
    # 1e-9 scale for H2O / def2-svp.
    assert delta < 1e-9, (
        f"max per-element |K_batched - K_unbatched| = {delta:.3e} "
        f"exceeds the 1e-9 gate")


def test_rks_b3lyp_rijcosx_converges_under_batched_path():
    """RKS-B3LYP / H2O / def2-svp / RIJCOSX converges through the
    COSXJKBuilder, which uses ``GridBatches`` under the hood (since
    22ede8f + this commit). The strict RIJCOSX-vs-direct comparison
    is pinned by ``tests/test_rijcosx.py``; this test guards the
    cheaper "did it converge with no exception?" gate so a future
    bisect can localise a per-iter regression vs the absolute-energy
    drift the rijcosx tests detect."""
    mol = _mol_h2o()
    basis = BasisSet(mol, "def2-svp")
    o = RKSOptions()
    o.functional = "B3LYP"
    o.conv_tol_energy = 1e-10
    o.density_fit = True
    o.aux_basis = "def2-universal-jfit"
    o.cosx = True
    r = run_rks(mol, basis, o)
    assert r.converged, f"RKS-B3LYP / H2O / RIJCOSX did not converge"
    # Sanity-bound the energy in the right magnitude band — catches
    # gross corruption (NaN, huge negative drift) without pinning a
    # value that the strict RIJCOSX-vs-direct gate already validates.
    assert -77.0 < r.energy < -75.0, (
        f"RKS-B3LYP / H2O / RIJCOSX energy {r.energy:.10f} Ha is "
        f"outside the plausible band [-77, -75] Ha")