wedb_embed 0.1.0

pub mod conf;
pub mod meta;

pub use conf::{
    DEFAULT_COMPRESSION, MAX_CAPACITY, MAX_COMPRESSION, MIN_CAPACITY, MIN_COMPRESSION,
    TDigestCreate, TDigestInfo, TDigestMerge, calculate_capacity,
};
pub use meta::{TDigestMeta, decode_double_from_u64, encode_double_to_u64};
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::f64::consts::PI;
use std::mem::take;
use std::str;

use crate::db::WeDb;
use crate::error::{Error, Result};
use crate::key_composer::KeyComposer;

/// 双精度浮点数比较精度阈值（对标 Apache Kvrocks DoubleCompare）
pub const REL_EPS: f64 = 1e-12;
pub const ABS_EPS: f64 = 1e-9;
pub const SINGLETON_BOUNDARY_WEIGHT: f64 = 1.0;
pub const HALF_SINGLETON_BOUNDARY_WEIGHT: f64 = 0.5;
pub const INV_TWO_PI: f64 = 1.0 / (2.0 * PI);

/// 浮点数安全比较（容忍微小精度误差，对标 Apache Kvrocks DoubleCompare）
#[inline]
pub fn double_compare(a: f64, b: f64, rel_eps: f64, abs_eps: f64) -> Ordering {
    if a.is_nan() || b.is_nan() {
        return a.total_cmp(&b);
    }
    let diff = a - b;
    let adiff = diff.abs();
    if adiff <= abs_eps {
        return Ordering::Equal;
    }
    let maxab = a.abs().max(b.abs());
    if adiff <= maxab * rel_eps {
        return Ordering::Equal;
    }
    if diff < 0.0 {
        Ordering::Less
    } else {
        Ordering::Greater
    }
}

/// 浮点数相等判断（对标 Apache Kvrocks DoubleEqual）
#[inline]
pub fn double_equal(a: f64, b: f64) -> bool {
    double_compare(a, b, REL_EPS, ABS_EPS).is_eq()
}

/// 线性插值（对标 Apache Kvrocks Lerp）
#[inline]
pub const fn lerp(a: f64, b: f64, t: f64) -> f64 {
    a + t * (b - a)
}

/// 单个质心节点 (Centroid，对标 Apache Kvrocks Centroid)
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct Centroid {
    pub mean: f64,
    pub weight: f64,
}

impl Centroid {
    #[inline]
    pub const fn new(mean: f64, weight: f64) -> Self {
        Self { mean, weight }
    }

    /// 与另一个质心就地加权合并（对标 Apache Kvrocks Centroid::Merge）
    #[inline]
    pub fn merge(&mut self, other: &Centroid) {
        self.weight += other.weight;
        self.mean += (other.mean - self.mean) * other.weight / self.weight;
    }

    #[inline]
    pub fn add(&mut self, mean: f64, weight: f64) {
        self.merge(&Centroid::new(mean, weight));
    }
}

/// 带压缩因子的质心集合（对标 Apache Kvrocks CentroidsWithDelta）
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CentroidsWithDelta {
    pub centroids: Vec<Centroid>,
    pub delta: u32,
    pub min: f64,
    pub max: f64,
    pub total_weight: f64,
}

/// K_1 比例缩放函数 (ScalerK1，对标 Apache Kvrocks / Apache Arrow ScalerK1)
#[derive(Debug, Clone, Copy)]
pub struct ScalerK1 {
    pub delta_norm: f64,
    pub inv_delta_norm: f64,
}

impl ScalerK1 {
    #[inline]
    pub fn new(delta: u32) -> Self {
        let norm = delta as f64 * INV_TWO_PI;
        Self {
            delta_norm: norm,
            inv_delta_norm: if norm != 0.0 { 1.0 / norm } else { 0.0 },
        }
    }

    #[inline]
    pub fn k(&self, q: f64) -> f64 {
        let q_clamped = q.clamp(0.0, 1.0);
        self.delta_norm * (2.0 * q_clamped - 1.0).asin()
    }

    #[inline]
    pub fn q(&self, k: f64) -> f64 {
        ((k * self.inv_delta_norm).sin() + 1.0) * 0.5
    }
}

/// T-Digest 质心流合并器（对标 Apache Kvrocks TDigestMerger）
pub struct TDigestMerger {
    scaler: ScalerK1,
    total_weight: f64,
    weight_so_far: f64,
    weight_limit: f64,
}

impl TDigestMerger {
    #[inline]
    pub fn new(delta: u32) -> Self {
        Self {
            scaler: ScalerK1::new(delta),
            total_weight: 0.0,
            weight_so_far: 0.0,
            weight_limit: -1.0,
        }
    }

    #[inline]
    pub fn reset(&mut self, total_weight: f64) {
        self.total_weight = total_weight;
        self.weight_so_far = 0.0;
        self.weight_limit = -1.0;
    }

    /// 向合并器添加单个质心
    #[inline]
    pub fn add(&mut self, output: &mut Vec<Centroid>, centroid: Centroid) {
        let weight = self.weight_so_far + centroid.weight;
        if weight <= self.weight_limit {
            if let Some(last) = output.last_mut() {
                last.merge(&centroid);
            } else {
                output.push(centroid);
            }
        } else {
            let quantile = if self.total_weight > 0.0 {
                self.weight_so_far / self.total_weight
            } else {
                0.0
            };
            let next_k = self.scaler.k(quantile) + 1.0;
            let next_weight_limit = self.total_weight * self.scaler.q(next_k);
            if next_weight_limit <= self.weight_limit {
                self.weight_limit = self.total_weight;
            } else {
                self.weight_limit = next_weight_limit;
            }
            output.push(centroid);
        }
        self.weight_so_far = weight;
    }

