use crate::core_crypto::commons::noise_formulas::noise_simulation::traits::{
    AllocateCenteredBinaryShiftedStandardModSwitchResult,
    AllocateDriftTechniqueStandardModSwitchResult, AllocateLweBootstrapResult,
    AllocateLweKeyswitchResult, AllocateStandardModSwitchResult,
    CenteredBinaryShiftedStandardModSwitch, DriftTechniqueStandardModSwitch,
    LweClassicFftBootstrap, LweKeyswitch, ScalarMul, StandardModSwitch,
};
use crate::core_crypto::commons::noise_formulas::noise_simulation::{
    NoiseSimulationLweFourier128Bsk, NoiseSimulationLweFourierBsk,
};
use crate::core_crypto::gpu::algorithms::lwe_keyswitch::cuda_keyswitch_lwe_ciphertext;
use crate::core_crypto::gpu::cuda_modulus_switch_ciphertext;
use crate::core_crypto::gpu::glwe_ciphertext_list::CudaGlweCiphertextList;
use crate::core_crypto::gpu::lwe_bootstrap_key::CudaModulusSwitchNoiseReductionConfiguration;
use crate::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
use crate::core_crypto::gpu::vec::CudaVec;
use crate::core_crypto::prelude::*;
use crate::integer::gpu::ciphertext::info::CudaBlockInfo;
use crate::integer::gpu::ciphertext::CudaRadixCiphertext;
use crate::integer::gpu::server_key::radix::{CudaNoiseSquashingKey, CudaRadixCiphertextInfo};
use crate::integer::gpu::server_key::{
    CudaBootstrappingKey, CudaDynamicKeyswitchingKey, CudaServerKey,
};
use crate::integer::gpu::{
    cuda_centered_modulus_switch_64, unchecked_small_scalar_mul_integer_async, CudaStreams,
};
use crate::shortint::server_key::tests::noise_distribution::utils::noise_simulation::NoiseSimulationModulusSwitchConfig;

/// Side resources for CUDA operations in noise simulation
#[derive(Clone)]
pub struct CudaSideResources {
    pub streams: CudaStreams,
    pub block_info: CudaBlockInfo,
}

impl CudaSideResources {
    pub fn new(streams: &CudaStreams, block_info: CudaBlockInfo) -> Self {
        Self {
            streams: streams.clone(),
            block_info,
        }
    }
}

/// GPU version of DynLwe for CUDA operations
#[derive(Clone)]
pub enum CudaDynLwe {
    U32(CudaLweCiphertextList<u32>),
    U64(CudaLweCiphertextList<u64>),
    U128(CudaLweCiphertextList<u128>),
}

impl CudaDynLwe {
    pub fn lwe_dimension(&self) -> LweDimension {
        match self {
            Self::U32(cuda_lwe) => cuda_lwe.lwe_dimension(),
            Self::U64(cuda_lwe) => cuda_lwe.lwe_dimension(),
            Self::U128(cuda_lwe) => cuda_lwe.lwe_dimension(),
        }
    }

    pub fn raw_modulus_float(&self) -> f64 {
        match self {
            Self::U32(cuda_lwe) => cuda_lwe.ciphertext_modulus().raw_modulus_float(),
            Self::U64(cuda_lwe) => cuda_lwe.ciphertext_modulus().raw_modulus_float(),
            Self::U128(cuda_lwe) => cuda_lwe.ciphertext_modulus().raw_modulus_float(),
        }
    }

    pub fn as_lwe_32(&self) -> &CudaLweCiphertextList<u32> {
        match self {
            Self::U32(cuda_lwe) => cuda_lwe,
            Self::U64(_) => panic!("Tried getting a u64 CudaLweCiphertextList as u32."),
            Self::U128(_) => panic!("Tried getting a u128 CudaLweCiphertextList as u32."),
        }
    }

    pub fn as_lwe_64(&self) -> &CudaLweCiphertextList<u64> {
        match self {
            Self::U32(_) => panic!("Tried getting a u32 CudaLweCiphertextList as u64."),
            Self::U64(cuda_lwe) => cuda_lwe,
            Self::U128(_) => panic!("Tried getting a u128 CudaLweCiphertextList as u64."),
        }
    }

