module Uniform.Coproduct {D : Set1} where

open import Prelude.Basic
open import Prelude.Equality
open import Prelude.Container

open import Data.Nat.Properties.Simple using (+-right-identity; +-suc)

open import Uniform.IR

-- "ranked" version of containerised uniform codes
mutual
  data UniR : ℕ -> Set1 where
    ιR : UniR 
    δσR : {n : ℕ} -> (G : UniR n )-> (InfoR G -> Cont) -> UniR (suc n)

  InfoR : {n : ℕ} -> UniR n -> Set1
  InfoR ιR = Lift ⊤
  InfoR (δσR G H) = Σ[ x ∈ InfoR G ] (cont (H x) D)

σR : {n : ℕ} -> (G : UniR n )-> (InfoR G -> Set) -> UniR (suc n)
σR G A = δσR G (λ γ → A γ , (λ _ → ⊥))

δR : {n : ℕ} -> (G : UniR n )-> (InfoR G -> Set) -> UniR (suc n)
δR G A = δσR G (λ γ → (⊤ , (λ _ → A γ)))

UFR : (E : Set1) -> ℕ -> Set1
UFR E n = Σ[ c ∈ UniR n ] (InfoR c -> E)

-- Going back and forth between ranked and unranked codes

length : Uni D -> ℕ
length ι = 
length (σ G A) = suc (length G)
length (δ G A) = suc (length G)

mutual
  containerise : (c : Uni D) -> UniR (length c)
  containerise ι = ιR
  containerise (σ G A) = σR (containerise G) (A ∘ containeriseInfo G)
  containerise (δ G A) = δR (containerise G) (A ∘ containeriseInfo G)

  containeriseInfo : (c : Uni D) -> InfoR (containerise c) -> Info c
  containeriseInfo ι = id
  containeriseInfo (σ G A) = map (containeriseInfo G) proj₁
  containeriseInfo (δ G A) = map (containeriseInfo G) proj₂

containeriseUF : {E : Set1} -> (R : UF D E) -> UFR E (length (proj₁ R))
containeriseUF (c , α)  = (containerise c , α ∘ containeriseInfo c)

mutual
  forget : {n : ℕ} -> (c : UniR n) -> Uni D
  forget ιR = ι
  forget (δσR G H)
    = δ (σ (forget G) (λ γ → proj₁ (H (forgetInfo G γ))))
            (λ { (γ , x) → proj₂ (H (forgetInfo G γ)) x})

  forgetInfo : {n : ℕ} -> (c : UniR n) -> Info (forget c) -> InfoR c
  forgetInfo ιR x = x
  forgetInfo (δσR G H) (( γ , x) , y) = forgetInfo G γ , x , y

forgetUFR : {E : Set1}{n : ℕ} -> UFR E n -> UF D E
forgetUFR (c , α) = (forget c , α ∘ forgetInfo c)

-- padding a ranked code to an arbitrary length

pad : {m : ℕ} -> (n : ℕ) -> (c : UniR m) -> UniR (suc (m + n))
pad {m} zero    c = subst (UniR ∘ suc) (sym (+-right-identity m))
                          (δσR c (λ _ → K1-Cont))
pad {m} (suc n) c = subst (UniR ∘ suc) (sym (+-suc m n))
                          (δσR (pad n c) (λ _ → K1-Cont))

padInfo : {m : ℕ} -> (n : ℕ)(c : UniR m) -> InfoR (pad n c) -> InfoR c
padInfo {m} zero c rewrite (+-right-identity m) = proj₁
padInfo {m} (suc n) c rewrite (+-suc m n) = padInfo n c ∘ proj₁

padUFR : {E : Set1}{m : ℕ} -> (n : ℕ) -> UFR E m -> UFR E (suc (m + n))
padUFR n (c , α) = pad n c , α ∘ padInfo n c

-- padding two codes to the same length (the maximum one)
padEqual : {m n : ℕ} -> (c : UniR m)(d : UniR n) -> Σ[ k ∈ ℕ ] (UniR k × UniR k)
padEqual {m} {n} c d with compare n m
... | less .n k    = _ ,  c , pad k d
... | equal .m     = _ , c , d
... | greater .m k = _ , pad k c , d

padEqualInfo₁ : {m n : ℕ} -> (c : UniR m)(d : UniR n) ->
              let (_ , c' , d') = padEqual c d in
              InfoR c' -> InfoR c
