-- 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: def triple (pre : l) (c : m α) (post : α -> l) : Prop :=
  pre ≤ wp c post

content: def wp (c : m α) (post : α -> l) : l := liftM (n := Cont l) c post

content: class LogicLift (l : outParam (Type u)) ( k : Type u) [CompleteLattice l] [CompleteLattice k] where
  [lift : MonadLift (Cont l) (Cont k)]
  lift_top {α : Type u} :
    monadLift (m := Cont l) (n := Cont k) (fun (_ : α -> l) => ⊤) = ⊤
  lift_bot {α : Type u} :
    monadLift (m := Cont l) (n := Cont k) (fun (_ : α -> l) => ⊥) = ⊥

content: notation "{" P "}" c "{" v "," Q "}" => triple P c (fun v => Q)

-- Repository lemmas used:
content: lemma wp_cons (x : m α) (post post' : α -> l) :
  (∀ y, post y ≤ post' y) ->
  wp x post ≤ wp x post'

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

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

-- Target theorem and proof:
theorem spmv_correct_triple (out: Array Int) (arr: Array Int) (spm: Array (SpV Int)):
  triple
    (∀ i < spm.size, ∀ j < spm[i]!.size, spm[i]!.ind[j]! < arr.size)
    (spmv out arr spv spm)
    fun ⟨_, outNew⟩ =>
      (∀ j < outNew.size, outNew[j]! = ∑ i ∈ Finset.range (arr.size), spm[j]![i] * arr[i]!)
:= by
      simp [triple]
      intro h_b
      apply wp_cons
        (spmv out arr spv spm)
        fun ⟨_, outNew⟩ =>
          ((∀ i < spm.size, outNew[i]! = sumUpTo spm[i]! arr spm[i]!.size) ∧ spm.size = outNew.size)
      { simp; rintro outNew;
        intro sum_eq sz_eq j h_j
        have single_elem : (Array.replicate 1 outNew[j]!)[0]! = outNew[j]! := by
          simp
          -- simp [getElem!] at replicate_get
          -- simp [getElem, replicate_get]
        have single_th := VSpV_correct_pure
          (Array.replicate 1 outNew[j]!)
          arr
          spm[j]!
          (h_b j (by rw [←sz_eq] at h_j; exact h_j))
          (by simp)
          (by simp; exact sum_eq j (by rw [←sz_eq] at h_j; exact h_j))
        simp at single_th
        simp [←single_th] }
      simp
      have triple_true := spmv_correct out arr spv spm
      simp [triple] at triple_true
      exact triple_true