use super::*;

use crate::core_crypto::commons::generators::{
    DeterministicSeeder, EncryptionRandomGenerator, SecretRandomGenerator,
};
use crate::core_crypto::commons::math::random::{ActivatedRandomGenerator, CompressionSeed};
use crate::core_crypto::commons::test_tools;

fn test_parallel_and_seeded_ggsw_encryption_equivalence<Scalar>(
    ciphertext_modulus: CiphertextModulus<Scalar>,
) where
    Scalar: UnsignedTorus + Sync + Send,
{
    // DISCLAIMER: these toy example parameters are not guaranteed to be secure or yield correct
    // computations
    // Define parameters for GgswCiphertext creation
    let glwe_size = GlweSize(2);
    let polynomial_size = PolynomialSize(1024);
    let decomp_base_log = DecompositionBaseLog(8);
    let decomp_level_count = DecompositionLevelCount(3);
    let glwe_modular_std_dev = StandardDev(0.00000000000000029403601535432533);

    // Create the PRNG
    let mut seeder = new_seeder();
    let seeder = seeder.as_mut();
    let main_seed = seeder.seed();
    let mut secret_generator =
        SecretRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed());

    const NB_TESTS: usize = 10;

    for _ in 0..NB_TESTS {
        // Create the GlweSecretKey
        let glwe_secret_key = allocate_and_generate_new_binary_glwe_secret_key(
            glwe_size.to_glwe_dimension(),
            polynomial_size,
            &mut secret_generator,
        );

        // Create the plaintext
        let encoded_msg: Scalar = test_tools::random_uint_between(Scalar::ZERO..Scalar::TWO.shl(2));
        let plaintext = Plaintext(encoded_msg);

        let compression_seed: CompressionSeed = seeder.seed().into();

        let mut ser_ggsw = GgswCiphertext::new(
            Scalar::ZERO,
            glwe_size,
            polynomial_size,
            decomp_base_log,
            decomp_level_count,
            ciphertext_modulus,
        );

        let mut deterministic_seeder =
            DeterministicSeeder::<ActivatedRandomGenerator>::new(main_seed);

        let mut encryption_generator = EncryptionRandomGenerator::<ActivatedRandomGenerator>::new(
            compression_seed.seed,
            &mut deterministic_seeder,
        );

        encrypt_constant_ggsw_ciphertext(
            &glwe_secret_key,
            &mut ser_ggsw,
            plaintext,
            glwe_modular_std_dev,
            &mut encryption_generator,
        );

        let mut par_ggsw = GgswCiphertext::new(
            Scalar::ZERO,
            glwe_size,
            polynomial_size,
            decomp_base_log,
            decomp_level_count,
            ciphertext_modulus,
        );

        let mut deterministic_seeder =
            DeterministicSeeder::<ActivatedRandomGenerator>::new(main_seed);

        let mut encryption_generator = EncryptionRandomGenerator::<ActivatedRandomGenerator>::new(
            compression_seed.seed,
            &mut deterministic_seeder,
        );

        par_encrypt_constant_ggsw_ciphertext(
            &glwe_secret_key,
            &mut par_ggsw,
            plaintext,
            glwe_modular_std_dev,
            &mut encryption_generator,
        );

        assert_eq!(ser_ggsw, par_ggsw);

        // Create a new GgswCiphertext
        let mut ser_seeded_ggsw = SeededGgswCiphertext::new(
            Scalar::ZERO,
            glwe_size,
            polynomial_size,
            decomp_base_log,
            decomp_level_count,
            compression_seed,
            ciphertext_modulus,
        );

        let mut deterministic_seeder =
            DeterministicSeeder::<ActivatedRandomGenerator>::new(main_seed);

        encrypt_constant_seeded_ggsw_ciphertext(
            &glwe_secret_key,
            &mut ser_seeded_ggsw,
            plaintext,
            glwe_modular_std_dev,
            &mut deterministic_seeder,
        );

        let mut par_seeded_ggsw = SeededGgswCiphertext::new(
            Scalar::ZERO,
            glwe_size,
            polynomial_size,
            decomp_base_log,
            decomp_level_count,
            compression_seed,
            ciphertext_modulus,
        );

        let mut deterministic_seeder =
            DeterministicSeeder::<ActivatedRandomGenerator>::new(main_seed);

        par_encrypt_constant_seeded_ggsw_ciphertext(
            &glwe_secret_key,
            &mut par_seeded_ggsw,
            plaintext,
            glwe_modular_std_dev,
            &mut deterministic_seeder,
        );

        assert_eq!(ser_seeded_ggsw, par_seeded_ggsw);

        let decompressed_ggsw = par_seeded_ggsw.decompress_into_ggsw_ciphertext();

        assert_eq!(ser_ggsw, decompressed_ggsw);
    }
}

