"""Parametric OpenCascade teaching model for an integrated hypersonic flowpath.

This script constructs editable boundary-representation solids in CadQuery.
All coordinates are arbitrary visualization units stored as millimetres only
for CAD interoperability; none are flight-vehicle or engine design values.
The component order is informed by cited public NASA descriptions, while the
specific shapes and proportions are illustrative and unvalidated.

Rebuild with: python build_hypersonic_flowpath_cad.py
Requires: cadquery==2.7.0, matplotlib
"""
from __future__ import annotations

import argparse
import csv
from pathlib import Path

import cadquery as cq
import matplotlib

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


# Arbitrary display geometry, in mm for CAD exchange only. The parameters do
# not represent a dimensioned aircraft, combustor or tested MHD device.
P = {
    "duct_half_width": 34.0,
    "sidewall_offset_y": -33.0,
    "sidewall_half_thickness": 2.0,
    "skin": 6.0,
    "flow_x": (0.0, 220.0, 335.0, 450.0, 600.0),
    "electrode_length": 42.0,
    "electrode_y_width": 42.0,
    "power_unit": (66.0, 38.0, 22.0),
    "wire_radius": 2.0,
}

COLORS = {
    "airframe": (0.22, 0.42, 0.63),
    "flowpath": (0.22, 0.74, 0.88),
    "combustor": (1.00, 0.43, 0.24),
    "nozzle": (0.78, 0.43, 0.95),
    "electrode": (0.95, 0.73, 0.27),
    "power": (1.00, 0.29, 0.75),
    "conditioning": (0.35, 0.55, 0.98),
}


def prism_xz(points: list[tuple[float, float]], half_width: float,
             y_offset: float = 0.0) -> cq.Shape:
    """Extrude a closed X-Z profile symmetrically about its Y centre."""
    return (
        cq.Workplane("XZ")
        .polyline(points)
        .close()
        .extrude(half_width, both=True)
        .translate((0, y_offset, 0))
        .val()
    )


def box_at(length: float, width: float, height: float,
           center: tuple[float, float, float]) -> cq.Shape:
    return cq.Workplane("XY").box(length, width, height).translate(center).val()


def rod_between(start: tuple[float, float, float],
                end: tuple[float, float, float], radius: float) -> cq.Shape:
    a, b = cq.Vector(*start), cq.Vector(*end)
    v = b - a
    return cq.Solid.makeCylinder(radius, v.Length, a, v.normalized())


def build_parts() -> list[tuple[str, str, cq.Shape, tuple[float, float, float]]]:
    w = P["duct_half_width"]
    skin = P["skin"]
    parts = [
        ("Forebody compression surface", "Airframe ramp", prism_xz(
            [(0, 0), (220, 34), (220, 40), (0, 6)], w), COLORS["airframe"]),
        ("Inlet cowl and upper ramp", "Airframe cowl", prism_xz(
            [(112, 88), (112, 94), (220, 59), (220, 53)], w), COLORS["airframe"]),
        ("Isolator floor", "Flowpath", prism_xz(
            [(220, 34), (335, 34), (335, 34 + skin), (220, 34 + skin)], w), COLORS["flowpath"]),
        ("Isolator roof", "Flowpath", prism_xz(
            [(220, 53), (335, 53), (335, 53 + skin), (220, 53 + skin)], w), COLORS["flowpath"]),
        ("Combustor floor", "Combustor", prism_xz(
            [(335, 34), (450, 27), (450, 27 + skin), (335, 34 + skin)], w), COLORS["combustor"]),
        ("Combustor roof", "Combustor", prism_xz(
            [(335, 53), (450, 65), (450, 65 + skin), (335, 53 + skin)], w), COLORS["combustor"]),
        ("Expansion nozzle lower wall", "Nozzle", prism_xz(
            [(450, 27), (600, 8), (600, 8 + skin), (450, 27 + skin)], w), COLORS["nozzle"]),
        ("Expansion nozzle upper wall", "Nozzle", prism_xz(
            [(450, 65), (600, 96), (600, 96 + skin), (450, 65 + skin)], w), COLORS["nozzle"]),
    ]

    # A far-side wall gives depth while leaving the near side open as a cutaway.
    side_profile = [(220, 39), (335, 39), (450, 33), (600, 14),
                    (600, 102), (450, 71), (335, 59), (220, 59)]
    parts.append(("Far-side duct wall (cutaway)", "Airframe", prism_xz(
        side_profile, P["sidewall_half_thickness"], P["sidewall_offset_y"]), COLORS["airframe"]))

    # Surface-mounted pairs are symbolic electrode regions, not hardware detail.
    gen_x0, gen_x1 = 250.0, 250.0 + P["electrode_length"]
    ew = P["electrode_y_width"] / 2
    parts.extend([
        ("MHD generator upper electrode region", "MHD concept", prism_xz(
            [(gen_x0, 49), (gen_x1, 49), (gen_x1, 52), (gen_x0, 52)], ew), COLORS["electrode"]),
        ("MHD generator lower electrode region", "MHD concept", prism_xz(
            [(gen_x0, 40), (gen_x1, 40), (gen_x1, 43), (gen_x0, 43)], ew), COLORS["electrode"]),
    ])

    acc_x0, acc_x1 = 475.0, 475.0 + P["electrode_length"]
    parts.extend([
        ("MHD accelerator upper electrode region", "MHD concept", prism_xz(
            [(acc_x0, 76), (acc_x1, 81), (acc_x1, 84), (acc_x0, 79)], ew), COLORS["electrode"]),
        ("MHD accelerator lower electrode region", "MHD concept", prism_xz(
            [(acc_x0, 22), (acc_x1, 18), (acc_x1, 21), (acc_x0, 25)], ew), COLORS["electrode"]),
    ])

    # A remote, simplified electrical link communicates energy transfer only;
    # it is not a realizable power-conditioning or control-system design.
    px, py, pz = 362.0, 0.0, 112.0
    pl, pw, ph = P["power_unit"]
    parts.append(("Power-conditioning placeholder", "MHD concept", box_at(
        pl, pw, ph, (px, py, pz)), COLORS["conditioning"]))
    gen_port = (gen_x1, 0.0, 96.0)
    unit_in = (px - pl / 2, 0.0, pz)
    unit_out = (px + pl / 2, 0.0, pz)
    acc_port = (acc_x0, 0.0, 104.0)
    parts.extend([
        ("Generator-to-conditioning electrical link", "Power path", rod_between(
            gen_port, unit_in, P["wire_radius"]), COLORS["power"]),
        ("Conditioning-to-accelerator electrical link", "Power path", rod_between(
            unit_out, acc_port, P["wire_radius"]), COLORS["power"]),
    ])
    return parts


