1use base2k::{
 2    Encoding, FFT64, MatZnxDft, MatZnxDftOps, Module, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps,
 3    ZnxInfos, ZnxLayout, alloc_aligned,
 4};
 5
 6fn main() {
 7    let log_n: i32 = 5;
 8    let n: usize = 1 << log_n;
 9
10    let module: Module<FFT64> = Module::<FFT64>::new(n);
11    let log_base2k: usize = 15;
12    let limbs_vec: usize = 5;
13    let log_k: usize = log_base2k * limbs_vec - 5;
14
15    let rows_mat: usize = limbs_vec;
16    let limbs_mat: usize = limbs_vec + 1;
17
18    // Maximum size of the byte scratch needed
19    let tmp_bytes: usize = module.vmp_prepare_tmp_bytes(rows_mat, 1, limbs_mat)
20        | module.vmp_apply_dft_tmp_bytes(limbs_vec, limbs_vec, rows_mat, limbs_mat);
21
22    let mut buf: Vec<u8> = alloc_aligned(tmp_bytes);
23
24    let mut a_values: Vec<i64> = vec![i64::default(); n];
25    a_values[1] = (1 << log_base2k) + 1;
26
27    let mut a: VecZnx = module.new_vec_znx(1, limbs_vec);
28    a.encode_vec_i64(0, log_base2k, log_k, &a_values, 32);
29    a.normalize(log_base2k, &mut buf);
30
31    a.print(n);
32    println!();
33
34    let mut mat_znx_dft: MatZnxDft<FFT64> = module.new_mat_znx_dft(rows_mat, 1, limbs_mat);
35
36    (0..a.size()).for_each(|row_i| {
37        let mut tmp: VecZnx = module.new_vec_znx(1, limbs_mat);
38        tmp.at_limb_mut(row_i)[1] = 1 as i64;
39        module.vmp_prepare_row(&mut mat_znx_dft, tmp.raw(), row_i, &mut buf);
40    });
41
42    let mut c_dft: VecZnxDft<FFT64> = module.new_vec_znx_dft(1, limbs_mat);
43    module.vmp_apply_dft(&mut c_dft, &a, &mat_znx_dft, &mut buf);
44
45    let mut c_big: VecZnxBig<FFT64> = c_dft.as_vec_znx_big();
46    module.vec_znx_idft_tmp_a(&mut c_big, &mut c_dft);
47
48    let mut res: VecZnx = module.new_vec_znx(1, limbs_vec);
49    module.vec_znx_big_normalize(log_base2k, &mut res, &c_big, &mut buf);
50
51    let mut values_res: Vec<i64> = vec![i64::default(); n];
52    res.decode_vec_i64(0, log_base2k, log_k, &mut values_res);
53
54    res.print(n);
55
56    module.free();
57
58    println!("{:?}", values_res)
59}