{-# OPTIONS --safe --experimental-lossy-unification #-}

open import Prelude
open import FinitenessConditions
open import Signatures

module Effects.NonDetState
  (S : Type)
  (choice : SplitQuotientedChoiceω S)
  (set : isSet S)
  where

import Effects.NonDet as ℕ𝔻
import Effects.State S choice set as π•Š
open import FreeMonad.Combinations.Tensor ℕ𝔻.ℕ𝔻-theory π•Š.π•Š-theory

open import FreeMonad.Quotiented combinedTheory


import FreeMonad.Syntax as Synt
open Synt (ℕ𝔻.ℕ𝔻 ⊞ π•Š.π•Š) using (Op⟦_⟧)
open Signature (ℕ𝔻.ℕ𝔻 ⊞ π•Š.π•Š)
open ℕ𝔻.NonDet.Op

put : S β†’ Term ⊀
put s = [ injβ‚‚ (π•Š.Synt.put s) ]

get : Term (S)
get = [ injβ‚‚ π•Š.Synt.get ]

fail : Term A
fail = [ Synt.op (inl `fail , absurd) ]


infixl 5 _<|>_
_<|>_ : Term A β†’ Term A β†’ Term A
xs <|> ys = opβ‚œ (inl `<|> , bool ys xs)

until : β„• β†’ Term Bool β†’ Term ⊀ β†’ Term ⊀
until zero    c t = fail
until (suc n) c t = c >>= bool (t >> until n c t) (return tt)

guard : Bool β†’ Term ⊀
guard = bool fail (return tt)

open import Data.Nat.Order using () renaming (_<ᡇ_ to _<_)
open import Data.Nat.Properties using (_+_)

<|>-alg : (p q : Term A) (k : A β†’ Term B) β†’ ((p <|> q) >>= k) ≑ ((p >>= k) <|> (q >>= k))
<|>-alg p q k = algebraic (inl `<|>) k (bool q p)

<|>-idem : (x : Term A) β†’ x <|> x ≑ x
<|>-idem = elimProp/ (Ξ» _ β†’ squash/ _ _ )
  Ξ» x β†’ cong ([_] ∘ Synt.op ∘ _,_ (inl `<|>)) (funExt (bool refl refl)) ΝΎ eq/ _ _ (eqβ‚œ (inl (inl ℕ𝔻.Synt.idem)) tt (const x))

<|>-assoc : (x y z : Term A) β†’ (x <|> y) <|> z ≑ x <|> (y <|> z)
<|>-assoc = elimProp3/ (Ξ» _ _ _ β†’ squash/ _ _ )
  Ξ» x y z β†’
    let h = eq/ _ _ (eqβ‚œ (inl (inl ℕ𝔻.Synt.assoc)) tt Ξ» { zero β†’ z ; (suc zero) β†’ y ; _ β†’ x })
    in cong ([_] ∘ Synt.op ∘ _,_ (inl `<|>)) (funExt (bool refl (cong (Synt.op ∘ _,_ (inl `<|>)) (funExt (bool refl refl)))))
     ΝΎ sym h
     ; cong ([_] ∘ Synt.op ∘ _,_ (inl `<|>)) (funExt (bool (cong (Synt.op ∘ _,_ (inl `<|>)) (funExt (bool refl refl))) refl))

<|>-comm : (x y : Term A) β†’ x <|> y ≑ y <|> x
<|>-comm = elimProp2/ (Ξ» _ _ β†’ squash/ _ _ )
  Ξ» x y β†’
    let h = eq/ _ _ (eqβ‚œ (inl (inl ℕ𝔻.Synt.comm)) tt Ξ» { zero β†’ y ; _ β†’ x })
    in cong ([_] ∘ Synt.op ∘ _,_ (inl `<|>)) (funExt (bool refl refl))
     ΝΎ h
     ; cong ([_] ∘ Synt.op ∘ _,_ (inl `<|>)) (funExt (bool refl refl))

