# pUtils.py
# Andrew Davison, ad@coe.psu.ac.th, Nov. 2024


import math, random, time
# import timeout   # see timeout.py

usingSieve = True
try:
  import pyprimesieve as pps
  import pyprimes          
except ImportError:
  print("pyprimesieve module not found")
  usingSieve = False
  print("Using pyprimes.py instead")
  import pyprimes as pps
 

GAMMA = 0.5772156649   # Euler-Mascheroni constant


# ---------- pyprimesieve decorators ----------
'''
pyprimesieve functions: 
  https://github.com/jaredks/pyprimesieve
  * primes(n): List of prime numbers up to n.
  * primes(start, n): 
    List of prime numbers from start up to n.
  * primes_sum(n):
  * primes_sum(start, n): 
  * primes_nth(n): 
    The nth prime number.
    e.g. pps.primes_nth(1)
    2
  * factorize(n): List of tuples in the form of 
    (prime, power) for the prime factorization of n.
'''

def primes(*argv, quiet=True):
  # return a list of prime numbers up to n, 
  # optionally with a starting value
  try:
    return pps.primes(*argv)
  except OverflowError as e:
    if usingSieve:
      if not quiet:
        print(e, "for", n, "so using pyprimes")
      return pyprimes.primes(*argv)
    else:
      print(e)


def primes_sum(*argv, quiet=True):
  # return the summation of prime numbers up to n, 
  # optionally with a starting value
  try:
    return pps.primes_sum(*argv)
  except OverflowError as e:
    if usingSieve:
      if not quiet:
        print(e, "for", n, "so using pyprimes")
      return pyprimes.primes_sum(*argv)
    else:
      print(e)


def prime(n, quiet=True):
  return primes_nth(n, quiet)
def primes_nth(n, quiet=True):
  # return the nth prime number
  try:
    return pps.primes_nth(n)
  except OverflowError as e:
    if usingSieve:
      if not quiet:
        print(e, "for", n, "so using pyprimes")
      return pyprimes.primes_nth(n)
    else:
      print(e)



def factorize(n, quiet=True):
  # factorize(120) --> [(2, 3), (3, 1), (5, 1)]
  try:
    return pps.factorize(n)
  except OverflowError as e:
    if usingSieve:
      if not quiet:
        print(e, "for", n, "so using pyprimes")
      return pyprimes.factorize(n)
    else:
      print(e)


# --------------- extra prime funcs ----------

def primesNum(n):
  # return the first n primes
  pn = primes_nth(n)
  return primes(pn+1)


def primes_nthEst(n):
  # return an estimated value for the nth prime
  # Dusart; https://t5k.org/howmany.html
  return int(round( n*(math.log(n) + math.log(math.log(n)) - 1) ))


def pi(n):
  # https://en.wikipedia.org/wiki/Prime-counting_function
  return len(primes(n))
def primesCount(n):
  return len(primes(n))


def primeGen(start=2):
  n = start
  while True:
    if isPrime(n):
      yield n
    n += 1

def nextPrime(n):
  pg = primeGen(n+1)
  return next(pg)


# ----------- testing for primality ----------

'''
def isPrimeT(n):
  # uses timeout.py function
  return timeout.timeout(isPrime, isPrimeMR, 
                     args=(n,), timeout=2)
'''

def isPrime(n):
  # is n a prime number?
  try:
    factors = pps.factorize(n)
    if factors == []:
      return False
    return (factors[0] == (n,1))
  except OverflowError as e:
    # print(e, "for", n, "so using isPrimeMR")
    return isPrimeMR(n)


def isPrimeSLOW(n):
  factors = factorize(n)
  if factors == []:
    return False
  return (factors[0] == (n,1))



def isPrimeMR(n, k=4):
  if n <= 1 or n == 4:
    return False
  if n <= 3:
    return True
  
  # Find d such that d * 2^r = n-1
  d = n - 1
  while d % 2 == 0:
    d //= 2
  
  # Iterate k times for accuracy
  for _ in range(k):
    if not millerRabin(d, n):
      return False
  return True


def millerRabin(d, n):
  '''
    It returns false if n is composite and returns
    true if n is probably prime. 
    d is an odd number such that d * 2^r = n-1
    for some r >= 1
  '''
  a = 2 + random.randint(1, n-4)  # random number in [2..n-2]
      # Corner cases like n <= 3 must be handled
  
  # Compute a^d % n
  x = pow(a, d, n)
  if x == 1 or x == n-1:
    return True
  
  # Keep squaring x while one of the following doesn't happen
  # (i) d does not reach n-1
  # (ii) (x^2) % n is not 1
  # (iii) (x^2) % n is not n-1
  while d != n-1:
    x = (x*x) % n
    d *= 2
    if x == 1:
      return False
    if x == n - 1:
      return True
  return False  # is composite



def isPrimes(n):
  # check the primality of every number up to n
  prims = primes(n)
  primsB = [False]*n
  for p in prims:
    primsB[p] = True
  return primsB


def isGaussianPrime(a, b):
  # is a Gaussian integer (a + bi) prime?
  if a == 0:  # Purely imaginary
    return isPrime(abs(b)) and (abs(b)%4 == 3)
  elif b == 0:  # Purely real
    return isPrime(abs(a)) and (abs(a)%4 == 3)
  else:   # Complex
    return isPrime(a**2 + b**2)


