stealth.py

#!/usr/bin/env python3
"""
stealth.py -- single-file faceted flying-wing "stealth bomber" designer.

Drag the sliders (pygame) to shape the wing. The trailing edge is a true B-2
style sawtooth built by the planform-alignment rule: every TE segment is
parallel to one of the two leading edges, so radar spikes collapse onto a few
azimuths (shown on the RCS polar plot). OUTPUT STEALTH lofts the sections in
CadQuery (ruled, so the skin stays faceted) and writes stealth.step.

deps: pip install pygame cadquery
"""
import math
import os
import pygame

# ------------------------------ constants ---------------------------------
RHO_SL = 1.225                    # sea-level air density, kg/m^3
B2_PAYLOAD_KG = 23_000.0          # published B-2 max payload
B2_WING_AREA_M2 = 478.0           # published B-2 wing area
PAYLOAD_PER_M2 = B2_PAYLOAD_KG / B2_WING_AREA_M2   # ~48 kg of bombs per m^2
MK84_KG = 907.0                   # one Mk-84 GP bomb
LAMBDA_X = 0.03                   # X-band radar wavelength, m (10 GHz)
RCS_FLOOR = 0.01                  # diffuse background, m^2 (-20 dBsm)
STEP_PATH = "stealth.step"
G = 9.81

BG, PANEL, EDGE = (16, 18, 22), (28, 31, 38), (90, 200, 160)
INK, DIM, HULL = (225, 228, 232), (130, 136, 146), (52, 56, 64)
AMBER = (235, 170, 90)

# ------------------------------ geometry ----------------------------------
def planform(p):
    """Right-half vertices (y_span, x_aft): nose, pointed tip, angled tip edge,
    sawtooth TE back to the center stinger. Every TE segment AND the tip edge
    are parallel to a leading edge (planform alignment) -- no edge is ever
    normal to the nose or abeam, which is what makes the RCS spikes collapse."""
    h, t = p["span"] / 2.0, math.tan(math.radians(p["sweep"]))
    cr, n = p["root"], int(round(p["serr"]))
    e = min(p["tip"], 0.5 * h)                   # spanwise run of the tip edge
    ys, xs = h - e, h * t + t * e                # tip edge parallel to opposite LE
    a = (ys - (cr - xs) / t) / 2.0               # forward-going TE run
    a = min(max(a, 0.0), ys)                     # degenerate -> straight TE
    b = ys - a                                   # aft-going TE run
    pts, y, x = [(0.0, 0.0), (h, h * t), (ys, xs)], ys, xs
    for i in range(2 * n):                       # tip -> center zigzag
        if i % 2 == 0:
            y, x = y - a / n, x - t * a / n      # parallel to same-side LE
        else:
            y, x = y - b / n, x + t * b / n      # parallel to opposite LE
        pts.append((y, x))
    pts[-1] = (0.0, cr)                          # center stinger, exact
    return pts

def full_polygon(p):
    right = planform(p)
    return right + [(-y, x) for y, x in reversed(right[1:-1])]

def wing_sections(p):
    """(y, x_le, chord) at every spanwise station, center -> tip. Chord goes
    to ~0 at the pointed tip (clamped so the loft wire stays non-degenerate)."""
    t = math.tan(math.radians(p["sweep"]))
    return sorted((y, y * t, max(x - y * t, 0.05)) for y, x in planform(p)[1:])

def pod_specs(p):
    """Top-side bump-outs, B-2 layout: center cockpit hump plus two engine
    nacelles. Each is (x_start, x_end, y_center, half_width, half_height,
    z_center) in meters, proportioned off root chord and span."""
    cr, half = p["root"], p["span"] / 2.0
    yn = 0.18 * half
    return [(0.10 * cr, 0.62 * cr, 0.0, 0.065 * cr, 0.080 * cr, 0.008 * cr),
            (0.20 * cr, 0.80 * cr, yn, 0.048 * cr, 0.060 * cr, 0.012 * cr),
            (0.20 * cr, 0.80 * cr, -yn, 0.048 * cr, 0.060 * cr, 0.012 * cr)]

def pod_scale(s):
    """Cross-section scale along a pod: blunt front (intake), tapered tail."""
    return max(math.sin(math.pi * (0.12 + 0.88 * s)) ** 0.8, 0.06)

