"""Part 7: a floor with five squares at every corner.

On a flat floor four squares meet at every corner. Make it five, everywhere, and the floor can't
lie flat any more: it ruffles like a lettuce leaf, and it is the hyperbolic plane. Nothing about
a single field changes, it is still a square with four neighbours, but the number of fields
within r steps of the start grows exponentially instead of like r squared.

The floor is built exactly, ring by ring, with nothing but whole numbers. The fields within n
steps of the start form a patch whose boundary is a ring of corners, each touched by some of the
five squares it will eventually have. The next ring of fields is the fields across the patch's
boundary edges, and the corners say how they fit together:

- a boundary corner missing just one square gets one new field that covers both boundary edges
  on either side of it (a notch fills in);
- a corner missing two gets two new fields that meet along a new edge from it, which completes
  it;
- a corner missing three or more gets two new fields that don't meet, and stays on the boundary.

Each new field has one or two fields of the previous ring beside it: its parents. No coordinates,
no rounding, and nothing is ever found twice, because every field is made exactly once, in the
ring where it first appears.

For drawing (and to show why rounding is a trap), `float_layers` builds the same floor the usual
way, with 3 x 3 matrices on the hyperboloid and fields recognised by their rounded centres.
"""
import math

Q = 5                               # squares at every corner


def rings(depth, q=Q):
    """Build the floor out to `depth` steps. Returns (tiles, layer):
    tiles[i] is field i's four corners in order round it, and layer[i] its distance from field 0.
    Also returns parents: the fields of the previous ring that each field shares an edge with."""
    tiles = [(0, 1, 2, 3)]
    layer = [0]
    parents = [()]
    deg = {0: 1, 1: 1, 2: 1, 3: 1}
    boundary = [0, 1, 2, 3]
    edge_tile = {(0, 1): 0, (1, 2): 0, (2, 3): 0, (3, 0): 0}   # boundary edge -> field inside
    fresh = 4
    for n in range(1, depth + 1):
        gap = {v: q - deg[v] for v in boundary}
        assert all(g >= 1 for g in gap.values())
        # Start the ring at a corner that is missing at least two squares.
        i0 = next(i for i, v in enumerate(boundary) if gap[v] >= 2)
        ring = boundary[i0:] + boundary[:i0]
        m = len(ring)
        # Runs: stretches of boundary edges that one new field covers, split at corners missing
        # two or more squares.
        cut = [i for i in range(m) if gap[ring[i]] >= 2]
        runs = []
        for a, b in zip(cut, cut[1:] + [cut[0] + m]):
            runs.append([ring[i % m] for i in range(a, b + 1)])
        assert len(runs) >= 2 and all(len(r) - 1 <= 2 for r in runs), 'unexpected notch'
        # Spoke ends: at each corner between runs, the new edge leaving it on the run before
        # (prev) and on the run after (after). One shared corner when the gap is two.
        uf = {}

        def root(x):
            while uf[x] != x:
                uf[x] = uf[uf[x]]
                x = uf[x]
            return x
        prev, after = [], []
        for run in runs:
            w = run[0]
            after_id = fresh
            fresh += 1
            uf[after_id] = after_id
            if gap[w] == 2:
                prev_id = after_id
            else:
                prev_id = fresh
                fresh += 1
                uf[prev_id] = prev_id
            prev.append(prev_id)
            after.append(after_id)
        for j, run in enumerate(runs):              # a field covering two edges: one far corner
            if len(run) == 3:
                a, b = root(after[j]), root(prev[(j + 1) % len(runs)])
                uf[b] = a
        new_edges = {}
        for j, run in enumerate(runs):
            s0, s1 = root(after[j]), root(prev[(j + 1) % len(runs)])
            if len(run) == 2:
                u0, u1 = run
                tile = (u1, u0, s0, s1)
                parent = (edge_tile[u0, u1],)
            else:
                u0, u1, u2 = run
                assert s0 == s1
                tile = (u2, u1, u0, s0)
                parent = (edge_tile[u0, u1], edge_tile[u1, u2])
            t = len(tiles)
            tiles.append(tile)
            layer.append(n)
            parents.append(parent)
            for v in tile:
                deg[v] = deg.get(v, 0) + 1
            for k in range(4):                      # its edges, seen from outside, for next time
                new_edges[tile[k], tile[(k + 1) % 4]] = t
        # The new boundary, in order.
        out = []
        for j, run in enumerate(runs):
            w = run[0]
            p, a = root(prev[j]), root(after[j])
            piece = [p, w, a] if gap[w] >= 3 else [a]
            for v in piece:
                if not out or out[-1] != v:
                    out.append(v)
        if out[0] == out[-1]:
            out.pop()
        assert len(set(out)) == len(out), 'the boundary touches itself'
        boundary = out
        # Every edge of the new boundary belongs to a field made in this ring. The patch is on
        # the boundary's left and each field's inside on the left of its own edges, so the
        # field runs along the boundary edge the same way.
        edge_tile = {}
        for k in range(len(boundary)):
            u, v = boundary[k], boundary[(k + 1) % len(boundary)]
            edge_tile[u, v] = new_edges[u, v]
    return tiles, layer, parents


def shells(depth):
    """How many fields are exactly n steps from the start, for n = 0 .. depth."""
    _, layer, _ = rings(depth)
    counts = [0] * (depth + 1)
    for n in layer:
        counts[n] += 1
    return counts


def flat_ball(r):
    """Fields within r steps on a flat floor: 1 + 4 + 8 + ... + 4r = 2r^2 + 2r + 1."""
    return 2 * r * r + 2 * r + 1


# ---- The usual way: coordinates and rounding --------------------------------------------------

def float_layers(depth, digits=6, q=Q):
    """The same floor from coordinates. Each field's centre is a point on the hyperboloid
    x^2 + y^2 - z^2 = -1, and stepping to a neighbour is a fixed 3 x 3 matrix. Fields are
    recognised by their centres rounded to `digits` places. Returns the shell counts, and the
    matrices of the fields (for drawing)."""
    a = math.acosh(math.cos(math.pi / q) / math.sin(math.pi / 4))   # centre to middle of a side
    steps = [mat_mul(rot(k * math.pi / 2), mat_mul(boost(2 * a), rot(math.pi))) for k in range(4)]
    start = ((1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0))
    seen = {key(start, digits)}
    front, counts, fields = [start], [1], [start]
    for _ in range(depth):
        nxt = []
        for m in front:
            for s in steps:
                n = mat_mul(m, s)
                k = key(n, digits)
                if k not in seen:
                    seen.add(k)
                    nxt.append(n)
        counts.append(len(nxt))
        fields.extend(nxt)
        front = nxt
    return counts, fields


def rot(t):
    c, s = math.cos(t), math.sin(t)
    return ((c, -s, 0.0), (s, c, 0.0), (0.0, 0.0, 1.0))


def boost(t):
    c, s = math.cosh(t), math.sinh(t)
    return ((c, 0.0, s), (0.0, 1.0, 0.0), (s, 0.0, c))


def mat_mul(a, b):
    return tuple(tuple(sum(a[i][k] * b[k][j] for k in range(3)) for j in range(3)) for i in range(3))


def key(m, digits):
    return round(m[0][2], digits), round(m[1][2], digits)


def centre(m):
    """A field's centre in the Poincare disc: the hyperboloid point (x, y, z) goes to
    (x, y) / (1 + z)."""
    x, y, z = m[0][2], m[1][2], m[2][2]
    return x / (1 + z), y / (1 + z)