    pub fn as_lwe_128(&self) -> &CudaLweCiphertextList<u128> {
        match self {
            Self::U32(_) => panic!("Tried getting a u32 CudaLweCiphertextList as u128."),
            Self::U64(_) => panic!("Tried getting a u64 CudaLweCiphertextList as u128."),
            Self::U128(cuda_lwe) => cuda_lwe,
        }
    }

    pub fn into_lwe_32(self) -> CudaLweCiphertextList<u32> {
        match self {
            Self::U32(cuda_lwe) => cuda_lwe,
            Self::U64(_) => panic!("Tried converting a u64 CudaLweCiphertextList to u32."),
            Self::U128(_) => panic!("Tried converting a u128 CudaLweCiphertextList to u32."),
        }
    }

    pub fn into_lwe_64(self) -> CudaLweCiphertextList<u64> {
        match self {
            Self::U32(_) => panic!("Tried converting a u32 CudaLweCiphertextList to u64."),
            Self::U64(cuda_lwe) => cuda_lwe,
            Self::U128(_) => panic!("Tried converting a u128 CudaLweCiphertextList to u64."),
        }
    }

    pub fn into_lwe_128(self) -> CudaLweCiphertextList<u128> {
        match self {
            Self::U32(_) => panic!("Tried converting a u32 CudaLweCiphertextList to u128."),
            Self::U64(_) => panic!("Tried converting a u64 CudaLweCiphertextList to u128."),
            Self::U128(cuda_lwe) => cuda_lwe,
        }
    }
    pub fn as_ct_64_cpu(&self, streams: &CudaStreams) -> LweCiphertext<Vec<u64>> {
        match self {
            Self::U32(_) => panic!("Tried getting a u32 CudaLweCiphertextList as u64."),
            Self::U64(_cuda_lwe) => {
                let cpu_lwe_list = self.as_lwe_64().to_lwe_ciphertext_list(streams);
                LweCiphertext::from_container(
                    cpu_lwe_list.clone().into_container(),
                    cpu_lwe_list.ciphertext_modulus(),
                )
            }
            Self::U128(_) => panic!("Tried getting a u128 CudaLweCiphertextList as u64."),
        }
    }

    pub fn from_lwe_32(cuda_lwe: CudaLweCiphertextList<u32>) -> Self {
        Self::U32(cuda_lwe)
    }

    pub fn from_lwe_64(cuda_lwe: CudaLweCiphertextList<u64>) -> Self {
        Self::U64(cuda_lwe)
    }

    pub fn from_lwe_128(cuda_lwe: CudaLweCiphertextList<u128>) -> Self {
        Self::U128(cuda_lwe)
    }
}

impl ScalarMul<u64> for CudaDynLwe {
    type Output = Self;
    type SideResources = CudaSideResources;

    fn scalar_mul(&self, scalar: u64, side_resources: &mut Self::SideResources) -> Self::Output {
        match self {
            Self::U32(_cuda_lwe) => {
                panic!("U32 scalar mul not implemented for CudaDynLwe - only U64 is supported")
            }
            Self::U64(cuda_lwe) => {
                // Use the block info from side_resources for proper modulus values
                let mut cuda_radix = CudaRadixCiphertext::new(
                    cuda_lwe.duplicate(&side_resources.streams),
                    CudaRadixCiphertextInfo {
                        blocks: vec![side_resources.block_info],
                    },
                );
                unsafe {
                    unchecked_small_scalar_mul_integer_async(
                        &side_resources.streams,
                        &mut cuda_radix,
                        scalar,
                        side_resources.block_info.message_modulus,
                        side_resources.block_info.carry_modulus,
                    );
                    side_resources.streams.synchronize();
                }

                Self::U64(cuda_radix.d_blocks)
            }
            Self::U128(_cuda_lwe) => {
                panic!("U128 scalar mul not implemented for CudaDynLwe - only U64 is supported")
            }
        }
    }
}

