AM/FM Receiver Molecule

"""Simulated-annealing placer for footprints inside rectangular regions.

Geometry is in mm, board coordinates (+y down).  Rotation follows KiCad:
positive angle = counter-clockwise on screen: (x, y) -> (x cos + y sin, -x sin + y cos).
"""
import math, random


def rot_pt(x, y, deg):
    d = deg % 360
    if d == 0:
        return x, y
    if d == 90:
        return y, -x
    if d == 180:
        return -x, -y
    if d == 270:
        return -y, x
    a = math.radians(deg)
    return x * math.cos(a) + y * math.sin(a), -x * math.sin(a) + y * math.cos(a)


class Item:
    def __init__(self, ref, crt, pads, region, rotations=(0, 90, 180, 270), fixed=None, margin=0.25):
        self.ref = ref
        self.crt = crt              # courtyard box at rot 0 relative to origin (x0, y0, x1, y1)
        self.pads = pads            # list of (net, lx, ly)
        self.region = region        # (x0, y0, x1, y1)
        self.rotations = rotations
        self.fixed = fixed          # (x, y, rot) or None
        self.margin = margin
        self.lattice = None         # (x0, y0, pitch): position must stay on this grid
        self.pref = None            # preferred orientation (mod 180) for a tidy look, or None
        self.x = self.y = 0.0
        self.r = 0
        if fixed:
            self.x, self.y, self.r = fixed
        self._rel = {}
        for r in (0, 90, 180, 270):
            x0, y0, x1, y1 = crt
            pts = [rot_pt(px, py, r) for px, py in ((x0, y0), (x1, y0), (x0, y1), (x1, y1))]
            m = margin
            self._rel[r] = (min(p[0] for p in pts) - m, min(p[1] for p in pts) - m,
                            max(p[0] for p in pts) + m, max(p[1] for p in pts) + m)
        self._padrel = {r: [rot_pt(lx, ly, r) for (_, lx, ly) in pads] for r in (0, 90, 180, 270)}
        self.update()

    def snap(self):
        if self.lattice:
            x0, y0, p = self.lattice
            self.x = round(x0 + round((self.x - x0) / p) * p, 4)
            self.y = round(y0 + round((self.y - y0) / p) * p, 4)

    def rebuild(self):
        x0, y0, x1, y1 = self.crt
        for r in (0, 90, 180, 270):
            pts = [rot_pt(px, py, r) for px, py in ((x0, y0), (x1, y0), (x0, y1), (x1, y1))]
            m = self.margin
            self._rel[r] = (min(q[0] for q in pts) - m, min(q[1] for q in pts) - m,
                            max(q[0] for q in pts) + m, max(q[1] for q in pts) + m)

    def update(self):
        rb = self._rel[self.r % 360]
        self.b = (self.x + rb[0], self.y + rb[1], self.x + rb[2], self.y + rb[3])

    def box(self, x=None, y=None, r=None):
        x = self.x if x is None else x
        y = self.y if y is None else y
        r = self.r if r is None else r
        rb = self._rel[r % 360]
        return (x + rb[0], y + rb[1], x + rb[2], y + rb[3])

    def pad_pos(self, i):
        dx, dy = self._padrel[self.r % 360][i]
        return self.x + dx, self.y + dy


def overlap(a, b):
    w = min(a[2], b[2]) - max(a[0], b[0])
    if w <= 0:
        return 0.0
    h = min(a[3], b[3]) - max(a[1], b[1])
    if h <= 0:
        return 0.0
    return w * h


def outside(box, reg):
    x0, y0, x1, y1 = box
    rx0, ry0, rx1, ry1 = reg
    ex = max(0.0, rx0 - x0) + max(0.0, x1 - rx1)
    ey = max(0.0, ry0 - y0) + max(0.0, y1 - ry1)
    w, h = x1 - x0, y1 - y0
    return ex * h + ey * w + ex * ey


def grow(b, d):
    return (b[0] - d, b[1] - d, b[2] + d, b[3] + d)


