
use nalgebra::{Rotation3, Unit, Vector3, ComplexField};

use crate::CircularBuffer;

type Sample = Vector3<f32>; // [ax, ay, az] in m/s^2

#[derive(Debug, Clone, Copy, Default)]
pub struct SensorAlignment {
    /// 3x3 transform matrix: raw -> body
    /// includes permutation, sign, and bias subtraction.
    pub transform: Rotation3<f32>,
    /// bias vector in sensor frame
    pub bias: Vector3<f32>
}

#[derive(Debug, Clone, Copy)]
pub enum Direction {
    UpZ,
    DnZ,
    UpX,
    DnX,
    UpY,
    DnY
}

impl From<Direction> for Unit<Vector3<f32>> {
    fn from(value: Direction) -> Self {
        Unit::new_normalize(match value {
            Direction::UpZ => Vector3::z(),
            Direction::DnZ => -Vector3::z(),
            Direction::UpX => Vector3::x(),
            Direction::DnX => -Vector3::x(),
            Direction::UpY => Vector3::y(),
            Direction::DnY => -Vector3::y(),
        })
    }
}

impl SensorAlignment {
    pub fn apply(&self, raw: &Sample) -> Vector3<f32> {
        self.transform * (raw - self.bias)
    }

    pub fn from_samples(accel_samples: &[Sample], down_direction: Direction) -> SensorAlignment {
        let bias: Vector3<f32> = accel_samples.iter().sum::<Vector3<f32>>() / (accel_samples.len() as f32);
        // Figure out the unit vector pointing 'down' at rest
        let down_body = Unit::new_normalize(bias);
        // Create our 'true' down vector
        let down_world: Unit<Vector3<f32>> = down_direction.into();//Unit::new_normalize(Vector3::new(0.0, 0.0, -1.0));
        // Determine how the axes are rotated between the two
        let axis = down_body.cross(&down_world);
        let axis_unit = Unit::new_normalize(axis);
        let angle = down_body.dot(&down_world).clamp(-1.0, 1.0).acos();
        // Create a rotation matrix from the angles
        let transform = Rotation3::from_axis_angle(&axis_unit, angle);

        Self {
            transform,
            bias,
        }
    }
}


#[derive(Debug)]
pub struct DirectionEstimator {
    buf: CircularBuffer<Vector3<f32>>,
    direction: Direction,
    pub alignment: SensorAlignment
}

impl DirectionEstimator {
    pub fn new(direction: Direction) -> Self {
        Self {
            buf: Default::default(),
            alignment: Default::default(),
            direction,
        }
    }
}

impl DirectionEstimator {
    pub fn update(&mut self, accel: Vector3<f32>) {
        self.buf.insert(accel);

        if self.buf.next_index == 0 {
            self.alignment = SensorAlignment::from_samples(&self.buf.data, self.direction);
        }
    }

    pub fn apply(&self, vec: &Vector3<f32>) -> Vector3<f32> {
        self.alignment.apply(vec)
    }
}