"""Finite H2 BIPOLE operator experiment (#1003); run with an installed vibeqc.

The printed signed differences are public minus bounded. This is a diagnostic,
not a post-HF provider. All lengths and energies are in atomic units.
"""

from __future__ import annotations

import hashlib
import itertools
import json
import math
from pathlib import Path
from unittest.mock import patch

import numpy as np
from vibeqc import _vibeqc_core as core
from vibeqc import pbc_bipole as rhf
from vibeqc import pbc_bipole_fock as fock
from vibeqc._bipole_physical_source import make_physical_source
from vibeqc.symmetry_shared import Budget

LENGTH = np.array([4.4, 6.0, 6.5])
XYZ = np.array([[0.0, 0.0, 0.0], [1.4, 0.0, 0.0]])
ALPHA, OMEGA = 0.8, 0.45
PAIR, MID, PREC = 3.15, 3.8, 1e-6
TOL, FORBIDDEN = 1e-10, 1e-12
BUDGET = Budget(512 << 20, 10**12)


def geometry(xyz):
    system = core.PeriodicSystem(3, np.diag(LENGTH), [core.Atom(1, p) for p in xyz])
    shells = [
        core.ShellInfo(i, 0, True, [ALPHA], [(2 * ALPHA / np.pi) ** 0.75], list(p))
        for i, p in enumerate(xyz)
    ]
    basis = core.BasisSet(system.unit_cell_molecule(), shells, "H2-single-s", True)
    return system, basis


def rhf_fixture(n):
    system, basis = geometry(XYZ)
    mesh = core.monkhorst_pack(system, [n, 1, 1], [0, 0, 0], False)
    opts = core.PeriodicRHFOptions()
    opts.lattice_opts.cutoff_bohr = PAIR
    opts.lattice_opts.nuclear_cutoff_bohr = 12.0
    opts.lattice_opts.schwarz_threshold = 0.0
    opts.lattice_opts.jk_permutational_symmetry = False
    opts.conv_tol_energy, opts.conv_tol_grad = 1e-11, 1e-9
    opts.max_iter = 80
    return (
        system,
        basis,
        mesh,
        opts,
        {
            "ewald_omega": OMEGA,
            "ewald_precision": PREC,
            "use_ewald_j_split": True,
            "use_exchange_ewald_split": True,
            "use_incremental_fock": False,
            "use_multipole_far_field": False,
            "use_fock_symmetry": False,
            "use_fock_symmetry_reduce": False,
            "sr_image_extent_bohr": MID,
            "exchange_exxdiv": "ewald",
            "progress": False,
        },
    )


def terminal(n):
    system, basis, mesh, opts, kwargs = rhf_fixture(n)
    seen = {}

    def observe(ctx, density, **kw):
        fb = fock.build_bipole_restricted_fock(ctx, density, **kw)
        seen.update(ctx=ctx, cold=not kw["use_incremental"])
        return fb

    # Observe the context of the unmodified terminal build. Evaluate each
    # ordered reciprocal pair explicitly (the mirror is only a speedup).
    from vibeqc.bipole_fock_ewald import KExchangeLongRangeCache

    with (
        patch.object(rhf, "build_bipole_restricted_fock", observe),
        patch.object(KExchangeLongRangeCache, "_mesh_mirror_pairs", return_value=(set(), {})),
    ):
        res = rhf.run_pbc_bipole_rhf(system, basis, mesh, opts, **kwargs)
    assert seen["cold"] and res.converged
    return res, seen["ctx"]


