\import Equiv.Fiber
\import Equiv.Path
\import Function
\import HLevel
\import Homotopy.Fibration
\import Logic
\import Logic.Meta
\import Paths
\import Paths.Meta

-- # Definition of equivalences

\protected \record Map {A B : \Type} (\coerce f : A -> B)

\record Section \extends Map
  | ret : B -> A
  | ret_f : \Pi (x : A) -> ret (f x) = x
  \where {
    \func =Fiber {A B : \Type} (f : A -> B) : Section f = Fib (-o f) id
      => path (iso (\lam (s : Section f) => Fib.make s.ret (path (\lam i x => s.ret_f x @ i)))
                   (\lam z => \new Section f z.1 (\lam x => path ((z.2 @ __) x)))
                   idpe
                   idpe)

    \lemma isContr (e : Equiv) : Contr (Section e)
      => rewrite (=Fiber e) (Equiv=>contrFibers (-o_Equiv e) id)

    \lemma levelProp (r : Retraction) : isProp (Section r)
      => \lam s => isContr=>isProp (isContr (\new Equiv { | Section => s | Retraction => r })) s
  }

\record Retraction \extends Map
  | sec : B -> A
  | f_sec : \Pi (y : B) -> f (sec y) = y
  \where {
    \func =Fiber {A B : \Type} (f : A -> B) : Retraction f = Fib (f `o-) id
      => path (iso (\lam (s : Retraction f) => Fib.make s.sec (path (\lam i x => s.f_sec x @ i)))
                   (\lam z => \new Retraction f z.1 (\lam x => path ((z.2 @ __) x)))
                   idpe
                   idpe)

    \lemma isContr (e : Equiv) : Contr (Retraction e)
      => rewrite (=Fiber e) (Equiv=>contrFibers (o-_Equiv e) id)

    \lemma levelProp (s : Section) : isProp (Retraction s)
      => \lam r => isContr=>isProp (isContr (\new Equiv { | Section => s | Retraction => r })) r
  }

\record Equiv \extends Section, Retraction {
  \func f_ret (y : B) : f (ret y) = y
    => pmap (\lam y => f (ret y)) (inv (f_sec y)) *> pmap f (ret_f (sec y)) *> f_sec y

  \func isInj {a a' : A} (p : f a = f a') : a = a'
    => inv (ret_f a) *> pmap ret p *> ret_f a'

  \lemma isSurj : Function.isSurj f
    => \lam y => inP (ret y, f_ret y)

  \func adjoint {a : A} {b : B} (p : f a = b) : a = ret b
    => inv (ret_f a) *> pmap ret p

  \func adjointInv {a : A} {b : B} (p : a = ret b) : f a = b
    => pmap f p *> f_ret b
}
  \where {
    \use \level levelProp {A B : \Type} {f : A -> B} (e e' : Equiv f) : e = e' =>
      path (\lam i => \new Equiv {
        | Section => isContr=>isProp {Section f} (Section.isContr e) e e' @ i
        | Retraction => isContr=>isProp {Retraction f} (Retraction.isContr e) e e' @ i
      })

    \func equals {A B : \Type} {e e' : Equiv {A} {B}} (p : e.f = e'.f) : e = e' =>
      path (\lam i => \new Equiv {
        | Section => pathOver (isContr=>isProp (Section.isContr e')
                                               (rewriteI p (\new Section e.f e.ret e.ret_f))
                                               (\new Section e'.f e'.ret e'.ret_f)) @ i
        | Retraction => pathOver (isContr=>isProp (Retraction.isContr e')
                                                  (rewriteI p (\new Retraction e.f e.sec e.f_sec))
                                                  (\new Retraction e'.f e'.sec e'.f_sec)) @ i
      })

    \lemma fromInjSurj {A B : \Set} (f : A -> B) (inj : Function.isInj f) (surj : Function.isSurj f) : Equiv {A} {B} f
      => \new ESEquiv {
        | Embedding => Embedding.fromInjection inj
        | isSurjMap => surj
      }
  }

\record QEquiv \extends Equiv { | sec => ret }
  \where {
    \use \coerce fromEquiv (e : Equiv) : QEquiv \cowith
      | Section => e
      | f_sec => e.f_ret

    \func fromEquiv' (e : Equiv) : QEquiv \cowith
      | Map => e
      | ret => e.sec
      | ret_f a => inv (e.ret_f (e.sec (e a))) *> pmap e.ret (e.f_sec (e a)) *> e.ret_f a
      | f_sec => e.f_sec
  }

-- # Examples of equialences

\func idEquiv {A : \Type} : QEquiv \cowith
  | A => A
  | B => A
  | f x => x
  | ret x => x
  | ret_f _ => idp
  | f_sec _ => idp

\func symQEquiv {A B : \Type} (e : QEquiv {A} {B}) : QEquiv {B} {A} \cowith
  | f => e.ret
  | ret => e.f
  | ret_f => e.f_sec
  | f_sec => e.ret_f

\func \fixr 3 transQEquiv {A B C : \Type} (e1 : QEquiv {A} {B}) (e2 : QEquiv {B} {C}) : QEquiv {A} {C} \cowith
  | f x => e2 (e1 x)
  | ret y => e1.ret (e2.ret y)
  | ret_f x => pmap e1.ret (e2.ret_f (e1 x)) *> e1.ret_f x
  | f_sec y => pmap e2 (e1.f_sec (e2.sec y)) *> e2.f_sec y

\module TwoOutOfThree \where {
  \func leftFactor {A B C : \Type} (e1 : Retraction {A} {B}) (g : B -> C) (e3 : QEquiv (g `o` e1)) : QEquiv g \cowith
    | ret c => e1 (e3.ret c)
    | ret_f b =>
        e1 (e3.ret (g b))               ==< pmap (e1 `o` e3.ret `o` g) (inv (e1.f_sec b)) >==
        e1 (e3.ret (g (e1 (e1.sec b)))) ==< pmap e1 (e3.ret_f (e1.sec b)) >==
        e1 (e1.sec b)                   ==< e1.f_sec b >==
        b                               `qed
    | f_sec => e3.f_sec

  \func leftFactorPath {A B C : \Type} (e1 : Retraction {A} {B}) {g : B -> C} (e3 : QEquiv {A} {C}) (p : g `o` e1 = e3) : QEquiv g
    => leftFactor e1 g (transportInv (QEquiv __) p e3)

  \func rightFactor {A B C : \Type} (f : A -> B) (e2 : Section {B} {C}) (e3 : QEquiv (\lam a => e2 (f a))) : QEquiv f \cowith
    | ret b => e3.ret (e2 b)
    | ret_f => e3.ret_f
    | f_sec b =>
        f (e3.ret (e2 b))               ==< inv (e2.ret_f (f (e3.ret (e2 b)))) >==
        e2.ret (e2 (f (e3.ret (e2 b)))) ==< pmap e2.ret (e3.f_sec (e2 b)) >==
        e2.ret (e2 b)                   ==< e2.ret_f b >==
        b                               `qed

  \func rightFactorPath {A B C : \Type} {f : A -> B} (e2 : Section {B} {C}) (e3 : QEquiv {A} {C}) (p : e2 `o` f = e3) : QEquiv f
    => rightFactor f e2 (transportInv (QEquiv __) p e3)

  \lemma rightEmbedding {A B C : \Type} (f : A -> B) (e2 : B >-> C) (e3 : QEquiv (e2 `o` f)) : Equiv f
    => rightFactor f (\new ESEquiv e2 {
      | isEmb => e2.isEmb
      | isSurjMap c => inP (f (e3.ret c), e3.f_sec c)
    }) e3

  \lemma parallelEquiv {A B C D : \Type} (ab : A -> B) (cd : C -> D) (ac : Equiv {A} {C}) (bd : Equiv {B} {D}) (h : \Pi (a : A) -> bd (ab a) = cd (ac a)) : Equiv ab = Equiv cd
    => propExt (\lam e => leftFactor ac cd (coe (\lam i => QEquiv (h __ @ i)) (transQEquiv e bd) right))
               (\lam e => rightFactor ab bd (coe (\lam i => QEquiv (\lam a => inv (h a) @ i)) (transQEquiv ac e) right))

  \lemma leftEquiv {A B C : \Type} (e : Equiv {A} {B}) (g : B -> C) : Equiv g = Equiv (g `o` e)
    => propExt (\lam e2 => transEquiv e e2) (\lam e3 => leftFactor e g e3)

  \lemma rightEquiv {A B C : \Type} (f : A -> B) (e : Equiv {B} {C}) : Equiv f = Equiv (e `o` f)
    => propExt (\lam e1 => transEquiv e1 e) (\lam e3 => rightFactor f e e3)

  \lemma compositeEquiv {A B C : \Type} (f : A -> B) (g : B -> C) (e : Equiv (g `o` f)) : Equiv f = Equiv g
    => propExt (\lam e1 => leftFactor e1 g e) (\lam e2 => rightFactor f e2 e)
}

\func \fixr 3 transEquiv {A B C : \Type} (e1 : Equiv {A} {B}) (e2 : Equiv {B} {C}) : Equiv {A} {C} \cowith
  | f x => e2 (e1 x)
  | ret y => e1.ret (e2.ret y)
  | ret_f x => pmap e1.ret (e2.ret_f (e1 x)) *> e1.ret_f x
  | sec y => e1.sec (e2.sec y)
  | f_sec y => pmap e2 (e1.f_sec (e2.sec y)) *> e2.f_sec y

\func -o_Equiv {C : \Type} (e : Equiv) : Equiv {e.B -> C} {e.A -> C} (-o e) \cowith
  | ret => -o e.sec
  | ret_f g => path (\lam i x => g (e.f_sec x @ i))
  | sec => -o e.ret
  | f_sec g => path (\lam i x => g (e.ret_f x @ i))

\func o-_Equiv {C : \Type} (e : Equiv) : Equiv {C -> e.A} (e `o-) \cowith
  | ret => e.ret `o-
  | ret_f g => path (\lam i x => e.ret_f (g x) @ i)
  | sec => e.sec `o-
  | f_sec g => path (\lam i x => e.f_sec (g x) @ i)

\func piEquiv {A : \Type} (B : A -> \Type) (f f' : \Pi (a : A) -> B a) : QEquiv \cowith
  | A => f = f'
  | B => \Pi (a : A) -> f a = f' a
  | f p a => path ((p @ __) a)
  | ret g => path (\lam i a => g a @ i)
  | ret_f _ => idp
  | f_sec _ => idp

\func sigmaEquiv {A : \Type} (B : A -> \Type) (p p' : \Sigma (a : A) (B a)) : QEquiv \cowith
  | A => p = p'
  | B => \Sigma (s : p.1 = p'.1) (transport B s p.2 = p'.2)
  | f q => (pmap __.1 q, pmapd __.2 q)
  | ret q => ext q
  | ret_f q => Jl (\lam p'' q' => ext (pmap __.1 q', pmapd __.2 q') = q') idp q
  | f_sec q =>
    ext (idp, \case p'.1 \as p'1        , q.1 \as q1 : p.1 = p'1,
                    p'.2 \as p'2 : B p'1, q.2 \as q2 : transport B q1 p.2 = p'2
              \return pmapd __.2 {p} {p'1, p'2} (ext (q1,q2)) = q2 \with {
                | p'1, idp, p'2, idp => idp
              })

\func piSigmaEquiv {A A' : \Type} (h : A -> A') (B : A' -> \Type)
  : QEquiv {\Pi (a : A) -> B (h a)}
           {\Sigma (g : A -> \Sigma (a' : A') (B a')) ((\lam a => (g a).1) = h)}
  \cowith
  | f g => (\lam a => (h a, g a), idp)
  | ret p => \lam a => transport B (path ((p.2 @ __) a)) (p.1 a).2
  | ret_f g => idp
  | f_sec p => Jl (\lam (h' : A -> A') q => (\lam a => (h' a, transport B (path ((q @ __) a)) (p.1 a).2), idp)
                                            = {\Sigma (g : A -> \Sigma (a' : A') (B a')) ((\lam a => (g a).1) = h')}
                                            (p.1, q))
                  idp
                  p.2

\func piSigmaIdEquiv {A : \Type} (B : A -> \Type)
  : QEquiv {\Pi (a : A) -> B a}
           {\Sigma (g : A -> \Sigma (a : A) (B a)) ((\lam a => (g a).1) = id)}
  => piSigmaEquiv id B

\func emptyEquiv {A B : \Type} (Ae : A -> Empty) (Be : B -> Empty) : QEquiv {A} {B} \cowith
  | f a => absurd (Ae a)
  | ret b => absurd (Be b)
  | ret_f a => absurd (Ae a)
  | f_sec b => absurd (Be b)

-- # Embeddings and surjections

\record Embedding \extends Map {
  | isEmb : \Pi (a a' : A) -> Retraction (pmap f {a} {a'})
  \func pmap-isEquiv {a a' : A} : QEquiv (pmap f {a} {a'}) => pathEquiv (f __ = f __) (\lam {a} {a'} => isEmb a a')
} \where {
  \use \level levelProp {A B : \Type} {f : A -> B} (e e' : Embedding f) : e = e'
    => path (\lam i => \new Embedding f (\lam a a' => Retraction.levelProp (e.pmap-isEquiv {a} {a'}) (e.isEmb a a') (e'.isEmb a a') @ i))

  \func fromInjection {A : \Type} {B : \Set} {f : A -> B} (inj : \Pi {a a' : A} -> f a = f a' -> a = a') : Embedding f \cowith
    | isEmb a a' => \new Retraction {
      | sec => inj
      | f_sec p => prop-pi {f a = f a'}
    }

  -- | The diagonal of an embedding is an equivalence
  \lemma diag-equiv (e : Embedding) : Equiv {e.A} {\Sigma (a1 a2 : e.A) (e a1 = e a2)} (\lam a => (a, a, idp)) =>
    \let | S => \Sigma (x : e.A) (\Sigma (y : e.A) (x = y))
         | T => \Sigma (x y : e.A) (e x = e y)
         | q => \new QEquiv {e.A} {S} (\lam x => (x,(x,idp))) {
                  | ret t => t.1
                  | ret_f x => idp
                  | f_sec t => pmap (\lam r => ((t.1,r) : S)) (Jl (\lam x p => (t.1,idp) = {\Sigma (x : e.A) (t.1 = x)} (x,p)) idp t.2.2)
                }
         | pe {a} {a'} : QEquiv (pmap e {a} {a'}) => e.pmap-isEquiv {a} {a'}
         | s => \new QEquiv {S} {T} (\lam t => (t.1, t.2.1, pmap e t.2.2)) {
                  | ret t => (t.1, (t.2, pe.ret t.3))
                  | ret_f t => pmap (\lam r => ((t.1, (t.2.1, r)) : S)) (pe.ret_f t.2.2)
                  | f_sec t => pmap (\lam r => ((t.1, t.2, r) : T)) (pe.f_sec t.3)
                }
    \in transQEquiv q s

  \lemma projection {A : \Type} (B : A -> \Prop) : Embedding {\Sigma (a : A) (B a)} __.1 \cowith
    | isEmb p p' => \new Retraction {
      | sec q => ext q
      | f_sec => idpe
    }

  -- | If the fibers of a map are propositions, then the map is an embedding
  \lemma fibers {A B : \Type} (f : A -> B) (p : \Pi (b : B) -> isProp (Fib f b)) : Embedding f \cowith
    | isEmb a a' =>
      \have e (q : f a = f a') => Fib.equiv (f a') (a,q) (a',idp) (p (f a') (a,q) (a',idp))
      \in \new Retraction {
        | sec q => (e q).1
        | f_sec q => (e q).2
      }

  -- | The fibers of an embedding are propositions.
  \lemma embeddingFiber-isProp (e : Embedding) (b : e.B) : isProp (Fib e b) => \lam x x' =>
    \let | eq : Retraction (pmap e {x.1} {x'.1}) => e.isEmb x.1 x'.1
         | fx=fx' => x.2 *> inv x'.2
    \in Fib.ext b x x' (eq.sec fx=fx') (
      pmap e (eq.sec fx=fx') *> x'.2 ==< pmap (`*> x'.2) (eq.f_sec fx=fx') >==
      fx=fx' *> x'.2                   ==< *>-assoc _ _ _ >==
      x.2 *> inv x'.2 *> x'.2          ==< pmap (x.2 *>) (inv_*> _) >==
      x.2                              `qed
    )
}

\func transEmbedding {A B C : \Type} (e1 : Embedding {A} {B}) (e2 : Embedding {B} {C}) : Embedding {A} {C} \cowith
  | f a => e2 (e1 a)
  | isEmb a a' => \new Retraction {
    | sec p => sec {e1.isEmb a a'} (sec {e2.isEmb (e1 a) (e1 a')} p)
    | f_sec p => pmap (pmap e2) (f_sec {e1.isEmb a a'} (sec {e2.isEmb (e1 a) (e1 a')} p)) *> f_sec {e2.isEmb (e1 a) (e1 a')} p
  }

\func \infixr 1 >-> (A B : \Type) => Embedding {A} {B}

\record Surjection \extends Map
  | isSurjMap (y : B) :  (x : A) (f x = y)

\func \infixr 1 ->> (A B : \Type) => Surjection {A} {B}

\record ESEquiv \extends Embedding, Surjection, Equiv
  | ret y => (aux \this isSurjMap y).1
  | sec y => (aux \this isSurjMap y).1
  | ret_f x => sec {isEmb _ _} (aux \this isSurjMap (f x)).2
  | f_sec y => (aux \this isSurjMap y).2
  \where {
    \use \level levelProp {A B : \Type} {f : A -> B} (e e' : ESEquiv f) : e = e' =>
      path (\lam i => \new ESEquiv { | Embedding  => Embedding.levelProp  e e' @ i
                                     | Surjection => prop-isProp {Surjection f} e e' @ i })

    \use \coerce fromEquiv (e : Equiv) : ESEquiv e.f \cowith
      | isEmb a a' => pmapEquiv e {a} {a'}
      | isSurjMap y => inP (e.sec y, e.f_sec y)

    \lemma aux (e : Embedding) (s : \Pi (y : e.B) ->  (x : e.A) (f x = y)) (y : e.B)
      => TruncP.remove (Embedding.embeddingFiber-isProp e y) (s y)
  }