generate_modules.py

"""
Generate only the 3D models for various SoM, DevKits, etc.
"""

from os import path

import cadquery as cq

from cadquery_helpers import StepAssembly, StepColor, StepConstants
from common import init_cache

CU_THICKNESS = 0.05


# Initialize UUID caches
connectors_uuid_cache = init_cache('uuid_cache_connectors.csv')


def get_connector_pkg_uuid(kind: str, pin_count: int, rows: int, drill: float, obj: str) -> str:
    """
    Get the UUID of a connector package item. See generate_connectors.py for details.
    """
    key = f'pkg-{kind}-{rows}x{pin_count // rows}-d{drill:.1f}-{obj}'
    return connectors_uuid_cache[key]


def load_connector_step_model(kind: str, pin_count: int, rows: int, drill: float) -> cq.Assembly:
    """
    Load the STEP model of a connector generated by generate_connectors.py
    """
    uuid_pkg = get_connector_pkg_uuid(kind, pin_count, rows, drill, 'pkg')
    uuid_3d = get_connector_pkg_uuid(kind, pin_count, rows, drill, '3d')
    path = f'out/LibrePCB_Connectors.lplib/pkg/{uuid_pkg}/{uuid_3d}.step'
    return cq.Assembly.load(path)


def generate_arduino_uno_r3(name: str, with_board: bool) -> None:
    print(f'generate_arduino_uno_r3: {name}')
    assembly = StepAssembly(name)

    # Pin headers
    header_6p = load_connector_step_model('pinheader', 6, 1, 1.0)
    header_8p = load_connector_step_model('pinheader', 8, 1, 1.0)
    header_10p = load_connector_step_model('pinheader', 10, 1, 1.0)
    assembly.add_body(header_6p, 'H1', location=cq.Location((-24.13, 22.86, 0)))
    assembly.add_body(header_8p, 'H2', location=cq.Location((-24.13, 2.54, 0)))
    assembly.add_body(header_8p, 'H3', location=cq.Location((24.13, 20.32, 0)))
    assembly.add_body(header_10p, 'H4', location=cq.Location((24.13, -4.064, 0)))

    if with_board:
        # PCB
        pcb_standoff = 11.0
        pcb_outline = [
            (-26.67, 31.75),
            (-24.13, 31.75),
            (-21.59, 34.29),
            (11.176, 34.29),
            (13.716, 31.75),
            (26.67, 31.75),
            (26.67, -34.29),
            (-26.67, -34.29),
        ]
        holes = [
            (-19.05, 31.75),
            (8.89, 31.75),
            (-24.13, -20.32),
            (24.13, -19.05),
        ]
        pcb_body = (
            cq.Workplane('XY')
            .polyline(pcb_outline)
            .close()
            .pushPoints(holes)
            .circle(1.6)
            .extrude(1.6)
        )
        assembly.add_body(
            pcb_body, 'PCB', cq.Color('turquoise3'), location=cq.Location((0, 0, pcb_standoff))
        )

        # Pin sockets
        socket_standoff = 3.0
        socket_height = pcb_standoff - socket_standoff
        socket_l = cq.Workplane('XY').box(2.54, 38.1, socket_height, centered=(True, True, False))
        assembly.add_body(
            socket_l,
            'S1',
            StepColor.IC_BODY,
            location=cq.Location((-24.13, 11.43, socket_standoff)),
        )
        socket_r = cq.Workplane('XY').box(2.54, 47.2, socket_height, centered=(True, True, False))
        assembly.add_body(
            socket_r, 'S2', StepColor.IC_BODY, location=cq.Location((24.13, 6.858, socket_standoff))
        )

        # Power connector
        pwr_height = 10.0  # guessed
        pwr_standoff = pcb_standoff - pwr_height
        pwr = cq.Workplane('XY').box(9.0, 13.0, pwr_height, centered=(True, False, False))
        assembly.add_body(
            pwr, 'PWR', StepColor.IC_BODY, location=cq.Location((-18.87, -36.09, pwr_standoff))
        )

        # USB connector
        usb_height = 10.0  # guessed
        usb_standoff = pcb_standoff - usb_height
        usb = cq.Workplane('XY').box(12.0, 16.0, usb_height, centered=(True, False, False))
        assembly.add_body(
            usb, 'USB', StepColor.LEAD_SMT, location=cq.Location((11.43, -40.63, usb_standoff))
        )

    out_path = path.join('out_3d', 'modules', 'arduino_uno_r3', f'{name}.step')
    assembly.save(out_path, fused=False)


