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 .density_field import evaluate_density_grid
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_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 = np.asarray(geometry["outer_boundary"])
    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", []):
        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) -> None:
    verts = triangulation["vertices"]
    tris = triangulation["triangles"]

    fig, ax = plt.subplots(figsize=(8, 6), dpi=160)
    ax.triplot(verts[:, 0], verts[:, 1], tris, color="#1f77b4", lw=0.4, alpha=0.85)

    outer = np.asarray(geometry["outer_boundary"])
    ax.plot(np.r_[outer[:, 0], outer[0, 0]], np.r_[outer[:, 1], outer[0, 1]], "k-", lw=1.6)
    for hole in geometry.get("holes", []):
        hb = np.asarray(hole["boundary"])
        ax.plot(np.r_[hb[:, 0], hb[0, 0]], np.r_[hb[:, 1], hb[0, 1]], "r-", lw=1.0)

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

    outer = np.asarray(geometry["outer_boundary"])
    ax.plot(np.r_[outer[:, 0], outer[0, 0]], np.r_[outer[:, 1], outer[0, 1]], "k-", lw=1.8, label="Outer boundary")

    for pocket in pockets:
        polyline = pocket.get("polyline", pocket.get("vertices", []))
        pv = np.asarray(polyline)
        if len(pv) >= 3:
            ax.fill(pv[:, 0], pv[:, 1], color="#88ccee", alpha=0.35, lw=0.0)

    if ribbed_plate.geom_type == "Polygon":
        geoms = [ribbed_plate]
    else:
        geoms = list(ribbed_plate.geoms)

    for g in geoms:
        x, y = g.exterior.xy
        ax.plot(x, y, color="#228833", lw=1.2)
        for i in g.interiors:
            ix, iy = i.xy
            ax.plot(ix, iy, color="#cc6677", lw=0.9)

    ax.set_aspect("equal", adjustable="box")
    ax.set_title("Final Ribbed Plate Profile with Pockets")
    ax.set_xlabel("x [mm]")
    ax.set_ylabel("y [mm]")
    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]:
    geometry = _load_json(geometry_path)
    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")
    _plot_final_profile(geometry, pockets, ribbed_plate, output_dir / f"{stem}_final_profile.png")

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