RustPhysicsMQ/src/simulation.rs

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

use crate::{
    drone::{
        controller::DroneController,
        input::{InputRecording, JoystickInput},
        Drone,
    },
    engine::World,
    rendering::Renderer,
};

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

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>,
}

#[derive(Default)]
pub struct SimulationState {
    pub data_results: Vec<DataResultRecord>,
}

pub struct Simulation {
    pub drone: Drone,
    pub world: World,
    pub mode: SimMode,
    results_file: Option<String>,
    drone_tick_rate: u64,
    state: SimulationState,
}

impl Simulation {
    pub fn new(
        drone: Drone,
        world: World,
        mode: SimMode,
        results_file: Option<String>,
        drone_tick_rate: u64,
    ) -> Self {
        let mut s = Self {
            drone,
            world,
            mode,
            results_file,
            drone_tick_rate,
            state: SimulationState::default(),
        };

        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);

        let mut current_input: JoystickInput;

        loop {
            renderer.update_camera(&self.world);

            // --- Input handling ---
            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) {
                        self.shutdown()?;
                        return Ok(());
                    }
                }
                SimMode::Playback(recording, start_time) => {
                    let playback_time = current_time - *start_time;
                    current_input = recording.get_input(playback_time);

                    if recording.ended(playback_time) {
                        println!("Playback ended.");
                        break;
                    }
                }
            }

            // --- 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::pidcontroller::PIDController>()
                {
                    cont.set_input(current_input);
                }
                self.drone.process_tick(&mut self.world);
            }

            let target_angular_vel: na::Vector3<f32>;
            let desired_motor_diff: na::Vector3<f32>;
            let applied_motor_diff: na::Vector3<f32>;
            if let Some(cont) = self
                .drone
                .controller
                .as_mut_any()
                .downcast_mut::<crate::drone::pidcontroller::PIDController>()
            {
                desired_motor_diff = cont.get_desired_motor_diffs();
                target_angular_vel = cont.get_desired_angular_velocity();

                /*
                let t = [
                    throttle - pitch + yaw + roll,
                    throttle - pitch - yaw - roll,
                    throttle + pitch + yaw - roll,
                    throttle + pitch - yaw + roll,
                ];
                */
                let m = cont.get_motor_throttles();
                applied_motor_diff = na::vector![
                    (m[2] + m[3] - m[0] - m[1]),
                    (m[0] + m[2] - m[1] - m[3]),
                    (m[0] + m[3] - m[1] - m[2])
                ] / 4.0;
            } else {
                target_angular_vel = na::vector![0.0, 0.0, 0.0];
                desired_motor_diff = na::vector![0.0, 0.0, 0.0];
                applied_motor_diff = na::vector![0.0, 0.0, 0.0];
            }

            self.state.data_results.push(DataResultRecord {
                time: self.world.get_time(),
                current_angular_velocity: self
                    .drone
                    .get_rot(&self.world)
                    .inverse()
                    .transform_vector(&self.drone.get_angvel(&self.world)),
                target_angular_velocity: target_angular_vel,
                applied_motor_offset: applied_motor_diff,
                desired_motor_offset: desired_motor_diff,
            });

            // --- Rendering ---
            if self.world.tick % ((self.world.integration_parameters.inv_dt() / 60.0) as u64) == 0 {
                renderer.draw(&mut self.world);
                mq::next_frame();
            }
        }

        self.shutdown()?;
        Ok(())
    }

    pub fn run(&mut self) -> Result<(), Box<dyn Error>> {
        let mut current_input: JoystickInput;

        loop {
            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) {
                        self.shutdown()?;
                        return Ok(());
                    }
                }
                SimMode::Playback(recording, start_time) => {
                    let playback_time = current_time - *start_time;
                    current_input = recording.get_input(playback_time);

                    if recording.ended(playback_time) {
                        println!("Playback ended.");
                        break;
                    }
                }
            }

            // --- 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::pidcontroller::PIDController>()
                {
                    cont.set_input(current_input);
                }
                self.drone.process_controller_tick(&mut self.world);
            }
            self.drone.process_tick(&mut self.world);

            let target_angular_vel: na::Vector3<f32>;
            let desired_motor_diff: na::Vector3<f32>;
            let applied_motor_diff: na::Vector3<f32>;
            if let Some(cont) = self
                .drone
                .controller
                .as_mut_any()
                .downcast_mut::<crate::drone::pidcontroller::PIDController>()
            {
                desired_motor_diff = cont.get_desired_motor_diffs();
                target_angular_vel = cont.get_desired_angular_velocity();

                /*
                let t = [
                    throttle - pitch + yaw + roll,
                    throttle - pitch - yaw - roll,
                    throttle + pitch + yaw - roll,
                    throttle + pitch - yaw + roll,
                ];
                */
                let m = self.drone.current_throttles;
                applied_motor_diff = na::vector![
                    (m[2] + m[3] - m[0] - m[1]),
                    (m[0] + m[2] - m[1] - m[3]),
                    (m[0] + m[3] - m[1] - m[2])
                ] / 2.0;
            } else {
                target_angular_vel = na::vector![0.0, 0.0, 0.0];
                desired_motor_diff = na::vector![0.0, 0.0, 0.0];
                applied_motor_diff = na::vector![0.0, 0.0, 0.0];
            }

            self.state.data_results.push(DataResultRecord {
                time: self.world.get_time(),
                current_angular_velocity: self
                    .drone
                    .get_rot(&self.world)
                    .inverse()
                    .transform_vector(&self.drone.get_angvel(&self.world)),
                target_angular_velocity: target_angular_vel,
                applied_motor_offset: applied_motor_diff,
                desired_motor_offset: desired_motor_diff,
            });
        }

        self.shutdown()?;
        Ok(())
    }

    pub fn save_logs_to_csv(&self, path: &str) -> std::io::Result<()> {
        use std::fs::File;
        use std::io::{BufWriter, Write};
        let file = File::create(path)?;
        let mut writer = BufWriter::with_capacity(1 << 20, file); // 1 MB buffer

        writeln!(
        writer,
        "time,target_x,target_y,target_z,current_x,current_y,current_z,error_x,error_y,error_z,mot_x,mot_y,mot_z,dmot_x,dmot_y,dmot_z"
    )?;

        for entry in &self.state.data_results {
            let tg = entry.target_angular_velocity;
            let cur = entry.current_angular_velocity;
            let err = tg - cur;
            let mot = entry.applied_motor_offset;
            let dmot = entry.desired_motor_offset;

            writeln!(
                writer,
                "{},{},{},{},{},{},{},{},{},{},{},{},{},{},{},{}",
                entry.time,
                tg.x,
                tg.y,
                tg.z,
                cur.x,
                cur.y,
                cur.z,
                err.x,
                err.y,
                err.z,
                mot.x,
                mot.y,
                mot.z,
                dmot.x,
                dmot.y,
                dmot.z
            )?;
        }

        Ok(())
    }

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

        Ok(())
    }
}