def generate_esp32_c6_wroom_1() -> None:
    name = 'ESP32-C6-WROOM-1'
    print(f'generate_esp32_c6_wroom_1: {name}')
    assembly = StepAssembly(name)

    # PCB
    x0 = -18.0 / 2
    x1 = 18.0 / 2
    pads_y = [1.5 + 16.51 - i * 1.27 for i in range(14)]
    pads_xy = [(x0, y) for y in pads_y] + [(x1, y) for y in pads_y]
    pcb_body = cq.Workplane('XY').box(18.0, 25.5, 0.8, centered=(True, False, False))
    pcb_body = pcb_body.pushPoints(pads_xy).rect(0.9, 0.9).cutThruAll()
    assembly.add_body(
        pcb_body, 'PCB', cq.Color('gray42'), location=cq.Location((0, -9.755, CU_THICKNESS))
    )

    # Antenna trace
    pts = [
        (-7.9375, 10.16),
        (-7.9375, 14.605),
        (-4.7625, 14.605),
        (-4.7625, 12.065),
        (-2.8575, 12.065),
        (-2.8575, 14.605),
        (0.3175, 14.605),
        (0.3175, 12.065),
        (2.2225, 12.065),
        (2.2225, 14.605),
        (7.9375, 14.605),
        (7.9375, 9.525),
        (7.3025, 9.525),
        (7.3025, 13.97),
        (5.3975, 13.97),
        (5.3975, 9.525),
        (4.7625, 9.525),
        (4.7625, 13.97),
        (2.8575, 13.97),
        (2.8575, 11.43),
        (-0.3175, 11.43),
        (-0.3175, 13.97),
        (-2.2225, 13.97),
        (-2.2225, 11.43),
        (-5.3975, 11.43),
        (-5.3975, 13.97),
        (-7.3025, 13.97),
        (-7.3025, 10.16),
        (-7.9375, 10.16),
    ]
    antenna = cq.Workplane('XY').polyline(pts).close().extrude(CU_THICKNESS)
    assembly.add_body(
        antenna, 'ANTENNA', cq.Color('gray56'), location=cq.Location((0, 0, 0.8 + CU_THICKNESS))
    )

    # Pads
    pad = cq.Workplane('XY').box(0.45, 0.9, 0.8 + 2 * CU_THICKNESS, centered=(False, True, False))
    pad = pad.pushPoints([(0, 0)]).circle(0.55 / 2).cutThruAll()
    for i in range(14):
        assembly.add_body(
            pad, f'PAD{i + 1}', StepColor.COPPER, location=cq.Location((x0, pads_y[i] - 9.755, 0))
        )
        assembly.add_body(
            pad,
            f'PAD{28 - i}',
            StepColor.COPPER,
            location=cq.Location((x1, pads_y[i] - 9.755, 0), (0, 0, 180)),
        )

    # Metal cover
    cover = (
        cq.Workplane('XY')
        .box(15.8, 17.6, 2.3, centered=(True, True, False))
        .edges('>Z or |Z')
        .fillet(0.3)
    )
    assembly.add_body(
        cover, 'COVER', cq.Color('gainsboro'), location=cq.Location((0, 0, 0.8 + CU_THICKNESS))
    )

    out_path = path.join('out_3d', 'modules', 'esp32_c6_wroom_1', f'{name}.step')
    assembly.save(out_path, fused=False)


