pyGUI/src/pygui/backend/wafer/wafer_layouts.py

"""Load wafer sensor layouts from bundled YAML files.

YAML schema mirrors the replay app's wafer_desc.py format:
  X/Y      — sensor positions in mm relative to x_origin/y_origin
  size     — wafer diameter/edge in mm
  shape    — "round" or "square"
  start_sn — first sensor number (usually 1)
  x_origin — "left" | "right" | "center"
  y_origin — "bottom" | "top" | "center"

Returned Sensor coords are center-origin mm (negative values toward the edge).
"""
from __future__ import annotations

from pathlib import Path

import yaml

from pygui.backend.wafer.zwafer_models import Sensor

_LAYOUTS_DIR = Path(__file__).parent.parent.parent / "assets" / "layouts"


class WaferLayout(list):
    """Subclass of list that carries wafer shape and size metadata."""
    def __init__(self, sensors: list[Sensor], shape: str = "round", size: float = 300.0):
        super().__init__(sensors)
        self.shape = shape
        self.size = size


def resolve_shape_and_size(sensors: list[Sensor], wafer_id: str = "") -> tuple[str, float]:
    """Helper to determine the shape and size of a sensor layout.

    Tries matching by wafer_id prefix, falling back to the number of sensors.
    """
    if wafer_id:
        prefix = wafer_id[0].upper()
        for path in _LAYOUTS_DIR.glob("*.yaml"):
            try:
                data = _load_yaml(path)
                if prefix in data.get("wafers", []):
                    return data.get("shape", "round"), float(data.get("size", 300.0))
            except Exception:
                continue

    # Fallback to sensor count mapping
    n = len(sensors)
    if n == 80:
        return "square", 310.0
    elif n == 65:
        return "round", 300.0
    elif n == 48:
        return "round", 300.0
    elif n == 29:
        return "round", 200.0
    elif n == 22:
        return "round", 200.0

    return "round", 300.0


def _family_name(raw_name: str) -> str:
    return raw_name.replace("wafer", "").strip("_")


def _load_yaml(path: Path) -> dict:
    with path.open(encoding="utf-8") as f:
        return yaml.safe_load(f)


def available_families() -> list[str]:
    return [_family_name(_load_yaml(p)["name"]) for p in _LAYOUTS_DIR.glob("*.yaml")]


def load_layout(family: str) -> WaferLayout:
    for path in _LAYOUTS_DIR.glob("*.yaml"):
        data = _load_yaml(path)
        if _family_name(data["name"]) == family:
            return _to_sensors(data)
    raise KeyError(f"Unknown wafer family: {family!r}")


def load_layout_for_wafer_id(wafer_id: str) -> WaferLayout:
    """Match 'B00108' → bcdwafer by looking up the first char in each YAML's 'wafers' list."""
    prefix = wafer_id[0].upper() if wafer_id else ""
    for path in _LAYOUTS_DIR.glob("*.yaml"):
        data = _load_yaml(path)
        if prefix in data.get("wafers", []):
            return _to_sensors(data)
    raise KeyError(f"No layout found for wafer ID prefix {prefix!r}")


def _to_sensors(data: dict) -> WaferLayout:
    xs: list[float] = data["X"]
    ys: list[float] = data["Y"]
    size: float = float(data["size"])
    shape: str = data.get("shape", "round")
    start_sn: int = data.get("start_sn", 1)
    reverse_x: bool = data.get("reverse_x", False)
    reverse_y: bool = data.get("reverse_y", False)
    x_orig: str = data.get("x_origin", "left")
    y_orig: str = data.get("y_origin", "bottom")
    label_exceptions = data.get("label_exceptions", {}) or {}

    x_shift = {"left": size / 2, "right": -(size / 2), "center": 0.0}.get(x_orig, 0.0)
    y_shift = {"bottom": size / 2, "top": -(size / 2), "center": 0.0}.get(y_orig, 0.0)

    sensors: list[Sensor] = []
    for i, (x_mm, y_mm) in enumerate(zip(xs, ys)):
        if reverse_x:
            x_mm = -x_mm
        if reverse_y:
            y_mm = -y_mm

        x_coord = x_mm - x_shift
        y_coord = y_mm - y_shift

        # Determine label position defaults based on quadrants
        sn = start_sn + i
        side = "right"
        offset_x = 0.0
        offset_y = 0.0

        # Check explicit label exception from layout config
        if sn in label_exceptions:
            exc = label_exceptions[sn]
            if exc:
                side_key = list(exc.keys())[0]
                val = exc[side_key]
                side = side_key
                if isinstance(val, list) and len(val) >= 2:
                    offset_x = float(val[0])
                    offset_y = float(val[1])
        else:
            # Smart default quadrant positioning to keep labels pointing inward
            if x_coord > 30.0 and y_coord > 30.0:
                side = "left"
                offset_y = -0.5
            elif x_coord < -30.0 and y_coord > 30.0:
                side = "right"
                offset_y = -0.5
            elif x_coord < -30.0 and y_coord < -30.0:
                side = "right"
                offset_y = 0.5
            elif x_coord > 30.0 and y_coord < -30.0:
                side = "left"
                offset_y = 0.5
            elif x_coord > 15.0:
                side = "left"
            elif x_coord < -15.0:
                side = "right"

        sensors.append(Sensor(
            label=str(sn),
            x=x_coord,
            y=y_coord,
            side=side,
            offset_x=offset_x,
            offset_y=offset_y,
        ))
    return WaferLayout(sensors, shape=shape, size=size)