RITHESH LED Display

"""Stage R2b: finish the connections Freerouting left open, with a small grid maze router.

For every net whose copper is still in more than one piece ("island"), the nearest two pieces are joined
by an A* search on a 0.1 mm grid (front layer preferred; back layer costs 4x, each via 2.5 mm extra).
Clearances are the same as the rest of the board (0.25 mm to power nets, 0.2 mm otherwise). The back layer
stays out of the USB-pair corridor and from under the ESP32; the front stays out of the band along the
bottom edge of the +5V fill. Copper fills are ignored while routing (they are re-poured afterwards) but
count as copper when deciding which pieces are connected.

Usage: python3 route_fix.py [net ...]      (default: every net with open connections except +5V/GND fills)
"""
import heapq
import math
import os
import sys

import numpy as np
import pcbnew
import shapely
from shapely.geometry import LineString, Point, Polygon, box
from shapely.ops import unary_union

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import make_board as mb  # noqa: E402
import route_board as rb  # noqa: E402

GRID = 0.1
BACK_COST, VIA_COST = 4.0, 2.5
WIDTH = {"+3V3": 0.4, "VBUS": 0.5, "+5V": 0.5, "GND": 0.4}


def zone_polys(board, netname):
    out = []
    for z in board.Zones():
        if z.GetIsRuleArea() or str(z.GetNetname()) != netname:
            continue
        for layer, ly in ((rb.F, "F"), (rb.B, "B")):
            if not z.IsOnLayer(layer):
                continue
            fp = z.GetFilledPolysList(layer)
            for i in range(fp.OutlineCount()):
                o = fp.Outline(i)
                outer = [rb.xy(o.CPoint(k)) for k in range(o.PointCount())]
                holes = [[rb.xy(fp.Hole(i, h).CPoint(k)) for k in range(fp.Hole(i, h).PointCount())]
                         for h in range(fp.HoleCount(i))]
                out.append((Polygon(outer, holes).buffer(0), {ly}))
    return out


def islands(cu, board, netname):
    """Connected pieces of a net: list of (geoms_by_layer {'F': geom, 'B': geom})."""
    items = [(g, ly) for (g, n, ly, kind) in cu.items if str(n) == netname and kind != "hole"]
    items += zone_polys(board, netname)
    n = len(items)
    parent = list(range(n))

    def find(a):
        while parent[a] != a:
            parent[a] = parent[parent[a]]
            a = parent[a]
        return a
    from shapely.strtree import STRtree
    tree = STRtree([g for g, _ in items])
    for i, (g, ly) in enumerate(items):
        for j in tree.query(g):
            if j <= i:
                continue
            if (items[j][1] & ly) and g.intersects(items[j][0]):
                parent[find(i)] = find(j)
    groups = {}
    for i in range(n):
        groups.setdefault(find(i), []).append(i)
    out = []
    for idxs in groups.values():
        by = {"F": [], "B": []}
        for i in idxs:
            for ly in items[i][1]:
                by[ly].append(items[i][0])
        out.append({k: unary_union(v) if v else None for k, v in by.items()})
    return out


