module Uniform.Examples where


open import Prelude.Basic
open import Prelude.Equality
open import Prelude.Equivalences
open import Prelude.Fam

open import Data.Fin hiding (_+_; lift)

open import Uniform.IR
open import Uniform.Initial
open import Uniform.Coproduct
open import Uniform.Coproduct.Maps
open import Uniform.Composition using (_○_)

-- W-types

module _ (S : Set)(P : S -> Set) where

  cW₀ : {D : Set1} -> Uni D
  cW₀ = δ (σ ι (λ _ → S)) (λ { (_ , s) → P s })

  -- ordinary W-type: choose D = ⊤

  cW⊤ : UF {lsuc lzero} (Lift ⊤) (Lift ⊤)
  cW⊤ = (cW₀ , λ _ → _)

  W : Set
  W = U {c = cW⊤}

  sup : (s : S) -> (P s -> W) -> W
  sup s f = intro ((tt , s) , f)

  -- "upgraded" W-type where decoding component applies T everywhere

  cW : UF {lsuc lzero} Set Set
  cW = cW₀ , (λ { ((_ , s) , Tf) → (x : P s) -> Tf x })

-- A universe closed under sigma- and pi-, and W-types

cUΣΠW : UF Set Set
cUΣΠW = (cUB UF-+ cUΣ) UF-+ (cUΠ UF-+ cUW)
  module UU where
   cUB = (ι , λ _ → Bool)
   cUΣ = (δ (δ ι (λ _ → ⊤)) (λ { (_ , A) → A _ }) ,
          (λ { ((_ , A) , B) → Σ (A _) B }))
   cUΠ = (δ (δ ι (λ _ → ⊤)) (λ { (_ , A) → A _ }) ,
          (λ { ((_ , A) , B) → (x : A _) -> B x }))
   cUW = (δ (δ ι (λ _ → ⊤)) (λ { (_ , A) → A _ }) ,
          (λ { ((_ , A) , B) → W (A _) B }))

UΣΠW : Set
UΣΠW = U {c = cUΣΠW}

TΣΠW : UΣΠW -> Set
TΣΠW x = T x


-- one part of the ismorphism ⟦ c + d ⟧ ≃ ⟦ c ⟧ + ⟦ d ⟧
iso :   (⟦ UU.cUB ⟧₀ (UΣΠW , TΣΠW) ⊎ ⟦ UU.cUΣ ⟧₀ (UΣΠW , TΣΠW))
      ⊎ (⟦ UU.cUΠ ⟧₀ (UΣΠW , TΣΠW) ⊎ ⟦ UU.cUW ⟧₀ (UΣΠW , TΣΠW))
                                     -> ⟦ cUΣΠW ⟧₀ (UΣΠW , TΣΠW)
iso =   right (UU.cUB UF-+ UU.cUΣ) (UU.cUΠ UF-+ UU.cUW)
      ∘ ⊎-map (right UU.cUB UU.cUΣ) (right UU.cUΠ UU.cUW)
  where right : (R Q : UF Set Set) ->
                ⟦ R ⟧₀ (UΣΠW , TΣΠW) ⊎ ⟦ Q ⟧₀ (UΣΠW , TΣΠW) ->
                ⟦ R UF-+ Q ⟧₀ (UΣΠW , TΣΠW)
        right R Q = (_Fam⇒_.indmor (_≃_.right (coprodEq R Q {Z = (UΣΠW , TΣΠW)})))


bool : UΣΠW
bool = intro (iso (inj₁ (inj₁ tt)))

sig : (a : UΣΠW) -> (b : TΣΠW a -> UΣΠW) -> UΣΠW
sig a b = intro (iso (inj₁ (inj₂ ((tt , (λ _ → a)) , b))))

pi : (a : UΣΠW) -> (b : TΣΠW a -> UΣΠW) -> UΣΠW
pi a b = intro (iso (inj₂ (inj₁ ((tt , (λ _ → a)) , b))))

w : (a : UΣΠW) -> (b : TΣΠW a -> UΣΠW) -> UΣΠW
w a b = intro (iso (inj₂ (inj₂ ((tt , (λ _ → a)) , b))))

-- A language of sums and products

sumFin : (n : ℕ) -> (Fin n -> ℕ) -> ℕ
sumFin zero f = 
sumFin (suc n) f = f (fromℕ n) + sumFin n (f ∘ inject₁)

