# Toy Balog–Wooley A = G·P in Z[phi] (d = 2), after Bloom–Sawin–Schildkraut–Zhelezov §2.
# Elements a + b*phi stored as integer pairs (a, b). phi^2 = phi + 1.
import math, sys
PHI = (1 + 5 ** 0.5) / 2
PSI = (1 - 5 ** 0.5) / 2

def emb(e):  # Minkowski embedding (sigma1, sigma2)
    a, b = e
    return (a + b * PHI, a + b * PSI)

def mul(x, y):
    a, b = x; c, d = y
    return (a * c + b * d, a * d + b * c + b * d)

def add(x, y):
    return (x[0] + y[0], x[1] + y[1])

def unit(k):  # phi^k, k may be negative; phi^-1 = phi - 1
    u, base = (1, 0), ((0, 1) if k >= 0 else (-1, 1))
    for _ in range(abs(k)):
        u = mul(u, base)
    return u

# --- Lemma 3.4 sharpness: phi's embeddings sit exactly on the boundary.
s1, s2 = emb(unit(1))
print(f"phi embeddings: {s1:.6f}, {s2:.6f}  |.| = {abs(s1):.6f}, {abs(s2):.6f}")
for x in [0.62, 0.9, 1.0, 1.3, 1.61]:
    print(f"  x={x}: (x-1/x)^2 = {(x - 1/x)**2:.4f}")

# --- the construction
X, c, K = int(sys.argv[1]) if len(sys.argv) > 1 else 120, 0.1, int(sys.argv[2]) if len(sys.argv) > 2 else 3
lo, hi = X - c * X, X + c * X
P = []
for b in range(-int(2 * c * X) - 2, int(2 * c * X) + 3):
    for a in range(int(lo) - 3 * abs(b) - 3, int(hi) + 3 * abs(b) + 4):
        x, y = emb((a, b))
        if lo <= x <= hi and lo <= y <= hi:
            P.append((a, b))
G = [unit(k) for k in range(-K, K + 1)]
A = list({mul(u, p) for u in G for p in P})
n = len(A)
print(f"|P|={len(P)}  |G|={len(G)}  |A|={n}  (direct product: {n == len(P) * len(G)})")
S = {add(x, y) for i, x in enumerate(A) for y in A[i:]}
M = {mul(x, y) for i, x in enumerate(A) for y in A[i:]}
print(f"|A+A|={len(S)}  ratio to |A|^2 = {len(S)/n**2:.4f}   log|A+A|/log|A| = {math.log(len(S))/math.log(n):.4f}")
print(f"|AA| ={len(M)}  ratio to |A|^2 = {len(M)/n**2:.4f}   log|AA| /log|A| = {math.log(len(M))/math.log(n):.4f}")
print(f"n(n+1)/2 = {n*(n+1)//2}  (max possible for either)")

# --- SVG: A in the Minkowski plane
import json
pts = [emb(e) for e in A]
json.dump({"pts": pts, "X": X, "K": K}, open(sys.argv[3] if len(sys.argv) > 3 else "pts.json", "w"))