class Maze:
    def __init__(self, cu, netname, width, window):
        self.cu, self.net, self.w = cu, netname, width
        x0, y0, x1, y1 = window
        self.x0, self.y0 = x0, y0
        self.nx, self.ny = int((x1 - x0) / GRID) + 1, int((y1 - y0) / GRID) + 1
        xs = x0 + np.arange(self.nx) * GRID
        ys = y0 + np.arange(self.ny) * GRID
        self.X, self.Y = np.meshgrid(xs, ys)
        self.block = {"F": np.zeros(self.X.shape, bool), "B": np.zeros(self.X.shape, bool)}
        self.vblock = np.zeros(self.X.shape, bool)
        win = box(x0, y0, x1, y1)
        for idx in cu.tree().query(win.buffer(1.0)):
            g, n, ly, kind = cu.items[idx]
            if kind != "hole" and str(n) == netname:
                continue
            clr = 0.25 if kind == "hole" else rb.clearance(netname, n)
            for L in ly:
                self._mark(self.block[L], g, clr + width / 2 + 0.03)
            self._mark(self.vblock, g, clr + rb.VIA_D / 2 + 0.03 if kind != "hole" else 0.25 + 0.45)
        for poly, ly in cu.keepouts:
            for L in ly:
                self._mark(self.block[L], poly, width / 2 + 0.02)
            self._mark(self.vblock, poly, rb.VIA_D / 2 + 0.02)
        for (a, b, c, d) in rb.BACK_KEEPOUT:
            self._mark(self.block["B"], box(a, b, c, d), width / 2)
            self._mark(self.vblock, box(a, b, c, d), rb.VIA_D / 2)
        for (a, b, c, d) in rb.FRONT_KEEPOUT[1:]:
            self._mark(self.block["F"], box(a, b, c, d), width / 2)
        inner = box(rb.EDGE_CLR + width / 2, rb.EDGE_CLR + width / 2, mb.W - rb.EDGE_CLR - width / 2,
                    mb.H - rb.EDGE_CLR - width / 2)
        out = ~shapely.contains_xy(inner, self.X, self.Y)
        self.block["F"] |= out
        self.block["B"] |= out
        self.vblock |= out | self.block["F"] | self.block["B"]

    def _mark(self, grid, geom, dist):
        g = geom.buffer(dist, 8)
        a, b, c, d = g.bounds
        j0, j1 = max(int((a - self.x0) / GRID), 0), min(int((c - self.x0) / GRID) + 1, self.nx)
        i0, i1 = max(int((b - self.y0) / GRID), 0), min(int((d - self.y0) / GRID) + 1, self.ny)
        if j0 >= j1 or i0 >= i1:
            return
        grid[i0:i1, j0:j1] |= shapely.contains_xy(g, self.X[i0:i1, j0:j1], self.Y[i0:i1, j0:j1])

    def cells_in(self, geom):
        if geom is None:
            return np.zeros(self.X.shape, bool)
        return shapely.contains_xy(geom.buffer(-0.02), self.X, self.Y)

    def search(self, src, dst, allow_back=True):
        """src/dst: {'F': bool grid, 'B': bool grid}. Returns list of (layer, i, j)."""
        L = ("F", "B") if allow_back else ("F",)
        free = {l: ~self.block[l] for l in L}
        for l in L:
            free[l] |= src[l] | dst[l]
        goal = dst["F"] | (dst["B"] if allow_back else False)
        if not goal.any():
            return None
        from scipy.ndimage import distance_transform_edt
        hmap = distance_transform_edt(~goal) * GRID * 0.95       # admissible straight-line estimate

        def h(i, j):
            return hmap[i, j]
        dist, prev, pq = {}, {}, []
        for l in L:
            si, sj = np.nonzero(src[l])
            for i, j in zip(si, sj):
                dist[(l, i, j)] = 0.0
                heapq.heappush(pq, (h(i, j), 0.0, (l, i, j)))
        moves = [(0, 1, 1), (1, 0, 1), (0, -1, 1), (-1, 0, 1), (1, 1, 1.414), (1, -1, 1.414), (-1, 1, 1.414),
                 (-1, -1, 1.414)]
        n = 0
        while pq:
            f, d, s = heapq.heappop(pq)
            if d > dist.get(s, 1e18):
                continue
            l, i, j = s
            if dst[l][i, j]:
                path = [s]
                while s in prev:
                    s = prev[s]
                    path.append(s)
                return path[::-1]
            n += 1
            if n > 1500000:
                return None
            cost = 1.0 if l == "F" else BACK_COST
            for di, dj, m in moves:
                ii, jj = i + di, j + dj
                if not (0 <= ii < self.ny and 0 <= jj < self.nx) or not free[l][ii, jj]:
                    continue
                if di and dj and not (free[l][i, jj] and free[l][ii, j]):
                    continue
                nd = d + m * GRID * cost
                t = (l, ii, jj)
                if nd < dist.get(t, 1e18):
                    dist[t] = nd
                    prev[t] = s
                    heapq.heappush(pq, (nd + h(ii, jj), nd, t))
            if allow_back and not self.vblock[i, j]:
                o = "B" if l == "F" else "F"
                if free[o][i, j]:
                    nd = d + VIA_COST
                    t = (o, i, j)
                    if nd < dist.get(t, 1e18):
                        dist[t] = nd
                        prev[t] = s
                        heapq.heappush(pq, (nd + h(i, j), nd, t))
        return None

    def xy(self, i, j):
        return (round(float(self.x0 + j * GRID), 3), round(float(self.y0 + i * GRID), 3))


