src/core/robot_design.py

from dataclasses import dataclass
from typing import List, Dict, Any
import json

@dataclass
class WheelSpec:
    radius: float
    width: float
    position: List[float]
    friction: float

@dataclass
class MotorSpec:
    max_force: float
    max_velocity: float

@dataclass
class BodySpec:
    mass: float
    dimensions: List[float]

class RobotDesign:
    def __init__(self):
        self.body: BodySpec = None
        self.wheels: List[WheelSpec] = []
        self.motors: List[MotorSpec] = []

    def to_dict(self) -> Dict[str, Any]:
        return {
            "body": {
                "mass": self.body.mass,
                "dimensions": self.body.dimensions
            },
            "wheels": [
                {
                    "radius": w.radius,
                    "width": w.width,
                    "position": w.position,
                    "friction": w.friction
                } for w in self.wheels
            ],
            "motors": [
                {
                    "max_force": m.max_force,
                    "max_velocity": m.max_velocity
                } for m in self.motors
            ]
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict())

class RobotDesignBuilder:
    def __init__(self):
        self._design = RobotDesign()

    def set_body(self, mass: float, dimensions: List[float]) -> 'RobotDesignBuilder':
        self._design.body = BodySpec(mass=mass, dimensions=dimensions)
        return self

    def add_wheel(self, radius: float, width: float, position: List[float], friction: float) -> 'RobotDesignBuilder':
        self._design.wheels.append(WheelSpec(radius, width, position, friction))
        return self

    def add_motor(self, max_force: float, max_velocity: float) -> 'RobotDesignBuilder':
        self._design.motors.append(MotorSpec(max_force, max_velocity))
        return self

    def build(self) -> RobotDesign:
        if not self._design.body:
            raise ValueError("Robot design must have a body")
        if not self._design.wheels:
            raise ValueError("Robot design must have at least one wheel")
        if not self._design.motors:
            raise ValueError("Robot design must have at least one motor")
        return self._design

    @staticmethod
    def get_default_design() -> RobotDesign:
        builder = RobotDesignBuilder()
        builder.set_body(10.0, [1.0, 1.0, 0.5])
        
        # Add four wheels
        wheel_positions = [
            [0.5, 0.5, 0.2],
            [-0.5, 0.5, 0.2],
            [0.5, -0.5, 0.2],
            [-0.5, -0.5, 0.2]
        ]
        for pos in wheel_positions:
            builder.add_wheel(0.2, 0.1, pos, 0.8)
        
        # Add four motors
        for _ in range(4):
            builder.add_motor(100.0, 10.0)
        
        return builder.build()