# ------------------------------ aero core ---------------------------------
def naca_half_thickness(x, tc):
    """Half-thickness of a symmetric NACA 00XX section at chord fraction x."""
    return 5.0 * tc * (0.2969 * math.sqrt(x) - 0.1260 * x - 0.3516 * x ** 2
                       + 0.2843 * x ** 3 - 0.1036 * x ** 4)  # -0.1036 closes TE

def best_fit_naca(tc, n=40):
    """Least-squares fit of the design thickness onto the NACA 00XX family."""
    xs = [(i + 0.5) / n for i in range(n)]
    target = [naca_half_thickness(x, tc) for x in xs]
    best, best_rms = 1, float("inf")
    for code in range(1, 31):
        cand = [naca_half_thickness(x, code / 100.0) for x in xs]
        rms = math.sqrt(sum((u - v) ** 2 for u, v in zip(target, cand)) / n)
        if rms < best_rms:
            best, best_rms = code, rms
    return best, best_rms          # rms in fraction-of-chord units

def rcs_spikes(p, lam=LAMBDA_X):
    """Physical-optics edge scattering, azimuth cut. Each straight edge throws
    a specular spike along its normal, peak sigma ~ L^2 / lambda, angular
    width ~ lambda / L. Azimuth 0 = nose-on. Returns [(az_deg, dbsm, L)]."""
    poly, spikes = full_polygon(p), []
    for i in range(len(poly)):
        (y0, x0), (y1, x1) = poly[i], poly[(i + 1) % len(poly)]
        dy, dx, L = y1 - y0, x1 - x0, math.hypot(y1 - y0, x1 - x0)
        if L < 0.05:
            continue
        for ny, nx in ((dx, -dy), (-dx, dy)):    # both normals (mono-static)
            az = math.degrees(math.atan2(ny, -nx)) % 360.0
            spikes.append((az, 10 * math.log10(L * L / lam), L))
    return spikes

def rcs_stats(spikes, lam=LAMBDA_X):
    sigma = RCS_FLOOR                            # nose-on: sum sinc^2 tails
    for az, _, L in spikes:
        u = math.pi * L * math.sin(math.radians(min(az, 360 - az))) / lam
        sigma += (L * L / lam) * (math.sin(u) / u) ** 2 if abs(u) > 1e-9 else L * L / lam
    dirs = {round(az) % 360 for az, _, _ in spikes}
    margin = min(min(az, 360 - az) for az, _, _ in spikes)
    return 10 * math.log10(sigma), len(dirs), margin

