fsd_model/config.py

"""
Configurable sensor and vehicle configuration for the FSD model.
Supports arbitrary sensor counts, types, and placements.
"""

import math
import json
from dataclasses import dataclass, field, asdict
from typing import List, Optional, Tuple, Dict, Any
from enum import Enum


class SensorType(Enum):
    CAMERA = "camera"
    ULTRASONIC = "ultrasonic"
    LIDAR = "lidar"
    RADAR = "radar"


class CameraPosition(Enum):
    FRONT_LEFT = "front_left"
    FRONT_RIGHT = "front_right"
    REAR_LEFT = "rear_left"
    REAR_RIGHT = "rear_right"
    LEFT_MIRROR = "left_mirror"
    RIGHT_MIRROR = "right_mirror"
    FRONT_CENTER = "front_center"
    REAR_CENTER = "rear_center"
    ROOF_FRONT = "roof_front"
    ROOF_REAR = "roof_rear"
    BUMPER_LEFT = "bumper_left"
    BUMPER_RIGHT = "bumper_right"


class UltrasonicZone(Enum):
    FRONT_LEFT_CORNER = "front_left_corner"
    FRONT_LEFT = "front_left"
    FRONT_CENTER_LEFT = "front_center_left"
    FRONT_CENTER = "front_center"
    FRONT_CENTER_RIGHT = "front_center_right"
    FRONT_RIGHT = "front_right"
    FRONT_RIGHT_CORNER = "front_right_corner"
    LEFT_FRONT = "left_front"
    LEFT_CENTER = "left_center"
    LEFT_REAR = "left_rear"
    RIGHT_FRONT = "right_front"
    RIGHT_CENTER = "right_center"
    RIGHT_REAR = "right_rear"
    REAR_LEFT_CORNER = "rear_left_corner"
    REAR_LEFT = "rear_left"
    REAR_CENTER_LEFT = "rear_center_left"
    REAR_CENTER = "rear_center"
    REAR_CENTER_RIGHT = "rear_center_right"
    REAR_RIGHT = "rear_right"
    REAR_RIGHT_CORNER = "rear_right_corner"


@dataclass
class SensorPlacement:
    x: float
    y: float
    z: float
    yaw: float = 0.0
    pitch: float = 0.0
    roll: float = 0.0

    def to_transform_matrix(self):
        import numpy as np
        yaw_r = math.radians(self.yaw)
        pitch_r = math.radians(self.pitch)
        roll_r = math.radians(self.roll)
        Rz = np.array([
            [math.cos(yaw_r), -math.sin(yaw_r), 0],
            [math.sin(yaw_r), math.cos(yaw_r), 0],
            [0, 0, 1]
        ])
        Ry = np.array([
            [math.cos(pitch_r), 0, math.sin(pitch_r)],
            [0, 1, 0],
            [-math.sin(pitch_r), 0, math.cos(pitch_r)]
        ])
        Rx = np.array([
            [1, 0, 0],
            [0, math.cos(roll_r), -math.sin(roll_r)],
            [0, math.sin(roll_r), math.cos(roll_r)]
        ])
        R = Rz @ Ry @ Rx
        T = np.eye(4)
        T[:3, :3] = R
        T[:3, 3] = [self.x, self.y, self.z]
        return T


@dataclass
class CameraSensorConfig:
    name: str
    position: CameraPosition
    placement: SensorPlacement
    resolution: Tuple[int, int] = (640, 480)
    fov_horizontal: float = 120.0
    fov_vertical: float = 90.0
    fps: int = 30
    encoding: str = "rgb"
    distortion_model: str = "pinhole"
    fx: Optional[float] = None
    fy: Optional[float] = None
    cx: Optional[float] = None
    cy: Optional[float] = None

    def __post_init__(self):
        if self.fx is None:
            self.fx = self.resolution[0] / (2 * math.tan(math.radians(self.fov_horizontal / 2)))
        if self.fy is None:
            self.fy = self.resolution[1] / (2 * math.tan(math.radians(self.fov_vertical / 2)))
        if self.cx is None:
            self.cx = self.resolution[0] / 2
        if self.cy is None:
            self.cy = self.resolution[1] / 2

    def get_intrinsic_matrix(self):
        import numpy as np
        return np.array([
            [self.fx, 0, self.cx],
            [0, self.fy, self.cy],
            [0, 0, 1]
        ])


@dataclass
class UltrasonicSensorConfig:
    name: str
    zone: UltrasonicZone
    placement: SensorPlacement
    max_range: float = 5.0
    min_range: float = 0.02
    beam_angle: float = 30.0
    frequency: float = 40000.0
    update_rate: float = 20.0
    accuracy: float = 0.01
    noise_std: float = 0.005


