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

-- # Definition of equivalences

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

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

    \lemma levelProp {A B : \Type} (r : Retraction {A} {B}) : isProp (Section r)
      => \lam s => isContr=>isProp (Contr.Contr_QEquiv (symQEquiv Section~Fiber) $
        IsEquiv=>contrFibers.fromSection (\new Section (-o r) {
          | ret => -o r.sec
          | ret_f g => path \lam i x => g (r.f_sec x i)
        }) (\new Retraction {
          | sec => -o s.ret
          | f_sec g => path \lam i x => g (s.ret_f x i)
        }) id) s
  }

\record Retraction \extends TypeMap
  | sec : B -> A
  | f_sec : \Pi (y : B) -> f (sec y) = y
  \where {
    \func Retraction~Fiber {A B : \Type} {f : A -> B} : QEquiv {Retraction f} {Fib (f `o-) id} \cowith
      | f r => Fib.make r.sec $ path \lam i x => r.f_sec x i
      | ret z => \new Retraction f z.1 \lam x => path (z.2 __ x)
      | ret_f => idpe
      | f_sec => idpe

    \lemma levelProp {A B : \Type} (s : Section {A} {B}) : isProp (Retraction s)
      => \lam r => isContr=>isProp (Contr.Contr_QEquiv (symQEquiv Retraction~Fiber) $
        IsEquiv=>contrFibers.fromSection (\new Section (s `o-) {
          | ret => s.ret `o-
          | ret_f g => path \lam i x => s.ret_f (g x) i
        }) (\new Retraction {
          | sec => r.sec `o-
          | f_sec g => path \lam i x => r.f_sec (g x) i
        }) id) r

    \lemma isContr (e : QEquiv) : Contr (Retraction e)
      => isProp=>isContr (levelProp e) e
  }