padEqualInfo₁ {m} {n} c d with compare n m
... | less .n k    = id
... | equal .m     = id
... | greater .m k = padInfo k c


padEqualInfo₂ : {m n : ℕ} -> (c : UniR m)(d : UniR n) ->
              let (_ , c' , d') = padEqual c d in
              InfoR d' -> InfoR d
padEqualInfo₂ {m} {n} c d with compare n m
... | less .n k    = padInfo k d
... | equal .m     = id
... | greater .m k = id

{-
padEqualInfo : {m n : ℕ} -> (c : UniR m)(d : UniR n) ->
              let (_ , c' , d') = padEqual c d in
              InfoR c' ⊎ InfoR d' -> InfoR c ⊎ InfoR d
padEqualInfo c d = ⊎-map (padEqualInfo₁ c d) (padEqualInfo₂ c d)
-}

pe₁ : ∀ {E}{n m} -> (R : UFR E n)(R' : UFR E m) ->
      UFR E (proj₁ (padEqual (proj₁ R) (proj₁ R')))
pe₁ (c , α) (d , β)
  = let (k , c' , d') = padEqual c d in c' , α ∘ padEqualInfo₁ c d

pe₂ : ∀ {E}{n m} -> (R : UFR E n)(R' : UFR E m) ->
      UFR E (proj₁ (padEqual (proj₁ R) (proj₁ R')))
pe₂ (c , α) (d , β)
  = let (k , c' , d') = padEqual c d in d' , β ∘ padEqualInfo₂ c d

-- coproduct of two codes with the same rank
mutual
  choice : {n : ℕ} -> (c d : UniR n) -> UniR (suc n)
  choice ιR ιR = σR ιR (λ _ → Bool)
  choice (δσR G A) (δσR G' A')
    = δσR (choice G G') ([ A , A' ]′ ∘ choiceInfo G G')

  choiceInfo : {n : ℕ}(c d : UniR n) ->
              InfoR (choice c d) -> InfoR c ⊎ InfoR d
  choiceInfo ιR ιR (_ , true , _) = inj₁ (lift tt)
  choiceInfo ιR ιR (_ , false , _) = inj₂ (lift tt)
  choiceInfo (δσR G A) (δσR G' A') (γ , x) with choiceInfo G G' γ
  choiceInfo _ _ (γ , x)  | inj₁ γG = inj₁ (γG , x)
  choiceInfo _ _ (γ , x') | inj₂ γG' = inj₂ (γG' , x')

choiceUFR : {E : Set1}{n : ℕ} -> (R R' : UFR E n) -> UFR E (suc n)
choiceUFR (c , α) (d , β) = (choice c d , [ α , β ] ∘ choiceInfo c d)

-- finally the coproduct of two arbitrary codes

{-
-- Would like to define

_+Uni_ : Uni D -> Uni D -> Uni D
c  +Uni d = let (n , c' , d') = padEqual (containerise c)
                                         (containerise d)
            in forget (choice c' d')

+Info : (c d : Uni D) -> Info D (c +Uni d) -> Info D c ⊎ Info D d
+Info c d
  = let (n , c' , d') = padEqual (containerise c) (containerise d)
    in ⊎-map (containeriseInfo c) (containeriseInfo d) ∘
        (padEqualInfo (containerise c) (containerise d)) ∘
                                        (choiceInfo c' d') ∘
                                (forgetInfo (choice c' d'))

_UF+_ : ∀ {a}{E} -> UF {a} D E -> UF D E -> UF D E
(c , α) UF+ (d , β) = (c +Uni d) , [ α , β ] ∘ +Info c d

-- but the fusion [ α , β ] ∘ ⊎-map f g == [ α ∘ f , β ∘ g]
-- is not definitional, so the definition below in one go is slightly
-- preferable. The definitions are extensionally equal of course.
-}

_UF-+_ : ∀ {a}{E} -> UF {a} D E -> UF D E -> UF D E
(c , α) UF-+ (d , β)
  = let cc = containerise c
        dd = containerise d
        (n , c' , d') = padEqual cc dd
    in (forget (choice c' d') ,
        [ α ∘ (containeriseInfo c) ∘ (padEqualInfo₁ cc dd) ,
          β ∘ (containeriseInfo d) ∘ (padEqualInfo₂ cc dd) ]
                      ∘ (choiceInfo c' d') ∘ (forgetInfo (choice c' d')))

infixl  _UF-+_