1use rug::{
 2    Float,
 3    float::Round,
 4    ops::{AddAssignRound, DivAssignRound, SubAssignRound},
 5};
 6
 7use crate::layouts::{Backend, DataRef, VecZnx, VecZnxBig, VecZnxBigToRef, ZnxInfos};
 8
 9pub struct Stats {
10    max: f64,
11    std: f64,
12}
13
14impl Stats {
15    pub fn max(&self) -> f64 {
16        self.max
17    }
18
19    pub fn std(&self) -> f64 {
20        self.std
21    }
22}
23
24impl<D: DataRef> VecZnx<D> {
25    pub fn stats(&self, base2k: usize, col: usize) -> Stats {
26        let prec: u32 = (self.size() * base2k) as u32;
27        let mut data: Vec<Float> = (0..self.n()).map(|_| Float::with_val(prec, 0)).collect();
28        self.decode_vec_float(base2k, col, &mut data);
29        // std = sqrt(sum((xi - avg)^2) / n)
30        let mut avg: Float = Float::with_val(prec, 0);
31        let mut max: Float = Float::with_val(prec, 0);
32
33        data.iter().for_each(|x| {
34            avg.add_assign_round(x, Round::Nearest);
35            max.max_mut(&Float::with_val(53, x.abs_ref()));
36        });
37        avg.div_assign_round(Float::with_val(prec, data.len()), Round::Nearest);
38        data.iter_mut().for_each(|x| {
39            x.sub_assign_round(&avg, Round::Nearest);
40        });
41        let mut std: Float = Float::with_val(prec, 0);
42        data.iter().for_each(|x| std += x * x);
43        std.div_assign_round(Float::with_val(prec, data.len()), Round::Nearest);
44        std = std.sqrt();
45        Stats {
46            std: std.to_f64(),
47            max: max.to_f64(),
48        }
49    }
50}
51
52impl<D: DataRef, B: Backend + Backend<ScalarBig = i64>> VecZnxBig<D, B> {
53    pub fn stats(&self, base2k: usize, col: usize) -> Stats {
54        let self_ref: VecZnxBig<&[u8], B> = self.to_ref();
55        let znx: VecZnx<&[u8]> = VecZnx {
56            data: self_ref.data,
57            n: self_ref.n,
58            cols: self_ref.cols,
59            size: self_ref.size,
60            max_size: self_ref.max_size,
61        };
62        znx.stats(base2k, col)
63    }
64}