    /// 校验质心序列的 K-Size 约束（对标 Apache Kvrocks TDigestMerger::Validate）
    pub fn validate(&self, tdigest: &[Centroid], total_weight: f64) -> Result<()> {
        let mut q_prev = 0.0;
        let mut k_prev = self.scaler.k(0.0);
        for i in tdigest {
            let q = q_prev + i.weight / total_weight;
            let k = self.scaler.k(q);
            if i.weight != 1.0 && (k - k_prev) > 1.001 {
                return Err(Error::invalid_data(format!(
                    "oversized centroid: {}",
                    k - k_prev
                )));
            }
            k_prev = k;
            q_prev = q;
        }
        Ok(())
    }
}

/// 合并多个质心集合（对标 Apache Kvrocks TDigestMerge）
pub fn tdigest_merge_centroids_list(
    centroids_list: &[CentroidsWithDelta],
    delta: u32,
) -> CentroidsWithDelta {
    tdigest_merge_buffer_and_centroids(&[], centroids_list, delta)
}

/// 合并原始缓冲区与质心集合（对标 Apache Kvrocks TDigestMerge(buffer, centroids)）
pub fn tdigest_merge_buffer_and_centroids(
    buffer: &[f64],
    centroids_list: &[CentroidsWithDelta],
    delta: u32,
) -> CentroidsWithDelta {
    let total_len = buffer.len()
        + centroids_list
            .iter()
            .map(|l| l.centroids.len())
            .sum::<usize>();
    let mut all_centroids = Vec::with_capacity(total_len);
    let mut total_w = 0.0;
    let mut min_val = f64::MAX;
    let mut max_val = -f64::MAX;

    for &v in buffer {
        if !v.is_nan() {
            all_centroids.push(Centroid::new(v, 1.0));
            total_w += 1.0;
            if v < min_val {
                min_val = v;
            }
            if v > max_val {
                max_val = v;
            }
        }
    }

    for list in centroids_list {
        if list.centroids.is_empty() {
            continue;
        }
        total_w += list.total_weight;
        if list.min < min_val {
            min_val = list.min;
        }
        if list.max > max_val {
            max_val = list.max;
        }
        all_centroids.extend_from_slice(&list.centroids);
    }

    if all_centroids.is_empty() {
        return CentroidsWithDelta {
            centroids: Vec::new(),
            delta,
            min: f64::MAX,
            max: -f64::MAX,
            total_weight: 0.0,
        };
    }

    all_centroids.sort_unstable_by(|a, b| double_compare(a.mean, b.mean, REL_EPS, ABS_EPS));

    let effective_delta = delta.max(10);
    let mut merger = TDigestMerger::new(effective_delta);
    merger.reset(total_w);

    let mut output = Vec::with_capacity(effective_delta as usize);
    for c in all_centroids {
        merger.add(&mut output, c);
    }

    CentroidsWithDelta {
        centroids: output,
        delta: effective_delta,
        min: min_val,
        max: max_val,
        total_weight: total_w,
    }
}

/// 分位数估算算法（对标 Apache Kvrocks TDigestQuantile）
pub fn tdigest_quantile_calc(
    centroids: &[Centroid],
    min: f64,
    max: f64,
    total_weight: f64,
    q: f64,
) -> f64 {
    if q.is_nan() || !(0.0..=1.0).contains(&q) || centroids.is_empty() || total_weight <= 0.0 {
        return f64::NAN;
    }

    let index = q * total_weight;
    if index <= 1.0 {
        return min;
    } else if index >= total_weight - 1.0 {
        return max;
    }

    let mut weight_sum = 0.0;
    let mut ci = 0;
    while ci < centroids.len() {
        weight_sum += centroids[ci].weight;
        if index <= weight_sum {
            break;
        }
        ci += 1;
    }
    if ci >= centroids.len() {
        ci = centroids.len() - 1;
    }

    let centroid = centroids[ci];
    let mut diff = index + centroid.weight * 0.5 - weight_sum;

    if centroid.weight == 1.0 && diff.abs() < 0.5 {
        return centroid.mean;
    }

    let mut ci_left = ci;
    let mut ci_right = ci;
    if diff > 0.0 {
        if ci_right == centroids.len() - 1 {
            let c = centroids[ci_right];
            return lerp(c.mean, max, diff / (c.weight * 0.5));
        }
        ci_right += 1;
    } else {
        if ci_left == 0 {
            let c = centroids[0];
            return lerp(min, c.mean, index / (c.weight * 0.5));
        }
        ci_left -= 1;
        diff += centroids[ci_left].weight * 0.5 + centroids[ci_right].weight * 0.5;
    }

    let lc = centroids[ci_left];
    let rc = centroids[ci_right];
    diff /= lc.weight * 0.5 + rc.weight * 0.5;
    lerp(lc.mean, rc.mean, diff)
}

