"""333 (p3): every generating pair, up to the group's own symmetries, checked exactly.

Coordinates: the lattice basis e1 = (1, 0), e2 = (1/2, sqrt3/2), nearest clocks 1 apart, so a point
p = (p1, p2) has squared length p1^2 + p1 p2 + p2^2. R = the 1/3 turn anticlockwise, R e1 = e2 - e1.
The 3-fold centres form the finer lattice L' = Z^2 + Z (1/3, 1/3) in three kinds:
  clocks (lattice points), up-triangle centres (1/3, 1/3) + Z^2, down-triangle centres (2/3, 2/3) + Z^2.
Elements of 333: translations t(u), u in Z^2, and the turns by +-1/3 about the points of L'.

Claims checked (the page's section 1, "333"):
  A. x, y, z = the anticlockwise 1/3 turns about 0, (1/3,1/3), (-1/3,2/3) satisfy x^3 = y^3 = z^3 = xyz = 1,
     and t(e1) = y x^-1, t(e2) = y^-1 x: so x and y generate.
  B. On a sample of 86 turns and 30 translations near the origin, a pair {g, h} generates 333 exactly when
       - g, h are turns about two centres 1/sqrt3 apart (the nearest possible), either directions; or
       - one is a turn and the other a translation by one of the six shortest lattice vectors.
     Every other pair fails, and the subgroup it makes is found:
       turns about one centre: the 3 turns about it; two translations: no turn at all;
       turns about centres d apart: a smaller 333 of index 3 d^2; turn and translation u: index |u|^2.
     Generation is proved by explicit words for t(e1), t(e2), x; failure by a proper image in G/nZ^2.
  C. Up to the affine symmetries of the pattern (the automorphisms of the group, by Bieberbach), every
     generating pair in the sample is one of exactly three kinds:
       (x, y)       two turns the same way about neighbouring centres;
       (x, y^-1)    two opposite turns about neighbouring centres;
       (x, t(e1))   a turn and a shortest translation.
  C2. Those are unordered pairs. As ordered pairs there are 4 kinds ((x, t(e1)) and (t(e1), x) differ); up to the
     symmetries that keep orientation, 5 unordered kinds and 8 ordered ones.
  D. Nielsen moves (swap, invert one, multiply one by the other or its inverse) join every ordered generating
     pair of the sample to (x, y), by a breadth-first search kept inside a box around the sample.
"""
from fractions import Fraction as F
from itertools import combinations

from wallpaper import (Aff, T, rot_about, I2, mm, mv, vsub, WGroup, closure_mod, find_words, evaluate,
                       word_str, fixed_point, inv_unimodular)

R = ((-1, -1), (1, 0))
R2 = mm(R, R)
assert mm(R2, R) == I2 and R2 != I2
G = WGroup("333", "p3", [Aff(I2), Aff(R), Aff(R2)])
G.check_is_group()


def n2(v):
    return v[0] * v[0] + v[0] * v[1] + v[1] * v[1]


def kind(p):
    k = (3 * F(p[0])) % 3
    return {0: "clock", 1: "up", 2: "down"}[int(k)]


assert kind((F(1, 3), F(1, 3))) == "up" and kind((F(2, 3), F(2, 3))) == "down" and kind((F(-1, 3), F(2, 3))) == "down"

# ---------------------------------------------------------------- A. the canonical presentation's generators
x = rot_about(R, (0, 0))
y = rot_about(R, (F(1, 3), F(1, 3)))
z = rot_about(R, (F(-1, 3), F(2, 3)))
e = Aff(I2)
assert x ** 3 == e and y ** 3 == e and z ** 3 == e and x * y * z == e
assert y == Aff(R, (1, 0)) and z == Aff(R, (0, 1))
gens = {"x": x, "y": y}
assert evaluate([("y", 1), ("x", -1)], gens) == T((1, 0))
assert evaluate([("y", -1), ("x", 1)], gens) == T((0, 1))
assert all(G.contains(g) for g in (x, y, z))
print("A. x, y, z (anticlockwise thirds about a clock, an up-centre, a down-centre) satisfy "
      "x^3 = y^3 = z^3 = xyz = 1;  t(e1) = y x^-1,  t(e2) = y^-1 x,  so {x, y} generates 333")

