Source code for simulation.utils.road.sections.road_element

"""Road elements are simple individual components of the road that have a frame.

Examples are traffic signs, obstacles or surface markings (e.g. turn arrow on the ground.
"""

from dataclasses import dataclass
from typing import Union

from kitcar_utils.geometry import Line, Point, Polygon, Transform

from simulation.utils.road.sections.transformable import Transformable


@dataclass
class RoadElement(Transformable):
    normalize_x: bool = True
    """If true, all x-values are substracted by the lowest x-value."""

    _frame: Polygon = Polygon(
        [Point(-0.1, -0.2), Point(0.1, -0.2), Point(0.1, 0.2), Point(-0.1, 0.2)]
    )
    """Polygon: Frame of the element in local coordinates."""

    def set_transform(self, obj: Union[Line, Transform]):
        """Calculate the correct transform to this element.

        Depending on :attr:`self.normalize_x` the positional behavior is different. If
        :attr:`self.normalize_x` is True, the element is aligned along the provided line.
        """
        if type(obj) is Line:
            pose = obj.interpolate_pose(
                arc_length=self._center.x if self.normalize_x else 0
            )
            obj = Transform(pose, pose.get_angle())
        super().set_transform(obj)

    @property
    def frame(self) -> Polygon:
        """Polygon: Frame of the element in global coordinates."""
        tf = (
            Transform([-self._center.x, 0], 0) if self.normalize_x else Transform([0, 0], 0)
        )
        return self.transform * (tf * self._frame)

    @property
    def center(self) -> Point:
        """Point: Center point of the element in global coordinates."""
        return Point(self.frame.centroid)

    @property
    def _center(self) -> Point:
        """Point: Center point of the element in local coordinates."""
        return Point(self._frame.centroid)


@dataclass
class RoadElementRect(RoadElement):
    """Generic element of the road that has a frame.

    Examples of road elements are obstacles and traffic signs.

    Args:
        arc_length: x coordinate of the element along the road.
        y: y coordinate of the element. (Perpendicular to the road.)
        width: Width of the element.
        depth: Depth of the element. Component of the size in the direction of the road.
        angle: Angle [radian] between the middle line and the element
            (measured at the center).
    """

    width: float = 0.2
    """Width of the element."""
    depth: float = 0.2
    """Depth (length) of the element."""
    angle: float = 0
    """Angle [radian] between the middle line and the element (measured at the center)."""

    def __init__(
        self,
        arc_length: float = 0.4,
        y: float = -0.2,
        width: float = width,
        depth: float = depth,
        angle: float = angle,
        normalize_x: bool = True,
        z: float = 0,
        height: float = 0,
    ):
        """Initialize a retangular road element."""
        for obj in arc_length, y, width, depth, angle:
            assert isinstance(obj, float) or isinstance(
                obj, int
            ), f"Should be a number but is {obj}"

        self.width = width
        self.depth = depth
        self.angle = angle
        self.height = height
        super().__init__(
            normalize_x=normalize_x,
            _frame=Transform(Point(arc_length, y, z), self.angle)
            * Polygon(
                [
                    [-self.depth / 2, self.width / 2],
                    [self.depth / 2, self.width / 2],
                    [self.depth / 2, -self.width / 2],
                    [-self.depth / 2, -self.width / 2],
                ]
            ),
        )

    @property
    def orientation(self) -> float:
        """float: Orientation of the element in global coordinates in radians."""
        return self.transform.get_angle() + self.angle