Source code for vibesketch.document.surfaces

"""Editable schematic surface partners and labelled grafts (#140)."""

from __future__ import annotations

import copy
import math
from typing import Any

from pydantic import ValidationError

from vibesketch.document.model import Document, Surface
from vibesketch.render.nanostructures import DepictionRefused
from vibesketch.render.surfaces import anchor_points, canvas_image, surface_of

SurfaceRefused = DepictionRefused




def records(document: Document) -> list[Surface]:
    """Re-derive surfaces and their attached generic identities from the current drawing."""
    from vibesketch.document.generics import attachment_labels, cxsmiles
    from vibesketch.document.ket import read_scene

    generics = {g.cxsmiles: g.id for g in document.generics if g.kind == "labelled"}
    found: list[Surface] = []
    canvas: Any = document.canvas.data
    objects = [
        (n, s)
        for n in canvas.get("root", {}).get("nodes", [])
        if isinstance(n, dict) and (s := surface_of(n)) is not None
    ]
    if not objects:
        return []
    scene = read_scene(document.canvas.data)
    for node, surface in objects:
        if surface.id in {s.id for s in found}:
            raise SurfaceRefused(
                f"duplicate surface id {surface.id}; give copied objects distinct ids"
            )
        for name, (x, y, _, _) in anchor_points(node).items():
            attached = []
            for fragment in scene.fragments:
                mol = fragment.figure
                if mol is None or not attachment_labels(mol):
                    continue
                gid = generics.get(cxsmiles(mol))
                if gid is None:
                    continue
                conf = mol.GetConformer()
                if any(
                    a.HasProp("attachmentLabel")
                    and math.hypot(
                        conf.GetAtomPosition(a.GetIdx()).x - x,
                        conf.GetAtomPosition(a.GetIdx()).y - y,
                    )
                    < 1e-4
                    for a in mol.GetAtoms()
                ):
                    attached.append(gid)
            surface.anchors[name].generics = sorted(set(attached))
        found.append(surface)
    return found





def with_surface(
    document: Document,
    parameters: Surface | dict[str, Any],
    *,
    at: tuple[float, float] | None = None,
    width: float = 8,
    height: float = 3,
) -> Document:
    """Add or replace an object by id. Geometry is drawing units, never a physical model."""
    from vibesketch.document.derive import derive
    from vibesketch.document.formats import _points

    try:
        surface = Surface.model_validate(parameters)
    except ValidationError as error:
        raise SurfaceRefused(str(error)) from None
    canvas: Any = copy.deepcopy(document.canvas.data) or {"root": {"nodes": []}}
    nodes = canvas.setdefault("root", {}).setdefault("nodes", [])
    old = next(
        (
            n
            for n in nodes
            if isinstance(n, dict) and (s := surface_of(n)) is not None and s.id == surface.id
        ),
        None,
    )
    if at is None:
        if old is not None:
            at = (old["boundingBox"]["x"], old["boundingBox"]["y"])
        else:
            points = _points(canvas)
            at = (max(p[0] for p in points) + 2, max(p[1] for p in points)) if points else (0, 0)
    node = canvas_image(surface, {s.id: s for s in document.species}, at, width, height)
    if old is not None:
        nodes[nodes.index(old)] = node
    else:
        nodes.append(node)
    canvas.setdefault("ket_version", "2.0.0")
    updated = document.model_copy(
        update={"canvas": document.canvas.model_copy(update={"data": canvas})}
    )
    return derive(updated).document





def attach_to_surface(
    document: Document, generic: str, surface: str, anchor: str = "top"
) -> Document:
    """Join a labelled generic to a surface anchor using the nanostructure grafting path."""
    from vibesketch.document.nanostructures import attach_to_picture

    return attach_to_picture(document, generic, surface, anchor)