# ----------- additional factoring ----------

def factorizeDict(n):
  # factorizeDict(120) --> {2:3, 3:1, 5:1}
  return dict( factorize(n))

def primeFactors(n):
  # primeFactors(120) --> [2, 3, 5]
  return list(factorizeDict(n).keys())

def factorizeList(n):
  # factorizeList(120) --> [2, 2, 2, 3, 5]
  factors = factorize(n)
  fs = []
  for factor in factors:
    fs.extend([factor[0]]*factor[1])
  return fs

def factorsCount(n):
  # factorsCount(120) --> 3+1+1 
  return sum(dict( factorize(n)).values())


def divisors(n, isProper = False):
  '''
  Finds all the divisors of n using the prime factorisation of n 
      print(divisors(15)) #[1, 3, 5, 15]
      print(divisors(15, isProper = True)) #[1, 3, 5]
  '''
  pf = factorizeDict(n)
  primes = [x for x in pf]
  primes = primeFactors(n)
  div = genDivs(0, primes, pf)
  if isProper:
    div.pop(-1)
    return div
  return div


def genDivs(n, primes, pf):
  if n == len(primes):
    return [1]
  else:
    pows = [1]
    p = primes[n]
    for _ in range(pf[p]):
      pows.append(pows[-1] * p)
    
    div = []
    for q in genDivs(n+1, primes, pf):
      for p in pows:
        div.append(q * p)
    return div

def d(n):
  return numDivisors(n)
def numDivisors(n): 
  facts = factorize(n)
  return math.prod([(count+1) for (f, count) in facts])


def sigma(n, power=1):
  # the sum of all the divisors, optionally to a power
  # https://en.wikipedia.org/wiki/Divisor_function
  divs = divisors(n)
  return sum([d**power for d in divs])



# ------------- utils ------------

def sumDigits(n):
  return sum([int(ch) for ch in str(n)])

def prodDigits(n):
  return math.prod([int(ch) for ch in str(n)])


def modInverse(a, m):
  # get x in  a*x = 1 (mod m)
  # or a*x + m*y = 1
  gcd, x, y = egcd(a, m)
  if gcd != 1:
    raise ValueError(f"No mod inverse for a={a} and m={m}")
  return x % m


'''   # use math.gcd() instead
def gcd(a,b):
  while b > 0:
    a, b = b, a%b
  return a
'''

def egcd(a, b):
  # Extended GCD function: 
  # returns gcd, x, y where ax + by = gcd(a, b)
  # https://www.youtube.com/watch?v=mgvA3z-vOzc
  x, x1 = 1, 0
  y, y1 = 0, 1
  gcd, g1 = a, b
  while g1 > 0:
    q = gcd // g1
    x, x1  =  x1, x - q*x1
    y, y1  =  y1, y - q*y1
    gcd, g1  =  g1, gcd - q*g1
  return gcd, x, y


def totient(n):
  return phi(n)
def phi(n):
  '''  Euler's totient function counts the number 
    of integers less than n that are coprime to n.
  '''
  tot = 0
  for i in range(1,n+1):
    if math.gcd(n,i) == 1:
      tot += 1
  return tot

  
# ===========================================

if __name__ == "__main__":
  start = time.time()
  prims = primes(10**8)
  print(len(prims)) #  5761455
  print(f"Elapsed time: {(time.time() - start):.2f} secs")  # 0.12-0.22  vs. 15-27 secs
  
  start = time.time()
  prims = primes(10**9)
  print(len(prims)) #  50847534
  print(f"Elapsed time: {(time.time() - start):.2f} secs")  # 1.11-2.18

  prims =  primes(10**6, 10**6 + 100)
  print(prims)
  
  print(primes_sum(10**8))
  
  print(primes_nth(81749371)) # 1648727417
  print(primes_nth(10001))  # 104743
  print("est:", primes_nthEst(81749371), primes_nthEst(10001))  # 104743
   
  print(factorize(600851475143)) 
     # [(71, 1), (839, 1), (1471, 1), (6857, 1)]
  print(factorize(1648727417))  # [(1648727417, 1)]
  print(isPrime(1648727417))  # True
  print(isPrimeMR(1648727417))
  
  print(factorize(1648727418))
     # [(2, 1), (3, 1), (13, 1), (21137531, 1)]
  print(isPrime(1648727418))    # False
  print(isPrimeMR(1648727418))
  
  print(factorizeList((2**10)*(3**5)))
    # [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3]
  print(numDivisors((2**10)*(3**5))) # (10+1)*(5+1) = 66


  fd = factorizeDict(1089)
  for k in fd.keys():   # 3 2
    print(k, fd[k])     # 11 2
  
  print(divisors(15)) # [1, 3, 5, 15]
  print(divisors(15, isProper = True)) # [1, 3, 5]

  print(primesNum(5))

  print("sigma(12):", sigma(12))  # 28
  print("sigma(12, 2):", sigma(12, 2))  # 210
  print("sigma(12, 0):", sigma(12, 0))  # 6

  print(25662, sumDigits(25662))
  print(25662, prodDigits(25662))

  pg = primeGen(1000000)   # start number
  for i in range(10):
    print(next(pg), end=' ')
  print()

  n = 100000000
  print("Next prime after", n, ":", nextPrime(n))