number-theory 0.0.22

use crate::data::nt_data::LIOUVILLE_LUT;
use crate::data::primes::PRIMELIST;
use crate::traits::NumberTheory;
use crate::traits::Reduction;
use crate::result::NTResult;

use crate::arithmetic::inlineops::*;

impl NumberTheory for u16 {
    fn rng() -> Self {
        (rng_32() >> 16) as u16
    }
    
    fn residue(&self, ring: &Self) -> Self{
      if ring == &0{
        return *self
      }
        *self % *ring
    } 

    fn euclidean_div(&self, other: &Self) -> (Self, Self) {
        (*self / *other, *self % *other)
    }

    fn is_sprp(&self, base: &Self) -> bool {
        let pminus = *self - 1;
        let zeroes = pminus.trailing_zeros(); // Breaks number down to p= d*2^n -1
        let d = pminus >> zeroes;
        let mut x = base.exp_residue(&d, self);

        if x == 1 || x == pminus {
            return true;
        }

        for _ in 1..zeroes {
            x = x.quadratic_residue(self);
            if x == pminus {
                return true;
            }
        }
        false // otherwise it fails
    }

    fn is_prime(&self) -> bool {
        if self.reducible(){
            return (*self as u8).is_prime();
        }

        for i in PRIMELIST[..54].iter() {
            if self == i {
                return true;
            }
            if self % i == 0 {
                return false;
            }
        }
        true
    }

