import Base
import ./math.bend as Math

# Num.bend — computational number theory on U32.
#
# All loops are fuel-bounded and freeze state with picks (same discipline
# as bits.bend). Costs are explicit: trial division and phi pay one
# U32.mod (~32 steps) per candidate; mod_pow pays one mod per bit.

# divides: d | n. divides(0, n) is (n == 0) by convention.
def divides(+d: U32, +n: U32) -> Bool:
  Bool.pick(Bool, U32.is_zero(d), U32.is_zero(n),
    U32.is_zero(U32.mod(n, d)))

# is_prime: trial division up to isqrt(n). Fixed (sqrt(n)-1) steps.
def prime_go(fuel: Nat, +d: U32, +n: U32, bad: Bool) -> Bool:
  match fuel:
    case 0n:
      Bool.not(bad)
    case 1n+p:
      prime_go(p, (d + 1 : U32), n, Bool.or(bad, divides(d, n)))

def is_prime(+n: U32) -> Bool:
  Bool.pick(Bool, U32.is_lt(n, 2), False{},
    prime_go(U32.to_nat(U32.sub(Math.isqrt(n), 1)), 2, n, False{}))

# prime_pi: count of primes 2 <= k <= n.
def pi_go(fuel: Nat, +k: U32, acc: U32) -> U32:
  match fuel:
    case 0n:
      acc
    case 1n+p:
      pi_go(p, (k + 1 : U32),
        (acc + Bool.pick(U32, is_prime(k), 1, 0) : U32))

def prime_pi(+n: U32) -> U32:
  pi_go(U32.to_nat(U32.sub(n, 1)), 2, 0)

# mod_pow: base^exp mod m (m >= 1). Binary method, fixed 32 steps.
def modpow_go(fuel: Nat, +base: U32, +exp: U32, +acc: U32, +m: U32) -> U32:
  match fuel:
    case 0n:
      acc
    case 1n+p:
      +bit = U32.and(exp, 1)
      +even = U32.is_zero(bit)
      modpow_go(p,
        U32.mod(U32.mul(base, base), m),
        U32.shr(exp),
        Bool.pick(U32, even, acc, U32.mod(U32.mul(acc, base), m)),
        m)

def mod_pow(+base: U32, +exp: U32, +m: U32) -> U32:
  modpow_go(32n, U32.mod(base, m), exp, U32.mod(1, m), m)

# coprime: gcd(a, b) == 1.
def coprime(+a: U32, +b: U32) -> Bool:
  U32.is_eq(Math.gcd(a, b), 1)

# phi: Euler totient, count of 1 <= k <= n with gcd(k, n) == 1.
def phi_go(fuel: Nat, +k: U32, +n: U32, acc: U32) -> U32:
  match fuel:
    case 0n:
      acc
    case 1n+p:
      phi_go(p, (k + 1 : U32), n,
        (acc + Bool.pick(U32, coprime(k, n), 1, 0) : U32))

def phi(+n: U32) -> U32:
  phi_go(U32.to_nat(n), 1, n, 0)

# is_square: n is a perfect square.
def is_square(+n: U32) -> Bool:
  U32.is_eq(U32.mul(Math.isqrt(n), Math.isqrt(n)), n)

# icbrt: floor(cbrt(n)). Integer Newton, fixed 32 steps, frozen when stable.
def icbrt_go(fuel: Nat, +n: U32, +x: U32) -> U32:
  match fuel:
    case 0n:
      x
    case 1n+p:
      +q = U32.div(n, U32.mul(x, x))
      +nx = U32.div((U32.mul(2, x) + q : U32), 3)
      icbrt_go(p, n, Bool.pick(U32, U32.is_eq(nx, x), x, nx))

def icbrt(+n: U32) -> U32:
  icbrt_go(32n, n, n)

# choose: n over k (0 when k > n). Multiplicative, exact for small values.
def choose_go(fuel: Nat, +i: U32, +n: U32, +k: U32, acc: U32) -> U32:
  match fuel:
    case 0n:
      acc
    case 1n+p:
      choose_go(p, (i + 1 : U32), n, k,
        U32.div(U32.mul(acc, (U32.sub(n, k) + i : U32)), i))

def choose(+n: U32, +k: U32) -> U32:
  +kk = U32.min(k, U32.sub(n, k))
  Bool.pick(U32, U32.is_gt(k, n), 0,
    choose_go(U32.to_nat(kk), 1, n, kk, 1))

# catalan: choose(2n, n) / (n + 1). Exact for small n.
def catalan(+n: U32) -> U32:
  U32.div(choose((n + n : U32), n), (n + 1 : U32))

law catalan_4:
  {catalan(4) == 14 : U32}

def catalan_4():
  {==}

# perm: falling product n * (n-1) * ... * (n-k+1), 0 when k > n.
def perm_go(fuel: Nat, +i: U32, +n: U32, +k: U32, acc: U32) -> U32:
  match fuel:
    case 0n:
      acc
    case 1n+p:
      perm_go(p, (i + 1 : U32), n, k,
        U32.mul(acc, (U32.sub(n, k) + i : U32)))

def perm(+n: U32, +k: U32) -> U32:
  Bool.pick(U32, U32.is_gt(k, n), 0,
    perm_go(U32.to_nat(k), 1, n, k, 1))

law perm_5_2:
  {perm(5, 2) == 20 : U32}

def perm_5_2():
  {==}

# prime_nth: value of the `want`-th prime (1-indexed), fuel-bounded search.
def nth_go(fuel: Nat, +c: U32, +want: U32, cnt: U32, +best: U32) -> U32:
  match fuel:
    case 0n:
      best
    case 1n+p:
      +isp = is_prime(c)
      +ncnt = (cnt + Bool.pick(U32, isp, 1, 0) : U32)
      nth_go(p, (c + 1 : U32), want, ncnt,
        Bool.pick(U32, Bool.and(isp, U32.is_eq(ncnt, want)), c, best))

def prime_nth(fuel: Nat, +want: U32) -> U32:
  nth_go(fuel, 2, want, 0, 2)

law prime_nth_6:
  {prime_nth(25n, 6) == 13 : U32}

def prime_nth_6():
  {==}

law prime_11:
  {is_prime(11) == True{} : Bool}

def prime_11():
  {==}

law prime_9_false:
  {is_prime(9) == False{} : Bool}

def prime_9_false():
  {==}