inire/inire/geometry/components.py

from __future__ import annotations

from typing import NamedTuple

import numpy as np
from shapely.geometry import Polygon

from .primitives import Port

# Search Grid Snap (1.0 µm)
SEARCH_GRID_SNAP_UM = 1.0


def snap_search_grid(value: float) -> float:
    """Snap a coordinate to the nearest 1µm."""
    return round(value / SEARCH_GRID_SNAP_UM) * SEARCH_GRID_SNAP_UM


class ComponentResult(NamedTuple):
    """The result of a component generation: geometry and the final port."""

    geometry: list[Polygon]
    end_port: Port


class Straight:
    @staticmethod
    def generate(start_port: Port, length: float, width: float) -> ComponentResult:
        """Generate a straight waveguide segment."""
        # Calculate end port position
        rad = np.radians(start_port.orientation)
        dx = length * np.cos(rad)
        dy = length * np.sin(rad)

        end_port = Port(start_port.x + dx, start_port.y + dy, start_port.orientation)

        # Create polygon (centered on port)
        half_w = width / 2.0
        # Points relative to start port (0,0)
        points = [(0, half_w), (length, half_w), (length, -half_w), (0, -half_w)]

        # Transform points
        cos_val = np.cos(rad)
        sin_val = np.sin(rad)
        poly_points = []
        for px, py in points:
            tx = start_port.x + px * cos_val - py * sin_val
            ty = start_port.y + px * sin_val + py * cos_val
            poly_points.append((tx, ty))

        return ComponentResult(geometry=[Polygon(poly_points)], end_port=end_port)


def _get_num_segments(radius: float, angle_deg: float, sagitta: float = 0.01) -> int:
    """Calculate number of segments for an arc to maintain a maximum sagitta."""
    if radius <= 0:
        return 1
    # angle_deg is absolute angle turned
    # s = R(1 - cos(theta/2)) => cos(theta/2) = 1 - s/R
    # theta = 2 * acos(1 - s/R)
    # n = total_angle / theta
    ratio = max(0.0, min(1.0, 1.0 - sagitta / radius))
    theta_max = 2.0 * np.arccos(ratio)
    if theta_max == 0:
        return 16
    num = int(np.ceil(np.radians(abs(angle_deg)) / theta_max))
    return max(4, num)


class Bend90:
    @staticmethod
    def generate(start_port: Port, radius: float, width: float, direction: str = "CW", sagitta: float = 0.01) -> ComponentResult:
        """Generate a 90-degree bend."""
        # direction: 'CW' (-90) or 'CCW' (+90)
        turn_angle = -90 if direction == "CW" else 90

        # Calculate center of the arc
        rad_start = np.radians(start_port.orientation)
        center_angle = rad_start + (np.pi / 2 if direction == "CCW" else -np.pi / 2)
        cx = start_port.x + radius * np.cos(center_angle)
        cy = start_port.y + radius * np.sin(center_angle)

        # Center to start is radius at center_angle + pi
        theta_start = center_angle + np.pi
        theta_end = theta_start + (np.pi / 2 if direction == "CCW" else -np.pi / 2)

        ex = cx + radius * np.cos(theta_end)
        ey = cy + radius * np.sin(theta_end)

        # End port orientation
        end_orientation = (start_port.orientation + turn_angle) % 360

        snapped_ex = snap_search_grid(ex)
        snapped_ey = snap_search_grid(ey)

        end_port = Port(snapped_ex, snapped_ey, float(end_orientation))

        # Generate arc geometry
        num_segments = _get_num_segments(radius, 90, sagitta)
        angles = np.linspace(theta_start, theta_end, num_segments + 1)

        inner_radius = radius - width / 2.0
        outer_radius = radius + width / 2.0

        inner_points = [(cx + inner_radius * np.cos(a), cy + inner_radius * np.sin(a)) for a in angles]
        outer_points = [(cx + outer_radius * np.cos(a), cy + outer_radius * np.sin(a)) for a in reversed(angles)]

        return ComponentResult(geometry=[Polygon(inner_points + outer_points)], end_port=end_port)


class SBend:
    @staticmethod
    def generate(start_port: Port, offset: float, radius: float, width: float, sagitta: float = 0.01) -> ComponentResult:
        """Generate a parametric S-bend (two tangent arcs). Only for offset < 2*radius."""
        if abs(offset) >= 2 * radius:
            raise ValueError(f"SBend offset {offset} must be less than 2*radius {2 * radius}")

        # Analytical length: L = 2 * sqrt(O * (2*R - O/4)) is for a specific S-bend type.
        # Standard S-bend with two equal arcs:
        # Offset O = 2 * R * (1 - cos(theta))
        # theta = acos(1 - O / (2*R))
        theta = np.arccos(1 - abs(offset) / (2 * radius))

        # Length of one arc = R * theta
        # Total length of S-bend = 2 * R * theta (arc length)
        # Horizontal distance dx = 2 * R * sin(theta)

        dx = 2 * radius * np.sin(theta)
        dy = offset

        # End port
        rad_start = np.radians(start_port.orientation)
        ex = start_port.x + dx * np.cos(rad_start) - dy * np.sin(rad_start)
        ey = start_port.y + dx * np.sin(rad_start) + dy * np.cos(rad_start)

        end_port = Port(ex, ey, start_port.orientation)

        # Geometry: two arcs
        # First arc center
        direction = 1 if offset > 0 else -1
        center_angle1 = rad_start + direction * np.pi / 2
        cx1 = start_port.x + radius * np.cos(center_angle1)
        cy1 = start_port.y + radius * np.sin(center_angle1)

        # Second arc center
        center_angle2 = rad_start - direction * np.pi / 2
        cx2 = ex + radius * np.cos(center_angle2)
        cy2 = ey + radius * np.sin(center_angle2)

        # Generate points for both arcs
        num_segments = _get_num_segments(radius, float(np.degrees(theta)), sagitta)
        # Arc 1: theta_start1 to theta_end1
        theta_start1 = center_angle1 + np.pi
        theta_end1 = theta_start1 - direction * theta

        # Arc 2: theta_start2 to theta_end2
        theta_start2 = center_angle2
        theta_end2 = theta_start2 + direction * theta

        def get_arc_points(cx: float, cy: float, r_inner: float, r_outer: float, t_start: float, t_end: float) -> list[tuple[float, float]]:
            angles = np.linspace(t_start, t_end, num_segments + 1)
            inner = [(cx + r_inner * np.cos(a), cy + r_inner * np.sin(a)) for a in angles]
            outer = [(cx + r_outer * np.cos(a), cy + r_outer * np.sin(a)) for a in reversed(angles)]
            return inner + outer

        poly1 = Polygon(get_arc_points(cx1, cy1, radius - width / 2, radius + width / 2, theta_start1, theta_end1))
        poly2 = Polygon(get_arc_points(cx2, cy2, radius - width / 2, radius + width / 2, theta_end2, theta_start2))

        return ComponentResult(geometry=[poly1, poly2], end_port=end_port)