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

-- Repository definitions used:
content: ...

-- Repository lemmas used:
content: theorem sum_Icc_split {α : Type*} [AddCommMonoid α] (f : ℕ → α) (a b c : ℕ)
    (h₁ : a ≤ b) (h₂ : b ≤ c):
    ∑ i ∈ Finset.Icc a c, f i = ∑ i ∈ Finset.Icc a b, f i + ∑ i ∈ Finset.Icc (b+1) c, f i

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

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

-- Target theorem and proof:
theorem getBit_repr {ℓ : Nat} : ∀ j, j < 2^ℓ →
  j = ∑ k ∈ Finset.Icc 0 (ℓ-1), (getBit k j) * 2^k
:= by
  induction ℓ with
  | zero =>
    -- Base case : ℓ = 0
    intro j h_j
    have h_j_zero : j = 0 := by exact Nat.lt_one_iff.mp h_j
    subst h_j_zero
    simp only [zero_tsub, Finset.Icc_self, Finset.sum_singleton, pow_zero, mul_one]
    unfold getBit
    rw [Nat.shiftRight_zero, Nat.and_one_is_mod]
  | succ ℓ₁ ih =>
    by_cases h_ℓ₁ : ℓ₁ = 0
    · simp only [h_ℓ₁, zero_add, pow_one, tsub_self, Finset.Icc_self, Finset.sum_singleton,
      pow_zero, mul_one];
      intro j hj
      interval_cases j
      · simp only [getBit, Nat.shiftRight_zero, Nat.and_one_is_mod, Nat.zero_mod]
      · simp only [getBit, Nat.shiftRight_zero, Nat.and_one_is_mod]
    · push_neg at h_ℓ₁
      set ℓ := ℓ₁ + 1
      have h_ℓ_eq : ℓ = ℓ₁ + 1 := by rfl
      intro j h_j
      -- Inductive step : assume theorem holds for ℓ₁ = ℓ - 1
      -- => show j = ∑ k ∈ Finset.range (ℓ + 1), (getBit k j) * 2^k
      -- Split j into lowBits (b) and higher getLowBits (m) &
      -- reason inductively from the predicate of (m, ℓ₁)
      set b := getBit 0 j -- Least significant getBit : j % 2
      set m := j >>> 1 -- Higher getLowBits : j / 2
      have h_b_eq : b = getBit 0 j := by rfl
      have h_m_eq : m = j >>> 1 := by rfl
      have h_getBit_shift : ∀ k, getBit (k+1) j = getBit k m := by
        intro k
        rw [h_m_eq]
        exact (getBit_of_shiftRight (n := j) (p := 1) k).symm
      have h_j_eq : j = b + 2 * m := by
        calc
          _ = 2 * m + b := by
            have h_m_eq : m = j/2 := by rfl
            have h_b_eq : b = j%2 := by
              rw [h_b_eq]; unfold getBit; rw [Nat.shiftRight_zero]; rw [Nat.and_one_is_mod];
            rw [h_m_eq, h_b_eq];
            rw [Nat.div_add_mod (m := j) (n := 2)]; -- n * (m / n) + m % n = m := by
          _ = b + 2 * m := by omega;
      have h_m : m < 2^ℓ₁ := by
        by_contra h_m_ge_2_pow_ℓ
        push_neg at h_m_ge_2_pow_ℓ
        have h_j_ge : j ≥ 2^ℓ := by
          calc _ = 2 * m + b := by rw [h_j_eq]; omega
            _ ≥ 2 * (2^ℓ₁) + b := by omega
            _ = 2^ℓ + b := by rw [h_ℓ_eq]; omega;
            _ ≥ 2^ℓ := by omega;
        exact Nat.not_lt_of_ge h_j_ge h_j -- contradiction
      have h_m_repr := ih (j := m) h_m
      have getBit_shift : ∀ k, getBit (k + 1) j = getBit k m := by
        intro k
        rw [h_m_eq]
        exact (getBit_of_shiftRight (n := j) (p := 1) k).symm
      -- ⊢ j = ∑ k ∈ Finset.range ℓ, getBit k j * 2 ^ k
      have h_sum : ∑ k ∈ Finset.Icc 0 (ℓ-1), getBit k j * 2 ^ k
        = (∑ k ∈ Finset.Icc 0 0, getBit k j * 2 ^ k)
        + (∑ k ∈ Finset.Icc 1 (ℓ-1), getBit k j * 2 ^ k) := by
        apply sum_Icc_split
        omega
        omega
      rw [h_sum]
      rw [h_j_eq]
      rw [Finset.Icc_self, Finset.sum_singleton, pow_zero, mul_one]

      have h_sum_2 : ∑ k ∈ Finset.Icc 1 (ℓ-1), getBit k (b + 2 * m) * 2 ^ k
        = ∑ k ∈ Finset.Icc 0 (ℓ₁-1), getBit k (m) * 2 ^ (k+1) := by
        apply Finset.sum_bij' (fun i _ => i - 1) (fun i _ => i + 1)
        · -- left inverse
          intro i hi
          simp only [Finset.mem_Icc] at hi ⊢
          exact Nat.sub_add_cancel hi.1
        · -- right inverse
          intro i hi
          norm_num
        · -- function value match
          intro i hi
          rw [←h_j_eq]
          rw [getBit_of_shiftRight]
          have ⟨left_bound, right_bound⟩ := Finset.mem_Icc.mp hi
          rw [Nat.sub_add_cancel left_bound]
        · -- left membership preservation
          intro i hi -- hi : i ∈ Finset.Icc 1 (ℓ - 1)
          rw [Finset.mem_Icc]
          have ⟨left_bound, right_bound⟩ := Finset.mem_Icc.mp hi
          -- ⊢ 0 ≤ i - 1 ∧ i - 1 ≤ ℓ₁ - 1
          apply And.intro
          · exact Nat.pred_le_pred left_bound
          · exact Nat.pred_le_pred right_bound
        · -- right membership preservation
          intro j hj
          rw [Finset.mem_Icc]
          have ⟨left_bound, right_bound⟩ := Finset.mem_Icc.mp hj -- (0 ≤ j ∧ j ≤ ℓ₁ - 1)
          -- ⊢ 1 ≤ j + 1 ∧ j + 1 ≤ ℓ - 1
          apply And.intro
          · exact Nat.le_add_of_sub_le left_bound
          · rw [h_ℓ_eq]; rw [Nat.add_sub_cancel]; -- ⊢ j + 1 ≤ ℓ₁
            have h_j_add_1_le_ℓ₁ : j + 1 ≤ ℓ₁ := by
              calc j + 1 ≤ (ℓ₁ - 1) + 1 := by apply Nat.add_le_add_right; exact right_bound;
              _ = ℓ₁ := by rw [Nat.sub_add_cancel]; omega;
            exact h_j_add_1_le_ℓ₁
      rw [h_sum_2]

      have h_sum_3 : ∑ k ∈ Finset.Icc 0 (ℓ₁-1), getBit k (m) * 2 ^ (k+1)
        = 2 * ∑ k ∈ Finset.Icc 0 (ℓ₁-1), getBit k (m) * 2 ^ k := by
        calc
          _ = ∑ k ∈ Finset.Icc 0 (ℓ₁-1), ((getBit k (m) * 2^k) * 2) := by
            apply Finset.sum_congr rfl (fun k hk => by
              rw [Finset.mem_Icc] at hk -- hk : 0 ≤ k ∧ k ≤ ℓ₁ - 1
              have h_res : getBit k (m) * 2 ^ (k+1) = getBit k (m) * 2 ^ k * 2 := by
                rw [Nat.pow_succ, ←mul_assoc]
              exact h_res
            )
        _ = (∑ k ∈ Finset.Icc 0 (ℓ₁-1), getBit k (m) * 2 ^ k) * 2 := by rw [Finset.sum_mul]
        _ = 2 * ∑ k ∈ Finset.Icc 0 (ℓ₁-1), getBit k (m) * 2 ^ k := by rw [mul_comm]
      rw [h_sum_3]
      rw [←h_m_repr]
      conv =>
        rhs
        rw [←h_j_eq]