/// 累积分布函数 CDF 计算（对标 Apache Kvrocks TDigestCDF 与 RedisBloom）
pub fn tdigest_cdf_calc(
    centroids: &[Centroid],
    centroids_min: f64,
    centroids_max: f64,
    total_weight: f64,
    inputs: &[f64],
) -> Vec<f64> {
    if centroids.is_empty() || total_weight <= 0.0 {
        return vec![f64::NAN; inputs.len()];
    }

    let mut indexed: Vec<(f64, usize)> = inputs
        .iter()
        .copied()
        .enumerate()
        .map(|(i, v)| (v, i))
        .collect();

    indexed.sort_unstable_by(|(a, _), (b, _)| double_compare(*a, *b, REL_EPS, ABS_EPS));

    let mut result = vec![f64::NAN; inputs.len()];
    if indexed.is_empty() {
        return result;
    }

    let mut i = 0;
    let n_inputs = indexed.len();

    if centroids.len() == 1 {
        let width = centroids_max - centroids_min;
        while i < n_inputs {
            let val = indexed[i].0;
            if val.is_nan() {
                break;
            }
            let start = i;
            while i < n_inputs && double_equal(indexed[i].0, val) {
                i += 1;
            }
            let weight = if val < centroids_min {
                0.0
            } else if val > centroids_max {
                total_weight
            } else if val - centroids_min <= width {
                total_weight * 0.5
            } else {
                (val - centroids_min) / width * total_weight
            };
            let prob = (weight / total_weight).clamp(0.0, 1.0);
            for item in &indexed[start..i] {
                result[item.1] = prob;
            }
        }
        return result;
    }

    // 1. 小于 centroids_min
    while i < n_inputs && indexed[i].0 < centroids_min {
        let val = indexed[i].0;
        if val.is_nan() {
            break;
        }
        let start = i;
        while i < n_inputs && double_equal(indexed[i].0, val) {
            i += 1;
        }
        for item in &indexed[start..i] {
            result[item.1] = 0.0;
        }
    }

    // 2. 在 [centroids_min, centroids[0].mean) 区间
    let first_mean = centroids[0].mean;
    let first_weight = centroids[0].weight;
    let first_width = first_mean - centroids_min;

    while i < n_inputs && indexed[i].0 < first_mean {
        let val = indexed[i].0;
        if val.is_nan() {
            break;
        }
        let start = i;
        while i < n_inputs && double_equal(indexed[i].0, val) {
            i += 1;
        }
        let weight = if first_width > 0.0 {
            if double_equal(val, centroids_min) {
                HALF_SINGLETON_BOUNDARY_WEIGHT
            } else {
                lerp(
                    HALF_SINGLETON_BOUNDARY_WEIGHT,
                    first_weight * 0.5,
                    (val - centroids_min) / first_width,
                )
            }
        } else {
            0.0
        };
        let prob = (weight / total_weight).clamp(0.0, 1.0);
        for item in &indexed[start..i] {
            result[item.1] = prob;
        }
    }

    // 3. 在中间质心区间 [centroids[0].mean, centroids[last].mean]
    let mut c_idx = 0;
    let mut weight_so_far = 0.0;
    let n_centroids = centroids.len();

    while c_idx < n_centroids - 1 && i < n_inputs {
        let val = indexed[i].0;
        if val.is_nan() {
            break;
        }
        let current_c = centroids[c_idx];
        let next_c = centroids[c_idx + 1];

        if double_equal(val, current_c.mean) {
            let start = i;
            while i < n_inputs && double_equal(indexed[i].0, val) {
                i += 1;
            }
            let mut dw = 0.0;
            let mut same_idx = c_idx;
            while same_idx < n_centroids && double_equal(centroids[same_idx].mean, current_c.mean) {
                dw += centroids[same_idx].weight;
                same_idx += 1;
            }
            let weight = weight_so_far + dw * 0.5;
            let prob = (weight / total_weight).clamp(0.0, 1.0);
            for item in &indexed[start..i] {
                result[item.1] = prob;
            }
            continue;
        }

        if current_c.mean < val && val < next_c.mean {
            let start = i;
            while i < n_inputs && double_equal(indexed[i].0, val) {
                i += 1;
            }
            let mean_diff = next_c.mean - current_c.mean;
            let weight = if mean_diff > 0.0 {
                let mut left_exclude = 0.0;
                let mut right_exclude = 0.0;
                if current_c.weight == SINGLETON_BOUNDARY_WEIGHT {
                    if next_c.weight == SINGLETON_BOUNDARY_WEIGHT {
                        weight_so_far + SINGLETON_BOUNDARY_WEIGHT
                    } else {
                        left_exclude = HALF_SINGLETON_BOUNDARY_WEIGHT;
                        let dw = (current_c.weight + next_c.weight) * 0.5;
                        let dw_no_singleton = dw - left_exclude - right_exclude;
                        let base = weight_so_far + current_c.weight * 0.5 + left_exclude;
                        lerp(
                            base,
                            base + dw_no_singleton,
                            (val - current_c.mean) / mean_diff,
                        )
                    }
                } else {
                    if next_c.weight == SINGLETON_BOUNDARY_WEIGHT {
                        right_exclude = HALF_SINGLETON_BOUNDARY_WEIGHT;
                    }
                    let dw = (current_c.weight + next_c.weight) * 0.5;
                    let dw_no_singleton = dw - left_exclude - right_exclude;
                    let base = weight_so_far + current_c.weight * 0.5 + left_exclude;
                    lerp(
                        base,
                        base + dw_no_singleton,
                        (val - current_c.mean) / mean_diff,
                    )
                }
            } else {
                weight_so_far + current_c.weight * 0.5
            };
            let prob = (weight / total_weight).clamp(0.0, 1.0);
            for item in &indexed[start..i] {
                result[item.1] = prob;
            }
            continue;
        }

        if val >= next_c.mean {
            weight_so_far += current_c.weight;
            c_idx += 1;
        }
    }

    // 4. 在 (centroids[last].mean, centroids_max] 区间
    let last_c = centroids[n_centroids - 1];
    let last_width = centroids_max - last_c.mean;

    while i < n_inputs && indexed[i].0 <= centroids_max {
        let val = indexed[i].0;
        if val.is_nan() {
            break;
        }
        let start = i;
        while i < n_inputs && double_equal(indexed[i].0, val) {
            i += 1;
        }
        let weight = if double_equal(val, last_c.mean) {
            total_weight - last_c.weight * 0.5
        } else if val > last_c.mean {
            if last_width > 0.0 {
                if double_equal(val, centroids_max) {
                    total_weight - HALF_SINGLETON_BOUNDARY_WEIGHT
                } else {
                    lerp(
                        total_weight - last_c.weight * 0.5,
                        total_weight - HALF_SINGLETON_BOUNDARY_WEIGHT,
                        (val - last_c.mean) / last_width,
                    )
                }
            } else {
                total_weight
            }
        } else {
            total_weight
        };
        let prob = (weight / total_weight).clamp(0.0, 1.0);
        for item in &indexed[start..i] {
            result[item.1] = prob;
        }
    }

    // 5. 大于 centroids_max
    while i < n_inputs {
        let val = indexed[i].0;
        if val.is_nan() {
            break;
        }
        let start = i;
        while i < n_inputs && double_equal(indexed[i].0, val) {
            i += 1;
        }
        for item in &indexed[start..i] {
            result[item.1] = 1.0;
        }
    }

    result
}

