import torch
import kernels

physarum = kernels.get_kernel("phanerozoic/physarum", version=1, trust_remote_code=True)


def test_network_develops():
    sim = physarum.Physarum(width=256, height=256, agents=20000, seed=1)
    sim.step(120)
    t = sim.trail
    assert torch.isfinite(t).all()
    assert float(t.max()) > 0.0
    # a network is high-contrast, not a uniform wash
    assert float(t.std()) > 0.4 * float(t.mean())


def test_image_shape():
    sim = physarum.Physarum(width=128, height=128, agents=5000, seed=2)
    sim.step(60)
    img = sim.image()
    assert tuple(img.shape) == (128, 128, 3)
    assert img.dtype == torch.uint8


def test_deterministic():
    a = physarum.Physarum(width=128, height=128, agents=5000, seed=7); a.step(40)
    b = physarum.Physarum(width=128, height=128, agents=5000, seed=7); b.step(40)
    assert torch.equal(a.trail, b.trail)


def _bfs_len(net, a, b):
    import collections
    H, W = net.shape
    seen = torch.zeros(H, W, dtype=torch.bool); seen[a[1], a[0]] = True
    q = collections.deque([(a[0], a[1], 0)])
    while q:
        x, y, d = q.popleft()
        if (x, y) == b:
            return d
        for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            nx, ny = x + dx, y + dy
            if 0 <= nx < W and 0 <= ny < H and bool(net[ny, nx]) and not seen[ny, nx]:
                seen[ny, nx] = True; q.append((nx, ny, d + 1))
    return None


def test_flow_solves_maze():
    mk = physarum.maze(65, 65, seed=3)
    src, goal = (2, 2), (62, 62)
    f = physarum.PhysarumFlow(mk).solve([src, goal], [1, -1], iters=150, cg_iters=100)
    path = f.path(src, goal)
    assert path is not None                                  # terminals connected
    assert len(path) - 1 == _bfs_len(mk.bool(), src, goal)   # equals the shortest path


def test_flow_network_connects():
    open_mask = torch.ones(80, 80); open_mask[0] = 0; open_mask[-1] = 0
    open_mask[:, 0] = 0; open_mask[:, -1] = 0
    terms = [(40, 40), (12, 12), (68, 12), (12, 68), (68, 68)]
    f = physarum.PhysarumFlow(open_mask).solve(terms, [4, -1, -1, -1, -1], iters=200)
    assert all(f.path(terms[0], t) is not None for t in terms[1:])