// Extensions for NoiseSimulationLweFourierBsk to support GPU operations
impl NoiseSimulationLweFourierBsk {
    pub fn matches_actual_bsk_gpu(&self, lwe_bsk: &CudaBootstrappingKey<u64>) -> bool {
        let input_lwe_dimension = self.input_lwe_dimension();
        let glwe_size = self.output_glwe_size();
        let polynomial_size = self.output_polynomial_size();
        let decomp_base_log = self.decomp_base_log();
        let decomp_level_count = self.decomp_level_count();

        match lwe_bsk {
            CudaBootstrappingKey::Classic(cuda_bsk) => {
                let bsk_input_lwe_dimension = cuda_bsk.input_lwe_dimension();
                let bsk_glwe_size = cuda_bsk.glwe_dimension().to_glwe_size();
                let bsk_polynomial_size = cuda_bsk.polynomial_size();
                let bsk_decomp_base_log = cuda_bsk.decomp_base_log();
                let bsk_decomp_level_count = cuda_bsk.decomp_level_count();

                input_lwe_dimension == bsk_input_lwe_dimension
                    && glwe_size == bsk_glwe_size
                    && polynomial_size == bsk_polynomial_size
                    && decomp_base_log == bsk_decomp_base_log
                    && decomp_level_count == bsk_decomp_level_count
            }
            CudaBootstrappingKey::MultiBit(cuda_mb_bsk) => {
                let bsk_input_lwe_dimension = cuda_mb_bsk.input_lwe_dimension();
                let bsk_glwe_size = cuda_mb_bsk.glwe_dimension().to_glwe_size();
                let bsk_polynomial_size = cuda_mb_bsk.polynomial_size();
                let bsk_decomp_base_log = cuda_mb_bsk.decomp_base_log();
                let bsk_decomp_level_count = cuda_mb_bsk.decomp_level_count();

                input_lwe_dimension == bsk_input_lwe_dimension
                    && glwe_size == bsk_glwe_size
                    && polynomial_size == bsk_polynomial_size
                    && decomp_base_log == bsk_decomp_base_log
                    && decomp_level_count == bsk_decomp_level_count
            }
        }
    }
}

// Extensions for NoiseSimulationLweFourier128Bsk to support GPU operations (for u128 noise
// squashing)
impl NoiseSimulationLweFourier128Bsk {
    pub fn matches_actual_bsk_gpu(&self, lwe_bsk: &CudaBootstrappingKey<u128>) -> bool {
        let input_lwe_dimension = self.input_lwe_dimension();
        let glwe_size = self.output_glwe_size();
        let polynomial_size = self.output_polynomial_size();
        let decomp_base_log = self.decomp_base_log();
        let decomp_level_count = self.decomp_level_count();

        match lwe_bsk {
            CudaBootstrappingKey::Classic(cuda_bsk) => {
                let bsk_input_lwe_dimension = cuda_bsk.input_lwe_dimension();
                let bsk_glwe_size = cuda_bsk.glwe_dimension().to_glwe_size();
                let bsk_polynomial_size = cuda_bsk.polynomial_size();
                let bsk_decomp_base_log = cuda_bsk.decomp_base_log();
                let bsk_decomp_level_count = cuda_bsk.decomp_level_count();

                input_lwe_dimension == bsk_input_lwe_dimension
                    && glwe_size == bsk_glwe_size
                    && polynomial_size == bsk_polynomial_size
                    && decomp_base_log == bsk_decomp_base_log
                    && decomp_level_count == bsk_decomp_level_count
            }
            CudaBootstrappingKey::MultiBit(cuda_mb_bsk) => {
                let bsk_input_lwe_dimension = cuda_mb_bsk.input_lwe_dimension();
                let bsk_glwe_size = cuda_mb_bsk.glwe_dimension().to_glwe_size();
                let bsk_polynomial_size = cuda_mb_bsk.polynomial_size();
                let bsk_decomp_base_log = cuda_mb_bsk.decomp_base_log();
                let bsk_decomp_level_count = cuda_mb_bsk.decomp_level_count();

                input_lwe_dimension == bsk_input_lwe_dimension
                    && glwe_size == bsk_glwe_size
                    && polynomial_size == bsk_polynomial_size
                    && decomp_base_log == bsk_decomp_base_log
                    && decomp_level_count == bsk_decomp_level_count
            }
        }
    }
}

impl AllocateStandardModSwitchResult for CudaDynLwe {
    type Output = Self;
    type SideResources = CudaSideResources;

