1use base2k::{Matrix3D, Module, VecZnx, VecZnxBig, VecZnxDft, VmpPMat, FFT64};
 2
 3fn main() {
 4    let log_n = 5;
 5    let n = 1 << log_n;
 6
 7    let module: Module = Module::new::<FFT64>(n);
 8    let log_base2k: usize = 15;
 9    let log_q: usize = 60;
10    let limbs: usize = (log_q + log_base2k - 1) / log_base2k;
11
12    let rows: usize = limbs + 1;
13    let cols: usize = limbs + 1;
14
15    // Maximum size of the byte scratch needed
16    let tmp_bytes: usize = module.vmp_prepare_contiguous_tmp_bytes(rows, cols)
17        | module.vmp_apply_dft_tmp_bytes(limbs, limbs, rows, cols);
18
19    let mut buf: Vec<u8> = vec![0; tmp_bytes];
20
21    let mut a_values: Vec<i64> = vec![i64::default(); n];
22    a_values[1] = (1 << log_base2k) + 1;
23
24    let mut a: VecZnx = module.new_vec_znx(log_base2k, log_q);
25    a.from_i64(&a_values, 32);
26    a.normalize(&mut buf);
27
28    (0..a.limbs()).for_each(|i| println!("{}: {:?}", i, a.at(i)));
29
30    let mut b_mat: Matrix3D<i64> = Matrix3D::new(rows, cols, n);
31
32    (0..rows).for_each(|i| {
33        (0..cols).for_each(|j| {
34            b_mat.at_mut(i, j)[0] = (i * cols + j) as i64;
35            b_mat.at_mut(i, j)[0] = (i * cols + j) as i64;
36        })
37    });
38
39    //b_mat.data.iter_mut().enumerate().for_each(|(i, xi)| *xi = i as i64);
40
41    println!();
42    (0..rows).for_each(|i| {
43        (0..cols).for_each(|j| println!("{} {}: {:?}", i, j, b_mat.at(i, j)));
44        println!();
45    });
46
47    let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
48    module.vmp_prepare_contiguous(&mut vmp_pmat, &b_mat.data, &mut buf);
49
50    (0..cols).for_each(|i| {
51        (0..rows).for_each(|j| println!("{} {}: {:?}", i, j, vmp_pmat.at(i, j)));
52        println!();
53    });
54
55    println!("{:?}", vmp_pmat.as_f64());
56
57    let mut c_dft: VecZnxDft = module.new_vec_znx_dft(limbs);
58    module.vmp_apply_dft(&mut c_dft, &a, &vmp_pmat, &mut buf);
59
60    let mut c_big: VecZnxBig = c_dft.as_vec_znx_big();
61    module.vec_znx_idft_tmp_a(&mut c_big, &mut c_dft, limbs);
62
63    let mut res: VecZnx = module.new_vec_znx(log_base2k, log_q);
64    module.vec_znx_big_normalize(&mut res, &c_big, &mut buf);
65
66    let mut values_res: Vec<i64> = vec![i64::default(); n];
67    res.to_i64(&mut values_res);
68
69    (0..res.limbs()).for_each(|i| println!("{}: {:?}", i, res.at(i)));
70
71    module.delete();
72    c_dft.delete();
73    vmp_pmat.delete();
74
75    println!("{:?}", values_res)
76}