const SENSOR_ROWS: usize = 12;
const SENSOR_COLS: usize = 7;
const SENSOR_COUNT: usize = SENSOR_ROWS * SENSOR_COLS;

const TOTAL_PRESSURE_LOW_THRESHOLD: f32 = 500.0;
const COP_STABILITY_FRAMES_REQUIRED: usize = 15;

const POST_INIT_WINDOW_CNT: usize = 50;
const POST_INIT_STABLE_CNT: usize = 50;
const POST_INIT_STABLE_THRESH: f32 = 0.1;

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

pub struct PztProcessor {
    first_frame: Option<Vec<f32>>,

    first_contact_cop_x: Option<f32>,
    first_contact_cop_y: Option<f32>,
    contact_initialized: bool,

    total_pressure_low_counter: usize,

    cop_init_x_buf: Vec<f32>,
    cop_init_y_buf: Vec<f32>,

    post_init_frame_cnt: usize,
    post_stable_cnt: usize,
    post_refined_flag: bool,
    post_cand_x: Option<f32>,
    post_cand_y: Option<f32>,
}

impl PztProcessor {
    pub fn new() -> Self {
        Self {
            first_frame: None,
            first_contact_cop_x: None,
            first_contact_cop_y: None,
            contact_initialized: false,
            total_pressure_low_counter: 0,
            cop_init_x_buf: Vec::with_capacity(COP_STABILITY_FRAMES_REQUIRED),
            cop_init_y_buf: Vec::with_capacity(COP_STABILITY_FRAMES_REQUIRED),
            post_init_frame_cnt: 0,
            post_stable_cnt: 0,
            post_refined_flag: false,
            post_cand_x: None,
            post_cand_y: None,
        }
    }

    fn subtract_baseline(&mut self, current_frame: &[f32]) -> Vec<f32> {
        if self.first_frame.is_none() {
            self.first_frame = Some(current_frame.to_vec());
        }

        let baseline = self.first_frame.as_ref().unwrap();
        current_frame
            .iter()
            .zip(baseline.iter())
            .map(|(current, baseline)| (current - baseline).max(0.0))
            .collect()
    }

    fn reset_cop_state(&mut self) {
        self.first_contact_cop_x = None;
        self.first_contact_cop_y = None;
        self.contact_initialized = false;
        self.total_pressure_low_counter = 0;
        self.cop_init_x_buf.clear();
        self.cop_init_y_buf.clear();
        self.post_init_frame_cnt = 0;
        self.post_stable_cnt = 0;
        self.post_refined_flag = false;
        self.post_cand_x = None;
        self.post_cand_y = None;
    }

    fn compute_median(sorted: &[f32]) -> f32 {
        let n = sorted.len();
        if n == 0 {
            return 0.0;
        }

        if n % 2 == 0 {
            (sorted[n / 2 - 1] + sorted[n / 2]) / 2.0
        } else {
            sorted[n / 2]
        }
    }