/// 排名估算算法（对标 Apache Kvrocks TDigestRank）
pub fn tdigest_rank_calc(
    centroids: &[Centroid],
    min: f64,
    max: f64,
    total_weight: f64,
    inputs: &[f64],
    reverse: bool,
) -> Vec<i64> {
    if centroids.is_empty() || total_weight <= 0.0 {
        return vec![-2; inputs.len()];
    }

    let mut result = vec![-2; inputs.len()];
    let mut indexed: Vec<(f64, usize)> = inputs
        .iter()
        .copied()
        .enumerate()
        .map(|(i, v)| (v, i))
        .collect();

    if reverse {
        indexed.sort_unstable_by(|(a, _), (b, _)| double_compare(*b, *a, REL_EPS, ABS_EPS));

        let mut it_idx = 0;
        let n_inputs = indexed.len();

        while it_idx < n_inputs && indexed[it_idx].0 > max {
            let val = indexed[it_idx].0;
            let start = it_idx;
            while it_idx < n_inputs && double_equal(indexed[it_idx].0, val) {
                it_idx += 1;
            }
            for item in &indexed[start..it_idx] {
                result[item.1] = -1;
            }
        }

        let mut cumulative_weight = 0.0;
        let mut c_idx = centroids.len();

        while c_idx > 0 && it_idx < n_inputs {
            c_idx -= 1;
            let centroid = centroids[c_idx];
            let input_val = indexed[it_idx].0;

            if double_equal(centroid.mean, input_val) {
                let current_mean = centroid.mean;
                let mut current_mean_cum_w = cumulative_weight + centroid.weight * 0.5;
                cumulative_weight += centroid.weight;

                while c_idx > 0 && double_equal(centroids[c_idx - 1].mean, current_mean) {
                    c_idx -= 1;
                    current_mean_cum_w += centroids[c_idx].weight * 0.5;
                    cumulative_weight += centroids[c_idx].weight;
                }

                let start = it_idx;
                while it_idx < n_inputs && double_equal(indexed[it_idx].0, input_val) {
                    it_idx += 1;
                }
                for item in &indexed[start..it_idx] {
                    result[item.1] = current_mean_cum_w as i64;
                }
            } else if double_compare(centroid.mean, input_val, REL_EPS, ABS_EPS).is_gt() {
                cumulative_weight += centroid.weight;
            } else {
                let start = it_idx;
                while it_idx < n_inputs && double_equal(indexed[it_idx].0, input_val) {
                    it_idx += 1;
                }
                for item in &indexed[start..it_idx] {
                    result[item.1] = cumulative_weight as i64;
                }
                c_idx += 1; // 重新处理当前质心
            }
        }

        while it_idx < n_inputs {
            let val = indexed[it_idx].0;
            let start = it_idx;
            while it_idx < n_inputs && double_equal(indexed[it_idx].0, val) {
                it_idx += 1;
            }
            for item in &indexed[start..it_idx] {
                result[item.1] = total_weight as i64;
            }
        }
    } else {
        indexed.sort_unstable_by(|(a, _), (b, _)| double_compare(*a, *b, REL_EPS, ABS_EPS));

        let mut it_idx = 0;
        let n_inputs = indexed.len();

        while it_idx < n_inputs && indexed[it_idx].0 < min {
            let val = indexed[it_idx].0;
            let start = it_idx;
            while it_idx < n_inputs && double_equal(indexed[it_idx].0, val) {
                it_idx += 1;
            }
            for item in &indexed[start..it_idx] {
                result[item.1] = -1;
            }
        }

        let mut cumulative_weight = 0.0;
        let mut c_idx = 0;
        let n_centroids = centroids.len();

        while c_idx < n_centroids && it_idx < n_inputs {
            let centroid = centroids[c_idx];
            let input_val = indexed[it_idx].0;

            if double_equal(centroid.mean, input_val) {
                let current_mean = centroid.mean;
                let mut current_mean_cum_w = cumulative_weight + centroid.weight * 0.5;
                cumulative_weight += centroid.weight;

                while c_idx + 1 < n_centroids
                    && double_equal(centroids[c_idx + 1].mean, current_mean)
                {
                    c_idx += 1;
                    current_mean_cum_w += centroids[c_idx].weight * 0.5;
                    cumulative_weight += centroids[c_idx].weight;
                }

                let start = it_idx;
                while it_idx < n_inputs && double_equal(indexed[it_idx].0, input_val) {
                    it_idx += 1;
                }
                for item in &indexed[start..it_idx] {
                    result[item.1] = current_mean_cum_w as i64;
                }
                c_idx += 1;
            } else if double_compare(centroid.mean, input_val, REL_EPS, ABS_EPS).is_lt() {
                cumulative_weight += centroid.weight;
                c_idx += 1;
            } else {
                let start = it_idx;
                while it_idx < n_inputs && double_equal(indexed[it_idx].0, input_val) {
                    it_idx += 1;
                }
                for item in &indexed[start..it_idx] {
                    result[item.1] = cumulative_weight as i64;
                }
            }
        }

        while it_idx < n_inputs {
            let val = indexed[it_idx].0;
            let start = it_idx;
            while it_idx < n_inputs && double_equal(indexed[it_idx].0, val) {
                it_idx += 1;
            }
            for item in &indexed[start..it_idx] {
                result[item.1] = total_weight as i64;
            }
        }
    }

    result
}

