-- Loom/MonadAlgebras/NonDetT'/Extract.lean

import Mathlib.Logic.Function.Basic

import Mathlib.Order.CompleteBooleanAlgebra

import Mathlib.Order.Lattice

import Mathlib.Order.Basic

import Mathlib.Data.W.Basic

import Mathlib.Data.FinEnum

import Loom.MonadAlgebras.WP.Gen

import Loom.MonadAlgebras.WP.Liberal

import Loom.MonadAlgebras.NonDetT'.Basic

open Lean.Order

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

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

class WeakFindable {α : Type u} (p : α -> Prop) where
  find : Unit -> Option α
  find_some_p : find () = some x -> p x

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

instance WeakFindable.of_Findable {α : Type u} (p : α -> Prop) [Findable p] : WeakFindable p where
  find := Findable.find p
  find_some_p := Findable.find_some_p

instance {p : α -> Prop} [Encodable α] [DecidablePred p] : Findable p where
  find := fun _ => find p
  find_none := find_none p
  find_some_p := find_some_p p _

@[instance high]
instance {p : α -> Prop} [FinEnum α] [DecidablePred p] : Findable p where
  find := fun _ => FinEnum.toList α |>.find? p
  find_none := by admit /- proof elided -/

inductive ExtractNonDet (findable : {τ : Type u} -> (τ -> Prop) -> Type u) {m} : {α : Type u} -> NonDetT m α -> Type _ where
  | pure {α} : ∀ (x : α), ExtractNonDet findable (NonDetT.pure x)
  | vis {α} {β} (x : m β) (f : β → NonDetT m α) :
      (∀ y, ExtractNonDet findable (f y)) → ExtractNonDet findable (.vis x f)
  | pickSuchThat {α} (τ : Type u) (p : τ -> Prop) (f : τ → NonDetT m α)
    {_ : findable p}
     : (∀ x, ExtractNonDet findable (f x)) → ExtractNonDet findable (.pickCont τ p f)
  | assume {α} (p : PUnit -> Prop) (f : PUnit → NonDetT m α) {_ : Decidable (p .unit)} :
    (∀ x, ExtractNonDet findable (f x)) → ExtractNonDet findable (.pickCont PUnit p f)

instance ExtractNonDet.pure' : ExtractNonDet findable (Pure.pure (f := NonDetT m) x) :=

instance ExtractNonDet.liftM (x : m α) :
  ExtractNonDet findable (liftM (n := NonDetT m) x) :=

instance ExtractNonDet.assume' {p : Prop} [Decidable p] : ExtractNonDet findable (MonadNonDet.assume (m :=  NonDetT m) p) :=

instance ExtractNonDet.pickSuchThat' {τ : Type u} (p : τ -> Prop) [Findable p] :
  ExtractNonDet Findable (MonadNonDet.pickSuchThat (m := NonDetT m) τ p) :=

instance ExtractNonDet.pickSuchThat_weak {τ : Type u} (p : τ -> Prop) [WeakFindable p] :
  ExtractNonDet WeakFindable (MonadNonDet.pickSuchThat (m := NonDetT m) τ p) :=

instance ExtractNonDet.if {p : Prop} {dec : Decidable p} {x y : NonDetT m α}
  (_ : ExtractNonDet findable x) (_ : ExtractNonDet findable y) :
  ExtractNonDet findable (if p then x else y) :=

instance ExtractNonDet.ForIn_list {xs : List α} {init : β} {f : α → β → NonDetT m (ForInStep β)}
  (_ : ∀ a b, ExtractNonDet findable (f a b)) :
  ExtractNonDet findable (forIn xs init f) :=

variable [Monad m] [CCPOBot m] [CompleteBooleanAlgebra l] [MAlgOrdered m l] [MAlgDet m l] [LawfulMonad m]

