1use poulpy_hal::{
 2    api::{
 3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
 4        VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalizeTmpBytes,
 5        VecZnxSubScalarInplace,
 6    },
 7    layouts::{Backend, DataRef, Module, ScalarZnx, ScratchOwned, ZnxZero},
 8    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl},
 9};
10
11use crate::layouts::{GGLWE, GGLWEInfos, GLWE, GLWEPlaintext, LWEInfos, prepared::GLWESecretPrepared};
12
13impl<D: DataRef> GGLWE<D> {
14    pub fn assert_noise<B, DataSk, DataWant>(
15        &self,
16        module: &Module<B>,
17        sk: &GLWESecretPrepared<DataSk, B>,
18        pt_want: &ScalarZnx<DataWant>,
19        max_noise: f64,
20    ) where
21        DataSk: DataRef,
22        DataWant: DataRef,
23        Module<B>: VecZnxDftBytesOf
24            + VecZnxBigBytesOf
25            + VecZnxDftApply<B>
26            + SvpApplyDftToDftInplace<B>
27            + VecZnxIdftApplyConsume<B>
28            + VecZnxBigAddInplace<B>
29            + VecZnxBigAddSmallInplace<B>
30            + VecZnxBigNormalize<B>
31            + VecZnxNormalizeTmpBytes
32            + VecZnxSubScalarInplace,
33        B: Backend + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
34    {
35        let dsize: usize = self.dsize().into();
36        let base2k: usize = self.base2k().into();
37
38        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWE::decrypt_scratch_space(module, self));
39        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, self);
40
41        (0..self.rank_in().into()).for_each(|col_i| {
42            (0..self.dnum().into()).for_each(|row_i| {
43                self.at(row_i, col_i)
44                    .decrypt(module, &mut pt, sk, scratch.borrow());
45
46                module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, (dsize - 1) + row_i * dsize, pt_want, col_i);
47
48                let noise_have: f64 = pt.data.std(base2k, 0).log2();
49
50                println!("noise_have: {noise_have}");
51
52                assert!(
53                    noise_have <= max_noise,
54                    "noise_have: {noise_have} > max_noise: {max_noise}"
55                );
56
57                pt.data.zero();
58            });
59        });
60    }
61}