#!/usr/bin/env python3
"""Construct and verify a Lempel-Golomb Costas permutation over GF(32)."""

from __future__ import annotations

import argparse
from pathlib import Path


# x^5 + x^2 + 1, represented with the x^5 term included.
PRIMITIVE_POLYNOMIAL = 0b100101
FIELD_SIZE = 32
MULTIPLICATIVE_ORDER = FIELD_SIZE - 1


def gf_multiply(left: int, right: int) -> int:
    """Multiply two polynomial-basis GF(2^5) elements."""
    product = 0
    while right:
        if right & 1:
            product ^= left
        right >>= 1
        left <<= 1
        if left & FIELD_SIZE:
            left ^= PRIMITIVE_POLYNOMIAL
    return product


def powers(base: int) -> list[int]:
    """Return base**0 through base**30 in GF(32)."""
    result = [1]
    for _ in range(1, MULTIPLICATIVE_ORDER):
        result.append(gf_multiply(result[-1], base))
    return result


def require_primitive(base: int) -> list[int]:
    base_powers = powers(base)
    if len(set(base_powers)) != MULTIPLICATIVE_ORDER:
        raise ValueError(f"field element {base} is not primitive")
    if gf_multiply(base_powers[-1], base) != 1:
        raise ValueError(f"field element {base} does not have order 31")
    return base_powers


def construct(alpha: int = 0b00010, beta: int = 0b00100) -> list[int]:
    """Return p where alpha**i + beta**p[i] = 1, for i=1..30."""
    alpha_powers = require_primitive(alpha)
    beta_powers = require_primitive(beta)
    beta_log = {value: exponent for exponent, value in enumerate(beta_powers)}

    permutation = []
    for i in range(1, FIELD_SIZE - 1):
        target = 1 ^ alpha_powers[i]  # Addition in characteristic two is XOR.
        j = beta_log[target]
        if not 1 <= j <= FIELD_SIZE - 2:
            raise AssertionError(f"unexpected exponent {j} for row {i}")
        permutation.append(j)
    return permutation


def verify_costas(permutation: list[int]) -> int:
    n = len(permutation)
    if sorted(permutation) != list(range(1, n + 1)):
        raise AssertionError("output is not a 1-based permutation")

    displacements: set[tuple[int, int]] = set()
    for first_row in range(n):
        for second_row in range(first_row + 1, n):
            vector = (
                second_row - first_row,
                permutation[second_row] - permutation[first_row],
            )
            if vector in displacements:
                raise AssertionError(f"duplicate displacement vector: {vector}")
            displacements.add(vector)

    expected = n * (n - 1) // 2
    if len(displacements) != expected:
        raise AssertionError("wrong displacement count")
    return expected


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, help="optional output artifact path")
    args = parser.parse_args()

    permutation = construct()
    displacement_count = verify_costas(permutation)
    line = " ".join(map(str, permutation)) + "\n"
    if args.output:
        args.output.write_text(line, encoding="ascii")
    print(line, end="")
    print(
        f"verified order {len(permutation)} with {displacement_count} distinct vectors",
        flush=True,
    )


if __name__ == "__main__":
    main()