"""Checks and a report for Part 7: five squares at every corner.

Run from this folder: python test_hyperbolic.py            (about a minute)
                      python test_hyperbolic.py --slow     (adds the 15-ring float comparison)
"""
import math
import random
import statistics
import sys
from collections import deque

import berlekamp_massey
import bidirectional
import box
import growth
import hyperbolic
from surface import Surface, corner_counts, depth_first, euler_characteristic, simulate


def test_rings(depth=9):
    """The exact floor: every corner inside it has five squares, every edge at most two fields,
    it is one disc, and a field's ring is its distance from the start."""
    tiles, layer, parents = hyperbolic.rings(depth)
    s = Surface(tiles)
    counts = corner_counts(s)
    inside = [n for n, arcs in counts.values() if not arcs]
    assert set(inside) == {5}
    assert euler_characteristic(s) == 1
    dist = {0: 0}
    queue = deque([0])
    while queue:
        f = queue.popleft()
        for g in s.glue[f]:
            if g is not None and g[0] not in dist:
                dist[g[0]] = dist[f] + 1
                queue.append(g[0])
    assert all(dist[f] == layer[f] for f in range(len(tiles)))
    for f in range(1, len(tiles)):
        assert 1 <= len(parents[f]) <= 2
        nbrs = {g[0] for g in s.glue[f] if g is not None}
        assert set(parents[f]) == {g for g in nbrs if layer[g] == layer[f] - 1}
    print(f'rings ................... ok ({len(tiles):,} fields to ring {depth}, '
          f'{len(inside):,} corners inside, all with five squares)')


def test_floats():
    exact = hyperbolic.shells(10)
    floats, _ = hyperbolic.float_layers(10, digits=6)
    assert floats == exact
    print('floats to ring 10 ....... ok (rounded to 6 places, they agree)')


def test_matrix():
    """The four rules, the matrix, the recurrence and the growth rate all agree."""
    degrees = growth.boundary_degrees(12)
    v = growth.START
    for n in range(12):
        assert degrees[n] == v
        v = growth.step(v)
    exact = hyperbolic.shells(14)
    assert [growth.shell(r) for r in range(15)] == exact
    assert growth.characteristic() == [1, -2, 0, -2, 1]
    rec = berlekamp_massey.berlekamp_massey(exact[1:10])     # nine shells, from ring 1
    assert rec == [2, 0, 2, -1]
    assert berlekamp_massey.extend(exact[1:10], rec, 14) == exact[1:15]
    lam = growth.growth_rate()
    assert abs(lam ** 4 - 2 * lam ** 3 - 2 * lam + 1) < 1e-12
    # The boundary's mix of corner types settles to the eigenvector for lambda.
    v = growth.START
    for _ in range(12):
        v = growth.step(v)
    shares = [x / sum(v) for x in v]
    assert all(abs(a - b) < 1e-4 for a, b in zip(shares, (0.58563, 0.11103, 0.25500, 0.04834)))
    assert v == (70460, 13360, 30680, 5816)
    assert abs(growth.shell(60) / growth.shell(59) - lam) < 1e-12
    print(f'the matrix .............. ok (rules hold on 12 rings; recurrence s(n) = 2s(n-1) + '
          f'2s(n-3) - s(n-4); growth {lam:.6f} per step)')


def test_robot(rng):
    """Part 1's walk cleans a round room on the hyperbolic floor, with furniture."""
    for _ in range(40):
        depth = rng.randint(2, 6)
        tiles, _, _ = hyperbolic.rings(depth)
        blocked = [f for f in range(1, len(tiles)) if rng.random() < 0.3]
        s = Surface(tiles, blocked)
        frame = (rng.randrange(4), rng.random() < 0.5)
        cmds = depth_first(s, 0, frame)
        cleaned, _, _ = simulate(s, 0, cmds, frame)
        reach = reachable(s, 0)
        assert cleaned == reach and len(cmds) == 2 * (len(reach) - 1)
    print('the walk ................ ok (40 round rooms with furniture)')


def reachable(s, start):
    seen, todo = {start}, [start]
    while todo:
        f = todo.pop()
        for g in s.glue[f]:
            if g is not None and g[0] not in seen:
                seen.add(g[0])
                todo.append(g[0])
    return seen


def test_corner_loops():
    """Drive one field, turn left: on this floor every corner has five squares, so every such
    loop closes after five stretches, with the robot a quarter turn round."""
    tiles, layer, _ = hyperbolic.rings(6)
    s = Surface(tiles)
    count = 0
    for f in range(len(tiles)):
        if layer[f] <= 3:
            for h in '^>v<':
                assert box.square_loop(s, f, 1, h, most=12) == (5, (1, False))
                count += 1
    print(f'five left turns ......... ok ({count} loops, every one home after five stretches)')


def test_search(rng):
    tiles, layer, _ = hyperbolic.rings(11)
    nbrs = bidirectional.tile_neighbours(tiles)
    for _ in range(100):
        d = rng.randint(0, 7)
        t = rng.choice([i for i, l in enumerate(layer) if l == d])
        s = rng.choice([i for i, l in enumerate(layer) if l <= 3])
        a, _ = bidirectional.one_sided(lambda x: nbrs[x], s, t)
        b, _ = bidirectional.two_sided(lambda x: nbrs[x], s, t)
        assert a == b
    print('search from both ends ... ok (100 pairs: same distance as plain search)')


