"""Build an editable, dimensionless Hall-effect thruster teaching cutaway.

All numeric lengths are arbitrary visualization proportions in mm for CAD
exchange. They are not taken from NASA Figure 4.30 and are not engineering
dimensions. Field and ion-flow arrows exist only in the PNG illustrations;
they are not CAD hardware or computed field/particle solutions.
"""
from __future__ import annotations

import argparse
import csv
import json
import math
import sys
from pathlib import Path

try:
    import cadquery as cq
except ModuleNotFoundError:  # Allow rebuild in the supplied workspace runtime.
    deps = Path(__file__).resolve().parents[2] / "tmp" / "cad-prep-cadquery-2.7-deps"
    sys.path.insert(0, str(deps))
    import cadquery as cq

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from mpl_toolkits.mplot3d.art3d import Poly3DCollection


P = {
    "axis": "+Z downstream along the annular discharge channel",
    "removed_sector_degrees": 90,
    "channel_inner_radius": 12.0,
    "channel_outer_radius": 20.0,
    "channel_axial_start": 0.0,
    "channel_axial_exit": 25.0,
    "inner_wall_inner_radius": 10.5,
    "inner_wall_outer_radius": 12.0,
    "outer_wall_inner_radius": 20.0,
    "outer_wall_outer_radius": 21.5,
    "wall_back_z": -2.0,
    "wall_exit_z": 27.0,
    "outer_case_inner_radius": 21.5,
    "outer_case_outer_radius": 25.0,
    "case_back_z": -7.0,
    "case_exit_z": 27.0,
    "manifold_inner_radius": 11.0,
    "manifold_outer_radius": 21.5,
    "manifold_back_z": -7.0,
    "manifold_front_z": -2.0,
    "anode_inner_radius": 12.0,
    "anode_outer_radius": 20.0,
    "anode_back_z": -2.0,
    "anode_front_z": 0.0,
    "cathode_stem_radius": 1.15,
    "cathode_stem_back_z": -5.0,
    "cathode_stem_front_z": 14.0,
    "cathode_body_radius": 3.1,
    "cathode_body_back_z": 14.0,
    "cathode_body_front_z": 33.0,
    "cathode_emitter_radius": 2.0,
    "cathode_emitter_back_z": 33.0,
    "cathode_emitter_front_z": 37.0,
}

COLORS = {
    "case": (0.50, 0.61, 0.72),
    "inner_wall": (0.74, 0.80, 0.85),
    "outer_wall": (0.42, 0.53, 0.65),
    "manifold": (0.12, 0.68, 0.69),
    "anode": (0.95, 0.57, 0.22),
    "cathode": (0.76, 0.39, 0.80),
    "cathode_tip": (0.96, 0.79, 0.33),
}

LIMIT = "Arbitrary teaching geometry; not an engineering dimension or hardware specification."


def annular_sector(r_outer: float, r_inner: float, z0: float, z1: float) -> cq.Shape:
    """Annular solid with an explanatory 90-degree sector removed."""
    body = (cq.Workplane("XY").circle(r_outer).circle(r_inner)
            .extrude(z1 - z0).translate((0, 0, z0)).val())
    reach = 100.0
    wedge = (cq.Workplane("XY")
             .polyline([(0, 0), (reach, 0), (0, reach)]).close()
             .extrude((z1 - z0) + 2.0)
             .translate((0, 0, z0 - 1.0)).val())
    return body.cut(wedge)


def cylinder(radius: float, z0: float, z1: float) -> cq.Shape:
    return cq.Solid.makeCylinder(radius, z1 - z0, cq.Vector(0, 0, z0), cq.Vector(0, 0, 1))