prodFin : (n : ℕ) -> (Fin n -> ℕ) -> ℕ
prodFin zero f = 
prodFin (suc n) f = f (fromℕ n) * prodFin n (f ∘ inject₁)

cArith : UF {lsuc lzero} (Lift ℕ) (Lift ℕ)
cArith = cFin UF-+ cSum UF-+ cProd
  module Ar where
   cFin = (σ ι (λ _ → ℕ)) , (λ { (_ , n) → lift n })
   cSum = δ (δ ι (λ _ → ⊤)) (λ { (_ , n) → Fin (lower (n _)) }) ,
          (λ { ((_ , n) , f) → lift (sumFin (lower (n _)) (lower ∘ f)) })
   cProd = δ (δ ι (λ _ → ⊤)) (λ { (_ , n) → Fin (lower (n _)) }) ,
          (λ { ((_ , n) , f) → lift (prodFin (lower (n _)) (lower ∘ f)) })

Arith : Set
Arith = U {c = cArith}

eval : Arith -> ℕ
eval = lower ∘ T

isoA :    (⟦ Ar.cFin ⟧₀ (Arith , T)
       ⊎  ⟦ Ar.cSum ⟧₀ (Arith , T))
       ⊎  ⟦ Ar.cProd ⟧₀ (Arith , T) -> ⟦ cArith ⟧₀ (Arith , T)
isoA = right (Ar.cFin UF-+ Ar.cSum) Ar.cProd ∘ ⊎-map (right Ar.cFin Ar.cSum) id
  where right : (R Q : UF (Lift ℕ) (Lift ℕ)) ->
                ⟦ R ⟧₀ (Arith , T) ⊎ ⟦ Q ⟧₀ (Arith , T) ->
                ⟦ R UF-+ Q ⟧₀ (Arith , T)
        right R Q = (_Fam⇒_.indmor (_≃_.right (coprodEq R Q {Z = (Arith , T)})))


fin : ℕ -> Arith
fin n = intro (isoA (inj₁ (inj₁ (tt , n))))

sum : (n : Arith) -> (Fin (eval n) -> Arith) -> Arith
sum n f = intro (isoA (inj₁ (inj₂ ((tt , (λ _ → n)) , f) )))

prod : (n : Arith) -> (Fin (eval n) -> Arith) -> Arith
prod n f = intro (isoA (inj₂ ((tt , (λ _ → n)) , f) ))

fac : ℕ -> Arith
fac n = prod (fin n) (λ x → fin (suc (toℕ x)))

_ : eval (fac ) ≡ 
_ = refl

-- List of arguments for free: composing with W ℕ Fin

cUList : UF Set Set
cUList = (UU.cUB UF-+ UU.cUΣ) ○ (cW ℕ Fin)

-- the isomorphism (a + b) ∘ c ≃ (a ∘ c) + (b ∘ c) makes it a bit
-- awkward to find the exact code cUList reduces to. Instead, we find
-- the code it morally reduces to using this isomorphism, presented as
-- a function of decodings: note that this is clearly invertible as it
-- is just shuffling trivial `lift tt`s and `⊥-elim`s around.

compCode : let
       cUΣ-composed
         = (δ (σ (cW₀ ℕ Fin)
                         (λ { ((_ , s) , Tf) → ((x : Fin s) → Tf x) → ℕ}))
                  ((λ { (((_ , s) , Tf) , e) →
                                 Σ[ y ∈ ((x : Fin s) → Tf x) ] (Fin (e y))}))) ,
           (λ { ((((_ , s) , Tf) , e) , B) →
                Σ[ y ∈ ((x : Fin s) -> Tf x) ] ((w : Fin (e y)) -> B (y , w)) })
    in {Z : Fam Set} ->
       ⟦ cUList ⟧₀ Z ->
       ⟦ UU.cUB UF-+ cUΣ-composed ⟧₀ Z
compCode (((((((((_ , true) , _) , _) , _) , _) , _) , _) , _) , _)
   = ((((((((((tt , true) , ⊥-elim) , lift tt) , ⊥-elim) , lift tt) , ⊥-elim) , lift tt) , ⊥-elim) , lift tt) , ⊥-elim)
compCode (((((((((_ , false) , _) , _) , _) , n) , f) , _) , e) , y)
   = ((((((((((tt , false) , ⊥-elim) , n _) , ⊥-elim) , tt) , (λ x → f (_ , x))) , e) , ⊥-elim) , tt) , y)