module combinators where
open import Data.Product
open import Data.Nat
open import Relation.Binary.PropositionalEquality

infixr 5 _⇒_
infixl 5 _·_

data Ty : Set where
  Nat : Ty
  _⇒_ : Ty → Ty → Ty

data Exp : Ty → Set where
  _·_ : {a b : Ty} → Exp (a ⇒ b) → Exp a → Exp b
  K : {a b : Ty} → Exp (a ⇒ b ⇒ a)
  S : {a b c : Ty} → Exp ((a ⇒ b ⇒ c) ⇒ (a ⇒ b) ⇒ a ⇒ c)
  zero : Exp Nat
  suc : Exp (Nat ⇒ Nat)
  rec : {a : Ty} → Exp (a ⇒ (Nat ⇒ a ⇒ a) ⇒ Nat ⇒ a)

data conv : {a : Ty} → Exp a → Exp a → Set where
  rec-zero : ∀ {a e f}  → conv {a} (rec · e · f · zero) e
  rec-suc : ∀ {a n e f} → conv {a} (rec · e · f · (suc · n)) (f · n · (rec · e · f · n)) 
  K : ∀ {a b x} {y : Exp b} → conv {a} (K · x · y) x
  S : ∀ {a b c x } {y : Exp (a ⇒ b)} {z : Exp a} → conv {c} (S · x · y · z) (x · z · (y · z))