def controls():
    """Set explicit finite ceilings for this two-s-shell experiment."""

    def capped(cls):
        obj = cls()
        for key in dir(obj):
            if key.startswith("maximum_"):
                bound = 65536
                if "bytes" in key:
                    bound = 64 << 20
                elif "work_units" in key:
                    bound = 10**9
                setattr(obj, key, bound)
        return obj

    panel = capped(core._BipoleFinitePanelCaps)
    panel.source = capped(core._BipoleFiniteSourceCaps)
    panel.source.maximum_images = panel.source.maximum_cells = 64
    panel.short_range = capped(core._BipoleErfcBlochCaps)
    panel.short_range.raw = capped(core._BipoleErfcPanelCaps)
    panel.long_range = capped(core._BipoleEwaldGramCaps)
    panel.zero_mode = capped(core._AOPairFourierCellPanelCaps)
    panel.long_range.ao_panel = panel.zero_mode
    stream = capped(core._BipoleProductJKStreamCaps)
    stream.maximum_panel_work_units = 10**8
    stream.maximum_work_units = 10**11
    stream.maximum_node_bytes = 512 << 20
    inv = core._BipoleFinitePanelInventory()
    inv.numerical_replicas, inv.reciprocal_block_size = 1, 32
    inv.backend_margin_bytes_per_replica = 8 << 20
    return panel, stream, inv


def physical(ctx, kmax):
    mesh = core._RegularKMesh([ctx.n_k, 1, 1], [0, 0, 0])
    options = core._BipoleFiniteSourceOptions()
    options.omega, options.reciprocal_energy_cutoff = OMEGA, kmax**2 / 2
    options.zero_mode = core._BipoleFiniteZeroMode.G0Omitted
    options.require_reciprocal_conjugacy = True
    panel, stream, inv = controls()
    owner = make_physical_source(
        ctx.basis,
        ctx.system,
        mesh,
        options,
        panel.source,
        pair_radius=PAIR,
        midpoint_radius=MID,
        maximum_candidates=100000,
        maximum_images=64,
        maximum_cells=64,
        maximum_quartets=16,
        budget=BUDGET,
    )
    return mesh, owner, inv, panel, stream


def pair(a, b, g):
    """Normalized equal-exponent s products: overlap and Gaussian centre."""
    x, y = XYZ[a], XYZ[b] + np.asarray(g) * LENGTH
    return np.exp(-ALPHA * np.dot(x - y, x - y) / 2), (x + y) / 2


def pair_cells(a, b):
    # All y/z images are excluded by PAIR. At |gx| >= 2 the x
    # separation alone exceeds PAIR for every ordered AO pair.
    return [g for g in (-1, 0, 1) if abs(XYZ[a, 0] - XYZ[b, 0] - g * LENGTH[0]) <= PAIR]


def images(a, b, c, d):
    # The product midpoints differ by < 2 cells before applying p.
    for g, h, p in itertools.product(pair_cells(a, b), pair_cells(c, d), range(-2, 3)):
        l, x = pair(a, b, [g, 0, 0])
        r, y = pair(c, d, [h, 0, 0])
        dist = np.linalg.norm(x - y - [p * LENGTH[0], 0, 0])
        if dist <= MID:
            # Izmaylov et al. (2006), Eq. 41 including its R -> 0 limit.
            u = math.sqrt(ALPHA)
            v = 1 / math.sqrt(1 / ALPHA + 1 / OMEGA**2)
            scale = (
                (math.erf(u * dist) - math.erf(v * dist)) / dist
                if dist
                else 2 * (u - v) / math.sqrt(math.pi)
            )
            yield (g, p, p + h), l * r * scale


def reciprocal(q, kmax):
    bound = np.ceil(kmax * LENGTH / (2 * np.pi) + 1).astype(int)
    labels = np.array(list(itertools.product(*(range(-n, n + 1) for n in bound))))
    vectors = (labels + [q, 0, 0]) * 2 * np.pi / LENGTH
    p2 = np.einsum("gi,gi->g", vectors, vectors)
    take = (p2 > 0) & (p2 <= kmax**2)
    p2 = p2[take]
    return (
        labels[take],
        vectors[take],
        4 * np.pi / np.prod(LENGTH) * np.exp(-p2 / (4 * OMEGA**2)) / p2,
    )


