module Prelude.Fam where

-- Families of Ds

open import Prelude.Basic
open import Prelude.Equality

{----- Basic definitions -------------------------------}

record Fam (D : Set₁) : Set₁ where
  constructor _,_
  field
    ind : Set       -- index set
    fib : ind -> D  -- fibres in D
open Fam public

Fam-map : {D D' : Set₁} -> (D -> D') -> Fam D -> Fam D'
Fam-map g (A , P) = (A , g ∘ P)

{-------------------------------------------------------}

{----- (Cartesian) morphisms in Fam --------------------}

record _Fam⇒_ {D : Set₁} (X Y : Fam D) : Set₁ where
  constructor _,_
  field
    indmor  : ind X -> ind Y
    indmor= : (x : ind X) -> fib X x ≡ fib Y (indmor x)
open _Fam⇒_

-- The identity morphism
FamId : {D : Set₁} -> (X : Fam D) -> X Fam⇒ X
FamId X = id , (λ x → refl)

-- Composition
_Fam∘_ : {D : Set₁} -> {X Y Z : Fam D} ->
         (f : Y Fam⇒ Z) -> (g : X Fam⇒ Y) -> X Fam⇒ Z
(f , f=) Fam∘ (g , g=) = f ∘ g , (λ x → trans (g= x) (f= (g x)))

{-------------------------------------------------------}

{----- Equality of families and Fam morphisms ----------}

-- Since Fam is a sigma type in disguise, this is how one proves
-- families equal
Fam-≡ : {D : Set₁} -> {X Y : Fam D} ->
        (ind= : ind X ≡ ind Y) ->
        (fib= : (x : ind X) -> fib X x ≡ fib Y (subst id ind= x)) ->
        X ≡ Y
Fam-≡ refl p = Fam-≡' refl (ext p)
  where Fam-≡' : {D : Set₁} -> {X Y : Fam D} ->
                 (ind= : ind X ≡ ind Y) ->
                 (dec= : subst (λ z → z -> D) ind= (fib X) ≡ fib Y) ->
                 X ≡ Y
        Fam-≡' refl refl = refl

-- Similarly, since Fam⇒ is a sigma type in disguise, this is how one
-- proves morphisms to be equal
⇒-≡ : {D : Set₁} -> {X Y : Fam D} ->
        {f g : X Fam⇒ Y} ->
        (ind= : indmor f ≡ indmor g) ->
        (fib= : subst (λ z → (x : ind X) → fib X x ≡ fib Y (z x))
                      ind= (indmor= f) ≡ indmor= g) ->
        f ≡ g
⇒-≡ refl refl = refl

-- How does one prove that a composition is an identity?
comp-is-id : {D : Set₁} -> {X Y : Fam D} ->
             (f : X Fam⇒ Y) -> (g : Y Fam⇒ X) ->
             (ind= : id ≡ (indmor f) ∘ (indmor g)) ->
             (coh : (y : ind Y) -> trans (indmor= g y) (indmor= f (indmor g y))
                                  ≡ cong (fib Y) (happly ind= y)) ->
             (f Fam∘ g) ≡ FamId Y
comp-is-id {X = X} {Y} f g ind= coh = sym (⇒-≡ ind= fib=)
  where fib= = begin
                 subst (λ z → (x : ind Y) → fib Y x ≡ fib Y (z x))
                       ind=
                       (λ x → refl)
               ≡⟨ subst-cod ind= ⟩
                 (λ x → subst (λ z → fib Y x ≡ fib Y (z x))
                               ind=
                               refl)
               ≡⟨ ext (λ x → subst-cong {B = λ z → fib Y x ≡ z} ind= {refl}) ⟩
                 (λ x → subst (λ z → fib Y x ≡ z)
                               (cong (λ z → fib Y (z x)) ind=)
                               refl)
               ≡⟨ ext (λ x → subst-path (cong (λ z → fib Y (z x)) ind=) {refl}) ⟩
                 (λ x → cong (λ z → fib Y (z x)) ind=)
               ≡⟨ ext (λ x → cong-cong ind=) ⟩
                 (λ x → cong (fib Y) (cong (λ z → z x) ind=))
               ≡⟨ refl ⟩
                 (λ x → cong (fib Y) (happly ind= x))
               ≡⟨ sym (ext coh) ⟩
                 (λ x → trans (indmor= g x) (indmor= f (indmor g x)))
               ∎ where open ≡-Reasoning

comp-is-id-ext : {D : Set₁} -> {X Y : Fam D} ->
                 (f : X Fam⇒ Y) -> (g : Y Fam⇒ X) ->
                 (ind= : (x : _) -> x ≡ (indmor f) (indmor g x)) ->
                 (coh : (y : ind Y) -> trans (indmor= g y) (indmor= f (indmor g y))
                                          ≡ cong (fib Y) (ind= y)) ->
                 (f Fam∘ g) ≡ FamId Y
comp-is-id-ext {X = X} {Y} f g ind= coh
 = comp-is-id f g (ext ind=) (λ y → trans (coh y) (cong (cong (fib Y))
                                                         (sym (ext-β ind= y))))

{-------------------------------------------------------}

{----- Fam is a monad ----------------------------------}

-- Unit

Fam-η : {I : Set₁} -> I -> Fam I
Fam-η i = (⊤ , const i)

-- Multiplication

Fam-μ : {D : Set₁} -> Fam (Fam D) → Fam D
Fam-μ {D} (A , B) = (Σ A (ind ∘ B) , g)
                    where g : Σ A (ind ∘ B) -> D
                          g (a , k) = fib (B a) k

-- Bind

_Fam->>=_ : {X Y : Set₁} -> Fam X → (X → Fam Y) -> Fam Y
P Fam->>= f = Fam-μ (Fam-map f P)

-- Dependent bind

_Fam->>==_ : {D : Set₁} -> {E : D -> Set₁} -> Fam D ->
             ((d : D) -> Fam (E d)) -> Fam (Σ D E)
P Fam->>== f = P Fam->>= (λ x → Fam-map (λ p → (x , p)) (f x))

{-------------------------------------------------------}

{----- Constructions on Fam ----------------------------}

-- Products (in the comma category Σ[ D ∈ Set1 ] (Fam D))

πFam : (S : Set){D : S -> Set1} -> (∏ S (Fam ∘ D)) -> Fam (∏ S D)
πFam S T = ((s : S) -> ind (T s)) , \ f s -> fib (T s) (f s)

-- Powers as degenerate products

_Fam⟶_ : (S : Set){D : Set1} -> Fam D -> Fam (S -> D)
S Fam⟶ T = πFam S (λ _ → T)

-- Coproduct of families

⋃ : {D : Set1} -> (A : Set) -> (A -> Fam D) -> Fam D
⋃ A f = (Σ[ a ∈ A ] ind (f a)) , (λ { (a , x) → fib (f a) x })

-- Binary coproduct
_Fam+_ : {D : Set1} -> Fam D -> Fam D -> Fam D
(A , P) Fam+ (B , Q) = (A ⊎ B) , [ P , Q ]′

{-------------------------------------------------------}