/// 按排名查询值（对标 Apache Kvrocks TDigestByRank）
pub fn tdigest_by_rank_calc(
    centroids: &[Centroid],
    total_weight: f64,
    inputs: &[i64],
    reverse: bool,
) -> Vec<f64> {
    if centroids.is_empty() || total_weight <= 0.0 {
        return vec![f64::NAN; inputs.len()];
    }

    let mut result = vec![f64::NAN; inputs.len()];
    let mut indexed: Vec<(i64, usize)> = inputs
        .iter()
        .copied()
        .enumerate()
        .map(|(i, r)| (r, i))
        .collect();

    indexed.sort_unstable_by_key(|&(r, _)| r);

    let mut it_idx = 0;
    let n_inputs = indexed.len();
    let mut cumulative_weight = 0.0;
    let inf_val = if reverse {
        -f64::INFINITY
    } else {
        f64::INFINITY
    };

    if reverse {
        for c in centroids.iter().rev() {
            cumulative_weight += c.weight;
            let target_w = cumulative_weight as i64;
            while it_idx < n_inputs && indexed[it_idx].0 < target_w {
                result[indexed[it_idx].1] = c.mean;
                it_idx += 1;
            }
        }
    } else {
        for c in centroids {
            cumulative_weight += c.weight;
            let target_w = cumulative_weight as i64;
            while it_idx < n_inputs && indexed[it_idx].0 < target_w {
                result[indexed[it_idx].1] = c.mean;
                it_idx += 1;
            }
        }
    }

    let total_w_int = total_weight as i64;
    while it_idx < n_inputs {
        if indexed[it_idx].0 >= total_w_int {
            result[indexed[it_idx].1] = inf_val;
        }
        it_idx += 1;
    }

    result
}

/// 截断均值计算（对标 Apache Kvrocks TDigestTrimmedMean）
pub fn tdigest_trimmed_mean_calc(
    centroids: &[Centroid],
    total_weight: f64,
    low_cut: f64,
    high_cut: f64,
) -> f64 {
    if centroids.is_empty()
        || low_cut.is_nan()
        || high_cut.is_nan()
        || !(0.0..=1.0).contains(&low_cut)
        || !(0.0..=1.0).contains(&high_cut)
        || low_cut >= high_cut
        || total_weight <= 0.0
    {
        return f64::NAN;
    }

    let leftmost_weight = (total_weight * low_cut).floor();
    let rightmost_weight = (total_weight * high_cut).ceil();

    let mut count_done = 0.0;
    let mut trimmed_sum = 0.0;
    let mut trimmed_count = 0.0;

    for c in centroids {
        let n_weight = c.weight;
        let mut count_add = n_weight;

        count_add -= (leftmost_weight - count_done).max(0.0).min(count_add);
        count_add = (rightmost_weight - count_done).max(0.0).min(count_add);

        count_done += n_weight;

        trimmed_sum += c.mean * count_add;
        trimmed_count += count_add;

        if count_done >= rightmost_weight {
            break;
        }
    }

    if trimmed_count == 0.0 {
        f64::NAN
    } else {
        trimmed_sum / trimmed_count
    }
}

/// T-Digest 核心算法结构
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TDigest {
    pub compression: f64,
    pub capacity: usize,
    pub centroids: Vec<Centroid>,
    pub unmerged_buffer: Vec<f64>,
    pub total_weight: f64,
    pub merged_weight: f64,
    pub total_observations: u64,
    pub merge_times: u64,
    pub min: f64,
    pub max: f64,
}

impl TDigest {
    #[inline]
    pub fn new(compression: f64) -> Self {
        let comp = if compression <= 0.0 {
            DEFAULT_COMPRESSION as f64
        } else {
            compression
        };
        let cap = calculate_capacity(comp as u32);
        Self {
            compression: comp,
            capacity: cap,
            centroids: Vec::with_capacity(cap),
            unmerged_buffer: Vec::with_capacity(cap),
            total_weight: 0.0,
            merged_weight: 0.0,
            total_observations: 0,
            merge_times: 0,
            min: f64::MAX,
            max: -f64::MAX,
        }
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.total_observations == 0 && self.total_weight <= 0.0
    }

