use core::cell::RefCell;

use embassy_embedded_hal::shared_bus::asynch::i2c::I2cDevice;
use embassy_sync::blocking_mutex::raw::NoopRawMutex;
use embassy_sync::channel::DynamicSender;
use embassy_time::{Delay, Timer, Instant};
use esp_hal::{i2c::master::I2c, Async};
use log::*;
use mpu6050_dmp::accel::AccelFullScale;
use mpu6050_dmp::gyro::GyroFullScale;
use mpu6050_dmp::{address::Address, error_async::Error, sensor_async::Mpu6050};
use nalgebra::Vector3;
use crate::events::SensorSource;

use crate::{backoff::Backoff, events::Measurement};

const G: f32 = 9.80665;
const GYRO_SCALE: GyroFullScale = GyroFullScale::Deg2000;
const ACCEL_SCALE: AccelFullScale = AccelFullScale::G2;

#[embassy_executor::task]
pub async fn mpu_task(events: DynamicSender<'static, Measurement>, bus: I2cDevice<'static, NoopRawMutex, I2c<'static, Async>>) {
    let backoff = Backoff::from_millis(5);
    let busref = RefCell::new(Some(bus));

    backoff.forever().attempt::<_, (), ()>(async || {
        events.send(Measurement::SensorHardwareStatus(SensorSource::IMU, crate::events::SensorState::Offline)).await;
        let mut sensor = backoff.forever().attempt(async || {
            warn!("Initializing connection to MPU");
            match Mpu6050::new(busref.replace(None).unwrap(), Address::default()).await.map_err(|e| { e.i2c }) {
                Err(i2c) => {
                    busref.replace(Some(i2c));
                    Err(())
                },
                Ok(mut sensor) => {
                    events.send(Measurement::SensorHardwareStatus(SensorSource::IMU, crate::events::SensorState::AcquiringFix)).await;
                    match backoff.attempt(async || { mpu_init(&mut sensor).await }).await {
                        Err(_) => {
                            busref.replace(Some(sensor.release()));
                            Err(())
                        },
                        Ok(_) => Ok(sensor)
                    }
                }
            }
        }).await?;

        let sensor_ref = &mut sensor;

        events.send(Measurement::SensorHardwareStatus(SensorSource::IMU, crate::events::SensorState::Online)).await;
        //TODO: Need to read in a constant buffer of accelerometer measurements, which we can then use to determine where "forward" points in the body frame when converting from the sensor frame.
        // From there, we can rotate the body frame into the world frame using gps headings to generate a compass north
        fn lowpass(prev: f32, current: f32, alpha: f32) -> f32 {
            prev + alpha * (current - prev) // alpha in (0,1), small alpha = heavy smoothing
        }
        let mut prev_accel = Vector3::default();
        loop {
            match backoff.attempt(async || { sensor_ref.motion6().await }).await {
                Ok((accel_data, gyro_data)) => {

                    // Apply the initial alignment correction to put it into the body frame where X is forward
                    let scaled = accel_data.scaled(ACCEL_SCALE);
                    let adjusted = Vector3::new(
                        scaled.x(),
                        scaled.y(),
                        scaled.z()
                    ) * G;

                    let scaled = gyro_data.scaled(GYRO_SCALE);
                    let adjusted_gyro = Vector3::new(
                        scaled.x(),
                        scaled.y(),
                        scaled.z()
                    );

                    let filtered = Vector3::new(
                        lowpass(prev_accel.x, adjusted.x, 0.3),
                        lowpass(prev_accel.y, adjusted.y, 0.3),
                        lowpass(prev_accel.z, adjusted.z, 0.3),
                    );
                    prev_accel = adjusted;

                    events.send(
                        Measurement::IMU {
                            accel: filtered,
                            gyro: adjusted_gyro
                        }
                    ).await;
                    Timer::after_millis(10).await;
                },
                Err(e) => {
                    error!("Failed to read MPU motion data! {e:?}");
                    busref.replace(Some(sensor.release()));
                    return Err(());
                }
            };
        }
    }).await.ok();
}

async fn mpu_init(sensor: &mut Mpu6050<I2cDevice<'static, NoopRawMutex, I2c<'static, Async>>>) -> Result<(), Error<I2cDevice<'static, NoopRawMutex, I2c<'static, Async>>>> {
    if cfg!(feature="wokwi") {
        warn!("Initializing simulated MPU");
        Ok(())
    } else {
        let mut delay = Delay;
        let backoff = Backoff::from_millis(10);
        info!("Initializing DMP");
        let start = Instant::now();
        backoff.attempt(async || { sensor.initialize_dmp(&mut delay).await }).await?;
        info!("DMP initialized in {}ms", start.as_millis());
        Ok(())
    }
}