    fn allocate_standard_mod_switch_result(
        &self,
        side_resources: &mut Self::SideResources,
    ) -> Self::Output {
        match self {
            Self::U32(cuda_lwe) => {
                let new_cuda_lwe = CudaLweCiphertextList::new(
                    cuda_lwe.lwe_dimension(),
                    cuda_lwe.lwe_ciphertext_count(),
                    cuda_lwe.ciphertext_modulus(),
                    &side_resources.streams,
                );
                Self::U32(new_cuda_lwe)
            }
            Self::U64(cuda_lwe) => {
                let new_cuda_lwe = CudaLweCiphertextList::new(
                    cuda_lwe.lwe_dimension(),
                    cuda_lwe.lwe_ciphertext_count(),
                    cuda_lwe.ciphertext_modulus(),
                    &side_resources.streams,
                );
                Self::U64(new_cuda_lwe)
            }
            Self::U128(cuda_lwe) => {
                let new_cuda_lwe = CudaLweCiphertextList::new(
                    cuda_lwe.lwe_dimension(),
                    cuda_lwe.lwe_ciphertext_count(),
                    cuda_lwe.ciphertext_modulus(),
                    &side_resources.streams,
                );
                Self::U128(new_cuda_lwe)
            }
        }
    }
}

impl StandardModSwitch<Self> for CudaDynLwe {
    type SideResources = CudaSideResources;

    fn standard_mod_switch(
        &self,
        output_modulus_log: CiphertextModulusLog,
        output: &mut Self,
        side_resources: &mut Self::SideResources,
    ) {
        match (self, output) {
            (Self::U32(_input), Self::U32(_output_cuda_lwe)) => {
                panic!("U32 modulus switch not implemented for CudaDynLwe - only U64 is supported");
            }
            (Self::U64(input), Self::U64(output_cuda_lwe)) => {
                let internal_output = input.duplicate(&side_resources.streams);
                cuda_modulus_switch_ciphertext(
                    &mut output_cuda_lwe.0.d_vec,
                    output_modulus_log.0 as u32,
                    &side_resources.streams,
                );
                let mut cpu_lwe = internal_output.to_lwe_ciphertext_list(&side_resources.streams);
                let shift_to_map_to_native = u64::BITS - output_modulus_log.0 as u32;
                for val in cpu_lwe.as_mut_view().into_container().iter_mut() {
                    *val <<= shift_to_map_to_native;
                }
                let d_after_ms = CudaLweCiphertextList::from_lwe_ciphertext_list(
                    &cpu_lwe,
                    &side_resources.streams,
                );

                *output_cuda_lwe = d_after_ms;
            }
            (Self::U128(_input), Self::U128(_output_cuda_lwe)) => {
                panic!("U128 modulus switch not implemented for CudaDynLwe - only U64 is supported")
            }
            _ => panic!("Inconsistent inputs/outputs for CudaDynLwe StandardModSwitch"),
        }
    }
}

impl AllocateCenteredBinaryShiftedStandardModSwitchResult for CudaDynLwe {
    type Output = Self;
    type SideResources = CudaSideResources;

    fn allocate_centered_binary_shifted_standard_mod_switch_result(
        &self,
        side_resources: &mut Self::SideResources,
    ) -> Self::Output {
        self.allocate_standard_mod_switch_result(side_resources)
    }
}

impl CenteredBinaryShiftedStandardModSwitch<Self> for CudaDynLwe {
    type SideResources = CudaSideResources;

    fn centered_binary_shifted_and_standard_mod_switch(
        &self,
        output_modulus_log: CiphertextModulusLog,
        output: &mut Self,
        side_resources: &mut Self::SideResources,
    ) {
        match (self, output) {
            (Self::U32(_input), Self::U32(_output_cuda_lwe)) => {
                panic!("U32 centered binary shifted modulus switch not implemented for CudaDynLwe - only U64 is supported")
            }
            (Self::U64(input), Self::U64(output_cuda_lwe)) => unsafe {
                let mut internal_output = output_cuda_lwe.duplicate(&side_resources.streams);
                cuda_centered_modulus_switch_64(
                    side_resources.streams.ptr[0],
                    0u32,
                    internal_output.0.d_vec.as_mut_c_ptr(0),
                    input.0.d_vec.as_c_ptr(0),
                    input.lwe_dimension().0 as u32,
                    output_modulus_log.0 as u32,
                );
                side_resources.streams.synchronize();
                let cpu_lwe = internal_output.into_lwe_ciphertext(&side_resources.streams);
                let mut cpu_ct = LweCiphertext::from_container(
                    cpu_lwe.clone().into_container(),
                    cpu_lwe.ciphertext_modulus(),
                );
                let shift_to_map_to_native = u64::BITS - output_modulus_log.0 as u32;
                for val in cpu_ct.as_mut() {
                    *val <<= shift_to_map_to_native;
                }
                let d_after_ms =
                    CudaLweCiphertextList::from_lwe_ciphertext(&cpu_ct, &side_resources.streams);
                *output_cuda_lwe = d_after_ms;
            },
            (Self::U128(_input), Self::U128(_output_cuda_lwe)) => {
                panic!("U128 centered binary shifted modulus switch not implemented for CudaDynLwe - only U64 is supported")
            }
            _ => panic!("Inconsistent inputs/outputs for CudaDynLwe StandardModSwitch"),
        }
    }
}

