-- Clean/Gadgets/Addition8/Theorems.lean

import Clean.Utils.Field

import Clean.Gadgets.Boolean

namespace Gadgets.Addition8.Theorems

variable {p : ℕ} [Fact p.Prime]

variable [p_large_enough: Fact (p > 512)]

-- Repository definitions used:
content: syntax "field_to_nat" : tactic

content: macro_rules
  | `(tactic|field_to_nat) =>
    `(tactic|(
      intros
      repeat rw [ZMod.val_add] 
      repeat rw [ZMod.val_mul] 
      repeat rw [val_eq_256]
      try simp only [Nat.add_mod_mod, Nat.mod_add_mod, Nat.mul_mod_mod, Nat.mod_mul_mod]
      rw [Nat.mod_eq_of_lt _]
      repeat linarith [‹Fact (_ > 512)›.elim]))

example [Fact (p > 512)] (x y : F p) (hx : x.val < 256) (hy : y.val < 2) :
    (x + y * 256).val = x.val + y.val * 256 := by admit /- proof elided -/

content: def F p := ZMod p

content: def IsBool {α : Type*} [Zero α] [One α] (x : α) : Prop := x = 0 ∨ x = 1

-- Repository lemmas used:
content: theorem byte_sum_and_bit_do_not_wrap (x y b : F p) [Fact (p > 512)]:
    x.val < 256 -> y.val < 256 -> b.val < 2 -> (b + x + y).val = b.val + x.val + y.val

content: theorem byte_sum_and_bit_do_not_wrap' (x y b : F p) [Fact (p > 512)]:
    x.val < 256 -> y.val < 256 -> b.val < 2 -> (x + y + b).val = x.val + y.val + b.val

content: theorem byte_sum_and_bit_lt_512 (x y b : F p) [Fact (p > 512)]:
    x.val < 256 -> y.val < 256 -> b.val < 2 -> (x + y + b).val < 512

content: theorem byte_plus_256_do_not_wrap (x : F p) [Fact (p > 512)]:
    x.val < 256 -> (x + 256).val = x.val + 256

content: theorem byte_sum_do_not_wrap (x y : F p) [Fact (p > 512)]:
    x.val < 256 -> y.val < 256 -> (x + y).val = x.val + y.val

content: theorem byte_sum_le_bound (x y : F p) [Fact (p > 512)]:
    x.val < 256 -> y.val < 256 -> (x + y).val < 511

content: theorem val_lt_two {p : ℕ} [Fact p.Prime] {x : F p} (h : IsBool x) : x.val < 2

content: lemma val_eq_256 [p_large_enough: Fact (p > 512)] : (256 : F p).val = 256

content: theorem ext {x y : F p} (h : x.val = y.val) : x = y

content: theorem ext_iff {x y : F p} : x = y ↔ x.val = y.val

content: theorem val_lt_p {p : ℕ} (x : ℕ) : (x < p) → (x : F p).val = x

-- Local definitions used:

-- Local lemmas used:
content: theorem soundness_zero_carry (x y out carry_in : F p):
    x.val < 256 -> y.val < 256 -> out.val < 256  -> carry_in.val < 2 ->
    (carry_in + x + y - out = 0 -> (out.val = (carry_in.val + x.val + y.val) % 256
    ∧ (carry_in.val + x.val + y.val) / 256 = 0))

content: theorem soundness_one_carry (x y out carry_in : F p):
    x.val < 256 -> y.val < 256 -> out.val < 256 -> carry_in.val < 2 ->
    carry_in + x + y - out - 256 = 0 -> (out.val = (carry_in.val + x.val + y.val) % 256
    ∧ (carry_in.val + x.val + y.val) / 256 = 1)

-- Target theorem and proof:
theorem soundness (x y out carry_in carry_out : F p):
    x.val < 256 -> y.val < 256 ->
    out.val < 256 ->
    IsBool carry_in ->
    IsBool carry_out ->
    (x + y + carry_in + -out + -(carry_out * 256) = 0) ->
    (out.val = (x.val + y.val + carry_in.val) % 256
    ∧ carry_out.val = (x.val + y.val + carry_in.val) / 256)
:= by
  intros hx hy hout carry_in_bool carry_out_bool h
  have carry_in_bound := IsBool.val_lt_two carry_in_bool

  rcases carry_out_bool with zero_carry | one_carry
  -- case with zero carry
  · rw [zero_carry] at h
    simp only [zero_mul, neg_zero, add_zero] at h
    rw [←sub_eq_add_neg] at h
    have h_spec : carry_in + x + y - out = 0 := by
      rw [add_comm (x + y), ←add_assoc] at h
      assumption

    have thm := soundness_zero_carry x y out carry_in hx hy hout carry_in_bound h_spec
    apply_fun ZMod.val at zero_carry

    -- now it is just a matter of shuffling terms around
    have shuffle_terms : carry_in.val + x.val + y.val = x.val + y.val + carry_in.val := by
      zify; ring
    rw [ZMod.val_zero, ← thm.right] at zero_carry
    rw [shuffle_terms] at thm
    rw [shuffle_terms] at zero_carry
    constructor
    · exact thm.left
    · exact zero_carry

  -- case with one carry
  · rw [one_carry] at h
    simp only [one_mul] at h
    -- rw [←sub_eq_add_neg, ←sub_eq_add_neg (carry_in + x + y)] at h
    have h_spec : carry_in + x + y - out - 256 = 0 := by
      rw [add_comm (x + y), ←add_assoc] at h
      ring_nf at h; ring_nf
      assumption

    -- instantiate the sub-theorem
    have thm := soundness_one_carry x y out carry_in hx hy hout carry_in_bound h_spec
    apply_fun ZMod.val at one_carry

    have shuffle_terms : carry_in.val + x.val + y.val = x.val + y.val + carry_in.val := by
      zify; ring
    rw [ZMod.val_one, ← thm.right] at one_carry
    rw [shuffle_terms] at thm
    rw [shuffle_terms] at one_carry
    constructor
    · exact thm.left
    · exact one_carry