@dataclass
class SensorConfig:
    cameras: List[CameraSensorConfig] = field(default_factory=list)
    ultrasonics: List[UltrasonicSensorConfig] = field(default_factory=list)

    @property
    def num_cameras(self) -> int:
        return len(self.cameras)

    @property
    def num_ultrasonics(self) -> int:
        return len(self.ultrasonics)

    @property
    def total_sensors(self) -> int:
        return self.num_cameras + self.num_ultrasonics

    def validate(self):
        issues = []
        if self.num_cameras == 0:
            issues.append("WARNING: No cameras configured")
        if self.num_ultrasonics == 0:
            issues.append("WARNING: No ultrasonic sensors configured")
        yaw_angles = sorted([c.placement.yaw for c in self.cameras])
        if len(yaw_angles) >= 2:
            gaps = []
            for i in range(len(yaw_angles)):
                next_i = (i + 1) % len(yaw_angles)
                gap = (yaw_angles[next_i] - yaw_angles[i]) % 360
                if gap == 0 and next_i != 0:
                    continue
                gaps.append(gap)
            max_gap = max(gaps) if gaps else 360
            if max_gap > 120:
                issues.append(f"WARNING: Camera coverage gap of {max_gap:.0f} degrees")
        return issues

    def get_sensor_summary(self) -> str:
        lines = [f"Sensor Configuration Summary:"]
        lines.append(f"  Total sensors: {self.total_sensors}")
        lines.append(f"  Cameras: {self.num_cameras}")
        for cam in self.cameras:
            lines.append(f"    - {cam.name}: {cam.position.value} | {cam.resolution[0]}x{cam.resolution[1]} | FOV: {cam.fov_horizontal}")
        lines.append(f"  Ultrasonic sensors: {self.num_ultrasonics}")
        for us in self.ultrasonics:
            lines.append(f"    - {us.name}: {us.zone.value} | Range: {us.min_range}-{us.max_range}m")
        issues = self.validate()
        if issues:
            lines.append("  Validation issues:")
            for issue in issues:
                lines.append(f"    {issue}")
        return "\n".join(lines)

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    def save(self, path: str):
        with open(path, 'w') as f:
            json.dump(self.to_dict(), f, indent=2, default=str)


