\import Function
\import Logic
\import Logic.Classical
\import Logic.Meta
\import Meta
\import Paths
\import Set

\class Transitive (A : \Set)
  | \infix  ~ : A -> A -> \Prop
  | ~-transitive {x y z : A} : x ~ y -> y ~ z -> x ~ z
  \where {
    \truncated \data Closure {A : \Type} (R : A -> A -> \Type) (x y : A) : \Prop
      | cin (R x y)
      | ctrans {z : A} (Closure R x z) (Closure R z y)
      \where {
        \func isTransitive {A : \Set} (R : A -> A -> \Type) : Transitive A \cowith
          | ~ x y => Closure R x y
          | ~-transitive => ctrans

        \lemma isSymmetric {A : \Set} {R : A -> A -> \Type} (sym : \Pi {a a' : A} -> R a a' -> R a' a) {a a' : A} (c : Closure R a a') : Closure R a' a \elim c
          | cin r => cin (sym r)
          | ctrans c1 c2 => ctrans (isSymmetric sym c2) (isSymmetric sym c1)

        \lemma univ {A : \Set} {R : A -> A -> \Type} (E : Transitive A) (f : \Pi {x y : A} -> R x y -> x ~ y) {x y : A} (c : Closure R x y) : x ~ y \elim c
          | cin r => f r
          | ctrans p q => ~-transitive (univ E f p) (univ E f q)

        \lemma ofTransitive (E : Transitive) {x y : E} (c : Closure (~) x y) : x ~ y => univ E (\lam p => p) c

        \lemma toEquality {A B : \Set} (f : A -> B) {R : A -> A -> \Type} (p : \Pi {a a' : A} -> R a a' -> f a = f a') {a a' : A} (c : Closure R a a') : f a = f a'
          => univ (\new Transitive A (f __ = f __) (*>)) p c
      }
  }

\class Equivalence \extends Transitive
  | ~-reflexive {x : A} : x ~ x
  | ~-symmetric {x y : A} : x ~ y -> y ~ x
  \where {
    \func map {A B : \Set} (f : A -> B) (E : Equivalence B) : Equivalence A \cowith
      | ~ a a' => f a ~ f a'
      | ~-transitive => ~-transitive
      | ~-reflexive => ~-reflexive
      | ~-symmetric => ~-symmetric

    \truncated \data Closure {A : \Type} (R : A -> A -> \Type) (x y : A) : \Prop
      | cin (R x y)
      | crefl (x = y)
      | csym (Closure R y x)
      | ctrans {z : A} (Closure R x z) (Closure R z y)
      \where {
        \func isEquivalence {A : \Set} (R : A -> A -> \Type) : Equivalence A \cowith
          | ~ x y => Closure R x y
          | ~-reflexive {x} => crefl idp
          | ~-symmetric {x} {y} x~y => csym x~y
          | ~-transitive {x} {y} {z} x~y y~z => ctrans x~y y~z

        \lemma univ {A : \Set} {R : A -> A -> \Type} (E : Equivalence A) (f : \Pi {x y : A} -> R x y -> x ~ y) {x y : A} (c : Closure R x y) : x ~ y \elim c
          | cin r => f r
          | crefl x=y => transport (x ~) x=y ~-reflexive
          | csym p => ~-symmetric (univ E f p)
          | ctrans {z} p q => ~-transitive (univ E f p) (univ E f q)

        \lemma ofEquivalence (E : Equivalence) {x y : E} (c : Closure (~) x y) : x ~ y => univ E (\lam p => p) c
      }
  }

\truncated \data Quotient {A : \Type} (R : A -> A -> \Type) : \Set
  | in~ A
  | ~-equiv (x y : A) (R x y) : in~ x = in~ y
  \where {
    \open Equivalence

    \lemma in-surj {A : \Type} {R : A -> A -> \Type} : isSurj (in~ {A} {R})
      => \lam (in~ a) => inP (a,idp)

    \lemma equality {A : \Type} {R : A -> A -> \Type} {x y : A} (p : in~ x = {Quotient R} in~ y) : Closure R x y =>
      \let | R' ([x] [y] : Quotient R) : \Prop =>
              \case [x], [y] \with {
                | in~ x, in~ y => Closure R x y
                | in~ x, ~-equiv y y' y~y' => propExt (ctrans __ (cin y~y')) (ctrans __ (csym (cin y~y')))
                | ~-equiv x x' x~x', in~ y => propExt (ctrans (csym (cin x~x')) __) (ctrans (cin x~x') __)
              }
           | R'-refl ([x] : Quotient R) =>
               \case \elim [x] \return R' [x] [x] \with {
                 | in~ x => crefl idp
               }
           | =-implies-R' {[x] [y] : Quotient R} (p : [x] = [y]) => transport (R' [x]) p (R'-refl [x])
      \in =-implies-R' p

    \lemma equalityEquiv (E : Equivalence) {x y : E} (p : in~ x = {Quotient (~)} in~ y) : x ~ y
      => Closure.ofEquivalence E (equality p)

    \lemma equalityClosure (E : Equivalence) {R : E -> E -> \Type} (s : \Pi {x y : E} -> R x y -> x ~ y) {x y : E} (p : in~ x = {Quotient R} in~ y) : x ~ y
      => Closure.univ E s (equality p)

    \lemma fromEquality {A : \Type} {R : A -> A -> \Type} {x y : A} (p : Closure R x y) : in~ x = {Quotient R} in~ y \elim p
      | cin r => path (~-equiv x y r)
      | crefl p => pmap in~ p
      | csym p => inv (fromEquality p)
      | ctrans p1 p2 => fromEquality p1 *> fromEquality p2

    \func map {A B : \Type} (f : A -> B) {R : A -> A -> \Type} {Q : B -> B -> \Type} (p : \Pi {a a' : A} -> R a a' -> Q (f a) (f a')) (x : Quotient R) : Quotient Q \elim x
      | in~ a => in~ (f a)
      | ~-equiv x y r => ~-pequiv (p r)

    \func liftArray {k : Nat} {A : Fin k -> \Type} (R : \Pi {j : Fin k} -> A j -> A j -> \Type) (refl : \Pi {j : Fin k} (x : A j) -> R x x) (ts : DArray (\lam j => Quotient (R {j})))
      : Quotient {DArray A} (\lam f g => \Pi (j : Fin ts.len) -> R (f j) (g j)) \elim k, ts
      | , nil => in~ nil
      | suc k, :: a l => \case \elim a, liftArray (\lam {j} => R {suc j}) (\lam {j} => refl {suc j}) l \with {
        | in~ a, in~ a1 => in~ (a :: a1)
        | in~ a, ~-equiv x y r i => ~-equiv _ _ (later \case \elim __ \with {
          |  => refl _
          | suc j => r j
        }) i
        | ~-equiv x y r i, in~ a => ~-equiv _ _ (later \case \elim __ \with {
          |  => r
          | suc j => refl _
        }) i
      }

    \func liftArrayFun {k : Nat} {A : Fin k -> \Type} (R : \Pi {j : Fin k} -> A j -> A j -> \Type) (refl : \Pi {j : Fin k} (x : A j) -> R x x) {B : \Type} (S : B -> B -> \Type) (h : DArray A -> B)
                       (p : \Pi (l l' : DArray A) -> (\Pi (j : Fin k) -> R (l j) (l' j)) -> S (h l) (h l'))
                       (ts : DArray (\lam j => Quotient (R {j}))) : Quotient S
      => \case liftArray R refl ts \with {
        | in~ l => in~ (h l)
        | ~-equiv x y r i => ~-equiv (h x) (h y) (p x y r) i
      }

    \lemma liftChoice {A : Choice} (B : A -> \Set) (R : \Pi {a : A} -> B a -> B a -> \Type) (f : \Pi (a : A) -> Quotient (\lam b b' => R {a} b b'))
      : ∃ (g : \Pi (a : A) -> B a) ∀ a (in~ (g a) = f a)
      => TruncP.map (choice {_} {\lam a => \Sigma (b : B a) (in~ b = f a)} (\lam a => cases (f a) \with { | in~ b => inP (b,idp) }))
          (\lam g => (\lam a => (g a)., \lam a => (g a).))
  }

\func ~-pequiv {A : \Type} {R : A -> A -> \Type} {x y : A} (r : R x y) : in~ x = in~ y => path (~-equiv x y r)

\type SubQuotient {A : \Type} (R : A -> A -> \Type)
  => Quotient {\Sigma (a : A) (R a a)} (\lam p q => R p. q.)

\func QuotientDec {R : Equivalence} (D : \Pi (a a' : R) -> Dec (a ~ a')) : DecSet (Quotient (~)) \cowith
  | decideEq (in~ a) (in~ a') => \case D a a' \with {
    | yes e => yes (~-pequiv e)
    | no q => no \lam r => q (Quotient.equalityEquiv R r)
  }