def analyze(p):
    poly = full_polygon(p)
    area = 0.5 * abs(sum(poly[i][0] * poly[(i + 1) % len(poly)][1]
                         - poly[(i + 1) % len(poly)][0] * poly[i][1]
                         for i in range(len(poly))))
    ar = p["span"] ** 2 / area
    code, rms = best_fit_naca(p["tc"] / 100.0)
    cla = 2.0 * math.pi * ar / (ar + 2.0)        # Helmbold finite-wing slope
    cl = cla * math.radians(p["aoa"])
    lift = 0.5 * RHO_SL * p["speed"] ** 2 * area * cl
    bombs = area * PAYLOAD_PER_M2
    spikes = rcs_spikes(p)
    nose_db, ndirs, margin = rcs_stats(spikes)
    return dict(area=area, ar=ar, naca=code, rms=rms, cl=cl, cla=cla,
                lift=lift, bombs=bombs, mk84=int(bombs // MK84_KG),
                vs_b2=100.0 * bombs / B2_PAYLOAD_KG, spikes=spikes,
                nose_db=nose_db, ndirs=ndirs, margin=margin)

# ------------------------------ cadquery ----------------------------------
def section_points(chord, tc, x_le, n=40):
    """Closed NACA outline in the XZ plane: LE->TE upper, TE->LE lower."""
    xs = [0.5 * (1.0 - math.cos(math.pi * i / n)) for i in range(n + 1)]
    up = [(x_le + x * chord, naca_half_thickness(x, tc) * chord) for x in xs]
    lo = [(x_le + x * chord, -naca_half_thickness(x, tc) * chord) for x in xs]
    return up + lo[-2:0:-1]        # duplicate LE/TE points dropped

def export_step(p, path=STEP_PATH):
    import cadquery as cq          # deferred: the import is slow, UI shouldn't pay
    mm = 1000.0                    # model in mm so STEP consumers read real size
    tc = best_fit_naca(p["tc"] / 100.0)[0] / 100.0
    wp, prev = cq.Workplane("XZ"), 0.0
    for y, x_le, chord in wing_sections(p):
        if y > prev:
            wp = wp.workplane(offset=(y - prev) * mm)
        wp = wp.polyline(section_points(chord * mm, tc, x_le * mm)).close()
        prev = y
    side = wp.loft(ruled=True)     # ruled = flat facets between stations
    bomber = side.union(side.mirror(mirrorPlane="XZ"))
    for x0, x1, yc, w, h, zc in pod_specs(p):    # cockpit + nacelle bump-outs
        pod, prev = cq.Workplane("YZ", origin=(x0 * mm, yc * mm, zc * mm)), x0
        for k in range(6):
            x = x0 + (x1 - x0) * k / 5.0
            f = pod_scale(k / 5.0)
            if x > prev:
                pod = pod.workplane(offset=(x - prev) * mm)
            pod = pod.ellipse(w * f * mm, h * f * mm)
            prev = x
        bomber = bomber.union(pod.loft())
    cq.exporters.export(bomber, path)
    return os.path.abspath(path)

# ------------------------------ wireframe ---------------------------------
def wireframe(p):
    """3D polylines [(x_aft, y_span, z_up), ...] for the whole aircraft:
    NACA rings + stringers on both wings, ellipse rings + stringers per pod."""
    tc = best_fit_naca(p["tc"] / 100.0)[0] / 100.0
    lines, secs = [], wing_sections(p)
    for sign in (1, -1):
        rings = [[(px, sign * y, pz) for px, pz in section_points(c, tc, x_le, n=8)]
                 for y, x_le, c in secs]
        lines += [r + [r[0]] for r in rings]
        lines += [[r[i] for r in rings] for i in range(len(rings[0]))]
    circ = [i * 2.0 * math.pi / 10 for i in range(10)]
    for x0, x1, yc, w, h, zc in pod_specs(p):
        rings = []
        for k in range(6):
            x, f = x0 + (x1 - x0) * k / 5.0, pod_scale(k / 5.0)
            rings.append([(x, yc + w * f * math.cos(a), zc + h * f * math.sin(a))
                          for a in circ])
        lines += [r + [r[0]] for r in rings]
        lines += [[r[i] for r in rings] for i in range(10)]
    return lines

def draw_wireframe(surf, font, rect, p, ang):
    pygame.draw.rect(surf, PANEL, rect, border_radius=8)
    xc = max(x for _, x in planform(p)) / 2.0    # spin about mid-depth
    cy, sy = math.cos(ang), math.sin(ang)
    ce, se = math.cos(math.radians(58)), math.sin(math.radians(58))
    s = (min(rect.w, rect.h) - 36) / p["span"]
    for line in wireframe(p):
        pts = []
        for x, y, z in line:
            a, b = (x - xc) * cy - y * sy, (x - xc) * sy + y * cy
            pts.append((rect.centerx + b * s, rect.centery + 12 + (a * se - z * ce) * s))
        pygame.draw.lines(surf, EDGE, False, pts, 1)
    surf.blit(font.render("WIREFRAME", True, DIM), (rect.x + 12, rect.y + 10))

# ------------------------------ ui widgets --------------------------------
class Slider:
    def __init__(self, key, label, lo, hi, val, unit="", fmt="{:.1f}"):
        self.key, self.label, self.lo, self.hi = key, label, lo, hi
        self.val, self.unit, self.fmt = val, unit, fmt
        self.rect, self.drag = pygame.Rect(0, 0, 0, 0), False

    def handle(self, ev):
        if ev.type == pygame.MOUSEBUTTONDOWN and self.rect.inflate(0, 14).collidepoint(ev.pos):
            self.drag = True
        elif ev.type == pygame.MOUSEBUTTONUP:
            self.drag = False
        if self.drag and ev.type in (pygame.MOUSEMOTION, pygame.MOUSEBUTTONDOWN):
            f = (ev.pos[0] - self.rect.x) / max(1, self.rect.w)
            self.val = self.lo + (self.hi - self.lo) * min(1.0, max(0.0, f))

    def draw(self, surf, font, x, y, w):
        self.rect = pygame.Rect(x, y + 22, w, 6)
        txt = f"{self.label}  {self.fmt.format(self.val)}{self.unit}"
        surf.blit(font.render(txt, True, INK), (x, y))
        pygame.draw.rect(surf, HULL, self.rect, border_radius=3)
        f = (self.val - self.lo) / (self.hi - self.lo)
        pygame.draw.rect(surf, EDGE, (x, y + 22, int(w * f), 6), border_radius=3)
        pygame.draw.circle(surf, INK, (x + int(w * f), y + 25), 8)

# ------------------------------ drawing -----------------------------------
def draw_planform(surf, font, rect, p):
    pygame.draw.rect(surf, PANEL, rect, border_radius=8)
    poly = full_polygon(p)
    depth = max(x for _, x in poly)
    s = min((rect.w - 50) / p["span"], (rect.h - 50) / depth)
    cx, top = rect.centerx, rect.y + (rect.h - depth * s) / 2
    pts = [(cx + y * s, top + x * s) for y, x in poly]
    pygame.draw.polygon(surf, HULL, pts)
    pygame.draw.polygon(surf, EDGE, pts, 2)
    pygame.draw.line(surf, DIM, (cx, rect.y + 12), (cx, rect.bottom - 12))
    surf.blit(font.render("PLANFORM -- every edge parallel to a LE",
                          True, DIM), (rect.x + 12, rect.y + 10))

def draw_rcs(surf, font, rect, st):
    pygame.draw.rect(surf, PANEL, rect, border_radius=8)
    cx, cy = rect.centerx, rect.centery + 10
    R = min(rect.w, rect.h) / 2 - 34
    db_to_r = lambda db: R * max(0.05, min(1.0, (db + 20) / 70.0))
    for db in (-20, 0, 20, 40):
        pygame.draw.circle(surf, HULL, (cx, cy), int(db_to_r(db)), 1)
        surf.blit(font.render(f"{db}", True, DIM), (cx + 3, cy - int(db_to_r(db))))
    for az, db, _ in st["spikes"]:               # 0 deg = nose = up
        r = db_to_r(db)
        end = (cx + r * math.sin(math.radians(az)), cy - r * math.cos(math.radians(az)))
        pygame.draw.line(surf, AMBER, (cx, cy), end, 2)
    surf.blit(font.render("RCS AZIMUTH SPIKES (dBsm, X-band 10 GHz)", True, DIM),
              (rect.x + 12, rect.y + 10))
    surf.blit(font.render("NOSE", True, DIM), (cx - 16, cy - int(R) - 16))

def draw_airfoil(surf, font, rect, code):
    pygame.draw.rect(surf, PANEL, rect, border_radius=8)
    n, w, cy = 60, rect.w - 40, rect.centery + 8
    up, lo = [], []
    for i in range(n + 1):
        x = i / n
        t = naca_half_thickness(x, code / 100.0) * w
        up.append((rect.x + 20 + x * w, cy - t))
        lo.append((rect.x + 20 + x * w, cy + t))
    pygame.draw.polygon(surf, HULL, up + lo[::-1])
    pygame.draw.lines(surf, EDGE, False, up, 2)
    pygame.draw.lines(surf, EDGE, False, lo, 2)
    surf.blit(font.render(f"SECTION  NACA 00{code:02d} (best fit)", True, DIM),
              (rect.x + 12, rect.y + 10))

def draw_stats(surf, font, rect, p, st):
    pygame.draw.rect(surf, PANEL, rect, border_radius=8)
    lines = [
        ("FLIGHT + WEAPONS DEPT", DIM),
        (f"wing area  {st['area']:7.1f} m^2      AR {st['ar']:5.2f}", INK),
        (f"NACA fit   00{st['naca']:02d}  (rms {st['rms'] * 100:.2f}%c)", INK),
        (f"CL @ {p['aoa']:.1f} deg  {st['cl']:6.3f}  (CLa {st['cla']:.2f}/rad)", INK),
        (f"lift @ {p['speed']:3.0f} m/s  {st['lift'] / 1000:6.0f} kN ({st['lift'] / G / 1000:.1f} t)", INK),
        (f"bomb load  {st['bombs']:7.0f} kg = {st['mk84']} x Mk-84 ({st['vs_b2']:.0f}% of B-2)", EDGE),
        (f"  (B-2 ratio: {PAYLOAD_PER_M2:.1f} kg per m^2 of wing)", DIM),
        ("RADAR", DIM),
        (f"nose-on RCS  {st['nose_db']:+.1f} dBsm", AMBER),
        (f"spike azimuths: {st['ndirs']}   nearest to nose: {st['margin']:.0f} deg", AMBER),
    ]
    y = rect.y + 12
    for txt, color in lines:
        surf.blit(font.render(txt, True, color), (rect.x + 14, y))
        y += 24

# ------------------------------ main loop ---------------------------------
def main(max_frames=None):
    pygame.init()
    screen = pygame.display.set_mode((1280, 768))
    pygame.display.set_caption("stealth.py -- faceted flying-wing designer")
    clock = pygame.time.Clock()
    font = pygame.font.SysFont("menlo,consolas,monospace", 14)
    big = pygame.font.SysFont("menlo,consolas,monospace", 19, bold=True)

    sliders = [
        Slider("span", "Wingspan (width)", 10, 60, 52.4, " m"),
        Slider("root", "Root chord (length)", 4, 30, 21.0, " m"),
        Slider("tip", "Tip edge run", 0.5, 8, 2.5, " m"),
        Slider("tc", "Thickness", 5, 20, 11.0, " %c"),
        Slider("sweep", "LE sweep", 15, 60, 33.0, " deg", "{:.0f}"),
        Slider("serr", "TE serrations/side", 1, 4, 2.0, "", "{:.0f}"),
        Slider("speed", "Airspeed", 60, 300, 150.0, " m/s", "{:.0f}"),
        Slider("aoa", "Angle of attack", 0, 10, 4.0, " deg"),
    ]
    button = pygame.Rect(20, 560, 380, 48)
    status, frames, running = "shape the wing, then OUTPUT STEALTH", 0, True

    while running:
        clicked = False
        for ev in pygame.event.get():
            if ev.type == pygame.QUIT:
                running = False
            for s in sliders:
                s.handle(ev)
            if ev.type == pygame.MOUSEBUTTONDOWN and button.collidepoint(ev.pos):
                clicked = True

        p = {s.key: s.val for s in sliders}
        st = analyze(p)

        screen.fill(BG)
        screen.blit(big.render("STEALTH.PY", True, EDGE), (20, 18))
        screen.blit(font.render("faceted flying-wing designer", True, DIM), (160, 23))
        y = 60
        for s in sliders:
            s.draw(screen, font, 20, y, 380)
            y += 56
        hover = button.collidepoint(pygame.mouse.get_pos())
        pygame.draw.rect(screen, EDGE if hover else HULL, button, border_radius=8)
        label = big.render("OUTPUT STEALTH", True, BG if hover else INK)
        screen.blit(label, label.get_rect(center=button.center))
        screen.blit(font.render(status[:52], True, DIM), (20, 622))

        draw_planform(screen, font, pygame.Rect(420, 16, 528, 330), p)
        draw_wireframe(screen, font, pygame.Rect(960, 16, 300, 330), p, frames * 0.02)
        draw_rcs(screen, font, pygame.Rect(420, 358, 330, 394), st)
        draw_airfoil(screen, font, pygame.Rect(762, 358, 498, 120), st["naca"])
        draw_stats(screen, font, pygame.Rect(762, 490, 498, 262), p, st)

        if clicked:
            msg = big.render("LOFTING IN CADQUERY...", True, EDGE)
            screen.blit(msg, msg.get_rect(center=(210, 680)))
            pygame.display.flip()
            try:
                status = "wrote " + export_step(p)
            except Exception as exc:
                status = f"export failed: {exc}"

        pygame.display.flip()
        clock.tick(60)
        frames += 1
        if max_frames and frames >= max_frames:
            running = False
    pygame.quit()

if __name__ == "__main__":
    main()