@[simp, inline]
def NonDetT.extractGen {findable : {τ : Type u} -> (τ -> Prop) -> Type u} {α : Type u}
  (findOf : ∀ {τ : Type u} (p : τ -> Prop), findable p -> Unit -> Option τ)
  : (s : NonDetT m α) -> (ex : ExtractNonDet findable s := by admit /- proof elided -/
  ) -> m α
  | .pure x, _ => Pure.pure x
  | .vis x f, .vis _ _ _ => liftM x >>= (fun x => extractGen findOf (f x))
  | .pickCont _ p f, .pickSuchThat _ _ _ _ =>
    match findOf p ‹_› () with
    | none => CCPOBot.compBot
    | some x =>  extractGen findOf (f x)
  | .pickCont _ p f, .assume _ _ _ =>
    if p .unit then
      extractGen findOf (f .unit)
    else CCPOBot.compBot

@[inline]
def NonDetT.extract {α : Type u} (s : NonDetT m α) (ex : ExtractNonDet Findable s := by admit /- proof elided -/
) : m α :=
  NonDetT.extractGen Findable.find s

abbrev NonDetT.prop {α : Type u} : (s : NonDetT m α) -> Cont l α
  | .pure x => Pure.pure x
  | .vis x f => fun post => wlp x fun y => NonDetT.prop (f y) post
  | .pickCont _ p f => fun post =>
    (⨅ t, ⌜p t⌝ ⇨ NonDetT.prop (f t) post) ⊓ (⨆ t, ⌜p t⌝)

structure Extractable (x : NonDetT m α) where
  cond : Cont l α
  prop : ∀ post, cond post <= x.prop post

def ExtractNonDet.prop {α : Type u} (s : NonDetT m α) :  ExtractNonDet WeakFindable s -> l
  | .pure x => ⊤
  | .vis x f ex => wlp x fun y => (ex y).prop
  | .pickSuchThat _ p f ex => ⨅ t ∈ WeakFindable.find p (), (ex t).prop
  | .assume p f ex =>
      if p .unit then
        (ex .unit).prop
      else ⊤

namespace DemonicChoice

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

content: class MonadNonDet (m : Type u → Type v) where
  pick : (τ : Type u) → [Inhabited τ] → m τ
  
  pickSuchThat : (τ : Type u) → (p : τ → Prop) → [Findable p] → m τ
  assume : (as : Prop) → [Decidable as] → m PUnit.{u+1}
  
  rep {α : Type u} : α → (α → m (ForInStep α)) → m α

content: def wlp (c : m α) (post : α -> l) : l := iwp c post ⊔ wp c post

content: abbrev iwp (c : m α) : Cont l α := Cont.inv (wp c)

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

