fixed-bigint 0.6.1

use crate::fixeduint::FixedUInt;
use crate::machineword::MachineWord;
use const_num_traits::{CheckedPow, Nct};
use num_integer::Roots;
use num_traits::{FromPrimitive, One, Zero};

impl<T: MachineWord, const N: usize> Roots for FixedUInt<T, N, Nct> {
    fn nth_root(&self, n: u32) -> Self {
        if n == 0 {
            panic!("nth_root: n must be non-zero");
        }

        if self.is_zero() {
            return Self::zero();
        }

        if self.is_one() || n == 1 {
            return *self;
        }

        let bit_len = self.bit_length();
        if n > bit_len {
            return Self::one();
        }

        // Initial guess: use ceiling(bit_len / n) for overestimate
        let initial_exp = bit_len.div_ceil(n).max(1);
        let mut x = Self::one() << (initial_exp as usize);

        // Constants using FromPrimitive
        let n_val = Self::from_u32(n).expect("n too large for FixedUInt");
        let n_minus_1 = Self::from_u32(n - 1).expect("n too large for FixedUInt");

        // Newton's method iteration. Power evaluations go through `checked_pow`:
        // an over-width `x^k` would panic (Nct multiply overflow), so overflow
        // is treated as "arbitrarily large" — the quotient `self / x^(n-1)` is 0
        // and the upper probe `(x+1)^n` counts as greater than `self`.
        loop {
            let quotient = match CheckedPow::checked_pow(x, n - 1) {
                Some(p) if p.is_zero() => break,
                Some(p) => *self / p,
                None => Self::zero(),
            };

            let numerator = x * n_minus_1 + quotient;
            let x_new = numerator / n_val;

            if x_new >= x {
                break;
            }

            x = x_new;
        }

        // Final adjustment to ensure r^n <= self < (r+1)^n. `self` fits the
        // width, so an over-width `x^n` (None) is necessarily greater.
        while CheckedPow::checked_pow(x, n).is_none_or(|p| p > *self) {
            x -= Self::one();
        }

        let mut x_plus_one = x + Self::one();
        while CheckedPow::checked_pow(x_plus_one, n).is_some_and(|p| p <= *self) {
            x += Self::one();
            x_plus_one = x + Self::one();
        }

        x
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use num_integer::Roots;
    use num_traits::One;

    #[test]
    fn test_sqrt_basic() {
        type TestInt = FixedUInt<u32, 2>;

        assert_eq!(TestInt::from(0u8).sqrt(), TestInt::from(0u8));
        assert_eq!(TestInt::from(1u8).sqrt(), TestInt::from(1u8));
        assert_eq!(TestInt::from(4u8).sqrt(), TestInt::from(2u8));
        assert_eq!(TestInt::from(9u8).sqrt(), TestInt::from(3u8));
        assert_eq!(TestInt::from(16u8).sqrt(), TestInt::from(4u8));
        assert_eq!(TestInt::from(25u8).sqrt(), TestInt::from(5u8));

        // Test non-perfect squares
        assert_eq!(TestInt::from(2u8).sqrt(), TestInt::from(1u8));
        assert_eq!(TestInt::from(3u8).sqrt(), TestInt::from(1u8));
        assert_eq!(TestInt::from(8u8).sqrt(), TestInt::from(2u8));
        assert_eq!(TestInt::from(15u8).sqrt(), TestInt::from(3u8));
        assert_eq!(TestInt::from(24u8).sqrt(), TestInt::from(4u8));
    }

    #[test]
    fn test_cbrt_basic() {
        type TestInt = FixedUInt<u32, 2>;

        assert_eq!(TestInt::from(0u8).cbrt(), TestInt::from(0u8));
        assert_eq!(TestInt::from(1u8).cbrt(), TestInt::from(1u8));
        assert_eq!(TestInt::from(8u8).cbrt(), TestInt::from(2u8));
        assert_eq!(TestInt::from(27u8).cbrt(), TestInt::from(3u8));
        assert_eq!(TestInt::from(64u8).cbrt(), TestInt::from(4u8));
        assert_eq!(TestInt::from(125u8).cbrt(), TestInt::from(5u8));

        // Test non-perfect cubes
        assert_eq!(TestInt::from(2u8).cbrt(), TestInt::from(1u8));
        assert_eq!(TestInt::from(7u8).cbrt(), TestInt::from(1u8));
        assert_eq!(TestInt::from(26u8).cbrt(), TestInt::from(2u8));
        assert_eq!(TestInt::from(63u8).cbrt(), TestInt::from(3u8));
    }

    #[test]
    fn test_nth_root_basic() {
        type TestInt = FixedUInt<u32, 2>;

        // Test 4th roots
        assert_eq!(TestInt::from(16u8).nth_root(4), TestInt::from(2u8));
        assert_eq!(TestInt::from(81u8).nth_root(4), TestInt::from(3u8));
        assert_eq!(TestInt::from(15u8).nth_root(4), TestInt::from(1u8));
        assert_eq!(TestInt::from(80u8).nth_root(4), TestInt::from(2u8));

        // Test 5th roots
        assert_eq!(TestInt::from(32u8).nth_root(5), TestInt::from(2u8));
        assert_eq!(TestInt::from(243u8).nth_root(5), TestInt::from(3u8));
        assert_eq!(TestInt::from(31u8).nth_root(5), TestInt::from(1u8));

        // Test n=1 (should return self)
        assert_eq!(TestInt::from(42u8).nth_root(1), TestInt::from(42u8));
        assert_eq!(TestInt::from(123u8).nth_root(1), TestInt::from(123u8));
    }

    #[test]
    fn test_nth_root_edge_cases() {
        type TestInt = FixedUInt<u32, 2>;

        // Test with 0 and 1
        assert_eq!(TestInt::from(0u8).nth_root(2), TestInt::from(0u8));
        assert_eq!(TestInt::from(1u8).nth_root(2), TestInt::from(1u8));
        assert_eq!(TestInt::from(0u8).nth_root(10), TestInt::from(0u8));
        assert_eq!(TestInt::from(1u8).nth_root(10), TestInt::from(1u8));

        // Test with large n (should return 1 for numbers > 1)
        assert_eq!(TestInt::from(2u8).nth_root(100), TestInt::from(1u8));
        assert_eq!(TestInt::from(1000u16).nth_root(50), TestInt::from(1u8));
    }

    #[test]
    #[should_panic(expected = "nth_root: n must be non-zero")]
    fn test_nth_root_zero_n() {
        let x = FixedUInt::<u32, 2>::from(16u8);
        x.nth_root(0);
    }

    #[test]
    fn test_root_correctness() {
        type TestInt = FixedUInt<u32, 2>;

        // Test that r^n <= x < (r+1)^n for various cases
        for x in 1..=100u16 {
            let x_int = TestInt::from(x);

            // Test square root
            let sqrt_x = x_int.sqrt();
            assert!(sqrt_x.pow(2) <= x_int);
            assert!((sqrt_x + TestInt::one()).pow(2) > x_int);

            // Test cube root
            let cbrt_x = x_int.cbrt();
            assert!(cbrt_x.pow(3) <= x_int);
            assert!((cbrt_x + TestInt::one()).pow(3) > x_int);

            // Test 4th root
            let root4_x = x_int.nth_root(4);
            assert!(root4_x.pow(4) <= x_int);
            assert!((root4_x + TestInt::one()).pow(4) > x_int);
        }
    }

    // sqrt of the max value when the type width == value width: the upper
    // probe (x+1)^2 = 2^BITS overflows and used to panic in the Nct multiply.
    // checked_pow now reads that overflow as "> self".
    #[test]
    fn sqrt_of_max_does_not_overflow_panic() {
        assert_eq!(
            FixedUInt::<u32, 1>::from(u32::MAX).sqrt(),
            FixedUInt::<u32, 1>::from(0xFFFFu32)
        );
        assert_eq!(
            FixedUInt::<u8, 1>::from(255u8).sqrt(),
            FixedUInt::<u8, 1>::from(15u8)
        );
        // A perfect square at the top of the range, and its neighbours.
        assert_eq!(
            FixedUInt::<u16, 1>::from(65025u16).sqrt(),
            FixedUInt::<u16, 1>::from(255u16)
        ); // 255^2
        assert_eq!(
            FixedUInt::<u16, 1>::from(65535u16).sqrt(),
            FixedUInt::<u16, 1>::from(255u16)
        );
    }
}