Source code for suboptimumg.sweep.models

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray
from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator

from ..compsim.models import CompetitionData




class SweepParamConfig(BaseModel):
    """Fields needed to specify a sweep variable."""

    model_config = ConfigDict(frozen=True, extra="forbid")

    name: str = Field(description="Name of the parameter")
    min: float | int = Field(description="Lower bound")
    max: float | int = Field(description="Upper bound")
    steps: int = Field(description="Number of steps to sweep across")



    @model_validator(mode="after")
    def check_valid(self) -> SweepParamConfig:
        """Ensure `min` is strictly less than `max`."""
        if self.min >= self.max:
            raise ValueError(f"`min` = {self.min} must be less than `max` = {self.max}")
        return self






class SweepProcessInput1D(BaseModel):
    """Input model for 1D sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    comp_data: CompetitionData = Field(description="Competition configuration data model")
    var_1_name: str = Field(description="Name of the variable being swept")
    var_1_value: float = Field(description="Value of the variable for this iteration")
    dep_vals: dict[str, float] = Field(
        description="Dictionary of dependent parameter values",
        default_factory=dict,
    )
    idx: int = Field(description="Index of this sweep iteration")





class SweepProcessInput2D(BaseModel):
    """Input model for 2D sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    comp_data: CompetitionData = Field(description="Competition configuration data model")
    var_1_name: str = Field(description="Name of the first variable being swept")
    var_1_value: float = Field(description="Value of the first variable for this iteration")
    var_2_name: str = Field(description="Name of the second variable being swept")
    var_2_value: float = Field(description="Value of the second variable for this iteration")
    dep_vals: dict[str, float] = Field(
        description="Dictionary of dependent parameter values",
        default_factory=dict,
    )
    x_idx: int = Field(description="X-axis index of this sweep iteration")
    y_idx: int = Field(description="Y-axis index of this sweep iteration")





class SweepProcessOutput1D(BaseModel):
    """Output model for 1D sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    idx: int = Field(description="Index of this sweep iteration")
    accel_pts: float = Field(description="Acceleration event points")
    skidpad_pts: float = Field(description="Skidpad event points")
    autoX_pts: float = Field(description="Autocross event points")
    endurance_pts: float = Field(description="Endurance event points")
    efficiency_pts: float = Field(description="Efficiency points")
    accel_t: float = Field(description="Acceleration event time")
    skidpad_t: float = Field(description="Skidpad event time")
    autoX_t: float = Field(description="Autocross event time")
    endurance_t: float = Field(description="Endurance event time")
    warnings: str = Field(description="Warning messages from the simulation", default="")
    error: str | None = Field(description="Error message if simulation failed", default=None)





class SweepProcessOutput2D(BaseModel):
    """Output model for 2D sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    x_idx: int = Field(description="Index of first variable of this sweep iteration")
    y_idx: int = Field(description="Index of second variable of this sweep iteration")
    accel_pts: float = Field(description="Acceleration event points")
    skidpad_pts: float = Field(description="Skidpad event points")
    autoX_pts: float = Field(description="Autocross event points")
    endurance_pts: float = Field(description="Endurance event points")
    efficiency_pts: float = Field(description="Efficiency points")
    accel_t: float = Field(description="Acceleration event time")
    skidpad_t: float = Field(description="Skidpad event time")
    autoX_t: float = Field(description="Autocross event time")
    endurance_t: float = Field(description="Endurance event time")
    warnings: str = Field(description="Warning messages from the simulation", default="")
    error: str | None = Field(description="Error message if simulation failed", default=None)





class ArraySweepVariableInput(BaseModel):
    """Per-variable sweep values used to build array sweep process inputs."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    name: str = Field(description="Name of the variable being swept")
    sweep_values: NDArray[np.float64] = Field(
        description="Values of the variable across the sweep steps"
    )





class ArraySweepProcessInput(BaseModel):
    """Input model for array sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    comp_data: CompetitionData = Field(description="Competition configuration data model")
    var_name: str = Field(description="Name of the variable being swept")
    step: float = Field(description="Value of the variable for this iteration")
    var_idx: int = Field(description="Index of the variable in the sweep array")
    step_idx: int = Field(description="Index of the step in the sweep")





class ArraySweepProcessOutput(BaseModel):
    """Output model for array sweep process."""

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    var_idx: int = Field(description="Index of the variable in the sweep array")
    step_idx: int = Field(description="Index of the step in the sweep")
    plausible: bool = Field(description="Whether the parameter configuration is plausible")
    accel_pts: float = Field(description="Acceleration event points")
    skidpad_pts: float = Field(description="Skidpad event points")
    autoX_pts: float = Field(description="Autocross event points")
    endurance_pts: float = Field(description="Endurance event points")
    efficiency_pts: float = Field(description="Efficiency points")
    accel_t: float = Field(description="Acceleration event time")
    skidpad_t: float = Field(description="Skidpad event time")
    autoX_t: float = Field(description="Autocross event time")
    endurance_t: float = Field(description="Endurance event time")
    warnings: str = Field(description="Warning messages from the simulation", default="")
    error: str | None = Field(description="Error message if simulation failed", default=None)