content: def Cont.inv {t : Type v} {α : Type u} [BooleanAlgebra t] (wp : Cont t α) : Cont t α :=
  fun f => (wp fun x => (f x)ᶜ)ᶜ

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

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: class CCPOBot (m : Type u -> Type v)  where
  compBot {α} : m α

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

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

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.μ  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m l -> l := fun x => NonDetT.wp x id

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.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: private def _root_.Lean.EnvExtension.get [Inhabited σ] (ext : EnvExtension σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

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: def NonDetT.pickSuchThat (τ : Type u) (p : τ → Prop) [Findable p] : NonDetT m τ :=
  NonDetT.pickCont τ p pure

content: instance {l σ : Type u} : MonadLift (Cont l) (Cont (σ -> l)) where
  monadLift x := fun f s => x (f · s)

-- Repository 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 wlp_join_wp (c : m α) (post post' : α -> l) :
  wlp c post ⊓ wp c post' = wp c (fun x => post x ⊓ post' x)

content: lemma wlp_himp (c : m α) (post post' : α -> l) :
  wp c (fun x => post' x ⇨ post x) = wlp c post' ⇨ wp c post

content: omit [LawfulMonad m] in
lemma wp_wlp (c : m α) (post : α -> l) :
  wp c post <= wlp c post

content: @[simp]
lemma NonDetT.wp_vis {β : Type u} (x : m β) (f : β → NonDetT m α) post :
  _root_.wp (NonDetT.vis x f) post = _root_.wp x fun a => _root_.wp (f a) post

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 wp_cons (x : m α) (post post' : α -> l) :
  (∀ y, post y ≤ post' y) ->
  wp x post ≤ wp x post'

content: @[simp]
lemma NonDetT.wp_pickCont {τ : Type u} p (f : τ → NonDetT m α) post :
  _root_.wp (NonDetT.pickCont τ p f) post = ⨅ a, ⌜p a⌝ ⇨ _root_.wp (f a) post

content: @[simp]
lemma NonDetT.wp_pickCont {τ : Type u} p (f : τ → NonDetT m α) post :
  _root_.wp (NonDetT.pickCont τ p f) post = ⨆ a, ⌜p a⌝ ⊓ _root_.wp (f a) post

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

-- Local definitions used:
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: def find [Encodable α] (p : α -> Prop) [DecidablePred p] : Option α :=
  findNat (fun x => (Encodable.decode x).any (p ·)) |>.bind Encodable.decode

content: class WeakFindable {α : Type u} (p : α -> Prop) where
  find : Unit -> Option α
  find_some_p : find () = some x -> p x

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: instance WeakFindable.of_Findable {α : Type u} (p : α -> Prop) [Findable p] : WeakFindable p where
  find := Findable.find p
  find_some_p := Findable.find_some_p

content: instance {p : α -> Prop} [Encodable α] [DecidablePred p] : Findable p where
  find := fun _ => find p
  find_none := find_none p
  find_some_p := find_some_p p _

content: @[instance high]
instance {p : α -> Prop} [FinEnum α] [DecidablePred p] : Findable p where
  find := fun _ => FinEnum.toList α |>.find? p
  find_none := by admit /- proof elided -/

content: inductive ExtractNonDet (findable : {τ : Type u} -> (τ -> Prop) -> Type u) {m} : {α : Type u} -> NonDetT m α -> Type _ where
  | pure {α} : ∀ (x : α), ExtractNonDet findable (NonDetT.pure x)
  | vis {α} {β} (x : m β) (f : β → NonDetT m α) :
      (∀ y, ExtractNonDet findable (f y)) → ExtractNonDet findable (.vis x f)
  | pickSuchThat {α} (τ : Type u) (p : τ -> Prop) (f : τ → NonDetT m α)
    {_ : findable p}
     : (∀ x, ExtractNonDet findable (f x)) → ExtractNonDet findable (.pickCont τ p f)
  | assume {α} (p : PUnit -> Prop) (f : PUnit → NonDetT m α) {_ : Decidable (p .unit)} :
    (∀ x, ExtractNonDet findable (f x)) → ExtractNonDet findable (.pickCont PUnit p f)

content: instance ExtractNonDet.pure' : ExtractNonDet findable (Pure.pure (f := NonDetT m) x) :=

content: instance ExtractNonDet.liftM (x : m α) :
  ExtractNonDet findable (liftM (n := NonDetT m) x) :=

content: instance ExtractNonDet.assume' {p : Prop} [Decidable p] : ExtractNonDet findable (MonadNonDet.assume (m :=  NonDetT m) p) :=

content: instance ExtractNonDet.pickSuchThat' {τ : Type u} (p : τ -> Prop) [Findable p] :
  ExtractNonDet Findable (MonadNonDet.pickSuchThat (m := NonDetT m) τ p) :=

content: instance ExtractNonDet.pickSuchThat_weak {τ : Type u} (p : τ -> Prop) [WeakFindable p] :
  ExtractNonDet WeakFindable (MonadNonDet.pickSuchThat (m := NonDetT m) τ p) :=

content: instance ExtractNonDet.if {p : Prop} {dec : Decidable p} {x y : NonDetT m α}
  (_ : ExtractNonDet findable x) (_ : ExtractNonDet findable y) :
  ExtractNonDet findable (if p then x else y) :=

content: instance ExtractNonDet.ForIn_list {xs : List α} {init : β} {f : α → β → NonDetT m (ForInStep β)}
  (_ : ∀ a b, ExtractNonDet findable (f a b)) :
  ExtractNonDet findable (forIn xs init f) :=

content: @[simp, inline]
def NonDetT.extractGen {findable : {τ : Type u} -> (τ -> Prop) -> Type u} {α : Type u}
  (findOf : ∀ {τ : Type u} (p : τ -> Prop), findable p -> Unit -> Option τ)
  : (s : NonDetT m α) -> (ex : ExtractNonDet findable s := by admit /- proof elided -/
  ) -> m α
  | .pure x, _ => Pure.pure x
  | .vis x f, .vis _ _ _ => liftM x >>= (fun x => extractGen findOf (f x))
  | .pickCont _ p f, .pickSuchThat _ _ _ _ =>
    match findOf p ‹_› () with
    | none => CCPOBot.compBot
    | some x =>  extractGen findOf (f x)
  | .pickCont _ p f, .assume _ _ _ =>
    if p .unit then
      extractGen findOf (f .unit)
    else CCPOBot.compBot

content: @[inline]
def NonDetT.extract {α : Type u} (s : NonDetT m α) (ex : ExtractNonDet Findable s := by admit /- proof elided -/
) : m α :=
  NonDetT.extractGen Findable.find s

content: abbrev NonDetT.prop {α : Type u} : (s : NonDetT m α) -> Cont l α
  | .pure x => Pure.pure x
  | .vis x f => fun post => wlp x fun y => NonDetT.prop (f y) post
  | .pickCont _ p f => fun post =>
    (⨅ t, ⌜p t⌝ ⇨ NonDetT.prop (f t) post) ⊓ (⨆ t, ⌜p t⌝)

content: structure Extractable (x : NonDetT m α) where
  cond : Cont l α
  prop : ∀ post, cond post <= x.prop post

content: def ExtractNonDet.prop {α : Type u} (s : NonDetT m α) :  ExtractNonDet WeakFindable s -> l
  | .pure x => ⊤
  | .vis x f ex => wlp x fun y => (ex y).prop
  | .pickSuchThat _ p f ex => ⨅ t ∈ WeakFindable.find p (), (ex t).prop
  | .assume p f ex =>
      if p .unit then
        (ex .unit).prop
      else ⊤

-- Local lemmas used:

-- Target theorem and proof:
lemma ExtractNonDet.extract_refines_wp (s : NonDetT m α) (inst : ExtractNonDet Findable s) :
  wp s post ⊓ s.prop ⊤ <= wp s.extract post
:= by
  unhygienic induction inst
  { simp [wp_pure, NonDetT.extract] }
  { simp [NonDetT.wp_vis, NonDetT.prop]; rw [inf_comm, wlp_join_wp]
    simp [NonDetT.extract, wp_bind]
    apply wp_cons; aesop (add norm inf_comm) }
  { simp [NonDetT.wp_pickCont, NonDetT.prop, NonDetT.extract]; split
    { have := Findable.find_none (p := p) (by simpa);
      have : (∀ x, p x = False) := by simpa
      simp [this] }
    rw [<-inf_assoc]; refine inf_le_of_left_le ?_
    rw [← @iInf_inf_eq]; simp [meet_himp _ _ _ _ rfl]
    rename_i y _
    refine iInf_le_of_le y ?_
    have := Findable.find_some_p (p := p) (by assumption)
    simp [this]; apply a_ih }
  simp [NonDetT.wp_pickCont, NonDetT.prop, NonDetT.extract]
  have : ∀ a : PUnit.{u_1 + 1}, a = .unit := by simp
  simp [this, iInf_const]; split_ifs <;> simp [*, iSup_const]
  apply a_ih