"""FEA subprocess worker.

Invoked by fea.py as a separate Python process so gmsh's signal handlers
(which require the main thread of the interpreter) work correctly.

Usage:
    python fea_worker.py <stl_path> <load_N> <axis> <material>
Prints a single line of JSON to stdout with the result.
"""
from __future__ import annotations

import json
import sys
import traceback

import os
import numpy as np


def _mesh_with_gmsh(stl_path: str, mesh_size: float):
    import gmsh
    gmsh.initialize(["", "-noenv"])
    try:
        gmsh.option.setNumber("General.Terminal", 0)
        gmsh.option.setNumber("Mesh.MeshSizeMax", mesh_size)
        gmsh.option.setNumber("Mesh.MeshSizeMin", mesh_size * 0.3)
        gmsh.merge(stl_path)
        gmsh.model.mesh.classifySurfaces(np.pi / 6, True, False, np.pi / 6)
        gmsh.model.mesh.createGeometry()
        surfs = gmsh.model.getEntities(2)
        sl = gmsh.model.geo.addSurfaceLoop([e[1] for e in surfs])
        gmsh.model.geo.addVolume([sl])
        gmsh.model.geo.synchronize()
        gmsh.model.mesh.generate(3)
        node_tags, coords, _ = gmsh.model.mesh.getNodes()
        coords = np.array(coords).reshape(-1, 3)
        tag_to_idx = {int(t): i for i, t in enumerate(node_tags)}
        elem_types, elem_tags, node_ids = gmsh.model.mesh.getElements(3)
        if not elem_types:
            raise RuntimeError("gmsh produced no 3D elements")
        tets = np.array(node_ids[0]).reshape(-1, 4)
        tets = np.vectorize(lambda t: tag_to_idx[int(t)])(tets)
        return coords, tets
    finally:
        gmsh.finalize()


_MATERIALS = {
    "aluminum":  {"E": 69e9,  "nu": 0.33, "rho": 2700.0},
    "steel":     {"E": 210e9, "nu": 0.30, "rho": 7850.0},
    "stainless": {"E": 200e9, "nu": 0.30, "rho": 7950.0},
    "brass":     {"E": 100e9, "nu": 0.34, "rho": 8500.0},
    "titanium":  {"E": 116e9, "nu": 0.34, "rho": 4500.0},
    "pla":       {"E":  3.5e9, "nu": 0.36, "rho": 1240.0},
    "abs":       {"E":  2.3e9, "nu": 0.35, "rho": 1050.0},
    "default":   {"E": 69e9,  "nu": 0.33, "rho": 2700.0},
}


