#!/usr/bin/env python3
"""Regenerate or verify the exact examples; uses only the Python standard library.

This is a finite computational check, not a replacement for the proofs.
Run from any directory: python3 verify_examples.py [--write]
"""
from collections import Counter
from functools import lru_cache
from itertools import accumulate, product
from math import gcd
from pathlib import Path
import json
import sys

HERE = Path(__file__).resolve().parent
PATH = ["C", "NW", "W", "NE", "S", "N", "SW", "E", "SE"]
SOUTH_UP = [["SE", "S", "SW"], ["E", "C", "W"], ["NE", "N", "NW"]]

def sieve(degree, count, length=None):
    if degree < 1 or count < 1:
        raise ValueError("degree and count must be positive")
    initial = list(range(1, (degree * count if length is None else length) + 1))
    values = initial
    stages = []
    for stride in range(degree, 1, -1):
        removed = [position % stride == 0 for position in range(1, len(values) + 1)]
        kept = [value for value, drop in zip(values, removed) if not drop]
        sums = list(accumulate(kept))
        stages.append({"stride": stride, "input": values, "removed": removed,
                       "removedPositions": [i + 1 for i, drop in enumerate(removed) if drop],
                       "kept": kept, "sums": sums})
        values = sums
    return {"degree": degree, "count": count, "initial": initial,
            "stages": stages, "values": values}

def wrap9(value):
    return 1 + (value - 1) % 9

def chart(degree, centre=9, direction=1):
    sequence = sieve(degree, 9)["values"]
    offsets = [value - sequence[0] for value in sequence]
    stars = [wrap9(centre + direction * offset) for offset in offsets]
    placed = dict(zip(PATH, stars))
    counts = {str(star): stars.count(star) for star in range(1, 10)}
    return {"degree": degree, "centre": centre, "direction": direction,
            "sequence": [str(v) for v in sequence], "offsets": [str(v) for v in offsets],
            "powerResidues": [v % 9 for v in sequence], "stars": stars,
            "distinct": len(set(stars)), "counts": counts,
            "southUpRows": [[placed[palace] for palace in row] for row in SOUTH_UP]}

@lru_cache(None)
def row(m, q, block, column):
    if m == 0:
        return 1
    return (sum(row(m - 1, q + 1, block, a) for a in range(column + 1))
            + sum(row(m - 1, q + 1, j, a)
                  for j in range(block) for a in range(q + 1)))

def decimal_trace(trace):
    converted = dict(trace)
    for key in ["initial", "values"]:
        converted[key] = [str(v) for v in trace[key]]
    converted["stages"] = []
    for stage in trace["stages"]:
        item = dict(stage)
        for key in ["input", "kept", "sums"]:
            item[key] = [str(v) for v in stage[key]]
        converted["stages"].append(item)
    return converted

def difference_table(degree):
    rows = [[n ** degree for n in range(8)]]
    while len(rows[-1]) > 1:
        rows.append([b - a for a, b in zip(rows[-1], rows[-1][1:])])
    return {"degree": degree, "inputArguments": list(range(8)),
            "rows": [[str(v) for v in row_values] for row_values in rows]}

