#!/usr/bin/env python3
"""Exact finite checks for the exposition of arXiv:2610.06783v1.

Run checks (standard library only):
    python3 "content/codes/2026 Fall/sparse_lopsided_checks.py"

Also regenerate both SVG figures (requires matplotlib):
    python3 "content/codes/2026 Fall/sparse_lopsided_checks.py" --figures

The finite checks are not a verification of the asymptotic algorithm.
All algebra and enumeration use integers. Floating point is used only for
displaying exponents and plotting logarithms of exact combinatorial counts.
"""

from __future__ import annotations

import argparse
from collections import Counter
from fractions import Fraction
from itertools import combinations, product
import json
from math import comb, log, log10
from pathlib import Path
from random import Random


# Left alphabet: x1,x2,x3,p11,p12,p21,p22; right: y1,y2,y3,q11,q12,q21,q22.
# Output / term alphabet: z0/P0 followed by zij/Pij in row-major order.
P_HAT = [
    [{3: 1}, {4: 1}, {}],
    [{5: 1}, {6: 1}, {}],
    [{3: -1, 5: -1}, {4: -1, 6: -1}, {}],
]
Q_HAT = [
    [{3: 1}, {4: 1}, {3: -1, 4: -1}],
    [{5: 1}, {6: 1}, {5: -1, 6: -1}],
    [{}, {}, {}],
]
PHI = [{0: -1, 1: -1, 2: -1}]
PSI = [{0: 1, 1: 1, 2: 1}]
for i, j in product(range(3), repeat=2):
    PHI.append({i: 1, **P_HAT[i][j]})
    PSI.append({j: 1, **Q_HAT[i][j]})


def identity_coefficients() -> tuple[dict, dict]:
    lhs, rhs = Counter(), Counter()
    for term in range(10):
        outputs = [0] if term == 0 else [0, term]
        for s, a in PHI[term].items():
            for r, b in PSI[term].items():
                for z in outputs:
                    lhs[s, r, z] += a * b
    # G: outer product and four-term inner product.
    for i, j in product(range(3), repeat=2):
        z = 1 + 3 * i + j
        rhs[i, j, z] += 1
        # E: x*qhat + phat*y + phat*qhat at each outer output.
        for r, b in Q_HAT[i][j].items():
            rhs[i, r, z] += b
        for s, a in P_HAT[i][j].items():
            rhs[s, j, z] += a
            for r, b in Q_HAT[i][j].items():
                rhs[s, r, z] += a * b
    for k in range(3, 7):
        rhs[k, k, 0] += 1
    for s, r, z in product(range(7), range(7), range(10)):
        assert lhs[s, r, z] == rhs[s, r, z], (s, r, z)
    return dict(lhs), dict(rhs)


def check_fixed_layer_coefficients(gamma: dict) -> int:
    """Exhaust all coefficients for L=2,m=1, not just sample evaluations."""
    left = [u for u in product(range(7), repeat=2) if sum(s >= 3 for s in u) == 1]
    outputs = [w for w in product(range(10), repeat=2) if w.count(0) == 1]
    checked = 0
    for u, v, w in product(left, left, outputs):
        actual = gamma.get((u[0], v[0], w[0]), 0)
        actual *= gamma.get((u[1], v[1], w[1]), 0)
        expected = 1
        for s, r, z in zip(u, v, w):
            if z == 0:
                expected *= int(s == r and s >= 3)
            else:
                i, j = divmod(z - 1, 3)
                expected *= int(s == i and r == j)
        assert actual == expected, (u, v, w, actual, expected)
        checked += 1
    return checked


