Atomizer/tools/adaptive-isogrid/src/brain/__main__.py

"""CLI entry point for the Adaptive Isogrid Python Brain standalone pipeline."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from shapely.geometry import Polygon

from src.atomizer_study import DEFAULT_PARAMS
from src.shared.arc_utils import typed_segments_to_polyline
from .density_field import evaluate_density_grid
from .geometry_schema import normalize_geometry_schema
from .triangulation import generate_triangulation
from .pocket_profiles import generate_pockets
from .profile_assembly import assemble_profile, profile_to_json
from .validation import validate_profile


def _load_json(path: Path) -> Dict[str, Any]:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def _merge_params(geometry: Dict[str, Any], params_file: Path | None) -> Dict[str, Any]:
    params = dict(DEFAULT_PARAMS)
    if params_file is not None:
        user_params = _load_json(params_file)
        if not isinstance(user_params, dict):
            raise ValueError("--params must point to a JSON object")
        params.update(user_params)
    raw_thick = geometry.get("thickness")
    if raw_thick is None:
        raw_thick = params.get("thickness", 10.0)
    params["thickness"] = float(raw_thick)
    return params


def _plot_boundary_polyline(geometry: Dict[str, Any], arc_pts: int = 64) -> np.ndarray:
    """Get the sandbox boundary as a clean polyline for plotting.

    For v2 typed segments: densify arcs, keep lines exact.
    For v1: just use the raw outer_boundary vertices — no resampling needed
    since these are the true polygon corners from NX.
    """
    typed = geometry.get("outer_boundary_typed")
    if typed:
        pts = typed_segments_to_polyline(typed, arc_pts=arc_pts)
        if len(pts) >= 3:
            return np.asarray(pts, dtype=float)

    # v1 fallback: use raw boundary vertices directly — they ARE the geometry.
    outer = np.asarray(geometry["outer_boundary"], dtype=float)
    return outer


def _plot_density(geometry: Dict[str, Any], params: Dict[str, Any], out_path: Path, resolution: float = 3.0) -> None:
    X, Y, eta = evaluate_density_grid(geometry, params, resolution=resolution)

    fig, ax = plt.subplots(figsize=(8, 6), dpi=160)
    m = ax.pcolormesh(X, Y, eta, shading="auto", cmap="viridis", vmin=0.0, vmax=1.0)
    fig.colorbar(m, ax=ax, label="Density η")

    outer = _plot_boundary_polyline(geometry)
    ax.plot(np.r_[outer[:, 0], outer[0, 0]], np.r_[outer[:, 1], outer[0, 1]], "w-", lw=1.5)

    for hole in geometry.get("holes", []):
        if hole.get("is_circular", False) and "center" in hole and "diameter" in hole:
            cx, cy = hole["center"]
            r = float(hole["diameter"]) * 0.5
            a = np.linspace(0.0, 2.0 * np.pi, 90, endpoint=True)
            hb = np.column_stack([cx + r * np.cos(a), cy + r * np.sin(a)])
        else:
            hb = np.asarray(hole["boundary"])
        ax.plot(np.r_[hb[:, 0], hb[0, 0]], np.r_[hb[:, 1], hb[0, 1]], "k-", lw=0.9)

    ax.set_aspect("equal", adjustable="box")
    ax.set_title("Density Field Heatmap")
    ax.set_xlabel("x [mm]")
    ax.set_ylabel("y [mm]")
    fig.tight_layout()
    fig.savefig(out_path)
    plt.close(fig)


def _plot_triangulation(geometry: Dict[str, Any], triangulation: Dict[str, Any], out_path: Path, params: Dict[str, Any] = None) -> None:
    """
    Plot the triangulation with:
    - Green: original sandbox boundary
    - Red dashed: inset contour (w_frame offset) — this is what the mesh fills
    - Blue: triangle edges
    - Blue circles: hole boundaries
    - Orange dashed: hole keepout rings
    """
    from shapely.geometry import Polygon as ShapelyPolygon, Point as ShapelyPoint

    verts = triangulation["vertices"]
    tris = triangulation["triangles"]

    outer = _plot_boundary_polyline(geometry)
    plate_poly = ShapelyPolygon(outer)
    w_frame = (params or {}).get("w_frame", 8.0)
    d_keep = (params or {}).get("d_keep", 1.5)
    # Use the exact inner_plate from triangulation if available (same PSLG boundary)
    inner_plate = triangulation.get("inner_plate") or plate_poly.buffer(-w_frame)

    fig, ax = plt.subplots(figsize=(10, 8), dpi=160)

    # Draw all triangle edges
    drawn_edges = set()
    for tri in tris:
        for i in range(3):
            edge = tuple(sorted([tri[i], tri[(i + 1) % 3]]))
            if edge in drawn_edges:
                continue
            drawn_edges.add(edge)
            p1, p2 = verts[edge[0]], verts[edge[1]]
            ax.plot([p1[0], p2[0]], [p1[1], p2[1]],
                    color="#1f77b4", lw=0.5, alpha=0.85)

    # Green: original sandbox boundary
    ax.plot(np.r_[outer[:, 0], outer[0, 0]], np.r_[outer[:, 1], outer[0, 1]],
            color="#228833", lw=1.8, zorder=5, label="Sandbox boundary")

    # Red dashed: inset contour (w_frame)
    if not inner_plate.is_empty:
        polys = [inner_plate] if inner_plate.geom_type == 'Polygon' else list(inner_plate.geoms)
        for poly in polys:
            ix, iy = poly.exterior.xy
            ax.plot(ix, iy, color="#cc3333", lw=1.2, ls="--", zorder=4,
                    label=f"Inset contour (w_frame={w_frame}mm)")

    # Blue hole circles + Orange keepout rings
    for hole in geometry.get("holes", []):
        if hole.get("is_circular", False) and "center" in hole and "diameter" in hole:
            cx, cy = hole["center"]
            r = float(hole["diameter"]) * 0.5
            a = np.linspace(0.0, 2.0 * np.pi, 90, endpoint=True)
            # Hole circle
            hb = np.column_stack([cx + r * np.cos(a), cy + r * np.sin(a)])
            ax.plot(np.r_[hb[:, 0], hb[0, 0]], np.r_[hb[:, 1], hb[0, 1]],
                    "b-", lw=0.8, zorder=3)
            # Keepout ring
            kr = r * (1.0 + d_keep)
            kb = np.column_stack([cx + kr * np.cos(a), cy + kr * np.sin(a)])
            ax.plot(np.r_[kb[:, 0], kb[0, 0]], np.r_[kb[:, 1], kb[0, 1]],
                    color="#ff8800", lw=0.8, ls="--", zorder=3)
        else:
            hb = np.asarray(hole["boundary"])
            ax.plot(np.r_[hb[:, 0], hb[0, 0]], np.r_[hb[:, 1], hb[0, 1]],
                    "b-", lw=0.8, zorder=3)

    # De-duplicate legend entries
    handles, labels = ax.get_legend_handles_labels()
    seen = set()
    unique_handles, unique_labels = [], []
    for h, l in zip(handles, labels):
        if l not in seen:
            seen.add(l)
            unique_handles.append(h)
            unique_labels.append(l)
    ax.legend(unique_handles, unique_labels, loc="upper right", fontsize=8)

    ax.set_aspect("equal", adjustable="box")
    ax.set_title(f"Isogrid Triangulation ({len(tris)} triangles, {len(drawn_edges)} edges)")
    ax.set_xlabel("x [mm]")
    ax.set_ylabel("y [mm]")
    fig.tight_layout()
    fig.savefig(out_path)
    plt.close(fig)


def _plot_final_profile(geometry, pockets, ribbed_plate, out_path: Path, params: Dict[str, Any] = None) -> None:
    """
    Plot the rib profile: green sandbox boundary, pink pocket outlines (clipped),
    blue hole circles. Pockets must respect w_frame inset from boundary.
    Also produces a zoomed corner view as a second subplot.
    """
    from shapely.geometry import Polygon as ShapelyPolygon

    outer = _plot_boundary_polyline(geometry)
    plate_poly = ShapelyPolygon(outer)
    w_frame = (params or {}).get("w_frame", 2.0)
    inner_plate = plate_poly.buffer(-w_frame)
    if inner_plate.is_empty or not inner_plate.is_valid:
        inner_plate = plate_poly

    fig, (ax_full, ax_zoom) = plt.subplots(1, 2, figsize=(14, 6), dpi=160)

    for ax in (ax_full, ax_zoom):
        # Green sandbox boundary
        ax.plot(np.r_[outer[:, 0], outer[0, 0]], np.r_[outer[:, 1], outer[0, 1]],
                color="#228833", lw=1.8, zorder=5)

        # Draw pockets clipped to inner plate (w_frame inset) — outlines only
        for pocket in pockets:
            polyline = pocket.get("polyline", pocket.get("vertices", []))
            pv = np.asarray(polyline)
            if len(pv) < 3:
                continue
            pocket_poly = ShapelyPolygon(pv)
            if not pocket_poly.is_valid:
                pocket_poly = pocket_poly.buffer(0)
            if pocket_poly.is_empty:
                continue
            clipped = inner_plate.intersection(pocket_poly)
            if clipped.is_empty:
                continue
            clip_geoms = [clipped] if clipped.geom_type == "Polygon" else list(getattr(clipped, "geoms", []))
            for cg in clip_geoms:
                if cg.geom_type != "Polygon" or cg.is_empty:
                    continue
                cx, cy = cg.exterior.xy
                ax.fill(cx, cy, color="#ffcccc", alpha=0.25, lw=0.0)
                ax.plot(cx, cy, color="#cc6677", lw=0.9, zorder=4)

        # Blue hole circles
        for hole in geometry.get("holes", []):
            if hole.get("is_circular", False) and "center" in hole and "diameter" in hole:
                cx, cy = hole["center"]
                r = float(hole["diameter"]) * 0.5
                a = np.linspace(0.0, 2.0 * np.pi, 90, endpoint=True)
                hb = np.column_stack([cx + r * np.cos(a), cy + r * np.sin(a)])
            else:
                hb = np.asarray(hole["boundary"])
            ax.plot(np.r_[hb[:, 0], hb[0, 0]], np.r_[hb[:, 1], hb[0, 1]],
                    "b-", lw=0.8, zorder=3)

        ax.set_aspect("equal", adjustable="box")
        ax.set_xlabel("x [mm]")
        ax.set_ylabel("y [mm]")

    # Full view
    num_pockets_shown = sum(
        1 for p in pockets
        if len(p.get("polyline", p.get("vertices", []))) >= 3
        and not inner_plate.intersection(
            ShapelyPolygon(p.get("polyline", p.get("vertices", []))).buffer(0)
        ).is_empty
    )
    ax_full.set_title(f"Full view ({num_pockets_shown} pockets)")

    # Zoomed view on a corner area (top-right where crossovers were worst)
    bounds = plate_poly.bounds  # minx, miny, maxx, maxy
    mid_x = (bounds[0] + bounds[2]) / 2
    ax_zoom.set_xlim(mid_x - 30, bounds[2] + 10)
    ax_zoom.set_ylim(bounds[3] - 60, bounds[3] + 10)
    ax_zoom.set_title("Zoomed: corner area")

    fig.tight_layout()
    fig.savefig(out_path)
    plt.close(fig)


def run_pipeline(geometry_path: Path, params_path: Path | None, output_dir: Path, output_json_name: str = "rib_profile.json") -> Dict[str, Any]:
    raw_geometry = _load_json(geometry_path)
    geometry = normalize_geometry_schema(raw_geometry)
    params = _merge_params(geometry, params_path)

    triangulation = generate_triangulation(geometry, params)
    pockets = generate_pockets(triangulation, geometry, params)
    ribbed_plate = assemble_profile(geometry, pockets, params)
    is_valid, checks = validate_profile(ribbed_plate, params)

    profile_json = profile_to_json(ribbed_plate, pockets, geometry, params)
    profile_json["checks"] = checks
    profile_json["pipeline"] = {
        "geometry_file": str(geometry_path),
        "num_vertices": int(len(triangulation.get("vertices", []))),
        "num_triangles": int(len(triangulation.get("triangles", []))),
        "num_pockets": int(len(pockets)),
        "validation_ok": bool(is_valid),
    }

    output_dir.mkdir(parents=True, exist_ok=True)
    out_json = output_dir / output_json_name
    with out_json.open("w", encoding="utf-8") as f:
        json.dump(profile_json, f, indent=2)

    stem = geometry_path.stem
    _plot_density(geometry, params, output_dir / f"{stem}_density.png")
    _plot_triangulation(geometry, triangulation, output_dir / f"{stem}_triangulation.png", params=params)
    _plot_final_profile(geometry, pockets, ribbed_plate, output_dir / f"{stem}_final_profile.png", params=params)

    summary = {
        "geometry": geometry.get("plate_id", geometry_path.name),
        "output_json": str(out_json),
        "mass_g": checks.get("mass_estimate_g"),
        "num_pockets": len(pockets),
        "validation_ok": is_valid,
        "checks": checks,
    }
    return summary


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Adaptive Isogrid Python Brain standalone runner")
    parser.add_argument("--geometry", required=True, help="Path to geometry JSON")
    parser.add_argument("--params", default=None, help="Optional path to parameters JSON override")
    parser.add_argument("--output-dir", default=".", help="Directory for rib_profile.json and PNGs")
    parser.add_argument("--output-json", default="rib_profile.json", help="Output profile JSON file name")
    return parser.parse_args()


def main() -> None:
    args = _parse_args()
    geometry_path = Path(args.geometry)
    params_path = Path(args.params) if args.params else None
    output_dir = Path(args.output_dir)

    summary = run_pipeline(geometry_path, params_path, output_dir, output_json_name=args.output_json)

    print("=== Adaptive Isogrid Brain Summary ===")
    print(f"Geometry: {summary['geometry']}")
    print(f"Output JSON: {summary['output_json']}")
    print(f"Mass estimate [g]: {summary['mass_g']:.2f}")
    print(f"Number of pockets: {summary['num_pockets']}")
    print(f"Validation OK: {summary['validation_ok']}")
    print("Validation checks:")
    for k, v in summary["checks"].items():
        print(f"  - {k}: {v}")


if __name__ == "__main__":
    main()