#!/usr/bin/env python3
"""Generate KiCad 7 PCB for the standalone oscillator module (min board).

Flow: J1(coil) -> TVS1 -> T1 -> Colpitts(Q1+C1/C2/Cc2) -> LM393(U1) -> F_OUT
Layout: 46x28mm, VCC rail top (y=25.5), GND pour bottom, 2-layer.
Routing: grid A* router (0.5mm), 0.25mm tracks, 0.2mm clearance.
"""
import os, math, heapq, glob

import pcbnew
from pcbnew import (VECTOR2I, FromMM, F_Cu, B_Cu, F_SilkS, F_CrtYd,
                    Edge_Cuts, PAD_SHAPE_ROUNDRECT, PAD_ATTRIB_SMD)

BASE = os.path.dirname(os.path.abspath(__file__))
PCB = os.path.join(BASE, "osc-module.kicad_pcb")
FP = "/usr/share/kicad/footprints"

MM = lambda v: int(FromMM(v))
PX = lambda x, y: VECTOR2I(MM(x), MM(y))

# ---------------- board ----------------
board = pcbnew.BOARD()
ds = board.GetDesignSettings()
ds.m_MinClearance = MM(0.2)
ds.m_TrackMinWidth = MM(0.25)
ds.m_MinThroughDrill = MM(0.4)
ds.m_ViasMinSize = MM(0.8)

BOARD_W, BOARD_H = 46.0, 28.0

# ---------------- nets ----------------
def make_net(name):
    n = pcbnew.NETINFO_ITEM(board, name)
    board.Add(n)
    return n

NETS = {}
for nm in ["COIL_A", "COIL_B", "BASE", "COLLECTOR", "EMITTER",
           "COMP_IN", "VREF", "F_OUT", "VCC", "GND"]:
    NETS[nm] = make_net(nm)

def nc(nm):
    return NETS[nm].GetNetCode()

# ---------------- footprints ----------------
def load_fp(lib, name):
    fp = pcbnew.FootprintLoad(os.path.join(FP, lib), name)
    if fp is None:
        raise RuntimeError(f"footprint not found: {lib}:{name}")
    return fp

def place(fp, ref, value, x, y, rot=0):
    fp.SetReference(ref)
    fp.SetValue(value)
    fp.SetPosition(PX(x, y))
    fp.SetOrientationDegrees(rot)
    board.Add(fp)
    return fp

place(load_fp("Connector_PinHeader_2.54mm.pretty", "PinHeader_1x02_P2.54mm_Vertical"),
      "J1", "接线圈", 4.5, 21.5)
place(load_fp("Diode_SMD.pretty", "D_SOD-123"), "TVS1", "TVS", 9.0, 21.5)
place(load_fp("Capacitor_SMD.pretty", "C_0603_1608Metric"), "C1", "100~470p", 20.5, 20.0, 180)
place(load_fp("Package_TO_SOT_SMD.pretty", "SOT-23-3"), "Q1", "S9018", 25.5, 21.5)
place(load_fp("Capacitor_SMD.pretty", "C_0603_1608Metric"), "Cc2", "10p", 30.5, 20.0)
place(load_fp("Capacitor_SMD.pretty", "C_0603_1608Metric"), "C2", "100p", 27.5, 23.5, 90)
place(load_fp("Resistor_SMD.pretty", "R_0603_1608Metric"), "Re", "1k", 30.0, 23.5, 90)
place(load_fp("Resistor_SMD.pretty", "R_0603_1608Metric"), "Rb1", "100k", 20.5, 16.5, 180)
place(load_fp("Resistor_SMD.pretty", "R_0603_1608Metric"), "Rb2", "47k", 22.7, 16.5, 180)
place(load_fp("Resistor_SMD.pretty", "R_0603_1608Metric"), "Rc", "2.2k", 29.5, 16.5, 180)
place(load_fp("Package_SO.pretty", "SOIC-8_3.9x4.9mm_P1.27mm"), "U1", "LM393", 33.0, 6.0)
place(load_fp("Resistor_SMD.pretty", "R_0603_1608Metric"), "R3", "10k", 27.5, 2.0, 90)
place(load_fp("Resistor_SMD.pretty", "R_0603_1608Metric"), "R4", "10k", 29.0, 2.0, 90)
place(load_fp("Resistor_SMD.pretty", "R_0603_1608Metric"), "R7", "10k上拉", 38.5, 4.5)
place(load_fp("Connector_PinHeader_2.54mm.pretty", "PinHeader_1x03_P2.54mm_Vertical"),
      "J2", "VCC/GND/F_OUT", 40.5, 6.0)

