renderbug-bike/src/ego/engine.rs

use embassy_sync::pubsub::DynPublisher;
use embassy_time::{Duration, Instant, WithTimeout};
use nalgebra::{Rotation3, Vector2, Vector3, ComplexField, RealField};
use log::*;

use core::fmt::Debug;

use crate::{Breaker, CircularBuffer, ego::{heading::HeadingEstimator, kalman::Ekf2D, orientation::OrientationEstimator}, events::{Personality, Prediction, SensorSource, SensorState}, idle::IdleClock};

#[derive(PartialEq, Debug, Default, Clone, Copy)]
pub enum MotionState {
    #[default]
    Stationary,
    Steady,
    Accelerating,
    Decelerating
}

pub struct BikeStates {
    /// Prediction models
    orientation: OrientationEstimator,
    heading: HeadingEstimator,
    speedo: CircularBuffer<f32, 300>,
    kf: Ekf2D,
    
    last_stamp: Instant,
    last_fix: Vector2<f64>, // The most recent GPS value
    reference_fix: Option<Vector2<f64>>, // The first GPS value, which is used to make sense out of the EKF output which is in meters offset from this point

    // State switches
    has_motion_frame: Breaker<bool>,
    motion_state: Breaker<MotionState>,
    acquiring_data: Breaker<bool>,
    // FIXME: pub
    pub has_gps_fix: Breaker<bool>,
    predicted_velocity: Breaker<f32>,
    reported_velocity: Breaker<f32>,
    predicted_location: Breaker<Vector2<f64>>,
    steady_timer: IdleClock,

    parking_timer: IdleClock,
    sleep_timer: IdleClock,
}

/// Above this speed, the system should consider waking up from sleep mode
const BUMP_SPEED_NOISE_GATE: f32 = 0.1;

/// Above this speed, the system should be fully awake and ready to start moving around with purpose
const WAKEUP_SPEED_NOISE_GATE: f32 = 0.5;

/// Above this average speed, we are considered to be moving around with purpose
const MOVING_SPEED_NOISE_GATE: f32 = 1.0;

/// IMU readings of at least this value are considered valid motion. Below this value, the signal is considered noise while stationary
const MOTION_NOISE_GATE: f32 = 0.8;

/// If the system's calculated speed is increasing or decreasing by this value or more, we are accelerating/decelerating
const ACCELERATION_NOISE_GATE: f32 = 1.0;

/// When we get a heading via GPS, this determines how much weight that value has when blended into the system.
const GPS_HEADING_ALPHA_CORRECTION: f32 = 0.9;

const TIMEOUT: Duration = Duration::from_millis(3);

macro_rules! publish {
    ($predictions:expr, $prediction:expr) => {
        $predictions.publish($prediction).with_timeout(TIMEOUT).await.expect("Could not commit IMU data in time. Is the prediction bus stalled?")
    }
}

impl Debug for BikeStates {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("BikeStates")
            .field("has_orientation", &self.orientation.is_ready())
            .field("heading", &self.heading.heading())
            .field("motion_state", &self.motion_state)
            .field("predicted_location", &self.predicted_location)
            .field("predicted_velocity", &self.predicted_velocity)
            .field("steady_timer", &self.steady_timer)
            .field("parking_timer", &self.parking_timer)
            .field("sleep_timer", &self.sleep_timer)
            .finish()
    }
}

impl BikeStates {
    pub fn insert_gps(&mut self, gps_pos: Vector2<f64>) {
        self.acquiring_data.set(true);
        match self.reference_fix {
            None => {
                self.reference_fix = Some(gps_pos);
                self.last_fix = gps_pos;
                self.has_gps_fix.set(true);
            },
            Some(coords) => {
                if self.last_fix != coords {
                    let gps_heading = self.last_fix.angle(&coords);
                    self.heading.correct(gps_heading as f32, GPS_HEADING_ALPHA_CORRECTION);
                }

                let delta = gps_to_local_meters_haversine(&coords, &gps_pos);

                self.kf.update_gps(delta.x, delta.y);

                self.last_fix = coords;
            }
        }
    }