def expected_residue_counts(degree):
    if degree == 1:
        return {r: 1 for r in range(9)}
    g = gcd(degree, 6)
    image = {pow(2, g * j, 9) for j in range(6 // g)}
    return {r: (3 if r == 0 else g if r in image else 0) for r in range(9)}

def checks():
    count_cases = 0
    for degree in range(1, 13):
        for count in range(1, 26):
            trace = sieve(degree, count)
            expected = [n ** degree for n in range(1, count + 1)]
            assert trace["values"] == expected
            for stage in trace["stages"]:
                assert len(stage["input"]) == count * stage["stride"]
                assert len(stage["sums"]) == count * (stage["stride"] - 1)
            shortest = degree * (count - 1) + 1
            assert sieve(degree, count, shortest)["values"] == expected
            assert len(sieve(degree, count, shortest - 1)["values"]) == count - 1
            count_cases += 1
    row_cases = 0
    for m, q, block, column in product(range(5), range(4), range(5), range(5)):
        assert row(m + 1, q + 1, block, column) == (
            row(m + 1, q, block, column) + block * row(m, q + 1, block, column))
        row_cases += 1
    for degree in range(13):
        for index in range(8):
            assert row(degree, 0, index, 0) == (index + 1) ** degree
    chart_cases = 0
    for degree, centre, direction in product(range(1, 13), range(1, 10), [-1, 1]):
        current = chart(degree, centre, direction)
        assert current["stars"][0] == centre
        assert all(1 <= star <= 9 for star in current["stars"])
        predicted = 9 if degree == 1 else 1 + 6 // gcd(degree, 6)
        assert current["distinct"] == predicted
        chart_cases += 1
    for degree in range(1, 61):
        actual = Counter(pow(n, degree, 9) for n in range(1, 10))
        assert {r: actual[r] for r in range(9)} == expected_residue_counts(degree)
        if degree >= 2:
            assert [pow(n, degree, 9) for n in range(1, 10)] == [
                pow(n, degree + 6, 9) for n in range(1, 10)]
    assert sieve(9, 100)["values"][-1] == 10 ** 18
    return {"finiteSieveAndMinimalPrefixCases": count_cases,
            "blockInvariantCases": row_cases, "outputPowerCases": 104,
            "chartCases": chart_cases, "residueDegrees": 60,
            "largeIntegerCase": "100^9 = 1000000000000000000"}

def make_examples():
    return {
        "schemaVersion": "fengshui-computer/mathematics-examples-v1",
        "conventions": {
            "deletionPositions": "one-based current-row positions; restart counting each pass",
            "summation": "inclusive prefix sum across all retained blocks; do not reset at block boundaries",
            "integerEncoding": "unreduced arithmetic values are decimal strings",
            "flightPath": PATH, "southUpPalaces": SOUTH_UP,
            "projection": "1 + ((centre - 1 + direction * ((i + 1)^degree - 1)) mod 9)",
            "modulo": "least nonnegative remainder"
        },
        "finiteTables": [decimal_trace(sieve(degree, 4)) for degree in [2, 3, 4]],
        "prefixLengths": [{"degree": d, "count": 4, "completeBlockLength": 4 * d,
                           "minimalLength": 3 * d + 1} for d in range(1, 10)],
        "nineStarCases": [chart(d, 9, direction) for direction in [1, -1] for d in range(1, 10)],
        "residueClassification": [{"degree": d, "gcdWithSix": gcd(d, 6),
                                   "powerResidues": [pow(n, d, 9) for n in range(1, 10)],
                                   "counts": {str(r): v for r, v in expected_residue_counts(d).items()},
                                   "distinct": 9 if d == 1 else 1 + 6 // gcd(d, 6)}
                                  for d in range(1, 19)],
        "finiteDifferenceTables": [difference_table(d) for d in [2, 3, 4]],
        "functionSpaceExample": {"domainSize": 2, "codomainSize": 3,
                                 "tables": [list(v) for v in product(range(3), repeat=2)],
                                 "count": "9"},
        "verification": {"kind": "independent finite computational checks; universal proofs are in the notes",
                         "results": checks()}
    }

def main():
    if any(arg != "--write" for arg in sys.argv[1:]):
        raise SystemExit("usage: python3 verify_examples.py [--write]")
    examples = make_examples()
    destination = HERE / "worked-examples.json"
    if "--write" in sys.argv[1:]:
        destination.write_text(json.dumps(examples, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        print(f"Wrote {destination.name}")
    else:
        assert json.loads(destination.read_text(encoding="utf-8")) == examples, "Saved examples differ from regeneration"
        print(f"Verified {destination.name}")
    print(json.dumps(examples["verification"]["results"], indent=2))

if __name__ == "__main__":
    main()