def ft(a, b, k, vectors):
    out = np.zeros(len(vectors), complex)
    for g in pair_cells(a, b):
        s, x = pair(a, b, [g, 0, 0])
        out += s * np.exp(-2j * np.pi * k * g - 1j * (vectors @ x))
    return out * np.exp(-np.einsum("gi,gi->g", vectors, vectors) / (8 * ALPHA))


def parts(ao, q, l, r, vectors, weights):
    a, b, c, d = ao
    sr = sum(v * np.exp(2j * np.pi * (l * g + (q + r) * p - r * s)) for (g, p, s), v in images(*ao))
    bl, br = ft(a, b, l, vectors), ft(c, d, r, vectors)
    lr = (bl.conj() * weights) @ br
    zero = 0j
    if q == 0:
        z = np.zeros((1, 3))
        zero = -np.pi / (np.prod(LENGTH) * OMEGA**2) * (ft(a, b, l, z).conj() @ ft(c, d, r, z))
    return np.array([sr, lr, zero]), (bl * weights) @ br


def centered(i, n):
    return (i if 2 * i < n else i - n) / n


def near(actual, expected, tol=TOL):
    np.testing.assert_allclose(actual, expected, atol=tol, rtol=0)


def quartet_table(ctx, kmax, operands):
    _mesh, owner, inv, caps, _stream = operands
    n = ctx.n_k
    # axes: transfer, left k, right k, a,b,c,d, component
    table = np.empty((n, n, n, 2, 2, 2, 2, 3), complex)
    wrong = error = 0.0
    grids = [reciprocal(centered(i, n), kmax) for i in range(n)]
    for ao in itertools.product(range(2), repeat=4):
        a, b, c, d = ao
        domain = owner.build_domain(ao, budget=BUDGET)
        near(
            sorted(tuple(v.ravel()) for v in domain.images),
            sorted((g, 0, 0, p, 0, 0, s, 0, 0) for (g, p, s), _ in images(*ao)),
            0.0,
        )
        for cells, indices in (
            (domain.left_cells, (a, b)),
            (domain.right_cells, (c, d)),
        ):
            near(
                sorted(map(tuple, cells)),
                [(g, 0, 0) for g in pair_cells(*indices)],
                0.0,
            )
        select = core._BipoleEwaldGramSelection()
        select.left_pair_begin, select.right_pair_begin = 2 * a + b, 2 * c + d
        select.left_pair_count = select.right_pair_count = 1
        for q, l, r in itertools.product(range(n), repeat=3):
            select.q_index, select.left_k_index, select.right_k_index = q, l, r
            values, bad = parts(ao, centered(q, n), centered(l, n), centered(r, n), *grids[q][1:])
            table[q, l, r, a, b, c, d] = values
            native = core._make_bipole_finite_product_panel(
                domain.native_source,
                ctx.basis,
                domain.images,
                domain.left_cells,
                domain.right_cells,
                select,
                inv,
                caps,
            ).element(0, 0)
            error = max(error, abs(native - values.sum()))
            wrong = max(wrong, abs(bad - values[1]))
    near(error, 0.0, 5e-12)
    assert wrong > 1e-4
    for q, l, r in itertools.product(range(n), repeat=3):
        tile = table[q, l, r]
        near(tile, table[-q, -l, -r].conj(), 5e-12)
        near(tile, table[q, r, l].transpose(2, 3, 0, 1, 4).conj(), 5e-12)
    # A mixed complex product across AO pairs and k: x* Gram x equals the
    # positive Coulomb norm of the combined Fourier density on the same grid.
    x = np.arange(1, 4 * n + 1) + 1j * np.arange(4 * n, 0, -1) ** 2
    x /= np.linalg.norm(x)
    for qi, (_, ps, ws) in enumerate(grids):
        gram = table[qi, ..., 1].transpose(0, 2, 3, 1, 4, 5).reshape(4 * n, 4 * n)
        functions = np.array(
            [
                ft(a, b, centered(l, n), ps)
                for l, a, b in itertools.product(range(n), range(2), range(2))
            ]
        )
        near(x.conj() @ gram @ x, np.dot(ws, abs(x @ functions) ** 2), 5e-13)
    return table, error, wrong


