RustPhysicsMQ/src/simulation.rs

use macroquad::prelude as mq;
use nalgebra::{self as na, vector};
use rapier3d::prelude as rp;
use std::error::Error;

use crate::{
    drone::{
        controller::DroneController,
        input::{Input, InputRecording, ModeInput},
        Drone,
    },
    engine::World,
    logger::{Logger, SimLogRow},
    rendering::Renderer,
};

pub enum SimMode {
    Record(InputRecording, String),
    Playback(InputRecording, f32),
}

enum StepOutcome {
    Continue,
    Exit,
}

pub struct DataResultRecord {
    pub time: f32,
    pub current_angular_velocity: na::Vector3<f32>,
    pub target_angular_velocity: na::Vector3<f32>,
    pub applied_motor_offset: na::Vector3<f32>,
    pub desired_motor_offset: na::Vector3<f32>,
}

pub struct Simulation {
    pub drone: Drone,
    pub world: World,
    pub mode: SimMode,

    logger: Option<Logger>,
    drone_tick_rate: u64,
}

impl Simulation {
    pub fn new(
        drone: Drone,
        world: World,
        mode: SimMode,
        results_file: Option<String>,
        drone_tick_rate: u64,
    ) -> Self {
        let logger = match &results_file {
            Some(path) => Some(Logger::new(path).unwrap()),
            None => None,
        };

        let mut s = Self {
            drone,
            world,
            mode,
            logger,
            drone_tick_rate,
        };

        s.world.register_free_collider(
            rp::ColliderBuilder::cuboid(30.0, 1.0, 30.0)
                .translation(na::vector![0.0, -2.0, 0.0])
                .restitution(0.5)
                .sensor(true)
                .build(),
            None,
        );

        return s;
    }

    pub async fn run_and_render(&mut self) -> Result<(), Box<dyn Error>> {
        let mut renderer: Renderer = Renderer::new(Box::new(crate::camera::AttachedCamera::new(
            self.drone.rb_handle,
            na::vector![1.0, 0.0, 0.0],
            na::vector![0.5, 0.0, 0.0],
        )));

        renderer.light.set_location(
            na::vector![70.0, 150.0, -90.0].into(),
            na::vector![-0.4, -0.7, 0.6].into(),
        );
        renderer.update_light(&self.world);

        loop {
            match self.step()? {
                StepOutcome::Continue => {
                    if self.world.tick
                        % ((self.world.integration_parameters.inv_dt() / 60.0) as u64)
                        == 0
                    {
                        renderer.update_camera(&self.world);
                        renderer.draw(&mut self.world);
                        mq::next_frame().await;
                    }
                }
                StepOutcome::Exit => {
                    self.shutdown()?;
                    return Ok(());
                }
            }
        }
    }

    pub fn run(&mut self) -> Result<(), Box<dyn Error>> {
        loop {
            match self.step()? {
                StepOutcome::Continue => {}
                StepOutcome::Exit => {
                    self.shutdown()?;
                    return Ok(());
                }
            }
        }
    }

    fn step(&mut self) -> Result<StepOutcome, Box<dyn Error>> {
        let current_input: Input;
        let current_time = self.world.get_time() as f32;

        match &mut self.mode {
            SimMode::Record(recording, _) => {
                current_input = recording.add_input_from_keyboard(current_time);
                if mq::is_key_pressed(mq::KeyCode::Q) {
                    return Ok(StepOutcome::Exit);
                }
            }
            SimMode::Playback(recording, start_time) => {
                let playback_time = current_time - *start_time;
                current_input = recording.get_input(playback_time);

                if recording.has_ended(playback_time) {
                    println!("Playback ended.");
                    return Ok(StepOutcome::Exit);
                }
            }
        }

        // --- Physics ---
        self.world.step();
        if self.world.tick
            % ((self.world.integration_parameters.inv_dt() / self.drone_tick_rate as f32) as u64)
            == 0
        {
            if let Some(cont) = self
                .drone
                .controller
                .as_mut_any()
                .downcast_mut::<crate::drone::stacked::StackedController>()
            {
                cont.set_input(current_input);
            }
            self.drone.process_controller_tick(&mut self.world);
        }
        self.drone.process_tick(&mut self.world);

        let mut target_angular_vel: na::Vector3<f32> = na::vector![
            current_input.joystick.roll_input,
            current_input.joystick.pitch_input,
            current_input.joystick.yaw_input,
        ] * 3.14;

        let mut target_mot: na::Vector3<f32> = Default::default();
        if let Some(cont) = self
            .drone
            .controller
            .as_mut_any()
            .downcast_mut::<crate::drone::stacked::StackedController>()
        {
            if current_input.mode != ModeInput::Acro {
                target_angular_vel = cont.rotation_rt.last_output.unwrap_or_default().0;
            }
            target_mot = cont.rate_rt.last_output.unwrap_or_default().0;
        }
        /*
        # roll, pitch, yaw
        motor_map = [
            [1.0, -1.0, 1.0],   # Front Right
            [-1.0, -1.0, -1.0], # Front Left
            [-1.0, 1.0, 1.0],   # Rear Left
            [1.0, 1.0, -1.0],   # Rear Right
        ]
        */
        let applied_motor_diff: na::Vector3<f32> = na::vector![
            (self.drone.current_throttles[0] + self.drone.current_throttles[3]
                - self.drone.current_throttles[1]
                - self.drone.current_throttles[2]),
            (self.drone.current_throttles[2] + self.drone.current_throttles[3]
                - self.drone.current_throttles[0]
                - self.drone.current_throttles[1]),
            (self.drone.current_throttles[0] + self.drone.current_throttles[2]
                - self.drone.current_throttles[1]
                - self.drone.current_throttles[3])
        ] / 2.0;

        if let Some(logger) = &mut self.logger {
            logger
                .log(&SimLogRow {
                    time: self.world.get_time(),
                    angvel_target: target_angular_vel.into(),
                    angvel_current: self
                        .drone
                        .get_rot(&self.world)
                        .inverse()
                        .transform_vector(&self.drone.get_angvel(&self.world))
                        .into(),
                    rot_current: self.drone.get_rot(&self.world).scaled_axis().into(),
                    rot_target: na::UnitQuaternion::from_scaled_axis(vector![
                        current_input.rotation.roll,
                        current_input.rotation.pitch,
                        current_input.rotation.yaw
                    ])
                    .scaled_axis()
                    .into(),
                    mot_current: applied_motor_diff.into(),
                    mot_target: target_mot.into(),
                })
                .unwrap();
        }
        Ok(StepOutcome::Continue)
    }

    fn shutdown(&mut self) -> Result<(), Box<dyn Error>> {
        // Save input recording if needed
        if let SimMode::Record(recording, dest) = &self.mode {
            recording.save_to_file(dest)?;
            println!("Input recording saved to {}", dest);
        }
        if let Some(logger) = &mut self.logger {
            logger.flush()?;
        }

        Ok(())
    }
}