# custom T1: 4-pad SMD matching transformer
t1 = pcbnew.FOOTPRINT(board)
t1.SetReference("T1")
t1.SetValue("XFMR 1:1")
t1.SetPosition(PX(15.0, 21.5))
for i, (dx, dy) in enumerate([(-2.5, -1.5), (-2.5, 1.5), (2.5, -1.5), (2.5, 1.5)], 1):
    pad = pcbnew.PAD(t1)
    pad.SetName(str(i))
    pad.SetShape(PAD_SHAPE_ROUNDRECT)
    pad.SetSize(PX(1.6, 1.6))
    pad.SetPosition(PX(dx, dy))
    pad.SetLayerSet(pcbnew.PAD_SMD)
    pad.SetAttribute(PAD_ATTRIB_SMD)
    pad.SetRoundRectRadiusRatio(0.25)
    t1.Add(pad)
for layer, x1, y1, x2, y2 in [(F_CrtYd, -3.4, -2.4, 3.4, 2.4),
                              (F_SilkS, -3.0, -2.0, 3.0, 2.0)]:
    s = pcbnew.FP_SHAPE(t1)
    s.SetShape(pcbnew.SHAPE_T_RECT)
    s.SetStart(PX(x1, y1))
    s.SetEnd(PX(x2, y2))
    s.SetLayer(layer)
    s.SetStrokeWidth(MM(0.15))
    t1.Add(s)
board.Add(t1)

# mounting holes (2x, non-plated M3)
holes = [(6.0, 4.0), (42.0, 22.0)]
for hx, hy in holes:
    name = "MountingHole_3.2mm_M3"
    cand = glob.glob(os.path.join(FP, "MountingHole.pretty", name + ".kicad_mod"))
    if not cand:
        cand = glob.glob(os.path.join(FP, "MountingHole.pretty", "*3.2mm*.kicad_mod"))
    if not cand:
        raise RuntimeError("no mounting hole footprint")
    place(load_fp("MountingHole.pretty", os.path.basename(cand[0])[:-9]),
          "H", "M3", hx, hy)

# ---------------- pad net map ----------------
PADMAP = {
    "J1":  {"1": "COIL_A", "2": "COIL_B"},
    "TVS1": {"1": "COIL_A", "2": "COIL_B"},
    "T1":  {"1": "COIL_A", "2": "COIL_B", "3": "BASE", "4": "GND"},
    "C1":  {"1": "BASE", "2": "GND"},
    "Q1":  {"1": "BASE", "2": "EMITTER", "3": "COLLECTOR"},
    "Cc2": {"1": "COLLECTOR", "2": "COMP_IN"},
    "C2":  {"1": "EMITTER", "2": "GND"},
    "Re":  {"1": "EMITTER", "2": "GND"},
    "Rb1": {"1": "BASE", "2": "VCC"},
    "Rb2": {"1": "BASE", "2": "GND"},
    "Rc":  {"1": "VCC", "2": "COLLECTOR"},
    "U1":  {"1": "F_OUT", "2": "VREF", "3": "COMP_IN", "4": "GND",
            "5": "GND", "6": "GND", "7": "GND", "8": "VCC"},
    "R3":  {"1": "VREF", "2": "VCC"},
    "R4":  {"1": "GND", "2": "VREF"},
    "R7":  {"1": "F_OUT", "2": "VCC"},
    "J2":  {"1": "VCC", "2": "GND", "3": "F_OUT"},
}

