-- CaseStudies/Velvet/VelvetTheory.lean

import Loom.MonadAlgebras.NonDetT.Extract

notation "[totl|" t "]" => open TotalCorrectness TotalCorrectness.DemonicChoice in t

notation "[part|" t "]" => open PartialCorrectness PartialCorrectness.DemonicChoice in t

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

content: notation "[part|" t "]" => open PartialCorrectness PartialCorrectness.DemonicChoice in t

content: notation "[part|" t "]" => open ExceptionAsSuccess in t

content: notation "[totl|" t "]" => open TotalCorrectness TotalCorrectness.DemonicChoice in t

content: notation "[totl|" t "]" => open ExceptionAsFailure in t

content: noncomputable def LE.pure {l : Type u} [inst: LE l] [OrderTop l] [OrderBot l] : Prop -> l := fun p =>
  if p then ⊤ else ⊥

content: inductive DivM (α : Type u) where
  | res (x : α)
  | div

content: def spec (pre : l) (post : α -> l) : Cont l α :=
  fun p => pre ⊓ ⌜post ≤ p⌝

content: inductive NonDetT (m : Type u -> Type v) : (α : Type u) -> Type _ where
  | pure {α} (ret : α) : NonDetT m α
  | vis {α} {β} (x : m β) (f : β → NonDetT m α) : NonDetT m α
  | pickCont {α} (τ : Type u) (p : τ -> Prop) [Findable p] (f : τ → NonDetT m α) : NonDetT m α
  | repeatCont {α} {β} (init : β) (f : β -> NonDetT m (ForInStep β)) (cont : β -> NonDetT m α) : NonDetT m α

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: 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: def NonDetT.wp {l : Type u} {α : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m α -> Cont l α
  | .pure ret => pure ret
  | .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | .pickCont τ p f => fun post => let p : Set τ := p; ⨅ a ∈ (p : Set τ), wp (f a) post

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: def NonDetT.wp {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : {α : Type u} -> NonDetT m α -> Cont l α
  | _, .pure ret => pure ret
  | _, .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | _, @NonDetT.pickCont _ _ τ p _ f =>
  fun post =>
    let p : Set τ := p;
    ⨅ a ∈ (p : Set τ), wp (f a) post
  | _, @NonDetT.repeatCont _ _ β init f cont => fun post => ⨆ (inv : ForInStep β -> l),
      ⌜ ∀ b, (inv (ForInStep.yield b)) <= wp (f b) inv⌝ ⊓
      spec (inv (.yield init)) (fun b => inv (.done b)) (fun b => wp (cont b) post)

content: def NonDetT.μ  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m l -> l := fun x => NonDetT.wp x id

content: def NonDetT.bind (x : NonDetT m α) (f : α → NonDetT m β) : NonDetT m β :=
  match x with
  | pure ret => f ret
  | vis x f' => vis x fun y => bind (f' y) f
  | pickCont τ p f' => pickCont τ p fun t => bind (f' t) f
  | repeatCont init f' cont => repeatCont init f' fun t => bind (cont t) f

content: def WPGen.bind {x : m α} {f : α -> m β} (wpg : WPGen x) (wpgf : ∀ a, WPGen (f a)) :
  WPGen (x >>= f) where
  get := fun post => wpg.get (fun a => (wpgf a).get post)
  prop := by admit /- proof elided -/

content: private def _root_.Lean.SimpleScopedEnvExtension.get [Inhabited σ] (ext : SimpleScopedEnvExtension α σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

content: structure Context where
  ref         : Syntax
   
  m           : Syntax
   
  returnType  : Syntax
  mutableVars : VarSet := {}
  insideFor   : Bool := false

content: private def _root_.Lean.SimplePersistentEnvExtension.get [Inhabited σ] (ext : SimplePersistentEnvExtension α σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

content: structure WPGen (x : m α) where
  get : Cont l α
  
  prop : ∀ post, get post <= wp x post

content: private def _root_.Lean.EnvExtension.get [Inhabited σ] (ext : EnvExtension σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

content: class CCPOBotLawful (m : Type u -> Type v) [∀ α, Lean.Order.CCPO (m α)] [CCPOBot m] where
  prop {α} : CCPOBot.compBot (m := m) (α := α) = Lean.Order.bot

-- Repository lemmas used:
content: lemma TotalCorrectness.DivM.wp_eq (α : Type) (x : DivM α) (post : α -> Prop) :
  wp x post =
    match x with

content: lemma PartialCorrectness.DivM.wp_eq (α : Type) (x : DivM α) (post : α -> Prop) :
  wp x post =
    match x with

content: lemma NonDetT.wp_eq_wp {α : Type u} (x : NonDetT m α) (post : α -> l) :
  _root_.wp x post = NonDetT.wp x post

content: lemma NonDetT.wp_bind  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] [LawfulMonad m] {α β : Type u} (x : NonDetT m α) (f : α -> NonDetT m β)
  (post : β -> l):
  NonDetT.wp (x.bind f) post = NonDetT.wp x (fun x => NonDetT.wp (f x) post)

-- Local definitions used:

-- Local lemmas used:

-- Target theorem and proof:
@[local simp]
lemma wp_mono_part (x : NonDetT DivM α) (post₁ post₂ : α -> Prop) :
  (post₁ ≤ post₂) → ([totl|wp x post₁]) ≤ ([part| wp x post₂])
:= by
    intro le
    simp [loomLogicSimp]
    simp [loomLogicSimp] at le
    unhygienic induction x <;> try simp [loomLogicSimp]
    { exact le ret }
    { simp [[totl| DivM.wp_eq], [part| DivM.wp_eq]]
      intro wp1
      match x_1 with
      | .div => trivial
      | .res a => simp; simp at wp1; exact f_ih a post₁ post₂ le wp1 }
    { intro _ _ wp1 i hi
      exact f_ih i post₁ post₂ le (wp1 i hi) }
    intro x x_1 h1 sp1
    exists x
    constructor
    { intro b xb
      have hbx := h1 b xb
      simp [←[totl| NonDetT.wp_eq_wp]] at hbx
      simp [←[part| NonDetT.wp_eq_wp]]
      exact f_ih b (fun x_2 ↦
        match x_2 with
        | ForInStep.yield b' => x (ForInStep.yield b') ∧ x_1 b' < x_1 b
        | ForInStep.done b' => x (ForInStep.done b'))
        x
        (by
          intro x2 hx2
          match x2 with
          | .yield b' => simp at hx2; simp [hx2]
          | .done b' => simp at hx2; simp [hx2] )
        hbx }
    simp [spec, LE.pure] at sp1
    simp [spec, LE.pure, sp1]
    exact le_trans sp1.right (by
      simp [loomLogicSimp, ←[totl| NonDetT.wp_eq_wp], ←[part| NonDetT.wp_eq_wp]]
      intro x1; exact cont_ih x1 post₁ post₂ le )