AM/FM Receiver Molecule

"""KiCad 10 schematic writer: places symbols, wire stubs, labels, power ports.

Coordinates are in mm, schematic space (+y down).  All placement is snapped to
the 1.27 mm grid so that KiCad sees the connections.
"""
import math, uuid as _uuid
from collections import OrderedDict
from sexp import Sym, dumps, find, findall
import symlib

S = Sym
GRID = 1.27
NS = _uuid.UUID('6f1d3c2a-9a51-4a0e-bb52-2f0c0de4a7f0')


def uid(*parts):
    return str(_uuid.uuid5(NS, '/'.join(str(p) for p in parts)))


def snap(v):
    return round(round(v / GRID) * GRID, 4)


def rot(dx, dy, r):
    """Rotate a screen-space vector by r degrees counter-clockwise (visual)."""
    r %= 360
    if r == 0:
        return dx, dy
    if r == 90:
        return dy, -dx
    if r == 180:
        return -dx, -dy
    if r == 270:
        return -dy, dx
    raise ValueError(r)


def font(size=1.27, bold=False, italic=False):
    f = [S('font'), [S('size'), size, size]]
    if bold:
        f.append([S('bold'), S('yes')])
    if italic:
        f.append([S('italic'), S('yes')])
    return f


def effects(size=1.27, justify=None, bold=False):
    e = [S('effects'), font(size, bold)]
    if justify:
        e.append([S('justify')] + [S(j) for j in justify.split()])
    return e


def prop(name, value, x, y, angle=0, hide=False, justify=None, size=1.27):
    p = [S('property'), name, value, [S('at'), x, y, angle]]
    if hide:
        p.append([S('hide'), S('yes')])
    p.append(effects(size, justify))
    return p


POWER_UP = {'+5V', '+3V3', 'VBUS', '+5VA', '+3.3VA'}
POWER_DOWN = {'GND'}
POWER_LIB = {'+5V': 'power:+5V', '+3V3': 'power:+3V3', 'VBUS': 'power:VBUS',
             '+5VA': 'power:+5VA', '+3.3VA': 'power:+3.3VA', 'GND': 'power:GND'}



DIRV = {'R': (1, 0), 'L': (-1, 0), 'U': (0, -1), 'D': (0, 1)}


def plan_runs(part, pins_here, net_of_pin, stub):
    """Return (no_connect points, runs).  run = (side, members, stub_len, stub_ends, anchor, net)."""
    by_loc = OrderedDict()
    for p, ax, ay, d in pins_here:
        by_loc.setdefault((ax, ay), []).append((p, d))
    ends, ncs = [], []
    for (ax, ay), plist in by_loc.items():
        nets = set()
        for p, d in plist:
            n = net_of_pin.get(p['number'])
            if n is not None:
                nets.add(n)
        visible = [pd for pd in plist if not pd[0]['hidden']]
        if len(nets) > 1:
            raise ValueError('%s: stacked pins at same location on different nets %s' % (part.ref, nets))
        if not visible:
            if nets:
                raise ValueError('%s: hidden pin(s) %s carry net %s but have no visible partner'
                                 % (part.ref, [p['number'] for p, _ in plist], nets))
            continue
        p, d = visible[0]
        if not nets:
            if p['type'] != 'no_connect':
                ncs.append((ax, ay))
            continue
        ends.append((ax, ay, d, nets.pop()))
    runs = []
    for side in ('L', 'R', 'U', 'D'):
        sp = [e for e in ends if e[2] == side]
        if not sp:
            continue
        along = (lambda e: e[1]) if side in ('L', 'R') else (lambda e: e[0])
        perp = (lambda e: e[0]) if side in ('L', 'R') else (lambda e: e[1])
        sp.sort(key=lambda e: (perp(e), along(e)))
        cur = []
        for e in sp:
            if cur and cur[-1][3] == e[3] and perp(cur[-1]) == perp(e) and \
                    abs(along(e) - along(cur[-1])) <= 2.54 + 1e-6:
                cur.append(e)
            else:
                if cur:
                    runs.append([side, cur])
                cur = [e]
        if cur:
            runs.append([side, cur])
    out = []
    for side in ('L', 'R', 'U', 'D'):
        rs = [r for r in runs if r[0] == side]
        along = (lambda e: e[1]) if side in ('L', 'R') else (lambda e: e[0])
        rs.sort(key=lambda r: min(along(e) for e in r[1]))
        prev_end, prev_long = None, False
        for side_, members in rs:
            net = members[0][3]
            is_pwr = net in POWER_LIB
            lo = min(along(e) for e in members)
            hi = max(along(e) for e in members)
            long_ = bool(is_pwr and prev_end is not None and lo - prev_end <= 5.08 + 1e-6 and not prev_long)
            L = stub * 3 if long_ else stub
            if is_pwr:
                prev_end, prev_long = hi, long_
            else:
                prev_end, prev_long = None, False
            dx, dy = DIRV[side]
            stub_ends = [(snap(m[0] + dx * L), snap(m[1] + dy * L)) for m in members]
            if side in ('L', 'R'):
                anchor = min(stub_ends, key=lambda t: t[1])
            else:
                anchor = min(stub_ends, key=lambda t: t[0])
            out.append((side, members, L, stub_ends, anchor, net))
    return ncs, out


