RustPhysicsMQ/src/drone/input.rs

#![allow(unused)]
use macroquad::prelude as mq;

#[derive(Default, Clone, Copy, PartialEq)]
pub struct JoystickInput {
    // Value should be between 0 and 1
    pub throttle_input: f32,

    // Rotation Directions: https://upload.wikimedia.org/wikipedia/commons/c/c1/Yaw_Axis_Corrected.svg
    /*
     * Values should be between -1 and 1.
     */
    pub yaw_input: f32,
    pub pitch_input: f32,
    pub roll_input: f32,
}

impl JoystickInput {
    pub fn from_keyboard() -> Self {
        let mut input = Self::default();
        if mq::is_key_down(mq::KeyCode::W) {
            input.throttle_input = 1.0;
        }
        if mq::is_key_down(mq::KeyCode::S) {
            input.throttle_input = 0.0;
        }

        if mq::is_key_down(mq::KeyCode::A) {
            input.yaw_input = 1.0;
        } else if mq::is_key_down(mq::KeyCode::D) {
            input.yaw_input = -1.0;
        } else {
            input.yaw_input = 0.0;
        }

        if mq::is_key_down(mq::KeyCode::Up) {
            input.pitch_input = -1.0;
        } else if mq::is_key_down(mq::KeyCode::Down) {
            input.pitch_input = 1.0;
        } else {
            input.pitch_input = 0.0;
        }

        if mq::is_key_down(mq::KeyCode::Left) {
            input.roll_input = -1.0;
        } else if mq::is_key_down(mq::KeyCode::Right) {
            input.roll_input = 1.0;
        } else {
            input.roll_input = 0.0;
        }
        input = input.clamp();
        return input;
    }
}

#[derive(Default, Clone, Copy)]
pub struct InputRecord {
    pub input: JoystickInput,
    pub timestamp: f32,
}

impl InputRecord {
    pub fn record_from_keyboard(timestamp: f32) -> Self {
        Self {
            input: JoystickInput::from_keyboard(),
            timestamp,
        }
    }
}

#[derive(Default, Clone)]
pub struct InputRecording {
    pub input_records: Vec<InputRecord>,
}

impl InputRecording {
    pub fn new() -> Self {
        Self {
            input_records: Vec::new(),
        }
    }
    pub fn save_inputs_to_csv(&self, path: &str) -> std::io::Result<()> {
        use std::fs::File;
        use std::io::Write;

        let mut file = File::create(path)?;
        writeln!(file, "time,throttle,yaw,pitch,roll")?;

        for record in self.input_records.iter() {
            let inp = record.input;
            writeln!(
                file,
                "{},{},{},{},{}",
                record.timestamp,
                inp.throttle_input,
                inp.yaw_input,
                inp.pitch_input,
                inp.roll_input
            )?;
        }

        Ok(())
    }

    pub fn load_inputs_from_csv(path: &str) -> std::io::Result<Self> {
        use std::fs::File;
        use std::io::{BufRead, BufReader};

        let file = File::open(path)?;
        let reader = BufReader::new(file);

        let mut input_records = Vec::new();

        for (i, line) in reader.lines().enumerate() {
            let line = line?;
            if i == 0 {
                continue;
            }

            let parts: Vec<&str> = line.split(',').collect();
            if parts.len() != 5 {
                continue; // or return an error if you prefer strict parsing
            }

            let timestamp: f32 = parts[0]
                .parse()
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;

            let throttle: f32 = parts[1]
                .parse()
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
            let yaw: f32 = parts[2]
                .parse()
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
            let pitch: f32 = parts[3]
                .parse()
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
            let roll: f32 = parts[4]
                .parse()
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;

            let input = JoystickInput {
                throttle_input: throttle,
                yaw_input: yaw,
                pitch_input: pitch,
                roll_input: roll,
            };

            input_records.push(InputRecord { input, timestamp });
        }

        Ok(Self { input_records })
    }

    pub fn get_input(&self, time: f32) -> JoystickInput {
        /*
         * Binary search returns index to element as OK, or where the element could be placed to
         * keep order as Err, so if result is Ok return that input, if its Err, return the previous
         * input if it exists, if it doesn't (because time is before the first action in the
         * recorded sequence, return an empty input)
         */
        let res = self
            .input_records
            .binary_search_by(|probe| probe.timestamp.total_cmp(&time));
        match res {
            Ok(res) => {
                return self
                    .input_records
                    .get(res)
                    .unwrap_or(&InputRecord::default())
                    .input;
            }
            Err(mut res) => {
                if res > 0 {
                    res -= 1;
                }
                return self
                    .input_records
                    .get(0.max(res))
                    .unwrap_or(&InputRecord::default())
                    .input;
            }
        }
    }

    pub fn ended(&self, time: f32) -> bool {
        match self.input_records.last() {
            Some(record) => {
                return record.timestamp < time;
            }
            None => {
                return true;
            }
        }
    }

    /*
     * Current time should always be larger thant the last input records time.
     * This method is made for recording inputs in real time, not for retroactively adding
     *
     * Returns the addded Joystick Input
     */
    pub fn add_input_from_keyboard(&mut self, current_time: f32) -> JoystickInput {
        let input = JoystickInput::from_keyboard();
        let last_input = self.input_records.last();
        match last_input {
            Some(last_record) => {
                if last_record.input != input {
                    self.input_records.push(InputRecord {
                        input,
                        timestamp: current_time,
                    });
                }
            }
            None => {
                self.input_records.push(InputRecord {
                    input,
                    timestamp: current_time,
                });
            }
        }
        return input;
    }
}