-- 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)

-- 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)

-- 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 getBit_two_pow {i k : ℕ} : (getBit k (2^i) = if i == k then 1 else 0)

content: lemma and_two_pow_eq_zero_of_getBit_0 {n i : ℕ} (h_getBit : getBit i n = 0)
    : n &&& (2 ^ i) = 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 sum_of_and_eq_zero_is_xor {n m : ℕ} (h_n_AND_m : n &&& m = 0) : n + m = 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 getBit_of_or {n m k: ℕ} : getBit k (n ||| m) = getBit k n ||| getBit k m

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

-- Target theorem and proof:
lemma getLowBits_succ {n: ℕ} (numLowBits: ℕ) :
    getLowBits (numLowBits + 1) n = getLowBits numLowBits n
    + (getBit numLowBits n) <<< numLowBits
:= by
  apply eq_iff_eq_all_getBits.mpr;
  intro k
  have h_getBit_lt_numLowBits: getBit numLowBits n < 2 := by exact getBit_lt_2
  interval_cases h_getBit: getBit numLowBits n
  · rw [Nat.zero_shiftLeft]
    simp only [add_zero]
    -- ⊢ getLowBits n (numLowBits + 1) >>> k &&& 1 = getLowBits n numLowBits >>> k &&& 1
    change getBit k (getLowBits (numLowBits + 1) n) = getBit k (getLowBits numLowBits n)
    have getBit_right := getBit_of_lowBits (n := n) (numLowBits := numLowBits) k
    have getBit_left := getBit_of_lowBits (n := n) (numLowBits := numLowBits + 1) k
    rw [getBit_right, getBit_left]
    if h_k: k < numLowBits then
      simp only [h_k, ↓reduceIte]
      have h_k_lt: k < numLowBits + 1 := by omega
      simp only [h_k_lt, ↓reduceIte]
    else if h_k_eq: k = numLowBits then
      simp only [h_k_eq]
      simp only [lt_add_iff_pos_right, zero_lt_one, ↓reduceIte, lt_self_iff_false]
      omega
    else
      have k_ne_lt: ¬(k < numLowBits) := by omega
      have k_ne_lt_add_1: ¬(k < numLowBits + 1) := by omega
      simp only [k_ne_lt_add_1, ↓reduceIte, k_ne_lt]
  · change getBit k (getLowBits (numLowBits + 1) n)
      = getBit k (getLowBits numLowBits n + 1 <<< numLowBits)
    have getBit_left := getBit_of_lowBits (n := n) (numLowBits := numLowBits + 1) k
    have getBit_right := getBit_of_lowBits (n := n) (numLowBits := numLowBits) k
    rw [getBit_left]

    have h_and_eq_0 := and_two_pow_eq_zero_of_getBit_0 (n:=getLowBits numLowBits n)
      (i:=numLowBits) (by
      simp only [getBit_of_lowBits (n := n) (numLowBits := numLowBits) numLowBits,
        lt_self_iff_false, ↓reduceIte]
    )
    rw [←one_mul (a:=2 ^ numLowBits)] at h_and_eq_0
    rw [←Nat.shiftLeft_eq (a:=1) (b:=numLowBits)] at h_and_eq_0
    have h_sum_eq_xor := sum_of_and_eq_zero_is_xor (n:=getLowBits numLowBits n)
      (m:=1 <<< numLowBits) (h_n_AND_m:=h_and_eq_0)
    have h_sum_eq_or := xor_of_and_eq_zero_is_or (n:=getLowBits numLowBits n)
      (m:=1 <<< numLowBits) (h_n_AND_m:=h_and_eq_0)
    rw [h_sum_eq_or] at h_sum_eq_xor
    rw [h_sum_eq_xor]
    rw [getBit_of_or]
    rw [getBit_of_lowBits]
    conv_rhs =>
      enter [2, 2]; rw [Nat.shiftLeft_eq, one_mul]
    rw [getBit_two_pow]

    if h_k: k < numLowBits then
      have h_k_lt: k < numLowBits + 1 := by omega
      simp only [h_k_lt, ↓reduceIte, h_k, beq_iff_eq]
      have h_k_ne_eq: numLowBits ≠ k := by omega
      simp only [h_k_ne_eq, ↓reduceIte, Nat.or_zero]
    else if h_k_eq: k = numLowBits then
      simp only [h_k_eq, lt_add_iff_pos_right, zero_lt_one, ↓reduceIte, lt_self_iff_false, BEq.rfl,
        Nat.zero_or]
      omega
    else
      have k_ne_lt: ¬(k < numLowBits) := by omega
      have k_ne_lt_add_1: ¬(k < numLowBits + 1) := by omega
      simp only [k_ne_lt_add_1, ↓reduceIte, k_ne_lt, beq_iff_eq, Nat.zero_or, right_eq_ite_iff,
        zero_ne_one, imp_false, ne_eq]
      omega