use crate::Uint;

impl<const BITS: usize, const LIMBS: usize> Uint<BITS, LIMBS> {
    #[must_use]
    pub fn checked_pow(self, exp: usize) -> Option<Self> {
        match self.overflowing_pow(exp) {
            (x, false) => Some(x),
            (_, true) => None,
        }
    }

    /// # Examples
    ///
    /// ```
    /// # use ruint::{Uint, uint};
    /// # uint!{
    /// assert_eq!(36_U64.overflowing_pow(12), (0x41c21cb8e1000000_U64, false));
    /// assert_eq!(36_U64.overflowing_pow(13), (0x3f4c09ffa4000000_U64, true));
    /// assert_eq!(36_U68.overflowing_pow(13), (0x093f4c09ffa4000000_U68, false));
    /// assert_eq!(16_U65.overflowing_pow(32), (0_U65, true));
    /// # }
    /// ```
    /// Small cases:
    /// ```
    /// # use ruint::{Uint, uint};
    /// # uint!{
    /// assert_eq!(0_U0.overflowing_pow(0), (0_U0, false));
    /// assert_eq!(0_U1.overflowing_pow(0), (1_U1, false));
    /// assert_eq!(0_U1.overflowing_pow(1), (0_U1, false));
    /// assert_eq!(1_U1.overflowing_pow(0), (1_U1, false));
    /// assert_eq!(1_U1.overflowing_pow(1), (1_U1, false));
    /// # }
    /// ```
    #[must_use]
    pub fn overflowing_pow(mut self, mut exp: usize) -> (Self, bool) {
        if BITS == 0 {
            return (self, false);
        }

        // Exponentiation by squaring
        let mut overflow = false;
        let mut base_overflow = false;
        let mut result = Self::from(1);
        while exp > 0 {
            // Multiply by base
            if exp & 1 == 1 {
                let (r, o) = result.overflowing_mul(self);
                result = r;
                overflow |= o | base_overflow;
            }

            // Square base
            let (s, o) = self.overflowing_mul(self);
            self = s;
            base_overflow |= o;
            exp >>= 1;
        }
        (result, overflow)
    }

    #[must_use]
    pub fn pow(self, exp: usize) -> Self {
        self.wrapping_pow(exp)
    }

    #[must_use]
    pub fn saturating_pow(self, exp: usize) -> Self {
        match self.overflowing_pow(exp) {
            (x, false) => x,
            (_, true) => Self::MAX,
        }
    }

    #[must_use]
    pub fn wrapping_pow(self, exp: usize) -> Self {
        self.overflowing_pow(exp).0
    }

    /// Construct from double precision binary logarithm.
    ///
    /// # Examples
    ///
    /// ```
    /// # use ruint::{Uint, uint, aliases::*};
    /// # uint!{
    /// assert_eq!(U64::approx_pow2(-2.0), Some(0_U64));
    /// assert_eq!(U64::approx_pow2(-1.0), Some(1_U64));
    /// assert_eq!(U64::approx_pow2(0.0), Some(1_U64));
    /// assert_eq!(U64::approx_pow2(1.0), Some(2_U64));
    /// assert_eq!(U64::approx_pow2(1.6), Some(3_U64));
    /// assert_eq!(U64::approx_pow2(2.0), Some(4_U64));
    /// assert_eq!(U64::approx_pow2(64.0), None);
    /// assert_eq!(U64::approx_pow2(10.385), Some(1337_U64));
    /// # }
    /// ```
    #[must_use]
    pub fn approx_pow2(exp: f64) -> Option<Self> {
        const LN2_1P5: f64 = 0.584_962_500_721_156_2_f64;
        const EXP2_63: f64 = 9_223_372_036_854_775_808_f64;

        // FEATURE: Round negative to zero.
        #[allow(clippy::cast_precision_loss)] // Self::BITS ~< 2^52 and so fits f64.
        if exp < LN2_1P5 {
            if exp < -1.0 {
                return Some(Self::ZERO);
            }
            return Self::try_from(1).ok();
        }
        #[allow(clippy::cast_precision_loss)]
        if exp > Self::BITS as f64 {
            return None;
        }

        // Since exp < BITS, it has an integer and fractional part.
        #[allow(clippy::cast_possible_truncation)] // exp <= BITS <= usize::MAX.
        #[allow(clippy::cast_sign_loss)] // exp >= 0.
        let shift = exp.trunc() as usize;
        let fract = exp.fract();

        // Compute the leading 64 bits
        // Since `fract < 1.0` we have `fract.exp2() < 2`, so we can rescale by
        // 2^63 and cast to u64.
        #[allow(clippy::cast_possible_truncation)] // fract < 1.0
        #[allow(clippy::cast_sign_loss)] // fract >= 0.
        let bits = (fract.exp2() * EXP2_63) as u64;
        // Note: If `fract` is zero this will result in `u64::MAX`.

        if shift >= 63 {
            // OPT: A dedicated function avoiding full-sized shift.
            Some(Self::try_from(bits).ok()?.checked_shl(shift - 63)?)
        } else {
            let shift = 63 - shift;
            // Divide `bits` by `2^shift`, rounding to nearest.
            let bits = (bits >> shift) + ((bits >> (shift - 1)) & 1);
            Self::try_from(bits).ok()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{const_for, nlimbs};
    use proptest::proptest;
    use std::iter::repeat;

    #[test]
    fn test_pow2_shl() {
        const_for!(BITS in NON_ZERO if (BITS >= 2) {
            const LIMBS: usize = nlimbs(BITS);
            type U = Uint<BITS, LIMBS>;
            proptest!(|(e in 0..=BITS+1)| {
                assert_eq!(U::from(2).pow(e), U::from(1) << e);
            });
        });
    }

    #[test]
    fn test_pow_product() {
        const_for!(BITS in NON_ZERO if (BITS >= 64) {
            const LIMBS: usize = nlimbs(BITS);
            type U = Uint<BITS, LIMBS>;
            proptest!(|(b in 2_u64..100, e in 0_usize..100)| {
                let b = U::from(b);
                let prod = repeat(b).take(e).product();
                assert_eq!(b.pow(e), prod);
            });
        });
    }
}

#[cfg(feature = "bench")]
#[doc(hidden)]
pub mod bench {
    use super::*;
    use crate::{const_for, nlimbs};
    use ::proptest::{
        arbitrary::Arbitrary,
        strategy::{Strategy, ValueTree},
        test_runner::TestRunner,
    };
    use criterion::{black_box, BatchSize, Criterion};

    pub fn group(criterion: &mut Criterion) {
        const_for!(BITS in BENCH {
            const LIMBS: usize = nlimbs(BITS);
            bench_pow::<BITS, LIMBS>(criterion);
        });
    }

    fn bench_pow<const BITS: usize, const LIMBS: usize>(criterion: &mut Criterion) {
        let input = (Uint::<BITS, LIMBS>::arbitrary(), usize::arbitrary());
        let mut runner = TestRunner::deterministic();
        criterion.bench_function(&format!("pow/{}", BITS), move |bencher| {
            bencher.iter_batched(
                || input.new_tree(&mut runner).unwrap().current(),
                |(b, e)| black_box(black_box(b).pow(black_box(e))),
                BatchSize::SmallInput,
            );
        });
    }
}