def run(stl_path: str, load_N: float, axis: str, material: str) -> dict:
    # bbox-based mesh size from STL
    with open(stl_path, "rb") as f:
        head = f.read(80); nb = f.read(4)
        if len(nb) < 4:
            return {"error": "STL too short / empty"}
        import struct
        n = struct.unpack("<I", nb)[0]
        pts = []
        for _ in range(min(n, 5000)):
            d = f.read(50)
            if len(d) < 50: break
            v = struct.unpack("<12fH", d)
            pts.extend([v[3:6], v[6:9], v[9:12]])
        pts = np.array(pts, dtype=np.float32)
    bb = np.ptp(pts, axis=0)
    longest = float(bb.max()) if bb.size else 10.0
    mesh_size = max(0.5, longest / 18.0)

    pts3d, tets = _mesh_with_gmsh(stl_path, mesh_size=mesh_size)
    if len(tets) < 10:
        return {"error": "mesh has <10 tetrahedral elements; part too small"}

    from skfem import (MeshTet, Basis, ElementVector, ElementTetP1, asm,
                       condense, solve)
    from skfem.models.elasticity import linear_elasticity, lame_parameters

    mesh = MeshTet(pts3d.T, tets.T)
    e = ElementVector(ElementTetP1())
    basis = Basis(mesh, e)

    mat = _MATERIALS.get(material.lower(), _MATERIALS["default"])
    # Use consistent mm-N-MPa units: geometry is in mm and loads in N, so
    # E must be in MPa (= N/mm²) for the stiffness matrix to come out in
    # N/mm. Displacements then emerge in mm, stresses in MPa directly.
    E_MPa = mat["E"] / 1e6
    lam, mu = lame_parameters(E_MPa, mat["nu"])

    K = asm(linear_elasticity(lam, mu), basis)

    n_nodes = pts3d.shape[0]
    ax = {"X": 0, "Y": 1, "Z": 2}[axis.upper()]
    bb_min = pts3d.min(0); bb_max = pts3d.max(0)
    tol = max((bb_max[ax] - bb_min[ax]) * 0.02, 0.01)
    top_mask = np.abs(pts3d[:, ax] - bb_max[ax]) < tol
    bot_mask = np.abs(pts3d[:, ax] - bb_min[ax]) < tol
    if top_mask.sum() == 0 or bot_mask.sum() == 0:
        return {"error": "could not identify top/bottom faces"}

    f = np.zeros(3 * n_nodes)
    top_idx = np.where(top_mask)[0]
    per_node = -float(load_N) / max(len(top_idx), 1)
    f[3 * top_idx + ax] = per_node

    bot_idx = np.where(bot_mask)[0]
    fixed = np.concatenate([3 * bot_idx, 3 * bot_idx + 1, 3 * bot_idx + 2])

    u = solve(*condense(K, f, D=fixed))

    disp = u.reshape(-1, 3)
    disp_mag = np.linalg.norm(disp, axis=1)
    # mm-N-MPa convention: displacement is in mm directly.
    max_disp_mm = float(disp_mag.max())

    # von Mises stress per tet
    vm = np.zeros(len(tets))
    for i, tet in enumerate(tets):
        P = pts3d[tet]; U = disp[tet]
        A = (P[:3] - P[3]).T
        B = (U[:3] - U[3]).T
        try:
            grad_u = B @ np.linalg.inv(A)
        except np.linalg.LinAlgError:
            continue
        eps = 0.5 * (grad_u + grad_u.T)
        sigma = lam * np.trace(eps) * np.eye(3) + 2 * mu * eps
        s = sigma - np.trace(sigma) / 3.0 * np.eye(3)
        vm[i] = np.sqrt(1.5 * np.sum(s * s))

    # Spread per-tet stress onto nodes (average of incident tets) so we
    # can colour the surface in the viewport.
    node_stress = np.zeros(n_nodes); node_count = np.zeros(n_nodes)
    for ti, tet in enumerate(tets):
        node_stress[tet] += vm[ti]; node_count[tet] += 1
    node_count[node_count == 0] = 1
    node_stress /= node_count

    # Export a colored OBJ of the SURFACE (outer faces only). Each surface
    # vertex gets an RGB colour from a viridis-style colormap based on its
    # von Mises stress. Three.js will load this and render it like a real
    # COMSOL stress plot.
    out_dir = os.path.dirname(os.path.abspath(stl_path))
    base = os.path.splitext(os.path.basename(stl_path))[0]
    obj_path = os.path.join(out_dir, f"{base}_stress.obj")
    try:
        _export_colored_surface(obj_path, pts3d, tets, node_stress, disp)
    except Exception as _e:
        obj_path = None

    return {
        "ok": True,
        "n_nodes": int(n_nodes),
        "n_elems": int(len(tets)),
        "material": material,
        "load_N": float(load_N),
        "axis": axis.upper(),
        "max_disp_mm": float(max_disp_mm),
        "max_stress_MPa": float(vm.max()),
        "min_stress_MPa": float(vm.min()),
        "E_GPa": mat["E"] / 1e9,
        "colored_obj_path": obj_path,
    }


