-- Loom/MonadAlgebras/NonDetT/Basic.lean

import Mathlib.Logic.Function.Basic

import Mathlib.Order.CompleteBooleanAlgebra

import Mathlib.Order.Lattice

import Mathlib.Order.Basic

import Loom.MonadAlgebras.WP.Basic

import Loom.MonadAlgebras.WP.Tactic

import Loom.MonadAlgebras.WP.Gen

import Loom.MonadAlgebras.NonDetT.Findable

section NonDeterministicTransformer

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 α

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

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

variable [CompleteBooleanAlgebra l] [MAlgOrdered m l]

namespace PartialCorrectness

namespace DemonicChoice

/- WP for NonDetT -/
noncomputable

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)

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

variable [LawfulMonad m]

/- Ordered Monad Algebra instance for NonDetT -/
noncomputable
scoped

-- Repository definitions used:
content: def WPGen.assert {l : Type u} {m : Type u -> Type v} [Monad m] [LawfulMonad m] [CompleteBooleanAlgebra l] [MAlgOrdered m l] (h : l) : WPGen (assertGadget (m := m) h) where
  get := fun post => h ⊓ (h ⇨ post .unit)
  prop := by admit /- proof elided -/

content: def WPGen.if {l : Type u} {m : Type u -> Type v} [Monad m] [LawfulMonad m] [CompleteBooleanAlgebra l] [MAlgOrdered m l]
  {hd : Decidable h} {x y : m α}
  (wpgx : h → WPGen x) (wpgy : ¬h → WPGen y)
  : WPGen (if h then x else y) where
  get := fun post =>
    (⨅ hc : WithName h (Lean.Name.anonymous.mkStr "if_pos"), (wpgx hc).get post) ⊓
    (⨅ hc : WithName (¬h) (Lean.Name.anonymous.mkStr "if_neg"), (wpgy hc).get post)
  prop := by admit /- proof elided -/

content: def WPGen.let  {l : Type u} {m : Type u -> Type v} [Monad m] [LawfulMonad m] [CompleteBooleanAlgebra l] [MAlgOrdered m l]
  (y : β) {x : β -> m α} (wpgx : ∀ y, WPGen (x y)) : WPGen (let z := y; x z) where
  get := fun post => ⨅ z, ⌜z = y⌝ ⇨ (wpgx z).get post
  prop := by admit /- proof elided -/
