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

-- Repository definitions used:
content: class MAlgPartial (m : Type u -> Type v) [Monad m] [∀ α, Lean.Order.CCPO (m α)]
  [CompleteLattice l] [MAlgOrdered m l] where
  csup_lift {α : Type u} (xc : Set (m α)) (post : α -> l) :
    Lean.Order.chain xc ->
    ⨅ x ∈ xc, MAlg.lift x post <= MAlg.lift (Lean.Order.CCPO.csup xc) 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: 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

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

content: omit [MonoBind m] [LawfulMonad m] in
lemma wp_csup (xc : Set (m α)) (post : α -> l) :
  Lean.Order.chain xc ->
  ⨅ c ∈ xc, wp c post ≤ wp (Lean.Order.CCPO.csup xc) post

-- Target theorem and proof:
lemma repeat_inv (f : Unit -> β -> m (ForInStep β))
  (inv : ForInStep β -> l)
  init :
   (∀ b, triple (inv (.yield b)) (f () b) (inv)) ->
   triple (inv (.yield init)) (Loop.forIn.loop f init) (fun b => inv (.done b))
:= by
  intro hstep
  revert init
  apply Loop.forIn.loop.fixpoint_induct (f := f) (motive :=
    fun loop => ∀ init, triple (inv (.yield init)) (loop init) (fun b =>inv (.done b)))
  { apply Lean.Order.admissible_pi_apply
      (P := fun init loop => triple (inv (.yield init)) (loop) (fun b =>inv (.done b)))
    simp [admissible, triple]; intro init loops cl h
    apply le_trans'; apply wp_csup; solve_by_elim
    simp; solve_by_elim }
  intro loop ih init; simp [triple, wp_bind]; apply le_trans; apply hstep
  apply wp_cons; rintro (_|_); simp [wp_pure]
  apply ih