data Cong (R : {a : Ty} → Exp a → Exp a → Set) : {a : Ty} → Exp a → Exp a → Set where
  cong-here : ∀ {a e e'} → R {a} e e' → Cong R e e'
  cong-refl : ∀ {a e} → Cong R {a} e e
  cong-sym : ∀ {a e e'} → Cong R {a} e e' → Cong R {a} e' e
  cong-trans : ∀ {a e e'} m → Cong R {a} e m → Cong R m e' → Cong R e e'
  cong-left : ∀ {a b e e'} r → Cong R {a ⇒ b} e e' → Cong R (e · r) (e' · r)
  cong-right : ∀ {a b} l {e e'} → Cong R {a} e e' → Cong R {b} (l · e) (l · e')

cong-cong₂ : {R : {a : Ty} → Exp a → Exp a → Set} {a b c : Ty} → (f : Exp (a ⇒ b ⇒ c)) →
    {ea ea' : Exp a} → Cong R ea ea' → {eb eb' : Exp b} → Cong R eb eb' → Cong R (f · ea · eb) (f · ea' · eb')
cong-cong₂ f equiva equivb = cong-trans _ (cong-right _ equivb) (cong-left _ (cong-right f equiva))

_~_ : {a : Ty} → Exp a → Exp a → Set
_~_ = Cong conv

⟦_⟧ : Ty → Set 
⟦ Nat ⟧ = ℕ
⟦ a ⇒ b ⟧ = Exp (a ⇒ b) × (⟦ a ⟧ → ⟦ b ⟧)

reify : {a : Ty} → ⟦ a ⟧ → Exp a
reify {Nat} zero = zero
reify {Nat} (suc n) = suc · reify n
reify {a ⇒ b} (exp , _) = exp

appsem : {a b : Ty} → ⟦ a ⇒ b ⟧ → ⟦ a ⟧ → ⟦ b ⟧
appsem (_ , f) q = f q

recs : {A : Set} → A → (ℕ → A → A) → ℕ → A
recs e f zero = e
recs e f (suc n) = f n (recs e f n)

recsem : {a : Ty} → ⟦ a ⟧ → ⟦ Nat ⇒ a ⇒ a ⟧ → ⟦ Nat ⟧ → ⟦ a ⟧
recsem e f zero = e
recsem e f (suc n) = appsem (appsem f n) (recsem e f n)

norm : (a : Ty) → Exp a → ⟦ a ⟧
norm a (y · y') = appsem (norm _ y) (norm _ y')
norm ._ K = K , (λ x → (K · reify x) , (λ y → x))
norm ._ S = S , (λ x → (S · reify x) , (λ y → S · reify x · reify y , (λ z → appsem (appsem x z) (appsem y z))))
norm ._ zero = zero
norm ._ suc = suc , suc
norm ._ (rec {a}) = rec , (λ e → rec · reify e
                        , (λ f → (rec · reify e · reify f)
                        , (λ n → recsem e f n)))

forward-lem : {a : Ty} {e e' : Exp a} → e ~ e' → norm _ e ≡ norm _ e'
forward-lem (cong-here rec-zero) = refl
forward-lem (cong-here rec-suc) = refl
forward-lem (cong-here K) = refl
forward-lem (cong-here S) = refl
forward-lem cong-refl = refl
forward-lem (cong-sym e'~e) = sym (forward-lem e'~e)
forward-lem (cong-trans m e~m m~e') = trans (forward-lem e~m) (forward-lem m~e')
forward-lem (cong-left r l~l') = cong (λ x → appsem x (norm _ r)) (forward-lem l~l')
forward-lem (cong-right l r~r') = cong (appsem (norm _ l)) (forward-lem r~r')

nbe : {a : Ty} → Exp a → Exp a
nbe e = reify (norm _ e)

forward-thm : {a : Ty} {e e' : Exp a} → e ~ e' → nbe e ≡ nbe e'
forward-thm equiv = cong reify (forward-lem equiv)

-- this is where it all goes wrong.
-- to prove e ~ nbe e,
-- I end up reintroducing all the proof components the slides say can be stripped from ⟦ a ⟧

-- a ⟦_⟧ soundly represents an expression if the normal form is equivalent to the expression
-- and the semantic function (if any), takes a ⟦_⟧ which soundly represents an argument to
-- a result which soundly represents the application. expression.
Sound : {a : Ty} → Exp a → ⟦ a ⟧ → Set
Sound {Nat} e k = e ~ reify k
Sound {a ⇒ b} e (nfe , sem) = e ~ nfe × ((a : Exp a) (ra : ⟦ a ⟧) → Sound a ra → Sound (e · a) (sem a))

-- soundness can be transported along convertibility
sound-conv : {a : Ty} (e e' : Exp a) (r : ⟦ a ⟧) → e ~ e' → Sound e' r → Sound e r
sound-conv {Nat} e e' r conv s = cong-trans e' conv s
sound-conv {y ⇒ y'} e e' r conv (eqv , sem) = (cong-trans e' conv eqv) ,
  (λ ea rea x → sound-conv (e · ea) (e' · ea) (appsem r rea) (cong-left _ conv) (sem ea rea x))

sound-reify : {a : Ty} {e : Exp a} {r : ⟦ a ⟧} → Sound e r → e ~ reify r
sound-reify {Nat} x = x
sound-reify {a ⇒ b} (conv , _) = conv

-- to apply 
sound-app : {a b : Ty} {e : Exp (a ⇒ b)} {r : ⟦ a ⇒ b ⟧} (se : Sound e r)
                       {ea : Exp a} {ra : ⟦ a ⟧} (sa : Sound ea ra) → Sound (e · ea) (appsem r ra)
sound-app (_ , fun) sa = fun _ _ sa


-- break out the soundness of recsem separately
rec-sound : {a : Ty}
  (ea : Exp a) (ra : ⟦ a ⟧) (x : Sound ea ra)
  (ea' : Exp (Nat ⇒ a ⇒ a)) (ra' : ⟦ Nat ⇒ a ⇒ a ⟧) (x' : Sound ea' ra')
  (k : ℕ) →
     Sound (rec · ea · ea' · reify k) (recsem ra ra' k)
rec-sound ea ra x ea' ra' x' zero =
   sound-conv (rec · ea · ea' · zero) ea ra (cong-here rec-zero) x
rec-sound ea ra x ea' ra' x' (suc n) =
   sound-conv (rec · ea · ea' · (suc · reify n)) (ea' · reify n · (rec · ea · ea' · reify n))
       (appsem (appsem ra' n) (recsem ra ra' n)) (cong-here rec-suc) (sound-app (sound-app x' cong-refl) (rec-sound ea ra x ea' ra' x' n))

-- soundness of the normalization function
norm-sound : {a : Ty} (e : Exp a) → Sound e (norm a e)
norm-sound (y · y') = sound-app (norm-sound y) (norm-sound y')
norm-sound K = cong-refl , (λ ea rea x → cong-right K (sound-reify x) , (λ ea' rea' x' →
   sound-conv (K · ea · ea') ea rea (cong-here K) x))
norm-sound S = cong-refl , (λ ea rea x → (cong-right S (sound-reify x)) , (λ ea' rea' x' →
   cong-cong₂ S (sound-reify x) (sound-reify x') , (λ ea0 rea0 x0 →
   sound-conv (S · ea · ea' · ea0) (ea · ea0 · (ea' · ea0)) (appsem (appsem rea rea0) (appsem rea' rea0))
    (cong-here S) (sound-app (sound-app x x0) (sound-app x' x0)))))
norm-sound zero = cong-refl
norm-sound suc = cong-refl , (λ x x' → cong-right suc)
norm-sound rec = cong-refl , (λ ea rea x → (cong-right rec (sound-reify x)) , (λ ea' rea' x' →
   (cong-cong₂ rec (sound-reify x) (sound-reify x')) , (λ ea0 rea0 x0 →
     sound-conv (rec · ea · ea' · ea0) (rec · ea · ea' · reify rea0) (recsem rea rea' rea0) (cong-right _ x0) 
     (rec-sound ea rea x ea' rea' x' rea0))))

-- now the proof goes through
backward-lem : {a : Ty} (e : Exp a) → e ~ nbe e
backward-lem e = sound-reify (norm-sound e)

backward-thm : {a : Ty} {e e' : Exp a} → nbe e ≡ nbe e' → e ~ e'
backward-thm {_} {e} {e'} eq = cong-trans (nbe e) (backward-lem _)
  (subst (λ x → x ~ e') (sym eq) (cong-sym (backward-lem _)))