import Base
import ./nat.bend as N

# Math.bend — fast integer arithmetic on U32 (+ tiny F32 helpers).
#
# Rules of thumb:
# - Base U32.pow is O(n) repeated multiply. pow_u32 below is O(log n).
# - U32.div/mod are ~32 steps each: pow_u32/gcd/isqrt budget them explicitly.
# - Wraps mod 2^32 on overflow (fact/fib document exact ranges).

# pow_u32: binary exponentiation, fixed 32 steps.
# pow_u32(base, exp) == base^exp (mod 2^32).
def pow_go(fuel: Nat, +base: U32, +exp: U32, +acc: U32) -> U32:
  match fuel:
    case 0n:
      acc
    case 1n+p:
      +bit = U32.and(exp, 1)
      +even = U32.is_zero(bit)
      pow_go(p,
        U32.mul(base, base),
        U32.shr(exp),
        Bool.pick(U32, even, acc, U32.mul(acc, base)))

def pow_u32(+base: U32, +exp: U32) -> U32:
  pow_go(32n, base, exp, 1)

# gcd: Euclidean algorithm, fixed 64 steps (> worst case ~47 for 32-bit).
# gcd(a, 0) == a, gcd(0, 0) == 0.
def gcd_go(fuel: Nat, +a: U32, +b: U32) -> U32:
  match fuel:
    case 0n:
      a
    case 1n+p:
      +z = U32.is_zero(b)
      gcd_go(p,
        Bool.pick(U32, z, a, b),
        Bool.pick(U32, z, b, U32.mod(a, b)))

def gcd(+a: U32, +b: U32) -> U32:
  gcd_go(64n, a, b)

# lcm: a / gcd(a,b) * b. Returns 0 when either side is 0.
def lcm_nonzero(+g: U32, a: U32, b: U32) -> U32:
  U32.mul(U32.div(a, g), b)

def lcm(+a: U32, +b: U32) -> U32:
  Bool.pick(U32, Bool.or(U32.is_zero(a), U32.is_zero(b)), 0,
    lcm_nonzero(gcd(a, b), a, b))

# isqrt: floor(sqrt(n)). Integer Newton, fixed 32 steps, frozen when stable.
# Never divides by zero: div(n, 0) == 0 by Base, and x == 0 only when n == 0.
def isqrt_go(fuel: Nat, +n: U32, +x: U32) -> U32:
  match fuel:
    case 0n:
      x
    case 1n+p:
      +nx = U32.shr((x + U32.div(n, x) : U32))
      isqrt_go(p, n, Bool.pick(U32, U32.is_eq(nx, x), x, nx))

def isqrt(+n: U32) -> U32:
  isqrt_go(32n, n, n)

# fib: nth Fibonacci (0, 1, 1, 2, 3, 5, ...). Tail loop over Nat fuel.
# Exact for n <= 47 (wraps mod 2^32 past that).
def fib_go(n: Nat, a: U32, +b: U32) -> U32:
  match n:
    case 0n:
      a
    case 1n+p:
      fib_go(p, b, (a + b : U32))

def fib(n: Nat) -> U32:
  fib_go(n, 0, 1)

# fact: n! . Tail loop. Exact for n <= 12 (wraps past that).
def fact_go(n: Nat, +i: U32, acc: U32) -> U32:
  match n:
    case 0n:
      acc
    case 1n+p:
      fact_go(p, (i + 1 : U32), (acc * i : U32))

def fact(n: Nat) -> U32:
  fact_go(n, 1, 1)

# tri: n*(n+1)/2 without overflow bias: (n/2)*((n|1)) style is overkill;
# plain version, wraps for huge n like everything else here.
def tri(+n: U32) -> U32:
  U32.shr(U32.mul(n, (n + 1 : U32)))

# fib_fast: doubling method, fixed 32 steps from bit 31 to bit 0.
# State (F(k), F(k+1)); bit 0 -> (F(2k), F(2k+1)); bit 1 -> (F(2k+1), F(2k+2)).
# 2*F(k+1) >= F(k) always, so the subtraction is exact (never wraps).
def ff_step(+bit: U32, +a: U32, +b: U32) -> U32 & U32:
  +t1 = U32.mul(a, U32.sub(U32.mul(2, b), a))
  +t2 = (U32.mul(a, a) + U32.mul(b, b) : U32)
  +one = U32.is_zero(bit)
  (Bool.pick(U32, one, t1, t2), Bool.pick(U32, one, t2, (t1 + t2 : U32)))

def ff_fin0(q: U32 & U32) -> U32:
  (a2, b2) = q
  a2