def _export_colored_surface(obj_path, pts3d, tets, node_stress, disp,
                             deform_scale: float = 0.0):
    """Write an OBJ (.obj) file of the outer surface with per-vertex colors
    encoded as 'v X Y Z R G B' (Three.js OBJLoader reads these as vertex
    colours when MeshPhysicalMaterial.vertexColors = true).
    """
    import collections
    # Find outer triangles: a face is "outer" if it appears in only one tet.
    # Each tet has 4 triangular faces.
    face_count = collections.Counter()
    for tet in tets:
        for tri in [(tet[0], tet[1], tet[2]),
                    (tet[0], tet[1], tet[3]),
                    (tet[0], tet[2], tet[3]),
                    (tet[1], tet[2], tet[3])]:
            face_count[tuple(sorted(tri))] += 1
    outer_faces = [f for f, c in face_count.items() if c == 1]
    if not outer_faces:
        return

    # viridis-ish colormap for stress
    s_min = float(node_stress.min()); s_max = float(node_stress.max())
    s_range = max(s_max - s_min, 1e-9)
    def color_for(s):
        t = (s - s_min) / s_range
        t = max(0.0, min(1.0, t))
        # 5-stop perceptual colormap: dark blue -> cyan -> green -> yellow -> red
        stops = [(0.00, (0.15, 0.10, 0.55)),
                 (0.25, (0.10, 0.55, 0.85)),
                 (0.50, (0.20, 0.80, 0.30)),
                 (0.75, (0.95, 0.85, 0.10)),
                 (1.00, (0.90, 0.15, 0.10))]
        for i in range(len(stops) - 1):
            t0, c0 = stops[i]; t1, c1 = stops[i + 1]
            if t <= t1:
                u = (t - t0) / (t1 - t0) if t1 > t0 else 0
                return (c0[0] + u * (c1[0] - c0[0]),
                        c0[1] + u * (c1[1] - c0[1]),
                        c0[2] + u * (c1[2] - c0[2]))
        return stops[-1][1]

    # only export vertices touched by outer faces (keeps file small)
    used = sorted({v for tri in outer_faces for v in tri})
    remap = {old: new + 1 for new, old in enumerate(used)}    # OBJ is 1-based
    with open(obj_path, "w") as f:
        f.write("# chat_cad FEA stress plot — vertex colors encode von Mises stress\n")
        f.write(f"# stress range: {s_min:.3f} - {s_max:.3f} MPa\n")
        for v in used:
            p = pts3d[v]
            # optionally apply deformed-shape scaling
            if deform_scale > 0:
                p = p + disp[v] * deform_scale
            r, g, b = color_for(node_stress[v])
            f.write(f"v {p[0]:.4f} {p[1]:.4f} {p[2]:.4f} {r:.3f} {g:.3f} {b:.3f}\n")
        for tri in outer_faces:
            f.write(f"f {remap[tri[0]]} {remap[tri[1]]} {remap[tri[2]]}\n")


def run_thermal(stl_path: str, t_hot: float, t_cold: float,
                axis: str = "Z") -> dict:
    """Steady-state heat conduction: T_hot on +axis face, T_cold on -axis face.
    Returns max/min temperature and the maximum gradient magnitude.
    """
    with open(stl_path, "rb") as f:
        f.read(80); nb = f.read(4)
        if len(nb) < 4:
            return {"error": "STL too short / empty"}
        import struct
        n = struct.unpack("<I", nb)[0]
        sample_pts = []
        for _ in range(min(n, 5000)):
            d = f.read(50)
            if len(d) < 50: break
            v = struct.unpack("<12fH", d)
            sample_pts.extend([v[3:6], v[6:9], v[9:12]])
        sample_pts = np.array(sample_pts, dtype=np.float32)
    bb = np.ptp(sample_pts, axis=0)
    longest = float(bb.max()) if bb.size else 10.0
    mesh_size = max(0.5, longest / 18.0)

    pts3d, tets = _mesh_with_gmsh(stl_path, mesh_size=mesh_size)
    if len(tets) < 10:
        return {"error": "mesh too small for thermal analysis"}

    from skfem import MeshTet, Basis, ElementTetP1, asm, condense, solve
    from skfem.helpers import dot, grad
    from skfem.models.poisson import laplace

    mesh = MeshTet(pts3d.T, tets.T)
    basis = Basis(mesh, ElementTetP1())
    K = asm(laplace, basis)

    ax = {"X": 0, "Y": 1, "Z": 2}[axis.upper()]
    bb_min = pts3d.min(0); bb_max = pts3d.max(0)
    tol = max((bb_max[ax] - bb_min[ax]) * 0.02, 0.01)
    hot_idx  = np.where(np.abs(pts3d[:, ax] - bb_max[ax]) < tol)[0]
    cold_idx = np.where(np.abs(pts3d[:, ax] - bb_min[ax]) < tol)[0]
    if len(hot_idx) == 0 or len(cold_idx) == 0:
        return {"error": "couldn't identify hot/cold faces"}

    n = pts3d.shape[0]
    T = np.zeros(n)
    T[hot_idx] = float(t_hot)
    T[cold_idx] = float(t_cold)
    fixed_dofs = np.concatenate([hot_idx, cold_idx])

    T_sol = solve(*condense(K, x=T, D=fixed_dofs))

    # Compute per-element gradient magnitude
    grads = np.zeros(len(tets))
    for i, tet in enumerate(tets):
        P = pts3d[tet]; Tv = T_sol[tet]
        A = (P[:3] - P[3]).T
        b = Tv[:3] - Tv[3]
        try:
            g = np.linalg.solve(A.T, b)
        except np.linalg.LinAlgError:
            continue
        grads[i] = np.linalg.norm(g)

    return {
        "ok": True,
        "n_nodes": int(n),
        "n_elems": int(len(tets)),
        "axis": axis.upper(),
        "t_max": float(T_sol.max()),
        "t_min": float(T_sol.min()),
        "grad_max": float(grads.max()),
    }