#[test]
fn test_parallel_and_seeded_ggsw_encryption_equivalence_u32_native_mod() {
    test_parallel_and_seeded_ggsw_encryption_equivalence::<u32>(CiphertextModulus::new_native());
}

#[test]
fn test_parallel_and_seeded_ggsw_encryption_equivalence_u32_custom_mod() {
    test_parallel_and_seeded_ggsw_encryption_equivalence::<u32>(
        CiphertextModulus::try_new_power_of_2(31).unwrap(),
    );
}

#[test]
fn test_parallel_and_seeded_ggsw_encryption_equivalence_u64_native_mod() {
    test_parallel_and_seeded_ggsw_encryption_equivalence::<u64>(CiphertextModulus::new_native());
}

#[test]
fn test_parallel_and_seeded_ggsw_encryption_equivalence_u64_custom_mod() {
    test_parallel_and_seeded_ggsw_encryption_equivalence::<u64>(
        CiphertextModulus::try_new_power_of_2(63).unwrap(),
    );
}

fn ggsw_encrypt_decrypt_custom_mod<Scalar: UnsignedTorus>(params: TestParams<Scalar>) {
    let glwe_dimension = params.glwe_dimension;
    let polynomial_size = params.polynomial_size;
    let glwe_modular_std_dev = params.glwe_modular_std_dev;
    let ciphertext_modulus = params.ciphertext_modulus;
    let decomposition_base_log = params.pbs_base_log;
    let decomposition_level_count = params.pbs_level;

    let mut rsc = TestResources::new();

    const NB_TESTS: usize = 10;

    let mut msg = Scalar::ONE << decomposition_base_log.0;

    while msg != Scalar::ZERO {
        // We are going to go faster if the base log is big
        if msg > (Scalar::ONE << 4) {
            msg /= Scalar::TWO;
        } else {
            // Then we can scan all values
            msg = msg.wrapping_sub(Scalar::ONE);
        }
        for _ in 0..NB_TESTS {
            let glwe_sk = allocate_and_generate_new_binary_glwe_secret_key(
                glwe_dimension,
                polynomial_size,
                &mut rsc.secret_random_generator,
            );

            let mut ggsw = GgswCiphertext::new(
                Scalar::ZERO,
                glwe_dimension.to_glwe_size(),
                polynomial_size,
                decomposition_base_log,
                decomposition_level_count,
                ciphertext_modulus,
            );

            // GGSW constants are seen as plaintext, the encoding is done by the encryption itself
            let plaintext = Plaintext(msg);

            encrypt_constant_ggsw_ciphertext(
                &glwe_sk,
                &mut ggsw,
                plaintext,
                glwe_modular_std_dev,
                &mut rsc.encryption_random_generator,
            );

            assert!(check_content_respects_mod(&ggsw, ciphertext_modulus));

            let decoded = decrypt_constant_ggsw_ciphertext(&glwe_sk, &ggsw);

            assert!(decoded.0 == msg);
        }
    }
}

create_parametrized_test!(ggsw_encrypt_decrypt_custom_mod);

fn ggsw_par_encrypt_decrypt_custom_mod<Scalar: UnsignedTorus + Send + Sync>(
    params: TestParams<Scalar>,
) {
    let glwe_dimension = params.glwe_dimension;
    let polynomial_size = params.polynomial_size;
    let glwe_modular_std_dev = params.glwe_modular_std_dev;
    let ciphertext_modulus = params.ciphertext_modulus;
    let decomposition_base_log = params.pbs_base_log;
    let decomposition_level_count = params.pbs_level;

    let mut rsc = TestResources::new();

    const NB_TESTS: usize = 10;

    let mut msg = Scalar::ONE << decomposition_base_log.0;

    while msg != Scalar::ZERO {
        // We are going to go faster if the base log is big
        if msg > (Scalar::ONE << 4) {
            msg /= Scalar::TWO;
        } else {
            // Then we can scan all values
            msg = msg.wrapping_sub(Scalar::ONE);
        }
        for _ in 0..NB_TESTS {
            let glwe_sk = allocate_and_generate_new_binary_glwe_secret_key(
                glwe_dimension,
                polynomial_size,
                &mut rsc.secret_random_generator,
            );

            let mut ggsw = GgswCiphertext::new(
                Scalar::ZERO,
                glwe_dimension.to_glwe_size(),
                polynomial_size,
                decomposition_base_log,
                decomposition_level_count,
                ciphertext_modulus,
            );

            // GGSW constants are seen as plaintext, the encoding is done by the encryption itself
            let plaintext = Plaintext(msg);

            par_encrypt_constant_ggsw_ciphertext(
                &glwe_sk,
                &mut ggsw,
                plaintext,
                glwe_modular_std_dev,
                &mut rsc.encryption_random_generator,
            );

            assert!(check_content_respects_mod(&ggsw, ciphertext_modulus));

            let decoded = decrypt_constant_ggsw_ciphertext(&glwe_sk, &ggsw);

            assert!(decoded.0 == msg);
        }
    }
}

