sim_robot.py

import asyncio
import logging
import sys
import math

# Add the current directory to path
sys.path.append('.')

from hirobo import Robot, Subsystem, Command, Motor, Button
from hirobo.dashboard.dashboard_server import DashboardServer

# Mock FastPID if not built
try:
    from hirobo._hirobo_native import FastPID
except ImportError:
    class FastPID:
        def __init__(self, p, i, d): 
            self.kP, self.kI, self.kD = p, i, d
            self.integral = 0
            self.prev_error = 0
        def calculate(self, target, current, dt):
            error = target - current
            self.integral += error * dt
            derivative = (error - self.prev_error) / dt
            self.prev_error = error
            return (self.kP * error) + (self.kI * self.integral) + (self.kD * derivative)
        def reset(self):
            self.integral = 0
            self.prev_error = 0

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("SimRobot")

class SimDrivetrain(Subsystem):
    def __init__(self):
        super().__init__()
        self.motor = Motor()
        self.pid = FastPID(0.5, 0.0, 0.05) # Tuned for simulation
        self.target_position = 0.0
        
    def periodic(self):
        # Simulate physics
        self.motor.update(0.02)

class SimRobot(Robot):
    def robot_init(self):
        logger.info("Initializing Simulation Robot...")
        self.drivetrain = SimDrivetrain()
        self.dashboard = DashboardServer()
        
        # Setup dashboard callbacks
        self.dashboard.button_states = {}
        self.dashboard.target_callback = self.on_target_change
        
        # Create a virtual button bound to the dashboard state
        self.virtual_btn = Button(lambda: self.dashboard.button_states.get('virtual_btn', False))
        
        # Bind button using decorator syntax
        @self.virtual_btn.while_held
        def spin_action():
            # Simple function-based command
            # Note: For continuous action like spinning, a proper Command class is better 
            # because InstantCommand runs once. But for 'while_held', the scheduler 
            # will cancel it when released. If it's Instant, it finishes immediately anyway.
            # So for 'while_held' with a function, it's tricky if it's not a loop.
            # Let's use 'when_pressed' for the function demo, and keep SpinCommand for held.
            logger.info("Virtual Button Pressed (Function)!")

        # Re-bind the SpinCommand for the actual behavior
        self.virtual_btn.while_held(SpinCommand(self.drivetrain))
        
        # Schedule startup command
        self.scheduler.schedule(DashboardStartupCommand(self))
        
        # Default command
        self.scheduler.schedule(AutoMoveCommand(self.drivetrain))

    def on_target_change(self, value):
        logger.info(f"Target changed to {value}")
        # Update the AutoMoveCommand's target if it's running, 
        # or we could set a global target variable
        pass

    async def telemetry_loop(self):
        while True:
            data = {
                'position': self.drivetrain.motor.get_position(),
                'velocity': self.drivetrain.motor.get_velocity(),
                'target': self.drivetrain.target_position
            }
            await self.dashboard.broadcast(data)
            await asyncio.sleep(0.05) # 20Hz update for UI

class AutoMoveCommand(Command):
    def __init__(self, drivetrain: SimDrivetrain):
        super().__init__()
        self.drivetrain = drivetrain
        self.add_requirements(drivetrain)
        self.timer = 0
        
    def execute(self):
        self.timer += 0.02
        # Sine wave target
        target = 50.0 * math.sin(self.timer * 0.5)
        self.drivetrain.target_position = target
        
        current = self.drivetrain.motor.get_position()
        output = self.drivetrain.pid.calculate(target, current, 0.02)
        self.drivetrain.motor.set(output)

class SpinCommand(Command):
    def __init__(self, drivetrain: SimDrivetrain):
        super().__init__()
        self.drivetrain = drivetrain
        self.add_requirements(drivetrain)
        
    def initialize(self):
        logger.info("SpinCommand Initialized!")
        
    def execute(self):
        # Spin the motor at full power
        self.drivetrain.motor.set(1.0)
        
    def end(self, interrupted):
        logger.info("SpinCommand Ended")
        self.drivetrain.motor.set(0.0)

class DashboardStartupCommand(Command):
    def __init__(self, robot: SimRobot):
        super().__init__()
        self.robot = robot
        
    def initialize(self):
        logger.info("Starting Dashboard Server...")
        loop = asyncio.get_running_loop()
        loop.create_task(self.robot.dashboard.start())
        loop.create_task(self.robot.telemetry_loop())

    def is_finished(self):
        return True

if __name__ == "__main__":
    robot = SimRobot()
    try:
        robot.start()
    except KeyboardInterrupt:
        pass