def ff_go(j: Nat, +n: U32, q: U32 & U32) -> U32:
  match j:
    case 0n:
      (a, b) = q
      ff_fin0(ff_step(U32.and(n, 1), a, b))
    case 1n++p:
      (a, b) = q
      ff_go(p, n, ff_step(U32.and(U32.shrn(n, Nat.add(1n, p)), 1), a, b))

def fib_fast(+n: U32) -> U32:
  ff_go(31n, n, (0, 1))

# collatz: steps to reach 1 (fuel-capped; frozen once x == 1).
def collatz_go(fuel: Nat, +x: U32, acc: Nat) -> Nat:
  match fuel:
    case 0n:
      acc
    case 1n+p:
      +done = U32.is_eq(x, 1)
      +even = U32.is_zero(U32.and(x, 1))
      collatz_go(p,
        Bool.pick(U32, done, x,
          Bool.pick(U32, even, U32.shr(x), (x * 3 + 1 : U32))),
        Nat.add(acc, U32.to_nat(Bool.to_u32(Bool.not(done)))))

def collatz(fuel: Nat, +x: U32) -> Nat:
  collatz_go(fuel, x, 0n)

law collatz_27:
  {collatz(200n, 27) == 111n : Nat}

def collatz_27():
  {==}

law fib_fast_20:
  {fib_fast(20) == 6765 : U32}

def fib_fast_20():
  {==}

law pow_2_10:
  {pow_u32(2, 10) == 1024 : U32}

def pow_2_10():
  {==}

law isqrt_16:
  {isqrt(16) == 4 : U32}

def isqrt_16():
  {==}

law isqrt_0:
  {isqrt(0) == 0 : U32}

def isqrt_0():
  {==}

# F32 helpers (each is one cheap axiom call; F32 itself is unprovable).
def finv(x: F32) -> F32:
  F32.div(1.0, x)

def favg(+a: F32, +b: F32) -> F32:
  F32.div(F32.add(a, b), 2.0)

def fclamp01(x: F32) -> F32:
  F32.clamp(x, 0.0, 1.0)

def fdist2(+x1: F32, +y1: F32, +x2: F32, +y2: F32) -> F32:
  F32.add(F32.square(F32.sub(x2, x1)), F32.square(F32.sub(y2, y1)))

def fdist(+x1: F32, +y1: F32, +x2: F32, +y2: F32) -> F32:
  F32.sqrt(fdist2(x1, y1, x2, y2))

# gcd_go_zero: frozen loop mirroring gcd_go(-, a, 0). Each gcd_go step with
# b == 0 has z = True, so both picks freeze (a, 0) -> (a, 0); the loop is
# definitionally the IH, and gcd(a, 0) normalizes to a (64n unrolls).
def gcd_go_zero(fuel: Nat, +a: U32) -> U32:
  match fuel:
    case 0n:
      a
    case 1n+p:
      gcd_go_zero(p, a)

law gcd_zero_r:
  for +a: U32
  {gcd(a, 0) == a : U32}

def gcd_zero_r(a):
  {==}

# add_le_mono: adding c to both sides preserves Nat order.
# Induction on c via N.add_zero_r (base) and N.add_succ_r (step, symmed
# since % rewrites RHS->LHS); 1n+_ is_le is definitional through Nat.cmp.
law add_le_mono:
  for +a: Nat
  for +b: Nat
  for +c: Nat
  {Nat.is_le(Nat.add(a, c), Nat.add(b, c)) == Nat.is_le(a, b) : Bool}

def add_le_mono(a, b, c):
  match c:
    case 0n:
      %Equal.sym(Nat, Nat.add(a, 0n), a, N.add_zero_r(a)) : {Nat.is_le(_, Nat.add(b, 0n)) == Nat.is_le(a, b) : Bool}
      %Equal.sym(Nat, Nat.add(b, 0n), b, N.add_zero_r(b)) : {Nat.is_le(a, _) == Nat.is_le(a, b) : Bool}
      {==}
    case 1n++p:
      %Equal.sym(Nat, Nat.add(a, 1n+p), 1n+Nat.add(a, p), N.add_succ_r(a, p)) : {Nat.is_le(_, Nat.add(b, 1n+p)) == Nat.is_le(a, b) : Bool}
      %Equal.sym(Nat, Nat.add(b, 1n+p), 1n+Nat.add(b, p), N.add_succ_r(b, p)) : {Nat.is_le(1n+Nat.add(a, p), _) == Nat.is_le(a, b) : Bool}
      %add_le_mono(a, b, p) : {Nat.is_le(1n+Nat.add(a, p), 1n+Nat.add(b, p)) == _ : Bool}
      {==}