\record QEquiv \extends Section, Retraction {
  | sec => ret

  \func f_ret (y : B) : f (ret y) = 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 : 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 {
    \sfunc fromInjSurj {A B : \Set} {f : A -> B} (inj : IsInj f) (surj : IsSurj f) : QEquiv f
      => \have s b => TruncP.remove (\lam t t' => ext $ inj $ t.2 *> inv t'.2) (surj b)
         \in \new QEquiv {
           | ret b => (s b).1
           | ret_f a => inj (s (f a)).2
           | f_sec b => (s b).2
         }

    \lemma equals {A B : \Set} {e e' : QEquiv {A} {B}} (p : \Pi (a : A) -> e a = e' a) : e = e'
      => exts (p, \lam b => e.isInj $ e.f_ret b *> inv (p _ *> e'.f_ret b))

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

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

\func QEquiv_= {A B : \Type} (e : QEquiv {A} {B}) (i : I) : \Type
  => iso e.f e.ret e.ret_f e.f_sec i

\type IsEquiv {A B : \Type} (f : A -> B) : \Prop
  => TruncP (QEquiv f)
  \where {
    \lemma toBiEquiv (e : IsEquiv f) : BiEquiv f \level BiEquiv.levelProp {A} {B} {f}
      | inP e => e

    \sfunc toQEquiv (e : IsEquiv f) : QEquiv f
      => toBiEquiv e

    \sfunc ret (e : IsEquiv f) (b : B) : A
      => (toQEquiv e).ret b

    \sfunc ret_f (e : IsEquiv f) {a : A} : ret e (f a) = a
      => (\peval ret e (f a)) *> (toQEquiv e).ret_f a

    \sfunc f_ret (e : IsEquiv f) {b : B} : f (ret e b) = b
      => pmap f (\peval ret e b) *> (toQEquiv e).f_ret b

    \func splitSurj (e : IsEquiv f) (b : B) : \Sigma (a : A) (f a = b)
      => (ret e b, f_ret e)

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

    \lemma ret-equiv (e : IsEquiv f) : IsEquiv (ret e)
      => inP \new QEquiv {
        | ret => f
        | ret_f b => f_ret e
        | f_sec a => ret_f e
      }

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

    \lemma isSurj (e : IsEquiv f) : IsSurj f \elim e
      | inP e => e.isSurj

    \lemma fromEmbSurj {A B : \Type} {f : A -> B} (emb : Embedding f) (surj : IsSurj f) : IsEquiv f
      => inP \new QEquiv {
        | ret b => (emb.surj-split surj b).1
        | ret_f a => (emb.isEmb _ _).sec (emb.surj-split surj (f a)).2
        | f_sec b => (emb.surj-split surj b).2
      }

    \lemma fromInjSurj {A B : \Set} {f : A -> B} (inj : IsInj f) (surj : IsSurj f) : IsEquiv f
      => inP (QEquiv.fromInjSurj inj surj)

    \lemma fromInjInv {A B : \Set} {f : A -> B} (g : B -> A) (inj : IsInj g) (surj : \Pi (a : A) -> g (f a) = a) : IsEquiv f
      => inP \new QEquiv {
        | ret => g
        | ret_f => surj
        | f_sec b => inj (surj (g b))
      }

    \lemma trans {A B C : \Type} {f : A -> B} (e : IsEquiv f) {g : B -> C} (e' : IsEquiv g) : IsEquiv \lam a => g (f a) \elim e, e'
      | inP e, inP e' => inP (transQEquiv e e')

    \lemma leftFactor {A B C : \Type} {f : A -> B} {g : B -> C} (e2 : IsEquiv g) (e3 : IsEquiv \lam a => g (f a)) : IsEquiv f \elim e2, e3
      | inP e2, inP e3 => inP (QEquiv.leftFactor e2 e3)

    \lemma rightFactor {A B C : \Type} {f : A -> B} (e1 : IsEquiv f) {g : B -> C} (e3 : IsEquiv \lam a => g (f a)) : IsEquiv g \elim e1, e3
      | inP e1, inP e3 => inP (QEquiv.rightFactor e1 e3)

    \lemma parallelEquiv {A B C D : \Type} {ab : A -> B} {cd : C -> D} {ac : A -> C} {bd : B -> D}
                         (ace : IsEquiv ac) (bde : IsEquiv bd) (h : \Pi (a : A) -> bd (ab a) = cd (ac a)) : IsEquiv ab = IsEquiv cd
      => propExt (\lam abe => rightFactor ace $ transport IsEquiv (ext h) $ trans abe bde)
                 (\lam cde => leftFactor bde $ transportInv IsEquiv (ext h) $ trans ace cde)
  }

\data Equiv (A B : \Type)
  | mkEquiv {f : A -> B} (IsEquiv f)
  \where {
    \use \coerce toFunc (e : Equiv A B) : A -> B \elim e
      | mkEquiv {f} _ => f

    \func toQEquiv (e : Equiv A B) : QEquiv {A} {B} e \elim e
      | mkEquiv e => IsEquiv.toQEquiv e

    \protected \func equals {e e' : Equiv A B} (p : toFunc e = toFunc e') : e = e' \elim e, e'
      | mkEquiv e, mkEquiv e' => path \lam i => mkEquiv {_} {_} {p i} (prop-dpi _ _ _ i)

    \lemma toFunc-embedding : Embedding {Equiv A B} {A -> B} toFunc \cowith
      | isEmb => \case \elim __, \elim __ \with {
        | mkEquiv e, mkEquiv e' => \new Retraction {
          | sec => equals
          | f_sec => idpe
        }
      }
  }

-- # 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

\lemma -o_IsEquiv {A B C : \Type} {f : A -> B} (e : IsEquiv f) : IsEquiv {B -> C} {A -> C} (-o f)
  => inP \new QEquiv {
    | ret => -o (IsEquiv.ret e)
    | ret_f g => path \lam i x => g (IsEquiv.f_ret e {x} i)
    | f_sec g => path \lam i x => g (IsEquiv.ret_f e {x} i)
  }

\lemma o-_IsEquiv {A B C : \Type} {f : A -> B} (e : IsEquiv f) : IsEquiv {C -> A} {C -> B} (f `o-)
  => inP \new QEquiv {
    | ret => IsEquiv.ret e `o-
    | ret_f g => path (\lam i x => IsEquiv.ret_f e {g x} i)
    | f_sec g => path (\lam i x => IsEquiv.f_ret e {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 TypeMap {
  | 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')

  \lemma surj-split (s : IsSurj f) (y : B)
    => TruncP.remove (Embedding.embeddingFiber-isProp \this y) (s y)
} \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 (pathEquiv (e __ = e __) (\lam {a} {a'} => e.isEmb 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'}
    }

  \sfunc fromIsEquiv {A B : \Type} {f : A -> B} (e : IsEquiv f) : Embedding f \cowith
    | isEmb a a' => IsEquiv.toQEquiv (pmapIsEquiv e)

  -- | The diagonal of an embedding is an equivalence
  \lemma diag-equiv (e : Embedding) : IsEquiv {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 inP (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
  }

\lemma piEmbedding {A : \Type} {B C : A -> \Type} (e : \Pi (a : A) -> Embedding {B a} {C a}) : Embedding {\Pi (a : A) -> B a} {\Pi (a : A) -> C a} (\lam f a => e a (f a)) \cowith
  | isEmb f g => \new Retraction {
    | sec p => path \lam i a => ((e a).isEmb (f a) (g a)).sec (path \lam j => p j a) i
    | f_sec p => path \lam i j a => ((e a).isEmb (f a) (g a)).f_sec (path \lam j => p j a) i j
  }

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