Bearing clearance gauge generator

"""
Bearing clearance gauge generator.

Inputs a deep-groove ball bearing designation; emits a STEP file containing
a two-level gauge plate:
  - Section A (low surface): 9 cylinders rising 10 mm; ID-fit shaft test
  - Section B (high surface): 9 pockets; OD-fit housing test

Run:  python3 bearing-clearance.py
Output: <designation>-<unixtimestamp>.step
"""

import os
import sys
import time
import pygame
import cadquery as cq

# ---------------------------------------------------------------------------
# Bearing catalog (deep-groove ball bearings, metric).  All dimensions in mm.
# ---------------------------------------------------------------------------
BEARINGS = {
    "6804": {"ID": 20.0, "OD": 32.0, "W": 7.0},
    "624":  {"ID":  4.0, "OD": 13.0, "W": 5.0},
}

# Offsets vs nominal ID for the 9 A-cylinders.
# Smallest is 0.10 under nominal (very loose), then exact, then progressively
# tighter.  Increments: 0.10 up to +0.20, then 0.05.
A_OFFSETS = [-0.10, 0.00, 0.10, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45]

# Offsets vs nominal OD for the 9 B-pockets.
# -0.05 (tight), then 0.10 increments to +0.15, then +0.20, then 0.05 increments.
B_OFFSETS = [-0.05, 0.05, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45]

# ---------------------------------------------------------------------------
# Geometry constants
# ---------------------------------------------------------------------------
A_BASE_THICK    = 3.0    # A-section base plate thickness
A_CYL_HEIGHT    = 10.0   # cylinders rise 10 mm above A-section top
A_PITCH_GAP     = 2.0    # bearing-OD-to-bearing-OD clearance on A grid
B_RIM           = 4.0    # min wall between B pockets
B_FLOOR         = 3.0    # solid material under deepest pocket
B_DEPTH_OFFSET  = 0.05   # pockets are bearing_W + 0.05 deep
MAX_FOOTPRINT   = 220.0  # printer build plate constraint

# Embossed labels (raised text near each feature)
LABEL_BAND      = 4.0    # extra Y at the plate bottom so bottom-row labels sit on solid plate
LABEL_FONT_A    = 3.0    # A-side label height (mm) — A cells are roomy
LABEL_FONT_B    = 2.0    # B-side label height (mm) — B cells are tighter
LABEL_OFFSET    = 1.0    # gap from feature edge to label edge (mm)
EMBOSS_HEIGHT   = 0.12   # raised height of text above surface (~2-6 print layers)
LABEL_FONTS     = ("Menlo", "Monaco", "Courier New", "Courier", "Arial")

# ---------------------------------------------------------------------------
# Geometry computation
# ---------------------------------------------------------------------------
def compute_geometry(key):
    b = BEARINGS[key]
    ID, OD, W = b["ID"], b["OD"], b["W"]

    # A-section pitch: when seated, adjacent bearings must clear by A_PITCH_GAP.
    pitch_a = OD + A_PITCH_GAP

    # B-section pitch: governed by the largest pocket plus rim wall.
    max_b_d = OD + max(B_OFFSETS)
    pitch_b = max_b_d + B_RIM

    a_w = 3 * pitch_a
    b_w = 3 * pitch_b
    total_x = a_w + b_w
    # Add LABEL_BAND of extra plate at the bottom so bottom-row labels can be
    # embossed below their features without overhanging the plate edge.
    grid_h = max(a_w, b_w)
    total_y = grid_h + LABEL_BAND

    pocket_depth = W + B_DEPTH_OFFSET
    b_thickness  = pocket_depth + B_FLOOR

    return {
        "key": key, "ID": ID, "OD": OD, "W": W,
        "pitch_a": pitch_a, "pitch_b": pitch_b,
        "a_w": a_w, "b_w": b_w,
        "total_x": total_x, "total_y": total_y,
        "grid_h": grid_h, "label_band": LABEL_BAND,
        "pocket_depth": pocket_depth,
        "a_thickness": A_BASE_THICK,
        "b_thickness": b_thickness,
        "a_top_z": A_BASE_THICK,
        "b_top_z": b_thickness,
        "cyl_height": A_CYL_HEIGHT,
        "fits": total_x <= MAX_FOOTPRINT and total_y <= MAX_FOOTPRINT,
    }

