{-# OPTIONS --type-in-type #-}

module State where
open import Base
open import Monoidal
open import Concurrent

open import Posets

-- meet semilattice
record Semilattice : Set where
  infix  _∧_
  field
    car : Set
    Top : car
    _∧_ : car → car → car
    ∧sym : {x y : car} → x ∧ y ≡ y ∧ x
    ∧idem : {x : car} → x ∧ x ≡ x
    ∧assoc : {x y z : car} → (x ∧ y) ∧ z ≡ x ∧ (y ∧ z)
    ∧unital : {x : car} → x ∧ Top ≡ x
    -- this should just be an auxillary function
    ∧swap : {a b c d : car} → (a ∧ b) ∧ (c ∧ d) ≡ (a ∧ c) ∧ (b ∧ d)

-- Semilattice → Poset
asPoset : Semilattice → Poset
asPoset sl = let open Semilattice sl in record
  { car = car
  ; _≤_ = λ a b → a ∧ b ≡ a
  ; ≤refl = ∧idem
  ; ≤trans = λ {a b c} f g →
    begin
      a ∧ c
    ≡⟨ cong (λ x → x ∧ c) (sym f) ⟩
      (a ∧ b) ∧ c
    ≡⟨ ∧assoc ⟩
      a ∧ (b ∧ c)
    ≡⟨ cong (λ x → a ∧ x) g ⟩
      a ∧ b
    ≡⟨ f ⟩
      a
    ∎
  ; ≤antisym = λ {a b} f g →
    begin
      a
    ≡⟨ sym f ⟩
      a ∧ b
    ≡⟨ ∧sym ⟩
      b ∧ a
    ≡⟨ g ⟩
      b
    ∎
  }

MonotoneIsSLH : {S T : Semilattice}
              → (h : (asPoset S →ₘ asPoset T))
              → (s₀ s₁ : Semilattice.car S)
              → (sm : Poset._≤_ (asPoset S) s₀ s₁)
              → (_→ₘ_.fun h) ((Semilattice._∧_ S) s₀ s₁)
                ≡ Semilattice._∧_ T (_→ₘ_.fun h s₀) (_→ₘ_.fun h s₁)
MonotoneIsSLH {S} {T} (h ! hm) s₀ s₁ sm =
  let module S = Semilattice S
      module T = Semilattice T
  in begin
    h (s₀ S.∧ s₁)
  ≡⟨ cong h sm ⟩
    h s₀
  ≡⟨ sym (hm sm) ⟩
    h s₀ T.∧ h s₁
  ∎

a∧b≤a : {L : Semilattice} → {a b : Semilattice.car L}
  → Semilattice._∧_ L (Semilattice._∧_ L a b) a ≡ (Semilattice._∧_ L a b)
a∧b≤a {L} {a} {b} = let module L = Semilattice L in
  begin
    (a L.∧ b) L.∧ a
  ≡⟨ L.∧sym ⟩
    a L.∧ (a L.∧ b)
  ≡⟨ sym L.∧assoc  ⟩
    (a L.∧ a) L.∧ b
  ≡⟨ cong (λ x → x L.∧ b) L.∧idem ⟩
    (a L.∧ b) ∎

a∧b≤b : {L : Semilattice} → {a b : Semilattice.car L}
  → Semilattice._∧_ L (Semilattice._∧_ L a b) b ≡ (Semilattice._∧_ L a b)
a∧b≤b {L} {a} {b} = let module L = Semilattice L in
  begin
    (a L.∧ b) L.∧ b
  ≡⟨ L.∧assoc ⟩
    a L.∧ (b L.∧ b)
  ≡⟨ cong (λ x → a L.∧ x) L.∧idem ⟩
    a L.∧ b
  ∎

-- the state monad is an example of a concurrent monad
-- on the monoidal category MonPoset

-- if Poset were a predicate on a Set, then this would be something like
-- Poset A → Poset B → Poset (A × B)
-- but that does increase complexity of types, a lot of implicit stuff
infix  _×,≤_
_×,≤_ : Poset → Poset → Poset
_×,≤_ A B =
  let module A = Poset A
      module B = Poset B
  in record
    { car = A.car × B.car 
    ; _≤_ = λ (a1 , b1) (a2 , b2) → a1 A.≤ a2 × b1 B.≤ b2 
    ; ≤refl = A.≤refl , B.≤refl 
    ; ≤trans = λ (f1 , g1) (f2 , g2) → A.≤trans f1 f2 , B.≤trans g1 g2
    ; ≤antisym = λ (x₁ , y₁) (x₂ , y₂)
        → let xc = A.≤antisym x₁ x₂
              yc = B.≤antisym y₁ y₂
          in trans (cong (λ b → (b , _)) xc) (cong (λ b → (_ , b)) yc)  
    }

MonPoset : MonoidalCategory
MonPoset = record
  { c = Poset
  ; MonoidalProduct₀ = λ (A , B) → A ×,≤ B
  ; Identity₀ = λ tt → trivialPoset
  ; P-monoidal = record
    { P = PosetC
    ; MonoidalProduct = record
      { prefunctor = record
        { F₁ = λ ((f ! fm) , (g ! gm)) →
          (λ {(x , z) → f x , g z}) !
          λ (a , b) → fm a , gm b
        ; preserveOrder = λ (fg1 , fg2) (a , b) → fg1 a , fg2 b
        }
      ; isParamFunctor = record
        { strictIdentity = refl
        ; strictComposition = refl
        }
      }
    ; funcI = record
      { prefunctor = record
        { F₁ = λ _ → (λ _ → tt) ! (λ _ → tt)
        ; preserveOrder = λ _ _ → tt
        }
      ; isParamFunctor = record
        { strictIdentity = refl
        ; strictComposition = refl
        }
      }
    ; isParamMonoidalCategory = record
      { λ-at = λ X → (λ (i , x) → x) ! (λ (i , xm) → xm)
      ; λ-naturality = λ {A} {B} f → refl
      ; ρ-at = λ X → (λ x → x , tt) ! (λ x → x , tt)
      ; ρ-naturality = λ {A} {B} f → refl
      ; α-at = λ X Y Z →
                  (λ ((x , y) , z) → x , (y , z)) !
                  (λ ((x , y) , z) → x , (y , z))
      ; α-naturality = λ f g h → refl
      ; triangle = λ X Y → refl
      ; pentagon = λ X Y Z W → refl
      }
    }
  }

-- here we use the same method as in the Kleisli category code,
-- we move definitions outside of the record to encapsulate code and
-- give typing information
SM-F₀ : Semilattice → Poset → Poset
SM-F₀ S X = (asPoset S) →ₘ,≤ ((asPoset S) ×,≤ X)

-- here a let binding on the type level would be really nice,
-- asPoset L is just noise. we could call it typelet?
SM-ψ-at : (L : Semilattice) → (X : Poset) → (Y : Poset)
        →   (((asPoset L) →ₘ,≤ ((asPoset L) ×,≤ X)) ×,≤ ((asPoset L)
            →ₘ,≤ ((asPoset L) ×,≤ Y)))
          →ₘ
            ((asPoset L) →ₘ,≤ ((asPoset L) ×,≤ (X ×,≤ Y)))
SM-ψ-at L X Y =
 let module L = Semilattice L in
  (λ ((g ! gm) , (h ! hm))
   → (λ s → let (s₀ , x) = g s
                (s₁ , y) = h s
            in (s₀ L.∧ s₁) , (x , y))
   ! λ {a b} ab → let (x1 , _) = g a
                      (x2 , _) = g b
                      (x3 , _) = h a
                      (x4 , _) = h b
                      (xx , yy) = gm ab
                      (zz , ww) = hm ab
                  in (begin
                    (x1 L.∧ x3) L.∧ (x2 L.∧ x4)
                  ≡⟨ L.∧swap ⟩
                    (x1 L.∧ x2) L.∧ (x3 L.∧ x4)
                  ≡⟨ cong (λ x → x L.∧ (x3 L.∧ x4)) xx ⟩
                    x1 L.∧ (x3 L.∧ x4)
                  ≡⟨ cong (λ x → x1 L.∧ x) zz ⟩
                    x1 L.∧ x3 ∎) , yy , ww)
  ! λ {((m1 ! m1m) , (n1 ! n1m)) ((m2 ! m2m) , (n2 ! n2m))} (m , n) t
      → let (a , _) = m1 t
            (b , _) = n1 t
            (c , _) = m2 t
            (d , _) = n2 t
            (xx , yy) = m t
            (zz , ww) = n t
         in ((begin
                (a L.∧ b) L.∧ (c L.∧ d)
              ≡⟨ L.∧swap ⟩
                (a L.∧ c) L.∧ (b L.∧ d)
              ≡⟨ cong (λ x → x L.∧ (b L.∧ d)) xx ⟩
                a L.∧ (b L.∧ d)
              ≡⟨ cong (λ x → a L.∧ x) zz ⟩
                a L.∧ b ∎)) , yy , ww

StateMonad : Semilattice → ConcurrentMonad MonPoset
StateMonad L =
  let S = asPoset L
      module L = Semilattice L
      module S = Poset (asPoset L)
      open MonoidalCategory MonPoset using (α-at)
  in record
  { M = record
    { F₀ = SM-F₀ L
    ; PF = record
      { prefunctor = record
        { F₁ = λ {X} {Y} (f ! fm) → (λ (xstate ! xstatem) →
           (λ s → let (s , x) = xstate s in (s , f x))
              ! λ c → let (x , y) = xstatem c in x , fm y)
                ! λ c s → let (x , y) = c s in x , fm y
        ; preserveOrder = λ g (f ! fm) x →
            let (s , b) = f x in S.≤refl , g b
        }
      ; isParamFunctor = record
        { strictIdentity = refl
        ; strictComposition = refl
        }
      }
    }
  ; isConcurrentMonad = record
    { η-at = λ X → let module X = Poset X in
      (λ x → (λ s → s , x) ! λ x → x , X.≤refl) ! λ reflm _ → S.≤refl , reflm
    ; η-naturality = λ {C D bf@(f ! fm)} →
      begin
        (λ x → (λ s → s , f x) ! (λ c → c , Poset.≤refl D)) ! (λ p z → L.∧idem , fm p)
      -- we can set up some kind of proof homomorphism that preserves
      -- the axioms of poset, in this case ≤refl. we can also assume
      -- uniqueness of inequality proofs. or we can simply define a
      -- special equality on →ₘ, monotone functions, that only looks
      -- at the function part
      ≡⟨ →ₘeq (funext λ x → →ₘeq refl) ⟩
        (λ x → (λ s → s , f x) ! (λ c → c , fm (Poset.≤refl C))) ! (λ p s → L.∧idem , fm p)
      ∎
    ; μ-at = λ X → let module X = Poset X in (λ bull@(f ! fm) →
          (λ s → let gull@(s' , sull@(t ! tm)) = f s in t s')
            ! (λ {x x'} c →
              let (s₀ , (g₁ ! gm₁)) = f x
                  (s₁ , (h₁ ! hm₁)) = f x'
                  (s₀≤s₁ , l) = fm c in
            S.≤trans (l s₀ .proj₁) (hm₁ s₀≤s₁ .proj₁)
          , X.≤trans (gm₁ s₀≤s₁ .proj₂) (l s₁ .proj₂)))
                ! λ {bf@(f ! fm) bg@(g ! gm)} x s →
                    let (s₀ , (h₀ ! hm₀)) = f s
                        (s₁ , (h₁ ! hm₁)) = g s
                    in
                    (begin
                      h₀ s₀ .proj₁ L.∧ h₁ s₁ .proj₁
                    ≡⟨ S.≤trans (x s .proj₂ s₀ .proj₁)
                        (g s .proj₂ ._→ₘ_.monot (x s .proj₁) .proj₁) ⟩
                      h₀ s₀ .proj₁
                    ∎)
                    , X.≤trans (x s .proj₂ s₀ .proj₂)
                       (g s .proj₂ ._→ₘ_.monot (x s .proj₁) .proj₂)
    ; μ-naturality =  →ₘeq (funext λ x → →ₘeq refl)
    ; ψ⁰ = (λ _ → (λ x → L.Top , tt)
             ! λ _ → S.≤refl , tt)
               ! λ _ _ → S.≤refl , tt
    ; ψ-at = SM-ψ-at L
    ; ψ-naturality = →ₘeq (funext λ x → →ₘeq refl)
    ; diagram1 = →ₘeq (funext λ x → →ₘeq refl)
    ; diagram2 = →ₘeq (funext λ x → →ₘeq refl)
    ; diagram3 = →ₘeq (funext λ x → →ₘeq refl)
    ; diagram4 = →ₘeq (funext λ (tt , (x ! xm)) → →ₘeq (funext λ y →
      let (th₁ , th₂) = x y in
      begin
         (th₁ , th₂)
      ≡⟨ cong (λ b → (b , th₂)) (trans (sym L.∧unital) L.∧sym ) ⟩
        (L.Top L.∧ th₁ , th₂)
      ∎))
    ; diagram5 = →ₘeq (funext λ (x ! xm)
        → →ₘeq (funext λ s → cong (λ b → (b , x s .proj₂ , tt)) L.∧unital))
    ; diagram6 = →ₘeq (funext λ (x) → →ₘeq (funext λ s
        → cong (λ b → (b , _ , α-at (SM-F₀ L _) (SM-F₀ L _) (SM-F₀ L _)
                               ._→ₘ_.fun x .proj₂ .proj₁ ._→ₘ_.fun s .proj₂
                               ,
                               α-at (SM-F₀ L _) (SM-F₀ L _) (SM-F₀ L _)
                               ._→ₘ_.fun x .proj₂ .proj₂ ._→ₘ_.fun s .proj₂)) (sym L.∧assoc)))
    ; inequation1 = λ _ _ → L.∧unital , tt
    ; inequation2 = λ _ _ → S.≤refl , tt
    ; inequation3 = λ {X Y} s₀ s₁ →
      trans (cong (λ x → s₁ L.∧ x) L.∧idem) L.∧idem
      , (Poset.≤refl X , Poset.≤refl Y)
    ; inequation4 = 
      λ {X Y} gg@((f ! fm) , (g ! gm)) s
      → let (s₁ , (ph₁ ! phm₁)) = f s
            (s₂ , (ph₂ ! phm₂)) = g s
            h₁ = λ x → ph₁ x .proj₁
            h₂ = λ x → ph₂ x .proj₁
            bing = (λ w → ph₁ w .proj₁) ! (λ w → phm₁ w .proj₁)
            bing2 = (λ w → ph₂ w .proj₁) ! (λ w → phm₂ w .proj₁)

      in (begin
        (h₁ (s₁ L.∧ s₂) L.∧ h₂ (s₁ L.∧ s₂)) L.∧ (h₁ s₁ L.∧ h₂ s₂)
      ≡⟨ L.∧swap ⟩
        (h₁ (s₁ L.∧ s₂) L.∧ h₁ s₁) L.∧ (h₂ (s₁ L.∧ s₂) L.∧ h₂ s₂)
      ≡⟨ cong (λ x → x L.∧ (h₂ (s₁ L.∧ s₂) L.∧ h₂ s₂))
          (sym (MonotoneIsSLH {L} {L} bing (s₁ L.∧ s₂) s₁ (a∧b≤a {L} {s₁} {s₂} ))) ⟩
        h₁ ((s₁ L.∧ s₂) L.∧ s₁) L.∧ (h₂ (s₁ L.∧ s₂) L.∧ h₂ s₂)
      ≡⟨ cong (λ x → h₁ x L.∧ (h₂ (s₁ L.∧ s₂) L.∧ h₂ s₂)) (a∧b≤a {L} {s₁} {s₂}) ⟩
        h₁ (s₁ L.∧ s₂) L.∧ (h₂ (s₁ L.∧ s₂) L.∧ h₂ s₂)
      ≡⟨ cong (λ x → h₁ (s₁ L.∧ s₂) L.∧ x)
          (sym (MonotoneIsSLH {L} {L} bing2 (s₁ L.∧ s₂) s₂ (a∧b≤b {L} {s₁} {s₂} ))) ⟩
        h₁ (s₁ L.∧ s₂) L.∧ h₂ ((s₁ L.∧ s₂) L.∧ s₂)
      ≡⟨ cong (λ x → h₁ (s₁ L.∧ s₂) L.∧ h₂ x) (a∧b≤b {L} {s₁} {s₂}) ⟩
         h₁ (s₁ L.∧ s₂) L.∧ h₂ (s₁ L.∧ s₂)
      ∎)
      , (let open Poset-Reasoning X
         in (≤begin
              ph₁ (s₁ L.∧ s₂) .proj₂
            ≤⟨ phm₁ (a∧b≤a {L} {s₁} {s₂}) .proj₂ ⟩
              ph₁ s₁ .proj₂
            ≤∎))
      , (let open Poset-Reasoning Y
         in ≤begin
              ph₂ (s₁ L.∧ s₂) .proj₂
            ≤⟨ phm₂ (a∧b≤b {L} {s₁} {s₂}) .proj₂ ⟩
              ph₂ s₂ .proj₂
            ≤∎)
    }
  }