use crate::core_crypto::backward_compatibility::commons::math::random::TUniformVersions;

use super::*;
use crate::core_crypto::commons::dispersion::Variance;
use serde::{Deserialize, Serialize};
use tfhe_versionable::Versionize;

/// The distribution $TUniform(1, -2^b, 2^b)$ is defined as follows, any value in the interval
/// $\left[-2^b, 2^b\right]$ is selected with probability $\frac{1}{2^{b+1}}$, with the two end
/// points $-2^b$ and $2^b$ being selected with probability $\frac{1}{2^{b+2}}$.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Versionize)]
#[versionize(TUniformVersions)]
pub struct TUniform<T: UnsignedInteger> {
    bound_log2: u32,
    _phantom: std::marker::PhantomData<T>,
}

impl<T: UnsignedInteger> TUniform<T> {
    /// Construct a TUniform distribution see [`TUniform`] for behavior of randomly generated
    /// values.
    ///
    /// # Panics
    ///
    /// Panics if `bound_log2` is greater than the Scalar type number of bits minus two.
    ///
    /// The reason for this is that with a higher `bound_log2` it is impossible to distinguish
    /// between $-2^{bound\_log2}$ and $2^{bound\_log2}$ because of the 2's complement
    /// representation of integers.
    pub const fn new(bound_log2: u32) -> Self {
        match Self::try_new(bound_log2) {
            Ok(result) => result,
            Err(e) => panic!("{}", e),
        }
    }

    /// Construct a TUniform distribution see [`TUniform`] for behavior of randomly generated
    /// values.
    ///
    /// # Note
    ///
    /// Returns an `Err` variant if `bound_log2` is greater than the Scalar type number of bits
    /// minus two.
    ///
    /// The reason for this is that with a higher `bound_log2` it is impossible to distinguish
    /// between $-2^{bound\_log2}$ and $2^{bound\_log2}$ because of the 2's complement
    /// representation of integers.
    pub const fn try_new(bound_log2: u32) -> Result<Self, &'static str> {
        if (bound_log2 + 2) as usize > T::BITS {
            return Err("Cannot create TUniform: \
            bound_log2 + 2 is greater than the current type's bit width");
        }

        Ok(Self {
            bound_log2,
            _phantom: std::marker::PhantomData,
        })
    }

    pub const fn bound_log2(&self) -> u32 {
        self.bound_log2
    }

    pub(crate) fn distinct_value_count(self) -> T {
        (T::ONE << ((self.bound_log2 + 1) as usize)) + T::ONE
    }

    pub fn min_value_inclusive(&self) -> T::Signed {
        -(T::Signed::ONE << self.bound_log2 as usize)
    }

    pub fn max_value_inclusive(&self) -> T::Signed {
        T::Signed::ONE << self.bound_log2 as usize
    }

    pub fn variance(&self, modulus: f64) -> Variance {
        // (2^{2 * b + 1} + 1) / 6
        Variance(
            ((2.0f64.powi(2 * (self.bound_log2 as i32) + 1) + 1.0f64) / 6.0f64) * modulus.powi(-2),
        )
    }
}

macro_rules! implement_t_uniform_uint {
    ($T:ty) => {
        impl RandomGenerable<TUniform<$T>> for $T {
            type CustomModulus = $T;
            #[allow(unused)]
            fn generate_one<G: ByteRandomGenerator>(
                generator: &mut RandomGenerator<G>,
                TUniform { bound_log2, .. }: TUniform<$T>,
            ) -> Self {
                let mut buf = [0; std::mem::size_of::<$T>()];
                let required_bits = bound_log2 + 2;
                let required_bytes = required_bits.div_ceil(u8::BITS) as usize;
                let mod_mask = <$T>::MAX >> (<$T>::BITS - required_bits);

                // For small moduli compared to the native type allows to avoid wasting too much
                // random bytes generated by the CSPRNG.
                buf.iter_mut()
                    .take(required_bytes)
                    .for_each(|a| *a = generator.generate_next());
                // We use from_le_bytes as most platforms are low endian, this avoids endianness
                // issues
                let native_int_random = <$T>::from_le_bytes(buf);
                let mut candidate_for_random = native_int_random & mod_mask;
                let bit_b_p_1 = candidate_for_random & 1;
                candidate_for_random >>= 1;
                candidate_for_random = candidate_for_random.wrapping_add(bit_b_p_1);
                candidate_for_random.wrapping_sub(1 << bound_log2)
            }

            fn generate_one_custom_modulus<G: ByteRandomGenerator>(
                generator: &mut RandomGenerator<G>,
                distribution: TUniform<$T>,
                custom_modulus: Self::CustomModulus,
            ) -> Self {
                // We have to assert otherwise random generation can return invalid values
                let tuniform_unique_values = distribution.distinct_value_count();
                // for b = 1 we have 2^{b + 1} + 1 distinct values == 5
                // If q = 5, we have 5 representable values, which means all TUniform values can
                // be represented, 0 == 0, 1 == 1, 2 == 2, 3 == -2, 4 == -1
                // So we need tuniform_unique_values <= q
                assert!(
                    tuniform_unique_values <= custom_modulus,
                    "Tried to generate TUniform for a modulus {custom_modulus:?} \
                    that cannot represent all values for bound {}",
                    distribution.bound_log2()
                );

                let tuniform_native_mod = Self::generate_one(generator, distribution);
                if tuniform_native_mod.into_signed() >= 0 {
                    tuniform_native_mod
                } else {
                    // tuniform_native_mod is negative, so adding it, even in unsigned integer will
                    // work as arithmetic uses 2's complement
                    custom_modulus.wrapping_add(tuniform_native_mod)
                }
            }

            fn single_sample_success_probability(
                _distribution: TUniform<$T>,
                _modulus: Option<Self::CustomModulus>,
            ) -> f64 {
                // The modulus and parameters of the distribution do not impact generation success
                1.0
            }

            fn single_sample_required_random_byte_count(
                TUniform { bound_log2, .. }: TUniform<$T>,
                _modulus: Option<Self::CustomModulus>,
            ) -> usize {
                // The modulus does not impact the amount of byte required
                let required_bits = bound_log2 + 2;
                required_bits.div_ceil(u8::BITS) as usize
            }
        }
    };
}

implement_t_uniform_uint!(u8);
implement_t_uniform_uint!(u16);
implement_t_uniform_uint!(u32);
implement_t_uniform_uint!(u64);
implement_t_uniform_uint!(u128);