def contract(table, density):
    n = len(density)
    out = np.zeros((2, n, 2, 2, 3), complex)
    for t, k in itertools.product(range(n), repeat=2):
        out[0, t] += np.einsum("abcdx,cd->abx", table[0, t, k], density[k]) / n
        out[1, t] += np.einsum("acbdx,cd->abx", table[(t - k) % n, k, k], density[k]) / n
    return out


def probes(n):
    """A real basis of all Hermitian 2x2 matrices at each independent k."""
    for k, h in itertools.product(
        range(n),
        ([[1, 0], [0, 0]], [[0, 0], [0, 1]], [[0, 1], [1, 0]], [[0, 1j], [-1j, 0]]),
    ):
        d = np.zeros((n, 2, 2), complex)
        d[k] = h
        yield d / np.linalg.norm(d)


def public(ctx, res, density):
    """Take the SR blocks before the restricted builder reuses the J buffer."""
    points = np.array(ctx.k_points)
    cells = res.density.cells
    rg = np.array([c.r_cart for c in cells])
    blocks = np.einsum("gk,k,kij->gij", np.exp(-1j * rg @ points.T), ctx.weights, density)
    near(blocks.imag, 0.0, 1e-12)
    lat = core.make_lattice_matrix_set(2, cells, list(blocks.real))
    sr = []

    def fold(lat):
        return np.array(fock._bloch_sum_blocks_multi_k(lat.blocks, lat.cells, points))

    build = fock._sr_image_padded_jk

    def observe(*args, **kw):
        jk = build(*args, **kw)
        sr.extend([fold(jk.J), fold(jk.K)])
        return jk

    from vibeqc.bipole_fock_ewald import KExchangeLongRangeCache

    with (
        patch.object(fock, "_sr_image_padded_jk", observe),
        patch.object(KExchangeLongRangeCache, "_mesh_mirror_pairs", return_value=(set(), {})),
    ):
        fb = fock.build_bipole_restricted_fock(
            ctx, lat, coeffs_for_rho=None, _chi_density_k=density, use_incremental=False
        )
    s = np.array(res.overlap)
    c0 = np.pi / (np.prod(LENGTH) * OMEGA**2)
    sds = s @ density @ s
    probe = ctx.xi_madelung * sds
    kz = -c0 / len(points) * sds
    jz = -c0 * ctx.n_elec * s
    klr = np.array(fb.k_corr_per_k) - probe - kz
    jlr = fold(fb.f2e_real) - sr[0] + 0.5 * sr[1] - jz
    return np.stack((np.stack(sr), np.stack([jlr, klr]), np.stack([jz, kz])), axis=-1), probe


