-- 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 VSpV_correct_triple (out: Array Int) (arr: Array Int) (spv: SpV Int):
  triple
    (∀ i < spv.size, spv.ind[i]! < arr.size)
    (VSpV out arr spv)
    fun ⟨_, outNew⟩ =>
      outNew[0]! = ∑ i ∈ Finset.range (arr.size), spv[i] * arr[i]!
:= by
      simp [triple]
      intro h_b
      apply wp_cons (VSpV out arr spv)
        fun ⟨u, outNew⟩ =>
          (outNew[0]! = sumUpTo spv arr spv.size ∧ outNew.size = 1)
      { rintro ⟨u, outNew⟩; simp
        intro sum_eq sz
        exact VSpV_correct_pure outNew arr spv h_b sz sum_eq }
      simp
      have triple_true := VSpV_correct out arr spv
      simp [triple] at triple_true
      apply triple_true