Source code for suboptimumg.track.models

from typing import Annotated, Any, List, Tuple, Union

import numpy as np
import numpy.typing as npt
from pydantic import BaseModel, ConfigDict, Field, field_validator




class ArcData(BaseModel):
    """Data model for arc geometry. Stores parameters needed to render an arc segment of a track."""

    center: Tuple[float, float] = Field(
        description="(x, y) coordinates of the arc's center in meters"
    )
    width: float = Field(description="Width of the arc ellipse (2 * radius) in meters")
    height: float = Field(
        description="Height of the arc ellipse (2 * radius) in meters"
    )
    angle: float = Field(description="Rotation angle of the arc ellipse in degrees")
    theta1: float = Field(
        description="Starting angle of the arc in degrees (sorted, theta1 < theta2)"
    )
    theta2: float = Field(
        description="Ending angle of the arc in degrees (sorted, theta1 < theta2)"
    )
    theta1_original: float = Field(
        description="Original starting angle before sorting (used to detect reversal)"
    )
    theta2_original: float = Field(
        description="Original ending angle before sorting (used to detect reversal)"
    )
    n_points: int = Field(
        description="Number of interpolation points for NewArc (for smooth rendering)"
    )





class CornerListInput(BaseModel):
    """Strongly-typed class defining a Track based on corners."""

    corners: List[Annotated[List, Field(min_length=2, max_length=2)]] = Field(
        description="A list of corners, each specified as [radius (m), length (decimeters)]"
    )
    shorten: float = Field(
        description="A ratio to scale each corner's length by. (1.0 = no change)"
    )
    distance_step: float = Field(
        description="The length dx to be used to quantize the track"
    )
    clean_corners: bool = Field(
        description="Whether or not to apply experimental modifications to tight corners"
    )
    ideal_rotation_angle: float = Field(
        description="The ideal angle to rotate the track by, so that it fits into a plot nicely"
    )



    @field_validator("corners", mode="before")
    @classmethod
    def wrap_single_corner(cls, v):
        # Automatically nest a single unnested corner [radius, length] into [[radius, length]]
        if isinstance(v, list) and len(v) == 2 and not isinstance(v[0], list):
            return [v]
        return v






class CoordinateListInput(BaseModel):
    """Strongly-typed class defining a Track based on GPS coordinates."""

    longitudes: List[float] = Field(description="A list of longitude coordinates")
    latitudes: List[float] = Field(
        description="A list of latitude coordinates, must match longitudes length"
    )
    distance_step: float = Field(
        description="The length dx to be used to quantize the track"
    )
    sample_dist: float = Field(
        default=1.0,
        description="Distance to sample before/after a point for Menger curvature calculation",
    )





class ContinuousTrackData(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    dx: npt.NDArray[np.float64] = Field(description="Array of distance steps")
    radius: npt.NDArray[np.float64] = Field(
        description="Array of radius of curvature values"
    )
    cumulative_dist: npt.NDArray[np.float64] = Field(
        description="Array of cumulative distances along the track"
    )
    x_m: npt.NDArray[np.float64] = Field(description="Array of x coordinates in meters")
    y_m: npt.NDArray[np.float64] = Field(description="Array of y coordinates in meters")
    distance_step: float = Field(description="Length of each step in meters")
    sample_dist: float = Field(
        description="Distance to sample for Menger curvature calculation"
    )
    tck: Any = Field(description="B-spline representation (tuple from scipy)")
    seed_idx: npt.NDArray[np.int32] = Field(
        description="Array of simulation seed indices"
    )





class DiscreteTrackData(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    dx: npt.NDArray[np.float64] = Field(description="Array of distance steps")
    radius: npt.NDArray[np.float64] = Field(
        description="Array of radius of curvature values"
    )
    cumulative_dist: npt.NDArray[np.float64] = Field(
        description="Array of cumulative distances along the track"
    )
    x_m: npt.NDArray[np.float64] = Field(description="Array of x coordinates in meters")
    y_m: npt.NDArray[np.float64] = Field(description="Array of y coordinates in meters")
    distance_step: float = Field(description="Length of each step in meters")
    original_corners: List[List] = Field(description="Original corner specifications")
    arcs: List[ArcData] = Field(description="Arc geometry data for each corner")
    seed_idx: npt.NDArray[np.int32] = Field(
        description="Array of simulation seed indices"
    )



TrackData = Union[ContinuousTrackData, DiscreteTrackData]