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