-- 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: def wp (c : m α) (post : α -> l) : l := liftM (n := Cont l) c post

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 => ⨆ a, ⌜p a⌝ ⊓ wp (f a) post
  | _, .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: 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 spec (pre : l) (post : α -> l) : Cont l α :=
  fun p => pre ⊓ ⌜post ≤ p⌝

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} {α : 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 => ⨆ a, ⌜p a⌝ ⊓ wp (f a) post

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

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.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: structure W (t : Type v) [Preorder t] (α : Type u) where
  wp : Cont t α
  wp_montone : wp.monotone

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; ⌜∃ h, p h⌝ ⊓  ⨅ a ∈ (p : Set τ), wp (f a) post
  | _, @NonDetT.repeatCont _ _ β init f cont => fun post => ⨆ (inv : ForInStep β -> l) (measure : β -> Nat),
      ⌜ ∀ b, (inv (ForInStep.yield b)) <= wp (f b) (fun | .yield b' => inv (.yield b') ⊓ ⌜ measure b' < measure b ⌝ | .done b' => inv (.done b'))⌝ ⊓
      spec (inv (.yield init)) (fun b => inv (.done b)) (fun b => wp (cont b) post)

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 => ⨆ a, ⌜p a⌝ ⊓ wp (f a) post
  | _, .repeatCont init f cont => fun post => ⨆ (inv : ForInStep _ -> l) (measure : _ -> Nat),
    ⌜ ∀ b, (inv (ForInStep.yield b)) <= wp (f b) (fun | .yield b' => inv (.yield b') ⊓ ⌜ measure b' < measure b ⌝ | .done b' => inv (.done b'))⌝ ⊓
    spec (inv (.yield init)) (fun b => inv (.done b)) (fun b => wp (cont b) post)

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

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.μ  {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)

content: lemma NonDetT.wp_mono [LawfulMonad m] {α : Type u} {l : Type u} [CompleteLattice l] [MAlgOrdered m l] (x : NonDetT m α) (f g : α -> l) :
  (∀ a, f a <= g a) ->
  NonDetT.wp x f <= NonDetT.wp x g

content: lemma spec_mono {α : Type u} {l : Type u} [CompleteLattice l] (pre : l) (post : α -> l) (f g : α -> l) :
  (∀ a, f a <= g a) ->
  spec pre post f <= spec pre post g

content: omit [MAlgOrdered m l] in
lemma spec_mono {α : Type u}  {l : Type u} [CompleteLattice l] (pre : l) (post : α -> l) (f g : α -> l) :
  (∀ a, f a <= g a) ->
  spec pre post f <= spec pre post g

content: lemma NonDetT.wp_mono  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] [LawfulMonad m] {α : Type u} (x : NonDetT m α) (f g : α -> l) :
  (∀ a, f a <= g a) ->
  NonDetT.wp x f <= NonDetT.wp x g

-- Local definitions used:

-- Local lemmas used:
content: @[local simp]
lemma DivM.total_decompose (α : Type) (x : DivM α) (post₁ post₂ : α -> Prop) :
  ([totl| wp x post₁] ∧ [part| wp x post₂]) = [totl| wp x (post₁ ⊓ post₂)]

content: @[local simp]
lemma wp_mono_part (x : NonDetT DivM α) (post₁ post₂ : α -> Prop) :
  (post₁ ≤ post₂) → ([totl|wp x post₁]) ≤ ([part| wp x post₂])

-- Target theorem and proof:
lemma VelvetM.total_decompose {α : Type} (x : VelvetM α) (post₁ post₂ : α -> Prop):
  [totl| wp x post₁] ⊓ [part| wp x post₂] = [totl| wp x (post₁ ⊓ post₂)]
:= by
    unhygienic induction x <;> try simp [loomLogicSimp]
    { simp [DivM.total_decompose]
      simp [[totl|DivM.wp_eq]]
      split
      { simp }
      rename_i arg
      have ind := f_ih arg post₁ post₂
      simp at ind
      rw [ind]
      trivial }
    { constructor <;> rintro hyp
      { constructor; aesop
        intro i hi
        have hl := hyp.left.right i hi
        have hr := hyp.right i hi
        have ind := f_ih i post₁ post₂
        simp [hl, hr] at ind
        exact ind }
      constructor; constructor; aesop
      all_goals
        intro i hi
        have conj := hyp.right i hi
        have ind := f_ih i post₁ post₂
        simp [loomLogicSimp] at ind
        rw [←ind] at conj
        simp [conj] }
    constructor
    { intro conj
      rcases conj with ⟨h, inv_spec, hspec⟩
      rcases h with ⟨inv, x1, hinv⟩
      rcases x1 with ⟨x, hx⟩
      exists inv ⊓ inv_spec
      constructor
      { exists x
        intro b hb
        simp [←[totl| NonDetT.wp_eq_wp]]
        simp [loomLogicSimp] at hb
        have hxb := hx b hb.left
        simp [←[totl| NonDetT.wp_eq_wp]] at hxb
        simp [spec, LE.pure, loomLogicSimp, ←[part| NonDetT.wp_eq_wp]] at hspec
        have hspecb := hspec.left b hb.right
        have ind := f_ih b
          (fun x_1 ↦
          match x_1 with
          | ForInStep.yield b' => inv (ForInStep.yield b') ∧ x b' < x b
          | ForInStep.done b' => inv (ForInStep.done b'))
          inv_spec
        simp [loomLogicSimp] at ind
        have indr := ind.mp (And.intro hxb hspecb)
        have v1 := [totl| NonDetT.wp_mono
          (f b)
          (fun x_1 ↦
            (match x_1 with
              | ForInStep.yield b' => inv (ForInStep.yield b') ∧ x b' < x b
              | ForInStep.done b' => inv (ForInStep.done b')) ∧
              inv_spec x_1)
          (fun x_1 ↦
            match x_1 with
            | ForInStep.yield b' => (inv (ForInStep.yield b') ∧ inv_spec (ForInStep.yield b')) ∧ x b' < x b
            | ForInStep.done b' => inv (ForInStep.done b') ∧ inv_spec (ForInStep.done b'))
          ]
        simp [loomLogicSimp, ←[totl| NonDetT.wp_eq_wp]] at v1
        exact v1
          (fun x => by
            match x with
            | .yield b' => intro hb hspec; simp [hb, hspec]
            | .done b' => intro hb hspec; simp [hb, hspec] )
          indr }
      simp [spec, LE.pure, loomLogicSimp, ←[totl| NonDetT.wp_eq_wp]] at hinv
      simp [spec, LE.pure, loomLogicSimp, ←[part| NonDetT.wp_eq_wp]] at hspec
      simp [spec, LE.pure, loomLogicSimp, ←[totl| NonDetT.wp_eq_wp]]
      simp [hinv, hspec]
      intro x inv_x inv_spec
      have h₁ := hinv.right x inv_x
      have h₂ := hspec.right.right x inv_spec
      have cont_ind := cont_ih x post₁ post₂
      simp [loomLogicSimp] at cont_ind
      exact cont_ind.mp (And.intro h₁ h₂) }
    intro hyp
    rcases hyp with ⟨inv, x_ex, h_inv⟩
    rcases x_ex with ⟨x, hx⟩
    simp [spec]
    simp [spec, LE.pure] at h_inv
    constructor <;> exists inv <;> constructor
    { exists x }
    { simp [h_inv, LE.pure]
      exact le_trans h_inv.right (by
        simp [← [totl|NonDetT.wp_eq_wp]]
        simp [loomLogicSimp]
        intro x and_wp
        have cont_ind := cont_ih x post₁ post₂
        simp [loomLogicSimp, and_wp] at cont_ind
        simp [cont_ind] ) }
    { intro b hb
      have hbx := hx b hb
      simp [←[totl| NonDetT.wp_eq_wp]] at hbx
      have hb_triv: True ≤ ([totl| wp (f b) fun x_1 ↦
        match x_1 with
        | ForInStep.yield b' => inv (ForInStep.yield b') ∧ x b' < x b
        | ForInStep.done b' => inv (ForInStep.done b')]) := by
        simp [loomLogicSimp]
        exact hbx
      have tr_intro: True ≤ ([part| NonDetT.wp (f b) inv]) := le_trans hb_triv (by
        simp [loomLogicSimp]
        intro wp_x
        simp [←[part| NonDetT.wp_eq_wp]]
        apply wp_mono_part (f b) (fun x_1 ↦
          match x_1 with
          | ForInStep.yield b' => inv (ForInStep.yield b') ∧ x b' < x b
          | ForInStep.done b' => inv (ForInStep.done b')) (fun x => inv x)
        { simp [loomLogicSimp]
          intro x1
          match x1 with
          | ForInStep.yield b' => simp; intros; simp [*]
          | ForInStep.done b' => simp }
        exact wp_x)
      simp at tr_intro
      simp [tr_intro] }
    simp [LE.pure, h_inv]
    exact le_trans h_inv.right (by
      simp [←[totl| NonDetT.wp_eq_wp], ←[part| NonDetT.wp_eq_wp]]
      simp [loomLogicSimp])