def build_parts() -> list[tuple[str, str, cq.Shape, tuple[float, float, float]]]:
    p = P
    parts = [
        ("Outer casing", "Outer housing", annular_sector(p["outer_case_outer_radius"], p["outer_case_inner_radius"], p["case_back_z"], p["case_exit_z"]), COLORS["case"]),
        ("Inner channel wall", "Annular discharge chamber", annular_sector(p["inner_wall_outer_radius"], p["inner_wall_inner_radius"], p["wall_back_z"], p["wall_exit_z"]), COLORS["inner_wall"]),
        ("Outer channel wall", "Annular discharge chamber", annular_sector(p["outer_wall_outer_radius"], p["outer_wall_inner_radius"], p["wall_back_z"], p["wall_exit_z"]), COLORS["outer_wall"]),
        ("Rear annular propellant manifold", "Rear manifold", annular_sector(p["manifold_outer_radius"], p["manifold_inner_radius"], p["manifold_back_z"], p["manifold_front_z"]), COLORS["manifold"]),
        ("Annular anode", "Rear electrode", annular_sector(p["anode_outer_radius"], p["anode_inner_radius"], p["anode_back_z"], p["anode_front_z"]), COLORS["anode"]),
        ("Central cathode support stem", "Central cathode configuration", cylinder(p["cathode_stem_radius"], p["cathode_stem_back_z"], p["cathode_stem_front_z"]), COLORS["cathode"]),
        ("Central cathode body", "Central cathode configuration", cylinder(p["cathode_body_radius"], p["cathode_body_back_z"], p["cathode_body_front_z"]), COLORS["cathode"]),
        ("Cathode emitter tip", "Central cathode configuration", cylinder(p["cathode_emitter_radius"], p["cathode_emitter_back_z"], p["cathode_emitter_front_z"]), COLORS["cathode_tip"]),
    ]
    return parts


def hex_color(rgb: tuple[float, float, float]) -> str:
    return "#" + "".join(f"{round(c * 255):02X}" for c in rgb)