class Placer:
    def __init__(self, items, net_weights, skip_nets, affinities, seed=1):
        self.items = {it.ref: it for it in items}
        self.net_weights = net_weights
        self.skip = set(skip_nets)
        self.aff = affinities         # list of (refA, padidxA, refB, padidxB, weight)
        self.rng = random.Random(seed)
        self.nets = {}
        for it in items:
            for i, (net, lx, ly) in enumerate(it.pads):
                if net is None or net in self.skip:
                    continue
                self.nets.setdefault(net, []).append((it, i))
        for n in list(self.nets):
            if len(self.nets[n]) < 2:
                del self.nets[n]
        self.item_nets = {r: [] for r in self.items}
        for n, pins in self.nets.items():
            for it in {id(p[0]): p[0] for p in pins}.values():
                self.item_nets[it.ref].append(n)
        self.item_aff = {r: [] for r in self.items}
        for k, (ra, ia, rb, ib, w) in enumerate(self.aff):
            self.item_aff[ra].append(k)
            self.item_aff[rb].append(k)
        self.W_OVL = 30.0
        self.W_OUT = 80.0
        self.W_ROT = 2.0            # cost of turning a part away from its preferred orientation

    def net_cost(self, n):
        pins = self.nets[n]
        w = self.net_weights.get(n, 1.0)
        if len(pins) <= 12:
            xs0 = ys0 = 1e9
            xs1 = ys1 = -1e9
            for it, i in pins:
                x, y = it.pad_pos(i)
                if x < xs0: xs0 = x
                if x > xs1: xs1 = x
                if y < ys0: ys0 = y
                if y > ys1: ys1 = y
            return w * ((xs1 - xs0) + (ys1 - ys0))
        pts = [it.pad_pos(i) for it, i in pins]
        s = 0.0
        for a in range(len(pts)):
            best = 1e9
            ax, ay = pts[a]
            for b in range(len(pts)):
                if a != b:
                    d = abs(ax - pts[b][0]) + abs(ay - pts[b][1])
                    if d < best:
                        best = d
            s += best
        return w * 0.5 * s

    def aff_cost(self, k):
        ra, ia, rb, ib, w = self.aff[k]
        pa = self.items[ra].pad_pos(ia)
        pb = self.items[rb].pad_pos(ib)
        return w * (abs(pa[0] - pb[0]) + abs(pa[1] - pb[1]))

    def geom_cost(self, it, neigh):
        b = it.b
        c = self.W_OUT * outside(b, it.region)
        if it.pref is not None and (it.r - it.pref) % 180:
            c += self.W_ROT
        W = self.W_OVL
        for o in neigh:
            if o is it:
                continue
            ob = o.b
            if ob[0] >= b[2] or ob[2] <= b[0] or ob[1] >= b[3] or ob[3] <= b[1]:
                continue
            c += W * overlap(b, ob)
        return c

    def anneal(self, refs, moves=40000, t0=None, t1=0.01, grid=0.25, log=None):
        movable = [self.items[r] for r in refs if self.items[r].fixed is None]
        if not movable:
            return 0
        reg = movable[0].region
        zone = grow((min(it.region[0] for it in movable), min(it.region[1] for it in movable),
                     max(it.region[2] for it in movable), max(it.region[3] for it in movable)), 3.0)
        neigh = [o for o in self.items.values() if overlap(o.b, zone) > 0 or o in movable]

        def local(its):
            nets = set()
            affs = set()
            for it in its:
                nets.update(self.item_nets[it.ref])
                affs.update(self.item_aff[it.ref])
            c = 0.0
            for n in sorted(nets):              # fixed order: float sums (and so the run) repeat exactly
                c += self.net_cost(n)
            for k in sorted(affs):
                c += self.aff_cost(k)
            for it in its:
                c += self.geom_cost(it, neigh)
            return c

        span = max(reg[2] - reg[0], reg[3] - reg[1])
        if t0 is None:
            deltas = []
            for _ in range(300):
                it = self.rng.choice(movable)
                before = local([it])
                ox, oy = it.x, it.y
                it.x += self.rng.uniform(-2, 2)
                it.y += self.rng.uniform(-2, 2)
                it.update()
                after = local([it])
                it.x, it.y = ox, oy
                it.update()
                deltas.append(abs(after - before))
            deltas.sort()
            t0 = max(0.5, deltas[len(deltas) // 2] * 1.5)
        alpha = (t1 / t0) ** (1.0 / max(1, moves))
        T = t0
        acc = 0
        base_w = self.W_OVL
        rng = self.rng
        for m in range(moves):
            frac = m / moves
            self.W_OVL = base_w * (1.0 + 9.0 * frac * frac)      # harden overlaps late
            rmax = max(0.4, span * 0.3 * (1 - frac) ** 1.5)
            it = rng.choice(movable)
            k = rng.random()
            if k < 0.12 and len(it.rotations) > 1:
                old = (it.x, it.y, it.r)
                before = local([it])
                it.r = rng.choice([r for r in it.rotations if r != it.r])
                it.update()
                d = local([it]) - before
                if d <= 0 or rng.random() < math.exp(-d / T):
                    acc += 1
                else:
                    it.x, it.y, it.r = old
                    it.update()
            elif k < 0.22:
                cands = [o for o in movable if o is not it and o.lattice == it.lattice]
                o = rng.choice(cands) if cands else None
                if o is None:
                    continue
                before = local([it, o])
                ix, iy, ox, oy = it.x, it.y, o.x, o.y
                it.x, it.y, o.x, o.y = ox, oy, ix, iy
                it.update(); o.update()
                d = local([it, o]) - before
                if d <= 0 or rng.random() < math.exp(-d / T):
                    acc += 1
                else:
                    it.x, it.y, o.x, o.y = ix, iy, ox, oy
                    it.update(); o.update()
            else:
                before = local([it])
                ox, oy = it.x, it.y
                it.x = round((it.x + rng.uniform(-rmax, rmax)) / grid) * grid
                it.y = round((it.y + rng.uniform(-rmax, rmax)) / grid) * grid
                it.snap()
                it.update()
                d = local([it]) - before
                if d <= 0 or rng.random() < math.exp(-d / T):
                    acc += 1
                else:
                    it.x, it.y = ox, oy
                    it.update()
            T *= alpha
        self.W_OVL = base_w
        return acc

    def overlaps(self):
        its = list(self.items.values())
        res = []
        for a in range(len(its)):
            for b in range(a + 1, len(its)):
                ov = overlap(its[a].b, its[b].b)
                if ov > 1e-6:
                    res.append((its[a].ref, its[b].ref, ov))
        return res

    def align(self, refs, tol=0.8, passes=3):
        """Tidy-up: pull parts whose centres almost line up (within tol) onto a common row / column."""
        its = [self.items[r] for r in refs if self.items[r].fixed is None and self.items[r].lattice
               and self.items[r].lattice[2] < 1.0]
        others = list(self.items.values())
        moved = 0

        def legal(it):
            if outside(it.b, it.region) > 1e-9:
                return False
            return all(o is it or overlap(it.b, o.b) <= 1e-9 for o in others)
        for _ in range(passes):
            for axis in (1, 0):                         # rows first, then columns
                key = (lambda it: it.y) if axis else (lambda it: it.x)
                srt = sorted(its, key=key)
                groups, cur = [], [srt[0]] if srt else []
                for it in srt[1:]:
                    if key(it) - key(cur[0]) <= tol:
                        cur.append(it)
                    else:
                        groups.append(cur)
                        cur = [it]
                if cur:
                    groups.append(cur)
                for g in groups:
                    if len(g) < 2:
                        continue
                    vals = sorted(key(it) for it in g)
                    target = vals[len(vals) // 2]
                    for it in g:
                        if abs(key(it) - target) < 1e-6:
                            continue
                        old = (it.x, it.y)
                        if axis:
                            it.y = target
                        else:
                            it.x = target
                        it.update()
                        if legal(it):
                            moved += 1
                        else:
                            it.x, it.y = old
                            it.update()
        return moved

    def legalize(self, step=0.25, max_radius=12.0):
        """Move each overlapping movable item to the nearest free spot (spiral search)."""
        its = list(self.items.values())
        # first on each part's lattice; parts on the 0.5 mm tidy grid that still overlap then use the finer step
        for fine in (False, True):
            if fine and not self.overlaps():
                break
            self._legalize_passes(its, step, max_radius, fine)
        return self.overlaps()

    def _legalize_passes(self, its, step, max_radius, fine):
        moved_any = True
        passes = 0
        while moved_any and passes < 6:
            moved_any = False
            passes += 1
            for it in its:
                if it.fixed is not None:
                    continue
                others = [o for o in its if o is not it]
                if all(overlap(it.b, o.b) <= 1e-9 for o in others) and outside(it.b, it.region) <= 1e-9:
                    continue
                best = None
                ox, oy = it.x, it.y
                r = step
                while r <= max_radius and best is None:
                    n = max(8, int(2 * math.pi * r / step))
                    cands = []
                    for k in range(n):
                        a = 2 * math.pi * k / n
                        cx = round((ox + r * math.cos(a)) / step) * step
                        cy = round((oy + r * math.sin(a)) / step) * step
                        if it.lattice and not (fine and it.lattice[2] < 1.0):
                            lx0, ly0, lp = it.lattice
                            cx = round(lx0 + round((cx - lx0) / lp) * lp, 4)
                            cy = round(ly0 + round((cy - ly0) / lp) * lp, 4)
                        b = it.box(cx, cy)
                        if outside(b, it.region) > 1e-9:
                            continue
                        if any(overlap(b, o.b) > 1e-9 for o in others):
                            continue
                        cands.append((abs(cx - ox) + abs(cy - oy), cx, cy))
                    if cands:
                        cands.sort()
                        best = cands[0]
                    r += step
                if best:
                    it.x, it.y = best[1], best[2]
                    it.update()
                    moved_any = True