    #[inline]
    pub fn reset(&mut self) {
        self.centroids.clear();
        self.unmerged_buffer.clear();
        self.total_weight = 0.0;
        self.merged_weight = 0.0;
        self.total_observations = 0;
        self.merge_times = 0;
        self.min = f64::MAX;
        self.max = -f64::MAX;
    }

    pub fn add(&mut self, val: f64, weight: f64) {
        if val.is_nan() || weight <= 0.0 {
            return;
        }
        self.unmerged_buffer.push(val);
        self.total_weight += weight;
        self.total_observations += 1;
        if val < self.min {
            self.min = val;
        }
        if val > self.max {
            self.max = val;
        }
        if self.unmerged_buffer.len() >= self.capacity {
            self.ensure_merged();
        }
    }

    #[inline]
    pub fn add_batch(&mut self, values: &[f64]) {
        for &val in values {
            self.add(val, 1.0);
        }
    }

    #[inline]
    pub fn ensure_merged(&mut self) {
        if self.unmerged_buffer.is_empty() {
            return;
        }
        self.compress();
    }

    pub fn compress(&mut self) {
        let delta = self.compression as u32;
        let existing = CentroidsWithDelta {
            centroids: take(&mut self.centroids),
            delta,
            min: self.min,
            max: self.max,
            total_weight: self.merged_weight,
        };

        let merged = tdigest_merge_buffer_and_centroids(&self.unmerged_buffer, &[existing], delta);
        self.unmerged_buffer.clear();
        self.centroids = merged.centroids;
        self.merged_weight = merged.total_weight;
        self.total_weight = merged.total_weight;
        self.merge_times += 1;
        self.min = merged.min;
        self.max = merged.max;
    }

    pub fn merge_from(&mut self, other: &mut TDigest) {
        other.ensure_merged();
        if other.is_empty() {
            return;
        }
        self.ensure_merged();

        let delta = (self.compression as u32).max(other.compression as u32);
        let list = [
            CentroidsWithDelta {
                centroids: take(&mut self.centroids),
                delta: self.compression as u32,
                min: self.min,
                max: self.max,
                total_weight: self.merged_weight,
            },
            CentroidsWithDelta {
                centroids: take(&mut other.centroids),
                delta: other.compression as u32,
                min: other.min,
                max: other.max,
                total_weight: other.merged_weight,
            },
        ];

        let merged = tdigest_merge_centroids_list(&list, delta);
        self.compression = delta as f64;
        self.capacity = calculate_capacity(delta);
        self.centroids = merged.centroids;
        self.merged_weight = merged.total_weight;
        self.total_weight = merged.total_weight;
        self.total_observations += other.total_observations;
        self.merge_times += 1;
        self.min = merged.min;
        self.max = merged.max;
    }

    pub fn merge_with_options(&mut self, sources: &mut [TDigest], options: TDigestMerge) {
        self.ensure_merged();
        for s in sources.iter_mut() {
            s.ensure_merged();
        }

        let mut max_source_comp = DEFAULT_COMPRESSION;
        for s in sources.iter() {
            max_source_comp = max_source_comp.max(s.compression as u32);
        }

        let final_comp = if let Some(c) = options.compression {
            c
        } else if options.override_dest || self.is_empty() {
            max_source_comp
        } else {
            self.compression as u32
        };

        let mut list = Vec::with_capacity(sources.len() + 1);
        let mut total_obs = 0u64;

        if !options.override_dest && !self.is_empty() {
            total_obs += self.total_observations;
            list.push(CentroidsWithDelta {
                centroids: take(&mut self.centroids),
                delta: self.compression as u32,
                min: self.min,
                max: self.max,
                total_weight: self.merged_weight,
            });
        }

        for s in sources.iter_mut() {
            if !s.is_empty() {
                total_obs += s.total_observations;
                list.push(CentroidsWithDelta {
                    centroids: take(&mut s.centroids),
                    delta: s.compression as u32,
                    min: s.min,
                    max: s.max,
                    total_weight: s.merged_weight,
                });
            }
        }

        if list.is_empty() {
            self.reset();
            self.compression = final_comp as f64;
            self.capacity = calculate_capacity(final_comp);
            return;
        }

        let merged = tdigest_merge_centroids_list(&list, final_comp);
        self.compression = final_comp as f64;
        self.capacity = calculate_capacity(final_comp);
        self.centroids = merged.centroids;
        self.merged_weight = merged.total_weight;
        self.total_weight = merged.total_weight;
        self.total_observations = total_obs;
        self.merge_times += 1;
        self.min = merged.min;
        self.max = merged.max;
    }

    #[inline]
    pub fn min(&mut self) -> f64 {
        if self.is_empty() { f64::NAN } else { self.min }
    }

    #[inline]
    pub fn max(&mut self) -> f64 {
        if self.is_empty() { f64::NAN } else { self.max }
    }

    #[inline]
    pub fn quantile(&mut self, q: f64) -> f64 {
        self.ensure_merged();
        tdigest_quantile_calc(&self.centroids, self.min, self.max, self.total_weight, q)
    }

    #[inline]
    pub fn cdf(&mut self, val: f64) -> f64 {
        self.ensure_merged();
        let res = tdigest_cdf_calc(
            &self.centroids,
            self.min,
            self.max,
            self.total_weight,
            &[val],
        );
        res.first().copied().unwrap_or(f64::NAN)
    }