impl DriftTechniqueStandardModSwitch<Self, Self, Self> for CudaDynLwe {
    type SideResources = CudaSideResources;

    fn drift_technique_and_standard_mod_switch(
        &self,
        _output_modulus_log: CiphertextModulusLog,
        _input: &Self,
        _after_drift_technique: &mut Self,
        _after_mod_switch: &mut Self,
        _side_resources: &mut Self::SideResources,
    ) {
        panic!("Drift technique is being deprecated, use other flavors of mod switch instead")
    }
}

impl AllocateLweKeyswitchResult for CudaServerKey {
    type Output = CudaDynLwe;
    type SideResources = CudaSideResources;

    fn allocate_lwe_keyswitch_result(
        &self,
        side_resources: &mut Self::SideResources,
    ) -> Self::Output {
        let output_lwe_dimension = match &self.key_switching_key {
            CudaDynamicKeyswitchingKey::Standard(std_key) => {
                std_key.output_key_lwe_size().to_lwe_dimension()
            }
            CudaDynamicKeyswitchingKey::KeySwitch32(ks32_key) => {
                ks32_key.output_key_lwe_size().to_lwe_dimension()
            }
        };
        let lwe_ciphertext_count = LweCiphertextCount(1);
        let ciphertext_modulus = self.ciphertext_modulus;

        let cuda_lwe = CudaLweCiphertextList::new(
            output_lwe_dimension,
            lwe_ciphertext_count,
            ciphertext_modulus,
            &side_resources.streams,
        );
        CudaDynLwe::U64(cuda_lwe)
    }
}

impl LweKeyswitch<CudaDynLwe, CudaDynLwe> for CudaServerKey {
    type SideResources = CudaSideResources;

    fn lwe_keyswitch(
        &self,
        input: &CudaDynLwe,
        output: &mut CudaDynLwe,
        side_resources: &mut Self::SideResources,
    ) {
        match (input, output) {
            (CudaDynLwe::U64(input_cuda_lwe), CudaDynLwe::U32(output_cuda_lwe)) => {
                let CudaDynamicKeyswitchingKey::KeySwitch32(computing_ks_key) =
                    &self.key_switching_key
                else {
                    panic!("Expecting 32b KSK in Cuda noise simulation tests when LWE is 32b");
                };

                let input_indexes = CudaVec::new(1, &side_resources.streams, 0);
                let output_indexes = CudaVec::new(1, &side_resources.streams, 0);

                cuda_keyswitch_lwe_ciphertext(
                    computing_ks_key,
                    input_cuda_lwe,
                    output_cuda_lwe,
                    &input_indexes,
                    &output_indexes,
                    false,
                    &side_resources.streams,
                    false,
                );
            }
            (CudaDynLwe::U64(input_cuda_lwe), CudaDynLwe::U64(output_cuda_lwe)) => {
                let CudaDynamicKeyswitchingKey::Standard(computing_ks_key) =
                    &self.key_switching_key
                else {
                    panic!("Expecting 64b KSK in Cuda noise simulation tests when LWE is 64b");
                };

                let input_indexes = CudaVec::new(1, &side_resources.streams, 0);
                let output_indexes = CudaVec::new(1, &side_resources.streams, 0);

                cuda_keyswitch_lwe_ciphertext(
                    computing_ks_key,
                    input_cuda_lwe,
                    output_cuda_lwe,
                    &input_indexes,
                    &output_indexes,
                    false,
                    &side_resources.streams,
                    false,
                );
            }
            (CudaDynLwe::U32(_), CudaDynLwe::U32(_)) => {
                panic!("U32 keyswitch not implemented for CudaServerKey - only U64 is supported");
            }
            (CudaDynLwe::U128(_), CudaDynLwe::U128(_)) => {
                panic!("U128 keyswitch not implemented for CudaServerKey - only U64 is supported");
            }
            _ => panic!("Inconsistent input/output types for CudaDynLwe keyswitch"),
        }
    }
}

