module
public import Mathlib.Logic.Equiv.Fin.Basic
public import Mathlib.Algebra.BigOperators.Group.Finset.Basic
public import Mathlib.Tactic

@[expose] public section
noncomputable section
namespace VibeQCCore
open scoped BigOperators

variable {AO Aux : Type*} [Fintype AO] [Fintype Aux]

def gramERI (B : Aux → AO → AO → ℝ) (i j k l : AO) : ℝ :=
  ∑ p, B p i j * B p k l

def denseJ (B : Aux → AO → AO → ℝ) (D : AO → AO → ℝ) (i j : AO) : ℝ :=
  ∑ k, ∑ l, D k l * gramERI B i j k l

def stagedJ (B : Aux → AO → AO → ℝ) (D : AO → AO → ℝ) (i j : AO) : ℝ :=
  ∑ p, B p i j * (∑ k, ∑ l, B p k l * D k l)

/-- Full finite-sum interchange, without positivity or ERI symmetry hypotheses. -/
theorem df_staged_equals_dense (B : Aux → AO → AO → ℝ) (D : AO → AO → ℝ)
    (i j : AO) : stagedJ B D i j = denseJ B D i j := by
  simp only [stagedJ, denseJ, gramERI, Finset.mul_sum]
  rw [Finset.sum_comm]
  apply Finset.sum_congr rfl
  intro k hk
  rw [Finset.sum_comm]
  apply Finset.sum_congr rfl
  intro l hl
  apply Finset.sum_congr rfl
  intro p hp
  ring

/-- Last-axis-fast pair packing used by a logical n by m matrix. -/
def packPair {n m : ℕ} (p : Fin n × Fin m) : Fin (n * m) := finProdFinEquiv p

def unpackPair {n m : ℕ} (p : Fin (n * m)) : Fin n × Fin m := finProdFinEquiv.symm p

theorem unpack_pack_pair {n m : ℕ} (p : Fin n × Fin m) :
    unpackPair (packPair p) = p := finProdFinEquiv.symm_apply_apply p

theorem pack_unpack_pair {n m : ℕ} (p : Fin (n * m)) :
    packPair (unpackPair p) = p := finProdFinEquiv.apply_symm_apply p

theorem pack_pair_offset {n m : ℕ} (p : Fin n × Fin m) :
    (packPair p).val = p.1.val * m + p.2.val := by
  simp [packPair, finProdFinEquiv, Nat.add_comm, Nat.mul_comm]

end VibeQCCore