# shortest words for the three representative pairs, by breadth-first search
reps = {
    "same-way turns (x, y)": {"x": x, "y": y},
    "opposite turns (x, w = y^-1)": {"x": x, "w": y.inv()},
    "turn and shortest slide (x, s = t(e1))": {"x": x, "s": T((1, 0))},
}
for lab, gg in reps.items():
    targets = {"t(e1)": T((1, 0)), "t(e2)": T((0, 1)), "x": x}
    found = find_words(gg, targets)
    assert set(found) == set(targets), (lab, found)
    for t, w in found.items():
        assert evaluate(w, gg) == targets[t]
    print(f"   {lab}: " + "; ".join(f"{t} = {word_str(found[t])}" for t in targets))

# ---------------------------------------------------------------- B. the sample, and every pair in it
centres = []
for i in range(-9, 10):
    for j in range(-9, 10):
        if (i - j) % 3 == 0:
            p = (F(i, 3), F(j, 3))
            if n2(p) <= 4:
                centres.append(p)
turns = []
for p in centres:
    for Mx in (R, R2):
        g = rot_about(Mx, p)
        assert G.contains(g), p
        turns.append(("turn", p, Mx, g))
slides = []
for a in range(-4, 5):
    for b in range(-4, 5):
        if 0 < n2((a, b)) <= 7:
            slides.append(("slide", (F(a), F(b)), I2, T((a, b))))
sample = turns + slides
print(f"B. sample: {len(turns)} turns about the {len(centres)} centres within 2 of a clock, "
      f"and {len(slides)} translations of squared length <= 7")


def solve_basis(u, v):
    """Integers (p, q, r, s) with e1 = p u + q v, e2 = r u + s v, for a unimodular pair (u, v)."""
    D = u[0] * v[1] - u[1] * v[0]
    assert D in (1, -1)
    p, q = v[1] * D, -u[1] * D
    r, s = -v[0] * D, u[0] * D
    assert (p * u[0] + q * v[0], p * u[1] + q * v[1]) == (1, 0)
    assert (r * u[0] + s * v[0], r * u[1] + s * v[1]) == (0, 1)
    return p, q, r, s


def power_word(name, k):
    return [(name, 1 if k > 0 else -1)] * abs(k)


def prove_generates(g, h):
    """Explicit words in g, h for t(e1), t(e2) and x; returns their total length."""
    gg = {"g": g, "h": h}
    # a translation tau and a turn r inside <g, h>, as words
    if g.M != I2 and h.M != I2:
        if g.M == h.M:
            tau_w = [("g", 1), ("h", -1)]
        else:
            tau_w = [("g", 1), ("h", 1)]
        r_w, r = [("g", 1)], g
    else:
        if g.M == I2:
            g, h = h, g
            gg = {"g": g, "h": h}
        tau_w, r_w, r = [("h", 1)], [("g", 1)], g
    tau = evaluate(tau_w, gg)
    assert tau.M == I2
    u = tau.t
    tau2_w = r_w + tau_w + [(n, -s) for (n, s) in reversed(r_w)]  # r tau r^-1 = t(R^{+-1} u)
    tau2 = evaluate(tau2_w, gg)
    assert tau2.M == I2
    v = tau2.t
    p, q, rr, s = solve_basis((int(u[0]), int(u[1])), (int(v[0]), int(v[1])))

    def comb(a, b):
        return (tau_w * a if a > 0 else [(n, -e2) for (n, e2) in reversed(tau_w)] * (-a)) + \
               (tau2_w * b if b > 0 else [(n, -e2) for (n, e2) in reversed(tau2_w)] * (-b))
    w1, w2 = comb(p, q), comb(rr, s)
    assert evaluate(w1, gg) == T((1, 0)) and evaluate(w2, gg) == T((0, 1))
    # x = t(-s0) r^k where r^k has linear part R
    k = 1 if r.M == R else 2
    rk = r ** k
    assert rk.M == R
    m = (-rk.t[0], -rk.t[1])
    assert m[0].denominator == 1 and m[1].denominator == 1
    wx = []
    for (cnt, base) in ((int(m[0]), w1), (int(m[1]), w2)):
        wx += base * cnt if cnt > 0 else [(n, -e2) for (n, e2) in reversed(base)] * (-cnt)
    wx += r_w * k
    assert evaluate(wx, gg) == x
    return len(w1) + len(w2) + len(wx)


