Source code for oowl.strings.base

"""Base representation for ordered turbine-string guide geometry."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import numpy as np
from shapely.affinity import scale as shp_scale
from shapely.affinity import translate as shp_translate
from shapely.geometry import LineString, MultiLineString

from .._repr import _ArgumentRepr

__all__ = ["StringGuide"]

if TYPE_CHECKING:  # pragma: no cover
    import matplotlib.axes

    from .composites import Composite
    from .modifiers import StringModifier


@dataclass(frozen=True)
class _LogicalString:
    """One ordered string, possibly split into segments by breaks."""

    guide_type: str
    segments: tuple[LineString, ...]

    @property
    def length(self) -> float:
        """Total arc length of this logical string."""
        return float(sum(segment.length for segment in self.segments))


def _segment_geometry(
    strings: tuple[_LogicalString, ...],
) -> LineString | MultiLineString:
    """Every segment of ``strings`` as one Shapely geometry."""
    segments = [segment for logical in strings for segment in logical.segments]
    if len(segments) == 1:
        return segments[0]
    return MultiLineString([segment.coords for segment in segments])




class StringGuide(_ArgumentRepr):
    """One or more ordered logical strings made from continuous segments.

    A guide only fixes the shape and relative size of turbine strings.
    :class:`~oowl.layouts.String` places turbines along it and sets the scale.
    Guides combine with :class:`~oowl.strings.Multi` and
    :class:`~oowl.strings.Parallel` and change with :meth:`modify`. A random
    guide or modifier given no ``seed`` draws one, and a
    :class:`~oowl.layouts.String` layout redraws it from its own seed, so
    every sampled farm gets its own guide. A given ``seed`` keeps its shape.
    """

    #: Whether every random part of the guide has a seed its caller gave.
    _pinned_seed = True

    def __init__(self, strings: tuple[_LogicalString, ...]):
        strings = tuple(strings)
        if not strings:
            raise ValueError("a string guide requires at least one logical string")
        for logical_string in strings:
            if not logical_string.segments:
                raise ValueError("each logical string requires at least one segment")
            for segment in logical_string.segments:
                coordinates = np.asarray(segment.coords, dtype=float)
                if segment.length <= 0:
                    raise ValueError("string segments must have positive length")
                if not np.all(np.isfinite(coordinates)):
                    raise ValueError("string coordinates must be finite")
        self._strings = strings

    def _reseeded(self, seed: int) -> "StringGuide":
        """Return the guide with each unseeded random part redrawn from ``seed``."""
        if self._pinned_seed:
            return self
        # Only a random guide, which takes a seed, is ever unpinned.
        guide: Any = type(self)
        return guide(**{**self._arguments(), "seed": seed})



    def metadata(self) -> dict:
        """JSON-compatible record of this guide's type and resolved seed."""
        record: dict = {"type": type(self).__name__}
        seed = getattr(self, "seed", None)
        if seed is not None:
            record["seed"] = int(seed)
        return record


    @property
    def geometry(self) -> LineString | MultiLineString:
        """Shapely geometry of every segment."""
        return _segment_geometry(self._strings)

    @property
    def n_strings(self) -> int:
        """Number of logical strings in this guide."""
        return len(self._strings)



    def modify(self, *modifiers: "StringModifier") -> "Composite":
        """Return a composite guide with modifiers applied in order."""
        from .composites import Composite

        return Composite(self, modifiers)




    def plot(
        self,
        ax: "matplotlib.axes.Axes | None" = None,
        **kwargs: Any,
    ) -> "matplotlib.axes.Axes":
        """Plot the guide and return its Matplotlib axes."""
        from ..plotting import _plot_string_guide

        return _plot_string_guide(self, ax=ax, **kwargs)


    def _transformed(self, scale: float, offset) -> "StringGuide":
        """A copy scaled about the origin by ``scale``, then moved by ``offset``."""
        strings = tuple(
            _LogicalString(
                logical_string.guide_type,
                tuple(
                    shp_translate(
                        shp_scale(segment, scale, scale, origin=(0, 0)),
                        xoff=offset[0],
                        yoff=offset[1],
                    )
                    for segment in logical_string.segments
                ),
            )
            for logical_string in self._strings
        )
        return _GeometryString(strings)



class _GeometryString(StringGuide):
    """Internal guide backed directly by realized logical strings."""

    def __repr__(self) -> str:
        # Realized strings have no constructor call, so summarize them.
        n_segments = sum(len(item.segments) for item in self._strings)
        return f"realized guide(n_strings={self.n_strings}, n_segments={n_segments})"

    def metadata(self) -> dict:
        """JSON-compatible record naming the public base class."""
        return {"type": "StringGuide"}