-- 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 getHighBits_no_shl (numLowBits : ℕ) (n : ℕ) : ℕ := n >>> numLowBits

def joinBits {n m : ℕ} (low : Fin (2 ^ n)) (high : Fin (2 ^ m)) : Fin (2 ^ (m+n)) :=
  ⟨(high.val <<< n) ||| low.val, by admit /- proof elided -/
  ⟩

-- 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 getHighBits_no_shl (numLowBits : ℕ) (n : ℕ) : ℕ := n >>> numLowBits

content: def joinBits {n m : ℕ} (low : Fin (2 ^ n)) (high : Fin (2 ^ m)) : Fin (2 ^ (m+n)) :=
  ⟨(high.val <<< n) ||| low.val, by admit /- proof elided -/
  ⟩

-- 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 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 add_shiftRight_distrib {n m k : ℕ} (h_and_zero : n &&& m = 0):
  (n + m) >>> k = (n >>> k) + (m >>> k)

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_lt_two_pow {n: ℕ} (a: Fin (2^n)) (k: ℕ) :
  getBit k a = if k < n then getBit k a else 0

content: lemma and_shl_eq_zero_of_lt_two_pow {a n b : ℕ} (hb : b < 2 ^ n) : (a <<< n) &&& b = 0

-- Target theorem and proof:
lemma getHighBits_no_shl_joinBits {n m : ℕ} (low : Fin (2 ^ n)) (high : Fin (2 ^ m)) :
  getHighBits_no_shl n (joinBits low high).val = high.val
:= by
  unfold joinBits getHighBits_no_shl
  dsimp
  have h_and_zero := and_shl_eq_zero_of_lt_two_pow (a := high.val) (b := low.val) (hb := low.isLt)
  rw [←Nat.sum_of_and_eq_zero_is_or h_and_zero]
  rw [Nat.add_shiftRight_distrib h_and_zero]
  rw [Nat.shiftLeft_shiftRight]
  rw [Nat.shiftRight_eq_div_pow]
  have h: low.val/2^n = 0 := by
    apply Nat.div_eq_zero_iff_lt (x:=low) (k:=2^n) (h:=by exact Nat.two_pow_pos n).mpr (by omega)
  simp only [h, add_zero]