-- ArkLib/Data/Nat/Bitwise.lean

import ArkLib.Data.Fin.BigOperators

import Mathlib.Algebra.BigOperators.Ring.Finset

import Mathlib.Algebra.Order.Ring.Star

import Mathlib.Data.Nat.Bitwise

import Mathlib.Data.Nat.Digits.Defs

import Mathlib.Data.Finsupp.Basic

import Mathlib.Algebra.Order.BigOperators.Group.Finset

import Mathlib.Algebra.BigOperators.Fin

namespace Nat

def getBit (k n : Nat) : Nat := (n >>> k) &&& 1

def getLowBits (numLowBits : ℕ) (n : ℕ) := n &&& ((1 <<< numLowBits) - 1)

def getHighBits_no_shl (numLowBits : ℕ) (n : ℕ) : ℕ := n >>> numLowBits

def getHighBits (numLowBits : ℕ) (n : ℕ) : ℕ :=
  (getHighBits_no_shl numLowBits n) <<< numLowBits

-- Repository definitions used:
content: ...

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: def getBit (k n : Nat) : Nat := (n >>> k) &&& 1

content: def getLowBits (numLowBits : ℕ) (n : ℕ) := n &&& ((1 <<< numLowBits) - 1)

content: def getHighBits_no_shl (numLowBits : ℕ) (n : ℕ) : ℕ := n >>> numLowBits

content: def getHighBits (numLowBits : ℕ) (n : ℕ) : ℕ :=
  (getHighBits_no_shl numLowBits n) <<< numLowBits

-- Local lemmas used:
content: lemma getBit_lt_2 {k n : Nat} : getBit k n < 2

content: lemma eq_iff_eq_all_getBits {n m : ℕ} : n = m ↔ ∀ k, getBit k n = getBit k m

content: lemma shiftRight_and_one_distrib {n m k : ℕ} :
    Nat.getBit k (n &&& m) = Nat.getBit k n &&& Nat.getBit k m

content: lemma and_eq_zero_iff_and_each_getBit_eq_zero {n m : ℕ} :
    n &&& m = 0 ↔ ∀ k, Nat.getBit k n &&& Nat.getBit k m = 0

content: lemma div_2_form {nD2 b : ℕ} (h_b : b < 2):
  (nD2 * 2 + b) / 2 = nD2

content: lemma and_by_split_lowBits {n m n1 m1 bn bm : ℕ} (h_bn : bn < 2) (h_bm : bm < 2)
  (h_n : n = n1 * 2 + bn) (h_m : m = m1 * 2 + bm):
  n &&& m = (n1 &&& m1) * 2 + (bn &&& bm)

content: lemma xor_by_split_lowBits {n m n1 m1 bn bm : ℕ} (h_bn : bn < 2) (h_bm : bm < 2)
  (h_n : n = n1 * 2 + bn) (h_m : m = m1 * 2 + bm):
  n ^^^ m = (n1 ^^^ m1) * 2 + (bn ^^^ bm)

content: lemma or_by_split_lowBits {n m n1 m1 bn bm : ℕ} (h_bn : bn < 2) (h_bm : bm < 2)
  (h_n : n = n1 * 2 + bn) (h_m : m = m1 * 2 + bm):
  n ||| m = (n1 ||| m1) * 2 + (bn ||| bm)

content: lemma sum_eq_xor_plus_twice_and (n : Nat) : ∀ m : ℕ, n + m = (n ^^^ m) + 2 * (n &&& m)

content: lemma xor_of_and_eq_zero_is_or {n m : ℕ} (h_n_AND_m : n &&& m = 0) : n ^^^ m = n ||| m

content: lemma sum_of_and_eq_zero_is_or {n m : ℕ} (h_n_AND_m : n &&& m = 0) : n + m = n ||| m

content: lemma getBit_of_multiple_of_power_of_two {n p : ℕ}: ∀ k,
  getBit (k) (2^p * n) = if k < p then 0 else getBit (k-p) n

content: lemma getBit_of_shiftLeft {n p : ℕ}:
  ∀ k, getBit (k) (n <<< p) = if k < p then 0 else getBit (k - p) n

content: lemma getBit_of_shiftRight {n p : ℕ}:
  ∀ k, getBit k (n >>> p) = getBit (k+p) n

content: lemma getBit_of_and {n m k: ℕ} : getBit k (n &&& m) = getBit k n &&& getBit k m

content: lemma getBit_of_two_pow_sub_one {i k: ℕ} : getBit k (2^i - 1) =
    if k < i then 1 else 0

content: lemma getBit_of_lowBits {n: ℕ} (numLowBits : ℕ) : ∀ k, getBit k (getLowBits numLowBits n) =
    if k < numLowBits then getBit k n else 0

content: theorem and_highBits_lowBits_eq_zero {n : ℕ} (numLowBits : ℕ) :
    getHighBits numLowBits n &&& getLowBits numLowBits n = 0

-- Target theorem and proof:
lemma num_eq_highBits_add_lowBits {n: ℕ} (numLowBits: ℕ) :
  n = getHighBits numLowBits n + getLowBits numLowBits n
:= by
  apply eq_iff_eq_all_getBits.mpr; unfold getBit
  intro k
  --- use 2 getBit extractions to get the condition for getLowBits of ((n >>> numLowBits) <<<
  --  numLowBits)
  set highBits_no_shl := n >>> numLowBits
  have h_getBit_highBits_shl := getBit_of_shiftLeft (n := highBits_no_shl) (p := numLowBits)
  have h_getBit_lowBits := getBit_of_lowBits (n := n) (numLowBits := numLowBits)
  -- AND of highBitss & lowBitss is 0 => we use this to convert the sum into OR
  have h_and := and_highBits_lowBits_eq_zero (n := n) (numLowBits := numLowBits)
  rw [sum_of_and_eq_zero_is_or h_and]
  --- now reason on bitwise operations only
  rw [Nat.shiftRight_or_distrib, Nat.and_distrib_right]
  change getBit k n = getBit k ((n >>> numLowBits) <<< numLowBits)
    ||| getBit k (getLowBits numLowBits n)
  rw [h_getBit_highBits_shl, h_getBit_lowBits]
  if h_k: k < numLowBits then
    simp only [h_k, ↓reduceIte, Nat.zero_or] at *
  else
    have h_ne: ¬(k < numLowBits) := by omega
    have h_num_le_k: numLowBits ≤ k := by omega
    simp only [h_ne, not_false_eq_true, ↓reduceIte, Nat.or_zero] at *
    rw [getBit_of_shiftRight]
    congr
    rw [Nat.sub_add_cancel (n:=k) (m:=numLowBits) (by omega)]