"""Native DF storage contracts; external SCF acceptance remains separate."""

import numpy as np
import pytest

import vibeqc as vq
from vibeqc import _vibeqc_core as core


@pytest.fixture(scope="module")
def basis_pair():
    mol = vq.Molecule([vq.Atom(8, [0., 0., 0.]),
                       vq.Atom(1, [0., 1.43, 1.11]),
                       vq.Atom(1, [0., -1.43, 1.11])])
    orbital = vq.BasisSet(mol, "def2-svp")
    auxiliary = vq.BasisSet(mol, "def2-universal-jkfit")
    return mol, orbital, auxiliary


@pytest.fixture(scope="module")
def storage(basis_pair):
    _, orbital, auxiliary = basis_pair
    return (core._DensityFittingStorage(orbital, auxiliary),
            int(orbital.nbasis), int(auxiliary.nbasis))


@pytest.mark.parametrize("method", ["build_J_streaming", "build_K_streaming"])
@pytest.mark.parametrize("block", ["default", "single", "partial", "full", "oversized"])
def test_retired_streaming_refuses_every_positive_block(storage, method, block):
    df, n, naux = storage
    assert naux != n and naux % 7 != 0
    # A signed asymmetric sentinel makes accidental dense interpretation
    # visible: unsupported streaming must refuse before touching its density.
    density = np.arange(n * n, dtype=float).reshape(n, n) - n * n / 2
    before = density.copy()
    sizes = {"single": 1, "partial": 7, "full": naux, "oversized": naux + 5}
    args = () if block == "default" else (sizes[block],)
    with pytest.raises(RuntimeError, match=rf"{method}: unsupported;.*retired"):
        getattr(df, method)(density, *args)
    np.testing.assert_array_equal(density, before)


@pytest.mark.parametrize("method", ["build_J_streaming", "build_K_streaming"])
@pytest.mark.parametrize("block", [0, -1, -64])
def test_retired_streaming_rejects_invalid_block_without_loop(storage, method, block):
    df, n, _ = storage
    with pytest.raises(ValueError, match="block_size must be positive"):
        getattr(df, method)(np.eye(n), block)


@pytest.mark.parametrize("method", ["build_J_streaming", "build_K_streaming"])
def test_retired_streaming_has_no_density_dependent_fallback(storage, method):
    df, n, _ = storage
    for density in (np.eye(n), np.zeros((0, 0)), np.full((1, 2), np.nan)):
        with pytest.raises(RuntimeError, match="resident DF storage is acceptable"):
            getattr(df, method)(density)


@pytest.fixture(scope="module")
def factors(storage):
    df, n, naux = storage
    # Identity MO transforms expose a temporary copy, never a resident flat B.
    return np.asarray(df.mo_transform(np.eye(n), np.eye(n))).reshape(naux, n, n)


def _roundoff_bound(terms, scale):
    """Small-fixture summation bound, not a GPU/scientific acceptance budget."""
    eps = np.finfo(float).eps
    return 16 * terms * eps / (1 - terms * eps) * max(float(scale), 1.)


def test_native_factors_match_coulomb_half_transform(basis_pair, factors):
    _, orbital, auxiliary = basis_pair
    raw = np.asarray(vq.compute_3c_eri(orbital, auxiliary))
    metric = np.asarray(vq.compute_2c_eri(auxiliary))
    lower = np.linalg.cholesky(metric)
    assert np.linalg.norm(np.tril(lower, -1)) > 1.
    expected = np.linalg.solve(lower, raw.reshape(len(raw), -1)).reshape(raw.shape)
    bound = _roundoff_bound(len(raw), np.linalg.cond(lower) * np.max(abs(expected)))
    assert np.max(abs(factors - expected)) <= bound