def simplify(pts):
    out = [pts[0]]
    for k in range(1, len(pts) - 1):
        (ax, ay), (bx, by), (cx, cy) = out[-1], pts[k], pts[k + 1]
        if abs((bx - ax) * (cy - by) - (by - ay) * (cx - bx)) > 1e-9:
            out.append(pts[k])
    out.append(pts[-1])
    return out


def join(R, board, netname, margin=6.0):
    isl = islands(R.cu, board, netname)
    if len(isl) < 2:
        return 0, []
    full = str(R.net(netname).GetNetname())
    width = WIDTH.get(rb.short(full), 0.25)
    joined, failed = 0, []
    while len(isl) > 1:
        # nearest pair of islands (front/back geometry combined)
        geo = [unary_union([g for g in (d["F"], d["B"]) if g is not None]) for d in isl]
        best = None
        for a in range(len(isl)):
            for b in range(a + 1, len(isl)):
                dd = geo[a].distance(geo[b])
                if best is None or dd < best[0]:
                    best = (dd, a, b)
        _, a, b = best
        from shapely.ops import nearest_points
        pa, pb = nearest_points(geo[a], geo[b])
        ok = False
        for m in (margin, margin * 2.5, margin * 6):
            win = (max(min(pa.x, pb.x) - m, 0), max(min(pa.y, pb.y) - m, 0),
                   min(max(pa.x, pb.x) + m, mb.W), min(max(pa.y, pb.y) + m, mb.H))
            z = Maze(R.cu, full, width, win)
            src = {k: z.cells_in(isl[a][k]) for k in ("F", "B")}
            dst = {k: z.cells_in(isl[b][k]) for k in ("F", "B")}
            path = z.search(src, dst)
            if path:
                ok = True
                break
        if not ok:
            failed.append(f"{rb.short(full)}: pieces near ({pa.x:.1f},{pa.y:.1f}) and ({pb.x:.1f},{pb.y:.1f})")
            isl.pop(b)
            continue
        runs, cur = [], [path[0]]
        for s in path[1:]:
            if s[0] != cur[-1][0]:
                runs.append(cur)
                cur = [s]
            else:
                cur.append(s)
        runs.append(cur)
        for k, run in enumerate(runs):
            pts = simplify([z.xy(i, j) for (_, i, j) in run])
            if k > 0:
                R.via(*pts[0], full, check=False)
            if len(pts) >= 2:
                R.track(pts, width, full, layer=rb.F if run[0][0] == "F" else rb.B, check=False)
        joined += 1
        isl = islands(R.cu, board, netname)
    return joined, failed


def main():
    board = pcbnew.LoadBoard(mb.PCB)
    R = rb.Router(board)
    nets = sys.argv[1:]
    if not nets:
        board.BuildConnectivity()
        nets = sorted({str(n) for n in board.GetNetsByName().keys() if str(n)})
    report = []
    for n in nets:
        if n.startswith("unconnected-") or n == "":
            continue
        j, f = join(R, board, n)
        if j or f:
            report.append((n, j, f))
            print(f"{rb.short(n)}: joined {j} gap(s)" + (f"; FAILED {f}" if f else ""), flush=True)
    pcbnew.SaveBoard(mb.PCB, board)
    print("saved; added", R.count)


if __name__ == "__main__":
    main()