create_parametrized_test!(ggsw_par_encrypt_decrypt_custom_mod);

fn ggsw_seeded_encrypt_decrypt_custom_mod<Scalar: UnsignedTorus>(params: TestParams<Scalar>) {
    let glwe_dimension = params.glwe_dimension;
    let polynomial_size = params.polynomial_size;
    let glwe_modular_std_dev = params.glwe_modular_std_dev;
    let ciphertext_modulus = params.ciphertext_modulus;
    let decomposition_base_log = params.pbs_base_log;
    let decomposition_level_count = params.pbs_level;

    let mut rsc = TestResources::new();

    const NB_TESTS: usize = 10;

    let mut msg = Scalar::ONE << decomposition_base_log.0;

    while msg != Scalar::ZERO {
        // We are going to go faster if the base log is big
        if msg > (Scalar::ONE << 4) {
            msg /= Scalar::TWO;
        } else {
            // Then we can scan all values
            msg = msg.wrapping_sub(Scalar::ONE);
        }
        for _ in 0..NB_TESTS {
            let glwe_sk = allocate_and_generate_new_binary_glwe_secret_key(
                glwe_dimension,
                polynomial_size,
                &mut rsc.secret_random_generator,
            );

            let mut seeded_ggsw = SeededGgswCiphertext::new(
                Scalar::ZERO,
                glwe_dimension.to_glwe_size(),
                polynomial_size,
                decomposition_base_log,
                decomposition_level_count,
                rsc.seeder.seed().into(),
                ciphertext_modulus,
            );

            // GGSW constants are seen as plaintext, the encoding is done by the encryption itself
            let plaintext = Plaintext(msg);

            encrypt_constant_seeded_ggsw_ciphertext(
                &glwe_sk,
                &mut seeded_ggsw,
                plaintext,
                glwe_modular_std_dev,
                rsc.seeder.as_mut(),
            );

            let ggsw = seeded_ggsw.decompress_into_ggsw_ciphertext();

            assert!(check_content_respects_mod(&ggsw, ciphertext_modulus));

            let decoded = decrypt_constant_ggsw_ciphertext(&glwe_sk, &ggsw);

            assert!(decoded.0 == msg);
        }
    }
}

create_parametrized_test!(ggsw_seeded_encrypt_decrypt_custom_mod);

fn ggsw_seeded_par_encrypt_decrypt_custom_mod<Scalar: UnsignedTorus + Sync + Send>(
    params: TestParams<Scalar>,
) {
    let glwe_dimension = params.glwe_dimension;
    let polynomial_size = params.polynomial_size;
    let glwe_modular_std_dev = params.glwe_modular_std_dev;
    let ciphertext_modulus = params.ciphertext_modulus;
    let decomposition_base_log = params.pbs_base_log;
    let decomposition_level_count = params.pbs_level;

    let mut rsc = TestResources::new();

    const NB_TESTS: usize = 10;

    let mut msg = Scalar::ONE << decomposition_base_log.0;

    while msg != Scalar::ZERO {
        // We are going to go faster if the base log is big
        if msg > (Scalar::ONE << 4) {
            msg /= Scalar::TWO;
        } else {
            // Then we can scan all values
            msg = msg.wrapping_sub(Scalar::ONE);
        }
        for _ in 0..NB_TESTS {
            let glwe_sk = allocate_and_generate_new_binary_glwe_secret_key(
                glwe_dimension,
                polynomial_size,
                &mut rsc.secret_random_generator,
            );

            let mut seeded_ggsw = SeededGgswCiphertext::new(
                Scalar::ZERO,
                glwe_dimension.to_glwe_size(),
                polynomial_size,
                decomposition_base_log,
                decomposition_level_count,
                rsc.seeder.seed().into(),
                ciphertext_modulus,
            );

            // GGSW constants are seen as plaintext, the encoding is done by the encryption itself
            let plaintext = Plaintext(msg);

            par_encrypt_constant_seeded_ggsw_ciphertext(
                &glwe_sk,
                &mut seeded_ggsw,
                plaintext,
                glwe_modular_std_dev,
                rsc.seeder.as_mut(),
            );

            let ggsw = seeded_ggsw.decompress_into_ggsw_ciphertext();

            assert!(check_content_respects_mod(&ggsw, ciphertext_modulus));

            let decoded = decrypt_constant_ggsw_ciphertext(&glwe_sk, &ggsw);

            assert!(decoded.0 == msg);
        }
    }
}

create_parametrized_test!(ggsw_seeded_par_encrypt_decrypt_custom_mod);