def render(parts: list[tuple[str, str, cq.Shape, tuple[float, float, float]]], out: Path,
           view: str) -> Path:
    fig = plt.figure(figsize=(12.0, 8.0), facecolor="#07152b")
    ax = fig.add_subplot(111, projection="3d", facecolor="#07152b")
    ax.set_proj_type("ortho")
    for _name, _role, shape, color in parts:
        vertices, triangles = shape.tessellate(0.28, 0.12)
        xyz = [(v.x, v.y, v.z) for v in vertices]
        faces = [[xyz[i] for i in tri] for tri in triangles]
        # A transparent outer skin keeps the editable assembly intact while
        # letting a static 2-D render expose the nested channel (Matplotlib's
        # 3-D painter does not perform per-pixel CAD depth buffering).
        alpha = 0.16 if _name == "Outer casing" else (0.55 if "channel wall" in _name.lower() else 0.98)
        ax.add_collection3d(Poly3DCollection(faces, facecolor=color,
                                             edgecolor="none", linewidth=0,
                                             alpha=alpha))

    # Direction glyphs are explanatory graphics only, not physical parts or
    # computed field maps. At the +X exit location, +Z x +X = +Y for E x B.
    glyph_colors = {"E": "#ff7b6c", "B": "#62c7ff", "drift": "#ffd166", "ion": "#f2f6ff"}
    ax.quiver(30.0, -28.0, 12.0, 0.0, 0.0, 10.0, color=glyph_colors["E"], linewidth=2.6,
              arrow_length_ratio=0.22, length=1.0)
    ax.quiver(30.0, -28.0, 24.0, 6.0, 0.0, 0.0, color=glyph_colors["B"], linewidth=2.6,
              arrow_length_ratio=0.22, length=1.0)
    ax.quiver(36.0, -28.0, 24.0, 0.0, 6.0, 0.0, color=glyph_colors["drift"], linewidth=2.6,
              arrow_length_ratio=0.22, length=1.0)

    # A few schematic ion-direction arrows outside the exit. They are not
    # trajectories and encode no speed, density, or divergence.
    for x, y in [(-16.0, -2.0), (2.0, -16.0), (-2.0, 16.0)]:
        ax.quiver(x, y, 25.0, 0.0, 0.0, 7.5, color=glyph_colors["ion"], linewidth=1.9,
                  arrow_length_ratio=0.25, length=1.0)

    # Look down into the open +X/+Y quarter so the annular channel, rear
    # anode/manifold and central cathode are visible together.
    ax.view_init(elev=58 if view == "isometric" else 43,
                 azim=45)
    ax.set_xlim(-40, 44)
    ax.set_ylim(-34, 38)
    ax.set_zlim(-11, 42)
    ax.set_box_aspect((1.0, 1.0, 0.95))
    ax.set_axis_off()
    ax.set_position([0.005, 0.07, 0.73, 0.82])

    title = "HALL-EFFECT THRUSTER · PARAMETRIC CUTAWAY"
    subtitle = "OPEN QUARTER SECTION · ARBITRARY VISUAL PROPORTIONS · NOT TO SCALE"
    fig.text(0.035, 0.955, title, color="#f1f5ff", fontsize=18, weight="bold")
    fig.text(0.037, 0.916, subtitle, color="#6fddf4", fontsize=10.3)

    legend = [
        Patch(facecolor=COLORS["case"], label="Outer casing"),
        Patch(facecolor=COLORS["inner_wall"], label="Inner channel wall"),
        Patch(facecolor=COLORS["outer_wall"], label="Outer channel wall"),
        Patch(facecolor=COLORS["manifold"], label="Rear annular manifold"),
        Patch(facecolor=COLORS["anode"], label="Annular anode"),
        Patch(facecolor=COLORS["cathode"], label="Central cathode · one layout"),
        Patch(facecolor=COLORS["cathode_tip"], label="Emitter tip"),
    ]
    fig.legend(handles=legend, loc="upper left", bbox_to_anchor=(0.75, 0.82),
               frameon=False, fontsize=9.4, labelcolor="#e2ecf7", handlelength=1.2)
    fig.text(0.755, 0.43, "FIELD / FLOW GLYPHS", color="#cbd8e9", fontsize=9.2, weight="bold")
    fig.text(0.755, 0.395, "Red: E axial (+Z)", color="#ffb4a9", fontsize=8.9)
    fig.text(0.755, 0.363, "Blue: B radial (+X example)", color="#62c7ff", fontsize=8.9)
    fig.text(0.755, 0.331, "Yellow: E×B drift (+Y)", color="#ffe2a1", fontsize=8.9)
    fig.text(0.755, 0.299, "White: ion outflow", color="#f2f6ff", fontsize=8.9)
    fig.text(0.755, 0.263, "Arrows are explanatory only;", color="#aebed3", fontsize=8.4)
    fig.text(0.755, 0.237, "not hardware or computed paths.", color="#aebed3", fontsize=8.4)
    fig.text(0.04, 0.127, "Walls are transparent here only to expose the internal parts.", color="#cad7e8", fontsize=8.6)
    fig.text(0.04, 0.095, "NASA Figure 4.30 informed the labels and arrangement; dimensions and geometry here are newly chosen teaching proportions.",
             color="#cad7e8", fontsize=8.6)
    fig.text(0.04, 0.062, "No magnet geometry, circuit, voltage, field strength, operating point, thrust, power, efficiency, or lifetime is modeled.",
             color="#aebed3", fontsize=8.2)
    fig.text(0.04, 0.033, "The 90° cut is an explanatory opening, not a physical aperture. No performance or flight validation is implied.",
             color="#aebed3", fontsize=8.2)

    path = out / ("hall-effect-thruster-isometric.png" if view == "isometric"
                  else "hall-effect-thruster-open-section.png")
    fig.savefig(path, dpi=180, facecolor=fig.get_facecolor(), bbox_inches="tight", pad_inches=0.22)
    plt.close(fig)
    return path


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parent)
    out = parser.parse_args().output_dir.resolve()
    out.mkdir(parents=True, exist_ok=True)
    parts = build_parts()

    assembly = cq.Assembly(name="Generic Hall-effect thruster educational cutaway")
    solids: list[cq.Shape] = []
    for name, role, shape, color in parts:
        if not shape.isValid() or shape.Volume() <= 0:
            raise RuntimeError(f"Invalid or zero-volume solid: {name}")
        solids.append(shape)
        assembly.add(shape, name=name.replace(" ", "_"), color=cq.Color(*color))

    step_path = out / "hall-effect-thruster-cutaway.step"
    stl_path = out / "hall-effect-thruster-cutaway.stl"
    assembly.export(str(step_path))
    cq.Compound.makeCompound(solids).exportStl(str(stl_path), tolerance=0.12,
                                                angularTolerance=0.08)

    # Re-import STEP and require all separately modeled positive-volume solids.
    roundtrip = cq.importers.importStep(str(step_path)).val()
    roundtrip_solids = roundtrip.Solids()
    if len(roundtrip_solids) != len(solids):
        raise RuntimeError(f"STEP roundtrip has {len(roundtrip_solids)} solids; expected {len(solids)}")
    for idx, shape in enumerate(roundtrip_solids):
        if not shape.isValid() or shape.Volume() <= 0:
            raise RuntimeError(f"STEP roundtrip solid {idx} invalid or zero volume")

    csv_path = out / "hall-effect-thruster-components.csv"
    with csv_path.open("w", newline="", encoding="utf-8") as stream:
        writer = csv.writer(stream)
        writer.writerow(["component", "role", "span_x_visual_mm", "span_y_visual_mm",
                         "span_z_visual_mm", "color_hex", "interpretation_limit"])
        for name, role, shape, color in parts:
            b = shape.BoundingBox()
            writer.writerow([name, role, f"{b.xlen:.3f}", f"{b.ylen:.3f}", f"{b.zlen:.3f}",
                             hex_color(color), LIMIT])

    assumptions = {
        "title": "Generic Hall-effect thruster teaching cutaway",
        "source_reference": {
            "document": "NASA SmallSat State of the Art 2026, In-Space Propulsion",
            "figure": "Figure 4.30, physical PDF page 36 / printed page 91",
            "sha256": "a5b2aa5843f4b8c09e2eb855f30f4038ad7ff2b588171dee2755932c8253e81c",
            "use": "Figure informed component roles and one central-cathode arrangement; dimensions and CAD geometry are original teaching proportions."
        },
        "coordinate_convention": {"axis": "+Z downstream through channel", "radial": "XY plane",
                                  "open_cutaway": "90-degree sector removed where X >= 0 and Y >= 0; explanatory only"},
        "units": "mm used solely for CAD exchange; all lengths are arbitrary visualization proportions, not scale, tolerances, or engineering dimensions.",
        "parameters": P,
        "component_count": len(parts),
        "included_components": [name for name, *_ in parts],
        "schematic_glyphs": {
            "stored_in_cad": False,
            "items": ["E axial +Z", "B radial +X at the +X exit location", "local E×B electron drift +Y there", "ion outflow +Z"],
            "limits": "PNG arrows are explanatory direction glyphs only; no field map, field magnitude, electron/ion trajectory, speed, or operating point is computed. The +Y E×B direction follows +Z cross +X at the indicated location."
        },
        "not_modeled": ["magnet hardware or its placement", "magnetic/electric field magnitude", "voltage or electrical circuit", "propellant feed internals beyond conceptual annular manifold", "plasma or ionization", "particle trajectories", "thrust", "power", "efficiency", "thermal behavior", "erosion", "lifetime", "structural strength", "manufacturing tolerances", "flight or performance validation"],
        "configuration_boundary": "Central cathode is one arrangement shown in the source example; other Hall-thruster designs may locate the cathode externally. The quarter cut is a visualization opening, not a physical aperture. No acceleration grids are included.",
        "validation": {"model_solids": len(solids), "all_original_solids_valid_positive_volume": True,
                       "step_roundtrip_solids": len(roundtrip_solids), "all_step_roundtrip_solids_valid_positive_volume": True,
                       "step_bytes": step_path.stat().st_size, "stl_bytes": stl_path.stat().st_size}
    }
    (out / "hall-effect-thruster-assumptions.json").write_text(
        json.dumps(assumptions, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    render(parts, out, "isometric")
    render(parts, out, "open-section")
    print(f"Built {len(solids)} valid CAD solids; STEP roundtrip validated {len(roundtrip_solids)} solids.")


if __name__ == "__main__":
    main()