-- CaseStudies/Velvet/VelvetExamples/SpMSpV_Example.lean

import Auto

import Lean

import Mathlib.Algebra.BigOperators.Intervals

import Mathlib.Algebra.Ring.Int.Defs

import Loom.MonadAlgebras.NonDetT.Extract

import Loom.MonadAlgebras.WP.Tactic

import Loom.MonadAlgebras.WP.DoNames'

import CaseStudies.Velvet.Std

section SpMV

structure SpV (valTyp : Type) where
  ind: Array Nat
  val: Array valTyp
  size: ℕ
  size_eq: ind.size = size ∧ val.size = size
  inc: ∀ (i j: Nat), i < size → j < size →  i < j → ind[i]! < ind[j]!

def sumUpTo
  (spv : SpV Int)
  (v : Array Int) (bound : ℕ) : Int := ∑ i ∈ Finset.range bound, ((spv.val)[i]! * v[(spv.ind)[i]!]!)

-- 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: structure SpV (valTyp : Type) where
  ind: Array Nat
  val: Array valTyp
  size: ℕ
  size_eq: ind.size = size ∧ val.size = size
  inc: ∀ (i j: Nat), i < size → j < size →  i < j → ind[i]! < ind[j]!

content: def sumUpTo
  (spv : SpV Int)
  (v : Array Int) (bound : ℕ) : Int := ∑ i ∈ Finset.range bound, ((spv.val)[i]! * v[(spv.ind)[i]!]!)

-- Local lemmas used:
content: theorem getValSpV_eq (spv: SpV Int) (j: ℕ) (h_ind: j < spv.size): spv[spv.ind[j]!] = (spv.val)[j]!

content: theorem getValSpV_empty (spv: SpV Int) (j: ℕ) (h_empty: ∀ i < spv.size, spv.ind[i]! ≠ j): spv[j] = 0

-- Target theorem and proof:
theorem VSpV_correct_pure (out: Array Int) (arr: Array Int)
  (spv: SpV Int)
  (h_b: ∀ i < spv.size, spv.ind[i]! < arr.size):
  out.size = 1 → out[0]! = sumUpTo spv arr spv.size →
    out[0]! = ∑ i ∈ Finset.range (arr.size), spv[i] * arr[i]!
:= by
      intro sz sum_eq
      rw [sum_eq, sumUpTo]
      have ind_lemma: ∀ k, k ≤ arr.size →
        ∑ i ∈ Finset.range (spv.size), (if spv.ind[i]! < k then spv.val[i]! * arr[spv.ind[i]!]! else 0) =
        ∑ i ∈ Finset.range k, spv[i] * arr[i]! := by
          intro k
          induction k with
          | zero =>
            simp
          | succ m hm =>
            intro lt_m
            simp [Finset.sum_range_succ]
            rw [←hm (by omega)]
            have splitted_sum:
              (∑ i ∈ Finset.range (spv.ind.size), if spv.ind[i]! < m then spv.val[i]! * arr[spv.ind[i]!]! else 0) +
              (∑ i ∈ Finset.range (spv.ind.size), if spv.ind[i]! = m then spv.val[i]! * arr[spv.ind[i]!]! else 0) =
              (∑ i ∈ Finset.range (spv.ind.size), if spv.ind[i]! < m + 1 then spv.val[i]! * arr[spv.ind[i]!]! else 0) := by
                rw [←Finset.sum_add_distrib]
                rw [Finset.sum_congr (by rfl)]
                intro x hx
                by_cases h_eq_m : spv.ind[x]! = m <;> simp [h_eq_m]
                have miff : spv.ind[x]! < m ↔ spv.ind[x]! < m + 1 := by
                  constructor <;> rintro h_lt <;> omega
                simp [miff]
            rw [←spv.size_eq.1, ←splitted_sum, add_left_cancel_iff.mpr]
            by_cases exists_i: ∃ i < spv.size, spv.ind[i]! = m
            { rcases exists_i with ⟨ind, h_ind⟩
              rw [← h_ind.right]
              have lemma_res := getValSpV_eq spv ind h_ind.left
              simp at lemma_res
              simp [lemma_res]
              have almost_zero : ∀ i ∈ Finset.range (spv.ind.size), i ≠ ind →
                ((if spv.ind[i]! = spv.ind[ind]! then spv.val[i]! * arr[spv.ind[i]!]! else 0) = 0) := by
                  intro i i_inb i_not_ind
                  by_cases vind_eq : spv.ind[i]! = spv.ind[ind]!
                  { simp at i_inb
                    have i_sz: i < spv.size := by
                      rw [spv.size_eq.1] at i_inb
                      exact i_inb
                    rcases lt_or_gt_of_ne i_not_ind with i_lt | inb_lt
                    { simp [getElem, Nat.lt_iff_le_and_ne.mp (spv.inc i ind i_sz h_ind.left i_lt)] }
                    have lt := Nat.lt_iff_le_and_ne.mp (spv.inc ind i h_ind.left i_sz inb_lt)
                    simp [getElem] at vind_eq
                    simp [vind_eq] at lt }
                  simp [vind_eq]
              have ind_inb: ind ∉ Finset.range (spv.size) →
                ((if spv.ind[ind]! = spv.ind[ind]! then spv.val[ind]! * arr[spv.ind[ind]!]! else 0) = 0) := by
                  intro ind_not_inb
                  simp at ind_not_inb
                  omega
              rw [←spv.size_eq.1] at ind_inb
              simp [Finset.sum_eq_single ind almost_zero ind_inb] }
            simp at exists_i
            have h_getVal := getValSpV_empty spv m exists_i
            simp at h_getVal
            simp [h_getVal]
            apply Finset.sum_eq_zero
            rintro x hx
            by_cases h_eq: spv.ind[x]! = m <;> simp [h_eq]
            simp [←h_eq] at h_getVal
            simp at hx
            rw [spv.size_eq.1] at hx
            simp [getValSpV_eq spv x hx] at h_getVal
            simp [h_getVal]
      have fin_lemma := ind_lemma (arr.size) (by rfl)
      rw [←fin_lemma]
      exact Finset.sum_congr (by rfl) fun i h_i => by aesop