def test_B_metric_is_frobenius_gram_and_psd(storage, factors, basis_pair):
    df, n, naux = storage
    expected = np.einsum("pmn,qmn->pq", factors, factors)
    actual = np.asarray(df.compute_B_metric())
    bound = _roundoff_bound(n*n, np.max(np.diag(expected)))
    assert np.max(abs(actual - expected)) <= bound
    np.testing.assert_array_equal(actual, actual.T)
    assert np.linalg.eigvalsh(actual).min() >= -_roundoff_bound(naux, np.linalg.norm(actual, 2))

    # Ensure this physical fixture distinguishes the retired transpose defect.
    _, orbital, auxiliary = basis_pair
    raw = np.asarray(vq.compute_3c_eri(orbital, auxiliary)).reshape(naux, -1)
    lower = np.linalg.cholesky(vq.compute_2c_eri(auxiliary))
    left = np.linalg.solve(lower, raw @ raw.T)
    defective = np.linalg.solve(lower.T, left.T).T
    assert np.max(abs(defective - expected)) > 100 * bound


@pytest.mark.parametrize("layout", ["c", "fortran", "strided", "transpose"])
def test_contract_B_uses_logical_pairs(storage, factors, layout):
    df, n, _ = storage
    sentinel = (np.arange(4*n*n).reshape(2*n, 2*n)[::2, ::2] - n*n) / (n*n)
    matrix = {"c": np.array(sentinel, order="C"),
              "fortran": np.array(sentinel, order="F"),
              "strided": sentinel, "transpose": sentinel.T}[layout]
    assert not np.array_equal(matrix, matrix.T)
    expected = np.einsum("pmn,mn->p", factors, matrix)
    actual = df.contract_B_with(matrix)
    bound = _roundoff_bound(n*n, np.max(abs(factors)) * np.max(abs(matrix)))
    assert np.max(abs(actual - expected)) <= bound


@pytest.mark.parametrize("shape", [(0, 0), (1, 2), (24, 23), (25, 24)])
def test_contract_B_rejects_wrong_shape(storage, shape):
    df, _, _ = storage
    with pytest.raises(ValueError, match="shape must match the AO basis"):
        df.contract_B_with(np.zeros(shape))


@pytest.mark.parametrize("kind", ["signed", "psd", "zero"])
def test_active_JK_matches_factor_oracle(storage, factors, kind):
    df, n, _ = storage
    rng = np.random.default_rng(382)
    a = rng.normal(size=(n, n))
    density = {"signed": (a + a.T) * .01,
               "psd": 2 * a[:, :5] @ a[:, :5].T * .01,
               "zero": np.zeros((n, n))}[kind]
    alpha = np.einsum("pmn,mn->p", factors, density)
    j = np.einsum("p,pmn->mn", alpha, factors)
    k = np.einsum("pml,ls,psn->mn", factors, density, factors)
    # Restricted spin-summed D: E_J = Tr(D J)/2, E_K = -Tr(D K)/4.
    for actual, expected, weight in [(df.build_J(density), j, .5),
                                     (df.build_K_density(density), k, -.25)]:
        bound = _roundoff_bound(factors.size, np.max(abs(factors))**2 * np.max(abs(density)))
        assert np.max(abs(actual - expected)) <= bound
        error = weight * np.sum(density * (actual - expected))
        assert abs(error) <= abs(weight) * bound * np.sum(abs(density))


def test_live_accounting_has_one_tensor_after_construction(basis_pair):
    mol, orbital, auxiliary = basis_pair
    opts = vq.RHFOptions()
    opts.density_fit = True
    opts.aux_basis = "def2-universal-jkfit"
    opts.stability_check = False
    opts.component_timing = True
    result = vq.run_rhf(mol, orbital, opts)
    assert result.converged
    memory = {m.label: m.bytes for m in result.component_timing.memory}
    n, naux = int(orbital.nbasis), int(auxiliary.nbasis)
    tensor = 8 * naux * n*n
    assert memory["df.tensor_extent_bytes"] == tensor
    assert memory["df.live_tracked_bytes"] == tensor + 8*naux*naux
    metric = 8*naux*naux
    # Samples follow actual container lifetimes, rather than treating every
    # allocation in the constructor as simultaneously resident. Backend
    # scratch/allocator retention are deliberately outside this contract.
    assert memory["df.metric_factor_tracked_bytes"] == 3*metric
    for boundary in ("pack", "transform", "unpack"):
        assert memory[f"df.{boundary}_tracked_bytes"] == 2*tensor + metric
    assert memory["df.construction_tracked_peak_bytes"] == max(
        3*metric, 2*tensor + metric)