def generate_esp32_pico_mini_02() -> None:
    name = 'ESP32-PICO-MINI-02'
    print(f'generate_esp32_pico_mini_02: {name}')
    assembly = StepAssembly(name)

    # PCB
    pcb_body = cq.Workplane('XY').box(13.2, 16.6, 0.8, centered=(True, False, False))
    assembly.add_body(pcb_body, 'PCB', cq.Color('gray42'), location=cq.Location((0, -11.2 / 2, 0)))

    # Antenna trace
    pts = [
        (-6.0325, 6.0325),
        (-6.0325, 10.4775),
        (-1.5875, 10.4775),
        (-1.5875, 7.9375),
        (-0.3175, 7.9375),
        (-0.3175, 10.4775),
        (2.2225, 10.4775),
        (2.2225, 7.9375),
        (3.4925, 7.9375),
        (3.4925, 10.4775),
        (6.0325, 10.4775),
        (6.0325, 6.0325),
        (5.3975, 6.0325),
        (5.3975, 9.8425),
        (4.1275, 9.8425),
        (4.1275, 7.3025),
        (1.5875, 7.3025),
        (1.5875, 9.8425),
        (0.3175, 9.8425),
        (0.3175, 7.3025),
        (-2.2225, 7.3025),
        (-2.2225, 9.8425),
        (-3.4925, 9.8425),
        (-3.4925, 6.0325),
        (-4.1275, 6.0325),
        (-4.1275, 9.8425),
        (-5.3975, 9.8425),
        (-5.3975, 6.0325),
        (-6.0325, 6.0325),
    ]
    antenna = cq.Workplane('XY').polyline(pts).close().extrude(CU_THICKNESS)
    assembly.add_body(antenna, 'ANTENNA', cq.Color('gray56'), location=cq.Location((0, 0, 0.8)))

    # Metal cover
    cover = (
        cq.Workplane('XY')
        .box(11.95, 9.95, 1.6, centered=(True, True, False))
        .edges('>Z or |Z')
        .fillet(0.3)
    )
    assembly.add_body(cover, 'COVER', cq.Color('gainsboro'), location=cq.Location((0, 0, 0.8)))

    out_path = path.join('out_3d', 'modules', 'esp32_pico_mini_02', f'{name}.step')
    assembly.save(out_path, fused=True)


