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

module Polynomial.IR where

mutual
   data Poly (D : Set₁) : Set₁ where
     idPN : Poly D
     con : (A : Set) -> Poly D
     sigma  : (S : Poly D) → (Info S → Poly D) → Poly D
     pi  : (A : Set) → (A → Poly D) → Poly D

   Info : {D : Set₁} -> Poly D -> Set₁
   Info {D} idPN = D
   Info (con A) = Lift A
   Info (sigma S T) = Σ[ x ∈ (Info S) ] Info (T x)
   Info (pi A T) = (x : A) -> Info (T x)

IR : (D E : Set1) -> Set1
IR D E = Σ[ c ∈ Poly D ] (Info c -> E)

module _ {D : Set1} where

  mutual

    Node : (T : Poly D) -> Fam D -> Set
    Node idPN F = ind F
    Node (con A) F = A
    Node (sigma S T) F = Σ[ s ∈ Node S F ] Node (T (info S F s)) F
    Node (pi A T) F = (x : A) -> Node (T x) F

    info : (T : Poly D) -> (F : Fam D) -> Node T F -> Info T
    info idPN F x = fib F x
    info (con A) F a = lift a
    info (sigma S T) F (s , x) = info S F s , info (T (info S F s)) F x
    info (pi A T) F g = λ a → info (T a) F (g a)

⟦_⟧ : {D E : Set1} -> IR D E -> Fam D -> Fam E
⟦ c , α ⟧ Z = Fam-map α (Node c Z , info c Z)

⟦_⟧₀ : {D E : Set1} -> IR D E -> Fam D -> Set
⟦ R ⟧₀ Z = ind (⟦ R ⟧ Z)

⟦_⟧₁ : {D E : Set1} -> (R : IR D E) -> (Z : Fam D) -> ⟦ R ⟧₀ Z -> E
⟦ R ⟧₁ Z = fib (⟦ R ⟧ Z)

-- IR D is functorial
IR-map : {D E E' : Set₁} -> (E -> E') -> IR D E  -> IR D E'
IR-map f (c , α) = (c , f ∘ α)

-- IR _ E is opfunctorial
mutual
  Poly-opmap : {D D' : Set₁} -> (D -> D') -> Poly D'  -> Poly D
  Poly-opmap f idPN = idPN
  Poly-opmap f (con A) = con A
  Poly-opmap f (sigma S T)
    = sigma (Poly-opmap f S)
            (λ s → Poly-opmap f (T (Info-opmap f S s)))
  Poly-opmap f (pi A T) = pi A (λ a → Poly-opmap f (T a))

  Info-opmap : {D D' : Set₁} -> (f : D -> D') -> (c : Poly D') ->
               Info (Poly-opmap f c) -> Info c
  Info-opmap f idPN = f
  Info-opmap f (con A) = id
  Info-opmap f (sigma S T)
   = map (Info-opmap f S) (λ {s} → Info-opmap f (T (Info-opmap f S s)))
  Info-opmap f (pi A T) = λ g a → Info-opmap f (T a) (g a)

IR-opmap : {D D' E : Set₁} -> (D -> D') -> IR D' E  -> IR D E
IR-opmap f (c , α) = Poly-opmap f c , α ∘ Info-opmap f c

-- Pi at the level of IR

πIR : {D : Set1 }(S : Set){E : S → Set1} ->
      ((s : S) -> (IR D (E s))) -> IR D ((s : S) -> E s)
πIR S f = (pi S (proj₁ ∘ f) , λ g s → proj₂ (f s) (g s))