def grid_centers(pitch, origin_x, origin_y_bot, grid_height):
    """3x3 grid, row-major from top-left.  Row 0 is the top.

    The grid is centered vertically inside [origin_y_bot, origin_y_bot+grid_height].
    """
    pts = []
    y_pad = (grid_height - 3 * pitch) / 2.0
    for row in range(3):
        for col in range(3):
            cx = origin_x + (col + 0.5) * pitch
            cy = origin_y_bot + grid_height - y_pad - (row + 0.5) * pitch
            pts.append((cx, cy))
    return pts


def emboss_text(workplane, text_str, fontsize, distance):
    """Build embossed text geometry; try a list of monospace fonts in order."""
    last_err = None
    for f in LABEL_FONTS:
        try:
            return workplane.text(
                text_str, fontsize, distance,
                font=f, kind="regular",
                halign="center", valign="center",
            )
        except Exception as e:
            last_err = e
    raise RuntimeError(f"no usable font for label embossing: {last_err}")


def offset_label(off):
    """Format a clearance offset as a fixed-width signed string, e.g. '+0.10', '-0.10'."""
    return f"{off:+.2f}"

# ---------------------------------------------------------------------------
# CadQuery model build
# ---------------------------------------------------------------------------
def build_model(key):
    g = compute_geometry(key)

    # A-side base slab (low level).
    plate = cq.Workplane("XY").box(
        g["total_x"], g["total_y"], g["a_thickness"],
        centered=(False, False, False),
    )

    # B-side raised slab (high level).  Sits beside A, full slab thickness so
    # the part is one continuous block with two distinct top elevations.
    b_block = (
        cq.Workplane("XY")
        .transformed(offset=cq.Vector(g["a_w"], 0, 0))
        .box(g["b_w"], g["total_y"], g["b_thickness"],
             centered=(False, False, False))
    )
    body = plate.union(b_block)

    # A cylinders (additive) — sit on the A-section top surface.
    a_pts = grid_centers(g["pitch_a"], 0.0, LABEL_BAND, g["grid_h"])
    for (cx, cy), off in zip(a_pts, A_OFFSETS):
        d = g["ID"] + off
        cyl = (
            cq.Workplane("XY")
            .transformed(offset=cq.Vector(cx, cy, g["a_top_z"]))
            .circle(d / 2.0)
            .extrude(g["cyl_height"])
        )
        body = body.union(cyl)

    # B pockets (subtractive) — sink down from the B-section top surface.
    b_pts = grid_centers(g["pitch_b"], g["a_w"], LABEL_BAND, g["grid_h"])
    for (cx, cy), off in zip(b_pts, B_OFFSETS):
        d = g["OD"] + off
        pocket = (
            cq.Workplane("XY")
            .transformed(offset=cq.Vector(cx, cy, g["b_top_z"]))
            .circle(d / 2.0)
            .extrude(-g["pocket_depth"])
        )
        body = body.cut(pocket)

    # Embossed labels — raised text below each feature on the section's top
    # surface.  Label positions use the LARGEST feature radius in the section
    # so labels align horizontally per row regardless of feature size.
    max_a_r = (g["ID"] + max(A_OFFSETS)) / 2.0
    a_label_dy = max_a_r + LABEL_OFFSET + LABEL_FONT_A / 2.0
    for (cx, cy), off in zip(a_pts, A_OFFSETS):
        wp = (cq.Workplane("XY")
              .workplane(offset=g["a_top_z"])
              .center(cx, cy - a_label_dy))
        label = emboss_text(wp, offset_label(off),
                            LABEL_FONT_A, EMBOSS_HEIGHT)
        body = body.union(label)

    max_b_r = (g["OD"] + max(B_OFFSETS)) / 2.0
    b_label_dy = max_b_r + LABEL_OFFSET + LABEL_FONT_B / 2.0
    for (cx, cy), off in zip(b_pts, B_OFFSETS):
        wp = (cq.Workplane("XY")
              .workplane(offset=g["b_top_z"])
              .center(cx, cy - b_label_dy))
        label = emboss_text(wp, offset_label(off),
                            LABEL_FONT_B, EMBOSS_HEIGHT)
        body = body.union(label)

    return body, g, a_pts, b_pts