class LibCache:
    def __init__(self, extra_libs=None):
        self.extra = extra_libs or {}
        self.syms = {}

    def get(self, lib_id):
        if lib_id not in self.syms:
            self.syms[lib_id] = symlib.get_symbol(lib_id, self.extra)
        return self.syms[lib_id]

    def pins(self, lib_id):
        return symlib.pins(self.get(lib_id))

    def lib_entry(self, lib_id):
        """Symbol definition for the schematic's lib_symbols section."""
        import copy
        s = copy.deepcopy(self.get(lib_id))
        s[1] = lib_id
        return s


class SheetWriter:
    def __init__(self, project, root_uuid, sheet_path, file_uuid, libs, title='', paper='A3',
                 page_label=''):
        self.project = project
        self.root_uuid = root_uuid
        self.sheet_path = sheet_path          # e.g. "/<root>/<sheet>" or "/<root>"
        self.file_uuid = file_uuid
        self.libs = libs
        self.items = []
        self.used = set()
        self.title = title
        self.paper = paper
        self.page_label = page_label
        self.pwr_count = 0
        self.flag_count = 0
        self.global_nets = set()
        self.pin_net_seen = {}                # debug: (ref, pin) -> net
        self.label_points = []

    # ---------------------------------------------------------------- primitives
    def wire(self, x1, y1, x2, y2):
        self.items.append([S('wire'), [S('pts'), [S('xy'), x1, y1], [S('xy'), x2, y2]],
                           [S('stroke'), [S('width'), 0], [S('type'), S('default')]],
                           [S('uuid'), uid(self.file_uuid, 'w', x1, y1, x2, y2)]])

    def no_connect(self, x, y):
        self.items.append([S('no_connect'), [S('at'), x, y], [S('uuid'), uid(self.file_uuid, 'nc', x, y)]])

    def text(self, s, x, y, size=1.27, bold=False, angle=0):
        self.items.append([S('text'), s, [S('exclude_from_sim'), S('no')], [S('at'), x, y, angle],
                           effects(size, 'left bottom', bold), [S('uuid'), uid(self.file_uuid, 't', s, x, y)]])

    def rect(self, x1, y1, x2, y2, dash=True):
        self.items.append([S('rectangle'), [S('start'), x1, y1], [S('end'), x2, y2],
                           [S('stroke'), [S('width'), 0.254], [S('type'), S('dash' if dash else 'default')]],
                           [S('fill'), [S('type'), S('none')]],
                           [S('uuid'), uid(self.file_uuid, 'r', x1, y1, x2, y2)]])

    def label(self, net, x, y, direction, is_global):
        """direction: outward screen direction 'R','L','U','D' of the stub end."""
        angle = {'R': 0, 'U': 90, 'L': 180, 'D': 270}[direction]
        just = {'R': 'left', 'U': 'left', 'L': 'right', 'D': 'right'}[direction]
        self.label_points.append((x, y, net))
        if is_global:
            self.global_nets.add(net)
            self.items.append([S('global_label'), net, [S('shape'), S('bidirectional')],
                               [S('at'), x, y, angle], [S('fields_autoplaced'), S('yes')],
                               effects(1.27, just), [S('uuid'), uid(self.file_uuid, 'gl', net, x, y)],
                               [S('property'), 'Intersheetrefs', '${INTERSHEET_REFS}', [S('at'), x, y, 0],
                                [S('hide'), S('yes')], effects(1.27, just)]])
        else:
            self.items.append([S('label'), net, [S('at'), x, y, angle], [S('fields_autoplaced'), S('yes')],
                               effects(1.27, just + ' bottom'), [S('uuid'), uid(self.file_uuid, 'l', net, x, y)]])

    def power_port(self, net, x, y, direction):
        lib_id = POWER_LIB[net]
        self.used.add(lib_id)
        self.pwr_count += 1
        ref = '#PWR%s%03d' % (self.page_label, self.pwr_count)
        # natural orientation (supply up, ground down) unless that would fold back onto the stub
        if net in POWER_DOWN:
            r = 180 if direction == 'U' else 0
        else:
            r = 180 if direction == 'D' else 0
        su = uid(self.file_uuid, 'pwr', ref)
        # value text placed beyond the symbol
        vx, vy = rot(0, -3.2 if net not in POWER_DOWN else 3.8, r)
        self.items.append([S('symbol'), [S('lib_id'), lib_id], [S('at'), x, y, r], [S('unit'), 1],
                           [S('exclude_from_sim'), S('no')], [S('in_bom'), S('yes')], [S('on_board'), S('yes')],
                           [S('dnp'), S('no')], [S('uuid'), su],
                           prop('Reference', ref, x, y, 0, hide=True),
                           prop('Value', net, snap(x + vx), snap(y + vy), 0),
                           prop('Footprint', '', x, y, 0, hide=True),
                           prop('Datasheet', '', x, y, 0, hide=True),
                           prop('Description', '', x, y, 0, hide=True),
                           [S('pin'), '1', [S('uuid'), uid(su, 'p1')]],
                           [S('instances'), [S('project'), self.project,
                                             [S('path'), self.sheet_path, [S('reference'), ref], [S('unit'), 1]]]]])

    def pwr_flag(self, net, x, y):
        """PWR_FLAG attached by a short stub + power port / label."""
        self.used.add('power:PWR_FLAG')
        self.flag_count += 1
        ref = '#FLG%s%02d' % (self.page_label, self.flag_count)
        su = uid(self.file_uuid, 'flg', ref)
        self.items.append([S('symbol'), [S('lib_id'), 'power:PWR_FLAG'], [S('at'), x, y, 0], [S('unit'), 1],
                           [S('exclude_from_sim'), S('no')], [S('in_bom'), S('yes')], [S('on_board'), S('yes')],
                           [S('dnp'), S('no')], [S('uuid'), su],
                           prop('Reference', ref, x, y - 5, 0, hide=True),
                           prop('Value', 'PWR_FLAG', x, y - 3.8, 0),
                           prop('Footprint', '', x, y, 0, hide=True),
                           prop('Datasheet', '', x, y, 0, hide=True),
                           prop('Description', '', x, y, 0, hide=True),
                           [S('pin'), '1', [S('uuid'), uid(su, 'p1')]],
                           [S('instances'), [S('project'), self.project,
                                             [S('path'), self.sheet_path, [S('reference'), ref], [S('unit'), 1]]]]])
        # stub down to the net connection
        self.wire(x, y, x, y + 2.54)
        self.attach_net(net, x, y + 2.54, 'D', is_global=True)

    def attach_net(self, net, x, y, direction, is_global):
        if net in POWER_LIB:
            self.power_port(net, x, y, direction)
        else:
            self.label(net, x, y, direction, is_global)

    # ---------------------------------------------------------------- symbols
    def pin_geometry(self, lib_id, unit, x, y, r):
        """Yield (pin dict, abs_x, abs_y, outward_direction)."""
        out = []
        for p in self.libs.pins(lib_id):
            if p['unit'] not in (0, unit) or p['convert'] not in (0, 1):
                continue
            dx, dy = rot(p['x'], -p['y'], r)
            ax, ay = snap(x + dx), snap(y + dy)
            a = math.radians(p['angle'] + 180)
            ox, oy = rot(round(math.cos(a)), -round(math.sin(a)), r)
            d = {(1, 0): 'R', (-1, 0): 'L', (0, -1): 'U', (0, 1): 'D'}[(ox, oy)]
            out.append((p, ax, ay, d))
        return out

    def place(self, part, unit, x, y, r, net_of_pin, is_global, stub=2.54, ref_pos=None, val_pos=None,
              show_value=True):
        """Place a symbol unit and wire every visible pin to its net (label/power port) or NC flag."""
        lib_id = part.lib_id
        self.used.add(lib_id)
        su = uid(self.file_uuid, 'sym', part.ref, unit)
        # Property text positions (default: to the right of the symbol)
        rx, ry, ra, rj = ref_pos if ref_pos else (x + 3.0, y - 1.3, 0, 'left')
        vx, vy, va, vj = val_pos if val_pos else (x + 3.0, y + 1.3, 0, 'left')
        props = [prop('Reference', part.ref, round(rx, 3), round(ry, 3), ra, justify=rj),
                 prop('Value', part.value, round(vx, 3), round(vy, 3), va, justify=vj, hide=not show_value),
                 prop('Footprint', part.footprint, x, y, 0, hide=True),
                 prop('Datasheet', part.datasheet or '', x, y, 0, hide=True),
                 prop('Description', part.description or '', x, y, 0, hide=True)]
        for k, v in part.fields.items():
            props.append(prop(k, v, x, y, 0, hide=True))
        pins_here = self.pin_geometry(lib_id, unit, x, y, r)
        pin_entries = []
        seen_numbers = set()
        for p, ax, ay, d in pins_here:
            if p['number'] not in seen_numbers:
                pin_entries.append([S('pin'), p['number'], [S('uuid'), uid(su, 'pin', p['number'])]])
                seen_numbers.add(p['number'])
        self.items.append([S('symbol'), [S('lib_id'), lib_id], [S('at'), x, y, r], [S('unit'), unit],
                           [S('exclude_from_sim'), S('no')], [S('in_bom'), S('yes' if part.in_bom else 'no')],
                           [S('on_board'), S('yes')], [S('dnp'), S('no')],
                           [S('uuid'), su]] + props + pin_entries +
                          [[S('instances'), [S('project'), self.project,
                                             [S('path'), self.sheet_path, [S('reference'), part.ref],
                                              [S('unit'), unit]]]]])
        self.connect_pins(part, pins_here, net_of_pin, is_global, stub)
        return pins_here

    def connect_pins(self, part, pins_here, net_of_pin, is_global, stub):
        ncs, runs = plan_runs(part, pins_here, net_of_pin, stub)
        for ax, ay in ncs:
            self.no_connect(ax, ay)
        for side, members, L, stub_ends, (ex, ey), net in runs:
            for (ax, ay, d, n), (sx, sy) in zip(members, [(snap(m[0] + DIRV[side][0] * L),
                                                           snap(m[1] + DIRV[side][1] * L)) for m in members]):
                self.wire(ax, ay, sx, sy)
            if len(stub_ends) > 1:
                se = sorted(stub_ends)
                for (x1, y1), (x2, y2) in zip(se, se[1:]):
                    self.wire(x1, y1, x2, y2)
                for jx, jy in se[1:-1]:
                    self.junction(jx, jy)
            if net in POWER_LIB:
                self.power_port(net, ex, ey, side)
            else:
                self.label(net, ex, ey, side if side in ('L', 'R') else 'R', is_global(net))

    def junction(self, x, y):
        self.items.append([S('junction'), [S('at'), x, y], [S('diameter'), 0], [S('color'), 0, 0, 0, 0],
                           [S('uuid'), uid(self.file_uuid, 'j', x, y)]])

    # ---------------------------------------------------------------- output
    def sheet_symbol(self, name, fname, x, y, w, h, sheet_uuid, page):
        self.items.append([S('sheet'), [S('at'), x, y], [S('size'), w, h],
                           [S('exclude_from_sim'), S('no')], [S('in_bom'), S('yes')], [S('on_board'), S('yes')],
                           [S('dnp'), S('no')], [S('fields_autoplaced'), S('yes')],
                           [S('stroke'), [S('width'), 0.1524], [S('type'), S('solid')]],
                           [S('fill'), [S('color'), 0, 0, 0, 0.0]],
                           [S('uuid'), sheet_uuid],
                           [S('property'), 'Sheetname', name, [S('at'), x, snap(y - 0.7) if False else y - 0.7112, 0],
                            effects(1.27, 'left bottom')],
                           [S('property'), 'Sheetfile', fname, [S('at'), x, y + h + 0.5938, 0],
                            effects(1.27, 'left top')],
                           [S('instances'), [S('project'), self.project,
                                             [S('path'), '/' + self.root_uuid, [S('page'), str(page)]]]]])

    def render(self, is_root=False, title_block=None):
        head = [S('kicad_sch'), [S('version'), 20251024], [S('generator'), 'eeschema'],
                [S('generator_version'), '10.0'], [S('uuid'), self.file_uuid], [S('paper'), self.paper]]
        if title_block:
            tb = [S('title_block')]
            for k in ('title', 'date', 'rev', 'company'):
                if title_block.get(k):
                    tb.append([S(k), title_block[k]])
            for i, c in enumerate(title_block.get('comments', []), 1):
                tb.append([S('comment'), i, c])
            head.append(tb)
        libsec = [S('lib_symbols')] + [self.libs.lib_entry(l) for l in sorted(self.used)]
        body = head + [libsec] + self.items
        if is_root:
            body.append([S('sheet_instances'), [S('path'), '/', [S('page'), '1']]])
        body.append([S('embedded_fonts'), S('no')])
        return dumps(body) + '\n'