1use std::time::Instant;
  2
  3use backend::{Encoding, Module, ScratchOwned, Stats, ZnxView, ZnxViewMut, FFT64};
  4use sampling::source::Source;
  5
  6use crate::{
  7    blind_rotation::{ccgi::{cggi_blind_rotate, cggi_blind_rotate_scratch_space, mod_switch_2n}, key::BlindRotationKeyCGGI}, lwe::{ciphertext::{LWECiphertextToMut, LWECiphertextToRef}, LWEPlaintext}, FourierGLWESecret, GLWECiphertext, GLWEOps, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWESecret
  8};
  9
 10#[test]
 11fn blind_rotation() {
 12    let module: Module<FFT64> = Module::<FFT64>::new(2048);
 13    let basek: usize = 20;
 14
 15    let n_lwe: usize = 1071;
 16
 17    let k_lwe: usize = 22;
 18    let k_brk: usize = 60;
 19    let rows_brk: usize = 2;
 20    let k_lut: usize = 60;
 21    let rank: usize = 1;
 22    let block_size: usize = 7;
 23
 24    let mut source_xs: Source = Source::new([0u8; 32]);
 25    let mut source_xe: Source = Source::new([0u8; 32]);
 26    let mut source_xa: Source = Source::new([0u8; 32]);
 27
 28    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
 29    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 30    let sk_glwe_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk_glwe);
 31
 32    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 33    sk_lwe.fill_binary_block(block_size, &mut source_xs);
 34
 35    let mut scratch: ScratchOwned = ScratchOwned::new(BlindRotationKeyCGGI::generate_from_sk_scratch_space(
 36        &module, basek, k_brk, rank,
 37    ) | cggi_blind_rotate_scratch_space(&module, basek, k_lut, k_brk, rows_brk, rank));
 38
 39    let start: Instant = Instant::now();
 40    let mut brk: BlindRotationKeyCGGI<FFT64> = BlindRotationKeyCGGI::allocate(&module, n_lwe, basek, k_brk, rows_brk, rank);
 41    brk.generate_from_sk(
 42        &module,
 43        &sk_glwe_dft,
 44        &sk_lwe,
 45        &mut source_xa,
 46        &mut source_xe,
 47        3.2,
 48        scratch.borrow(),
 49    );
 50    let duration: std::time::Duration = start.elapsed();
 51    println!("brk-gen: {} ms", duration.as_millis());
 52
 53    let mut lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe);
 54
 55    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe);
 56
 57    let x: i64 = 0;
 58    let bits: usize = 6;
 59
 60    pt_lwe.data.encode_coeff_i64(0, basek, bits, 0, x, bits);
 61
 62    println!("{}", pt_lwe.data);
 63
 64    lwe.encrypt_sk(&pt_lwe, &sk_lwe, &mut source_xa, &mut source_xe, 3.2);
 65
 66    lwe.decrypt(&mut pt_lwe, &sk_lwe);
 67
 68    println!("{}", pt_lwe.data);
 69
 70    let mut lut: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_lut);
 71
 72    
 73    lut.data.at_mut(0, 0)[0] = 0;
 74    (1..module.n()).for_each(|i|{
 75        lut.data.at_mut(0, 0)[i] = - ((module.n() as i64 - i as i64  - 1)<<(basek - module.log_n() - 1));
 76    });
 77
 78    
 79
 80
 81    let mut res: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_lut, rank);
 82
 83    let start: Instant = Instant::now();
 84    (0..1).for_each(|i|{
 85        cggi_blind_rotate(&module, &mut res, &lwe, &lut, &brk, scratch.borrow());
 86    });
 87    
 88    let duration: std::time::Duration = start.elapsed();
 89    println!("blind-rotate: {} ms", duration.as_millis());
 90
 91    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_lut);
 92
 93    res.decrypt(&module , &mut pt_have, &sk_glwe_dft, scratch.borrow());
 94
 95    let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
 96
 97    mod_switch_2n(&module, &mut lwe_2n, &lwe.to_ref());
 98
 99    let pt_want: i64 = (lwe_2n[0] + lwe_2n[1..].iter().zip(sk_lwe.data.at(0, 0)).map(|(x, y)| x * y).sum::<i64>()) % (module.n() as i64 * 2);
100
101    lut.rotate_inplace(&module, pt_want);
102
103    lut.sub_inplace_ab(&module, &pt_have);
104
105    let noise: f64 = lut.data.std(0, basek);
106
107    println!("noise: {}", noise);
108
109    
110}