def run_modal(stl_path: str, material: str = "aluminum",
              n_modes: int = 6) -> dict:
    """Modal analysis: free-free eigenvalue problem on the 3D tet mesh.
    Returns the first N natural frequencies of the part (rigid-body modes
    have very small omega^2 and are filtered out).
    """
    with open(stl_path, "rb") as f:
        f.read(80); nb = f.read(4)
        if len(nb) < 4:
            return {"error": "STL too short"}
        import struct
        n = struct.unpack("<I", nb)[0]
        sample = []
        for _ in range(min(n, 5000)):
            d = f.read(50)
            if len(d) < 50: break
            v = struct.unpack("<12fH", d)
            sample.extend([v[3:6], v[6:9], v[9:12]])
        sample = np.array(sample, dtype=np.float32)
    bb = np.ptp(sample, axis=0)
    longest = float(bb.max()) if bb.size else 10.0
    mesh_size = max(0.5, longest / 18.0)

    pts3d, tets = _mesh_with_gmsh(stl_path, mesh_size=mesh_size)
    if len(tets) < 10:
        return {"error": "mesh too small for modal analysis"}

    from skfem import (MeshTet, Basis, ElementVector, ElementTetP1, asm,
                      condense, BilinearForm)
    from skfem.helpers import dot as _dot
    from skfem.models.elasticity import linear_elasticity, lame_parameters
    from scipy.sparse.linalg import eigsh

    @BilinearForm
    def mass_form(u, v, w):
        return _dot(u, v)

    mesh = MeshTet(pts3d.T, tets.T)
    elem = ElementVector(ElementTetP1())
    basis = Basis(mesh, elem)

    mat = _MATERIALS.get(material.lower(), _MATERIALS["default"])
    lam, mu = lame_parameters(mat["E"], mat["nu"])
    rho = mat.get("rho", 2700.0)  # kg/m^3; aluminum default

    K = asm(linear_elasticity(lam, mu), basis)
    M = asm(mass_form, basis) * rho

    # Free-free modal: no boundary conditions. Skip the 6 rigid-body modes
    # by asking for n_modes + 6 and dropping the lowest 6.
    k = int(n_modes) + 6
    try:
        # 'SM' = smallest magnitude. sigma=0 makes shift-invert robust at 0.
        omega2, _ = eigsh(K, k=k, M=M, sigma=0.0, which="LM", maxiter=5000)
    except Exception as e:
        return {"error": f"eigenvalue solve failed: {e}"}

    omega2 = np.sort(np.real(omega2))
    # drop rigid-body modes (very small eigenvalues)
    elastic = omega2[omega2 > 1e-3][:int(n_modes)]
    freqs_Hz = (np.sqrt(np.abs(elastic)) / (2 * np.pi)).tolist()
    return {
        "ok": True,
        "n_nodes": int(pts3d.shape[0]),
        "n_elems": int(len(tets)),
        "material": material,
        "n_modes": len(freqs_Hz),
        "frequencies_Hz": [float(f) for f in freqs_Hz],
    }


