"""Exact finite certificates; not a proof of asymptotic optimality."""
import json
from fractions import Fraction
from math import comb
from pathlib import Path

def span(generators):
    code = [0]
    for g in generators:
        code += [c ^ g for c in code]
    return code

results = []
for n in (2, 4):
    rows = [((1 << n) - 1) << (i * n) for i in range(n)]
    corner = 1 << (n - 1)
    generators = [row ^ corner for row in rows]
    code, row_code = span(generators), span(rows)
    assert len(set(code)) == 1 << n
    totals = [0, 0]
    conditional = []
    for v in range(1 << n):
        scores = [0, 0]
        for u in range(1 << n):
            q = sum(v << (i * n) for i in range(n) if (u >> i) & 1)
            for j, candidates in enumerate((row_code, code)):
                scores[j] += q.bit_count() - min((q ^ c).bit_count() for c in candidates)
        totals = [a + b for a, b in zip(totals, scores)]
        actual = Fraction(scores[1] - scores[0], 1 << n)
        k, t = v.bit_count(), (v >> (n - 1)) & 1
        expected = (Fraction(0) if k < n // 2 else
                    Fraction(1, 2) - Fraction(t, 1 << n) if k == n // 2 else
                    -Fraction(t, 2))
        assert actual == expected
        conditional.append({"v": v, "row_savings_sum": scores[0], "new_savings_sum": scores[1]})
    means = [Fraction(t, 1 << (2 * n)) for t in totals]
    delta = Fraction(comb(n, n // 2), 1 << (n + 1)) * (1 - Fraction(1, 1 << n)) - Fraction(1, 8)
    assert means[1] - means[0] == delta > 0
    assert means[0] == Fraction(n * n * comb(n - 1, n // 2), 1 << (n + 1))
    results.append({"n": n, "generators": generators, "code_size": len(code),
                    "row_mean_savings": str(means[0]), "new_mean_savings": str(means[1]),
                    "new_mean_distance": str(Fraction(n * n, 4) - means[1]),
                    "improvement": str(delta), "conditional": conditional})

target = Path(__file__).with_name("row-parity-counterexample-output.json")
target.write_text(json.dumps({"status": "passed", "scope": "all queries for the two specified finite codes", "results": results}, indent=2) + "\n")
print(json.dumps({"status": "passed", "output": str(target), "results": [{k:v for k,v in r.items() if k != "conditional"} for r in results]}))
