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

open import NeilPrelude
open import Equational

module Monad (M : Set → Set)

       (_>>='_ : {A B : Set} → M A → (A → M B) → M B)
       (ret    : {A : Set} → A → M A)

-- monad laws

       (1stLaw : {A B : Set} → {a : A} → {f : A → M B} → ret a >>=' f ≡ f a)
       (2ndLaw : {A : Set} → {ma : M A} → Extensionality → ma >>=' ret ≡ ma)
       (3rdLaw : {A B C : Set} → {f : A → M B} → {g : B → M C} → Extensionality
                 → (ma : M A) → (ma >>=' f) >>=' g ≡ ma >>=' (λ b → f b >>=' g))

where

infixl  _>>=_ _>>_

_>>=_ : {A B : Set} → M A → (A → M B) → M B
_>>=_ = _>>='_

_=<<_ : {A B : Set} → (A → M B) → M A → M B
_=<<_ = flip _>>=_

return : {A : Set} → A → M A
return = ret

ap : {A B : Set} → M (A → B) → M A → M B
ap mf ma = mf >>= λ f → ma >>= return ∘ f

liftM : {A B : Set} → (A → B) → M A → M B
liftM f = ap (return f)

join : {A : Set} → M (M A) → M A
join ma = ma >>= id

_>>_ : {A B : Set} → M A → M B → M B
ma >> mb = ma >>= const mb

_<<_ : {A B : Set} → M B → M A → M B
_<<_ = flip _>>_


--- Applicative Functor Instance ------------------------------------------


applicativeIdentLaw : {A : Set} → {fa : M A} → Extensionality → return id >>= (λ f → fa >>= return ∘ f) ≡ fa
applicativeIdentLaw ext = trans 1stLaw (2ndLaw ext)


applicativeCompLaw : {A B C : Set} → {fa : M A} → {fbc : M (B → C)} → {fab : M (A → B)} → Extensionality →
       ((ret _∘_ >>= (λ f → fbc >>= return ∘ f)) >>= (λ f → fab >>= return ∘ f)) >>= (λ g → fa >>= return ∘ g)
     ≡
       fbc >>= (λ f → (fab >>= λ g → fa >>= λ a → return (g a)) >>= λ b → return (f b))

applicativeCompLaw {_} {_} {_} {fa} {fbc} {fab} ext =

  ((ret _∘_ >>= (λ f → fbc >>= return ∘ f)) >>= (λ f → fab >>= return ∘ f)) >>= (λ g → fa >>= return ∘ g)

              ≡⟨ 3rdLaw ext (ret _∘_ >>= λ f → fbc >>= return ∘ f) ⟩

  (ret _∘_ >>= (λ f → fbc >>= return ∘ f)) >>= (λ b → (fab >>= return ∘ b) >>= (λ g → fa >>= return ∘ g))

              ≡⟨ cong (flip _>>=_ (λ b → fab >>= return ∘ b >>= (λ g → fa >>= return ∘ g))) 1stLaw ⟩

  (fbc >>= return ∘ _∘_) >>= (λ b → (fab >>= return ∘ b) >>= (λ g → fa >>= return ∘ g))

              ≡⟨ 3rdLaw ext fbc ⟩

  fbc >>= (λ f → return (_∘_ f) >>= (λ b → (fab >>= return ∘ b) >>= (λ g → fa >>= return ∘ g)))

              ≡⟨ cong (_>>=_ fbc) (ext (λ f → 
                                   return (_∘_ f) >>= (λ b → fab >>= return ∘ b >>= (λ g → fa >>= return ∘ g))
                                              ≡⟨ 1stLaw ⟩
                                   fab >>= return ∘ (_∘_ f) >>= (λ g → fa >>= return ∘ g)
                                              ≡⟨ 3rdLaw ext fab ⟩
                                   fab >>= (λ h → return (f ∘ h) >>= (λ g → fa >>= return ∘ g) )
                                              ≡⟨ cong (_>>=_ fab) (ext (λ h →
                                                                             return (f ∘ h) >>= (λ g → fa >>= return ∘ g)
                                                                                         ≡⟨ 1stLaw ⟩
                                                                             fa >>= return ∘ f ∘ h
                                                                                         ≡⟨ cong (_>>=_ fa) (ext (λ _ → sym 1stLaw) ) ⟩
                                                                             fa >>= (λ g → return (h g) >>= return ∘ f)
                                                                                         ≡⟨ sym (3rdLaw ext fa) ⟩
                                                                             (fa >>= return ∘ h) >>= return ∘ f
                                                                                         QED
                                                                  )) ⟩
                                   fab >>= (λ h → (fa >>= return ∘ h) >>= return ∘ f)
                                              ≡⟨ sym (3rdLaw ext fab) ⟩
                                   fab >>= (λ h → fa >>= return ∘ h) >>= return ∘ f
                                              QED
                                  )) ⟩

  fbc >>= (λ f → (fab >>= (λ h → fa >>= return ∘ h)) >>= return ∘ f)

                                 QED


applicativeHomoLaw : {A B : Set} → {f : A → B} → {a : A} → (return f >>= λ g → return a >>= λ b → return (g b)) ≡ return (f a)
applicativeHomoLaw = trans 1stLaw 1stLaw


applicativeInterLaw : {A B : Set} → {fab : M (A → B)} → {a : A} → Extensionality →
       fab >>= (λ f → ret a >>= (λ b → ret (f b))) ≡ ret (applyTo a) >>= (λ f → fab >>= (λ b → ret (f b)))
applicativeInterLaw {_} {_} {fab} {a} ext = fab >>= (λ f → ret a >>= (λ g → ret (f g)))
                                                                                              ≡⟨ cong (_>>=_ fab) (ext (λ _ → 1stLaw)) ⟩
                                            fab >>= (λ g → ret (g a))
                                                                                              ≡⟨ sym 1stLaw ⟩
                                            ret (applyTo a) >>= (λ f → fab >>= (λ g → ret (f g)))
                                                                                              QED

import Applicative
open Applicative M liftM return ap refl applicativeIdentLaw applicativeCompLaw applicativeHomoLaw applicativeInterLaw public