use crate::serial_core::multi_dim_force::PztProcessor;

pub const FINGER_SAMPLE_COUNT: usize = 84;
pub const HAND_FINGERTIP_COUNT: usize = 5;
const MIN_TANGENTIAL_MAGNITUDE: f32 = 0.1;

#[derive(Debug, Clone, Copy)]
pub struct HudSpatialForce {
    pub angle_deg: f32,
    pub magnitude: f32,
}

pub struct ForceEstimatorState {
    pzt_processor: PztProcessor,
    fingertip_processors: [PztProcessor; HAND_FINGERTIP_COUNT],
}

impl ForceEstimatorState {
    pub fn new() -> Self {
        Self {
            pzt_processor: PztProcessor::new(),
            fingertip_processors: std::array::from_fn(|_| PztProcessor::new()),
        }
    }

    pub fn reset(&mut self) {
        self.pzt_processor.reset_baseline();
        self.fingertip_processors
            .iter_mut()
            .for_each(PztProcessor::reset_baseline);
    }

    pub fn analyze(&mut self, values: &[u32]) -> Option<HudSpatialForce> {
        if values.len() != FINGER_SAMPLE_COUNT {
            return None;
        }

        let pzt_values = values.iter().map(|value| *value as f32).collect::<Vec<_>>();

        self.pzt_processor
            .get_pzt_analysis(&pzt_values)
            .ok()
            .filter(|analysis| analysis.magnitude > MIN_TANGENTIAL_MAGNITUDE)
            .map(|analysis| HudSpatialForce {
                angle_deg: analysis.angle_deg,
                magnitude: analysis.magnitude,
            })
    }

    pub fn analyze_fingertips(
        &mut self,
        values: &[u32],
    ) -> [Option<HudSpatialForce>; HAND_FINGERTIP_COUNT] {
        self.analyze_fingertips_masked(values, &[true; HAND_FINGERTIP_COUNT])
    }

    pub fn analyze_fingertips_masked(
        &mut self,
        values: &[u32],
        enabled: &[bool; HAND_FINGERTIP_COUNT],
    ) -> [Option<HudSpatialForce>; HAND_FINGERTIP_COUNT] {
        let mut forces = [None; HAND_FINGERTIP_COUNT];

        for (tip_index, processor) in self.fingertip_processors.iter_mut().enumerate() {
            if !enabled[tip_index] {
                continue;
            }

            let start = tip_index * FINGER_SAMPLE_COUNT;
            let end = start + FINGER_SAMPLE_COUNT;
            let Some(segment) = values.get(start..end) else {
                break;
            };

            let pzt_values = segment
                .iter()
                .map(|value| *value as f32)
                .collect::<Vec<_>>();

            forces[tip_index] = processor
                .get_pzt_analysis(&pzt_values)
                .ok()
                .filter(|analysis| analysis.magnitude > MIN_TANGENTIAL_MAGNITUDE)
                .map(|analysis| HudSpatialForce {
                    angle_deg: analysis.angle_deg,
                    magnitude: analysis.magnitude,
                });
        }

        forces
    }
}

impl Default for ForceEstimatorState {
    fn default() -> Self {
        Self::new()
    }
}