impl AllocateDriftTechniqueStandardModSwitchResult for CudaServerKey {
    type AfterDriftOutput = CudaDynLwe;
    type AfterMsOutput = CudaDynLwe;
    type SideResources = CudaSideResources;

    fn allocate_drift_technique_standard_mod_switch_result(
        &self,
        _side_resources: &mut Self::SideResources,
    ) -> (Self::AfterDriftOutput, Self::AfterMsOutput) {
        panic!("Drift technique is being deprecated, use other flavors of mod switch instead")
    }
}

impl DriftTechniqueStandardModSwitch<CudaDynLwe, CudaDynLwe, CudaDynLwe> for CudaServerKey {
    type SideResources = CudaSideResources;

    fn drift_technique_and_standard_mod_switch(
        &self,
        _output_modulus_log: CiphertextModulusLog,
        _input: &CudaDynLwe,
        _after_drift_technique: &mut CudaDynLwe,
        _after_mod_switch: &mut CudaDynLwe,
        _side_resources: &mut Self::SideResources,
    ) {
        panic!("Drift technique is being deprecated, use other flavors of mod switch instead")
    }
}

impl CudaServerKey {
    pub fn br_input_modulus_log(&self) -> CiphertextModulusLog {
        match &self.bootstrapping_key {
            CudaBootstrappingKey::Classic(bsk) => {
                bsk.polynomial_size().to_blind_rotation_input_modulus_log()
            }
            CudaBootstrappingKey::MultiBit(mb_bsk) => mb_bsk
                .polynomial_size()
                .to_blind_rotation_input_modulus_log(),
        }
    }
    pub fn noise_simulation_modulus_switch_config(
        &self,
    ) -> NoiseSimulationModulusSwitchConfig<&Self> {
        match &self.bootstrapping_key {
            CudaBootstrappingKey::Classic(bsk) => match &bsk.ms_noise_reduction_configuration {
                None => NoiseSimulationModulusSwitchConfig::Standard,
                Some(CudaModulusSwitchNoiseReductionConfiguration::Centered) => {
                    NoiseSimulationModulusSwitchConfig::CenteredMeanNoiseReduction
                }
            },
            CudaBootstrappingKey::MultiBit(_) => {
                todo!()
            }
        }
    }
}

impl CudaNoiseSquashingKey {
    pub fn noise_simulation_modulus_switch_config(
        &self,
    ) -> NoiseSimulationModulusSwitchConfig<&Self> {
        match &self.bootstrapping_key {
            CudaBootstrappingKey::Classic(bsk) => match &bsk.ms_noise_reduction_configuration {
                None => NoiseSimulationModulusSwitchConfig::Standard,
                Some(CudaModulusSwitchNoiseReductionConfiguration::Centered) => {
                    NoiseSimulationModulusSwitchConfig::CenteredMeanNoiseReduction
                }
            },
            CudaBootstrappingKey::MultiBit(_) => {
                todo!()
            }
        }
    }
}

impl AllocateDriftTechniqueStandardModSwitchResult for CudaNoiseSquashingKey {
    type AfterDriftOutput = CudaDynLwe;
    type AfterMsOutput = CudaDynLwe;
    type SideResources = CudaSideResources;

    fn allocate_drift_technique_standard_mod_switch_result(
        &self,
        _side_resources: &mut Self::SideResources,
    ) -> (Self::AfterDriftOutput, Self::AfterMsOutput) {
        panic!("Drift technique is being deprecated, use other flavors of mod switch instead")
    }
}

impl DriftTechniqueStandardModSwitch<CudaDynLwe, CudaDynLwe, CudaDynLwe> for CudaNoiseSquashingKey {
    type SideResources = CudaSideResources;

    fn drift_technique_and_standard_mod_switch(
        &self,
        _output_modulus_log: CiphertextModulusLog,
        _input: &CudaDynLwe,
        _after_drift_technique: &mut CudaDynLwe,
        _after_mod_switch: &mut CudaDynLwe,
        _side_resources: &mut Self::SideResources,
    ) {
        panic!("Drift technique is being deprecated, use other flavors of mod switch instead")
    }
}