stats = {}
gen_pairs = []
for (A_, B_) in combinations(sample, 2):
    ta, pa, Ma, ga = A_
    tb, pb, Mb, gb = B_
    if ta == "turn" and tb == "turn":
        d2 = n2(vsub(pa, pb))
        if d2 == 0:
            label, pred_gen, pred_index = "turns, same centre", False, None
        else:
            same = kind(pa) == kind(pb)
            label = f"turns, centres at squared distance {d2}" + (" (same kind)" if same else " (different kinds)")
            pred_index = 3 * d2
            pred_gen = pred_index == 1
    elif ta == "slide" and tb == "slide":
        label, pred_gen, pred_index = "two translations", False, None
    else:
        u = pb if ta == "turn" else pa
        pred_index = n2(u)
        label = f"turn and translation, |u|^2 = {pred_index}"
        pred_gen = pred_index == 1
    assert pred_index is None or pred_index.denominator == 1
    if pred_gen:
        prove_generates(ga, gb)
        gen_pairs.append((A_, B_))
        verdict = "generates"
    else:
        n = 2 if pred_index is None else int(pred_index)
        img = closure_mod([ga.mod(n), gb.mod(n)], n)
        full = G.order_mod(n)
        assert len(img) < full, (label, pa, pb)
        if pred_index is not None:
            assert full == len(img) * n, (label, len(img), full)  # index exactly 3 d^2 or |u|^2
        verdict = "fails"
    key = (label, verdict)
    stats[key] = stats.get(key, 0) + 1

print(f"   all {sum(stats.values())} unordered pairs decided; every prediction held:")


def sort_key(k):
    lab = k[0]
    order = ["turns, same centre", "two translations"]
    if lab in order:
        return (0, order.index(lab), F(0))
    if lab.startswith("turns, centres"):
        return (1, 0, F(lab.split("distance ")[1].split(" ")[0]))
    return (2, 0, F(lab.split("= ")[1]))


for k in sorted(stats, key=sort_key):
    lab, verdict = k
    extra = ""
    if verdict == "fails":
        if lab.startswith("turns, centres"):
            d2 = F(lab.split("distance ")[1].split(" ")[0])
            extra = f" -> a 333 of index {3 * d2}"
        elif lab.startswith("turn and"):
            extra = f" -> a 333 of index {lab.split('= ')[1]}"
        elif lab == "turns, same centre":
            extra = " -> the 3 turns about that centre"
        else:
            extra = " -> translations only"
    print(f"     {stats[k]:5d}  {lab}: {verdict}{extra}")

# ---------------------------------------------------------------- C. three kinds up to symmetry
R6 = ((0, -1), (1, 1))      # the 1/6 turn: e1 -> e2, e2 -> e2 - e1
S = ((0, 1), (1, 0))        # the mirror swapping e1 and e2
assert mm(R6, R6) == R
D6 = []
Mx = I2
for k in range(6):
    D6.append(Mx)
    D6.append(mm(Mx, S))
    Mx = mm(R6, Mx)
assert len(set(D6)) == 12
for A in D6:
    Ai = inv_unimodular(A)
    assert mm(mm(A, R), Ai) in (R, R2)
# the normaliser's generators normalise G (so every composite does): the 1/6 turn, the mirror,
# and the translation from a clock to an up-centre
for nrm in (Aff(R6), Aff(S), Aff(I2, (F(1, 3), F(1, 3)))):
    for g in (x, T((1, 0)), T((0, 1))):
        assert G.contains(nrm * g * nrm.inv()) and G.contains(nrm.inv() * g * nrm)
rep_pairs = {1: (x, y), 2: (x, y.inv()), 3: (x, T((1, 0)))}


def classify(g, h):
    hits = set()
    for (a, b) in ((g, h), (h, g)):
        if a.M == I2:
            continue
        c = fixed_point(a)
        for A in D6:
            nrm = Aff(A, (0, 0))
            nrm = Aff(A, vsub((0, 0), mv(A, c)))  # sends the centre of a to the origin
            a2, b2 = nrm * a * nrm.inv(), nrm * b * nrm.inv()
            for kk, (r1, r2) in rep_pairs.items():
                if a2 == r1 and b2 == r2:
                    hits.add(kk)
    return hits


kinds = {1: 0, 2: 0, 3: 0}
for (A_, B_) in gen_pairs:
    h = classify(A_[3], B_[3])
    assert len(h) == 1, (A_, B_, h)
    kinds[h.pop()] += 1