def export_step(key, out_dir="."):
    body, g, _, _ = build_model(key)
    ts = int(time.time())
    fname = f"{key}-{ts}.step"
    path = os.path.join(out_dir, fname)
    cq.exporters.export(body, path)
    return path, g

# ---------------------------------------------------------------------------
# Pygame GUI
# ---------------------------------------------------------------------------
def run_gui():
    pygame.init()
    W, H = 1180, 760
    screen = pygame.display.set_mode((W, H))
    pygame.display.set_caption("Bearing Clearance Gauge — CadQuery generator")
    clock = pygame.time.Clock()

    font  = pygame.font.SysFont("monospace", 13)
    fontB = pygame.font.SysFont("monospace", 13, bold=True)
    head  = pygame.font.SysFont("monospace", 15, bold=True)
    title = pygame.font.SysFont("monospace", 22, bold=True)
    small = pygame.font.SysFont("monospace", 11)

    BG    = (18, 22, 30)
    PANEL = (28, 34, 44)
    LINE  = (60, 70, 85)
    TXT   = (215, 222, 232)
    DIM   = (140, 150, 165)
    OK    = (110, 210, 140)
    BAD   = (240, 100, 100)
    AMB   = (220, 170, 60)
    BLU   = (90, 160, 240)
    BTN   = (45, 60, 80)
    BTNH  = (60, 90, 130)
    A_FILL = (110, 150, 200)
    B_FILL = (90, 110, 140)
    POCKET = (32, 40, 52)
    SHADE  = (40, 48, 60)

    selected = "6804"
    last_msg = ("ready — press [G] or click GENERATE", time.time())

    def text(s, x, y, c=TXT, f=font):
        screen.blit(f.render(s, True, c), (x, y))

    def text_w(s, f=font):
        return f.size(s)[0]

    keys_list = list(BEARINGS.keys())

    running = True
    while running:
        mx, my = pygame.mouse.get_pos()
        click = False
        for ev in pygame.event.get():
            if ev.type == pygame.QUIT:
                running = False
            elif ev.type == pygame.MOUSEBUTTONDOWN and ev.button == 1:
                click = True
            elif ev.type == pygame.KEYDOWN:
                if ev.key == pygame.K_ESCAPE or ev.key == pygame.K_q:
                    running = False
                elif ev.key in (pygame.K_1,):
                    selected = "6804"
                elif ev.key in (pygame.K_2,):
                    selected = "624"
                elif ev.key in (pygame.K_g, pygame.K_RETURN, pygame.K_SPACE):
                    try:
                        path, _ = export_step(selected)
                        last_msg = (f"wrote {os.path.basename(path)}", time.time())
                    except Exception as e:
                        last_msg = (f"export failed: {e}", time.time())

        g = compute_geometry(selected)
        b = BEARINGS[selected]

        screen.fill(BG)

        # ---- header
        text(f"Bearing Clearance Gauge", 20, 14, TXT, title)
        text("design rules: A pitch = OD + 2.0  |  B pitch = max_pocket_OD + 4.0  |  "
             "pocket depth = W + 0.05", 20, 44, DIM, small)

        # ---- bearing selector + generate button
        bx, by = 20, 70
        for i, k in enumerate(keys_list):
            r = pygame.Rect(bx + i*120, by, 110, 30)
            hov = r.collidepoint(mx, my)
            if k == selected:
                col = BLU
            else:
                col = BTNH if hov else BTN
            pygame.draw.rect(screen, col, r, border_radius=4)
            pygame.draw.rect(screen, LINE, r, 1, border_radius=4)
            tcol = (10, 10, 14) if k == selected else TXT
            lab = f"[{i+1}] {k}"
            text(lab, r.x + (r.w - text_w(lab))//2, r.y + 7, tcol)
            if click and hov:
                selected = k

        gen = pygame.Rect(bx + len(keys_list)*120 + 20, by, 180, 30)
        ghov = gen.collidepoint(mx, my)
        pygame.draw.rect(screen, OK if ghov else (60, 130, 90), gen, border_radius=4)
        pygame.draw.rect(screen, LINE, gen, 1, border_radius=4)
        glab = "[G] GENERATE STEP"
        text(glab, gen.x + (gen.w - text_w(glab, fontB))//2, gen.y + 7, (10,14,10), fontB)
        if click and ghov:
            try:
                path, _ = export_step(selected)
                last_msg = (f"wrote {os.path.basename(path)}", time.time())
            except Exception as e:
                last_msg = (f"export failed: {e}", time.time())

        msg, mts = last_msg
        text(msg, gen.right + 16, by + 9, OK if msg.startswith("wrote") else AMB)

        # =====================================================================
        # LEFT PANEL — specs + geometry
        # =====================================================================
        L = pygame.Rect(20, 110, 360, 630)
        pygame.draw.rect(screen, PANEL, L, border_radius=6)
        pygame.draw.rect(screen, LINE, L, 1, border_radius=6)

        y = L.y + 12
        text("BEARING SPECS", L.x + 12, y, AMB, head); y += 22
        rows = [
            ("designation", f"{selected}"),
            ("ID  (bore)",  f"{b['ID']:.3f} mm"),
            ("OD  (outer)", f"{b['OD']:.3f} mm"),
            ("W   (width)", f"{b['W']:.3f} mm"),
        ]
        for k, v in rows:
            text(f"  {k:<12}: {v}", L.x + 12, y); y += 17
        y += 8

        text("DEVICE GEOMETRY", L.x + 12, y, AMB, head); y += 22
        rows = [
            ("footprint",    f"{g['total_x']:.2f} x {g['total_y']:.2f} mm"),
            ("fits 220x220", "YES" if g['fits'] else "NO"),
            ("A pitch",      f"{g['pitch_a']:.3f} mm"),
            ("B pitch",      f"{g['pitch_b']:.3f} mm"),
            ("A region",     f"{g['a_w']:.2f} x {g['total_y']:.2f}"),
            ("B region",     f"{g['b_w']:.2f} x {g['total_y']:.2f}"),
            ("label band",   f"{g['label_band']:.2f} mm (bottom)"),
            ("A base thick", f"{g['a_thickness']:.2f} mm"),
            ("B base thick", f"{g['b_thickness']:.2f} mm"),
            ("cyl height",   f"{g['cyl_height']:.2f} mm"),
            ("pocket depth", f"{g['pocket_depth']:.3f} mm  (W + 0.05)"),
            ("A top  Z",     f"{g['a_top_z']:.2f} mm"),
            ("B top  Z",     f"{g['b_top_z']:.2f} mm"),
            ("cyl  top Z",   f"{g['a_top_z'] + g['cyl_height']:.2f} mm"),
            ("emboss tall",  f"{EMBOSS_HEIGHT:.2f} mm  (~{round(EMBOSS_HEIGHT/0.02)} layers @0.02)"),
            ("emboss font",  f"A={LABEL_FONT_A:.1f} mm  B={LABEL_FONT_B:.1f} mm"),
        ]
        for k, v in rows:
            c = OK if (k == "fits 220x220" and v == "YES") else (BAD if v == "NO" else TXT)
            text(f"  {k:<12}: {v}", L.x + 12, y, c); y += 17
        y += 10

        text("PRINT NOTES", L.x + 12, y, AMB, head); y += 22
        notes = [
            "* recommended layer: 0.20 mm or finer",
            "* perimeters >= 4 on A cylinders",
            "* infill >= 30%, top layers >= 5",
            "* no supports needed (all features +Z)",
            "* orient with A/B faces UP on bed",
        ]
        for n in notes:
            text(n, L.x + 12, y, DIM); y += 16

        # =====================================================================
        # MIDDLE PANEL — A and B size tables
        # =====================================================================
        M = pygame.Rect(L.right + 16, 110, 360, 630)
        pygame.draw.rect(screen, PANEL, M, border_radius=6)
        pygame.draw.rect(screen, LINE, M, 1, border_radius=6)

        y = M.y + 12
        text("A — CYLINDERS  (shaft / inner-race fit)", M.x + 12, y, AMB, head); y += 22
        text("    grid order: row-major, top-left = #1", M.x + 12, y, DIM); y += 18
        text(f"    {'#':<3} {'offset':>7}  {'diam':>9}  {'note':<20}",
             M.x + 12, y, DIM, fontB); y += 16
        a_pts = grid_centers(g["pitch_a"], 0.0, LABEL_BAND, g["grid_h"])
        for i, off in enumerate(A_OFFSETS):
            d = b["ID"] + off
            sgn = "+" if off >= 0 else ""
            note = "exact" if off == 0 else ("loose" if off < 0 else f"+{off:.2f} press")
            text(f"    {i+1:<3} {sgn}{off:6.2f}  {d:7.3f}   {note}",
                 M.x + 12, y); y += 16
        y += 10

        text("B — POCKETS  (housing / outer-race fit)", M.x + 12, y, AMB, head); y += 22
        text("    grid order: row-major, top-left = #1", M.x + 12, y, DIM); y += 18
        text(f"    {'#':<3} {'offset':>7}  {'diam':>9}  {'note':<20}",
             M.x + 12, y, DIM, fontB); y += 16
        for i, off in enumerate(B_OFFSETS):
            d = b["OD"] + off
            sgn = "+" if off >= 0 else ""
            note = "press" if off < 0 else ("slip" if off == 0 else f"+{off:.2f} clear")
            text(f"    {i+1:<3} {sgn}{off:6.2f}  {d:7.3f}   {note}",
                 M.x + 12, y); y += 16
        y += 10

        text("OUTPUT", M.x + 12, y, AMB, head); y += 22
        text(f"  filename : {selected}-<unixts>.step", M.x + 12, y); y += 16
        text(f"  format   : ISO 10303-21 STEP AP214", M.x + 12, y); y += 16
        text(f"  cad      : cadquery {cq.__version__}", M.x + 12, y); y += 16

        # =====================================================================
        # RIGHT PANEL — top-down 2D preview
        # =====================================================================
        R = pygame.Rect(M.right + 16, 110, W - (M.right + 16) - 20, 630)
        pygame.draw.rect(screen, PANEL, R, border_radius=6)
        pygame.draw.rect(screen, LINE, R, 1, border_radius=6)

        text("TOP-DOWN PREVIEW (to scale)", R.x + 12, R.y + 12, AMB, head)

        pad = 30
        avail_w = R.w - 2 * pad
        avail_h = R.h - 70
        scale = min(avail_w / g["total_x"], avail_h / g["total_y"])
        view_w = g["total_x"] * scale
        view_h = g["total_y"] * scale
        ox = R.x + (R.w - view_w) / 2
        oy = R.y + 50

        def to_px(x, y):
            # plate model coords are bottom-left origin; flip Y for screen
            return (ox + x * scale, oy + (g["total_y"] - y) * scale)

        # B section shaded slightly darker (it's "raised")
        ax_px, ay_px = to_px(0, 0)
        bx_px, by_px = to_px(g["a_w"], 0)
        ex_px, ey_px = to_px(g["total_x"], g["total_y"])

        a_rect = pygame.Rect(ax_px, ey_px, g["a_w"]*scale, view_h)
        b_rect = pygame.Rect(bx_px, ey_px, g["b_w"]*scale, view_h)
        pygame.draw.rect(screen, SHADE, a_rect)
        pygame.draw.rect(screen, (50, 60, 75), b_rect)
        pygame.draw.rect(screen, LINE, pygame.Rect(ax_px, ey_px, view_w, view_h), 1)
        pygame.draw.line(screen, LINE, (bx_px, ey_px), (bx_px, ey_px + view_h), 1)

        # section labels
        text("A", a_rect.x + 6, a_rect.y + 4, AMB, head)
        text("B", b_rect.x + 6, b_rect.y + 4, AMB, head)

        max_a_r = (b["ID"] + max(A_OFFSETS)) / 2.0
        a_label_dy = max_a_r + LABEL_OFFSET + LABEL_FONT_A / 2.0
        max_b_r = (b["OD"] + max(B_OFFSETS)) / 2.0
        b_label_dy = max_b_r + LABEL_OFFSET + LABEL_FONT_B / 2.0

        # A cylinders + embossed label markers
        for i, ((cx, cy), off) in enumerate(zip(a_pts, A_OFFSETS)):
            px, py = to_px(cx, cy)
            d = b["ID"] + off
            r_px = (d/2.0) * scale
            pygame.draw.circle(screen, A_FILL, (int(px), int(py)), int(r_px))
            pygame.draw.circle(screen, TXT, (int(px), int(py)), int(r_px), 1)
            text(f"#{i+1}", px - 8, py - 6, (12,16,22), small)
            lpx, lpy = to_px(cx, cy - a_label_dy)
            text(offset_label(off), lpx - 16, lpy - 5, AMB, small)

        # B pockets + embossed label markers
        b_pts = grid_centers(g["pitch_b"], g["a_w"], LABEL_BAND, g["grid_h"])
        for i, ((cx, cy), off) in enumerate(zip(b_pts, B_OFFSETS)):
            px, py = to_px(cx, cy)
            d = b["OD"] + off
            r_px = (d/2.0) * scale
            pygame.draw.circle(screen, POCKET, (int(px), int(py)), int(r_px))
            pygame.draw.circle(screen, B_FILL, (int(px), int(py)), int(r_px), 1)
            text(f"#{i+1}", px - 8, py - 6, TXT, small)
            lpx, lpy = to_px(cx, cy - b_label_dy)
            text(offset_label(off), lpx - 16, lpy - 5, AMB, small)

        # scale bar
        bar_mm = 20.0
        bar_px = bar_mm * scale
        bar_y = R.bottom - 28
        pygame.draw.line(screen, TXT, (R.x + 20, bar_y), (R.x + 20 + bar_px, bar_y), 2)
        pygame.draw.line(screen, TXT, (R.x + 20, bar_y - 4), (R.x + 20, bar_y + 4), 2)
        pygame.draw.line(screen, TXT, (R.x + 20 + bar_px, bar_y - 4),
                          (R.x + 20 + bar_px, bar_y + 4), 2)
        text(f"{bar_mm:.0f} mm  (scale 1:{1/scale:.2f})", R.x + 20, bar_y + 6, DIM, small)

        # footer help
        text("[1]/[2] select bearing   [G]/[Enter]/[Space] generate STEP   [Q]/[Esc] quit",
             20, H - 24, DIM, small)

        pygame.display.flip()
        clock.tick(60)

    pygame.quit()


if __name__ == "__main__":
    # CLI fallback: `python3 bearing-clearance.py 6804` writes a STEP and exits.
    if len(sys.argv) > 1:
        key = sys.argv[1]
        if key not in BEARINGS:
            print(f"unknown bearing '{key}'.  known: {', '.join(BEARINGS)}")
            sys.exit(1)
        path, g = export_step(key)
        print(f"wrote {path}")
        print(f"  footprint: {g['total_x']:.2f} x {g['total_y']:.2f} mm  "
              f"(fits 220^2: {'YES' if g['fits'] else 'NO'})")
        sys.exit(0)
    run_gui()