side : Term Bool
side = Opβ‚œβŸ¦ inl `<|> ⟧



open import Hoare.Definition combinedTheory
open import Hoare.Lemmas     combinedTheory
open import Hoare.Calculus   combinedTheory

comm-<|>-op : βˆ€ Oβ‚› (k : Arity (inr Oβ‚›) Γ— Arity (inl `<|>) β†’ Term A) β†’ (do s ← Opβ‚œβŸ¦ inr Oβ‚› ⟧; n ← side; k (s , n)) ≑ (do n ← side; s ← Opβ‚œβŸ¦ inr Oβ‚› ⟧; k (s , n))
comm-<|>-op Oβ‚› =
  SplitQuotientedChoiceAt.elim~canonical
    (quotientedChoiceProd (π•Š.finiteOps Oβ‚›) quotientedChoiceBool)
    _
    (Ξ» _ β†’ squash/ _ _)
    Ξ» k β†’ sym (s ⇐ Opβ‚œβŸ¦ inr Oβ‚› ⟧ ΝΎ/ syntactic-bind (Ξ» n β†’ k (s , n)) Op⟦ inl `<|> ⟧)
        ΝΎ sym (syntactic-bind (Ξ» s β†’ Synt.op (inl `<|> , Ξ» n β†’ k (s , n))) (Op⟦ inr Oβ‚› ⟧))
        ΝΎ sym ( eq/ _ _ (commutesβ€² `<|> Oβ‚› (flip (curry k))) )
        ΝΎ syntactic-bind (Ξ» n β†’ Synt.op (inr Oβ‚› , Ξ» s β†’ k (s , n)))  Op⟦ inl `<|> ⟧
        ΝΎ (n ⇐ side ΝΎ/ (syntactic-bind (Ξ» s β†’ k (s , n)) Op⟦ inr Oβ‚› ⟧))

fail-exit : (k : A β†’ Term B) β†’ (fail >>= k) ≑ fail
fail-exit k = 
  fail >>= k β‰‘βŸ¨ algebraic (inl `fail) k absurd ⟩
  opβ‚œ (inl `fail , k ∘ absurd) β‰‘βŸ¨ cong {x = k ∘ absurd} {y = [_] ∘ absurd} (opβ‚œ ∘ _,_ (inl `fail)) (funExt (Ξ» ())) ⟩
  opβ‚œ (inl `fail , [_] ∘ absurd) β‰‘βŸ¨ opβ‚œ-com-[] absurd ⟩
  [ Synt.op (inl `fail , absurd) ] β‰‘βŸ¨βŸ©
  fail ∎

<|>-comm-dupe : (k : Bool β†’ Bool β†’ Term A) β†’ (do s₁ ← side; sβ‚‚ ← side; k s₁ sβ‚‚) ≑ (do s₁ ← side; sβ‚‚ ← side; k sβ‚‚ s₁)
<|>-comm-dupe k =
  (k true true <|> k true false) <|> (k false true <|> k false false) β‰‘βŸ¨ <|>-assoc (k true true) (k true false) (k false true <|> k false false) ΝΎ cong (k true true <|>_) (sym (<|>-assoc (k true false) (k false true) (k false false))) ⟩
  k true true <|> ((k true false <|> k false true) <|> k false false) β‰‘βŸ¨ cong (Ξ» e β†’ k true true <|> (e <|> k false false)) (<|>-comm (k true false) (k false true)) ⟩
  k true true <|> ((k false true <|> k true false) <|> k false false) β‰‘Λ˜βŸ¨ <|>-assoc (k true true) (k false true) (k true false <|> k false false) ΝΎ cong (k true true <|>_) (sym (<|>-assoc (k false true) (k true false) (k false false))) ⟩
  (k true true <|> k false true) <|> (k true false <|> k false false) ∎

sefβ€²-<|> : (k : A β†’ Bool β†’Β Term B) (p : Term A) β†’ 
           (do x ← p ; s ← side ; k x s) ≑ (do s ← side ; x ← p ; k x s)
sefβ€²-<|> k =
  elimβ‚œ-prop
    _
    (Ξ» _ β†’ squash/ _ _)
    Ξ» { (Synt.var x) β†’ refl
      ; (Synt.op (Oα΅’ , xs) P⟨xs⟩) β†’ algebraic Oα΅’ _ _ ΝΎ (i ⇐ Opβ‚œβŸ¦ Oα΅’ ⟧ ΝΎ/ P⟨xs⟩ i) ΝΎ lemma Oα΅’ xs k P⟨xs⟩ ΝΎ s ⇐ side ΝΎ/ (sym (algebraic Oα΅’ (Ξ» x β†’ k x s) xs))
      }
  where
  lemma : βˆ€ Oα΅’ (xs : Arity Oα΅’ β†’ Term A) (k : A β†’ Bool β†’ Term B)
        β†’ (βˆ€ i β†’ (xs i >>= Ξ» x β†’ side >>= k x) ≑ side >>= (Ξ» s β†’ xs i >>= Ξ» x β†’ k x s))
        β†’ (do i ← Opβ‚œβŸ¦ Oα΅’ ⟧; s ← side; x ← xs i; k x s) ≑ (do s ← side ; i ← Opβ‚œβŸ¦ Oα΅’ ⟧; x ← xs i; k x s)
  lemma (inl `<|>)  xs k x = <|>-comm-dupe (Ξ» s₁ sβ‚‚ β†’ xs s₁ >>= Ξ» x β†’ k x sβ‚‚)

  lemma (inr Oβ‚›)    xs k x = comm-<|>-op Oβ‚› (Ξ» { (i , s) β†’ xs i >>= flip k s })
  lemma (inl `fail) xs k x =

    (do f ← Opβ‚œβŸ¦ inl `fail ⟧
        s ← side
        x ← xs f
        k x s) β‰‘βŸ¨ fail-exit (Ξ» f β†’ side >>= Ξ» s β†’ xs f >>= flip k s) ⟩

    fail β‰‘Λ˜βŸ¨ <|>-idem fail ⟩

    (do s ← side
        fail)  β‰‘Λ˜βŸ¨ s ⇐ side ΝΎ/ fail-exit (xs >=> flip k s)  ⟩

    (do s ← side
        f ← Opβ‚œβŸ¦ inl `fail ⟧
        x ← xs f
        k x s) ∎
       

sef-<|> : (p : Term A) (q r : A β†’ Term B) β†’ (p >>= (Ξ» x β†’ q x <|> r x)) ≑ ((p >>= q) <|> (p >>= r))
sef-<|> p q r = sefβ€²-<|> (Ξ» x s β†’ bool (r x) (q x) s) p


module _ where
  open import Truth.MonoLevel β„“zero
  open import Truth using (|β†’|-id)
  guard-hoare : βˆ€ p β†’ βŸ…βŸ† guard p βŸ… return (|T| p) βŸ†
  guard-hoare false .proof k = refl
  guard-hoare true  .proof k = cong (k tt) (|β†’|-id _)

  fail-absurd : (Ο• : A β†’ Term Ξ©) β†’
    βŸ…βŸ† x ← fail βŸ… Ο• x βŸ†
  fail-absurd Ο• .proof k = refl

open import Truth

<|>-conj : {A : Type a}
     β†’                 (Ο• : Term (Ξ© a)) β†’
                       (ψ : A β†’ Term (Ξ© a))
     β†’                              (p q : Term A) β†’
  βŸ… Ο• βŸ† x ← p        βŸ… ψ x βŸ† β†’
  βŸ… Ο• βŸ† x ← q        βŸ… ψ x βŸ† β†’
  βŸ… Ο• βŸ† x ← p <|> q  βŸ… ψ x βŸ†
<|>-conj Ο• ψ p q lhs rhs .proof k =
  (do a ← Ο•
      x ← (p <|> q)
      b ← ψ x
      k x (a |β†’| b)) β‰‘βŸ¨ a ⇐ Ο• ΝΎ/ <|>-alg p q _ ⟩

  (do a ← Ο•;
        (do x ← p
            b ← ψ x
            k x (a |β†’| b))
          <|>
        (do x ← q
            b ← ψ x
            k x (a |β†’| b)))

          β‰‘βŸ¨ sef-<|> Ο• _ _ ⟩

  (do a ← Ο•
      x ← p
      b ← ψ x
      k x (a |β†’| b))
    <|>
  (do a ← Ο•
      x ← q
      b ← ψ x
      k x (a |β†’| b))
          β‰‘βŸ¨ congβ‚‚ _<|>_ (lhs .proof k) (rhs .proof k) ⟩

  (do a ← Ο•
      x ← p
      b ← ψ x
      k x True)
    <|>
  (do a ← Ο•
      x ← q
      b ← ψ x
      k x True)
        β‰‘Λ˜βŸ¨ sef-<|> Ο• _ _ ⟩

  (do a ← Ο•;
        (do x ← p
            b ← ψ x
            k x True)
          <|>
        (do x ← q
            b ← ψ x
            k x True))
          β‰‘Λ˜βŸ¨ a ⇐ Ο• ΝΎ/ <|>-alg p q _ ⟩

  (do a ← Ο•
      x ← (p <|> q)
      b ← ψ x
      k x True) ∎