def render(parts, output_dir: Path, view: str, camera: tuple[float, float]) -> Path:
    mobile = view == "mobile"
    fig = plt.figure(figsize=(5.2, 5.8) if mobile else (12.0, 8.0), facecolor="#08162e")
    ax = fig.add_subplot(111, projection="3d", facecolor="#08162e")
    ax.set_proj_type("ortho")
    for name, _group, shape, color in parts:
        # The CAD includes a far wall to complete the STEP assembly. Omit it in
        # these views so the cutaway exposes the flow channel and electrodes.
        if name == "Far-side duct wall (cutaway)":
            continue
        vertices, triangles = shape.tessellate(0.8, 0.25)
        xyz = [(v.x, v.y, v.z) for v in vertices]
        faces = [[xyz[i] for i in face] for face in triangles]
        collection = Poly3DCollection(faces, facecolor=color, edgecolor="none", alpha=0.97)
        ax.add_collection3d(collection)
    if view in {"section", "mobile"}:
        ax.set_xlim(-20, 620)
        ax.set_ylim(-60, 60)
        ax.set_zlim(-8, 130)
        ax.view_init(elev=0, azim=-90)
        ax.set_box_aspect((640, 2, 235 if mobile else 138))
        ax.set_position([0.025, 0.27, 0.95, 0.50] if mobile else [0.035, 0.20, 0.93, 0.62])
    else:
        ax.set_xlim(-15, 615)
        ax.set_ylim(-55, 55)
        ax.set_zlim(-8, 132)
        ax.view_init(elev=camera[0], azim=58)
        ax.set_box_aspect((630, 110, 140))
        ax.set_position([0.035, 0.20, 0.93, 0.62])
    ax.set_axis_off()
    if mobile:
        fig.text(0.06, 0.94, "HYPERSONIC FLOWPATH · CAD CUTAWAY", color="#f3f7ff",
                 fontsize=18, weight="bold", family="DejaVu Sans")
        fig.text(0.06, 0.89, "AIRFLOW  →   SCHEMATIC · NOT TO SCALE", color="#7bd5ee",
                 fontsize=12, family="DejaVu Sans")
        fig.text(0.06, 0.83, "MHD generator here: upstream, as in a NASA cycle study.",
                 color="#f3c66a", fontsize=10.5, weight="bold", family="DejaVu Sans")
        fig.text(0.06, 0.79, "HVEPS tested a different setup: hot exhaust downstream.",
                 color="#f3c66a", fontsize=10.5, family="DejaVu Sans")
    else:
        fig.text(0.055, 0.94, "INTEGRATED HYPERSONIC FLOWPATH", color="#f3f7ff",
                 fontsize=21, weight="bold", family="DejaVu Sans")
        fig.text(0.057, 0.895, "PARAMETRIC CAD CUTAWAY  /  CONCEPTUAL LAYOUT  /  NOT VALIDATED",
                 color="#7bd5ee", fontsize=11, family="DejaVu Sans")
        fig.text(0.057, 0.852,
                 "AIRFLOW  →     MHD GENERATOR SHOWN UPSTREAM FOR A NASA CYCLE STUDY; HVEPS TESTED HOT EXHAUST DOWNSTREAM",
                 color="#f3c66a", fontsize=9.4, weight="bold", family="DejaVu Sans")
    legend = [
        Patch(facecolor=COLORS["airframe"], label="Forebody & cowl"),
        Patch(facecolor=COLORS["flowpath"], label="Inlet / isolator"),
        Patch(facecolor=COLORS["combustor"], label="Combustor"),
        Patch(facecolor=COLORS["nozzle"], label="Expansion nozzle"),
        Patch(facecolor=COLORS["electrode"], label="MHD concept regions"),
        Patch(facecolor=COLORS["conditioning"], label="Power-unit placeholder"),
        Patch(facecolor=COLORS["power"], label="Illustrative electrical path"),
    ]
    if mobile:
        fig.text(0.06, 0.17, "Blue: body   ·   Teal: inlet   ·   Orange: combustor",
                 color="#d5e1f2", fontsize=9.6, family="DejaVu Sans")
        fig.text(0.06, 0.13, "Purple: nozzle   ·   Gold: MHD regions   ·   Pink: electrical path",
                 color="#d5e1f2", fontsize=9.2, family="DejaVu Sans")
        fig.text(0.06, 0.075, "Illustrative geometry. No performance calculated.",
                 color="#93a9c9", fontsize=9.2, family="DejaVu Sans")
    else:
        fig.legend(handles=legend, loc="lower center", bbox_to_anchor=(0.5, 0.115), ncol=2,
                   frameon=False, fontsize=11.0, labelcolor="#d5e1f2",
                   handlelength=1.15, columnspacing=1.25)
        fig.text(0.057, 0.061,
                 "Component order is source-informed; all geometry and proportions are illustrative. Not to scale.",
                 color="#c5d2e6", fontsize=11.0, family="DejaVu Sans")
        fig.text(0.057, 0.030,
                 "No aerodynamic, thermal, structural, electrical or vehicle performance is represented.",
                 color="#93a9c9", fontsize=10.5, family="DejaVu Sans")
    out = output_dir / f"hypersonic-flowpath-cad-{view}.png"
    fig.savefig(out, dpi=185, facecolor=fig.get_facecolor(), bbox_inches="tight", pad_inches=0.25)
    plt.close(fig)
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parent)
    args = parser.parse_args()
    out = args.output_dir.resolve()
    out.mkdir(parents=True, exist_ok=True)
    parts = build_parts()
    assembly = cq.Assembly(name="Hypersonic flowpath teaching layout")
    shapes = []
    for name, group, shape, rgb in parts:
        if not shape.isValid():
            raise RuntimeError(f"Invalid CAD solid: {name}")
        if shape.Volume() <= 0:
            raise RuntimeError(f"Zero-volume CAD solid: {name}")
        shapes.append(shape)
        assembly.add(shape, name=name.replace(" ", "_"), color=cq.Color(*rgb), loc=cq.Location())

    step_path = out / "hypersonic-flowpath-parametric-layout.step"
    assembly.export(str(step_path))
    compound = cq.Compound.makeCompound(shapes)
    stl_path = out / "hypersonic-flowpath-parametric-layout.stl"
    compound.exportStl(str(stl_path), tolerance=0.35, angularTolerance=0.18)

    rows = [(name, group) for name, group, _shape, _color in parts]
    with (out / "hypersonic-flowpath-cad-parts.csv").open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["component", "model_group", "interpretation"])
        for name, group in rows:
            status = "illustrative layout only" if group in {"MHD concept", "Power path"} else "source-informed component order; illustrative shape"
            writer.writerow([name, group, status])

    # Re-import the neutral exchange file to verify that a non-empty solid
    # assembly survived STEP serialization before producing derived views.
    imported = cq.importers.importStep(str(step_path)).val()
    imported_solids = imported.Solids()
    if len(imported_solids) < len(parts):
        raise RuntimeError(f"STEP round-trip lost solids: {len(imported_solids)} / {len(parts)}")
    if not all(s.isValid() and s.Volume() > 0 for s in imported_solids):
        raise RuntimeError("STEP round-trip contains an invalid or empty solid")

    render(parts, out, "isometric", (23, -58))
    render(parts, out, "section", (0, -90))
    render(parts, out, "mobile", (0, -90))
    print(f"CAD_CHECK_OK parts={len(parts)} step_solids={len(imported_solids)} step_bytes={step_path.stat().st_size} stl_bytes={stl_path.stat().st_size}")


if __name__ == "__main__":
    main()