class SweepData1D(BaseModel):
    """Strongly typed 1D sweep data container."""

    var_name: str = Field(description="Variable name being swept")
    sweep_values: NDArray[np.float64] = Field(
        description="Actual parameter values used in the sweep (x-values)"
    )

    accel_pts: NDArray[np.float64] = Field(description="Acceleration event points")
    skidpad_pts: NDArray[np.float64] = Field(description="Skidpad event points")
    autoX_pts: NDArray[np.float64] = Field(description="Autocross event points")
    endurance_pts: NDArray[np.float64] = Field(description="Endurance event points")
    efficiency_pts: NDArray[np.float64] = Field(description="Efficiency points")
    accel_t: NDArray[np.float64] = Field(description="Acceleration event time")
    skidpad_t: NDArray[np.float64] = Field(description="Skidpad event time")
    autoX_t: NDArray[np.float64] = Field(description="Autocross event time")
    endurance_t: NDArray[np.float64] = Field(description="Endurance event time")

    plausible: NDArray[np.bool_] = Field(
        description="Plausibility check result for each sweep step"
    )

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    @computed_field  # type: ignore[prop-decorator]
    # See vehicle_models.py: valid Pydantic V2 pattern, mypy plugin currently
    # incompatible with the installed mypy version.
    @property
    def total_pts(self) -> NDArray[np.float64]:
        """Compute total points as sum of all point categories."""
        return (
            self.accel_pts
            + self.skidpad_pts
            + self.autoX_pts
            + self.endurance_pts
            + self.efficiency_pts
        )



    @classmethod
    def create(cls, var_name: str, sweep_values: NDArray[np.float64]) -> SweepData1D:
        """Create a SweepData1D with pre-allocated numpy arrays of specified length."""
        length = len(sweep_values)
        return cls(
            var_name=var_name,
            sweep_values=sweep_values,
            accel_pts=np.zeros(length, dtype=float),
            skidpad_pts=np.zeros(length, dtype=float),
            autoX_pts=np.zeros(length, dtype=float),
            endurance_pts=np.zeros(length, dtype=float),
            efficiency_pts=np.zeros(length, dtype=float),
            accel_t=np.zeros(length, dtype=float),
            skidpad_t=np.zeros(length, dtype=float),
            autoX_t=np.zeros(length, dtype=float),
            endurance_t=np.zeros(length, dtype=float),
            plausible=np.zeros(length, dtype=bool),
        )






class SweepData2D(BaseModel):
    """Strongly typed 2D sweep data container."""

    var_name_1: str = Field(description="First variable name being swept")
    var_list_1: NDArray[np.float64] = Field(
        description="Actual values for first variable used in the sweep (x-values)"
    )
    var_name_2: str = Field(description="Second variable name being swept")
    var_list_2: NDArray[np.float64] = Field(
        description="Actual values for second variable used in the sweep (y-values)"
    )

    accel_pts: NDArray[np.float64] = Field(description="Acceleration event points")
    skidpad_pts: NDArray[np.float64] = Field(description="Skidpad event points")
    autoX_pts: NDArray[np.float64] = Field(description="Autocross event points")
    endurance_pts: NDArray[np.float64] = Field(description="Endurance event points")
    efficiency_pts: NDArray[np.float64] = Field(description="Efficiency points")
    accel_t: NDArray[np.float64] = Field(description="Acceleration event time")
    skidpad_t: NDArray[np.float64] = Field(description="Skidpad event time")
    autoX_t: NDArray[np.float64] = Field(description="Autocross event time")
    endurance_t: NDArray[np.float64] = Field(description="Endurance event time")

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)

    @computed_field  # type: ignore[prop-decorator]
    # See vehicle_models.py: valid Pydantic V2 pattern, mypy plugin currently
    # incompatible with the installed mypy version.
    @property
    def total_pts(self) -> NDArray[np.float64]:
        """Compute total points as sum of all point categories."""
        return (
            self.accel_pts
            + self.skidpad_pts
            + self.autoX_pts
            + self.endurance_pts
            + self.efficiency_pts
        )



    @classmethod
    def create(
        cls,
        var_name_1: str,
        var_list_1: NDArray[np.float64],
        var_name_2: str,
        var_list_2: NDArray[np.float64],
    ) -> SweepData2D:
        """Create a SweepData2D with pre-allocated numpy arrays of specified dimensions."""
        dim1 = len(var_list_1)
        dim2 = len(var_list_2)
        return cls(
            var_name_1=var_name_1,
            var_list_1=var_list_1,
            var_name_2=var_name_2,
            var_list_2=var_list_2,
            accel_pts=np.zeros((dim1, dim2), dtype=float),
            skidpad_pts=np.zeros((dim1, dim2), dtype=float),
            autoX_pts=np.zeros((dim1, dim2), dtype=float),
            endurance_pts=np.zeros((dim1, dim2), dtype=float),
            efficiency_pts=np.zeros((dim1, dim2), dtype=float),
            accel_t=np.zeros((dim1, dim2), dtype=float),
            skidpad_t=np.zeros((dim1, dim2), dtype=float),
            autoX_t=np.zeros((dim1, dim2), dtype=float),
            endurance_t=np.zeros((dim1, dim2), dtype=float),
        )






class ArraySweepData(BaseModel):
    """Strongly typed container for array sweep data."""

    percent_steps: NDArray[np.float64] = Field(
        description="Percent step values used in the sweep (shared across all variables)"
    )
    sweep_outputs: list[SweepData1D] = Field(
        description="List of sweep output data for each variable"
    )

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)





class MotorSweepData(BaseModel):
    """Strongly typed container for motor sweep data."""

    data_by_motor: dict[str, SweepData1D] = Field(description="Motor name to sweep data map")

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)