# Arithmetic for the proofs: addition, multiplication, equality, squares.

import Base
import ./nat.bend as N

def add(a: Nat, b: Nat) -> Nat:
  match a:
    case 0n:
      b
    case 1n+p:
      1n+add(p, b)

def mul(a: Nat, +b: Nat) -> Nat:
  match a:
    case 0n:
      0n
    case 1n+p:
      add(b, mul(p, b))

def nat_eq(a: Nat, b: Nat) -> Bool:
  match a:
    case 0n:
      match b:
        case 0n:
          True{}
        case 1n+q:
          False{}
    case 1n+p:
      match b:
        case 0n:
          False{}
        case 1n+q:
          nat_eq(p, q)

def sq(+a: Nat) -> Nat:
  mul(a, a)

def pow(+a: Nat, e: Nat) -> Nat:
  match e:
    case 0n:
      1n
    case 1n+p:
      mul(a, pow(a, p))

# Pythagoras, as a definition: a right triangle is a triple
# whose hypotenuse squared is the sum of the squared sides.
def pythag(a: Nat, b: Nat, c: Nat) -> Bool:
  nat_eq(add(sq(a), sq(b)), sq(c))

# 1 + 2 + … + n
def sum_to(n: Nat) -> Nat:
  match n:
    case 0n:
      0n
    case 1n+p:
      +p2 = p
      add(1n+p2, sum_to(p2))

# ── proof scaffolding, used by LAWS.bend and PROOF.bend ──────────────────
# sym flips an equation; add_zero and add_succ are the two
# little inductions that let addition compute on its second
# argument (add() itself only recurses on its first).

def sym(-A: Type, -a: A, -b: A, e: {a == b : A}) -> {b == a : A}:
  %e : {_ == a : A}
  {==}

def add_zero(a: Nat) -> {add(a, 0n) == a : Nat}:
  match a:
    case 0n:
      {==}
    case 1n+p:
      %sym(Nat, add(p, 0n), p, add_zero(p)) : {1n+_ == 1n+p : Nat}
      {==}

def add_succ(a: Nat, b: Nat) -> {add(a, 1n+b) == 1n+add(a, b) : Nat}:
  match a:
    case 0n:
      {==}
    case 1n+p:
      %sym(Nat, add(p, 1n+b), 1n+add(p, b), add_succ(p, b)) : {1n+_ == 1n+1n+add(p, b) : Nat}
      {==}

# ── primality ───────────────────────────────────────────────

def pick(-A: Type, c: Bool, a: A, b: A) -> A:
  match c:
    case True{}:
      a
    case False{}:
      b

def nat_le(a: Nat, b: Nat) -> Bool:
  match a:
    case 0n:
      True{}
    case 1n+p:
      match b:
        case 0n:
          False{}
        case 1n+q:
          nat_le(p, q)

# subtraction that matches its second argument first, so that
# msub(x, 0) and msub(2+x, 2) normalize even for symbolic x
def msub_go(b: Nat, a: Nat) -> Nat:
  match b:
    case 0n:
      a
    case 1n+q:
      match a:
        case 0n:
          0n
        case 1n+p:
          msub_go(q, p)

def msub(a: Nat, b: Nat) -> Nat:
  msub_go(b, a)

# n mod d, by repeated subtraction (fuel keeps it structural)
def mod_go(f: Nat, +n: Nat, +d: Nat) -> Nat:
  match f:
    case 0n:
      n
    case 1n+g:
      pick(Nat, nat_le(d, n), mod_go(g, msub(n, d), d), n)

def mod(+n: Nat, +d: Nat) -> Nat:
  mod_go(n, n, d)

def divides(+d: Nat, +n: Nat) -> Bool:
  nat_eq(mod(n, d), 0n)

# does any d with 2 ≤ d ≤ k divide n?
def has_divisor(+n: Nat, k: Nat) -> Bool:
  match k:
    case 0n:
      False{}
    case 1n+p:
      match p:
        case 0n:
          False{}
        case 1n+q:
          +q2 = q
          Bool.or(divides(1n+1n+q2, n), has_divisor(n, 1n+q2))