fps = {fp.GetReference(): fp for fp in board.GetFootprints()}
net_pads = {nm: [] for nm in NETS}
for ref, pm in PADMAP.items():
    for pname, net in pm.items():
        pad = fps[ref].FindPadByNumber(pname)
        pad.SetNetCode(nc(net))
        pos = pad.GetPosition()
        net_pads[net].append((ref, pname, pos.x, pos.y))

def mmpos(p):
    return (p.x / 1e6, p.y / 1e6)

# ---------------- grid A* router ----------------
GRID = 0.5
NX = int(BOARD_W / GRID)
NY = int(BOARD_H / GRID)
blocked = bytearray(NX * NY)

def cell(x, y):
    return (int(round(x / GRID)), int(round(y / GRID)))

def idx(i, j):
    return j * NX + i

def inb(i, j):
    return 0 <= i < NX and 0 <= j < NY

def block_circle(cx, cy, r):
    i0, j0 = cell(cx - r, cy - r)
    i1, j1 = cell(cx + r, cy + r)
    for j in range(max(0, j0), min(NY, j1 + 1)):
        for i in range(max(0, i0), min(NX, i1 + 1)):
            dx = i * GRID - cx
            dy = j * GRID - cy
            if dx * dx + dy * dy <= r * r:
                blocked[idx(i, j)] = 1

def block_rect(x0, y0, x1, y1):
    i0, j0 = cell(x0, y0)
    i1, j1 = cell(x1, y1)
    for j in range(max(0, j0), min(NY, j1 + 1)):
        for i in range(max(0, i0), min(NX, i1 + 1)):
            blocked[idx(i, j)] = 1

def dist_pt_seg(px, py, x0, y0, x1, y1):
    dx, dy = x1 - x0, y1 - y0
    L2 = dx * dx + dy * dy
    if L2 == 0:
        return math.hypot(px - x0, py - y0)
    t = max(0.0, min(1.0, ((px - x0) * dx + (py - y0) * dy) / L2))
    return math.hypot(px - (x0 + t * dx), py - (y0 + t * dy))

def block_segment(x0, y0, x1, y1, pad=0.4):
    i0, j0 = cell(min(x0, x1) - pad, min(y0, y1) - pad)
    i1, j1 = cell(max(x0, x1) + pad, max(y0, y1) + pad)
    for j in range(max(0, j0), min(NY, j1 + 1)):
        for i in range(max(0, i0), min(NX, i1 + 1)):
            if dist_pt_seg(i * GRID, j * GRID, x0, y0, x1, y1) <= pad:
                blocked[idx(i, j)] = 1

# block everything that already exists
for fp in board.GetFootprints():
    for pad in fp.Pads():
        if not (pad.IsOnLayer(F_Cu) or pad.IsOnLayer(B_Cu)):
            continue
        bb = pad.GetBoundingBox()
        block_rect(bb.GetLeft() / 1e6 - 0.4, bb.GetTop() / 1e6 - 0.4,
                   bb.GetRight() / 1e6 + 0.4, bb.GetBottom() / 1e6 + 0.4)

# board edge margin
block_rect(-5, -5, 1.2, BOARD_H + 5)
block_rect(BOARD_W - 1.2, -5, BOARD_W + 5, BOARD_H + 5)
block_rect(-5, -5, BOARD_W + 5, 1.2)
block_rect(-5, BOARD_H - 1.2, BOARD_W + 5, BOARD_H + 5)

for hx, hy in holes:
    block_circle(hx, hy, 2.0)

tracks = []

