mxpic_forge/build_images.py

import importlib
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

from generate_handbook import DEFAULT_SRC_ROOT


_PRIMITIVE_BUILDER_PATH = Path(__file__).resolve().parent / "tests" / "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)

PRIMITIVES_PACKAGE = _primitive_builder.PRIMITIVES_PACKAGE
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": ("✗", "red", "Failed"),
}
STATUS_FALLBACK_SYMBOLS = {
    "info": ">",
    "generated": "+",
    "skipped": "-",
    "failed": "x",
}


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


def can_print(text: str) -> bool:
    encoding = sys.stdout.encoding or "utf-8"
    if not encoding.lower().replace("-", "").startswith("utf"):
        return text.isascii()

    try:
        text.encode(encoding)
    except UnicodeEncodeError:
        return False
    return True


def status_item(status: str, count: Optional[int] = None) -> str:
    symbol, color, label = STATUS_STYLES[status]
    if not can_print(symbol):
        symbol = STATUS_FALLBACK_SYMBOLS[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 generated images while preserving Markdown and other docs files."""
    if not img_root.exists():
        return

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


def module_name_for(py_file: Path, src_root: Path, package_root: str) -> str:
    relative_path = py_file.relative_to(src_root).with_suffix("")
    return f"{package_root}.{'.'.join(relative_path.parts)}"


def generate_image_for_class(
    target_dir: Path,
    class_name: str,
    component_class: type,
    device_name: Optional[str] = None,
    context: Optional[dict[str, Any]] = None,
) -> str:
    """Instantiate one component class and save its cell image."""
    kwargs = build_kwargs(component_class, device_name, context) if context else {}
    instance = component_class(**kwargs)
    cell = get_cell(instance) if context else getattr(instance, "cell", None)
    if cell is None:
        return "skipped"

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


def generate_component_images(
    img_root: Path = Path("docs/source/images"),
    src_root: Path = DEFAULT_SRC_ROOT / "primitives",
    package_root: str = PRIMITIVES_PACKAGE,
) -> dict[str, int]:
    print(colorize("mxPIC primitive image generation", "bold"))

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

    if not src_root.exists():
        print_status("failed", f"Source directory not found: {src_root}")
        counts["failed"] += 1
        return counts

    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__
        module_suffix = module_name.removeprefix(f"{package_root}.")
        target_dir = img_root / Path(*module_suffix.split(".")[:-1])

        try:
            device_name = safe_name("IMG", index, class_name)
            result = generate_image_for_class(
                target_dir,
                class_name,
                component_class,
                device_name=device_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()