1mod test {
  2    use itertools::{izip, Itertools};
  3
  4    use crate::{
  5        backend::{ModularOpsU64, ModulusPowerOf2},
  6        bool::{
  7            evaluator::{BoolEncoding, BoolEvaluator, BooleanGates},
  8            keys::{
  9                InteractiveMultiPartyClientKey, PublicKey, ServerKeyEvaluationDomain,
 10                ShoupServerKeyEvaluationDomain,
 11            },
 12            parameters::{CiphertextModulus, SMALL_MP_BOOL_PARAMS},
 13        },
 14        evaluator::MultiPartyCrs,
 15        ntt::NttBackendU64,
 16        parameters::OPTIMISED_SMALL_MP_BOOL_PARAMS,
 17        random::DefaultSecureRng,
 18    };
 19
 20    #[test]
 21    fn mp_noise() {
 22        // set_parameter_set(&SMALL_MP_BOOL_PARAMS);
 23        let mut evaluator = BoolEvaluator::<
 24            Vec<Vec<u64>>,
 25            NttBackendU64,
 26            ModularOpsU64<CiphertextModulus<u64>>,
 27            ModulusPowerOf2<CiphertextModulus<u64>>,
 28            ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>,
 29        >::new(OPTIMISED_SMALL_MP_BOOL_PARAMS);
 30
 31        let parties = 2;
 32
 33        let cr_seed = MultiPartyCrs::random();
 34
 35        let cks = (0..parties)
 36            .into_iter()
 37            .map(|_| evaluator.client_key())
 38            .collect_vec();
 39
 40        // construct ideal rlwe sk for meauring noise
 41        let mut ideal_rlwe_sk = vec![0i32; evaluator.parameters().rlwe_n().0];
 42        cks.iter().for_each(|k| {
 43            izip!(
 44                ideal_rlwe_sk.iter_mut(),
 45                InteractiveMultiPartyClientKey::sk_rlwe(k)
 46            )
 47            .for_each(|(ideal_i, s_i)| {
 48                *ideal_i = *ideal_i + s_i;
 49            });
 50        });
 51        let mut ideal_lwe_sk = vec![0i32; evaluator.parameters().lwe_n().0];
 52        cks.iter().for_each(|k| {
 53            izip!(
 54                ideal_lwe_sk.iter_mut(),
 55                InteractiveMultiPartyClientKey::sk_lwe(k)
 56            )
 57            .for_each(|(ideal_i, s_i)| {
 58                *ideal_i = *ideal_i + s_i;
 59            });
 60        });
 61
 62        // round 1
 63        let pk_shares = cks
 64            .iter()
 65            .map(|c| evaluator.multi_party_public_key_share(&cr_seed, c))
 66            .collect_vec();
 67
 68        // public key
 69        let pk = PublicKey::<Vec<Vec<u64>>, DefaultSecureRng, ModularOpsU64<CiphertextModulus<u64>>>::from(
 70            pk_shares.as_slice(),
 71        );
 72
 73        // round 2
 74        let server_key_shares = cks
 75            .iter()
 76            .enumerate()
 77            .map(|(index, c)| {
 78                evaluator.multi_party_server_key_share(index, parties, &cr_seed, &pk.key(), c)
 79            })
 80            .collect_vec();
 81
 82        let server_key = evaluator.aggregate_multi_party_server_key_shares(&server_key_shares);
 83        let runtime_server_key = ShoupServerKeyEvaluationDomain::from(ServerKeyEvaluationDomain::<
 84            _,
 85            _,
 86            DefaultSecureRng,
 87            NttBackendU64,
 88        >::from(&server_key));
 89
 90        let mut m0 = false;
 91        let mut m1 = true;
 92
 93        let mut c_m0 = evaluator.pk_encrypt(pk.key(), m0);
 94        let mut c_m1 = evaluator.pk_encrypt(pk.key(), m1);
 95
 96        // let mut stats = Stats::new();
 97
 98        for _ in 0..1000 {
 99            let now = std::time::Instant::now();
100            let c_out = evaluator.xor(&c_m0, &c_m1, &runtime_server_key);
101            println!("Gate time: {:?}", now.elapsed());
102
103            // mp decrypt
104            let decryption_shares = cks
105                .iter()
106                .map(|c| evaluator.multi_party_decryption_share(&c_out, c))
107                .collect_vec();
108            let m_out = evaluator.multi_party_decrypt(&decryption_shares, &c_out);
109            let m_expected = (m0 ^ m1);
110            assert_eq!(m_expected, m_out, "Expected {m_expected} but got {m_out}");
111
112            // // find noise update
113            // {
114            //     let out = decrypt_lwe(
115            //         &c_out,
116            //         ideal_client_key.sk_rlwe().values(),
117            //         evaluator.pbs_info().modop_rlweq(),
118            //     );
119
120            //     let out_want = {
121            //         if m_expected == true {
122            //             true_el_encoded
123            //         } else {
124            //             false_el_encoded
125            //         }
126            //     };
127            //     let diff = evaluator.pbs_info().modop_rlweq().sub(&out, &out_want);
128
129            //     stats.add_more(&vec![evaluator
130            //         .pbs_info()
131            //         .rlwe_q()
132            //         .map_element_to_i64(&diff)]);
133            // }
134
135            m1 = m0;
136            m0 = m_expected;
137
138            c_m1 = c_m0;
139            c_m0 = c_out;
140        }
141
142        // println!("log2 std dev {}", stats.std_dev().abs().log2());
143    }
144}