mxpic_forge/runs/build_images.py

from __future__ import annotations

import importlib.util
import os
from pathlib import Path
import sys
from typing import Any, Optional

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import nazca as nd


ROOT = Path(__file__).resolve().parents[1]
RUNS_DIR = Path(__file__).resolve().parent
OUTPUT_IMAGE_ROOT = ROOT / "docs" / "source" / "images"
PACKAGE_ROOT = "mxpic.components.primitives"

sys.path.insert(0, str(ROOT))

_PRIMITIVE_BUILDER_PATH = RUNS_DIR / "build_all_primitives.py"
_PRIMITIVE_BUILDER_SPEC = importlib.util.spec_from_file_location(
    "mxpic_build_all_primitives",
    _PRIMITIVE_BUILDER_PATH,
)
if _PRIMITIVE_BUILDER_SPEC is None or _PRIMITIVE_BUILDER_SPEC.loader is None:
    raise ImportError(f"Could not load primitive builder: {_PRIMITIVE_BUILDER_PATH}")

_primitive_builder = importlib.util.module_from_spec(_PRIMITIVE_BUILDER_SPEC)
_PRIMITIVE_BUILDER_SPEC.loader.exec_module(_primitive_builder)

bootstrap_technology = _primitive_builder.bootstrap_technology
build_context = _primitive_builder.build_context
build_kwargs = _primitive_builder.build_kwargs
discover_primitive_classes = _primitive_builder.discover_primitive_classes
get_cell = _primitive_builder.get_cell
safe_name = _primitive_builder.safe_name


PALETTE = {
    "WG": "blue",
    "SLAB": "cyan",
    "M1": "#FFD700",
    "M2": "silver",
    "DEEP_TRENCH": "black",
    (1, 0): "darkred",
    (2, 0): "green",
    (1111, 0): "green",
    (63, 30): "#FFD700",
}


USE_COLOR = sys.stdout.isatty() and os.environ.get("NO_COLOR") is None
COLORS = {
    "bold": "\033[1m",
    "cyan": "\033[36m",
    "green": "\033[32m",
    "red": "\033[31m",
    "yellow": "\033[33m",
    "reset": "\033[0m",
}
STATUS_STYLES = {
    "info": (">", "cyan", "Info"),
    "generated": ("+", "green", "Generated"),
    "skipped": ("-", "yellow", "Skipped"),
    "failed": ("x", "red", "Failed"),
}
PIN_NAME_EXCLUDE = {"org"}


def colorize(text: str, color: str) -> str:
    if not USE_COLOR:
        return text
    return f"{COLORS[color]}{text}{COLORS['reset']}"


def status_item(status: str, count: Optional[int] = None) -> str:
    symbol, color, label = STATUS_STYLES[status]
    suffix = f": {count}" if count is not None else ""
    return colorize(f"{symbol} {label}{suffix}", color)


def print_status(status: str, message: str) -> None:
    print(f"{status_item(status)} {message}")


def apply_mxpic_colors() -> None:
    """Apply the mxPIC color palette to Nazca when the layers are available."""
    print_status("info", "Applying mxPIC layer colors...")
    for layer_id, color in PALETTE.items():
        try:
            nd.set_layercolor(layer=layer_id, color=color)
        except Exception:
            continue


def remove_generated_pngs(img_root: Path) -> None:
    """Remove previously generated image artifacts."""
    if not img_root.exists():
        return

    for png_file in img_root.rglob("*.png"):
        png_file.unlink()


def image_subdir_for_module(module_name: str, package_root: str) -> Path:
    module_suffix = module_name.removeprefix(f"{package_root}.")
    parts = module_suffix.split(".")[:-1]
    if not parts:
        return Path()
    return Path(*parts)


def visible_pin_names(cell: nd.Cell) -> list[str]:
    """Return all user-facing cell pins in deterministic cell order."""
    return [name for name in cell.pin if name not in PIN_NAME_EXCLUDE]


def build_image_top_cell(cell: nd.Cell, top_name: str) -> nd.Cell:
    """Build a plotting wrapper that exposes and draws every component pin."""
    pin_names = visible_pin_names(cell)
    with nd.Cell(name=top_name, instantiate=False) as top_cell:
        instance = cell.put(0, 0, 0)
        if pin_names:
            instance.raise_pins(pin_names, pin_names)
            nd.put_stub(pinname=pin_names, pinsize=0.8)

    return top_cell


def generate_image_for_class(
    target_dir: Path,
    class_name: str,
    component_class: type,
    device_name: str,
    top_name: str,
    context: dict[str, Any],
) -> str:
    """Instantiate one component class and save its cell image."""
    kwargs = build_kwargs(component_class, device_name, context)
    instance = component_class(**kwargs)
    top_cell = build_image_top_cell(get_cell(instance), top_name)

    target_dir.mkdir(parents=True, exist_ok=True)
    image_path = target_dir / f"{class_name}.png"
    nd.export_plt(path="", title=top_name, topcells=[top_cell])
    plt.savefig(image_path, bbox_inches="tight", dpi=300)
    plt.close()
    return "generated"


def generate_component_images(
    img_root: Path = OUTPUT_IMAGE_ROOT,
    package_root: str = PACKAGE_ROOT,
) -> dict[str, int]:
    print(colorize("mxPIC primitive image generation", "bold"))

    img_root = Path(img_root)
    counts = {"generated": 0, "skipped": 0, "failed": 0}

    img_root.mkdir(parents=True, exist_ok=True)
    remove_generated_pngs(img_root)
    nd.clear_layout()
    bootstrap_technology()
    apply_mxpic_colors()
    context = build_context()

    for index, component_class in enumerate(discover_primitive_classes(), start=1):
        class_name = component_class.__name__
        module_name = component_class.__module__
        target_dir = img_root / image_subdir_for_module(module_name, package_root)

        try:
            device_name = safe_name("IMG", index, class_name)
            top_name = safe_name("IMGTOP", index, class_name)
            result = generate_image_for_class(
                target_dir,
                class_name,
                component_class,
                device_name=device_name,
                top_name=top_name,
                context=context,
            )
        except Exception as error:
            plt.close()
            print_status("failed", f"{module_name}.{class_name}: {error}")
            counts["failed"] += 1
            continue

        counts[result] += 1
        if result == "generated":
            print_status("generated", str(target_dir / f"{class_name}.png"))
        else:
            print_status("skipped", f"{module_name}.{class_name} has no cell")

    print(
        "\n"
        + colorize("Image generation complete.", "bold")
        + " "
        + status_item("generated", counts["generated"])
        + " | "
        + status_item("skipped", counts["skipped"])
        + " | "
        + status_item("failed", counts["failed"])
    )
    if counts["failed"]:
        raise RuntimeError(f"{counts['failed']} primitive image(s) failed to build")

    return counts


if __name__ == "__main__":
    generate_component_images()