-- 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 binaryFinMapToNat {n : ℕ} (m : Fin n → ℕ) (h_binary : ∀ j: Fin n, m j ≤ 1) : Fin (2^n) :=

-- 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 binaryFinMapToNat {n : ℕ} (m : Fin n → ℕ) (h_binary : ∀ j: Fin n, m j ≤ 1) : Fin (2^n) :=

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

content: lemma getBit_eq_testBit (k n : Nat) : getBit k n = if n.testBit k then 1 else 0

content: lemma getBit_zero_eq_zero {k : Nat} : getBit k 0 = 0

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 getBit_of_xor {n m k: ℕ} : getBit k (n ^^^ m) = getBit k n ^^^ getBit k m

content: lemma getBit_of_lt_two_pow {n: ℕ} (a: Fin (2^n)) (k: ℕ) :
  getBit k a = if k < n then getBit k a else 0

-- Target theorem and proof:
lemma getBit_of_binaryFinMapToNat {n : ℕ} (m : Fin n → ℕ) (h_binary: ∀ j: Fin n, m j ≤ 1) :
    ∀ k: ℕ, Nat.getBit k (binaryFinMapToNat m h_binary).val
      = if h_k: k < n then m ⟨k, by omega⟩ else 0
:= by
  -- We prove this by induction on `n`.
  induction n with
  | zero =>
    intro k;
    simp only [Nat.pow_zero, Fin.val_eq_zero, not_lt_zero', ↓reduceDIte]
    exact getBit_zero_eq_zero
  | succ n ih =>
    -- Inductive step: Assume the property holds for `n`, prove it for `n+1`.
    have h_lt: 2^n - 1 < 2^n := by
      refine sub_one_lt ?_
      exact Ne.symm (NeZero.ne' (2 ^ n))
    intro k
    dsimp [binaryFinMapToNat]
    -- ⊢ (↑k).getBit (∑ j, 2 ^ ↑j * m j) = m k
    rw [Fin.sum_univ_castSucc] -- split the msb
    set prevSum := ∑ i: Fin n, (2 ^ i.castSucc.val) * (m i.castSucc)
    let mPrev := fun i: Fin n => m i.castSucc
    have h_getBit_prevSum := ih (m:=mPrev) (h_binary:=by exact fun j ↦ h_binary j.castSucc)
    have h_prevSum_eq: prevSum = binaryFinMapToNat mPrev
      (by exact fun j ↦ h_binary j.castSucc) := by rfl
    set msbTerm := 2 ^ ((Fin.last n).val) * m (Fin.last n)
    -- ⊢ (↑k).getBit (prevSum + msbTerm) = m k
    have h_m_at_last: m ⟨n, by omega⟩ ≤ 1 := by exact h_binary (Fin.last n)
    have h_sum_eq_xor: prevSum + msbTerm = prevSum ^^^ msbTerm := by
      rw [sum_of_and_eq_zero_is_xor]
      unfold msbTerm
      interval_cases h_m_last_val: m ⟨n, by omega⟩
      · simp only [Fin.last, h_m_last_val, mul_zero, Nat.and_zero]
      · simp only [Fin.last, h_m_last_val, mul_one]
        apply and_two_pow_eq_zero_of_getBit_0
        have h_getBit_prevSum_at_n := getBit_of_lt_two_pow (k:=n) (n:=n) (a:=⟨prevSum, by omega⟩)
        simp only [lt_self_iff_false, ↓reduceIte] at h_getBit_prevSum_at_n
        rw [h_getBit_prevSum_at_n]
    rw [h_sum_eq_xor, getBit_of_xor]
    if h_k_eq: k = n then
      rw [h_k_eq]
      simp only [lt_add_iff_pos_right, zero_lt_one, ↓reduceDIte]
      rw [h_prevSum_eq]
      rw [getBit_of_lt_two_pow]
      simp only [lt_self_iff_false, ↓reduceIte, zero_xor]
      unfold msbTerm
      -- ⊢ n.getBit (2 ^ ↑(Fin.last n) * m (Fin.last n)) = m ⟨n, ⋯⟩
      interval_cases h_m_last_val: m ⟨n, by omega⟩
      · -- ⊢ n.getBit (2 ^ ↑(Fin.last n) * m (Fin.last n)) = 0
        rw [Fin.val_last, Fin.last]
        rw [h_m_last_val, mul_zero]
        exact getBit_zero_eq_zero
      · -- ⊢ n.getBit (2 ^ ↑(Fin.last n) * m (Fin.last n)) = 1
        simp only [Fin.last]
        rw [h_m_last_val, mul_one]
        rw [Nat.getBit_two_pow]
        simp only [BEq.rfl, ↓reduceIte]
    else
      have hBitLhs := h_getBit_prevSum (k:=k)
      simp only at hBitLhs
      rw [h_prevSum_eq.symm] at hBitLhs
      rw [hBitLhs]
      if h_k_lt_n: k < n then
        have h_k_lt_n_add_1: k < n + 1 := by omega
        simp only [h_k_lt_n_add_1, ↓reduceDIte]
        push_neg at h_k_eq
        simp only [h_k_lt_n, ↓reduceDIte]
        unfold msbTerm
        interval_cases h_m_last_val: m ⟨n, by omega⟩
        · simp only [Fin.last, h_m_last_val, mul_zero]
          rw [Nat.getBit_zero_eq_zero, Nat.xor_zero]
          rfl
        · simp only [Fin.last, h_m_last_val, mul_one]
          rw [Nat.getBit_two_pow]
          simp only [beq_iff_eq]
          simp only [h_k_eq.symm, ↓reduceIte, xor_zero]
          rfl
      else
        have h_k_not_lt_n_add_1: ¬(k < n + 1) := by omega
        have h_k_not_lt_n: ¬(k < n) := by omega
        simp only [h_k_not_lt_n_add_1, h_k_not_lt_n, ↓reduceDIte, Nat.zero_xor]
        unfold msbTerm
        interval_cases h_m_last_val: m ⟨n, by omega⟩
        · simp only [Fin.last, h_m_last_val, mul_zero]
          exact getBit_zero_eq_zero
        · simp only [Fin.last, h_m_last_val, mul_one]
          rw [Nat.getBit_two_pow]
          simp only [beq_iff_eq]
          simp only [ite_eq_right_iff, one_ne_zero, imp_false, ne_eq]
          omega