# tinygrad/nat.bend -- arithmetic lemmas over Nat used by PROOF.bend.
# Textbook Peano arithmetic, structural inductions only.

import Base

# x + 0 == x
def add_zero_right(+a: Nat) -> {Nat.add(a, 0n) == a : Nat}:
  match a:
    case 0n:
      {==}
    case 1n+p:
      %add_zero_right(p) : {1n+Nat.add(p, 0n) == 1n+_ : Nat}
      {==}

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

# x - x == 0
def sub_diag(+a: Nat) -> {Nat.sub(a, a) == 0n : Nat}:
  match a:
    case 0n:
      {==}
    case 1n+p:
      %sub_diag(p) : {Nat.sub(p, p) == _ : Nat}
      {==}

# a + b == b + a
def succ1(n: Nat) -> Nat:
  (1n+n : Nat)

def add_comm(+a: Nat, +b: Nat) -> {Nat.add(a, b) == Nat.add(b, a) : Nat}:
  match a:
    case 0n:
      eq = add_zero_right(b)
      eq2 = Equal.sym(Nat, Nat.add(b, 0n), b, eq)
      Equal.trans(Nat, Nat.add(0n, b), b, Nat.add(b, 0n), {==}, eq2)
    case 1n+p:
      c1 = Equal.cong(Nat, Nat, succ1, Nat.add(p, b), Nat.add(b, p), add_comm(p, b))
      c2 = Equal.sym(Nat, Nat.add(b, 1n+p), (1n+Nat.add(b, p) : Nat), add_succ_right(b, p))
      t1 = Equal.trans(Nat, Nat.add(1n+p, b), (1n+Nat.add(p, b) : Nat), (1n+Nat.add(b, p) : Nat), {==}, c1)
      Equal.trans(Nat, Nat.add(1n+p, b), (1n+Nat.add(b, p) : Nat), Nat.add(b, 1n+p), t1, c2)

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

# (n + k) - k == n
def sub_add_cancel(+n: Nat, +k: Nat) -> {Nat.sub(Nat.add(n, k), k) == n : Nat}:
  match n k:
    case 0n, 0n:
      {==}
    case 0n, 1n+k2:
      sub_diag(k2)
    case 1n+n2, 0n:
      %add_zero_right(n2) : {1n+Nat.add(n2, 0n) == 1n+_ : Nat}
      {==}
    case 1n+n2, 1n+k2:
      e1 = add_succ_right(n2, k2)
      e2 = Equal.cong(Nat, Nat, x => (Nat.sub(x, k2) : Nat), Nat.add(n2, 1n+k2), (1n+Nat.add(n2, k2) : Nat), e1)
      e3 = sub_add_cancel(1n+n2, k2)
      Equal.trans(Nat, Nat.sub(Nat.add(n2, 1n+k2), k2), Nat.sub((1n+Nat.add(n2, k2) : Nat), k2), 1n+n2, e2, e3)