    pub fn insert_imu(&mut self, accel: Vector3<f32>, gyro: Vector3<f32>) {
        self.acquiring_data.set(true);
        self.orientation.insert(accel);

        if self.orientation.has_down() {
            let body_accel = self.orientation.apply(accel);
            let body_gyro = self.orientation.apply(gyro);

            let cur_stamp = Instant::now();
            let dt = (cur_stamp - self.last_stamp).as_millis() as f32 / 1000.0;
            self.last_stamp = cur_stamp;

            self.heading.update(body_gyro.z, dt);
            let heading_rotation = Rotation3::from_axis_angle(&Vector3::z_axis(), self.heading.heading());

            let enu_rotated = heading_rotation * body_accel;
            if body_accel.xy().magnitude() >= MOTION_NOISE_GATE {
                self.kf.predict(enu_rotated.xy(), body_gyro.z, dt);
            } else {
                // Otherwise, we are standing stationary and should insert accel=0 data into the model
                self.kf.update_zupt();
            }

            if self.orientation.is_ready() {
                self.has_motion_frame.set(true);
            }
        }
    }

    pub async fn commit(&mut self, predictions: &DynPublisher<'static, Prediction>) {
        let last_motion = self.motion_state.value;

        if let Some(true) = self.acquiring_data.read_tripped() {
            publish!(predictions, Prediction::SensorStatus(SensorSource::MotionFrame, SensorState::AcquiringFix));
            publish!(predictions, Prediction::SensorStatus(SensorSource::Location, SensorState::AcquiringFix));
        }

        if let Some(true) = self.has_motion_frame.read_tripped() {
            publish!(predictions, Prediction::SensorStatus(SensorSource::MotionFrame, SensorState::Online));
        }

        match self.has_gps_fix.read_tripped() {
            None => (),
            Some(true) => publish!(predictions, Prediction::SensorStatus(SensorSource::Location, SensorState::Online)),
            Some(false) => publish!(predictions, Prediction::SensorStatus(SensorSource::Location, SensorState::Degraded)),
        }

        let est = self.kf.x;

        let position = self.reference_fix.as_ref().map(|start_coords| local_to_gps(*start_coords, Vector2::new(est.x as f64, est.y as f64)));
        let velocity = Vector2::new(est[2], est[3]);

        // If we have a new location, update the predictions
        if let Some(pos) = position {
            self.predicted_location.set(pos);
            if let Some(pos) = self.predicted_location.read_tripped() {
                publish!(predictions, Prediction::Location(pos));
            }
        }

        // If we have a new predicted velocity, update the acceleration status
        self.predicted_velocity.set(velocity.norm());
        if let Some(current_prediction) = self.predicted_velocity.read_tripped() {
            // Add the prediction into the speedometer model
            self.speedo.insert(current_prediction);
            // If the model has enough samples to report useful data, we can start analyzing the motion trends
            if self.speedo.is_filled() {
                // Calculate if the velocity is increasing, decreasing, or mostly the same
                let trend = self.speedo.data().windows(2).map(|n| {
                    n[1] - n[0]
                }).sum::<f32>();
                // Also grab the average velocity of the last few sample periods
                let average_speed = self.speedo.mean();
                trace!("prediction={current_prediction:?} mean={average_speed}");

                // Reported velocity is kept only to the first decimal, so we aren't spamming the system with floating point noise
                self.reported_velocity.set((average_speed * 10.0).trunc() / 10.0);
                if let Some(reported) = self.reported_velocity.read_tripped() {
                    publish!(predictions, Prediction::Velocity(reported));
                }

                // We only want to wake up from sleep if our current velocity is something like a bump
                if average_speed > BUMP_SPEED_NOISE_GATE && self.sleep_timer.wake() {
                    warn!("Waking from sleep into idle mode");
                    publish!(predictions, Prediction::SetPersonality(Personality::Waking));
                    publish!(predictions, Prediction::SetPersonality(Personality::Parked));
                }
                // Here, we additionally release the parking brake if we are currently parked and reading substantial movement
                if average_speed > WAKEUP_SPEED_NOISE_GATE && self.parking_timer.wake() {
                    warn!("Disengaging parking brake");
                    publish!(predictions, Prediction::SetPersonality(Personality::Active));
                }

                // If the total slope is more upwards than not, we are accelerating.
                if trend >= ACCELERATION_NOISE_GATE {
                    self.motion_state.set(MotionState::Accelerating);
                    self.steady_timer.wake();
                } else if trend <= -ACCELERATION_NOISE_GATE {
                    self.steady_timer.wake();
                    self.motion_state.set(MotionState::Decelerating);
                } else if self.steady_timer.check() && average_speed > MOVING_SPEED_NOISE_GATE {
                    // If we haven't changed our acceleration for a while, and we still have speed, we are moving at a steady pace
                    self.motion_state.set(MotionState::Steady);
                } else if current_prediction <= 1.0 && average_speed <= MOVING_SPEED_NOISE_GATE {
                    // If the average and instantaneous speed is rather low, we are probably stationary!
                    self.motion_state.set(MotionState::Stationary);
                }
            }
        }
        // And if the motion status has changed, send it out
        if let Some(state) = self.motion_state.read_tripped() {
            //debug!("state={state:?} trend={trend} mean={mean} v={v}");
            //if state != MotionState::Stationary {
            //    warn!("Active due to motion");
            //    publish!(predictions, Prediction::SetPersonality(Personality::Active));
            //}
            publish!(predictions, Prediction::Motion { prev: last_motion, next: state });
        } else if self.motion_state.value == MotionState::Stationary {
            // Finally, if we are stationary, check our parking and sleep timers
            if self.parking_timer.check() {
                warn!("Engaging parking brake");
                publish!(predictions, Prediction::SetPersonality(Personality::Parked));
            }

            if self.sleep_timer.check() {
                warn!("Sleeping!");
                publish!(predictions, Prediction::SetPersonality(Personality::Sleeping));
            }
        }
    }
}

