"""Benchmark the Physarum flow solver against classical graph algorithms on a
4-connected grid: shortest path (vs breadth-first search), Steiner network (vs
the minimum spanning tree of the terminals and a Dijkstra shortest-path tree),
and fault tolerance (vs those trees). Classical baselines are pure Python.

    python benchmark.py
"""
import torch, heapq, collections, time

try:
    from kernels import get_kernel
    physarum = get_kernel("phanerozoic/physarum", version=1, trust_remote_code=True)
except Exception:
    import load_local
    physarum = load_local.load()

torch.set_grad_enabled(False)


def open_grid(H, W):
    m = torch.ones(H, W); m[0] = 0; m[-1] = 0; m[:, 0] = 0; m[:, -1] = 0
    return m

def bfs_len(net, a, b):
    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 dijkstra(mask, src):
    H, W = mask.shape; INF = 1 << 30; dist = {src: 0}; prev = {}; pq = [(0, src)]
    while pq:
        d, u = heapq.heappop(pq)
        if d > dist.get(u, INF): continue
        x, y = u
        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 mask[ny, nx]:
                v = (nx, ny)
                if d + 1 < dist.get(v, INF): dist[v] = d + 1; prev[v] = u; heapq.heappush(pq, (d + 1, v))
    return dist, prev

def sp_tree(mask, terms):                                # Dijkstra shortest-path tree
    _, prev = dijkstra(mask, terms[0]); E = set()
    for t in terms[1:]:
        cur = t
        while cur != terms[0] and cur in prev:
            p = prev[cur]; E.add((cur, p) if cur < p else (p, cur)); cur = p
    return E

def mst_net(mask, terms):                                # MST of terminals, drawn on the grid
    D = {}; P = {}
    for t in terms: D[t], P[t] = dijkstra(mask, t)
    n = len(terms); used = {0}; picked = []
    while len(used) < n:
        best = None
        for i in used:
            for j in range(n):
                if j not in used:
                    d = D[terms[i]].get(terms[j], 1 << 30)
                    if best is None or d < best[0]: best = (d, i, j)
        _, i, j = best; used.add(j); picked.append((i, j))
    E = set()
    for i, j in picked:
        a = terms[i]; cur = terms[j]
        while cur != a and cur in P[a]:
            p = P[a][cur]; E.add((cur, p) if cur < p else (p, cur)); cur = p
    return E

def net_edges(cE, cS, thr):
    E = set()
    for y, x in (cE[:, :-1] > thr).nonzero().tolist(): E.add(((x, y), (x + 1, y)))
    for y, x in (cS[:-1, :] > thr).nonzero().tolist(): E.add(((x, y), (x, y + 1)))
    return E

def connected(E, a, b):
    adj = collections.defaultdict(list)
    for u, v in E: adj[u].append(v); adj[v].append(u)
    seen = {a}; q = collections.deque([a])
    while q:
        u = q.popleft()
        if u == b: return True
        for v in adj[u]:
            if v not in seen: seen.add(v); q.append(v)
    return a == b

def all_conn(E, terms): return all(connected(E, terms[0], t) for t in terms[1:])

def edge_conn(E, a, b):                                  # min edge cut = max edge-disjoint paths
    adj = collections.defaultdict(dict)
    for u, v in E: adj[u][v] = adj[u].get(v, 0) + 1; adj[v][u] = adj[v].get(u, 0) + 1
    f = 0
    while True:
        prev = {a: None}; q = collections.deque([a]); ok = False
        while q:
            u = q.popleft()
            if u == b: ok = True; break
            for v, c in adj[u].items():
                if c > 0 and v not in prev: prev[v] = u; q.append(v)
        if not ok: break
        v = b
        while v != a: u = prev[v]; adj[u][v] -= 1; adj[v][u] += 1; v = u
        f += 1
    return f

def min_cut(E, terms): return min(edge_conn(E, terms[0], t) for t in terms[1:])

def backbone(cE, cS, terms):                             # leanest connected network
    for thr in [x / 1000 for x in range(950, 3, -4)]:
        E = net_edges(cE, cS, thr)
        if all_conn(E, terms): return E
    return None

def leanest_2conn(cE, cS, terms):                        # leanest network surviving any single cut
    for thr in [x / 100 for x in range(95, 4, -1)]:
        E = net_edges(cE, cS, thr)
        if all_conn(E, terms) and min_cut(E, terms) >= 2: return E
    return None


import random

print("physarum flow solver vs classical graph algorithms\n")

# 1. shortest path: the flow model converges to the exact shortest path.
mk = physarum.maze(127, 127, 5); src, goal = (2, 2), (124, 124)
t = time.time(); f = physarum.PhysarumFlow(mk).solve([src, goal], [1, -1], iters=200, cg_iters=120); ts = time.time() - t
plen = len(f.path(src, goal)) - 1; opt = bfs_len(mk.bool(), src, goal)
print(f"shortest path (127x127 maze): physarum {plen} vs BFS optimum {opt}  "
      f"[{'exact' if plen == opt else 'MISS'}, {ts:.1f}s]  (Dijkstra is far faster in time)")

# 2. connect N terminals: a heuristic, comparable to the MST across configs, not
# consistently better. Averaged over a few random terminal sets.
print("\nconnect N terminals -- physarum network length / MST length:")
mask = open_grid(90, 90); ratios = []
for seed in range(4):
    rng = random.Random(seed); terms = []
    while len(terms) < 7:
        p = (rng.randint(6, 83), rng.randint(6, 83))
        if all(abs(p[0] - q[0]) + abs(p[1] - q[1]) > 9 for q in terms): terms.append(p)
    Emst = mst_net(mask, terms)
    f = physarum.PhysarumFlow(mask).solve(terms, [6.0] + [-1.0] * 6, iters=300)
    Est = backbone(f.cE, f.cS, terms)
    if Est:
        r = len(Est) / len(Emst); ratios.append(r)
        print(f"  config {seed}: physarum {len(Est):4d}  MST {len(Emst):4d}  ratio {r:.2f}")
if ratios:
    print(f"  mean ratio {sum(ratios) / len(ratios):.2f} (>1 means longer than the MST)")

# 3. solve_robust yields a network that survives any single edge cut (min-cut 2) --
# a topology a tree lacks, though longer than a good classical 2-edge-connected
# construction. See PhysarumFlow.solve_robust.