@dataclass
class VehicleConfig:
    name: str = "FSD_Vehicle"
    length: float = 4.5
    width: float = 1.8
    height: float = 1.5
    wheelbase: float = 2.7
    track_width: float = 1.5
    max_speed_mph: float = 20.0
    max_speed_ms: float = 8.94
    max_steering_angle: float = 35.0
    max_acceleration: float = 3.0
    max_deceleration: float = 8.0
    sensor_config: Optional[SensorConfig] = None

    def __post_init__(self):
        self.max_speed_ms = self.max_speed_mph * 0.44704
        if self.sensor_config is None:
            self.sensor_config = self._default_sensor_config()

    def _default_sensor_config(self) -> SensorConfig:
        cameras = self._create_default_cameras()
        ultrasonics = self._create_default_ultrasonics()
        return SensorConfig(cameras=cameras, ultrasonics=ultrasonics)

    def _create_default_cameras(self) -> List[CameraSensorConfig]:
        half_l = self.length / 2
        half_w = self.width / 2
        mirror_h = 1.1
        cameras = [
            CameraSensorConfig(
                name="cam_front_left", position=CameraPosition.FRONT_LEFT,
                placement=SensorPlacement(x=half_l - 0.3, y=half_w, z=0.8, yaw=-45, pitch=-5),
                resolution=(640, 480), fov_horizontal=120,
            ),
            CameraSensorConfig(
                name="cam_front_right", position=CameraPosition.FRONT_RIGHT,
                placement=SensorPlacement(x=half_l - 0.3, y=-half_w, z=0.8, yaw=45, pitch=-5),
                resolution=(640, 480), fov_horizontal=120,
            ),
            CameraSensorConfig(
                name="cam_rear_left", position=CameraPosition.REAR_LEFT,
                placement=SensorPlacement(x=-half_l + 0.3, y=half_w, z=0.8, yaw=-135, pitch=-5),
                resolution=(640, 480), fov_horizontal=120,
            ),
            CameraSensorConfig(
                name="cam_rear_right", position=CameraPosition.REAR_RIGHT,
                placement=SensorPlacement(x=-half_l + 0.3, y=-half_w, z=0.8, yaw=135, pitch=-5),
                resolution=(640, 480), fov_horizontal=120,
            ),
            CameraSensorConfig(
                name="cam_left_mirror", position=CameraPosition.LEFT_MIRROR,
                placement=SensorPlacement(x=0.8, y=half_w + 0.1, z=mirror_h, yaw=-90, pitch=-10),
                resolution=(640, 480), fov_horizontal=90,
            ),
            CameraSensorConfig(
                name="cam_right_mirror", position=CameraPosition.RIGHT_MIRROR,
                placement=SensorPlacement(x=0.8, y=-(half_w + 0.1), z=mirror_h, yaw=90, pitch=-10),
                resolution=(640, 480), fov_horizontal=90,
            ),
        ]
        return cameras

    def _create_default_ultrasonics(self) -> List[UltrasonicSensorConfig]:
        half_l = self.length / 2
        half_w = self.width / 2
        bumper_h = 0.4
        zones_front = [
            UltrasonicZone.FRONT_LEFT_CORNER, UltrasonicZone.FRONT_LEFT,
            UltrasonicZone.FRONT_CENTER_LEFT, UltrasonicZone.FRONT_CENTER,
            UltrasonicZone.FRONT_CENTER_RIGHT, UltrasonicZone.FRONT_RIGHT,
            UltrasonicZone.FRONT_RIGHT_CORNER,
        ]
        front_y = [half_w, half_w*0.66, half_w*0.33, 0.0, -half_w*0.33, -half_w*0.66, -half_w]
        front_yaw = [-30, -15, -5, 0, 5, 15, 30]
        ultrasonics = []
        for i, (zone, y_pos, yaw) in enumerate(zip(zones_front, front_y, front_yaw)):
            ultrasonics.append(UltrasonicSensorConfig(
                name=f"us_front_{i}", zone=zone,
                placement=SensorPlacement(x=half_l, y=y_pos, z=bumper_h, yaw=yaw),
                max_range=5.0,
            ))
        zones_rear = [
            UltrasonicZone.REAR_LEFT_CORNER, UltrasonicZone.REAR_LEFT,
            UltrasonicZone.REAR_CENTER_LEFT, UltrasonicZone.REAR_CENTER,
            UltrasonicZone.REAR_CENTER_RIGHT, UltrasonicZone.REAR_RIGHT,
            UltrasonicZone.REAR_RIGHT_CORNER,
        ]
        rear_yaw = [-150, -165, -175, 180, 175, 165, 150]
        for i, (zone, y_pos, yaw) in enumerate(zip(zones_rear, front_y, rear_yaw)):
            ultrasonics.append(UltrasonicSensorConfig(
                name=f"us_rear_{i}", zone=zone,
                placement=SensorPlacement(x=-half_l, y=y_pos, z=bumper_h, yaw=yaw),
                max_range=5.0,
            ))
        zones_left = [UltrasonicZone.LEFT_FRONT, UltrasonicZone.LEFT_CENTER, UltrasonicZone.LEFT_REAR]
        left_x = [half_l * 0.5, 0.0, -half_l * 0.5]
        for i, (zone, x_pos) in enumerate(zip(zones_left, left_x)):
            ultrasonics.append(UltrasonicSensorConfig(
                name=f"us_left_{i}", zone=zone,
                placement=SensorPlacement(x=x_pos, y=half_w, z=bumper_h + 0.2, yaw=-90),
                max_range=3.0,
            ))
        zones_right = [UltrasonicZone.RIGHT_FRONT, UltrasonicZone.RIGHT_CENTER, UltrasonicZone.RIGHT_REAR]
        for i, (zone, x_pos) in enumerate(zip(zones_right, left_x)):
            ultrasonics.append(UltrasonicSensorConfig(
                name=f"us_right_{i}", zone=zone,
                placement=SensorPlacement(x=x_pos, y=-half_w, z=bumper_h + 0.2, yaw=90),
                max_range=3.0,
            ))
        return ultrasonics


def create_custom_config(
    num_cameras=6, num_ultrasonics=20,
    camera_placements=None, ultrasonic_placements=None,
    max_speed_mph=20.0, **vehicle_kwargs
) -> VehicleConfig:
    config = VehicleConfig(max_speed_mph=max_speed_mph, **vehicle_kwargs)
    if camera_placements is not None:
        cameras = []
        positions = list(CameraPosition)
        for i, cp in enumerate(camera_placements):
            pos = cp.get("position", positions[i % len(positions)])
            if isinstance(pos, str):
                pos = CameraPosition(pos)
            placement = SensorPlacement(**cp.get("placement", {"x": 0, "y": 0, "z": 1.0}))
            cameras.append(CameraSensorConfig(
                name=cp.get("name", f"cam_{i}"), position=pos, placement=placement,
                resolution=cp.get("resolution", (640, 480)),
                fov_horizontal=cp.get("fov_horizontal", 120),
                fov_vertical=cp.get("fov_vertical", 90),
            ))
        config.sensor_config.cameras = cameras
    if ultrasonic_placements is not None:
        ultrasonics = []
        zones = list(UltrasonicZone)
        for i, up in enumerate(ultrasonic_placements):
            zone = up.get("zone", zones[i % len(zones)])
            if isinstance(zone, str):
                zone = UltrasonicZone(zone)
            placement = SensorPlacement(**up.get("placement", {"x": 0, "y": 0, "z": 0.4}))
            ultrasonics.append(UltrasonicSensorConfig(
                name=up.get("name", f"us_{i}"), zone=zone, placement=placement,
                max_range=up.get("max_range", 5.0),
                beam_angle=up.get("beam_angle", 30.0),
            ))
        config.sensor_config.ultrasonics = ultrasonics
    return config