Source code for suboptimumg.compsim.models

"""
Pydantic models for competition simulation data structures.
"""

from typing import List, Optional

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

from ..track.models import TrackData
from ..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=f"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=f"Rules-defined baseline laptime (1.45 = 145% of fastest endurance finisher), used as constant in score calculations"
    )
    eff_min_co2_your: float = Field(
        description=f"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(default=None)
    acc_proposal: NDArray[np.float64] = Field(default=None)
    grown_forward_mask: NDArray[np.bool] = Field(default=None)
    slower_mask: NDArray[np.bool] = Field(default=None)
    v_max_post: NDArray[np.float64] = Field(default=None)
    acc_max_post: NDArray[np.float64] = Field(default=None)
    p_proposal: NDArray[np.float64] = Field(default=None)





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)"
    )

    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: Optional[InternalData] = Field(
        default=None, description="Optional internal data from the simulation"
    )





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"
    )

    @computed_field
    @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")