def is_prime(+n: Nat) -> Bool:
  Bool.and(nat_le(2n, n), Bool.not(has_divisor(n, msub(n, 1n))))

# the p-th Mersenne number, 2^p − 1
def mersenne(p: Nat) -> Nat:
  Nat.sub(pow(2n, p), 1n)

# ── scaffolding for the power-of-two theorem ────────────────
# or_true absorbs disjunctions; double/double_add/mul2_double
# relate ×2 to the structural by-twos builder; mod2_zero tames
# the fuel in mod; divides2_double and has2 push the witness
# divisor 2 through the trial-division chain; pow2_shape shows
# every 2^(2+n) has the form double(2+x).

def or_true(b: Bool) -> {Bool.or(b, True{}) == True{} : Bool}:
  match b:
    case True{}:
      {==}
    case False{}:
      {==}

def double(n: Nat) -> Nat:
  match n:
    case 0n:
      0n
    case 1n+p:
      1n+1n+double(p)

def double_add(y: Nat) -> {double(y) == add(y, y) : Nat}:
  match y:
    case 0n:
      {==}
    case 1n+p:
      +p2 = p
      %sym(Nat, add(p2, 1n+p2), 1n+add(p2, p2), add_succ(p2, p2)) : {1n+1n+double(p2) == 1n+_ : Nat}
      %sym(Nat, double(p2), add(p2, p2), double_add(p2)) : {1n+1n+_ == 1n+1n+add(p2, p2) : Nat}
      {==}

def mul2_double(+m: Nat) -> {mul(2n, m) == double(m) : Nat}:
  %sym(Nat, add(m, 0n), m, add_zero(m)) : {add(m, _) == double(m) : Nat}
  sym(Nat, double(m), add(m, m), double_add(m))

def mod2_zero(x: Nat, g: Nat) -> {mod_go(1n+add(x, g), double(x), 2n) == 0n : Nat}:
  match x:
    case 0n:
      {==}
    case 1n+p:
      mod2_zero(p, g)

def divides2_double(+y: Nat) -> {divides(2n, double(1n+y)) == True{} : Bool}:
  %sym(Nat, double(y), add(y, y), double_add(y)) : {nat_eq(mod_go(1n+_, double(y), 2n), 0n) == True{} : Bool}
  %sym(Nat, mod_go(1n+add(y, y), double(y), 2n), 0n, mod2_zero(y, y)) : {nat_eq(_, 0n) == True{} : Bool}
  {==}

def has2(k: Nat, +m: Nat, e: {divides(2n, m) == True{} : Bool}) -> {has_divisor(m, 1n+1n+k) == True{} : Bool}:
  match k:
    case 0n:
      %sym(Bool, divides(2n, m), True{}, e) : {Bool.or(_, False{}) == True{} : Bool}
      {==}
    case 1n+p:
      +p2 = p
      %sym(Bool, has_divisor(m, 1n+1n+p2), True{}, has2(p2, m, e)) : {Bool.or(divides(1n+1n+1n+p2, m), _) == True{} : Bool}
      or_true(divides(1n+1n+1n+p2, m))

def use_shape(-n: Nat, w: &x: Nat -> {pow(2n, 2n+n) == double(1n+1n+x) : Nat},
  -P: Type,
  k: @x: Nat -> @e: {pow(2n, 2n+n) == double(1n+1n+x) : Nat} -> P) -> P:
  (x, w1) = w
  k(x, w1)

# every 2^(2+n) is double(2+x) for some x
def pow2_shape(n: Nat) -> &x: Nat -> {pow(2n, 2n+n) == double(1n+1n+x) : Nat}:
  match n:
    case 0n:
      (0n, {==})
    case 1n+p:
      +p2 = p
      use_shape(p2, pow2_shape(p2),
        &x: Nat -> {pow(2n, 2n+1n+p2) == double(1n+1n+x) : Nat},
        x => e =>
          +x2 = x
          (1n+1n+double(x2),
            %sym(Nat, pow(2n, 2n+p2), double(1n+1n+x2), e) : {mul(2n, _) == double(1n+1n+1n+1n+double(x2)) : Nat}
            mul2_double(double(1n+1n+x2))))