def channel_witness(ctx, kmax, table, op):
    """Audit exact and executed vector lists, then isolate one crossed quartet."""
    records = []
    for qi in range(ctx.n_k):
        q = centered(qi, ctx.n_k)
        cache = ctx.x_lr_cache
        if qi:
            # Read the channel actually used for target qi and source Gamma.
            # Looking up a Cartesian difference here would select another path.
            ((key, qc),) = [
                (key, qc)
                for key, (qc, sources) in cache._mesh_groups.items()
                if qi in sources.get(0, ())
            ]
            channel = cache.q_channels[key]
            qp = qc[0] * LENGTH[0] / (2 * np.pi)
        else:
            qp = 0.0
        labels, ps, ws = reciprocal(q, kmax)
        public_vectors = ctx.j_lr_cache.K_vectors if not qi else channel.K_vectors
        actual_labels = np.rint(public_vectors * LENGTH / (2 * np.pi) - [qp, 0, 0]).astype(int)
        near(sorted(map(tuple, actual_labels)), sorted(map(tuple, labels)), 0.0)
        p = np.lexsort(actual_labels.T[::-1])
        near(public_vectors[p], (labels + [qp, 0, 0]) * 2 * np.pi / LENGTH, 1e-14)
        if not qi:
            continue
        from vibeqc.bipole_fock_ewald import _k_cache_key

        # The tiny fixture retains all panels (mirror reuse is disabled).
        b = channel.b_per_k[_k_cache_key(ctx.k_points[0])]
        public_terms = (abs(b[0, 0]) ** 2 * channel.kernel)[p]
        exact_terms = abs(ft(0, 0, 0.0, ps)) ** 2 * ws
        pub = float(public_terms.sum())
        bounded = table[qi, 0, 0, 0, 0, 0, 0, 1].real
        near(exact_terms.sum(), bounded, 5e-14)
        _, rounded_ps, rounded_ws = reciprocal(round(q, 9), kmax)
        rounded_terms = abs(ft(0, 0, 0.0, rounded_ps)) ** 2 * rounded_ws
        # Same public FT and kernel with its existing exact-address diagnostic.
        w_exact, b_exact = cache._channel_tables_on_mesh(op[0], qi, 0)
        near(np.dot(w_exact, abs(b_exact[0, 0]) ** 2), bounded, 5e-14)
        # Measure the LR child itself, not just the assembled panel.
        gm = core._BipoleEwaldGramOptions()
        gm.omega, gm.reciprocal_energy_cutoff = OMEGA, kmax**2 / 2
        gm.require_reciprocal_conjugacy, gm.reciprocal_block_size = True, 32
        gi = core._BipoleEwaldGramInventory()
        gi.numerical_replicas, gi.backend_margin_bytes_per_replica = 1, 8 << 20
        sel = core._BipoleEwaldGramSelection()
        sel.q_index = qi
        sel.left_pair_count = sel.right_pair_count = 1
        dom = op[1].build_domain((0, 0, 0, 0), budget=BUDGET)
        child = core._make_bipole_ewald_product_gram(
            ctx.basis,
            ctx.system,
            op[0],
            dom.left_cells,
            dom.right_cells,
            sel,
            gm,
            gi,
            op[3].long_range,
        )
        measured = child.element(0, 0).real
        near(measured, bounded, 5e-14)
        g = np.flatnonzero(np.all(labels == 0, axis=1)).item()
        records.append(
            {
                "q": q,
                "public_q": qp,
                "vectors": len(labels),
                "ao": [0, 0, 0, 0],
                "public_lr": pub,
                "bounded_lr": measured,
                "signed": pub - measured,
                "G000_public": public_terms[g],
                "G000_exact": exact_terms[g],
                "G000_signed": public_terms[g] - exact_terms[g],
                "rounded_G000_signed": rounded_terms[g] - exact_terms[g],
                "rounded_quartet_signed": rounded_terms.sum() - measured,
            }
        )
    return records


def covariance(ctx, table):
    """Actually reanchor H1 by one unit cell and reevaluate its native panel."""
    import copy

    n = ctx.n_k
    pos = XYZ.copy()
    pos[1, 0] += LENGTH[0]
    other = copy.copy(ctx)
    other.system, other.basis = geometry(pos)
    mesh, owner, inv, caps, _ = physical(other, ctx.resolved_domain.reciprocal_cutoff_bohr_inv)
    a, b, c, d = ao = (0, 1, 1, 0)
    domain = owner.build_domain(ao, budget=BUDGET)
    s = core._BipoleEwaldGramSelection()
    s.left_pair_begin, s.right_pair_begin = 2 * a + b, 2 * c + d
    s.left_pair_count = s.right_pair_count = 1
    forbidden = 0.0
    for q, l, r in itertools.product(range(n), repeat=3):
        s.q_index, s.left_k_index, s.right_k_index = q, l, r
        value = core._make_bipole_finite_product_panel(
            domain.native_source,
            other.basis,
            domain.images,
            domain.left_cells,
            domain.right_cells,
            s,
            inv,
            caps,
        ).element(0, 0)
        phase = np.exp(2j * np.pi * ((l + q) * a - l * b - (r + q) * c + r * d) / n)
        ref = table[q, l, r, a, b, c, d].sum()
        near(value, phase * ref, 5e-12)
        # A common translation of the four AOs gives the missing fourth-leg
        # character. Project the unfiltered amplitude on the finite BvK group.
        for first in range(n):
            mismatch = l + (r + q) - r - first
            projection = ref * np.exp(2j * np.pi * mismatch * np.arange(n) / n).mean()
            allowed = first == mesh.conserved_index(l, r, (r + q) % n)
            if allowed:
                near(projection, ref, 5e-12)
            else:
                forbidden = max(forbidden, abs(projection))
    near(forbidden, 0.0, FORBIDDEN)
    return forbidden


