{-# OPTIONS --safe #-}

module Problem20 where

open import Data.Maybe as Maybe hiding (_>>=_; map)
open import Function
open import Relation.Binary.PropositionalEquality

open ≡-Reasoning

data Expr (V : Set) : Set where
   λ′_ : Expr (Maybe V) → Expr V
   _·_ : Expr V → Expr V → Expr V
   var : V → Expr V

variable V A B C D : Set

map : (A → B) → Expr A → Expr B
map f (λ′ e)  = λ′ map (Maybe.map f) e
map f (e · a) = map f e · map f a
map f (var x) = var (f x)

_>>=_ : Expr A → (A → Expr B) → Expr B
(λ′ e)   >>= f = λ′ (e >>= maybe (map just ∘ f) (var nothing))
(e · e₁) >>= f = (e >>= f) · (e₁ >>= f)
var x    >>= f = f x

-- Lean users have the advantage of functional extensionality, but Agda users don't
-- So I have provided these congruence lemmas so that Agda users are at no axiomatic disadvantage.

map-cong : ∀{f g : A → B}(ma : Expr A) → (∀ x → f x ≡ g x) → map f ma ≡ map g ma
map-cong (λ′ ma) eq = cong λ′_ (map-cong ma λ { nothing → refl ; (just x) → cong just (eq x) })
map-cong (ma · ma₁) eq = cong₂ _·_ (map-cong ma eq) (map-cong ma₁ eq)
map-cong (var x) eq = cong var (eq x)

>>=-cong : ∀{f g : A → Expr B}(ma : Expr A) → (∀ x → f x ≡ g x) → ma >>= f ≡ ma >>= g
>>=-cong (λ′ ma) eq = cong λ′_ (>>=-cong ma λ { nothing → refl ; (just x) → cong (map just) (eq x) })
>>=-cong (ma · ma₁) eq = cong₂ _·_ (>>=-cong ma eq) (>>=-cong ma₁ eq)
>>=-cong (var x) eq = eq x

-- Functor Laws

map-id : (ma : Expr A) → map id ma ≡ ma
map-id (λ′ ma) = cong λ′_ (begin
  map (Maybe.map id) ma ≡⟨ map-cong ma (λ { (just a) → refl ; nothing → refl }) ⟩
  map id ma             ≡⟨ map-id ma ⟩
  ma                    ∎)
map-id (ma · ma₁) = cong₂ _·_ (map-id ma) (map-id ma₁)
map-id (var x) = refl

map-homo : (f : B → C) (g : A → B) (ma : Expr A) → map f (map g ma) ≡ map (f ∘ g) ma
map-homo f g (λ′ ma) = cong λ′_ (begin
  map (Maybe.map f) (map (Maybe.map g) ma) ≡⟨ map-homo (Maybe.map f) (Maybe.map g) ma ⟩
  map (Maybe.map f ∘ Maybe.map g) ma       ≡⟨ map-cong ma (λ { (just a) → refl ; nothing → refl }) ⟩
  map (Maybe.map (f ∘ g)) ma               ∎)
map-homo f g (ma · ma₁) = cong₂ _·_ (map-homo f g ma) (map-homo f g ma₁)
map-homo f g (var x) = refl

-- Bind-var is map

>>=-var : (f : A → B) (ma : Expr A) → ma >>= (var ∘ f) ≡ map f ma
>>=-var f (λ′ ma) = cong λ′_ (begin
  ma >>= maybe (var ∘ just ∘ f) (var nothing) ≡⟨ >>=-cong ma (λ { (just a) → refl ; nothing → refl }) ⟩
  ma >>= (var ∘ Maybe.map f)                  ≡⟨ >>=-var (Maybe.map f) ma ⟩
  map (Maybe.map f) ma                       ∎)
>>=-var f (ma · ma₁) = cong₂ _·_ (>>=-var f ma) (>>=-var f ma₁)
>>=-var f (var x) = refl

>>=-assoc : (g : B → C) (f : A → Expr B) (ma : Expr A) → map g (ma >>= f) ≡ ma >>= (map g ∘ f)
>>=-assoc g f (λ′ ma) = cong λ′_ (begin
  map (Maybe.map g) (ma >>= maybe (map just ∘ f) (var nothing))   ≡⟨  >>=-assoc (Maybe.map g) (maybe (map just ∘ f) (var nothing)) ma ⟩
  ma >>= (map (Maybe.map g) ∘ maybe (map just ∘ f) (var nothing)) ≡⟨  >>=-cong ma (λ { (just a) → map-homo (Maybe.map g) just (f a) ; nothing → refl }) ⟩
  ma >>= maybe (map (Maybe.map g ∘ just) ∘ f) (var nothing)       ≡˘⟨ >>=-cong ma (λ { (just a) → map-homo just g (f a) ; nothing → refl }) ⟩
  ma >>= maybe (map just ∘ map g ∘ f) (var nothing)               ∎)
>>=-assoc g f (ma · ma₁) = cong₂ _·_ (>>=-assoc g f ma) (>>=-assoc g f ma₁)
>>=-assoc g f (var x) = refl

>>=-lift : (f : A → B) (g : B → Expr C) (h : A → Expr C) → g ∘ f ≗ h → (ma : Expr A) → map f ma >>= g ≡ ma >>= h
>>=-lift f g h gf≗h (λ′ ma) = cong λ′_ (>>=-lift _ _ _ (λ { (just a) → cong (map just) (gf≗h a) ; nothing → refl }) ma)
>>=-lift f g h gf≗h (ma · ma₁) = cong₂ _·_ (>>=-lift f g h gf≗h ma) (>>=-lift f g h gf≗h ma₁)
>>=-lift f g h gf≗h (var x) = gf≗h x

assoc : ∀(g : A → Expr B)(h : B → Expr C)(a : Expr A)
      → a >>= (λ a′ → g a′ >>= h) ≡ (a >>= g) >>= h
assoc g h (λ′ ma) = cong λ′_ (begin
  ma >>= maybe (λ a → map just (g a >>= h)) (var nothing)                                      ≡⟨ >>=-cong ma (λ { (just a) → lemma a ; nothing → refl }) ⟩
  (ma >>= (λ a → maybe (map just ∘ g) (var nothing) a >>= maybe (map just ∘ h) (var nothing))) ≡⟨ assoc (maybe (map just ∘ g) (var nothing)) (maybe (map just ∘ h) (var nothing)) ma ⟩
  (ma >>= maybe (map just ∘ g) (var nothing)) >>= maybe (map just ∘ h) (var nothing)           ∎)
  where
  lemma : (a : _) → map just (g a >>= h) ≡ map just (g a) >>= maybe (map just ∘ h) (var nothing)
  lemma a = begin
    map just (g a >>= h)                                  ≡⟨  >>=-assoc just h (g a) ⟩
    g a >>= (map just ∘ h)                                ≡˘⟨ >>=-lift just (maybe (map just ∘ h) (var nothing)) (map just ∘ h) (λ _ → refl) (g a) ⟩
    map just (g a) >>= maybe (map just ∘ h) (var nothing) ∎
assoc g h (ma · ma₁) = cong₂ _·_ (assoc g h ma) (assoc g h ma₁)
assoc g h (var x) = refl

identₗ : ∀(f : A → Expr B)(x : A) → var x >>= f ≡ f x
identₗ _ _ = refl

identᵣ : ∀(e : Expr A) → e >>= var ≡ e
identᵣ e = begin
  e >>= var ≡⟨ >>=-var id e ⟩
  map id e  ≡⟨ map-id e ⟩
  e         ∎