def add_track(x0, y0, x1, y1, net):
    t = pcbnew.PCB_TRACK(board)
    t.SetStart(PX(x0, y0))
    t.SetEnd(PX(x1, y1))
    t.SetWidth(MM(0.25))
    t.SetLayer(F_Cu)
    t.SetNetCode(nc(net))
    board.Add(t)
    tracks.append(t)
    block_segment(x0, y0, x1, y1)

def unblock_net(net, r=0.45):
    for (ref, pname, x, y) in net_pads[net]:
        block_circle_free(x / 1e6, y / 1e6, r)

def block_circle_free(cx, cy, r):
    """unblock cells within r of (cx,cy)"""
    i0, j0 = cell(cx - r, cy - r)
    i1, j1 = cell(cx + r, cy + r)
    for j in range(max(0, j0), min(NY, j1 + 1)):
        for i in range(max(0, i0), min(NX, i1 + 1)):
            dx = i * GRID - cx
            dy = j * GRID - cy
            if dx * dx + dy * dy <= r * r:
                blocked[idx(i, j)] = 0

def astar(sx, sy, gx, gy):
    si, sj = cell(sx, sy)
    gi, gj = cell(gx, gy)
    if not inb(si, sj) or not inb(gi, gj):
        return None
    if blocked[idx(si, sj)] or blocked[idx(gi, gj)]:
        return None
    open_h = [(0.0, si, sj)]
    g = {(si, sj): 0.0}
    came = {}
    closed = set()
    while open_h:
        f, i, j = heapq.heappop(open_h)
        if (i, j) == (gi, gj):
            path = [(gi, gj)]
            while (i, j) in came:
                (i, j) = came[(i, j)]
                path.append((i, j))
            path.reverse()
            return path
        if (i, j) in closed:
            continue
        closed.add((i, j))
        for di, dj in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            ni, nj = i + di, j + dj
            if not inb(ni, nj) or blocked[idx(ni, nj)]:
                continue
            ng = g[(i, j)] + 1.0
            if (ni, nj) not in g or ng < g[(ni, nj)]:
                g[(ni, nj)] = ng
                came[(ni, nj)] = (i, j)
                heapq.heappush(open_h, (ng + abs(ni - gi) + abs(nj - gj), ni, nj))
    return None

def path_to_tracks(path, net):
    # merge collinear runs
    runs = []
    cur = [path[0]]
    for p in path[1:]:
        if cur and (p[0] == cur[-1][0] or p[1] == cur[-1][1]):
            cur.append(p)
        else:
            runs.append(cur)
            cur = [p]
    runs.append(cur)
    for run in runs:
        x0, y0 = run[0][0] * GRID, run[0][1] * GRID
        x1, y1 = run[-1][0] * GRID, run[-1][1] * GRID
        if (x0, y0) != (x1, y1):
            add_track(x0, y0, x1, y1, net)

def route_net(net, order=None):
    pads = [(p[2] / 1e6, p[3] / 1e6) for p in net_pads[net]]
    if order:
        pads = [pads[i] for i in order]
    unblock_net(net)
    cur = pads[0]
    for nxt in pads[1:]:
        path = astar(cur[0], cur[1], nxt[0], nxt[1])
        if path is None:
            print(f"!! A* FAILED: {net}  {cur} -> {nxt}")
            return False
        path_to_tracks(path, net)
        cur = nxt
    return True

# ---- VCC rail (top) ----
add_track(2.0, 25.5, 44.0, 25.5, "VCC")

# ---- signal nets ----
ok = True
ok &= route_net("COIL_A")
ok &= route_net("COIL_B")
ok &= route_net("EMITTER")
ok &= route_net("COLLECTOR")
ok &= route_net("BASE")
ok &= route_net("COMP_IN")
ok &= route_net("VREF")
ok &= route_net("F_OUT")

