use tfhe::prelude::*;
use tfhe::{set_server_key, generate_keys, ConfigBuilder, FheUint128};


fn mul_all(a: &FheUint128, b: &FheUint128, c: &FheUint128) -> FheUint128 {
    // Use the debug format ('{:?}'), if you don't want to unwrap()
    // and panic if the value is not a trivial.
    println!(
        "a: {:?}, b: {:?}, c: {:?}", 
        a.try_decrypt_trivial::<u128>(),
        b.try_decrypt_trivial::<u128>(),
        c.try_decrypt_trivial::<u128>(),
    );
    let tmp = a * b;
    
    println!("a * b = {:?}", tmp.try_decrypt_trivial::<u128>());

    tmp * c
}


fn main() {
    let (cks, sks) = generate_keys(ConfigBuilder::default().build());
    
    set_server_key(sks);
    
    let a = FheUint128::encrypt_trivial(1234u128);
    let b = FheUint128::encrypt_trivial(4567u128);
    let c = FheUint128::encrypt_trivial(89101112u128);
    
    // since all inputs are trivially encrypted, this is going to be
    // much faster
    let result = mul_all(&a, &b, &c);
}

a: Ok(1234), b: Ok(4567), c: Ok(89101112)
a * b = Ok(5635678)

a: Err(NotTrivialCiphertextError), b: Err(NotTrivialCiphertextError), c: Err(NotTrivialCiphertextError)
a * b = Err(NotTrivialCiphertextError)

use tfhe::prelude::*;
use tfhe::*;

pub fn main() {
    // Config and key generation
    let config = ConfigBuilder::default().build();

    let (cks, sks) = generate_keys(config);

    // Encryption
    let a = FheUint32::encrypt(42u32, &cks);
    let b = FheUint32::encrypt(16u32, &cks);

    // Set the server key
    set_server_key(sks);

    // Compute and get the PBS count for the 32 bits multiplication
    let c = &a * &b;
    let mul_32_count = get_pbs_count();

    // Reset the PBS count, and get the PBS count for a 32 bits bitwise AND
    reset_pbs_count();
    let d = &a & &b;
    let and_32_count = get_pbs_count();

    // Display the result
    println!("mul_32_count: {mul_32_count}");
    println!("and_32_count: {and_32_count}");
}

use std::ops::{Add, Mul};
use tfhe::prelude::*;
use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheUint32, FheUint64};

pub fn ex1<'a, FheType, ClearType>(ct: &'a FheType, pt: ClearType) -> FheType
    where
        &'a FheType: Add<ClearType, Output = FheType>,
{
    ct + pt
}

pub fn ex2<'a, FheType, ClearType>(a: &'a FheType, b: &'a FheType, pt: ClearType) -> FheType
    where
        &'a FheType: Mul<&'a FheType, Output = FheType>,
        FheType: Add<ClearType, Output = FheType>,
{
    (a * b) + pt
}

pub fn ex3<FheType, ClearType>(a: FheType, b: FheType, pt: ClearType) -> FheType
    where
            for<'a> &'a FheType: Add<&'a FheType, Output = FheType>,
            FheType: Add<FheType, Output = FheType> + Add<ClearType, Output = FheType>,
{
    let tmp = (&a + &b) + (&a + &b);
    tmp + pt
}

pub fn ex4<FheType, ClearType>(a: FheType, b: FheType, pt: ClearType) -> FheType
    where
        FheType: Clone + Add<FheType, Output = FheType> + Add<ClearType, Output = FheType>,
{
    let tmp = (a.clone() + b.clone()) + (a.clone() + b.clone());
    tmp + pt
}

fn main() {
    let config = ConfigBuilder::default()
        .build();

    let (client_key, server_keys) = generate_keys(config);

    set_server_key(server_keys);

    // Use FheUint32
    {
        let clear_a = 46546u32;
        let clear_b = 6469u32;
        let clear_c = 64u32;

        let a = FheUint32::try_encrypt(clear_a, &client_key).unwrap();
        let b = FheUint32::try_encrypt(clear_b, &client_key).unwrap();
        assert_eq!(
            ex1(&clear_a, clear_c),
            ex1(&a, clear_c).decrypt(&client_key)
        );
        assert_eq!(
            ex2(&clear_a, &clear_b, clear_c),
            ex2(&a, &b, clear_c).decrypt(&client_key)
        );
        assert_eq!(
            ex3(clear_a, clear_b, clear_c),
            ex3(a.clone(), b.clone(), clear_c).decrypt(&client_key)
        );
        assert_eq!(
            ex4(clear_a, clear_b, clear_c),
            ex4(a, b, clear_c).decrypt(&client_key)
        );
    }

    // Use FheUint64
    {
        let clear_a = 46544866u64;
        let clear_b = 6469446677u64;
        let clear_c = 647897u64;

        let a = FheUint64::try_encrypt(clear_a, &client_key).unwrap();
        let b = FheUint64::try_encrypt(clear_b, &client_key).unwrap();
        assert_eq!(
            ex1(&clear_a, clear_c),
            ex1(&a, clear_c).decrypt(&client_key)
        );
        assert_eq!(
            ex2(&clear_a, &clear_b, clear_c),
            ex2(&a, &b, clear_c).decrypt(&client_key)
        );
        assert_eq!(
            ex3(clear_a, clear_b, clear_c),
            ex3(a.clone(), b.clone(), clear_c).decrypt(&client_key)
        );
        assert_eq!(
            ex4(clear_a, clear_b, clear_c),
            ex4(a, b, clear_c).decrypt(&client_key)
        );
    }
}