    fn prime_proof(&self) -> (bool, Vec<Self>) {
        if *self == 2 {
            return (true, vec![3]);
        }

        let x_minus = *self - 1;
        let fctrs = x_minus
            .factor()
            .iter()
            .step_by(2)
            .map(|y| *y)
            .collect::<Vec<Self>>();
        let mut certificate = vec![2];

        certificate.extend_from_slice(&fctrs[..]);

        loop {
            // loops until it has either been shown to be prime or composite

            let mut witness = Self::rng() % (self - 2) + 2;

            'witness: loop {
                if witness.gcd(self) == 1 {
                    break 'witness;
                }
                witness += 1;
            }

            if witness.exp_residue(&x_minus, self) != 1 {
                // If any witness says it's composite then it is
                certificate[0] = witness;

                return (false, certificate);
            }

            'inner: for (idx, j) in fctrs.iter().enumerate() {
                if witness.exp_residue(&((*self - 1) / j), self) == 1 {
                    break 'inner;
                }
                if idx == fctrs.len() - 1 {
                    certificate[0] = witness;

                    return (true, certificate);
                }
            }
        }
    }

    fn prime_list(&self, sup: &Self) -> Vec<Self> {
        let inf = std::cmp::min(*self, *sup);
        let mut hi = std::cmp::max(*self, *sup);

        if hi < u16::MAX {
            hi += 1;
        }

        let mut primevector = vec![];

        for i in inf..hi {
            if i.is_prime() {
                primevector.push(i)
            }
        }
        primevector
    }

    fn nth_prime(&self) -> NTResult<Self> {
        
        if *self < 54{
          return (*self as u8).nth_prime().map(|y| y as u16)
        }
        
        if *self > 6542{
          return NTResult::Overflow
        }
        
        let mut count = 54;
        let mut start = 256;
         
        loop {
            start += 1;

            if start == Self::MAX {
                return NTResult::Overflow;
            }
            if start.is_prime() {
                count += 1;
            }
            if count == *self {
                return NTResult::Eval(start);
            }
        }
    }

    fn pi(&self) -> Self {
       if *self < 255{
         return (*self as u8).pi() as Self
       }
        let mut count = 0u16;
        for i in 0u16..*self {
            if i.is_prime() {
                count += 1;
            }
        }
        count
    }

    fn prime_gen(k: u32) -> NTResult<Self> {
        if k > 16 {
            return NTResult::Overflow;
        }
        if k < 8 {
            return u8::prime_gen(k).map(|x| x as u16);
        }
        let form = (1 << (k - 1)) + 1;
        let bitlength = form - 2;

        loop {
            let q = u64::rng();
            let p = unsafe { std::mem::transmute::<u64, (u16, u16, u16, u16)>(q) };

            if ((p.0 & bitlength) | form).is_prime() {
                return NTResult::Eval((p.0 & bitlength) | form);
            }
            if ((p.1 & bitlength) | form).is_prime() {
                return NTResult::Eval((p.1 & bitlength) | form);
            }
            if ((p.2 & bitlength) | form).is_prime() {
                return NTResult::Eval((p.2 & bitlength) | form);
            }
            if ((p.3 & bitlength) | form).is_prime() {
                return NTResult::Eval((p.3 & bitlength) | form);
            }
        }
    }

    fn factor(&self) -> Vec<Self> {
      
      if self < &255{
        return (*self as u8).factor().iter().map(|x| *x as u16).collect::<Vec<u16>>()
      }
        let mut n = *self;
        let twofactors = n.trailing_zeros();
        n >>= twofactors;

        let mut factors: Vec<u16> = vec![];

        if twofactors > 0 {
            factors.push(2);
            factors.push(twofactors as u16);
        }

        for i in PRIMELIST[1..54].iter() {
            // strips out small primes
            if n % i == 0 {
                factors.push(*i);
                let mut count = 0u16;
                while n % i == 0 {
                    count += 1;
                    n /= i;
                }
                factors.push(count);
            }
        }
        if n > 1 {
            factors.push(n);
            factors.push(1);
        }
        factors
    }
    
    fn checked_factor(&self) -> NTResult<Vec<Self>>{    
     if *self == 0{
       return NTResult::InfiniteSet
     }
     if *self == 1{
       return NTResult::DNE
     }
     
     NTResult::Eval(self.factor())
    }



    fn sqrt(&self) -> (Self, Self) {
        ((*self as f64).sqrt() as Self, 0)
    }

    fn nth_root(&self, n: &Self) -> (Self, Self) {
        if *n > 15 {
            return (1, 0);
        }

        if *n == 1 {
            return (*n, 0);
        }

        if *n == 0 {
            panic!("No integer is a zeroth factor ")
        }
        (((*self as f64).powf((*n as f64).recip())) as Self, 0)
    }

    fn max_exp(&self) -> (Self,Self){
    
      for i in 1..16{
      let p = 16-i;
      let base = self.nth_root(&p).0;
         if base.pow(p as u32) == *self{
           return(base,p)
         }
      }
     return (*self,1)
    }    
    
    fn radical(&self) -> NTResult<Self> {
     if self.reducible(){
        return (*self as u8).radical().map(|kishum| kishum as u16)
     }
        self.checked_factor().map(|y| y.iter().step_by(2).product::<Self>())
    }

    fn k_free(&self, k: &Self) -> bool {
       if self.reducible(){
         return (*self as u8).k_free(&(*k as u8))
       }
        let factors = self.factor();
        for (idx, el) in factors.iter().enumerate() {
            if el == k && idx != 0 {
                return false;
            }
        }
        true
    }

    fn gcd(&self, other: &Self) -> Self {
        let mut a = *self;
        let mut b = *other;
        if b == 0 {
            return a;
        } else if a == 0 {
            return b;
        }

        let self_two_factor = a.trailing_zeros();
        let other_two_factor = b.trailing_zeros();
        let min_two_factor = std::cmp::min(self_two_factor, other_two_factor);
        a >>= self_two_factor;
        b >>= other_two_factor;
        loop {
            if b > a {
                std::mem::swap(&mut b, &mut a);
            }
            a -= b;

            if a == 0 {
                return b << min_two_factor;
            }
            a >>= a.trailing_zeros();
        }
    }

    fn extended_gcd(&self, other: &Self) -> (Self, Self, Self) {
        let mut gcd: u16 = *self;
        let mut new_r: u16 = *other;
        let mut bezout_1: u16 = 1;
        let mut new_s: u16 = 0;
        let mut bezout_2: u16 = 0;
        let mut new_t: u16 = 1;

        while new_r != 0 {
            let quotient = gcd / new_r;
            let mut temp: u16 = new_r;
            new_r = gcd - quotient * temp;
            gcd = temp;

            temp = new_s;
            if bezout_1 < quotient.product_residue(&temp, other) {
                // First bezout coefficient is computed over Z[q]
                new_s = *other - (quotient.product_residue(&temp, other) - bezout_1)
            } else {
                new_s = bezout_1.wrapping_sub(quotient * temp);
            }

            bezout_1 = temp;

            temp = new_t;
            if bezout_2 < quotient.product_residue(&temp, self) {
                // Second bezout coefficient is computed over Z[p]
                new_t = *self - (quotient.product_residue(&temp, self) - bezout_2)
            } else {
                new_t = bezout_2.wrapping_sub(quotient.product_residue(&temp, self));
            }
            bezout_2 = temp
        }
        (gcd, bezout_1, bezout_2)
    }

    fn lcm(&self, other: &Self) -> Self {
       if self == &0 && other == &0{
       return 0
       }
        let cf = self.gcd(other);
        (*self / cf) * (*other)
    }

    fn checked_lcm(&self, other: &Self) -> NTResult<Self> {
        if self == &0 && other == &0{
         return NTResult::Eval(0)
        }
        let cf = self.gcd(other);
        let (v, flag) = (*self / cf).overflowing_mul(*other);
        if flag {
            return NTResult::Overflow;
        }
        NTResult::Eval(v)
    }

    fn euler_totient(&self) -> Self {
       if self.reducible(){
         return (*self as u8).euler_totient() as u16
       }
        let factors = self.factor();
        let numerator = factors.iter().step_by(2).map(|x| x - 1u16).product::<u16>();
        let denominator = factors.iter().step_by(2).product::<u16>();
        (self / denominator) * numerator
    }

    fn jordan_totient(&self, k: &Self) -> NTResult<Self> {

        if *self < 2{
           return NTResult::Eval(*self)
        }
        
        let (coef, flag) = self.overflowing_pow(*k as u32);
        
        if flag {
            return NTResult::CompOverflow;
        }

        let mut denom = 1u16;
        let mut numer = 1u16;

        for i in self.factor().iter().step_by(2) {
            let pow = i.pow(*k as u32);

            denom *= pow;

            numer *= pow - 1;
        }

        NTResult::Eval(numer * (coef / denom))
    }
    
     fn carmichael_totient(&self) -> NTResult<Self>{
       if  self.reducible(){
         return (*self as u8).carmichael_totient().map(|x| x as u16)
       }
       let fctr = self.factor();
       let base = fctr.iter().step_by(2).map(|z| *z).collect::<Vec<Self>>();
       let mut result = 1;
      for (idx,el) in base.iter().enumerate(){
        if el == &2 && fctr[1] > 2{
         let phi = ((el.pow(fctr[2*idx+1] as u32) /el) *(el-1)) /2;
          result = result.lcm(&phi);
        }
       else{
         let phi =  (el.pow(fctr[2*idx+1] as u32)/el)*(el-1);
         result = result.lcm(&phi);
       } 
      }
     NTResult::Eval(result)
    }

    fn dedekind_psi(&self, k: &Self) -> NTResult<Self> {
        if *self == 0{
         return NTResult::Infinite
       }
        let (k2, flag) = k.overflowing_shl(1);
        if flag {
            return NTResult::Overflow;
        }
        self.jordan_totient(&k2).map(|y| y/self.jordan_totient(k).unwrap())
    }

   fn quadratic_residue(&self, n: &Self) -> Self {
        if n == &0 {
            return self.wrapping_mul(*self)
        }
        ((*self as u32 * *self as u32) % *n as u32) as Self
    }
    
    fn checked_quadratic_residue(&self, n: &Self) -> NTResult<Self> {
        if n == &0 {
            return NTResult::from_option(self.checked_mul(*self),NTResult::Overflow)
        }
        NTResult::Eval(((*self as u32 * *self as u32) % *n as u32) as Self)
    }
    

    fn product_residue(&self, other: &Self, n: &Self) -> Self {
        if n == &0 {
            return self.wrapping_mul(*other)
        }
        ((*self as u32 * *other as u32) % *n as u32) as Self
    }
    
    fn checked_product_residue(&self, other: &Self, n: &Self) -> NTResult<Self> {
        if n == &0 {
            return NTResult::from_option(self.checked_mul(*self),NTResult::Overflow)
        }
        NTResult::Eval(((*self as u32 * *other as u32) % *n as u32) as Self)
    }

    fn exp_residue(&self, p: &Self, modulus: &Self) -> Self {
        if modulus == &0 {
            return self.wrapping_pow(*p as u32)
        }

        if modulus.is_power_of_two() {
            return self.wrapping_pow((*p) as u32) & (*modulus - 1);
        }

        let mut z = 1u32;
        let mut base = *self as u32;
        let n = *modulus as u32;
        let mut pow = *p;
        if pow == 0 {
            return z as u16;
        }

        while pow > 1 {
            if pow % 2 == 0 {
                base = base * base % n;
                pow >>= 1;
            } else {
                z = base * z % n;
                base = base * base % n;
                pow = (pow - 1) >> 1;
            }
        }

        (base * z % n) as u16
    }

    fn checked_exp_residue(&self, p: &Self, modulus: &Self) -> NTResult<Self> {
        if modulus == &0 {
            match self.checked_pow(*p as u32) {
                Some(x) => return NTResult::Eval(x),
                None => return NTResult::Overflow,
            };
        }

        if modulus.is_power_of_two() {
            return NTResult::Eval(self.wrapping_pow((*p) as u32) & (*modulus - 1));
        }

        let mut z = 1u32;
        let mut base = *self as u32;
        let n = *modulus as u32;
        let mut pow = *p;
        if pow == 0 {
            return NTResult::Eval(z as u16);
        }

        while pow > 1 {
            if pow % 2 == 0 {
                base = base * base % n;
                pow >>= 1;
            } else {
                z = base * z % n;
                base = base * base % n;
                pow = (pow - 1) >> 1;
            }
        }

        NTResult::Eval((base * z % n) as u16)
    }

    fn legendre(&self, p: &Self) -> i8 {
        let k = self.exp_residue(&((*p - 1) >> 1), p);
        if k == 1 {
            return 1;
        };
        if k == *p - 1 {
            return -1;
        };
        0i8
    }

    fn checked_legendre(&self, p: &Self) -> NTResult<i8> {
        if p == &2 || !p.is_prime() {
            return NTResult::Undefined;
        }
        NTResult::Eval(self.legendre(p))
    }

    fn liouville(&self) -> i8 {
        if (LIOUVILLE_LUT[(*self / 64) as usize] >> (*self % 64)) & 1 == 1 {
            return -1;
        }
        return 1;
    }
    
    fn derivative(&self) -> NTResult<Self> {
      if *self < 94 {
         return (*self as u8).derivative().map(|y| y as Self)
       }
       let fctr = self.factor();
       let mut sum : u16 = 0;
       
     for i in 0..fctr.len() / 2 {
        match sum.checked_add(fctr[2 * i + 1] * (*self / fctr[2 * i])){
          Some(x) => sum = x,
          None => return NTResult::Overflow,
        }
      }
    NTResult::Eval(sum)
    }

    fn mangoldt(&self) -> f64 {
      if self.reducible(){
        return (*self as u8).mangoldt()
      }
      let base = self.max_exp().0;
       if base.is_prime(){
         return (base as f64).ln()
       }
        return 0f64
    }
    
    fn mobius(&self) -> i8 {
      if self.reducible(){
        return (*self as u8).mobius()
      }
      let fctr = self.factor();
      if fctr.len() == 1{ // if only one factor then return -1
         return -1
      }
      for i in 0..fctr.len()/2{
        if fctr[2*i+1]  > 1{
         return 0
        }
      }
      let fctrsum = fctr[1..].iter().step_by(2).sum::<Self>();
      if fctrsum&1 == 1{// if odd number of factors and square free
        return -1
      }
      return 1
    }

    fn jacobi(&self, k: &Self) -> i8 {
        let mut n = *self;
        let mut p = *k;
        let mut t = 1i8;
        n %= p;

        while n != 0 {
            let zeros = n.trailing_zeros();
            n >>= zeros;

            if (p % 8 == 3 || p % 8 == 5) && (zeros % 2 == 1) {
                t = -t
            }

            std::mem::swap(&mut n, &mut p);
            if n % 4 == 3 && p % 4 == 3 {
                t = -t;
            }
            n %= p;
        }

        if p == 1 {
            t
        } else {
            0
        }
    }

    fn checked_jacobi(&self, k: &Self) -> NTResult<i8> {
        if k > &0 && *k % 2 == 1 {
            return NTResult::Eval(self.jacobi(k));
        }
        NTResult::Undefined
    }
    
     fn kronecker(&self, k: &Self) -> i8{
     let x = self.clone();
     if *k == 0{
      if x == 1{
         return 1
      }
     return 0
    }
   if *k == 1{
      return 1
   }
   let fctr = k.factor();
   let mut start = 0;
   let mut res = 1;
   
   if fctr[0] ==  2{
     start = 1;
     if x&1 == 0{
     res = 0;
     }
     else if x % 8 == 1 || x % 8 == 7{
      res=1
     }
     else{
       res = (-1i8).pow(fctr[1] as u32)
     }
   }
   if fctr[0] == 2 && fctr.len() == 2{
     return res
   }
   for i in start..fctr.len()/2{
     res*=self.legendre(&fctr[2*i]).pow(fctr[2*i+1] as u32);
   }
   return res
}

    fn smooth(&self) -> NTResult<Self> {
       if *self == 0{
         return NTResult::Infinite
       }
       if *self == 1{
        return NTResult::DNE
       }
        let k = self.factor();
        NTResult::Eval(k[k.len() - 2])
    }
    
    

    fn is_smooth(&self, b: &Self) -> bool {
     match self.smooth(){
      NTResult::Infinite => false,
      NTResult::Eval(x) => x <= *b, 
      _=> false,
     }
   }
}