# ---- VCC drops (chain from rail) ----
vcc_pads = [(p[2] / 1e6, p[3] / 1e6) for p in net_pads["VCC"]]
vcc_chain = [(43.0, 25.5)] + vcc_pads + [(3.0, 25.5)]
unblock_net("VCC")
cur = vcc_chain[0]
for nxt in vcc_chain[1:]:
    path = astar(cur[0], cur[1], nxt[0], nxt[1])
    if path is None:
        print(f"!! A* FAILED: VCC  {cur} -> {nxt}")
        ok = False
    else:
        path_to_tracks(path, "VCC")
    cur = nxt

# ---- GND vias + stubs ----
gnd_vias = []
for (ref, pname, x, y) in net_pads["GND"]:
    pad = fps[ref].FindPadByNumber(pname)
    if pad.GetAttribute() != PAD_ATTRIB_SMD:
        continue  # THT connects to pour through hole
    px, py = x / 1e6, y / 1e6
    # chain the three unused U1 comparator pads together first
    vx, vy = None, None
    for ox, oy in [(1.6, 0), (-1.6, 0), (0, 1.6), (0, -1.6)]:
        cx0, cy0 = px + ox, py + oy
        i0, j0 = cell(cx0, cy0)
        if not inb(i0, j0):
            continue
        free = True
        for j in range(max(0, j0 - 2), min(NY, j0 + 3)):
            for i in range(max(0, i0 - 2), min(NX, i0 + 3)):
                if blocked[idx(i, j)]:
                    free = False
                    break
            if not free:
                break
        if free:
            vx, vy = cx0, cy0
            break
    if vx is None:
        print(f"!! no via spot for GND {ref}.{pname}")
        ok = False
        continue
    via = pcbnew.PCB_VIA(board)
    via.SetPosition(PX(vx, vy))
    via.SetDrill(MM(0.4))
    via.SetWidth(MM(0.8))
    via.SetLayerPair(F_Cu, B_Cu)
    via.SetNetCode(nc("GND"))
    board.Add(via)
    block_circle(vx, vy, 0.7)
    # route stub pad -> via
    unblock_net("GND")
    path = astar(px, py, vx, vy)
    if path is None:
        print(f"!! A* FAILED: GND stub {ref}.{pname}")
        ok = False
    else:
        path_to_tracks(path, "GND")

# ---- GND pour (bottom) ----
z = pcbnew.ZONE(board)
z.SetLayer(B_Cu)
z.SetNetCode(nc("GND"))
z.SetIsFilled(False)
ol = z.Outline()
for (ox, oy) in [(1.0, 1.0), (BOARD_W - 1.0, 1.0), (BOARD_W - 1.0, BOARD_H - 1.0),
                 (1.0, BOARD_H - 1.0), (1.0, 1.0)]:
    ol.Append(PX(ox, oy))
board.Add(z)

# ---- edge cuts ----
e = pcbnew.PCB_SHAPE(board)
e.SetShape(pcbnew.SHAPE_T_RECT)
e.SetStart(PX(0, 0))
e.SetEnd(PX(BOARD_W, BOARD_H))
e.SetLayer(Edge_Cuts)
board.Add(e)

# ---- silk text ----
for txt, x, y, sz in [("OSC-MOD v1.0", 23.0, 27.2, 1.2),
                      ("J1: COIL", 4.5, 26.2, 0.8),
                      ("J2: 1=VCC 2=GND 3=F_OUT", 33.5, 12.5, 0.8)]:
    t = pcbnew.PCB_TEXT(board)
    t.SetText(txt)
    t.SetPosition(PX(x, y))
    t.SetLayer(F_SilkS)
    t.SetTextSize(PX(sz, sz))
    t.SetTextThickness(MM(0.15))
    board.Add(t)

board.Save(PCB)
print("saved", PCB)
print("nets routed OK" if ok else "!!! SOME NETS FAILED !!!")

# ---- zone fill ----
filler = pcbnew.ZONE_FILLER(board)
filler.Fill([z])
board.Save(PCB)
print("zone filled & re-saved")