def generate_rpi_pico(name: str, bottom_pads: bool, headers: bool) -> None:
    print(f'generate_rpi_pico: {name}')
    assembly = StepAssembly(name)

    pcb_standoff = 9.5 if headers else CU_THICKNESS

    # PCB
    y0 = 51.0 / 2
    y1 = -51.0 / 2
    x0 = -21.0 / 2
    x1 = 21.0 / 2
    pcb_outline = [
        (x0, y0),
        (x1, y0),
        (x1, y1),
        (x0, y1),
    ]
    holes = [
        (-11.4 / 2, 51.0 / 2 - 2.0),
        (11.4 / 2, 51.0 / 2 - 2.0),
        (11.4 / 2, -51.0 / 2 + 2.0),
        (-11.4 / 2, -51.0 / 2 + 2.0),
    ]
    pcb_body = (
        cq.Workplane('XY')
        .polyline(pcb_outline)
        .close()
        .pushPoints(holes)
        .circle(2.1 / 2)
        .extrude(1.6)
    )
    if not bottom_pads:
        holes = [
            (x1 - 7.3759 - 2.54, y1 + 19.8),
            (x1 - 7.3759, y1 + 19.8),
            (x1 - 7.3759 + 2.54, y1 + 19.8),
        ]
        pcb_body = pcb_body.pushPoints(holes).circle(0.5).cutThruAll()
    pads_y = [48.26 / 2 - i * 2.54 for i in range(20)]
    pads_xy = [(x0, y) for y in pads_y] + [(x1, y) for y in pads_y]
    pads_inner_xy = [(x0 + 1.61, y) for y in pads_y] + [(x1 - 1.61, y) for y in pads_y]
    pcb_body = pcb_body.pushPoints(pads_xy).rect(2 * 1.61, 1.6).cutThruAll()
    pcb_body = pcb_body.pushPoints(pads_inner_xy).circle(0.8).cutThruAll()
    if bottom_pads:
        pads_bottom_xy = [(-2.54, y1), (0, y1), (2.54, y1)]
        pads_bottom_inner_xy = [(x, y1 + 1.61) for x, y in pads_bottom_xy]
        pcb_body = pcb_body.pushPoints(pads_bottom_xy).rect(1.6, 2 * 1.61).cutThruAll()
        pcb_body = pcb_body.pushPoints(pads_bottom_inner_xy).circle(0.8).cutThruAll()
    assembly.add_body(
        pcb_body, 'PCB', cq.Color('green3'), location=cq.Location((0, 0, pcb_standoff))
    )

    # Pads
    pad = (
        cq.Workplane('XY')
        .lineTo(0, 0.8)
        .lineTo(1.61, 0.8)
        .threePointArc((1.61 + 0.8, 0), (1.61, -0.8))
        .lineTo(0, -0.8)
        .close()
        .extrude(1.6 + 2 * CU_THICKNESS)
    )
    pad = pad.pushPoints([(0, 0), (1.61, 0)]).circle(0.5).cutThruAll()
    for i in range(20):
        assembly.add_body(
            pad,
            f'PAD{i + 1}',
            StepColor.COPPER,
            location=cq.Location((x0, pads_y[i], pcb_standoff)),
        )
        assembly.add_body(
            pad,
            f'PAD{40 - i}',
            StepColor.COPPER,
            location=cq.Location((x1, pads_y[i], pcb_standoff), (0, 0, 180)),
        )
    if bottom_pads:
        for i, pos in enumerate(pads_bottom_xy):
            assembly.add_body(
                pad,
                f'DBGPAD{i + 1}',
                StepColor.COPPER,
                location=cq.Location((pos[0], pos[1], pcb_standoff), (0, 0, 90)),
            )

    # USB connector
    usb_height = 4.0  # guessed
    usb = cq.Workplane('XY').box(8.0, 5.8, usb_height, centered=(True, False, False))
    assembly.add_body(
        usb,
        'USB',
        StepColor.LEAD_SMT,
        location=cq.Location((0, 51.0 / 2 + 1.3 - 5.8, pcb_standoff + 1.6)),
    )

    if headers:
        # Pin sockets
        dx = 8.89
        dy = -23.9 if bottom_pads else y1 + 19.8
        bottom_x = 0 if bottom_pads else x1 - 7.3759
        header = cq.Workplane('XY').box(
            2.5, 2.5 + 48.26, pcb_standoff, centered=(True, True, False)
        )
        for x, n in [(-dx, 'HEADER1'), (dx, 'HEADER2')]:
            assembly.add_body(header, n, StepColor.IC_BODY, location=cq.Location((x, 0, 0)))
        header = cq.Workplane('XY').box(2.5 + 5.08, 2.5, pcb_standoff, centered=(True, True, False))
        assembly.add_body(
            header,
            'HEADER3',
            StepColor.IC_BODY,
            location=cq.Location((bottom_x, dy, 0)),
        )
        pin = cq.Workplane('XY').box(
            0.7, 0.7, pcb_standoff + 2 * StepConstants.THT_LEAD_SOLDER_LENGTH
        )

        # Pins
        for i in range(20):
            assembly.add_body(
                pin,
                f'PIN{i + 1}',
                StepColor.LEAD_THT,
                location=cq.Location((-dx, pads_y[i], pcb_standoff / 2)),
            )
            assembly.add_body(
                pin,
                f'PIN{40 - i}',
                StepColor.LEAD_THT,
                location=cq.Location((dx, pads_y[i], pcb_standoff / 2)),
            )
        for i, x in enumerate([bottom_x - 2.54, bottom_x, bottom_x + 2.54]):
            assembly.add_body(
                pin,
                f'DBGPIN{i + 1}',
                StepColor.LEAD_THT,
                location=cq.Location((x, dy, pcb_standoff / 2)),
            )

    out_path = path.join('out_3d', 'modules', 'rpi_pico', f'{name}.step')
    assembly.save(out_path, fused=False)


if __name__ == '__main__':
    # Arduino
    generate_arduino_uno_r3(name='With Uno', with_board=True)
    generate_arduino_uno_r3(name='Only Headers', with_board=False)

    # ESP32 Modules
    generate_esp32_c6_wroom_1()
    generate_esp32_pico_mini_02()

    # Raspberry Pi
    generate_rpi_pico('Pico (SMD)', bottom_pads=True, headers=False)
    generate_rpi_pico('Pico (THT)', bottom_pads=True, headers=True)
    generate_rpi_pico('Pico W (SMD)', bottom_pads=False, headers=False)
    generate_rpi_pico('Pico W (THT)', bottom_pads=False, headers=True)