def run_cfd_2d(stl_path: str, inlet_velocity: float = 1.0,
               viscosity: float = 1.0e-3, axis: str = "Z") -> dict:
    """2D steady Stokes flow around the part's XY silhouette.
    The part is treated as a no-slip obstacle in a rectangular channel;
    inlet on -X (u = inlet_velocity), outlet on +X (p = 0), no-slip on
    top/bottom. Returns the maximum velocity magnitude and the inlet-to-
    outlet pressure drop. Real PDE solve via Taylor-Hood elements.

    Limitations: 2D only (XY mid-plane of the part), Stokes regime only
    (Re << 1, no inertia / turbulence). For turbulent or 3D CFD use
    OpenFOAM or similar.
    """
    with open(stl_path, "rb") as f:
        f.read(80); nb = f.read(4)
        if len(nb) < 4: return {"error": "STL too short"}
        import struct
        n = struct.unpack("<I", nb)[0]
        all_verts = []
        for _ in range(n):
            d = f.read(50)
            if len(d) < 50: break
            v = struct.unpack("<12fH", d)
            all_verts.extend([v[3:6], v[6:9], v[9:12]])
        all_verts = np.array(all_verts, dtype=np.float32)

    # 2D silhouette in XY = convex hull of all (x,y) coords
    xy = all_verts[:, :2]
    bb_xy = np.ptp(xy, axis=0)
    longest = float(max(bb_xy.max(), 10.0))
    h_mesh = max(0.5, longest / 25.0)

    # Build a 2D channel: bounding-box + 1.5x padding in X, 1.0x padding in Y
    x0, y0 = xy.min(0); x1, y1 = xy.max(0)
    cx = (x0 + x1) / 2; cy = (y0 + y1) / 2
    dx = (x1 - x0); dy = (y1 - y0)
    chan_x0 = cx - dx * 1.5; chan_x1 = cx + dx * 2.0
    chan_y0 = cy - dy * 1.0; chan_y1 = cy + dy * 1.0
    obstacle_d = float(min(dx, dy)) * 0.5  # treat part as a disc obstacle

    import gmsh
    gmsh.initialize(["", "-noenv"])
    try:
        gmsh.option.setNumber("General.Terminal", 0)
        gmsh.option.setNumber("Mesh.MeshSizeMax", h_mesh)
        # rectangular channel
        p1 = gmsh.model.geo.addPoint(chan_x0, chan_y0, 0, h_mesh)
        p2 = gmsh.model.geo.addPoint(chan_x1, chan_y0, 0, h_mesh)
        p3 = gmsh.model.geo.addPoint(chan_x1, chan_y1, 0, h_mesh)
        p4 = gmsh.model.geo.addPoint(chan_x0, chan_y1, 0, h_mesh)
        l1 = gmsh.model.geo.addLine(p1, p2)
        l2 = gmsh.model.geo.addLine(p2, p3)
        l3 = gmsh.model.geo.addLine(p3, p4)
        l4 = gmsh.model.geo.addLine(p4, p1)
        cl = gmsh.model.geo.addCurveLoop([l1, l2, l3, l4])
        # circular obstacle in middle
        obs = gmsh.model.geo.addPoint(cx, cy, 0, h_mesh * 0.4)
        obs_pts = []
        n_obs = 24
        for i in range(n_obs):
            ang = 2 * math.pi * i / n_obs
            obs_pts.append(gmsh.model.geo.addPoint(
                cx + obstacle_d / 2 * math.cos(ang),
                cy + obstacle_d / 2 * math.sin(ang), 0, h_mesh * 0.4))
        obs_lines = []
        for i in range(n_obs):
            obs_lines.append(gmsh.model.geo.addLine(
                obs_pts[i], obs_pts[(i + 1) % n_obs]))
        obs_cl = gmsh.model.geo.addCurveLoop(obs_lines)
        gmsh.model.geo.addPlaneSurface([cl, obs_cl])
        gmsh.model.geo.synchronize()
        gmsh.model.mesh.generate(2)
        node_tags, coords, _ = gmsh.model.mesh.getNodes()
        coords = np.array(coords).reshape(-1, 3)[:, :2]
        tag_to_idx = {int(t): i for i, t in enumerate(node_tags)}
        elem_types, elem_tags, node_ids = gmsh.model.mesh.getElements(2)
        if not elem_types:
            return {"error": "2D mesh failed"}
        tris = np.array(node_ids[0]).reshape(-1, 3)
        tris = np.vectorize(lambda t: tag_to_idx[int(t)])(tris)
    finally:
        gmsh.finalize()

    if len(tris) < 50:
        return {"error": "2D mesh too small"}

    from skfem import MeshTri, Basis, ElementTriP2, ElementVector, ElementTriP1, asm, condense, solve
    from skfem.helpers import dot, ddot, sym_grad, div, grad
    from skfem.assembly import BilinearForm, LinearForm

    mesh = MeshTri(coords.T, tris.T)
    eu = ElementVector(ElementTriP2())  # velocity
    ep = ElementTriP1()                 # pressure (Taylor-Hood)
    bu = Basis(mesh, eu)
    bp = Basis(mesh, ep)

    @BilinearForm
    def stiff(u, v, w):
        return viscosity * ddot(sym_grad(u), sym_grad(v))

    @BilinearForm
    def coupling(p, v, w):
        return -p * div(v)

    K = asm(stiff, bu, bu)
    B = asm(coupling, bp, bu)
    # assemble combined saddle-point matrix [K B; B^T 0]
    from scipy.sparse import bmat, csr_matrix
    Z = csr_matrix((bp.N, bp.N))
    A = bmat([[K, B], [B.T, Z]]).tocsr()
    rhs = np.zeros(A.shape[0])

    # Boundary conditions on velocity
    # Inlet (x ~ chan_x0): u_x = inlet_velocity, u_y = 0
    # Top + bottom + obstacle: u = 0 (no-slip)
    # Outlet (x ~ chan_x1): natural (do nothing)
    inlet_dofs = bu.get_dofs(lambda x: np.abs(x[0] - chan_x0) < 1e-3 * longest)
    walls = bu.get_dofs(lambda x: (
        (np.abs(x[1] - chan_y0) < 1e-3 * longest) |
        (np.abs(x[1] - chan_y1) < 1e-3 * longest) |
        ((x[0] - cx) ** 2 + (x[1] - cy) ** 2 < (obstacle_d / 2 + h_mesh) ** 2)
    ))
    # pin one pressure dof to fix the constant
    pin_p = bp.get_dofs(lambda x: np.abs(x[0] - chan_x1) < 1e-3 * longest)

    # combine dofs into global numbering (vel dofs first, then pressure)
    nu = bu.N
    u_inlet = inlet_dofs.flatten()
    u_walls = walls.flatten()
    p_pin = pin_p.flatten() + nu

    # Set BC values
    x_bc = np.zeros(A.shape[0])
    # inlet: x-component = inlet_velocity (every-other index in vector basis)
    x_bc[u_inlet[::2]] = float(inlet_velocity)
    fixed = np.concatenate([u_inlet, u_walls, p_pin]).astype(int)
    free = np.setdiff1d(np.arange(A.shape[0]), fixed)

    A_ff = A[free][:, free]
    A_fc = A[free][:, fixed]
    rhs_f = rhs[free] - A_fc @ x_bc[fixed]

    from scipy.sparse.linalg import spsolve
    try:
        x_free = spsolve(A_ff.tocsc(), rhs_f)
    except Exception as e:
        return {"error": f"Stokes solve failed: {e}"}

    x_full = x_bc.copy()
    x_full[free] = x_free
    u_field = x_full[:nu]
    p_field = x_full[nu:]

    # vector basis has u_x and u_y interleaved
    u_mag = np.sqrt(u_field[::2] ** 2 + u_field[1::2] ** 2)
    dp = float(p_field.max() - p_field.min())
    return {
        "ok": True,
        "n_nodes": int(coords.shape[0]),
        "n_tris": int(len(tris)),
        "inlet_velocity": float(inlet_velocity),
        "viscosity": float(viscosity),
        "max_velocity": float(u_mag.max()),
        "pressure_drop": dp,
        "channel_x_range": [float(chan_x0), float(chan_x1)],
        "channel_y_range": [float(chan_y0), float(chan_y1)],
        "obstacle_diameter": float(obstacle_d),
    }


