\import Data.Sigma
\import Equiv
\import Equiv.Path
\import Equiv.Sigma
\import Equiv.Univalence
\import Function
\import HLevel
\import Homotopy.Square
\import Logic
\import Paths
\import Paths.Meta

\class Universe
  | isLocal : \hType -> \Prop

-- | The type of local types
\class Local {U : Universe} (S : \hType)
  | local : isLocal S

-- | The type of localizations of a type
\class Localization {U : Universe} (S : \hType) (S' : Local) {
  | inL : S -> S'
  | local-univ (Z : Local) : Equiv {S' -> Z} {S -> Z} (-o inL)

  \func lift {Z : Local {U}} (f : S -> Z) (x' : S') : Z => sec {local-univ Z} f x'
    \where
      -- | The lift of a constant map is constant
      \func const {Z : Local {U}} (z : Z) (x' : S') : lift (\lam _ => z) x' = z
        => remove_inL (lift (\lam _ => z)) (\lam _ => z) (lift-prop (\lam _ => z)) x'

  \func lift-prop {Z : Local {U}} (f : S -> Z) (x : S) : lift f (inL x) = f x => path (\lam i => (f_sec {local-univ Z} f @ i) x)

  \func remove_inL {Z : Local {U}} (f g : S' -> Z) (p : \Pi (x : S) -> f (inL x) = g (inL x)) (x : S') : f x = g x
    => path (\lam i => (sec {pmapEquiv (local-univ Z) {f} {g}} (path (\lam i x => p x @ i)) @ i) x)

  \func remove_inL-coh {Z : Local {U}} (f g : S' -> Z) (p : \Pi (x : S) -> f (inL x) = g (inL x)) (x : S) : remove_inL f g p (inL x) = p x
    => \have t => f_sec {pmapEquiv (local-univ Z) {f} {g}} (path (\lam i x => p x @ i))
       \in pmap (\lam q => path ((q @ __) x)) t
} \where {
  \use \level levelProp {U : Universe} (X : \hType) (l1 l2 : Localization X) : l1 = l2
    => \let p : l1.S' = {\hType} l2.S' => path (iso (l1.lift l2.inL)
                                                    (l2.lift l1.inL)
                                                    (remove_inL (l2.lift l1.inL `o` l1.lift l2.inL) id (\lam x =>
                                                      l2.lift l1.inL (l1.lift l2.inL (l1.inL x)) ==< pmap (l2.lift l1.inL) (l1.lift-prop l2.inL x) >==
                                                      l2.lift l1.inL (l2.inL x)                  ==< l2.lift-prop l1.inL x >==
                                                      l1.inL x                                   `qed))
                                                    (remove_inL (l1.lift l2.inL `o` l2.lift l1.inL) id (\lam x =>
                                                      l1.lift l2.inL (l2.lift l1.inL (l2.inL x)) ==< pmap (l1.lift l2.inL) (l2.lift-prop l1.inL x) >==
                                                      l1.lift l2.inL (l1.inL x)                  ==< l1.lift-prop l2.inL x >==
                                                      l2.inL x                                   `qed)))
       \in exts Localization {
         | S' => p
         | inL => l1.lift-prop l2.inL
       }
}

\open Localization (lift,lift-prop,remove_inL,remove_inL-coh)

\class ReflUniverse \extends Universe
  | localization (A : \hType) : Localization A

-- | The localization functor
\func LType {U : ReflUniverse} (A : \hType) => Localization.S' {localization A}

\func lEta {U : ReflUniverse} {A : \hType} (a : A) => Localization.inL {localization A} a

\func lmap {U : ReflUniverse} {A B : \hType} (f : A -> B) (la : LType A) : LType B
  => lift (inL `o` f) la
  \where {
    \func id-prop {U : ReflUniverse} {A : \hType} (la : LType A) : lmap id la = la
      => remove_inL (lmap id) id (lift-prop inL) la
  }

-- | If ``inL : X -> X'`` is a section, then it is an equivalence.
\lemma sectionIsEquiv {U : ReflUniverse} (X : \hType) (s : Section (inL {localization X})) : Equiv (inL {localization X})
  => lem (localization X) s
  \where
    \lemma lem {U : Universe} (L : Localization) (s : Section {L.S} {L.S'} inL) : Equiv {L.S} {L.S'} inL
      => \new QEquiv {
        | ret => s.ret
        | ret_f => s.ret_f
        | f_sec => remove_inL (inL `o` s.ret) id (\lam x => pmap inL (s.ret_f x))
      }

-- | If {X} is local, then ``inL : X -> X'`` is an equivalence.
\lemma localizationOfLocalType {U : ReflUniverse} (X : Local) : Equiv (inL {localization X})
  => lem X (localization X)
  \where
    \lemma lem {U : Universe} (X : Local) (L : Localization X) : Equiv {X} {L.S'} inL
      => sectionIsEquiv.lem L (\new Section {
           | ret => lift id
           | ret_f => lift-prop id
         })

-- | If ``inL : X -> X'`` is a section, then {X} is local.
\lemma localizationWithRetraction {U : ReflUniverse} (X : \hType) (s : Section (inL {localization X})) : Local X
  => lem (localization X) s
  \where
    \lemma lem {U : Universe} (L : Localization) (s : Section {L.S} {L.S'} inL) : Local L.S \cowith
      | local => transport isLocal (inv (Equiv-to-= (sectionIsEquiv.lem L s))) L.S'.local

-- | If ``inL : X -> X'`` factors through a local type {Y}, then the map ``Y -> X'`` is a retraction.
\func localizationFactor {U : ReflUniverse} {X : \hType} (Y : Local) (f : X -> Y) (g : Y -> LType X) (p : \Pi (x : X) -> g (f x) = inL {localization X} x) : Retraction g
  => lem (localization X) Y f g p
  \where
    \func lem {U : Universe} (L : Localization) (Y : Local) (f : L.S -> Y) (g : Y -> L.S') (p : \Pi (x : L.S) -> g (f x) = inL x) : Retraction g \cowith
      | sec => lift f
      | f_sec => remove_inL (g `o` lift f) id (\lam x => pmap g (lift-prop f x) *> p x)

-- | If ``inL : X -> X'`` factors through a section ``g : Y -> X'`` with {Y} local, then {g} is an equivalence.
\lemma localizationFactorSection {U : ReflUniverse} {X : \hType} (Y : Local) (f : X -> Y) (g : Section {Y} {LType X}) (p : \Pi (x : X) -> g (f x) = inL {localization X} x) : Equiv g
  => lem (localization X) Y f g p
  \where
    \func lem {U : Universe} (L : Localization) (Y : Local) (f : L.S -> Y) (g : Section {Y} {L.S'}) (p : \Pi (x : L.S) -> g (f x) = inL x) : Equiv g \cowith
      | Section => g
      | Retraction => localizationFactor.lem L Y f g p

-- | If ``inL : X -> X'`` factors through an embedding ``g : Y -> X'`` with {Y} local, then {g} is an equivalence.
\lemma localizationFactorEmbedding {U : ReflUniverse} {X : \hType} (Y : Local) (f : X -> Y) (g : Y >-> LType X) (p : \Pi (x : X) -> g (f x) = inL {localization X} x) : Equiv g
  => lem (localization X) Y f g p
  \where
    \lemma lem {U : ReflUniverse} (L : Localization) (Y : Local) (f : L.S -> Y) (g : Y >-> L.S') (p : \Pi (x : L.S) -> g (f x) = inL x) : Equiv g
      => \new ESEquiv {
        | Embedding => g
        | isSurjMap x =>\have r : Retraction g => localizationFactor.lem L Y f g p
                      \in inP (r.sec x, r.f_sec x)
      }

-- | Contractible types with a localization are local.
\lemma contrLocal {U : ReflUniverse} (C : Contr) : Local C
  => lem C (localization C)
  \where
    \lemma lem {U : Universe} (C : Contr) (L : Localization C) : Local C
      => localizationWithRetraction.lem L (\new Section { | ret _ => C.center | ret_f => C.contraction })

-- | The unit type is local if it has a localization.
\lemma unitLocal {U : ReflUniverse} : Local (\Sigma) => contrLocal unit-isContr
  \where
    \lemma lem {U : Universe} (L : Localization (\Sigma)) : Local (\Sigma) => contrLocal.lem unit-isContr L

{- | Pullbacks of local types with a localization are local.

 - Proof:
 - By {localizationWithRetraction}, we just need to construct a retraction of ``inL : U -> L U``.
 - By the universal property of the localization, we have maps from ``L U`` to the objects in the pullback square
 - such that their composition with ``inL`` equals corresponding maps in the pullback square.
 - So, we have a square with the corner ``L U``:
 - ```
 -        lux
 -    LU ----> X
 -     |       |
 - luv |       |
 -     |       |
 -     V ----> Y
 - ```
 - The [pullback]{Square.pull} of this square along ``inL`` is the original pullback square.
 - By the universal property of pullbacks, we have a map ``h : L U -> U`` such that ``h `o` inL`` is the identity map.
 -}
\lemma pullbackLocal {U : ReflUniverse} (P : Pullback \lp \oo) (VL : Local P.square.V) (XL : Local P.square.X) (YL : Local P.square.Y) : Local P.square.U
  => lem P (localization P.square.U) VL XL YL
  \where {
    \open Square (pull)

    \lemma lem {U : Universe} (P : Pullback \lp \oo) (L : Localization P.square.U) (VL : Local P.square.V) (XL : Local P.square.X) (YL : Local P.square.Y) : Local P.square.U =>
      \let | lux => lift P.square.ux
           | luv => lift P.square.uv
           | pux u => pmap P.square.xy (lift-prop P.square.ux u)
           | puv u => pmap P.square.vy (lift-prop P.square.uv u)
           | puy u =>
               P.square.vy (luv (inL u))   ==< puv u >==
               P.square.vy (P.square.uv u) ==< P.square.sqcomm u >==
               P.square.xy (P.square.ux u) ==< inv (pux u) >==
               P.square.xy (lux (inL u))   `qed
           | LU-square => \new P.square {
                            | U => L.S'
                            | ux => lux
                            | uv => luv
                            | sqcomm => remove_inL (P.square.vy `o` luv) (P.square.xy `o` lux) puy
                          }
           | g => P.pullback-univ.sec LU-square
      \in localizationWithRetraction.lem L (\new Section {
        | ret => g
        | ret_f u => path (\lam i => (ret {pmapSection P.pullback-univ {g `o` inL} {id}} (
            pull P (g `o` inL)  ==< idp >==
            pull (pull P g) inL ==< path (\lam i => pull (P.pullback-univ.f_sec LU-square @ i) inL) >==
            pull LU-square inL  ==< path (\lam i => \new P.square {
                                                      | ux u => lift-prop P.square.ux u @ i
                                                      | uv u => lift-prop P.square.uv u @ i
                                                      | sqcomm u => <*.concat (remove_inL-coh (P.square.vy `o` luv) (P.square.xy `o` lux) puy u)
                                                                              (pathOver (coe_path.alt (puv u) (P.square.sqcomm u) (pux u))) @ i
                                                    }) >==
            pull P id           `qed) @ i) u)
      })
  }

-- | Products of local types with a localization are local.
\lemma productLocal {U : ReflUniverse} (A B : Local) : Local (\Sigma A B)
  => pullbackLocal (productPullback A B) A B unitLocal

-- | Path types of local types with a localization are local.
\lemma pathLocal {U : ReflUniverse} {A : Local} (x y : A) : Local (x = y)
  => pullbackLocal (pathPullback x y) unitLocal unitLocal A

{- | Pi-types of local types with a localization are local.

 - Proof:
 - By {localizationWithRetraction}, we just need to construct a retraction of ``inL : (\Pi (x : A) -> B x) -> L (\Pi (x : A) -> B x)``.
 - Since ``B a`` is local, there is a map ``h a : L (\Pi (x : A) -> B x) -> B a`` such that ``h a `o` inL = ap a``, where ``ap a`` is ``__ a``.
 - Thus, we have a map ``r`` defined as ``\lam lx a => h a lx :  L (\Pi (x : A) -> B x) -> \Pi (x : A) -> B a`` such that ``r  `o` inL`` is the identity map.
 -}
\lemma piLocal {U : ReflUniverse} {A : \hType} (B : A -> Local) : Local (\Pi (x : A) -> B x)
  => lem B (localization (\Pi (x : A) -> B x))
  \where
    \lemma lem {U : Universe} {A : \hType} (B : A -> Local) (L : Localization (\Pi (x : A) -> B x)) : Local (\Pi (x : A) -> B x)
      => \let ap (a : A) (g : \Pi (x : A) -> B x) => g a
         \in localizationWithRetraction.lem L (\new Section {
           | ret lx a => lift (ap a) lx
           | ret_f g => path (\lam i a => lift-prop (ap a) g @ i)
         })

-- | The type of equivalences between local types is local
\lemma equivLocal {U : ReflUniverse} (A B : Local) : Local (Equiv {A} {B}) =>
  \let | T => \Sigma (A -> B) (\Sigma (B -> A) (B -> A))
       | F (t : T) => (t.2.1 `o` t.1, t.1 `o` t.2.2)
  \in \new Local { | local =>
    transport isLocal (inv (QEquiv-to-= (\new QEquiv {Equiv {A} {B}} {\Sigma (t : T) (\Sigma) (F t = (id,id))} {
      | f (e : Equiv) => ((e,(e.ret,e.sec)), (), path (\lam i => (e.ret_f __ @ i, e.f_sec __ @ i)))
      | ret p => \new Equiv {
        | f => p.1.1
        | ret => p.1.2.1
        | ret_f a => path ((p.3 @ __).1 a)
        | sec => p.1.2.2
        | f_sec b => path ((p.3 @ __).2 b)
      }
      | ret_f => idpe
      | f_sec => idpe
    }))) (local {pullbackLocal (sigmaPullback F (\lam (_ : \Sigma) => (id,id)))
                               (productLocal (piLocal (\lam _ => B)) (productLocal (piLocal (\lam _ => A)) (piLocal (\lam _ => A))))
                               unitLocal
                               (productLocal (piLocal (\lam _ => A)) (piLocal (\lam _ => B)))}) }

{- | An elimination principle for the localization in a reflective universe.

 - Proof:
 - The pi-types in the statement are equivalent to the types of functions ``LType A -> \Sigma (x : LType A) (B x)`` and ``A -> \Sigma (x : LType A) (B x)``
 - such that their composition with the first projection equal to {id} and {lEta}, respectively.
 - Thus, the required function is an equivalence iff the type of diagonal maps in the diagram below such that the bottom triangle commutes
 - is canonically equivalent to the type of the top maps such that the square commutes.
 - This follows from the universal property of the localization.
 - ```
 -      A --------> \Sigma (x : LType A) (B x)
 -      |             |
 - lEta |             | __.1
 -      |             |
 -    LType A ----> LType A
 -             id
 - ```
-}
\lemma universe-elim {U : ReflUniverse} {A : \hType} (B : LType A -> \hType) (totalIsLocal : Local (\Sigma (x : LType A) (B x)))
  : Equiv {\Pi (x : LType A) -> B x} {\Pi (a : A) -> B (lEta a)} (\lam f a => f (lEta a))
  => piSigmaIdEquiv B `transEquiv` lem1 B `transEquiv` lem2 B totalIsLocal `transEquiv` symQEquiv (piSigmaEquiv lEta B)
  \where {
    \lemma lem1 {U : ReflUniverse} {A : \hType} (B : LType A -> \hType)
      : Equiv {\Sigma (g : LType A -> \Sigma (x : LType A) (B x)) ((\lam x => (g x).1) = id)}
              {\Sigma (g : LType A -> \Sigma (x : LType A) (B x)) ((\lam a => (g (lEta a)).1) = lEta)}
              (\lam p => (p.1, pmap (-o lEta) p.2))
      => sigma-right {LType A -> \Sigma (x : LType A) (B x)}
                     {\lam g => (\lam x => (g x).1) = id}
                     (\lam g => pmapEquiv (local-univ {localization A} (LType A)))

    \lemma lem2 {U : ReflUniverse} {A : \hType} (B : LType A -> \hType) (totalIsLocal : Local (\Sigma (x : LType A) (B x)))
      : Equiv {\Sigma (g : LType A -> \Sigma (x : LType A) (B x)) ((\lam a => (g (lEta a)).1) = lEta)}
              {\Sigma (g : A -> \Sigma (x : LType A) (B x)) ((\lam a => (g a).1) = lEta)}
              (\lam p => (\lam a => p.1 (lEta a), p.2))
      => sigma-left (local-univ {localization A} totalIsLocal)
  }

\func localYoneda {U : ReflUniverse} {A : \hType} (a : A) (P : A -> Local)
  : Equiv {\Pi (a' : A) -> LType (a = a') -> P a'} {P a}
  => \let e a' => local-univ {localization (a = a')} (P a')
     \in transEquiv (\new Equiv {
       | A => \Pi (a' : A) -> LType (a = a') -> P a'
       | B => \Pi (a' : A) -> a = a' -> P a'
       | f h a' => e a' (h a')
       | ret h a' => ret {e a'} (h a')
       | ret_f h => path (\lam i a' => ret_f {e a'} (h a') @ i)
       | sec h a' => sec {e a'} (h a')
       | f_sec h => path (\lam i a' => f_sec {e a'} (h a') @ i)
     }) (pi-contr-left a (\lam a' _ => P a'))

\func localization-equiv {U : ReflUniverse} {A B : \hType}
                         (f : A -> LType B)
                         (g : B -> LType A)
                         (p : \Pi (a : A) -> lift g (f a) = lEta a)
                         (q : \Pi (b : B) -> lift f (g b) = lEta b)
                         : QEquiv {LType A} {LType B} \cowith
  | f => lift f
  | ret => lift g
  | ret_f => remove_inL (\lam x => lift g (lift f x)) (\lam x => x) (\lam a => pmap (lift g) (lift-prop f a) *> p a)
  | f_sec => remove_inL (\lam y => lift f (lift g y)) (\lam y => y) (\lam b => pmap (lift f) (lift-prop g b) *> q b)