# ── the bridge to the lemma pack ────────────────────────────
# nat.bend proves its laws about Base's Nat.add and Nat.mul.
# mathlib's add and mul are the same functions, definition for
# definition; these two inductions let those laws apply here.
# nat_eq_true / nat_eq_false say the Bool test nat_eq is faithful
# to the equation, both ways; sqrt2_carry rewrites the √2 claim
# from mathlib's vocabulary into Base's, where sqrt2.bend proves it.

def add_bridge(a: Nat, -b: Nat) -> {add(a, b) == Nat.add(a, b) : Nat}:
  match a:
    case 0n:
      {==}
    case 1n+p:
      Equal.cong(Nat, Nat, x => 1n+x, add(p, b), Nat.add(p, b), add_bridge(p, b))

def mul_bridge(a: Nat, +b: Nat) -> {mul(a, b) == Nat.mul(a, b) : Nat}:
  match a:
    case 0n:
      {==}
    case 1n+p:
      +p2 = p
      %Equal.sym(Nat, mul(p2, b), Nat.mul(p2, b), mul_bridge(p2, b)) : {add(b, _) == Nat.add(b, Nat.mul(p2, b)) : Nat}
      add_bridge(b, Nat.mul(p2, b))

def nat_eq_true(a: Nat, b: Nat, e: {a == b : Nat}) -> {nat_eq(a, b) == True{} : Bool}:
  match a b:
    case 0n 0n:
      {==}
    case 0n 1n+q:
      Empty.absurd({nat_eq(0n, 1n+q) == True{} : Bool}, N.Nat.zero_neq_succ(q, e))
    case 1n+p 0n:
      Empty.absurd({nat_eq(1n+p, 0n) == True{} : Bool}, N.Nat.succ_neq_zero(p, e))
    case 1n+p 1n+q:
      nat_eq_true(p, q, N.Nat.succ_inj(p, q, e))

def nat_eq_false(a: Nat, b: Nat, ne: {a == b : Nat} -> Empty) -> {nat_eq(a, b) == False{} : Bool}:
  match a b:
    case 0n 0n:
      Empty.absurd({True{} == False{} : Bool}, ne({==}))
    case 0n 1n+q:
      {==}
    case 1n+p 0n:
      {==}
    case 1n+p 1n+q:
      nat_eq_false(p, q, e => ne(Equal.cong(Nat, Nat, x => 1n+x, p, q, e)))

def sqrt2_carry(+p: Nat, +q: Nat, e: {mul(1n+p, 1n+p) == mul(2n, mul(1n+q, 1n+q)) : Nat}) -> {Nat.mul(1n+p, 1n+p) == Nat.mul(2n, Nat.mul(1n+q, 1n+q)) : Nat}:
  Equal.trans(Nat, Nat.mul(1n+p, 1n+p), mul(1n+p, 1n+p), Nat.mul(2n, Nat.mul(1n+q, 1n+q)),
    Equal.sym(Nat, mul(1n+p, 1n+p), Nat.mul(1n+p, 1n+p), mul_bridge(1n+p, 1n+p)),
    Equal.trans(Nat, mul(1n+p, 1n+p), mul(2n, mul(1n+q, 1n+q)), Nat.mul(2n, Nat.mul(1n+q, 1n+q)), e,
      Equal.trans(Nat, mul(2n, mul(1n+q, 1n+q)), Nat.mul(2n, mul(1n+q, 1n+q)), Nat.mul(2n, Nat.mul(1n+q, 1n+q)),
        mul_bridge(2n, mul(1n+q, 1n+q)),
        Equal.cong(Nat, Nat, x => Nat.mul(2n, x), mul(1n+q, 1n+q), Nat.mul(1n+q, 1n+q), mul_bridge(1n+q, 1n+q)))))