def recursive_product(a: list[int], b: list[int], levels: int) -> list[int]:
    """Full recursion; output 0 sums all ten children, others select one."""
    if levels == 0:
        return [a[0] * b[0]]
    stride = 7 ** (levels - 1)
    children = []
    for phi, psi in zip(PHI, PSI):
        encoded_a = [
            sum(c * a[s * stride + k] for s, c in phi.items())
            for k in range(stride)
        ]
        encoded_b = [
            sum(c * b[s * stride + k] for s, c in psi.items())
            for k in range(stride)
        ]
        children.append(recursive_product(encoded_a, encoded_b, levels - 1))
    result = [sum(c[k] for c in children) for k in range(10 ** (levels - 1))]
    for child in children[1:]:
        result.extend(child)
    return result


def digits(values, base: int) -> int:
    result = 0
    for value in values:
        result = base * result + value
    return result


def check_recursive_examples() -> int:
    rng = Random(261006783)
    checked = 0
    for levels, m in [(2, 1), (3, 1), (3, 2)]:
        n0, inner = 3 ** (levels - m), 4**m
        sets = list(combinations(range(levels), m))
        for _ in range(3):
            matrices_x = {
                q: [[rng.randrange(-3, 4) for _ in range(inner)] for _ in range(n0)]
                for q in sets
            }
            matrices_y = {
                q: [[rng.randrange(-3, 4) for _ in range(n0)] for _ in range(inner)]
                for q in sets
            }
            a, b = [], []
            for u in product(range(7), repeat=levels):
                q = tuple(i for i, s in enumerate(u) if s >= 3)
                if len(q) != m:
                    a.append(0)
                    b.append(0)
                    continue
                outer_index = digits((s for s in u if s < 3), 3)
                inner_index = digits((s - 3 for s in u if s >= 3), 4)
                a.append(matrices_x[q][outer_index][inner_index])
                b.append(matrices_y[q][inner_index][outer_index])
            results = recursive_product(a, b, levels)
            for w, actual in zip(product(range(10), repeat=levels), results):
                q = tuple(i for i, z in enumerate(w) if z == 0)
                if len(q) != m:
                    continue
                row = digits(((z - 1) // 3 for z in w if z), 3)
                col = digits(((z - 1) % 3 for z in w if z), 3)
                expected = sum(
                    matrices_x[q][row][k] * matrices_y[q][k][col]
                    for k in range(inner)
                )
                assert actual == expected, (levels, m, q, row, col)
                checked += 1
    return checked


def leaf_box(leaf: tuple[int, ...], t: int) -> tuple[int, ...]:
    """-1 denotes a wildcard; the tuple contains only the inner levels."""
    if t == 0:
        return (-1,) * len(leaf)
    nonzero_seen = 0
    for i in range(len(leaf) - 1, -1, -1):
        nonzero_seen += leaf[i] != 0
        if nonzero_seen == t:
            return (-1,) * i + leaf[i:]
    raise ValueError("The leaf order must be at least t")


def check_boxes() -> dict:
    examples = []
    for m in range(1, 5):
        leaves = list(product(range(10), repeat=m))
        for t in range(m + 1):
            eligible = {leaf for leaf in leaves if m - leaf.count(0) >= t}
            boxes = Counter(leaf_box(leaf, t) for leaf in eligible)
            assert len(boxes) == comb(m, t) * 9**t
            expanded = set()
            for box, count in boxes.items():
                assert count == 10 ** box.count(-1)
                assert sum(s in (-1, 0) for s in box) == m - t
                assert all(
                    i < j
                    for i, s in enumerate(box)
                    for j, r in enumerate(box)
                    if s == -1 and r == 0
                )
                for leaf in product(*(range(10) if s == -1 else [s] for s in box)):
                    assert leaf not in expanded
                    assert leaf_box(leaf, t) == box
                    expanded.add(leaf)
            assert expanded == eligible
            examples.append({"m": m, "t": t, "boxes": len(boxes), "leaves": len(eligible)})
    return next(item for item in examples if item["m"] == 4 and item["t"] == 2)


def check_leaf_counts() -> int:
    checked = 0
    for m in range(1, 31):
        levels = 19 * m
        total = comb(levels, m) * 9 ** (levels - m)
        previous = None
        for d in range(m + 1):
            alpha = comb(m, d) * 9**d
            beta = comb(levels, m - d) * 9 ** (levels - m + d)
            assert total * alpha == beta * comb(levels - m + d, d)
            if previous is not None:
                assert 2 * beta < previous
            previous = beta
            checked += 1
    assert 72 * Fraction(9, 8) ** 9 < 208 < 256
    # Raise the encoding inequality to its ninth power, avoiding real rounding.
    assert 10 ** (19 * 9) * 2 < (4**18 * 3**18) ** 9
    return checked


def figures() -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.patches import Circle, Rectangle

    target = Path(__file__).resolve().parents[2] / "images" / "2026 Fall"
    target.mkdir(parents=True, exist_ok=True)
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 12,
            "axes.spines.top": False,
            "axes.spines.right": False,
            "svg.hashsalt": "2610.06783v1",
            "savefig.facecolor": "#fafafa",
            "figure.facecolor": "#fafafa",
            "axes.facecolor": "#fafafa",
            "text.color": "#303446",
            "axes.labelcolor": "#303446",
            "xtick.color": "#303446",
            "ytick.color": "#303446",
        }
    )

    # A schematic graph, with precisely the selected four query entries shown.
    left_neighbors = [{0, 1}, {0}, {1}, {0, 1}]
    right_neighbors = [{0, 1}, {1}, {0}, set()]
    queries = {(0, 0), (1, 2), (2, 1), (3, 3)}
    fig, (graph, matrix) = plt.subplots(2, 1, figsize=(6.1, 7.0))
    graph.set_xlim(-0.35, 2.35)
    graph.set_ylim(-0.1, 1.55)
    graph.set_aspect("equal")
    graph.axis("off")
    outer_y = [0.15, 0.48, 0.81, 1.14]
    middle_y = [0.62, 1.02]
    for i in range(4):
        for k in left_neighbors[i]:
            graph.plot([0, 1], [outer_y[i], middle_y[k]],
                       color="#1e66a8" if i == 0 else "#c3c7ce",
                       lw=2.5 if i == 0 else 1, zorder=1)
        for k in right_neighbors[i]:
            graph.plot([1, 2], [middle_y[k], outer_y[i]],
                       color="#1e66a8" if i == 0 else "#c3c7ce",
                       lw=2.5 if i == 0 else 1, zorder=1)
    graph.plot([0, 2], [outer_y[0], outer_y[0]], color="#b45b13", lw=2.5,
               linestyle="--", zorder=1)
    graph.text(1, -0.07, "queried edge: two triangles", ha="center", fontsize=12)
    for x, letter in [(0, "a"), (2, "b")]:
        for i, y in enumerate(outer_y):
            graph.add_patch(Circle((x, y), 0.075, fc="#e7eef6", ec="#1e66a8", zorder=3))
            graph.text(x, y, f"{letter}{i+1}", ha="center", va="center", fontsize=10, zorder=4)
    for k, y in enumerate(middle_y):
        graph.add_patch(Circle((1, y), 0.075, fc="#e2efdf", ec="#3d7041", zorder=3))
        graph.text(1, y, f"m{k+1}", ha="center", va="center", fontsize=10, zorder=4)
    for x, label in [(0, "A: N vertices"), (1, "M: D vertices"), (2, "B: N vertices")]:
        graph.text(x, 1.42, label, ha="center", fontsize=12)
    graph.set_title("Common neighbors give one product entry", fontsize=14, pad=4)
    matrix.set_xlim(-0.6, 4.1)
    matrix.set_ylim(4.65, -0.55)
    matrix.set_aspect("equal")
    matrix.axis("off")
    for i, j in product(range(4), repeat=2):
        wanted = (i, j) in queries
        matrix.add_patch(Rectangle((j, i), 1, 1, ec="#ffffff", lw=2,
                                   fc="#dcebd5" if wanted else "#e8e9ed"))
        if wanted:
            matrix.text(j + 0.5, i + 0.5,
                        str(len(left_neighbors[i] & right_neighbors[j])),
                        ha="center", va="center", fontsize=18)
    for i in range(4):
        matrix.text(-0.18, i + 0.5, f"a{i+1}", ha="right", va="center")
        matrix.text(i + 0.5, -0.15, f"b{i+1}", ha="center", va="bottom")
    matrix.text(2, 4.35, "Green: requested entries   Gray: not requested",
                ha="center", fontsize=11)
    matrix.set_title("Only positions in W are returned", fontsize=14, pad=14)
    fig.subplots_adjust(hspace=0.4, top=0.94, bottom=0.05)
    fig.savefig(target / "lopsided-triangle-query.svg", bbox_inches="tight",
                metadata={"Date": None, "Description": "Schematic N=4, D=2 example; not an asymptotic benchmark."})
    plt.close(fig)

    # Plot logarithms of the exact integer quantities, with no enumeration of M entries.
    m, levels = 90, 19 * 90
    total = comb(levels, m) * 9 ** (levels - m)
    wanted = total // 2**m
    orders = list(range(m + 1))
    local_counts = [wanted * comb(m, d) * 9**d for d in orders]
    global_counts = [comb(levels, m - d) * 9 ** (levels - m + d) for d in orders]
    normalize = lambda values: [log10(v) - log10(total) for v in values]
    fig, ax = plt.subplots(figsize=(6.2, 4.8))
    ax.plot(orders, normalize(local_counts), color="#1e66a8", lw=2,
            label=r"$|U|\alpha_d/M$")
    ax.plot(orders, normalize(global_counts), color="#c23a52", lw=2,
            label=r"$\beta_d/M$")
    ax.plot(orders, normalize([min(a, b) for a, b in zip(local_counts, global_counts)]),
            color="#3b7351", lw=4, alpha=0.65, label="smaller bound")
    ax.axvline(m / 9, color="#6e7381", lw=1, linestyle="--")
    ax.text(m / 9 + 1.5, 45, r"$m/9$", fontsize=12)
    ax.axhline(0, color="#a5aab5", lw=0.8)
    ax.set_xlabel("Leaf order d")
    ax.set_ylabel("Log₁₀ of normalized count")
    ax.set_title("Count locally at low order,\nglobally at high order", fontsize=15)
    ax.legend(loc="lower left", frameon=True, framealpha=0.95,
              facecolor="#fafafa", edgecolor="none", fontsize=12)
    ax.grid(alpha=0.15)
    fig.tight_layout()
    fig.savefig(target / "sparse-product-leaf-counts.svg", metadata={"Date": None})
    plt.close(fig)
    for path in [target / "lopsided-triangle-query.svg", target / "sparse-product-leaf-counts.svg"]:
        # Matplotlib emits indentation-only blank lines; keep the tracked SVGs clean.
        path.write_text("\n".join(line.rstrip() for line in path.read_text().splitlines()) + "\n")
        print(f"Figure: {path}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--figures", action="store_true", help="Regenerate SVGs using matplotlib")
    args = parser.parse_args()
    lhs, _ = identity_coefficients()
    report = {
        "identity_coefficients": 7 * 7 * 10,
        "fixed_layer_coefficients": check_fixed_layer_coefficients(lhs),
        "recursive_matrix_entries": check_recursive_examples(),
        "box_example": check_boxes(),
        "incidence_and_decay_checks": check_leaf_counts(),
        "gamma0": log(20 / 9) / (9 * log(4)),
        "q0": (log(72) + 9 * log(9 / 8)) / (9 * log(4)),
        "epsilon_star": log(4) / (5 * log(10)),
    }
    print(json.dumps(report, indent=2))
    if args.figures:
        figures()
    print("All finite exact checks passed.")


if __name__ == "__main__":
    main()