\import Equiv
\import Equiv.BiEquiv
\import Equiv.Path
\import Equiv.Sigma
\import Function
\import Function.Meta
\import Logic.Unique
\import Homotopy.Square
\import Logic
\import Paths
\import Paths.Meta

\class Universe.{u}
  | isLocal : \Type u -> \Prop

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

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

  \func lift {Z : Local {U}} (f : S -> Z) (x' : S') : Z => IsEquiv.ret (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 => IsEquiv.f_ret (local-univ Z) 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 => IsEquiv.ret (pmapIsEquiv (local-univ Z)) (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
    => pmap (\lam q => path (q __ x)) $ IsEquiv.f_ret $ pmapIsEquiv (local-univ Z)
} \where {
  \sfunc levelProp.{u} {U : Universe.{u}} (X : \Type u) (l1 l2 : Localization X) : l1 = l2
    => \let p : l1.S' = {\Type u} 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.{u} \extends Universe.{u}
  | localization (A : \Type u) : Localization A

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

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

\func lmap.{u} {U : ReflUniverse.{u}} {A B : \Type u} (f : A -> B) (la : LType A) : LType B
  => lift (inL `o` f) la
  \where {
    \func id-prop.{u} {U : ReflUniverse.{u}} {A : \Type u} (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} {U : ReflUniverse.{u}} (X : \Type u) (s : Section (inL {localization X})) : IsEquiv (localization X).inL
  => inP $ lem (localization X) s
  \where
    \func lem.{u} {U : Universe.{u}} (L : Localization.{u}) (s : Section {L.S} {L.S'} inL) : QEquiv {L.S} {L.S'} inL \cowith
      | 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} {U : ReflUniverse.{u}} (X : Local.{u}) : IsEquiv (localization X).inL
  => inP $ lem X (localization X)
  \where
    \func lem.{u} {U : Universe.{u}} (X : Local.{u}) (L : Localization X) : QEquiv {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} {U : ReflUniverse.{u}} (X : \Type u) (s : Section (localization X).inL) : Local X
  => lem (localization X) s
  \where
    \lemma lem.{u} {U : Universe.{u}} (L : Localization.{u}) (s : Section {L.S} {L.S'} inL) : Local L.S \cowith
      | local => transport isLocal (inv (QEquiv_= (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} {U : ReflUniverse.{u}} {X : \Type u} (Y : Local.{u}) (f : X -> Y) (g : Y -> LType X) (p : \Pi (x : X) -> g (f x) = (localization X).inL x) : Retraction g
  => lem (localization X) Y f g p
  \where
    \func lem.{u} {U : Universe.{u}} (L : Localization.{u}) (Y : Local.{u}) (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} {U : ReflUniverse.{u}} {X : \Type u} (Y : Local.{u}) (f : X -> Y) (g : Section {Y} {LType X}) (p : \Pi (x : X) -> g (f x) = (localization X).inL x) : IsEquiv g
  => lem (localization X) Y f g p
  \where
    \func lem.{u} {U : Universe.{u}} (L : Localization.{u}) (Y : Local.{u}) (f : L.S -> Y) (g : Section {Y} {L.S'}) (p : \Pi (x : L.S) -> g (f x) = inL x) : IsEquiv g
      => inP $ BiEquiv.toQEquiv \new BiEquiv {
        | 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} {U : ReflUniverse.{u}} {X : \Type u} (Y : Local.{u}) (f : X -> Y) (g : Y >-> LType X) (p : \Pi (x : X) -> g (f x) = (localization X).inL x) : IsEquiv g
  => lem (localization X) Y f g p
  \where
    \lemma lem.{u} {U : ReflUniverse.{u}} (L : Localization.{u}) (Y : Local.{u}) (f : L.S -> Y) (g : Y >-> L.S') (p : \Pi (x : L.S) -> g (f x) = inL x) : IsEquiv g
      => IsEquiv.fromEmbSurj g \lam 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} {U : ReflUniverse.{u}} {X : \Type u} (C : Contr X) : Local X
  => lem C (localization X)
  \where
    \lemma lem.{u} {U : Universe.{u}} {X : \Type u} (C : Contr X) (L : Localization X) : Local X
      => 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} {U : ReflUniverse.{u}} : Local (\Sigma) => contrLocal unit-isContr
  \where
    \lemma lem.{u} {U : Universe.{u}} (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} {U : ReflUniverse.{u}} (P : Pullback.{u}) (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} {U : Universe.{u}} (P : Pullback.{u}) (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 => IsEquiv.ret P.pullback-univ LU-square
      \in localizationWithRetraction.lem L \new Section {
        | ret => g
        | ret_f u => path \lam i => ret {pmapSection (IsEquiv.toQEquiv P.pullback-univ) {g `o` inL} {id}} (
            pull P (g `o` inL)  ==< idp >==
            pull (pull P g) inL ==< path (\lam i => pull (IsEquiv.f_ret P.pullback-univ 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} {U : ReflUniverse.{u}} (A B : Local.{u}) : Local (\Sigma A B)
  => pullbackLocal (productPullback A B) A B unitLocal

-- | Path types of local types with a localization are local.
\lemma pathLocal.{u} {U : ReflUniverse.{u}} {A : Local.{u}} (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} {U : ReflUniverse.{u}} {A : \Type u} (B : A -> Local.{u}) : Local (\Pi (x : A) -> B x)
  => lem B (localization (\Pi (x : A) -> B x))
  \where
    \lemma lem.{u} {U : Universe.{u}} {A : \Type u} (B : A -> Local.{u}) (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} {U : ReflUniverse.{u}} (A B : Local.{u}) : 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 =>
    transportInv isLocal (inv (QEquiv_= BiEquiv~Equiv) *> QEquiv_= \new QEquiv {BiEquiv {A} {B}} {\Sigma (t : T) (\Sigma) (F t = (id,id))} {
      | f e => ((e,(e.ret,e.sec)), (), path (\lam i => (e.ret_f __ @ i, e.f_sec __ @ i)))
      | ret p => \new BiEquiv {
        | 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
    }) (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)))).local
  }

{- | 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} {U : ReflUniverse.{u}} {A : \Type u} (B : LType A -> \Type u) (totalIsLocal : Local (\Sigma (x : LType A) (B x)))
  : IsEquiv {\Pi (x : LType A) -> B x} {\Pi (a : A) -> B (lEta a)} (\lam f a => f (lEta a))
  => inP (piSigmaIdEquiv B) IsEquiv.`trans` lem1 B IsEquiv.`trans` lem2 B totalIsLocal IsEquiv.`trans` inP (symQEquiv (piSigmaEquiv lEta B))
  \where {
    \lemma lem1.{u} {U : ReflUniverse.{u}} {A : \Type u} (B : LType A -> \Type u)
      : IsEquiv {\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))
      => inP $ sigma-right {LType A -> \Sigma (x : LType A) (B x)}
                           {\lam g => (\lam x => (g x).1) = id}
                           \lam g => IsEquiv.toQEquiv $ pmapIsEquiv $ (localization A).local-univ (LType A)

    \lemma lem2.{u} {U : ReflUniverse.{u}} {A : \Type u} (B : LType A -> \Type u) (totalIsLocal : Local (\Sigma (x : LType A) (B x)))
      : IsEquiv {\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))
      => inP $ sigma-left $ IsEquiv.toQEquiv $ (localization A).local-univ totalIsLocal
  }

\lemma localYoneda.{u} {U : ReflUniverse.{u}} {A : \Type u} (a : A) (P : A -> Local)
  : IsEquiv {\Pi (a' : A) -> LType (a = a') -> P a'} {P a} (\lam f => f a (inL idp))
  => \let e a' => (localization (a = a')).local-univ (P a')
     \in inP $ transQEquiv (\new QEquiv {
       | A => \Pi (a' : A) -> LType (a = a') -> P a'
       | B => \Pi (a' : A) -> a = a' -> P a'
       | f h a' x => (h a') (inL x)
       | ret h a' => IsEquiv.ret (e a') (h a')
       | ret_f h => path \lam i a' => IsEquiv.ret_f (e a') {h a'} i
       | f_sec h => path \lam i a' => IsEquiv.f_ret (e a') {h a'} i
     }) (pi-contr-left a (\lam a' _ => P a'))

\func localization-equiv.{u}
  {U : ReflUniverse.{u}} {A B : \Type u}
  (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)