use core::cell::RefCell;

use embassy_embedded_hal::shared_bus::asynch::i2c::I2cDevice;
use embassy_sync::{blocking_mutex::raw::{CriticalSectionRawMutex, NoopRawMutex}, channel::{DynamicSender, Sender}};
use embassy_time::{Delay, Timer};
use esp_hal::{i2c::master::I2c, Async};
use log::{info, warn};
use mpu6050_dmp::{address::Address, calibration::CalibrationParameters, error_async::Error, sensor_async::Mpu6050};
use nalgebra::Vector3;
use core::f32::consts::PI;
use nalgebra::ComplexField;

use crate::{backoff::Backoff, ego::{alignment::SensorAlignment, tilt::TiltEstimator}, events::{Measurement, Notification}};

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

/// Rotate body accel into ENU and remove gravity; returns (east, north, up) linear accel in m/s^2
pub fn accel_body_to_enu_lin(body_acc: (f32,f32,f32), roll: f32, pitch: f32, yaw: f32) -> (f32,f32,f32) {
    let (ax, ay, az) = body_acc;

    let (sr, cr) = roll.sin_cos();
    let (sp, cp) = pitch.sin_cos();
    let (sy, cy) = yaw.sin_cos();

    // Build R = Rz(yaw) * Ry(pitch) * Rx(roll)
    // multiply R * a_body
    let r11 =  cy*cp;
    let r12 =  cy*sp*sr - sy*cr;
    let r13 =  cy*sp*cr + sy*sr;

    let r21 =  sy*cp;
    let r22 =  sy*sp*sr + cy*cr;
    let r23 =  sy*sp*cr - cy*sr;

    let r31 = -sp;
    let r32 =  cp*sr;
    let r33 =  cp*cr;

    // a_nav = R * a_body
    let ax_nav = r11*ax + r12*ay + r13*az; // east
    let ay_nav = r21*ax + r22*ay + r23*az; // north
    let az_nav = r31*ax + r32*ay + r33*az; // up

    // remove gravity (nav frame gravity points down in NED; in our ENU 'up' axis positive up)
    // If accel measured gravity as +9.8 on sensor when static and z down convention used above,
    // then az_nav will include +g in 'up' direction; subtract G.
    let ax_lin = ax_nav;
    let ay_lin = ay_nav;
    let az_lin = az_nav - G; // linear acceleration in up axis

    (ax_lin, ay_lin, az_lin)
}

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

    backoff.forever().attempt::<_, (), ()>(async || {
        let mut sensor = backoff.forever().attempt(async || {
            info!("Initializing 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) => {
                    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;

        let alignment = {
            info!("Locating earth...");
            let mut accel_samples = [Vector3::default(); 100];
            let mut gyro_samples = [Vector3::default(); 100];
            for (idx, (accel_sample, gyro_sample)) in accel_samples.iter_mut().zip(gyro_samples.iter_mut()).enumerate() {
                let (accel, gyro) = backoff.attempt(async || { sensor_ref.motion6().await }).await.unwrap();
                *accel_sample = Vector3::new(accel_raw_to_ms2(accel.x()), accel_raw_to_ms2(accel.y()), accel_raw_to_ms2(accel.z()));
                *gyro_sample = Vector3::new(gyro_raw_to_rads(gyro.x()), gyro_raw_to_rads(gyro.y()), gyro_raw_to_rads(gyro.z()));
                info!("{idx}/200 {accel_sample:?} {gyro_sample:?}");
                Timer::after_millis(5).await;
            }

            SensorAlignment::new(accel_samples.as_slice(), gyro_samples.as_slice())
        };

        info!("MPU is ready! earth={alignment:?}");
        sink.send(Notification::CalibrationComplete).await;
        let mut tilt = TiltEstimator::new(0.3);
        //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 = alignment.apply(&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 = alignment.apply_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;

                    // Apply current rotation and accel values to rotate our values back into the ENU frame
                    tilt.update(0.01, (filtered.x, filtered.y, filtered.z), (adjusted_gyro.x, adjusted_gyro.y, adjusted_gyro.z));
                    let (roll, pitch, yaw) = tilt.angles();
                    let (ax_e, ay_n, _az_up) = accel_body_to_enu_lin((filtered.x,filtered.y,filtered.z), roll, pitch, yaw);

                    let measured = Vector3::new(ax_e, ay_n, adjusted_gyro.z);
                    events.send(
                        Measurement::IMU(
                            measured
                        )
                    ).await;
                    Timer::after_millis(10).await;
                },
                Err(e) => {
                    warn!("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, CriticalSectionRawMutex, I2c<'static, Async>>>) -> Result<(), Error<I2cDevice<'static, CriticalSectionRawMutex, I2c<'static, Async>>>> {
    let mut delay = Delay;
    let backoff = Backoff::from_millis(10);
    info!("Initializing MPU");
    backoff.attempt(async || { sensor.initialize_dmp(&mut delay).await }).await?;
    backoff.attempt(async || { sensor.set_digital_lowpass_filter(mpu6050_dmp::config::DigitalLowPassFilter::Filter3).await }).await?;
    info!("Calibrating MPU");
    backoff.attempt(async || {
        // TODO: Store calibration data with set_accel_calibration, set_gyro_calibration
        sensor.calibrate(&mut delay, &CalibrationParameters::new(
            mpu6050_dmp::accel::AccelFullScale::G2,
            mpu6050_dmp::gyro::GyroFullScale::Deg2000,
            mpu6050_dmp::calibration::ReferenceGravity::Zero
        )).await.map(|_| { Ok(()) })
    }).await?
}