def report_counts():
    print('\nFields within r steps:')
    print('   r   flat floor   five squares at a corner   share in the outer ring')
    for r in (1, 2, 3, 5, 8, 10, 11, 12, 15, 20):
        b = growth.ball(r)
        print(f'  {r:>2}   {hyperbolic.flat_ball(r):>10,}   {b:>24,}   {growth.shell(r) / b:.3f}')
    lam = growth.growth_rate()
    print(f'  outer ring share tends to 1 - 1/{lam:.4f} = {1 - 1 / lam:.4f} (flat: 0)')
    budget = max(r for r in range(40) if 2 * (growth.ball(r) - 1) <= 50_000)
    flat = max(r for r in range(400) if 2 * (hyperbolic.flat_ball(r) - 1) <= 50_000)
    print(f'  Part 1\'s walk fits 50,000 commands up to radius {budget} here '
          f'({2 * (growth.ball(budget) - 1):,} commands), {flat} on a flat floor')
    print(f'  ring 1000 has a {len(str(growth.shell(1000)))}-digit number of fields')


def search_means(d, tiles=None, layer=None):
    """Fields expanded, one side and both sides, averaged over every goal exactly d steps from
    the start: on the hyperbolic floor (the start is field 0) and on a flat floor."""
    if tiles is None:
        tiles, layer, _ = hyperbolic.rings(max(13, d + d // 2 + 1))
    nbrs = bidirectional.tile_neighbours(tiles)
    goals = [i for i, l in enumerate(layer) if l == d]
    flat = [(a, d - abs(a)) for a in range(-d, d + 1)] + [(a, abs(a) - d) for a in range(-d + 1, d)]
    h1 = statistics.mean(bidirectional.one_sided(lambda x: nbrs[x], 0, t)[1] for t in goals)
    h2 = statistics.mean(bidirectional.two_sided(lambda x: nbrs[x], 0, t)[1] for t in goals)
    f1 = statistics.mean(bidirectional.one_sided(bidirectional.grid_neighbours, (0, 0), g)[1] for g in flat)
    f2 = statistics.mean(bidirectional.two_sided(bidirectional.grid_neighbours, (0, 0), g)[1] for g in flat)
    return h1, h2, f1, f2, len(goals), len(flat)


def report_search():
    tiles, layer, _ = hyperbolic.rings(13)
    print('\nFields expanded to find a goal d steps away (mean over every goal at that distance):')
    for d in range(2, 9):
        h1, h2, f1, f2, nh, nf = search_means(d, tiles, layer)
        print(f'  d = {d}: hyperbolic ({nh:>4} goals) {h1:>7.1f} vs {h2:>5.1f} ({h1 / h2:>5.2f}x)   '
              f'flat ({nf:>2} goals) {f1:>6.1f} vs {f2:>5.1f} ({f1 / f2:.2f}x)')


def report_floats(rng):
    """Out along the floor and straight back, in floating point: how far from home is 'home'?"""
    a = math.acosh(math.cos(math.pi / 5) / math.sin(math.pi / 4))
    out = [hyperbolic.mat_mul(hyperbolic.rot(k * math.pi / 2),
                              hyperbolic.mat_mul(hyperbolic.boost(2 * a), hyperbolic.rot(math.pi)))
           for k in range(4)]
    back = [hyperbolic.mat_mul(hyperbolic.rot(-math.pi),
                               hyperbolic.mat_mul(hyperbolic.boost(-2 * a), hyperbolic.rot(-k * math.pi / 2)))
            for k in range(4)]
    print('\nOut n steps (always away from the start) and back again, in double precision:')
    for n in (5, 10, 15, 20, 30):
        worst = 0
        for _ in range(20):
            m, word = ((1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0)), []
            for _ in range(n):
                k = max(range(4), key=lambda j: hyperbolic.mat_mul(m, out[j])[2][2] + rng.random() * 1e-9)
                word.append(k)
                m = hyperbolic.mat_mul(m, out[k])
            for k in reversed(word):
                m = hyperbolic.mat_mul(m, back[k])
            worst = max(worst, max(abs(m[i][j] - (i == j)) for i in range(3) for j in range(3)))
        print(f'  n = {n:>2}: largest error on the way home {worst:.2g}')


def report_slow():
    exact = hyperbolic.shells(15)
    for digits in (6, 9):
        floats, _ = hyperbolic.float_layers(15, digits=digits)
        wrong = [(r, f, e) for r, (f, e) in enumerate(zip(floats, exact)) if f != e]
        r, f, e = wrong[0]
        print(f'  rounded to {digits} places: first wrong at ring {r} ({f:,} fields instead of {e:,})')


if __name__ == '__main__':
    rng = random.Random(7)
    test_rings()
    test_floats()
    test_matrix()
    test_robot(rng)
    test_corner_loops()
    test_search(rng)
    report_counts()
    report_search()
    report_floats(rng)
    if '--slow' in sys.argv:
        print('\nThe float construction against the exact one, to ring 15:')
        report_slow()