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::{address::Address, calibration::CalibrationParameters, error_async::Error, sensor_async::Mpu6050};
use mpu6050_dmp::calibration::CalibrationActions;
use nalgebra::Vector3;
use core::f32::consts::PI;
use crate::events::SensorSource;

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

const G: f32 = 9.80665;
const SENS_2G: f32 = 16384.0;
const DEG2RAD: f32 = PI / 180.0;

/// Convert raw accel (i16) to m/s^2
fn accel_raw_to_ms2(raw: i16) -> f32 {
    (raw as f32 / SENS_2G) * G
}

/// Convert raw gyro (i16) to rad/s. For MPU6050 typical gyro scale at ±2000 deg/s -> 16.4 LSB/deg/s
/// Adjust if you configured different full-scale range.
fn gyro_raw_to_rads(raw: i16) -> f32 {
    // lsb_per_dps e.g. 16.4 for ±2000 dps
    (raw as f32 / 16.4) * (DEG2RAD)
}

#[esp_hal::ram(unstable(rtc_fast, persistent))]
static mut MPU_WAS_CALIBRATED: u8 = 0;

#[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;
        events.send(Measurement::SensorHardwareStatus(SensorSource::GravityReference, crate::events::SensorState::AcquiringFix)).await;
        events.send(Measurement::SensorHardwareStatus(SensorSource::ForwardsReference, crate::events::SensorState::AcquiringFix)).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 adjusted = Vector3::new(
                        accel_raw_to_ms2(accel_data.x()),
                        accel_raw_to_ms2(accel_data.y()),
                        accel_raw_to_ms2(accel_data.z())
                    );
                    let adjusted_gyro = Vector3::new(
                        gyro_raw_to_rads(gyro_data.x()),
                        gyro_raw_to_rads(gyro_data.y()),
                        gyro_raw_to_rads(gyro_data.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 calibration_data = sensor.get_accel_calibration().await;
        info!("Calibration data: {calibration_data:?}");
        if unsafe { MPU_WAS_CALIBRATED } == u8::MAX {
            warn!("Re-using MPU calibration data from last boot");
            //sensor.set_accel_calibration(&accel).await?;
            //sensor.set_gyro_calibration(&gyro).await?;
        } 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?;
            backoff.attempt(async || { sensor.set_digital_lowpass_filter(mpu6050_dmp::config::DigitalLowPassFilter::Filter1).await }).await?;
            info!("DMP initialized in {}ms", start.as_millis());
            warn!("Calibrating MPU from scratch");
            let start = Instant::now();
            let _calibration_data = backoff.attempt(async || {
                // TODO: Store calibration data with set_accel_calibration, set_gyro_calibration
                let mut actions = CalibrationActions::all();
                let mut accel;
                let mut gyro;
                let params = CalibrationParameters::new(
                    mpu6050_dmp::accel::AccelFullScale::G2,
                    mpu6050_dmp::gyro::GyroFullScale::Deg2000,
                    mpu6050_dmp::calibration::ReferenceGravity::Zero
                );
                loop {
                    (actions, accel, gyro) = sensor.calibration_loop(&mut delay, &params, actions).await?;
                    if actions.is_empty() {
                        return Ok((accel, gyro))
                    }
                }
                /*sensor.calibrate(&mut delay, &CalibrationParameters::new(
                    mpu6050_dmp::accel::AccelFullScale::G2,
                    mpu6050_dmp::gyro::GyroFullScale::Deg2000,
                    mpu6050_dmp::calibration::ReferenceGravity::Zero
                )).await*/
            }).await?;
            //unsafe { MPU_CALIBRATION = Some(calibration_data) };
            unsafe { MPU_WAS_CALIBRATED = u8::MAX };
            info!("MPU calibrated in {}ms", start.as_millis());
        }
        Ok(())
    }
}