def experiment(n):
    res, ctx = terminal(n)
    kmax = ctx.resolved_domain.reciprocal_cutoff_bohr_inv
    op = physical(ctx, kmax)
    mesh, owner, inv, caps, stream = op
    near(ctx.weights, np.ones(n) / n, 1e-15)
    near(
        np.asarray(ctx.k_points) * LENGTH / (2 * np.pi),
        [mesh.fractional_at(i) for i in range(n)],
        1e-14,
    )
    table, err, bad = quartet_table(ctx, kmax, op)
    forbidden = covariance(ctx, table)
    witnesses = channel_witness(ctx, kmax, table, op)
    overlap = np.array(res.overlap)
    base = np.linalg.inv(overlap)
    density = np.asarray(res.accepted_density_k)
    actions = np.zeros((2, 3))
    private_error = public_count = 0
    directions = []
    worst_probe = {"signed": [0.0, 0.0], "abs": -1.0}
    worst_private = {"signed": [0.0, 0.0], "abs": -1.0}

    def record_worst(record, difference, probe):
        idx = np.unravel_index(np.argmax(abs(difference)), difference.shape)
        z = difference[idx]
        if abs(z) > record["abs"]:
            record.update(
                probe=probe, index=list(map(int, idx)), signed=[z.real, z.imag], abs=abs(z)
            )

    for i, d in enumerate(probes(n)):
        # Private producer accepts the whole space (including imaginary Gamma).
        pred = contract(table, d).sum(axis=-1)
        actual = owner.contract_jk_all(
            d, inventory=inv, panel_caps=caps, stream_caps=stream, budget=BUDGET
        )
        jk = np.array([x.values_copy().reshape(2, 2, 2) for x in actual]).swapaxes(0, 1)
        record_worst(worst_private, jk - pred, i)
        private_error = max(private_error, np.max(abs(pred - jk)))
        # Public real-lattice RHF: project onto time reversal, then remove the
        # fixed-electron trace. Perturb a positive interior density.
        d = (d + d[(-np.arange(n)) % n].conj()) / 2
        d -= base * np.einsum("kab,kba->", overlap, d).real / (2 * n)
        if np.linalg.norm(d) > 1e-12:
            directions.append(d)
    near(private_error, 0.0, 5e-12)
    rank = np.linalg.matrix_rank(
        np.array(directions).view(float).reshape(len(directions), -1), tol=1e-10
    )
    assert rank == (2 if n == 1 else 6)
    for dp in [base, *(base + 0.1 * d for d in directions), density]:
        near(np.einsum("kab,kba->", overlap, dp) / n, 2, 1e-12)
        assert np.linalg.eigvalsh(dp).min() > -1e-12
        actual, probe = public(ctx, res, dp)
        predicted = contract(table, dp)
        record_worst(worst_probe, (actual - predicted).sum(axis=-1), public_count)
        actions = np.maximum(actions, abs(actual - predicted).max(axis=(1, 2, 3)))
        public_count += 1
    # The last call uses the accepted physical density and its actual Hcore.
    j, k = actual.sum(axis=-1)
    terminal_fock = np.asarray(res.hcore) + j - (k + probe) / 2
    near(terminal_fock, res.fock, 1e-12)
    coeff = np.array(res.mo_coeffs)
    ch = coeff.conj().transpose(0, 2, 1)
    near(ch @ overlap @ coeff, np.broadcast_to(np.eye(2), (n, 2, 2)), 1e-12)
    probe_mo = ch @ (-probe / 2) @ coeff
    near(probe_mo, np.broadcast_to(np.diag([-ctx.xi_madelung, 0]), (n, 2, 2)), TOL)
    eps = np.array(res.mo_energies)
    near(np.diagonal(ch @ terminal_fock @ coeff, axis1=1, axis2=2), eps, 1e-10)
    near(actions[0], 0.0, TOL)
    near(actions[1, [0, 2]], 0.0, TOL)
    dk = actual[1, ..., 1] - predicted[1, ..., 1]
    index = np.unravel_index(np.argmax(abs(dk)), dk.shape)
    delta = dk[index]
    bounded_j, bounded_k = predicted.sum(axis=-1)
    bounded_fock = np.asarray(res.hcore) + bounded_j - (bounded_k + probe) / 2
    f_error = np.asarray(res.fock) - bounded_fock
    total_k_error = k - bounded_k  # same separately tested probe on each side
    k_index = np.unravel_index(np.argmax(abs(total_k_error)), total_k_error.shape)
    f_index = np.unravel_index(np.argmax(abs(f_error)), f_error.shape)
    df = f_error[index]
    near(df, -delta / 2, 1e-12)
    assert eps[:, 1].min() > eps[:, 0].max()
    from vibeqc.dlpno.source import (
        CorrelationSourceError,
        refuse_unsupported_periodic_reference,
    )

    try:
        refuse_unsupported_periodic_reference(res)
    except CorrelationSourceError as exc:
        assert "BIPOLE" in str(exc)
    else:
        raise AssertionError("uncertified BIPOLE reference admitted for correlation")
    return {
        "n": n,
        "kmax": kmax,
        "terminal_K_LR_index": list(map(int, index)),
        "terminal_K_LR_signed": [delta.real, delta.imag],
        "terminal_F_signed": [df.real, df.imag],
        "terminal_total_K_index": list(map(int, k_index)),
        "terminal_total_K_signed": [total_k_error[k_index].real, total_k_error[k_index].imag],
        "terminal_F_worst_index": list(map(int, f_index)),
        "terminal_F_worst_signed": [f_error[f_index].real, f_error[f_index].imag],
        "terminal_F_max_abs": float(np.max(abs(f_error))),
        "original_K01_signed": [dk[0, 0, 1].real, dk[0, 0, 1].imag],
        "original_F01_signed": [f_error[0, 0, 1].real, f_error[0, 0, 1].imag],
        "worst_public_action": worst_probe,
        "worst_private_action": worst_private,
        "status": (
            "counterexample"
            if max(np.max(actions), worst_probe["abs"], np.max(abs(f_error))) > TOL
            or any(abs(w["signed"]) > TOL for w in witnesses)
            else "within fixture gates"
        ),
        "vectors": [len(reciprocal(centered(i, n), kmax)[0]) for i in range(n)],
        "energy": res.energy,
        "electrons": 2,
        "eps": eps.tolist(),
        "gap": eps[:, 1].min() - eps[:, 0].max(),
        "xi": ctx.xi_madelung,
        "probe_orbital_shift": probe_mo.diagonal(axis1=1, axis2=2).real.tolist(),
        "private_panel_error": err,
        "wrong_conjugation": bad,
        "private_jk_error": private_error,
        "forbidden": forbidden,
        "public_probe_count": public_count,
        "trace_tangent_rank": int(rank),
        "max_public_component_errors": actions.tolist(),
        "witness": witnesses,
    }


if __name__ == "__main__":
    print(
        json.dumps(
            {
                "native_sha256": hashlib.sha256(Path(core.__file__).read_bytes()).hexdigest(),
                "diagnostic_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
            }
        )
    )
    for nk in (1, 3):
        print(json.dumps(experiment(nk)), flush=True)