macro "⌜" p:term "⌝" : term => `(LE.pure $p)

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 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: class Findable {α : Type u} (p : α -> Prop) where
  find : Unit -> Option α
  find_none : (find ()).isNone -> ∀ x, ¬ p x
  find_some_p : find () = some x -> p x

content: def find [Encodable α] (p : α -> Prop) [DecidablePred p] : Option α :=
  findNat (fun x => (Encodable.decode x).any (p ·)) |>.bind Encodable.decode

content: def findNat (p : Nat -> Prop) [DecidablePred p] : Option Nat :=
  let rec aux i :=
    if p i then
      some i
    else
      aux (i + 1)
  partial_fixpoint
  aux 0

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 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: def spec (pre : l) (post : α -> l) : Cont l α :=
  fun p => pre ⊓ ⌜post ≤ p⌝

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: 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: @[aesop safe]
lemma emp_sum (x: Bal) (h: 0 ≤ x) : ∃ h: List Nat, h.sum ≤ x

content: omit [LawfulMonad m] in
lemma WPGen.intro (x : m α) (wpg : WPGen x) :
  pre <= wpg.get post ->
  -- pre <= wpg.sideCond ->
  triple pre x post

content: theorem triple_forIn_deacreasing {β} {measure : β -> ℕ}
  {init : β} {f : β → m (ForInStep β)}
  (inv : β → l)
  (hstep : ∀ b,
    measure b <= measure init ->
    triple
      (inv b)
      (f b)
      (fun | .yield b' => inv b' ⊓ ⌜measure b' < measure b⌝ | .done b' => ⌜ measure b' = 0 ⌝ ⊓ inv b')) :
  triple (inv init) (forIn [0:measure init] init (fun _ => f)) (fun b => inv b ⊓ ⌜measure b = 0⌝)

content: @[simp]
lemma trueE (l : Type v) [inst: LE l] [OrderTop l] [OrderBot l] : ⌜True⌝ = (⊤ : l)

content: @[simp]
lemma falseE (l : Type v) [inst: LE l] [OrderTop l] [OrderBot l] : ⌜False⌝ = (⊥ : l)

content: lemma LE.pure_imp {l : Type u} [inst: LE l] [OrderTop l] [OrderBot l]
  (p₁ p₂ : Prop) : (p₁ -> p₂) -> ⌜p₁⌝ <= (⌜p₂⌝ : l)

content: @[simp]
lemma LE.pure_intro {l : Type u} [inst: LE l] [OrderTop l] [OrderBot l]
  (p : Prop) (h : l) : (⌜p⌝ <= h) = (p -> ⊤ <= h)

content: @[simp]
lemma pure_intro_l {l : Type u} [CompleteLattice l] (x y : l) :
  (x ⊓ ⌜ p ⌝ <= y) = (p -> x <= y)

content: @[simp]
lemma pure_intro_r {l : Type u} [CompleteLattice l] (x y : l) :
  (⌜ p ⌝ ⊓ x <= y) = (p -> x <= y)

content: lemma MAlgOrdered.bind {α : Type u} {m} {l : Type u} [Monad m] [CompleteLattice l] [MAlgOrdered m l] :
    ∀ (x : m α) (f g : α -> m l), μ ∘ f = μ ∘ g ->
     μ (x >>= f) = μ (x >>= g)

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: @[simp]
lemma MAlg.μ_eq {m l} [Monad m] [CompleteLattice l] [MAlgOrdered m l] : MAlg.μ (m := m) = MAlgOrdered.μ (m := m)

content: lemma MAlg.monadLift_bind {α β} {l : Type u} {m : Type u -> Type v} [Monad m] [LawfulMonad m] [CompleteLattice l] [MAlgOrdered m l]
  (x : m α) (f g : α -> Cont l β) :
    f <= g ->
    (lift x >>= f) ≤ (lift x >>= g)

content: @[simp]
lemma lift_cont_eq {l σ : Type u} [CompleteLattice l] [CompleteLattice σ] (c : Cont l α) :
  liftM (m := Cont l) (n := Cont (σ -> l)) c = fun post s => c (post · s)

content: @[loomLogicSimp]
lemma leE (l : Type u) [PartialOrder l] (a b : α -> l) : a ≤ b ↔ ∀ x, a x ≤ b x

content: @[loomLogicSimp]
lemma lePropE (a b : Prop) : (a ≤ b) = (a → b)

content: @[loomLogicSimp]
lemma pureE (l : Type u) [CompleteLattice l] (a : Prop) : (⌜a⌝ : α -> l) = fun _ => ⌜a⌝

content: @[loomLogicSimp]
lemma purePropE  : (⌜a⌝ : Prop) = a

content: @[loomLogicSimp]
lemma infPropE (a b : Prop) : (a ⊓ b) = (a ∧ b)

content: @[loomLogicSimp]
lemma infE (l : Type u) [CompleteLattice l] (a b : α -> l) : (a ⊓ b) = fun x => a x ⊓ b x

content: @[loomLogicSimp]
lemma supE (l : Type u) [CompleteLattice l] (a b : α -> l) : (a ⊔ b) = fun x => a x ⊔ b x

content: @[loomLogicSimp]
lemma supPropE (a b : Prop) : (a ⊔ b) = (a ∨ b)

content: @[loomLogicSimp]
lemma iInfE (l : Type u) [CompleteLattice l] (a : ι -> α -> Prop) : (⨅ i, a i) = fun x => ⨅ i, a i x

content: @[loomLogicSimp]
lemma iSupE (l : Type u) [CompleteLattice l] (a : ι -> α -> Prop) : (⨆ i, a i) = fun x => ⨆ i, a i x

content: @[loomLogicSimp]
lemma himpE  (l : Type u) [CompleteBooleanAlgebra l] (a b : α -> l) :
  (a ⇨ b) = fun x => a x ⇨ b x

content: @[loomLogicSimp]
lemma himpPureE (a b : Prop) :
  (a ⇨ b) = (a -> b)

content: @[loomLogicSimp]
lemma topE (l : Type u) [CompleteLattice l] : (⊤ : α -> l) = fun _ => ⊤

content: @[loomLogicSimp]
lemma topPureE : (⊤ : Prop) = True

content: @[aesop safe]
theorem List.sum_lt (x: Bal) : x < y -> x < ([Int.toNat y]).sum

content: @[aesop safe]
theorem balance_lt (x: Bal) : x < x + 1

content: lemma Except.bind'_bind {m : Type u -> Type v} {ε α β} [Monad m] [LawfulMonad m] (i : m (Except ε α)) (f : α -> ExceptT ε m β) :
  (i >>= fun a => Except.bind' a f) = bind (m := ExceptT ε m) i f

content: lemma MAlg.lift_ExceptT ε (hd : ε -> Prop) [IsHandler hd] [CompleteLattice l] [inst: MAlgOrdered m l]
   (c : ExceptT ε m α) post :
  MAlg.lift c post = MAlg.lift (m := m) c (fun | .ok x => post x | .error e => ⌜hd e⌝)

content: lemma lift_map {α : Type u} {β : Type u} (f : α -> β) (x : m α)
  [Monad m] [Monad n] [LawfulMonad m] [LawfulMonad n] [MonadLiftT m n] [LawfulMonadLiftT m n] :
  liftM (f <$> x) = f <$> liftM (n := n) x

content: @[aesop unsafe]
theorem tail_length : ∀ q : List Nat, q.nonEmpty → q.tail.length < q.length

content: @[aesop norm]
theorem tail_sum (q: List Nat) (hnemp: q.nonEmpty): q.sum = q.tail.sum + q.head!

content: @[aesop unsafe]
theorem non_zero_length (q: List Nat) (hnemp: q.nonEmpty): 0 < q.length

content: @[aesop norm]
theorem sum_zero : ∀ q : List Nat, ¬q.nonEmpty → q.sum = 0

content: lemma meet_himp (x x' y z : l) :
  x = x' ->
  (x ⇨ y) ⊓ (x' ⇨ z) = x ⇨ (y ⊓ z)

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

content: @[ext]
lemma W_ext (t : Type v) (α : Type u) [Preorder t] (w w' : W t α) :
  w.wp = w'.wp → w = w'

-- Local definitions used:
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 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 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

-- Local lemmas used:
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_eq_wp {α : Type u} (x : NonDetT m α) (post : α -> l) :
  _root_.wp x post = NonDetT.wp x post

-- Target theorem and proof:
@[simp]
lemma NonDetT.wp_pickCont {τ : Type u} p (f : τ → NonDetT m α) [Findable p] post :
  _root_.wp (NonDetT.pickCont τ p f) post =  ⨅ a, ⌜p a⌝ ⇨ _root_.wp (f a) post
:= by
  simp [NonDetT.wp_eq_wp, NonDetT.wp]; congr; ext x
  simp [Membership.mem, Set.Mem]
  by_cases h: p x <;> simp [h]