-- Loom/MonadAlgebras/WP/Basic.lean

import Loom.MonadAlgebras.Defs

import Loom.MonadAlgebras.Instances.Basic

import Loom.MonadAlgebras.Instances.ExceptT

import Loom.MonadAlgebras.Instances.StateT

import Loom.MonadAlgebras.Instances.ReaderT

import Loom.MonadAlgebras.Instances.Gen

variable {m : Type u -> Type v} [Monad m] [LawfulMonad m] {α : Type u} {l : Type u}

section

variable [CompleteLattice l]

section

variable [mprop : MAlgOrdered m l]

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

def triple (pre : l) (c : m α) (post : α -> l) : Prop :=
  pre ≤ wp c post

end

variable [MAlgOrdered m l]

end

section

variable [CompleteLattice l] [MAlgOrdered m l]

noncomputable

end

section Determinism

variable [inst: CompleteLattice l] [MAlgOrdered m l]

end Determinism

section Loops

open Lean.Order

@[specialize, inline]
def Loop.forIn.loop {m : Type u -> Type v} [Monad m] [∀ α, CCPO (m α)] [MonoBind m] (f : Unit → β → m (ForInStep β)) (b : β) : m β := do
    match ← f () b with
      | ForInStep.done b  => pure b
      | ForInStep.yield b => loop f b
  partial_fixpoint

@[inline]
def Loop.forIn {β : Type u} [Monad m] [∀ α, CCPO (m α)] [MonoBind m]
  (_ : Lean.Loop) (init : β) (f : Unit → β → m (ForInStep β)) : m β :=
  Loop.forIn.loop f init

variable [inst: _root_.CompleteLattice l] [MAlgOrdered m l]

namespace PartialCorrectness

variable [∀ α, CCPO (m α)] [MonoBind m] [MAlgPartial m]

end PartialCorrectness

namespace TotalCorrectness

variable [∀ α, CCPO (m α)] [MonoBind m]

-- Repository definitions used:
content: macro "⌜" p:term "⌝" : term => `(LE.pure $p)

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: abbrev MAlg.lift {m : Type u -> Type v} {l : Type u} [Monad m] [MAlg m l] :
  {α : Type u} -> m α -> Cont l α := fun x f => μ $ f <$> x

content: class MAlg [Monad m] (l : outParam (Type v)) where
  μ : m l -> l
  pure : ∀ l, μ (pure l) = l
  bind : ∀ {α : Type v} (x : m α) (f g : α -> m l),
    μ ∘ f = μ ∘ g ->
    μ (x >>= f) = μ (x >>= g)

content: class MAlgOrdered (l : outParam (Type v)) [Monad m] [CompleteLattice l] where
  μ : m l -> l
  μ_ord_pure : ∀ l, μ (pure l) = l
  μ_ord_bind {α : Type v} :
    ∀ (f g : α -> m l), μ ∘ f ≤ μ ∘ g ->
      ∀ x : m α, μ (x >>= f) ≤ μ (x >>= g)

content: @[simp]
def Cont.monotone {t : Type v} {α : Type u} [Preorder t] (wp : Cont t α) :=
  ∀ (f f' : α -> t), (∀ a, f a ≤ f' a) → wp f ≤ wp f'

content: abbrev Cont (t : Type v) (α : Type u) := (α -> t) -> t

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

-- Repository lemmas used:
content: lemma Cont.monotone_lift {l : Type u} {m : Type u -> Type v} [Monad m] [LawfulMonad m] [CompleteLattice l] [MAlgOrdered m l] :
  ∀ {α : Type u} (x : m α), MAlg.lift x |>.monotone

content: lemma triple_bind {β} (pre : l) (x : m α) (cut : α -> l)
  (f : α -> m β) (post : β -> l) :
  triple pre x cut ->
  (∀ y, triple (cut y) (f y) post) ->
  triple pre (x >>= f) post

content: lemma triple_pure (pre : l) (x : α) (post : α -> l) :
  triple pre (pure (f := m) x) post <-> pre ≤ (post x)

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

content: def triple (pre : l) (c : m α) (post : α -> l) : Prop :=
  pre ≤ wp c post

content: @[specialize, inline]
def Loop.forIn.loop {m : Type u -> Type v} [Monad m] [∀ α, CCPO (m α)] [MonoBind m] (f : Unit → β → m (ForInStep β)) (b : β) : m β := do
    match ← f () b with
      | ForInStep.done b  => pure b
      | ForInStep.yield b => loop f b
  partial_fixpoint

content: @[inline]
def Loop.forIn {β : Type u} [Monad m] [∀ α, CCPO (m α)] [MonoBind m]
  (_ : Lean.Loop) (init : β) (f : Unit → β → m (ForInStep β)) : m β :=
  Loop.forIn.loop f init

-- Local lemmas used:
content: lemma wp_pure (x : α) (post : α -> l) : wp (m := m) (pure x) post = post x

content: lemma wp_bind {β} (x : m α) (f : α -> m β) (post : β -> l) :
    wp (x >>= f) post = wp x (fun x => wp (f x) post)

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

-- Target theorem and proof:
lemma repeat_inv (f : Unit -> β -> m (ForInStep β))
  (inv : ForInStep β -> l) (measure : β -> Nat)
  init :
   (∀ b, triple (inv (.yield b)) (f () b) (fun | .yield b' => inv (.yield b') ⊓ ⌜ measure b' < measure b ⌝ | .done b' => inv (.done b'))) ->
   triple (inv (.yield init)) (Loop.forIn.loop f init) (fun b => inv (.done b))
:= by
  intro hstep
  have induc (C : β → Prop) (a : β) (h : ∀ x, (∀ y, measure y < measure x → C y) → C x): C a := by
    have lem: ∀ n: Nat, ∀ x, measure x ≤ n →  C x := by
      intro n
      induction n with
      | zero =>
        intro x hx
        exact h x (fun y => by omega)
      | succ m ih =>
        intro x hx
        by_cases neq: measure x ≤ m
        { exact ih x neq }
        have eq: measure x = m + 1 := by omega
        exact h x (fun y hy => ih y (by omega))
    exact lem (measure a) a (by simp)
  apply induc
    (fun ini => triple (inv (.yield ini)) (Loop.forIn.loop f ( ini)) (fun b => inv (.done b))) init
  intro b ih; unfold Loop.forIn.loop; simp [triple, wp_bind]; apply le_trans
  apply hstep; apply wp_cons; rintro (_|_)
  { simp [wp_pure] }
  rename_i a
  match ForInStep.yield a with
  | .yield b' =>
    simp; intro m_lt
    have ihb := ih b' m_lt
    simp [triple] at ihb
    exact ihb
  | .done b' =>
    simp
    rw [wp_pure b' (fun s: β => inv (ForInStep.done s))]