RustPhysicsMQ/simulation/src/simulation.rs

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

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

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

enum StepOutcome {
    Continue,
    Exit,
}

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(
                    current_time,
                    Input {
                        joystick: controller::input_from_keyboard(),
                        mode: ModeInput::Acro,
                        ..Default::default()
                    },
                );
                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::controller::stacked::StackedController>()
            {
                cont.set_input(current_input);
            }
            self.drone.process_controller_tick(&mut self.world);
        }
        self.drone.process_tick(&mut self.world);

        /* Logging */

        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(cont) = self
            .drone
            .controller
            .as_mut_any()
            .downcast_mut::<crate::controller::stacked::StackedController>()
        {
            let target_mot = cont.rate_rt.last_output.unwrap_or_default().0;

            let target_angular_vel = if current_input.mode != ModeInput::Acro {
                cont.rotation_rt.last_output.unwrap_or_default().0
            } else {
                na::vector![
                    current_input.joystick.roll_input,
                    current_input.joystick.pitch_input,
                    current_input.joystick.yaw_input,
                ] * 3.14
            };

            let target_rot = if current_input.mode != ModeInput::Rotation {
                cont.linaccel_rt.last_output.unwrap_or_default().0 .0
            } else {
                na::vector![
                    current_input.rotation.roll,
                    current_input.rotation.pitch,
                    current_input.rotation.yaw,
                ] * 3.14
            };

            let accel_target = na::vector![
                current_input.acceleration.x,
                current_input.acceleration.y,
                current_input.acceleration.z
            ];

            let accel = self.drone.get_accel();

            if let Some(logger) = &mut self.logger {
                logger
                    .log(&SimLogRow {
                        time: self.world.get_time(),
                        linvel_current: (*self.drone.get_rb(&self.world).linvel()).into(),
                        accel_current: accel.into(),
                        accel_target: accel_target.into(),
                        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(target_rot)
                            .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(())
    }
}