    fn compute_pressure_direction(&mut self, frame: &[f32]) -> (f32, f32) {
        let total_pressure = frame.iter().sum::<f32>();
        if total_pressure < TOTAL_PRESSURE_LOW_THRESHOLD {
            self.total_pressure_low_counter += 1;
        } else {
            self.total_pressure_low_counter = 0;
        }

        if self.total_pressure_low_counter >= COP_STABILITY_FRAMES_REQUIRED {
            self.reset_cop_state();
            return (0.0, 0.0);
        }

        if total_pressure == 0.0 {
            return (0.0, 0.0);
        }

        let mut sum_x = 0.0;
        let mut sum_y = 0.0;
        for row in 0..SENSOR_ROWS {
            for col in 0..SENSOR_COLS {
                let value = frame[row * SENSOR_COLS + col];
                sum_x += value * col as f32;
                sum_y += value * row as f32;
            }
        }
        let cop_x = sum_x / total_pressure;
        let cop_y = sum_y / total_pressure;

        if !self.contact_initialized {
            self.cop_init_x_buf.push(cop_x);
            self.cop_init_y_buf.push(cop_y);

            if self.cop_init_x_buf.len() >= COP_STABILITY_FRAMES_REQUIRED {
                let mut xs = self.cop_init_x_buf.clone();
                xs.sort_by(|a, b| a.partial_cmp(b).unwrap());
                let mut ys = self.cop_init_y_buf.clone();
                ys.sort_by(|a, b| a.partial_cmp(b).unwrap());

                self.first_contact_cop_x = Some(Self::compute_median(&xs));
                self.first_contact_cop_y = Some(Self::compute_median(&ys));
                self.contact_initialized = true;
                self.cop_init_x_buf.clear();
                self.cop_init_y_buf.clear();
            }

            return (0.0, 0.0);
        }

        self.post_init_frame_cnt += 1;
        if !self.post_refined_flag && self.post_init_frame_cnt <= POST_INIT_WINDOW_CNT {
            if let (Some(cx), Some(cy)) = (self.post_cand_x, self.post_cand_y) {
                let dist = ((cop_x - cx).powi(2) + (cop_y - cy).powi(2)).sqrt();
                if dist <= POST_INIT_STABLE_THRESH {
                    self.post_stable_cnt += 1;
                } else {
                    self.post_cand_x = Some(cop_x);
                    self.post_cand_y = Some(cop_y);
                    self.post_stable_cnt = 1;
                }
            } else {
                self.post_cand_x = Some(cop_x);
                self.post_cand_y = Some(cop_y);
                self.post_stable_cnt = 1;
            }

            if self.post_stable_cnt >= POST_INIT_STABLE_CNT {
                self.first_contact_cop_x = self.post_cand_x;
                self.first_contact_cop_y = self.post_cand_y;
                self.post_refined_flag = true;
            }
        } else {
            self.post_refined_flag = true;
        }

        let base_x = self.first_contact_cop_x.unwrap_or(cop_x);
        let base_y = self.first_contact_cop_y.unwrap_or(cop_y);
        let delta_x = cop_x - base_x;
        let delta_y = base_y - cop_y;
        (delta_x, delta_y)
    }

    fn compute_vector_angle(x: f32, y: f32) -> (f32, f32) {
        let epsilon = 1e-8f32;
        let magnitude = (x * x + y * y).sqrt();
        let mut angle = y.atan2(x + epsilon).to_degrees();
        if angle < 0.0 {
            angle += 360.0;
        }
        (angle, magnitude)
    }

    pub fn get_pzt_analysis(
        &mut self,
        adc_data: &[f32],
    ) -> Result<PztSpatialAnalysis, &'static str> {
        if adc_data.len() != SENSOR_COUNT {
            return Err("ADC data length must be 84");
        }

        let baseline = self.subtract_baseline(adc_data);
        let (dx, dy) = self.compute_pressure_direction(&baseline);

        let planar_x = dx;
        let planar_y = -dy;
        let (angle_deg, magnitude) = Self::compute_vector_angle(planar_x, planar_y);

        Ok(PztSpatialAnalysis {
            angle_deg,
            magnitude,
            planar_x,
            planar_y,
        })
    }

    pub fn get_pzt_angle(&mut self, adc_data: &[f32]) -> Result<f32, &'static str> {
        Ok(self.get_pzt_analysis(adc_data)?.angle_deg)
    }

    pub fn reset_baseline(&mut self) {
        self.first_frame = None;
        self.reset_cop_state();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn idle_frame_returns_zero() {
        let mut processor = PztProcessor::new();
        let frame = [0.0f32; SENSOR_COUNT];
        let analysis = processor.get_pzt_analysis(&frame).unwrap();
        assert_eq!(analysis.magnitude, 0.0);
        assert_eq!(analysis.angle_deg, 0.0);
    }
}