    #[inline]
    pub fn rank(&mut self, val: f64) -> i64 {
        self.ensure_merged();
        let res = tdigest_rank_calc(
            &self.centroids,
            self.min,
            self.max,
            self.total_weight,
            &[val],
            false,
        );
        res.first().copied().unwrap_or(-2)
    }

    #[inline]
    pub fn revrank(&mut self, val: f64) -> i64 {
        self.ensure_merged();
        let res = tdigest_rank_calc(
            &self.centroids,
            self.min,
            self.max,
            self.total_weight,
            &[val],
            true,
        );
        res.first().copied().unwrap_or(-2)
    }

    #[inline]
    pub fn byrank(&mut self, r: u64) -> f64 {
        self.ensure_merged();
        let res = tdigest_by_rank_calc(&self.centroids, self.total_weight, &[r as i64], false);
        res.first().copied().unwrap_or(f64::NAN)
    }

    #[inline]
    pub fn byrevrank(&mut self, r: u64) -> f64 {
        self.ensure_merged();
        let res = tdigest_by_rank_calc(&self.centroids, self.total_weight, &[r as i64], true);
        res.first().copied().unwrap_or(f64::NAN)
    }

    #[inline]
    pub fn trimmed_mean(&mut self, low_cut: f64, high_cut: f64) -> f64 {
        self.ensure_merged();
        tdigest_trimmed_mean_calc(&self.centroids, self.total_weight, low_cut, high_cut)
    }

    #[inline]
    pub fn info(&mut self) -> TDigestInfo {
        self.ensure_merged();
        let unmerged_w = (self.total_weight - self.merged_weight).max(0.0);
        TDigestInfo {
            compression: self.compression as u32,
            capacity: self.capacity,
            merged_nodes: self.centroids.len(),
            unmerged_nodes: self.unmerged_buffer.len(),
            merged_weight: self.merged_weight,
            unmerged_weight: unmerged_w,
            total_weight: self.total_weight,
            observations: self.total_observations,
            total_compressions: self.merge_times,
            minimum: if self.is_empty() {
                None
            } else {
                Some(self.min)
            },
            maximum: if self.is_empty() {
                None
            } else {
                Some(self.max)
            },
        }
    }
}

/// TDigest 合并器工具
pub struct TDigestMergerTool;

impl TDigestMergerTool {
    pub fn merge(dest: &mut TDigest, sources: &mut [TDigest]) {
        for src in sources {
            dest.merge_from(src);
        }
    }
}

impl WeDb {
    /// TDIGEST.CREATE key [COMPRESSION compression]
    pub fn tdigest_create<K: AsRef<[u8]>>(&self, key: K, compression: f64) -> Result<()> {
        let comp = if compression <= 0.0 {
            DEFAULT_COMPRESSION
        } else {
            compression as u32
        };
        if !(MIN_COMPRESSION..=MAX_COMPRESSION).contains(&comp) {
            return Err(Error::invalid_data(format!(
                "ERR compression out of range [{MIN_COMPRESSION}, {MAX_COMPRESSION}]"
            )));
        }

        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        if self.data_ks.contains_key(raw_k.as_bytes())? {
            return Err(Error::invalid_data("ERR item exists"));
        }

        let td = TDigest::new(comp as f64);
        let serialized = sonic_rs::to_vec(&td)
            .map_err(|e| Error::invalid_data(format!("ERR tdigest serialize: {e}")))?;

        self.data_ks.insert(raw_k.as_bytes(), serialized)?;
        Ok(())
    }

    /// TDIGEST.ADD key value [value ...]
    pub fn tdigest_add<K: AsRef<[u8]>>(&self, key: K, values: &[f64]) -> Result<()> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        for &val in values {
            if !val.is_nan() {
                td.add(val, 1.0);
            }
        }
        td.ensure_merged();

        let serialized = sonic_rs::to_vec(&td)
            .map_err(|e| Error::invalid_data(format!("ERR tdigest serialize: {e}")))?;