if __name__ == "__main__":
    try:
        mode = sys.argv[1]
        if mode == "thermal":
            stl_path = sys.argv[2]; t_hot = float(sys.argv[3])
            t_cold = float(sys.argv[4])
            axis = sys.argv[5] if len(sys.argv) > 5 else "Z"
            out = run_thermal(stl_path, t_hot, t_cold, axis)
        elif mode == "modal":
            stl_path = sys.argv[2]
            material = sys.argv[3] if len(sys.argv) > 3 else "aluminum"
            n_modes = int(sys.argv[4]) if len(sys.argv) > 4 else 6
            out = run_modal(stl_path, material, n_modes)
        elif mode == "cfd":
            stl_path = sys.argv[2]
            U = float(sys.argv[3]) if len(sys.argv) > 3 else 1.0
            mu = float(sys.argv[4]) if len(sys.argv) > 4 else 1.0e-3
            axis = sys.argv[5] if len(sys.argv) > 5 else "Z"
            out = run_cfd_2d(stl_path, U, mu, axis)
        else:
            # Legacy positional: <stl> <load_N> <axis> [material] -> elasticity
            stl_path = sys.argv[1]
            load_N = float(sys.argv[2]); axis = sys.argv[3]
            material = sys.argv[4] if len(sys.argv) > 4 else "aluminum"
            out = run(stl_path, load_N, axis, material)
    except Exception as e:
        out = {"error": f"{type(e).__name__}: {e}",
               "trace": traceback.format_exc()[-500:]}
    sys.stdout.write(json.dumps(out))