# the three kinds are different: a translation in the pair or not; turns the same way or not
assert rep_pairs[1][0].M == rep_pairs[1][1].M and rep_pairs[2][0].M != rep_pairs[2][1].M
print(f"C. each of the {len(gen_pairs)} generating pairs is carried by a symmetry of the pattern to exactly one of:")
print(f"     {kinds[1]:4d}  (x, y)      two turns the same way about neighbouring centres")
print(f"     {kinds[2]:4d}  (x, y^-1)   two opposite turns about neighbouring centres")
print(f"     {kinds[3]:4d}  (x, t(e1))  a turn and one of the six shortest translations")
print("   and no symmetry carries one kind to another (a translation in the pair, or turns the same way, is kept)")


# ---------------------------------------------------------------- C2. ordered pairs, and orientation
# The affine symmetries of 333 are v -> A v + w with A in D6 and w in L' (the 3-fold centres). Those that send
# a centre c to the origin are v -> A (v - c): 12 of them. The orbit of an ordered pair (g, h) under all of them
# is then read off from its images whose first turn sits at the origin; its canonical form is the smallest key.
def canon(g, h, mats):
    a = g if g.M != I2 else h            # the pair's first turn
    c = fixed_point(a)
    best = None
    for A in mats:
        nrm = Aff(A, vsub((0, 0), mv(A, c)))
        k = ((nrm * g * nrm.inv()).key(), (nrm * h * nrm.inv()).key())
        if best is None or k < best:
            best = k
    return best


def det2x2(A):
    return A[0][0] * A[1][1] - A[0][1] * A[1][0]


D6_keep = [A for A in D6 if det2x2(A) == 1]
assert len(D6_keep) == 6
ordered = [(A_[3], B_[3]) for (A_, B_) in gen_pairs] + [(B_[3], A_[3]) for (A_, B_) in gen_pairs]
counts = {}
for lab, mats in (("all symmetries", D6), ("keeping orientation", D6_keep)):
    o_ord = {canon(g, h, mats) for (g, h) in ordered}
    o_un = {min(canon(g, h, mats), canon(h, g, mats)) for (g, h) in ordered}
    counts[lab] = (len(o_un), len(o_ord))
assert counts == {"all symmetries": (3, 4), "keeping orientation": (5, 8)}, counts
print(f"C2. the {len(ordered)} ordered generating pairs of the sample: up to all symmetries of the pattern, "
      f"{counts['all symmetries'][0]} kinds as unordered pairs and {counts['all symmetries'][1]} as ordered ones "
      "(a turn then a slide, or a slide then a turn); up to the symmetries that keep orientation, "
      f"{counts['keeping orientation'][0]} unordered and {counts['keeping orientation'][1]} ordered")


# ---------------------------------------------------------------- D. Nielsen moves join every generating pair
# Breadth-first search from (x, y) over ordered pairs, by Nielsen moves (swap; invert one; multiply one by the
# other or its inverse, on either side). Every move keeps a generating pair generating. The search keeps to
# turns about centres within 3 of a clock and slides of squared length <= 13, a box that holds the sample.
def small(g):
    if g.M == I2:
        return n2(g.t) <= 13
    return n2(fixed_point(g)) <= 9


def moves(g, h):
    gi, hi = g.inv(), h.inv()
    return [(h, g), (gi, h), (g, hi), (g * h, h), (g * hi, h), (h * g, h), (hi * g, h),
            (g, h * g), (g, h * gi), (g, g * h), (g, gi * h)]


start = (x, y)
seen = {(start[0].key(), start[1].key())}
frontier = [start]
while frontier:
    nxt = []
    for (g, h) in frontier:
        for (a, b) in moves(g, h):
            if not (small(a) and small(b)):
                continue
            k = (a.key(), b.key())
            if k not in seen:
                seen.add(k)
                nxt.append((a, b))
    frontier = nxt
want = {(g.key(), h.key()) for (g, h) in ordered}
assert want <= seen, len(want - seen)
print(f"D. Nielsen moves from (x, y), kept inside a box around the sample, reach {len(seen)} ordered pairs, among "
      f"them all {len(want)} ordered generating pairs of the sample: every one is joined to (x, y)")
print(f"HEADLINE 333: {sum(stats.values())} pairs from a sample of {len(sample)} elements; {len(gen_pairs)} generate, "
      f"all of 3 kinds up to symmetry ({kinds[1]} + {kinds[2]} + {kinds[3]}; 4 kinds as ordered pairs); turns generate "
      f"only about centres 1/sqrt3 apart, otherwise a 333 of index 3 d^2 (d = distance between centres); Nielsen moves "
      f"join all {len(want)} ordered generating pairs to (x, y)")
