use num_traits::{Float, FromPrimitive, One, Zero, ToPrimitive};

use crate::smooth::SmoothError;

pub fn median_filter<T>(
    data: &[T],
    period: usize,
    precision: f64,
) -> Result<Vec<T>, SmoothError> where 
    T: Float + FromPrimitive + Zero + Clone + One,
{
    let len = data.len();
    if len == 0 || period == 0 || precision <= 0.0 {
        return Err(SmoothError::MedianApplyError);
    }

    let half = period / 2;
    let mut result = Vec::with_capacity(len);

    // 统计范围
    let precision_t = T::from(precision).unwrap();
    let min_val = data.iter().cloned().fold(T::infinity(), T::min);
    let max_val = data.iter().cloned().fold(T::neg_infinity(), T::max);
    let bin_count = ((max_val - min_val) / precision_t).ceil().to_usize().unwrap() + 1;

    let to_bin = |v: T| -> usize {
        let idx = ((v - min_val) / precision_t).floor().to_usize().unwrap();
        idx.clamp(0, (bin_count.to_isize().unwrap() - 1).try_into().unwrap()).to_usize().unwrap()
    };

    let get_value = |idx: isize| -> T {
        if idx < 0 {
            data[0] // 延拓第一个值
        } else if (idx as usize) >= len {
            data[len - 1] // 延拓最后一个值
        } else {
            data[idx as usize]
        }
    };

    for i in 0..len {
        let mut hist = vec![0u32; bin_count];

        // 生成窗口并统计直方图
        for j in -(half as isize)..=(half as isize) {
            let v = get_value(i as isize + j);
            let bin = to_bin(v);
            hist[bin] += 1;
        }

        // 找中位数
        let mut count = 0;
        let median_target = (period + 1) / 2;
        let mut median_bin = 0;
        for (i, &h) in hist.iter().enumerate() {
            count += h;
            if count >= median_target as u32 {
                median_bin = i;
                break;
            }
        }

        let median = min_val + (T::from(median_bin).unwrap() + T::from(0.5).unwrap()) * precision_t;
        result.push(median);
    }

    Ok(result)
}