impl Default for BikeStates {
    fn default() -> Self {
        Self {
            motion_state: Default::default(),
            orientation: Default::default(),
            last_stamp: Instant::now(),
            speedo: Default::default(),
            heading: Default::default(),
            has_motion_frame: Default::default(),
            kf: Default::default(),
            steady_timer: IdleClock::new(Duration::from_secs(3)),
            last_fix: Default::default(),
            reference_fix: Default::default(),
            has_gps_fix: Default::default(),
            predicted_location: Default::default(),
            predicted_velocity: Default::default(),
            reported_velocity: Default::default(),
            acquiring_data: Default::default(),
            parking_timer: IdleClock::new(Duration::from_secs(10)),
            sleep_timer: IdleClock::new(Duration::from_secs(30))
        }
    }
}

/// Compute displacement vector (east, north) in meters between two GPS coordinates
/// Input: Vector2::new(lat_deg, lon_deg)
pub fn gps_to_local_meters_haversine(a: &Vector2<f64>, b: &Vector2<f64>) -> Vector2<f32> {
    let r = 6_371_000.0; // Earth radius in meters

    // Convert to radians
    let lat1 = a.x.to_radians();
    let lon1 = a.y.to_radians();
    let lat2 = b.x.to_radians();
    let lon2 = b.y.to_radians();

    // Differences
    let dlat = lat2 - lat1;
    let dlon = lon2 - lon1;

    // Haversine formula for great-circle distance
    let hav = (dlat / 2.0).sin().powi(2)
        + lat1.cos() * lat2.cos() * (dlon / 2.0).sin().powi(2);
    let c = 2.0 * hav.sqrt().atan2((1.0 - hav).sqrt());
    let distance = r * c;

    // Bearing from point A to B
    let y = dlon.sin() * lat2.cos();
    let x = lat1.cos() * lat2.sin() - lat1.sin() * lat2.cos() * dlon.cos();
    let bearing = y.atan2(x); // radians, from north, clockwise

    // Convert to local east/north vector
    let dx = distance * bearing.sin(); // east
    let dy = distance * bearing.cos(); // north

    Vector2::new(dx as f32, dy as f32)
}

/// Convert a local displacement (east, north) in meters back to GPS coordinates.
/// Input:
///   ref_gps = Vector2::new(lat_deg, lon_deg)
///   disp    = Vector2::new(dx_east_m, dy_north_m)
/// Output:
///   Vector2::new(lat_deg, lon_deg)
pub fn local_to_gps(ref_gps: Vector2<f64>, disp: Vector2<f64>) -> Vector2<f64> {
    let r = 6_371_000.0; // Earth radius in meters

    let lat0 = ref_gps.x.to_radians();
    let lon0 = ref_gps.y.to_radians();

    // Displacement in meters
    let dx = disp.x; // east
    let dy = disp.y; // north

    // Distance and bearing
    let distance = (dx.powi(2) + dy.powi(2)).sqrt();
    let bearing = dx.atan2(dy); // atan2(east, north)

    // Destination point on sphere
    let lat = (lat0.sin() * (distance / r).cos()
        + lat0.cos() * (distance / r).sin() * bearing.cos())
        .asin();

    let lon = lon0
        + (bearing.sin() * (distance / r).sin() * lat0.cos())
            .atan2((distance / r).cos() - lat0.sin() * lat.sin());

    Vector2::new(lat.to_degrees(), lon.to_degrees())
}