Source code for suboptimumg.track.models

from __future__ import annotations

from typing import Annotated

import numpy as np
from numpy.typing import NDArray
from perda.core_data_structures import SingleRunData
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_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[float], 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: object) -> object:
        """
        Nest a single unnested corner into a one-element list of corners.

        Parameters
        ----------
        v : object
            Raw input before validation: either one corner [radius, length],
            or a list of corners.

        Returns
        -------
        object
            The corners, guaranteed to be a list of [radius, length] pairs.
        """
        if isinstance(v, list) and len(v) == 2 and not isinstance(v[0], list):
            return [v]
        return v






class ContinuousTrackData(BaseModel):
    """Serializable data for a continuous (B-spline based) Track."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    dx: NDArray[np.float64] = Field(description="Array of distance steps")
    radius: NDArray[np.float64] = Field(description="Array of radius of curvature values")
    cumulative_dist: NDArray[np.float64] = Field(
        description="Array of cumulative distances along the track"
    )
    x_m: NDArray[np.float64] = Field(description="Array of x coordinates in meters")
    y_m: NDArray[np.float64] = Field(description="Array of y coordinates in meters")
    distance_step: float = Field(description="Length of each step in meters")
    tck: tuple[NDArray[np.float64], list[NDArray[np.float64]], int] = Field(
        description="B-spline representation from scipy: knots, coefficients, and degree"
    )
    seed_idx: NDArray[np.int32] = Field(description="Array of simulation seed indices")





class DiscreteTrackData(BaseModel):
    """Serializable data for a discrete (corner based) Track."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    dx: NDArray[np.float64] = Field(description="Array of distance steps")
    radius: NDArray[np.float64] = Field(description="Array of radius of curvature values")
    cumulative_dist: NDArray[np.float64] = Field(
        description="Array of cumulative distances along the track"
    )
    x_m: NDArray[np.float64] = Field(description="Array of x coordinates in meters")
    y_m: 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[float]] = Field(description="Original corner specifications")
    arcs: list[ArcData] = Field(description="Arc geometry data for each corner")
    seed_idx: NDArray[np.int32] = Field(description="Array of simulation seed indices")



TrackData = ContinuousTrackData | DiscreteTrackData

# Keys required in SingleRunData for from_perda_logs, mapped to the function that produces them.
REQUIRED_PERDA_FIELDS: dict[str, str] = {
    "body.curvature": "add_curvature() from suboptimumg.log_analysis.kinematics",
    "groundSpeed": "add_groundspeed() from suboptimumg.log_analysis.kinematics",
    "posX": "preprocess_gps_data_instances() from suboptimumg.log_analysis.preprocess_gps",
    "posY": "preprocess_gps_data_instances() from suboptimumg.log_analysis.preprocess_gps",
}




class PerdaLogInput(BaseModel):
    """Input for building a Track from a PERDA SingleRunData lap segment."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    data: SingleRunData = Field(description="Trimmed single-lap SingleRunData object")
    distance_step: float = Field(
        default=0.1,
        gt=0.0,
        description="Uniform distance grid spacing dx in meters",
    )
    cutoff_freq_per_meter: float = Field(
        default=0.18,
        gt=0.0,
        description="Spatial frequency cutoff for curvature lowpass filter (cycles/m)",
    )
    max_radius: float = Field(
        default=300.0,
        gt=0.0,
        description="Radius clamp for near-straight segments (m)",
    )
    filter_order: int = Field(
        default=4,
        ge=1,
        description="Butterworth filter order for spatial lowpass",
    )

    @model_validator(mode="after")
    def _check_required_fields(self) -> PerdaLogInput:
        """
        Check that the SingleRunData contains all fields from_perda_logs needs.

        Returns
        -------
        PerdaLogInput
            The validated input, unchanged.
        """
        missing = {k: hint for k, hint in REQUIRED_PERDA_FIELDS.items() if k not in self.data}
        if missing:
            lines = "\n".join(f"  '{k}': run {hint}" for k, hint in missing.items())
            raise ValueError(f"SingleRunData is missing required fields:\n{lines}")
        return self