/// Implementation for CudaGlweCiphertextList<u64> to return CudaDynLwe (for test compatibility)
impl AllocateLweBootstrapResult for CudaGlweCiphertextList<u64> {
    type Output = CudaDynLwe;
    type SideResources = CudaSideResources;

    fn allocate_lwe_bootstrap_result(
        &self,
        side_resources: &mut Self::SideResources,
    ) -> Self::Output {
        // For PBS result, we allocate LWE ciphertexts wrapped in CudaDynLwe
        let lwe_dimension = self
            .glwe_dimension()
            .to_equivalent_lwe_dimension(self.polynomial_size());

        let cuda_lwe = CudaLweCiphertextList::new(
            lwe_dimension,
            LweCiphertextCount(1),
            self.ciphertext_modulus(),
            &side_resources.streams,
        );
        CudaDynLwe::U64(cuda_lwe)
    }
}

// Implement LweClassicFftBootstrap for CudaServerKey
impl LweClassicFftBootstrap<CudaDynLwe, CudaDynLwe, CudaGlweCiphertextList<u64>> for CudaServerKey {
    type SideResources = CudaSideResources;

    fn lwe_classic_fft_pbs(
        &self,
        input: &CudaDynLwe,
        output: &mut CudaDynLwe,
        accumulator: &crate::core_crypto::gpu::glwe_ciphertext_list::CudaGlweCiphertextList<u64>,
        side_resources: &mut Self::SideResources,
    ) {
        use crate::core_crypto::gpu::algorithms::lwe_programmable_bootstrapping::cuda_programmable_bootstrap_lwe_ciphertext;
        use crate::core_crypto::gpu::vec::CudaVec;
        use crate::integer::gpu::server_key::CudaBootstrappingKey;
        use crate::integer::gpu::CastInto;

        match (input, output) {
            (CudaDynLwe::U64(input_cuda_lwe), CudaDynLwe::U64(output_cuda_lwe)) => {
                // Create indexes for PBS
                let num_ct_blocks = 1;
                let lwe_indexes: Vec<u64> = (0..num_ct_blocks)
                    .map(<usize as CastInto<u64>>::cast_into)
                    .collect();
                let mut d_lut_vector_indexes =
                    unsafe { CudaVec::<u64>::new_async(num_ct_blocks, &side_resources.streams, 0) };
                let mut d_input_indexes =
                    unsafe { CudaVec::<u64>::new_async(num_ct_blocks, &side_resources.streams, 0) };
                let mut d_output_indexes =
                    unsafe { CudaVec::<u64>::new_async(num_ct_blocks, &side_resources.streams, 0) };

                unsafe {
                    d_lut_vector_indexes.copy_from_cpu_async(
                        &lwe_indexes,
                        &side_resources.streams,
                        0,
                    );
                    d_input_indexes.copy_from_cpu_async(&lwe_indexes, &side_resources.streams, 0);
                    d_output_indexes.copy_from_cpu_async(&lwe_indexes, &side_resources.streams, 0);
                }

                match &self.bootstrapping_key {
                    CudaBootstrappingKey::Classic(d_bsk) => {
                        cuda_programmable_bootstrap_lwe_ciphertext(
                            input_cuda_lwe,
                            output_cuda_lwe,
                            accumulator,
                            &d_lut_vector_indexes,
                            &d_output_indexes,
                            &d_input_indexes,
                            d_bsk,
                            &side_resources.streams,
                        );
                    }
                    CudaBootstrappingKey::MultiBit(_d_multibit_bsk) => {
                        panic!("Can not execute MultiBit PBS from classic FFT PBS implementation");
                    }
                }
            }
            _ => panic!("Only U64 PBS is supported for CudaServerKey"),
        }
    }
}

impl AllocateLweBootstrapResult for CudaGlweCiphertextList<u128> {
    type Output = CudaDynLwe;
    type SideResources = CudaSideResources;

    fn allocate_lwe_bootstrap_result(
        &self,
        side_resources: &mut Self::SideResources,
    ) -> Self::Output {
        let lwe_dimension = self
            .glwe_dimension()
            .to_equivalent_lwe_dimension(self.polynomial_size());

        let cuda_lwe = CudaLweCiphertextList::<u128>::new(
            lwe_dimension,
            LweCiphertextCount(1),
            self.ciphertext_modulus(),
            &side_resources.streams,
        );
        CudaDynLwe::U128(cuda_lwe)
    }
}
