-- ArkLib/Data/MlPoly/Equiv.lean

import ArkLib.Data.MlPoly.Basic

import ArkLib.Data.MvPolynomial.Notation

open MvPolynomial

variable {R : Type*} [CommRing R] {n : ℕ}

noncomputable section

namespace MlPoly

noncomputable def monomialOfNat (i : ℕ) : (Fin n) →₀ ℕ :=
  Finsupp.onFinset (s:=Finset.univ (α:=Fin n)) (fun j => Nat.getBit j.val i) (by admit /- proof elided -/
    ) 

def toMvPolynomial (p : MlPoly R n) : R[X Fin n] :=
  ∑ i : Fin (2 ^ n), MvPolynomial.monomial (monomialOfNat i) (a:=p[i])

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

content: @[reducible]
def MlPoly (R : Type*) (n : ℕ) := Vector R (2 ^ n) 

content: def binaryFinMapToNat {n : ℕ} (m : Fin n → ℕ) (h_binary : ∀ j: Fin n, m j ≤ 1) : Fin (2^n) :=

-- Repository 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 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 getBit_eq_testBit (k n : Nat) : getBit k n = if n.testBit k then 1 else 0

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

content: lemma and_two_pow_eq_zero_of_getBit_0 {n i : ℕ} (h_getBit : getBit i n = 0)
    : n &&& (2 ^ i) = 0

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 shiftRight_and_one_distrib {n m k : ℕ} :
    Nat.getBit k (n &&& m) = Nat.getBit k n &&& Nat.getBit k m

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

content: lemma getBit_of_xor {n m k: ℕ} : getBit k (n ^^^ m) = getBit k n ^^^ getBit k m

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

content: lemma sum_of_and_eq_zero_is_xor {n m : ℕ} (h_n_AND_m : n &&& m = 0) : n + m = n ^^^ m

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

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 div_2_form {nD2 b : ℕ} (h_b : b < 2):
  (nD2 * 2 + b) / 2 = nD2

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

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)

-- Local definitions used:
content: noncomputable def monomialOfNat (i : ℕ) : (Fin n) →₀ ℕ :=
  Finsupp.onFinset (s:=Finset.univ (α:=Fin n)) (fun j => Nat.getBit j.val i) (by admit /- proof elided -/
    )

content: def toMvPolynomial (p : MlPoly R n) : R[X Fin n] :=
  ∑ i : Fin (2 ^ n), MvPolynomial.monomial (monomialOfNat i) (a:=p[i])

-- Local lemmas used:
content: theorem eq_monomialOfNat_iff_eq_bitRepr (m : Fin n →₀ ℕ)
  (h_binary : ∀ j : Fin n, m j ≤ 1) (i: Fin (2^n)) :
  monomialOfNat i = m ↔ i = Nat.binaryFinMapToNat m h_binary

content: theorem toMvPolynomial_is_multilinear (p : MlPoly R n) :
  (toMvPolynomial p) ∈ R⦃≤ 1⦄[X Fin n]

-- Target theorem and proof:
theorem coeff_of_toMvPolynomial_eq_coeff_of_MlPoly (p : MlPoly R n) (m : Fin n →₀ ℕ) :
  coeff m (toMvPolynomial p) =
    if h_binary: (∀ j: Fin n, m j ≤ 1) then
      let i_of_m: ℕ := Nat.binaryFinMapToNat (m:=m) (h_binary:=h_binary)
      p[i_of_m]
    else
      0
:= Nat.binaryFinMapToNat (m:=m) (h_binary:=h_binary)
      p[i_of_m]
    else
      0
  := by
  if h_binary: (∀ j: Fin n, m j ≤ 1) then
    unfold toMvPolynomial
    simp only [h_binary, implies_true, ↓reduceDIte]
    let i_of_m := Nat.binaryFinMapToNat m h_binary
    have h_mono_eq : monomialOfNat i_of_m = m := by
      ext j; simp only [monomialOfNat, Finsupp.onFinset_apply]
      have h_getBit := Nat.getBit_of_binaryFinMapToNat (n:=n) (m:=m)
        (h_binary:=h_binary) (k:=j)
      rw [h_getBit]
      simp only [j.isLt, ↓reduceDIte, Fin.eta]
    rw [MvPolynomial.coeff_sum]
    simp only [MvPolynomial.coeff_monomial]
    -- ⊢ (∑ x, if monomialOfNat ↑x = m then p[x] else 0) = p[↑(Nat.binaryFinMapToNat ⇑m ⋯)]
    set f := fun x: Fin (2^n) => if monomialOfNat x.val = m then p[x] else (0: R)
    -- ⊢ Finset.univ.sum f = p[↑(Nat.binaryFinMapToNat ⇑m ⋯)]
    rw [Finset.sum_eq_single (a:=⟨i_of_m, by omega⟩)]
    · -- Goal 1: Prove the main term is correct.
      simp only [h_mono_eq, ↓reduceIte, Fin.eta, Fin.getElem_fin];
      rfl
    · -- Goal 2: Prove all other terms are zero.
      intro j h_j_mem_univ h_ji_ne
      -- If `j ≠ i_of_m`, then `monomialOfNat j ≠ monomialOfNat i_of_m` (which is `m`).
      -- ⊢ (monomial (monomialOfNat ↑j)) p[j] = 0
      have h_mono_ne : monomialOfNat j.val ≠ m := by
        intro h_eq_contra
        have h_j_is_i_of_m := eq_monomialOfNat_iff_eq_bitRepr (m:=m)
          (h_binary:=h_binary) (i:=j).mp h_eq_contra
        exact h_ji_ne h_j_is_i_of_m
      simp only [h_mono_ne, ↓reduceIte]
    -- Goal 3: Prove `i` is in the summation set.
    · simp [Finset.mem_univ]
  else -- this case is similar to the proof of `right_inv` in `equivMvPolynomialDeg1`
    simp only [h_binary, ↓reduceDIte]
    -- ⊢ coeff m p.toMvPolynomial = 0
    have hv := toMvPolynomial_is_multilinear p
    let vMlPoly: R⦃≤ 1⦄[X Fin n] := ⟨p.toMvPolynomial, hv⟩
    have h_v_coeff_zero : vMlPoly.val.coeff m = 0 := by
      refine notMem_support_iff.mp ?_
      by_contra h_mem_support
      have hvMlPoly := vMlPoly.2
      rw [MvPolynomial.mem_restrictDegree] at hvMlPoly
      have h_deg_le_one: ∀ j: Fin n, (m j) ≤ 1 := by
        exact fun j ↦ hvMlPoly m h_mem_support j
      simp only [not_forall, not_le] at h_binary -- h_binary : ∃ x, 1 < m x
      obtain ⟨j, hj⟩ := h_binary
      have h_not_1_lt_m_j: ¬(1 < m j) := by exact Nat.not_lt.mpr (hv h_mem_support j)
      exact h_not_1_lt_m_j hj
    exact h_v_coeff_zero