\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)