Source code for suboptimumg.compsim.models

import csv
from pathlib import Path

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

from ..track.models import TrackData
from ..vehicle.vehicle_models import VehicleModel




class EventScoring(BaseModel):
    event_best_time: float = Field(description="Best time recorded for this event")
    worst_time_ratio: float = Field(
        description="The maximum race time for this event, expressed as a ratio of the track_best_time"
    )
    min_points: float = Field(description="The minimum amount of points awarded for this event")
    max_point_gain: float = Field(
        description="The maximum amount of point a car can gain on this event, if they meet or exceed track_best_time"
    )
    scale_time_ratio: float = Field(
        description="Exponentiate event time ratios for point calculations (1 = no effect)"
    )





class EfficiencyScoring(BaseModel):
    num_endurance_laps: int = Field(
        description="Number of laps around the track for endurance event"
    )
    max_eff_points: float = Field(description="Maximum amount of points awarded for efficiency")
    co2_scaling: float = Field(
        description="(kg CO2 / kWh) ratio to convert EV to IC (from FSAE Rules)"
    )

    fastest_lap_time: float = Field(description="Average lap time for fastest endurance finisher")
    co2_min: float = Field(
        description="Minimum co2 usage out of teams that were eligible for endurance score"
    )
    eff_max: float = Field(description="Best efficiency factor out of all teams")

    max_time_scale: float = Field(
        description="Rules-defined baseline laptime (1.45 = 145% of fastest endurance finisher), used as constant in score calculations"
    )
    eff_min_co2_your: float = Field(
        description="Rules-defined baseline CO2 value ((20.02 kg CO2 / 100km) * 22km), used as constant in score calculations"
    )





class CompetitionScoring(BaseModel):
    accel: EventScoring
    skidpad: EventScoring
    autoX: EventScoring
    endurance: EventScoring
    efficiency: EfficiencyScoring





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

    v_proposal: NDArray[np.float64] = Field(description="Proposed velocity profile")
    acc_proposal: NDArray[np.float64] = Field(description="Proposed acceleration profile")
    grown_forward_mask: NDArray[np.bool] = Field(
        description="Mask of indices grown forward from the seed"
    )
    slower_mask: NDArray[np.bool] = Field(
        description="Mask of indices where the proposal is slower than the current max"
    )
    v_max_post: NDArray[np.float64] = Field(
        description="Velocity profile after merging this seed's proposal"
    )
    acc_max_post: NDArray[np.float64] = Field(
        description="Acceleration profile after merging this seed's proposal"
    )
    p_proposal: NDArray[np.float64] = Field(description="Proposed power profile")





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

    initial_velocity: float = Field(description="Initial velocity at start of track (m/s)")
    v_max_profile: NDArray[np.float64] = Field(
        description="Initial maximum velocity profile for the track (m/s)"
    )
    seed_idx_list: NDArray[np.integer] = Field(
        description="List of seed indices used in the simulation, sorted by increasing max velocity"
    )
    cumulative_dist: NDArray[np.float64] = Field(
        description="Cumulative distance along the track (m)"
    )
    # NOTE: was not included before
    seed_v_list: NDArray[np.float64] = Field(
        description="List of seed velocities, corresponding to seed_idx_list"
    )
    per_seed: list[InternalDataSeedInfo] = Field(
        description="List of per-seed internal data, corresponding to seed_idx_list"
    )





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

    lap_t: NDArray[np.float64] = Field(description="List of lap time (s)")
    lap_dxs: NDArray[np.float64] = Field(description="List of distances (m)")
    lap_vels: NDArray[np.float64] = Field(description="List of velocities (m/s)")
    lap_accs: NDArray[np.float64] = Field(description="List of accelerations (m/s^2)")
    lap_powers: NDArray[np.float64] = Field(description="List of powers (W)")
    lap_eff_motor_torques: NDArray[np.float64] = Field(
        description="List of effective motor torques (N⋅m)"
    )
    internal_data: InternalData | None = Field(
        default=None, description="Optional internal data from the simulation"
    )



    def write_csv(self, filepath: str | Path) -> None:
        """Write the lap simulation time series to a CSV file.

        Parameters
        ----------
        filepath : str or Path
            Destination CSV file path.
        """
        columns = {
            "Time (s)": self.lap_t,
            "Distance (m)": self.lap_dxs,
            "Velocity (m/s)": self.lap_vels,
            "Acceleration (m/s^2)": self.lap_accs,
            "Power Draw (W)": self.lap_powers,
            "Effective Motor Torque (N⋅m)": self.lap_eff_motor_torques,
        }
        with open(filepath, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(columns.keys())
            writer.writerows(zip(*columns.values()))






class EventResults(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    points: float = Field(description="Points scored in the event")
    tyour: float = Field(description="Your time (s)")
    lapsim_results: LapsimResults = Field(description="Results from the lapsim")





class CompetitionResults(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    accel: EventResults = Field(description="Acceleration event results")
    skidpad: EventResults = Field(description="Skidpad event results")
    autoX: EventResults = Field(description="Autocross event results")
    endurance: EventResults = Field(description="Endurance event results")
    efficiency_points: float = Field(description="Efficiency points scored in the endurance event")
    endu_consumed_energy: float = Field(description="Energy consumed in the endurance event (kWh)")

    @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_points(self) -> float:
        return (
            self.accel.points
            + self.skidpad.points
            + self.autoX.points
            + self.endurance.points
            + self.efficiency_points
        )





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

    vehicle_model: VehicleModel = Field(description="Vehicle configuration")
    accel: TrackData = Field(description="Acceleration track data")
    skidpad: TrackData = Field(description="Skidpad track data")
    autoX: TrackData = Field(description="Autocross track data")
    endurance: TrackData = Field(description="Endurance track data")
    scoring: CompetitionScoring = Field(description="Competition scoring configuration")