module Uniform.posIR
  (C : Set1)
  (HomC : C -> C -> Set)
  (idC : (c : C) -> HomC c c)
  (_∘C_ : {a b c : C} -> HomC b c -> HomC a b -> HomC a c)
  where

-- A formalisation of uniform positive IR, inspired by work by Peter
-- Hancock

open import Data.Product as P
open import Function
open import Data.Unit
open import Level

import Prelude
open Prelude C HomC idC _∘C_
import posIR
open posIR C HomC idC _∘C_


-- uniform codes are defined in a context of what we have seen
-- previously. 
mutual
  data Ctxt  : Set1 where
    [] : Ctxt
    _,σ_ : (Γ : Ctxt) -> (Env Γ -> Set) -> Ctxt
    _,δ_ : (Γ : Ctxt) -> (Env Γ -> Set) -> Ctxt

  Env : Ctxt -> Set1
  Env [] =  Lift ⊤
  Env (Γ ,σ A) = Σ[ x ∈ (Env Γ) ]  (A x)
  Env (Γ ,δ A) = Σ[ x ∈ (Env Γ) ]  (A x -> C)

-- Uniform codes in context Γ. The codes we are really interested in
-- are in the empty context (hence the definition of UIR+ below).
mutual

  data IR' (Γ : Ctxt) : Set1 where
    ι : (Env Γ -> C) -> IR' Γ
    σ : (A : Env Γ -> Set) -> IR' (Γ ,σ A) -> IR' Γ
    δ : (A : Env Γ -> Set) -> (γ : IR' (Γ ,δ A)) ->
        (γ→ : ∀ {e f g} -> ((a : A e) -> HomC (f a) (g a))
                                  -> Hom' (e , f) (e , g) γ γ)
          -> IR' Γ

  data Hom' : {Γ Δ : Ctxt} -> Env Γ -> Env Δ -> IR' Γ -> IR' Δ -> Set1 where

    ι→ι : ∀ {Γ Δ} -> {c : Env Γ -> C}{c' : Env Δ -> C} ->
            {e : Env Γ}{e' : Env Δ} ->
            HomC (c e) (c' e') -> Hom' {Γ} {Δ} e e' (ι c) (ι c')
    σ→σ : ∀ {Γ Δ A B γ γ'}-> {e : Env Γ}{e' : Env Δ} -> (f : A e -> B e') ->
             ((x : A e) -> Hom' (e , x) (e' , f x) γ γ') ->
             Hom' {Γ} {Δ} e e' (σ A γ) (σ B γ')
    δ→δ : ∀ {Γ Δ e e' A B γ γ'} ->
            {γ→ : ∀ {e}{f g : A e -> C} -> ((a : A e) -> HomC (f a) (g a)) ->
                  Hom' {Γ ,δ A} {Γ ,δ A} (e , f) (e , g) γ γ} ->
            {γ'→ : ∀ {e f g} -> ((a : B e) -> HomC (f a) (g a)) ->
                  Hom' {Δ ,δ B} {Δ ,δ B} (e , f) (e , g) γ' γ'} ->
            (f : B e' -> A e) ->
            ((h : A e -> C) -> Hom' (e , h) (e' , h ∘ f) γ γ') ->
             Hom' {Γ} {Δ} e e' (δ A γ γ→) (δ B γ' γ'→)


-- In the empty context:

UIR+ : Set1
UIR+ = IR' []

UHom : UIR+ -> UIR+ -> Set1
UHom = Hom' (lift tt) (lift tt)

------------------------------------------------------------------

-- Identity ------------------------------------------------------

idIR' : {Γ : Ctxt} -> {e : Env Γ} -> (γ : IR' Γ) -> Hom' e e γ γ
idIR' (ι d) = ι→ι (idC _)
idIR' (σ A γ) = σ→σ id (λ x → idIR' γ)
idIR' (δ A γ γ→) = δ→δ id (λ h → idIR' γ)

-- In the empty context:

idUIR+ : (γ : UIR+) -> UHom γ γ
idUIR+ = idIR' 

------------------------------------------------------------------

-- Composition ---------------------------------------------------

_∘IR'_ : {Γ Γ' Γ'' : Ctxt} -> {e : Env Γ}{e' : Env Γ'}{e'' : Env Γ''} ->
         {γ : IR' Γ}{γ' : IR' Γ'}{γ'' : IR' Γ''} ->
         Hom' e' e'' γ' γ'' -> Hom' e e' γ γ' -> Hom' e e'' γ γ''
ι→ι f ∘IR' ι→ι f' = ι→ι (f ∘C f')
σ→σ f g ∘IR' σ→σ f' g' = σ→σ (f ∘ f') (λ a → g (f' a) ∘IR' g' a)
δ→δ f g ∘IR' δ→δ f' g' = δ→δ (f' ∘ f) (λ h → g (h ∘ f') ∘IR' g' h)

-- In the empty context:

_∘UIR+_ : {γ γ' γ'' : UIR+} -> UHom γ' γ'' -> UHom γ γ' -> UHom γ γ''
_∘UIR+_ = _∘IR'_ 

------------------------------------------------------------------

-- Decoding ------------------------------------------------------

-- Object part

⟦_⟧IR' : {Γ : Ctxt} -> IR' Γ -> Env Γ -> FamC -> FamC
⟦ ι d ⟧IR' e UT = ⊤ , (λ _ → d e)
⟦ σ A γ ⟧IR' e UT = ΣFamC (A e) (λ a → ⟦ γ ⟧IR' (e , a) UT)
⟦ δ A γ γ→ ⟧IR' e (U , T)
  = ΣFamC (A e -> U) (λ g → ⟦ γ ⟧IR' (e , T ∘ g) (U , T))

-- Morphism part

mutual
  ⟦_⟧→' : {Γ : Ctxt} -> (γ : IR' Γ) -> (e : Env Γ) ->
          {A B : FamC} -> HomFamC A B ->
          HomFamC (⟦ γ ⟧IR' e A) (⟦ γ ⟧IR' e B)
  ⟦ ι d ⟧→' e hk = id , (λ _ → idC (d e))
  ⟦ σ A γ ⟧→' e hk = [ inA a ∘FamC ⟦ γ ⟧→' _ hk ][ a ∈ (A e) ]
  ⟦_⟧→' (δ A γ γ→) e {X , P} {Y , Q} (h , k)
       = [ inA {B = λ g → ⟦ γ ⟧IR' _ (Y , Q)} (h ∘ g)
           ∘FamC (⟦ γ→ (k ∘ g) ⟧Hom' (Y , Q)
           ∘FamC ⟦ γ ⟧→' _ (h , k)) ][ g ∈ _ ]

  ⟦_⟧Hom' : {Γ Δ : Ctxt} -> {e : Env Γ}{e' : Env Δ} ->
            {γ : IR' Γ}{γ' : IR' Δ} ->
            Hom' e e' γ γ' ->
            (A : FamC) ->            
            HomFamC (⟦ γ ⟧IR' e A) (⟦ γ' ⟧IR' e' A)
  ⟦ ι→ι f ⟧Hom' UT = id , (λ _ → f)
  ⟦ σ→σ f g ⟧Hom' UT = ΣFamCmor f (λ a → ⟦ g a ⟧Hom' UT)
  ⟦ δ→δ f g ⟧Hom' (U , T)
    = ΣFamCmor (λ z → z ∘ f) (λ h → ⟦ g (T ∘ h) ⟧Hom' (U , T))

-- In the empty context:

⟦_⟧UIR+ : UIR+ -> FamC -> FamC
⟦ γ ⟧UIR+ = ⟦ γ ⟧IR' (lift tt)

⟦_⟧U→ : (γ : UIR+) -> {A B : FamC} ->
        HomFamC A B -> HomFamC (⟦ γ ⟧UIR+ A) (⟦ γ ⟧UIR+ B)
⟦ γ ⟧U→ = ⟦_⟧→' γ (lift tt)

⟦_⟧UHom : {γ γ' : UIR+} -> UHom γ γ' -> (X : FamC) ->
          HomFamC (⟦ γ ⟧UIR+ X) (⟦ γ' ⟧UIR+ X)
⟦ f ⟧UHom = ⟦ f ⟧Hom' 

------------------------------------------------------------------

-- Translation to ordinary positive codes ------------------------

mutual
  U2DS' : {Γ : Ctxt} -> Env Γ -> IR' Γ -> IR+
  U2DS' e (ι d) = ι (d e)
  U2DS' e (σ A γ) = σ (A e) (λ a → U2DS' (e , a) γ)
  U2DS' e (δ A γ γ→) = δ (A e) (λ g → U2DS' (e , g) γ) (Hom'2Hom ∘ γ→)

  Hom'2Hom : {Γ Δ : Ctxt} -> {e : Env Γ}{e' : Env Δ} ->
               {γ : IR' Γ}{γ' : IR' Δ} ->
               Hom' e e' γ γ' -> Hom (U2DS' e γ) (U2DS' e' γ')
  Hom'2Hom (ι→ι f) = ι→ι f
  Hom'2Hom (σ→σ f g) = σ→σ f (Hom'2Hom ∘ g)
  Hom'2Hom (δ→δ f g) = δ→δ f (Hom'2Hom ∘ g)



UIR+-to-IR+ : UIR+ -> IR+
UIR+-to-IR+ = U2DS' (lift tt)

UHom-to-Hom :  {γ γ' : UIR+} -> UHom γ γ' ->
               Hom (UIR+-to-IR+ γ) (UIR+-to-IR+ γ')
UHom-to-Hom = Hom'2Hom {[]}

-- Context morphisms

mutual
  IRcMor : ∀ {Γ Δ} -> (f : Env Γ -> Env Δ) -> IR' Δ -> IR' Γ
  IRcMor {Γ} f (ι d) = ι (d ∘ f)
  IRcMor {Γ} {Δ} f (σ A γ) = σ (A ∘ f) (IRcMor (P.map f id) γ)
  IRcMor {Γ} {Δ} f (δ A γ γ→)
   = δ (A ∘ f) (IRcMor (P.map f id) γ) (HomcMor (P.map f id) (P.map f id) ∘ γ→)

  HomcMor : ∀ {Γ Δ Γ' Δ'} -> (f : Env Γ -> Env Δ)(f' : Env Γ' -> Env Δ') ->
            {γ : IR' Δ}{γ' : IR' Δ'}{e : Env Γ}{e' : Env Γ'} -> 
            Hom' (f e) (f' e') γ γ' ->
            Hom' {Γ} {Γ'} e e' (IRcMor f γ) (IRcMor f' γ')
  HomcMor f f' (ι→ι g) = ι→ι g
  HomcMor f f' (σ→σ g h) = σ→σ g (HomcMor (P.map f id) (P.map f' id) ∘ h)
  HomcMor f f' (δ→δ g h) = δ→δ g (HomcMor (P.map f id) (P.map f' id) ∘ h)

-- in particular, weakening

IRwkσ : forall {Γ B} -> IR' Γ -> IR' (Γ ,σ B)
IRwkσ {Γ} = IRcMor proj₁

IRwkδ : forall {Γ B} -> IR' Γ -> IR' (Γ ,δ B)
IRwkδ {Γ} = IRcMor proj₁