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:
    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)

    outer = np.asarray(geometry["outer_boundary"], dtype=float)
    if len(outer) < 4:
        return outer
    # v1 fallback: dense polyline boundaries may encode fillets.
    # Resample uniformly for smoother plotting while preserving the polygon path.
    if np.allclose(outer[0], outer[-1]):
        ring = outer
    else:
        ring = np.vstack([outer, outer[0]])
    seg = np.diff(ring, axis=0)
    seg_len = np.hypot(seg[:, 0], seg[:, 1])
    nonzero = seg_len[seg_len > 1e-9]
    if len(nonzero) == 0:
        return outer
    step = max(float(np.median(nonzero) * 0.5), 0.5)
    cum = np.r_[0.0, np.cumsum(seg_len)]
    total = float(cum[-1])
    if total <= step:
        return outer
    samples = np.arange(0.0, total, step, dtype=float)
    if samples[-1] < total:
        samples = np.r_[samples, total]
    out = []
    j = 0
    for s in samples:
        while j < len(cum) - 2 and cum[j + 1] < s:
            j += 1
        den = max(cum[j + 1] - cum[j], 1e-12)
        t = (s - cum[j]) / den
        p = ring[j] + t * (ring[j + 1] - ring[j])
        out.append([float(p[0]), float(p[1])])
    return np.asarray(out, dtype=float)


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 Delaunay triangulation clipped to the inner frame (w_frame inset).
    Green boundary, blue rib edges, blue hole circles.
    """
    from shapely.geometry import Polygon as ShapelyPolygon, LineString, Point

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

    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 = plt.subplots(figsize=(8, 6), dpi=160)

    # Draw triangle edges clipped to inner plate
    drawn_edges = set()
    for tri in tris:
        centroid = Point(verts[tri].mean(axis=0))
        if not inner_plate.contains(centroid):
            continue
        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]]
            line = LineString([p1, p2])
            clipped = inner_plate.intersection(line)
            if clipped.is_empty:
                continue
            if clipped.geom_type == "LineString":
                cx, cy = clipped.xy
                ax.plot(cx, cy, color="#1f77b4", lw=0.5, alpha=0.85)
            elif clipped.geom_type == "MultiLineString":
                for seg in clipped.geoms:
                    cx, cy = seg.xy
                    ax.plot(cx, cy, color="#1f77b4", lw=0.5, alpha=0.85)

    # 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)

    # 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_title("Constrained Delaunay Triangulation / Rib Pattern")
    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()