        self.data_ks.insert(raw_k.as_bytes(), serialized)?;
        Ok(())
    }

    /// TDIGEST.MIN key
    pub fn tdigest_min<K: AsRef<[u8]>>(&self, key: K) -> Result<f64> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        Ok(td.min())
    }

    /// TDIGEST.MAX key
    pub fn tdigest_max<K: AsRef<[u8]>>(&self, key: K) -> Result<f64> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        Ok(td.max())
    }

    /// TDIGEST.QUANTILE key quantile [quantile ...]
    pub fn tdigest_quantile<K: AsRef<[u8]>>(
        &self,
        key: K,
        quantiles: &[f64],
    ) -> Result<Vec<Option<f64>>> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        let mut results = Vec::with_capacity(quantiles.len());
        for &q in quantiles {
            let val = td.quantile(q);
            if val.is_nan() {
                results.push(None);
            } else {
                results.push(Some(val));
            }
        }
        Ok(results)
    }

    /// TDIGEST.CDF key value [value ...]
    pub fn tdigest_cdf<K: AsRef<[u8]>>(&self, key: K, values: &[f64]) -> Result<Vec<Option<f64>>> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        let mut results = Vec::with_capacity(values.len());
        for &val in values {
            let cdf_val = td.cdf(val);
            if cdf_val.is_nan() {
                results.push(None);
            } else {
                results.push(Some(cdf_val));
            }
        }
        Ok(results)
    }

    /// TDIGEST.RANK key value [value ...]
    pub fn tdigest_rank<K: AsRef<[u8]>>(&self, key: K, values: &[f64]) -> Result<Vec<i64>> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        Ok(values.iter().map(|&v| td.rank(v)).collect())
    }

    /// TDIGEST.REVRANK key value [value ...]
    pub fn tdigest_revrank<K: AsRef<[u8]>>(&self, key: K, values: &[f64]) -> Result<Vec<i64>> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        Ok(values.iter().map(|&v| td.revrank(v)).collect())
    }

    /// TDIGEST.BYRANK key rank [rank ...]
    pub fn tdigest_byrank<K: AsRef<[u8]>>(
        &self,
        key: K,
        ranks: &[u64],
    ) -> Result<Vec<Option<f64>>> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        Ok(ranks
            .iter()
            .map(|&r| {
                let v = td.byrank(r);
                if v.is_nan() { None } else { Some(v) }
            })
            .collect())
    }

    /// TDIGEST.BYREVRANK key rank [rank ...]
    pub fn tdigest_byrevrank<K: AsRef<[u8]>>(
        &self,
        key: K,
        ranks: &[u64],
    ) -> Result<Vec<Option<f64>>> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        Ok(ranks
            .iter()
            .map(|&r| {
                let v = td.byrevrank(r);
                if v.is_nan() { None } else { Some(v) }
            })
            .collect())
    }

    /// TDIGEST.TRIMMED_MEAN key low_cut_quantile high_cut_quantile
    pub fn tdigest_trimmed_mean<K: AsRef<[u8]>>(
        &self,
        key: K,
        low_cut: f64,
        high_cut: f64,
    ) -> Result<Option<f64>> {
        if !low_cut.is_finite()
            || !high_cut.is_finite()
            || !(0.0..=1.0).contains(&low_cut)
            || !(0.0..=1.0).contains(&high_cut)
        {
            return Err(Error::invalid_data(
                "ERR low_cut_percentile and high_cut_percentile should be in [0,1]",
            ));
        }
        if low_cut >= high_cut {
            return Err(Error::invalid_data(
                "ERR low_cut_percentile should be lower than high_cut_percentile",
            ));
        }

        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        let mean = td.trimmed_mean(low_cut, high_cut);
        if mean.is_nan() {
            Ok(None)
        } else {
            Ok(Some(mean))
        }
    }

    /// TDIGEST.RESET key
    pub fn tdigest_reset<K: AsRef<[u8]>>(&self, key: K) -> Result<()> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };
        td.reset();

        let serialized = sonic_rs::to_vec(&td)
            .map_err(|e| Error::invalid_data(format!("ERR tdigest serialize: {e}")))?;

        self.data_ks.insert(raw_k.as_bytes(), serialized)?;
        Ok(())
    }

    /// TDIGEST.MERGE dest_key numkeys source_key [source_key ...] [COMPRESSION compression] [OVERRIDE]
    pub fn tdigest_merge<K: AsRef<[u8]>>(
        &self,
        dest_key: K,
        source_keys: &[K],
        options: TDigestMerge,
    ) -> Result<()> {
        if source_keys.is_empty() {
            return Err(Error::invalid_data(
                "ERR wrong number of arguments for 'tdigest.merge' command",
            ));
        }

        if let Some(comp) = options.compression
            && !(MIN_COMPRESSION..=MAX_COMPRESSION).contains(&comp)
        {
            return Err(Error::invalid_data(format!(
                "ERR compression out of range [{MIN_COMPRESSION}, {MAX_COMPRESSION}]"
            )));
        }

        let kc = KeyComposer::new("default");
        let dst_str = str::from_utf8(dest_key.as_ref()).unwrap_or("");
        let dst_raw_k = kc.tdigest_meta(dst_str);

        let mut dest_td = if !options.override_dest {
            match self.data_ks.get(dst_raw_k.as_bytes())? {
                Some(bytes) => sonic_rs::from_slice(&bytes)
                    .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
                None => TDigest::new(100.0),
            }
        } else {
            TDigest::new(100.0)
        };

        let mut source_tds = Vec::with_capacity(source_keys.len());
        for src in source_keys {
            let src_str = str::from_utf8(src.as_ref()).unwrap_or("");
            let src_raw_k = kc.tdigest_meta(src_str);
            if let Some(bytes) = self.data_ks.get(src_raw_k.as_bytes())? {
                let src_td: TDigest = sonic_rs::from_slice(&bytes)
                    .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?;
                source_tds.push(src_td);
            } else {
                return Err(Error::invalid_data(format!("ERR key not found: {src_str}")));
            }
        }

        dest_td.merge_with_options(&mut source_tds, options);

        let serialized = sonic_rs::to_vec(&dest_td)
            .map_err(|e| Error::invalid_data(format!("ERR tdigest serialize: {e}")))?;

        self.data_ks.insert(dst_raw_k.as_bytes(), serialized)?;
        Ok(())
    }

    /// TDIGEST.INFO key
    pub fn tdigest_info<K: AsRef<[u8]>>(&self, key: K) -> Result<TDigestInfo> {
        let kc = KeyComposer::new("default");
        let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
        let raw_k = kc.tdigest_meta(k_str);

        let mut td: TDigest = match self.data_ks.get(raw_k.as_bytes())? {
            Some(bytes) => sonic_rs::from_slice(&bytes)
                .map_err(|e| Error::invalid_data(format!("ERR tdigest deserialize: {e}")))?,
            None => return Err(Error::invalid_data("ERR key does not exist")),
        };

        Ok(td.info())
    }
}