(** Theory of finite probability distributions.  This file defines the {prob T}
type, which represents finite probability distributions with rational
coefficients over the type [T], which is required to be an ordered type.
Distributions form a monad, where the unit [dirac : T -> {prob T}] is the point
mass distribution, whereas the monadic bind [sample : {prob T} -> (T -> {prob
S}) -> {prob S}] samples from its first argument and feeds the sample to its
second argument. *)

Require Import Coq.Strings.String.
Require Import Coq.Unicode.Utf8.

From mathcomp Require Import ssreflect ssrfun ssrbool eqtype ssrnat choice seq
  ssrint rat ssralg ssrnum bigop path.

From extructures Require Import ord fset fmap ffun.

From deriving Require Import deriving.

From Netter Require Import Extra.

Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

Local Open Scope ring_scope.
Local Open Scope fset_scope.

Section Prob.

Variable T : ordType.
Implicit Types (x : T) (X : {fset T}) (f : T -> rat).

(** We define probability distributions in two stages.  First, we define a
distribution as a map into the positive rationals that is 0 almost everywhere.
Then, we define {prob T} as the type of distributions of mass 1. *)

Record distr := Distr {
  dval :> ffun (fun _ : T => 0 : rat);
  _    :  all (fun x => 0 <= dval x) (supp dval)
}.
Definition distr_of & phant T := distr.
Notation "{ 'distr' T }" := (distr_of (Phant T))
  (at level 0, format "{ 'distr'  T }") : form_scope.
Identity Coercion distr_of_distr : distr_of >-> distr.

Canonical distr_subType  := [subType for dval].
Definition distr_eqMixin := [eqMixin of distr by <:].
Canonical distr_eqType   := EqType distr distr_eqMixin.
Definition distr_choiceMixin := [choiceMixin of distr by <:].
Canonical distr_choiceType := Eval hnf in ChoiceType distr distr_choiceMixin.
Definition distr_ordMixin := [ordMixin of distr by <:].
Canonical distr_ordType := Eval hnf in OrdType distr distr_ordMixin.

Canonical distr_of_newType := [subType of {distr T}].
Canonical distr_of_eqType  := [eqType of {distr T}].
Canonical distr_of_choiceType := [choiceType of {distr T}].
Canonical distr_of_ordType := [ordType of {distr T}].

Implicit Types (d : {distr T}).

Lemma eq_distr d1 d2 : d1 =1 d2 <-> d1 = d2.
Proof.
by split=> [e|-> //]; apply/val_inj/eq_ffun.
Qed.

Lemma distr_ge0 d x : 0 <= d x.
Proof.
case: (boolP (x \in supp d))=> [|/suppPn -> //].
apply/allP: x; exact/(valP d).
Qed.

Definition mass d := \sum_(x <- supp d) d x.

Lemma massE d X :
  fsubset (supp d) X -> mass d = \sum_(x <- X) d x.
Proof.
rewrite /mass => sub.
have -> : supp d = fset_filter (mem (supp d)) X.
  apply/eq_fset=> x; rewrite in_fset_filter /=.
  by case: (boolP (x \in supp d))=> // /(fsubsetP sub) ->.
rewrite val_fset_filter big_filter big_mkcond /=.
by apply/eq_big=> // x _; case: ifPn=> // /suppPn ->.
Qed.

Fact mkdistr_key : unit. Proof. exact: tt. Qed.

Lemma mkdistr_subproof (X : {fset T}) (f : T -> rat) :
  (forall x, x \in X -> 0 <= f x) ->
  let ff := @mkffun _ _ (fun _ => 0) f X in
  all (fun x => 0 <= ff x) (supp ff).
Proof.
move=> /= pos; apply/allP=> x _; rewrite mkffunE.
by case: ifP=> // /pos.
Qed.

Definition mkdistr X f pos :=
  locked_with mkdistr_key
    (Distr (@mkdistr_subproof X f pos)).

Lemma mkdistrE X f pos x :
  @mkdistr X f pos x = if (x \in X) then f x else 0%R.
Proof. by rewrite /mkdistr unlock /= mkffunE. Qed.

Lemma supp_mkdistr X f pos : supp (@mkdistr X f pos) = fset_filter (fun x => f x != 0) X.
Proof.
apply/eq_fset=> x; rewrite mem_supp mkdistrE.
by rewrite -mkffunE -mem_supp supp_mkffun in_fset mem_filter.
Qed.

Record prob :=
  Prob {pval : {distr T}; _ : mass pval == 1%R}.
Definition prob_of & phant T := prob.
Notation "{ 'prob' T }" := (prob_of (Phant T))
  (at level 0, format "{ 'prob'  T }") : form_scope.
Identity Coercion prob_of_prob : prob_of >-> prob.
Coercion pval : prob >-> distr_of.

Canonical prob_subType := [subType for pval].
Definition prob_eqMixin := [eqMixin of prob by <:].
Canonical prob_eqType := EqType prob prob_eqMixin.
Definition prob_choiceMixin := [choiceMixin of prob by <:].
Canonical prob_choiceType := Eval hnf in ChoiceType prob prob_choiceMixin.
Definition prob_ordMixin := [ordMixin of prob by <:].
Canonical prob_ordType := Eval hnf in OrdType prob prob_ordMixin.

Canonical prob_of_subType := [subType of {prob T}].
Canonical prob_of_eqType  := [eqType  of {prob T}].
Canonical prob_of_choiceType := [choiceType of {prob T}].
Canonical prob_of_ordType := [ordType of {prob T}].

Implicit Types (p : {prob T}).

Lemma mkprob_subproof X f pos :
  \sum_(x <- X) f x = 1 ->
  mass (@mkdistr X f pos) == 1.
Proof.
move=> <-; apply/eqP.
rewrite [RHS](bigID (fun x => f x == 0)) /=.
have -> : \sum_(x <- X | f x == 0) f x = 0.
  by rewrite big1 // => x /eqP.
rewrite GRing.add0r /mass supp_mkdistr val_fset_filter big_filter.
rewrite big_seq_cond [RHS]big_seq_cond.
by apply: eq_big=> // x /andP [x_X _]; rewrite mkdistrE x_X.
Qed.

Fact mkprob_key : unit. Proof. exact: tt. Qed.

Definition mkprob X f pos e : {prob T} :=
  locked_with mkprob_key (Prob (@mkprob_subproof X f pos e)).

Lemma mkprobE X f pos e x :
  @mkprob X f pos e x = if x \in X then f x else 0.
Proof. by rewrite /mkprob unlock /= mkdistrE. Qed.

Lemma suppPrN0 p : supp p != fset0.
Proof.
by apply/eqP=> e; move: (eqP (valP p)); rewrite /mass e big_nil.
Qed.

Definition dirac_def x x' : rat :=
  if x == x' then 1 else 0.

Lemma dirac_subproof1 x x' : x' \in fset1 x -> 0 <= dirac_def x x'.
Proof. by rewrite /dirac_def; case: eq_op. Qed.

Lemma dirac_subproof2 x : \sum_(x' <- fset1 x) dirac_def x x' = 1.
Proof.
by rewrite /= big_seq1 /dirac_def eqxx.
Qed.

(** Point mass distribution *)

Definition dirac x :=
  mkprob (@dirac_subproof1 x) (dirac_subproof2 x).

Lemma diracE x x' : dirac x x' = if x' == x then 1 else 0.
Proof.
rewrite /dirac /dirac_def /= mkprobE in_fset1.
by rewrite eq_sym; case: (x == x').
Qed.

Lemma supp_dirac x : supp (dirac x) = fset1 x.
Proof.
apply/eq_fset=> x'.
by rewrite mem_supp in_fset1 diracE; case: ifP.
Qed.

Lemma supp_diracP x x' : reflect (x' = x) (x' \in supp (dirac x)).
Proof. rewrite supp_dirac; exact: fset1P. Qed.

Lemma dirac_inj : injective dirac.
Proof.
by move=> x y e; apply/fset1_inj; rewrite -!supp_dirac e.
Qed.

Lemma eq_prob p1 p2 : p1 =1 p2 <-> p1 = p2.
Proof. by split=> [/eq_distr/val_inj|-> //]. Qed.

Lemma in_eq_probL p1 p2 : {in supp p1, p1 =1 p2} -> p1 = p2.
Proof.
move=> e; apply/eq_prob=> x.
case: (boolP (x \in supp p1))=> xP; first exact: e.
rewrite (suppPn xP).
have: \sum_(y <- supp p1) p2 y = 1.
  rewrite -(eqP (valP p1)) /mass /= big_seq [RHS]big_seq.
  by apply/eq_big=> // ? /e ->.
rewrite -(eqP (valP p2)) /mass /= [RHS](bigID (mem (supp p1))) /=.
rewrite -[in RHS]big_filter -val_fset_filter.
have ->: fset_filter (mem (supp p1)) (supp p2) = supp p1.
  apply/eq_fset=> y; rewrite in_fset_filter /=.
  case: (boolP (y \in supp p1))=> //= y_p1.
  by rewrite mem_supp -(e _ y_p1) -mem_supp.
rewrite -[LHS]GRing.addr0 => /GRing.addrI/esym/eqP.
rewrite Num.Theory.psumr_eq0; last by move=> ? _; rewrite distr_ge0.
case: (boolP (x \in supp p2)) => x_p2; last by rewrite (suppPn x_p2).
by move=> /allP/(_ _ x_p2); rewrite xP => /eqP ->.
Qed.

Lemma in_eq_projR p1 p2 : {in supp p2, p1 =1 p2} -> p1 = p2.
Proof.
by move=> e; apply/esym/in_eq_probL=> x x_p2; rewrite e.
Qed.

Definition of_dirac p : option T :=
  if val (supp p) is [:: x] then Some x
  else None.

Lemma diracK : pcancel dirac of_dirac.
Proof. by move=> x; rewrite /of_dirac supp_dirac /=. Qed.

Lemma of_diracK : ocancel of_dirac dirac.
Proof.
rewrite /of_dirac => p.
case e: (val (supp p))=> [//|x[|//]] /=.
have {}e: supp p = fset1 x by rewrite fset1E -e fsvalK.
move/eqP: (valP p); rewrite /mass e /= big_seq1 => p_x.
apply/in_eq_projR=> y; rewrite e => /fset1P ->.
by rewrite p_x diracE eqxx.
Qed.

Lemma eq_supp_dirac p x : (supp p == fset1 x) = (p == dirac x).
Proof.
apply/(sameP eqP)/(iffP eqP)=> [->|e]; first exact: supp_dirac.
by move: (of_diracK p); rewrite /of_dirac e /= => <-.
Qed.

End Prob.

Notation "{ 'distr' T }" := (distr_of (Phant T))
  (at level 0, format "{ 'distr'  T }") : form_scope.
Notation "{ 'prob' T }" := (prob_of (Phant T))
  (at level 0, format "{ 'prob'  T }") : form_scope.

Arguments dirac {_} x.
Arguments of_dirac {_} p.
Arguments dirac_inj {_}.

Section Sample.

Variables T S : ordType.
Variable (p : {prob T}) (f : T -> {prob S}).
Implicit Types (x : T) (y : S).

Let Y   : {fset S} := \bigcup_(x <- supp p) supp (f x).
Let P y : rat      := \sum_(x <- supp p) p x * f x y.

Lemma sample_subproof1 y : y \in Y -> 0 <= P y.
Proof.
move=> _; apply: Num.Theory.sumr_ge0 => x _.
apply: Num.Theory.mulr_ge0; exact: distr_ge0.
Qed.

Lemma sample_subproof2 : \sum_(y <- Y) P y = 1.
Proof.
rewrite /P exchange_big /= -(eqP (valP p)).
apply/eq_big=> //= x _.
case: (boolP (x \in supp p)); last first.
  by move=> /suppPn ->; apply/big1=> y _; rewrite GRing.mul0r.
rewrite -GRing.mulr_sumr -[RHS]GRing.mulr1 => x_p; congr (_ * _).
rewrite -(eqP (valP (f x))) /=; symmetry; apply/massE.
apply/fsubsetP=> y; rewrite mem_supp => yP.
apply/bigcupP; exists x=> //.
by rewrite mem_supp.
Qed.

Definition sample :=
  mkprob sample_subproof1 sample_subproof2.

Lemma sample_defE0 y : (y \in Y) = (P y != 0).
Proof.
rewrite /P Num.Theory.psumr_eq0 -?has_predC /=; last first.
  move=> x _; apply: Num.Theory.mulr_ge0; exact: distr_ge0.
apply/(sameP bigcupP)/(iffP hasP).
- case=> /= x x_p n0; exists x=> //.
  by move: n0; rewrite GRing.mulf_eq0 negb_or mem_supp; case/andP.
- case=> /= x; rewrite !mem_supp => x_p _ y_f.
  by exists x; rewrite 1?mem_supp // GRing.mulf_neq0.
Qed.

Lemma sampleE y : sample y = P y.
Proof.
rewrite /sample mkprobE sample_defE0.
by case: (P y =P 0)=> [->|].
Qed.

Lemma supp_sample : supp sample = Y.
Proof.
apply/eq_fset=> x.
by rewrite mem_supp sample_defE0 sampleE.
Qed.

Lemma supp_sampleP y :
  reflect (exists2 x, x \in supp p & y \in supp (f x)) (y \in supp sample).
Proof.
rewrite supp_sample; apply/(iffP bigcupP).
- by case=> ????; eauto.
- by case=> ???; eexists; eauto.
Qed.

End Sample.

Arguments supp_sampleP {_ _ _ _ _}.

Declare Scope prob_scope.
Local Open Scope prob_scope.

(** A more convenient notation for sampling. *)

Notation "'sample:' x '<-' t1 ';' t2" :=
  (sample t1 (fun x => t2))
  (at level 20, t1 at level 100, t2 at level 200,
   right associativity, format "'[' 'sample:'  x  '<-'  '[' t1 ;  ']' ']' '/' t2")
  : prob_scope.

Section SampleProperties.

Variables T S : ordType.

Lemma sample_diracL (x : T) (f : T -> {prob S}) : sample (dirac x) f = f x.
Proof.
apply/eq_prob=> y; rewrite sampleE supp_dirac /= big_seq1.
by rewrite mkprobE in_fset1 eqxx /dirac_def eqxx GRing.mul1r.
Qed.

Lemma sample_diracR (p : {prob T}) : sample p dirac = p.
Proof.
apply/eq_prob=> x; rewrite sampleE.
transitivity (\sum_(x' <- supp p) if x == x' then p x' else 0).
  apply/eq_big=> // x' _.
  rewrite diracE /= eq_sym.
  case: eq_op; by rewrite ?GRing.mulr0 ?GRing.mulr1.
rewrite -big_mkcond /= -big_filter -val_fset_filter.
case: (boolP (x \in supp p))=> x_p.
  rewrite (_ : fset_filter _ _ = fset1 x) /= ?big_seq1 //.
  apply/eq_fset=> x'; rewrite in_fset_filter in_fset1 eq_sym.
  by case: (x' =P x)=> [->|].
rewrite (_ : fset_filter _ _ = fset0) /= ?big_nil ?(suppPn x_p) //.
apply/eq_fset=> x'; rewrite in_fset_filter eq_sym.
by case: (x' =P x)=> // ->; rewrite (negbTE x_p).
Qed.

Lemma eq_sample (p : {prob T}) (f g : T -> {prob S}) :
  f =1 g -> sample p f = sample p g.
Proof.
move=> efg; apply/eq_prob=> y.
by rewrite !sampleE; apply/eq_big=> // x _; rewrite efg.
Qed.

Lemma eq_in_sample (p : {prob T}) (f g : T -> {prob S}) :
  {in supp p, f =1 g} -> sample p f = sample p g.
Proof.
move=> efg; apply/eq_prob=> y.
rewrite !sampleE big_seq [in RHS]big_seq.
by apply/eq_big=> // x /efg ->.
Qed.

Lemma sample_const (px : {prob T}) (py : {prob S}) :
  (sample: _ <- px; py) = py.
Proof.
apply/eq_prob=> y.
rewrite sampleE -GRing.mulr_suml -[RHS]GRing.mul1r; congr *%R.
exact/eqP/(valP px).
Qed.

Lemma eq_sample_dirac (p : {prob T}) (f : T -> {prob S}) y :
  sample p f = dirac y ->
  forall x, x \in supp p -> f x = dirac y.
Proof.
move=> e x x_p.
have {}e: supp (sample p f) = supp (dirac y) by rewrite e.
rewrite supp_sample supp_dirac in e.
apply/eqP; rewrite -eq_supp_dirac eqEfsubset; apply/andP; split.
  rewrite -e; exact/bigcup_sup.
rewrite fsub1set; have /fset0Pn [z zP] := suppPrN0 (f x).
suff: z \in fset1 y by move=> /fset1P => <-.
by rewrite -e; apply/bigcupP; exists x.
Qed.

End SampleProperties.

Lemma sampleA (T S R : ordType) p (f : T -> {prob S}) (g : S -> {prob R}) :
  (sample: y <- (sample: x <- p; f x); g y) =
  (sample: x <- p; sample: y <- f x; g y).
Proof.
apply/eq_prob=> z.
transitivity (\sum_(y <- supp (sample: x <- p; f x))
                \sum_(x <- supp p) p x * f x y * g y z).
  rewrite sampleE; apply/eq_big=> // y _.
  by rewrite sampleE GRing.mulr_suml.
rewrite sampleE exchange_big /= big_seq [RHS]big_seq.
apply/eq_big=> // x px.
transitivity (\sum_(y <- supp (sample: x <- p; f x))
                 p x * (f x y * g y z)).
  by apply/eq_big=> ? // _; rewrite GRing.mulrA.
rewrite <- GRing.mulr_sumr; congr *%R; rewrite sampleE.
have /fsetIidPl <-: fsubset (supp (f x)) (supp (sample: x <- p; f x)).
  apply/fsubsetP=> y fxy; rewrite supp_sample.
  by apply/bigcupP; exists x.
rewrite /fsetI val_fset_filter big_filter [RHS]big_mkcond /=.
apply/eq_big=> // y _; rewrite mem_supp.
by case: eqP=> //= ->; rewrite GRing.mul0r.
Qed.

Lemma sampleC (T S R : ordType) (p1 : {prob T}) (p2 : {prob S}) (f : T -> S -> {prob R}) :
  (sample: x <- p1; sample: y <- p2; f x y) =
  (sample: y <- p2; sample: x <- p1; f x y).
Proof.
apply/eq_prob=> z; rewrite !sampleE.
under eq_big => [x|x _]; first over.
  rewrite sampleE GRing.mulr_sumr; over.
under [in RHS]eq_big=> [y|y _]; first over.
  rewrite sampleE GRing.mulr_sumr.
  under eq_big=> [x|x _]; first over.
    rewrite GRing.mulrA [p2 _ * _]GRing.mulrC -GRing.mulrA; over.
  over.
by rewrite /= exchange_big.
Qed.

Open Scope prob_scope.

Section Uniform.

Variable T : ordType.
Variable X : {fset T}.
Hypothesis Xn0 : X != fset0.

Definition unif_def (x : T) : rat := (size X)%:R^-1.

Lemma unif_subproof1 x : x \in X -> 0 <= unif_def x.
Proof.
move=> _; by rewrite Num.Theory.invr_ge0 Num.Theory.ler0n.
Qed.

Lemma unif_subproof2 : \sum_(x <- X) unif_def x = 1.
Proof.
rewrite big_const_seq count_predT -Monoid.Theory.iteropE.
rewrite -[RHS](@GRing.mulfV _ (size X)%:R) ?GRing.mulr_natl //.
by rewrite Num.Theory.pnatr_eq0 sizes_eq0.
Qed.

Definition unif := mkprob unif_subproof1 unif_subproof2.

Lemma unifE x : unif x = if x \in X then (size X)%:R^-1 else 0.
Proof. by rewrite mkprobE. Qed.

Lemma supp_unif : supp unif = X.
Proof.
apply/eq_fset => x; rewrite mem_supp unifE.
by case: ifP => // _; rewrite GRing.invr_eq0 Num.Theory.pnatr_eq0 sizes_eq0.
Qed.

End Uniform.

(** To simplify the reasoning about our optimizations, we use probabilistic
couplings.  A coupling is a way of lifting a relation between two sets to a
relation over distributions over these sets.  (NB: We define this relation in
[Type] to avoid issues with the axiom of choice when proving [coupling_sample]
below.)  The definition is useful because (1) it has good composition
properties, as seen by [coupling_dirac] and [coupling_sample], and because (2)
it is strong enough to establish the equality of two distributions, as seen in
[coupling_same]. *)

Variant coupling (T S : ordType) (R : T -> S -> Prop) pT pS : Type :=
| Coupling p of
  pT = sample p (dirac \o fst) &
  pS = sample p (dirac \o snd) &
  (forall xy, xy \in supp p -> R xy.1 xy.2).

Lemma coupling_dirac (T S : ordType) (R : T -> S -> Prop) x y :
  R x y -> coupling R (dirac x) (dirac y).
Proof.
move=> xy; exists (dirac (x, y)); rewrite ?sample_diracL //.
by move=> [??] /supp_diracP [-> ->].
Qed.

Lemma coupling_sample (T1 S1 T2 S2 : ordType) (R1 : T1 -> S1 -> Prop) (R2 : T2 -> S2 -> Prop) pT pS f g :
  coupling R1 pT pS ->
  (forall x y, x \in supp pT -> y \in supp pS -> R1 x y ->
    coupling R2 (f x) (g y)) ->
  coupling R2 (sample pT f) (sample pS g).
Proof.
case=> /= p eT eS R1P R12.
pose def xy := sample: x' <- f xy.1; sample: y' <- g xy.2; dirac (x', y').
have WT xy : xy \in supp p -> xy.1 \in supp pT.
  move=> xyp; rewrite eT; apply/supp_sampleP.
  exists xy=> //=; exact/supp_diracP.
have WS xy : xy \in supp p -> xy.2 \in supp pS.
  move=> xyp; rewrite eS; apply/supp_sampleP.
  exists xy=> //=; exact/supp_diracP.
pose draw xy := if insub xy is Some xy then
                  let xyP := svalP xy in
                  let xP := WT _ xyP in
                  let yP := WS _ xyP in
                  let: Coupling p _ _ _ := R12 _ _ xP yP (R1P _ xyP) in p
                else def xy.
exists (sample p draw).
- rewrite eT !sampleA; apply/eq_in_sample; case=> [x y] /= xy_supp.
  by rewrite sample_diracL insubT /=; case: (R12 _ _ _ _ _).
- rewrite eS !sampleA; apply/eq_in_sample; case=> [x y] /= xy_supp.
  by rewrite sample_diracL insubT /=; by case: (R12 _ _ _ _ _).
case=> x' y' /supp_sampleP [] [x y] xy_supp.
rewrite /draw insubT /=.
case: (R12 _ _ _ _ _)=> /= pxy eT' eS' R2P; exact: R2P.
Qed.

Lemma coupling_same (T : ordType) (p : {prob T}) : coupling eq p p.
Proof.
pose pp := sample: x <- p; dirac (x, x); exists pp; rewrite ?sampleA.
- under eq_sample do rewrite sample_diracL /=.
  by rewrite sample_diracR.
- under eq_sample do rewrite sample_diracL /=.
  by rewrite sample_diracR.
by move=> xx /supp_sampleP [] x x_p /supp_diracP ->.
Qed.

Lemma coupling_eq (T : ordType) (p q : {prob T}) : coupling eq p q -> p = q.
Proof.
by case=> pq -> -> {p q} e; apply/eq_in_sample; case=> x y /e /= ->.
Qed.

Lemma coupling_sample_same (T S1 S2 : ordType) (R : S1 -> S2 -> Prop) (p : {prob T}) f g :
  (forall x, x \in supp p -> coupling R (f x) (g x)) ->
  coupling R (sample p f) (sample p g).
Proof.
move=> e; apply: coupling_sample; first exact: coupling_same.
move=> x _ x_p _ <-; exact: e.
Qed.

Lemma coupling_trivial (T S : ordType) (R : T -> S -> Prop) (pT : {prob T}) (pS : {prob S}) :
  (forall x y, x \in supp pT -> y \in supp pS -> R x y) ->
  @coupling T S R pT pS.
Proof.
move=> RP.
exists (sample: x <- pT; sample: y <- pS; dirac (x, y)).
- rewrite sampleA; under eq_sample => x.
    rewrite sampleA; under eq_sample do rewrite sample_diracL.
    rewrite /= sample_const; over.
  by rewrite sample_diracR.
- rewrite sampleA; under eq_sample=> x.
    rewrite sampleA; under eq_sample do rewrite sample_diracL.
    rewrite /= sample_diracR; over.
  by rewrite sample_const.
case=> x y /supp_sampleP [] x' xP /supp_sampleP [] y' yP /=.
move=> /supp_diracP [-> ->]; exact: RP.
Qed.

Lemma couplingT T S pT pS : @coupling T S (fun _ _ => True) pT pS.
Proof. by apply: coupling_trivial. Qed.

Lemma coupling_sampleL (T S1 S2 : ordType) (R : S1 -> S2 -> Prop) (pT : {prob T}) f pS :
  (forall x, x \in supp pT -> coupling R (f x) pS) ->
  coupling R (sample: x <- pT; f x) pS.
Proof.
move=> RP; rewrite -[pS](sample_const pS).
apply: coupling_sample (couplingT _ _) _ => x _ xP _ _.
exact: RP.
Qed.

Lemma coupling_sampleR (T S1 S2 : ordType) (R : S1 -> S2 -> Prop) (pT : {prob T}) f pS :
  (forall x, x \in supp pT -> coupling R pS (f x)) ->
  coupling R pS (sample: x <- pT; f x).
Proof.
move=> RP; rewrite -[pS](sample_const pS).
apply: coupling_sample (couplingT _ _) _ => _ x _ xP _.
exact: RP.
Qed.

Lemma couplingW (T S : ordType) (R1 R2 : T -> S -> Prop) pT pS :
  (forall x y, R1 x y -> R2 x y) ->
  @coupling T S R1 pT pS ->
  coupling R2 pT pS.
Proof. by move=> R12 [p eT eS R1P]; exists p; eauto. Qed.

Lemma in_couplingW (T S : ordType) (R1 R2 : T -> S -> Prop) (pT : {prob T}) (pS : {prob S}) :
  (forall x y, x \in supp pT -> y \in supp pS -> R1 x y -> R2 x y) ->
  coupling R1 pT pS ->
  coupling R2 pT pS.
Proof.
move=> R12 [p eT eS R1P]; exists p; eauto.
case=> x y xyP; apply: R12; last exact: R1P.
- rewrite eT; apply/supp_sampleP.
  by exists (x, y)=> //; rewrite // supp_dirac in_fset1.
- rewrite eS; apply/supp_sampleP.
  by exists (x, y)=> //; rewrite // supp_dirac in_fset1.
Qed.

(** The following definitions allow us to apply probabilistic computations to
the elements of sequences, maps and other container types.  They will be used to
define the semantics of Imp programs. *)

Definition foldrM T (S : ordType) (f : T -> S -> {prob S}) (y : S) (xs : seq T) : {prob S} :=
  foldr (fun x p => sample p (f x)) (dirac y) xs.

Fixpoint map_p T (S : ordType) (f : T -> {prob S}) (xs : seq T) : {prob seq S} :=
  match xs with
  | [::] => dirac [::]
  | x :: xs =>
    sample: y  <- f x;
    sample: ys <- map_p f xs;
    dirac (y :: ys)
  end.

Lemma eq_map_p T (S : ordType) (f g : T -> {prob S}) :
  f =1 g -> map_p f =1 map_p g.
Proof. by move=> fg; elim=> //= x xs IH; rewrite fg IH. Qed.

Lemma map_p_dirac (T : ordType) (xs : seq T) : map_p dirac xs = dirac xs.
Proof.
elim: xs=> //= x xs IH.
by rewrite sample_diracL IH sample_diracL.
Qed.

Lemma map_p_comp T S (R : ordType) (f : T -> S) (g : S -> {prob R}) xs :
  map_p g [seq f x | x <- xs] = map_p (g \o f) xs.
Proof. by elim: xs=> //= x xs ->. Qed.

Lemma map_p_sample (T S R : ordType) (g : S -> {prob R}) (f : T -> {prob S}) (xs : seq T) :
  map_p (fun x => sample: y <- f x; g y) xs =
  sample: ys <- map_p f xs; map_p g ys.
Proof.
elim: xs=> [|x xs IH] /=; first by rewrite sample_diracL.
rewrite !sampleA; apply/eq_sample=> y.
rewrite sampleA {}IH.
under eq_sample=> z do rewrite sampleA.
under [in RHS]eq_sample=> zs do rewrite sample_diracL /=.
by rewrite sampleC.
Qed.

Lemma supp_map_p T (S : ordType) (f : T -> {prob S}) xs ys :
  ys \in supp (map_p f xs) =
  all2 (fun x y => y \in supp (f x)) xs ys.
Proof.
elim: xs ys=> [|x xs IH] [|y ys] /=.
- by rewrite supp_dirac.
- by rewrite supp_dirac.
- case: supp_sampleP=> //=.
  by case=> y' y'P /supp_sampleP [ys' _ /supp_diracP].
- rewrite -IH; apply/(sameP supp_sampleP)/(iffP andP).
  + case=> [yP ysP]; exists y=> //.
    apply/supp_sampleP; exists ys=> //.
    by apply/supp_diracP.
  + by case=> [y' y'P /supp_sampleP [ys' ys'P /supp_diracP [-> ->]]].
Qed.

Section MapMapProb.

Variable T : ordType.

Definition mapim_p (S : Type) (R : ordType)
  (f : T -> S -> {prob R}) (m : {fmap T -> S}) : {prob {fmap T -> R}} :=
  let do_pair p := sample: y <- f p.1 p.2; dirac (p.1, y) in
  sample: pairs <- map_p do_pair (val m);
  dirac (mkfmap pairs).

Lemma mapim_pE S R f m :
  @mapim_p S R f m =
  foldrM (fun p m => sample: z <- f p.1 p.2; dirac (setm m p.1 z)) emptym m.
Proof.
rewrite /mapim_p /=.
elim: (val m)=> {m} [|[x y] m /= IH] //=.
- by rewrite sample_diracL //.
- rewrite !sampleA [in RHS]sampleC; apply/eq_sample=> z.
  rewrite sample_diracL sampleA -IH sampleA.
  by apply/eq_sample=> pairs; rewrite !sample_diracL.
Qed.

Lemma eq_mapim_p (S : Type) (R : ordType)
  (f g : T -> S -> {prob R}) :
  f =2 g -> mapim_p f =1 mapim_p g.
Proof.
move=> fg m; rewrite /mapim_p.
by under eq_map_p => p do rewrite fg.
Qed.

Lemma mapim_p_dirac (S : ordType) (m : {fmap T -> S}) :
  mapim_p (fun _ => dirac) m = dirac m.
Proof.
rewrite /mapim_p.
under eq_map_p => p do rewrite sample_diracL -surjective_pairing.
by rewrite map_p_dirac sample_diracL fmvalK.
Qed.

Lemma mapim_p_comp (S R U : ordType)
  (g : T -> R -> {prob U}) (f : T -> S -> R) m :
  mapim_p g (mapim f m) =
  mapim_p (fun x y => g x (f x y)) m.
Proof. by rewrite /mapim_p /= map_p_comp. Qed.

Fact mapm_p_key : unit. Proof. exact: tt. Qed.

Definition mapm_p (S : Type) (R : ordType) (f : S -> {prob R}) :=
  locked_with mapm_p_key (mapim_p (fun (x : T) => f)).

Lemma mapm_pE S R f m :
  @mapm_p S R f m =
  if splitm m is Some (x, y, m) then
    sample: z <- f y;
    sample: m <- mapm_p f m;
    dirac (setm m x z)
  else dirac emptym.
Proof.
rewrite /splitm /= /mapm_p unlock /mapim_p /=.
move: (valP m) => /=; case: (val m)=> [|[x y] ps] //=.
  by rewrite sample_diracL.
move=> /path_sorted psP.
rewrite !sampleA; apply/eq_sample=> z.
rewrite sample_diracL !sampleA mkfmapK //.
by apply/eq_sample=> /= ps'; rewrite !sample_diracL.
Qed.

Lemma eq_mapm_p (S : Type) (R : ordType) (f g : S -> {prob R}) :
  f =1 g -> mapm_p f =1 mapm_p g.
Proof.
rewrite /mapm_p !unlock.
by move=> e; apply/eq_mapim_p=> ??; eauto.
Qed.

Lemma mapm_p_dirac (S : ordType) (m : {fmap T -> S}) :
  mapm_p dirac m = dirac m.
Proof. rewrite /mapm_p !unlock; exact/mapim_p_dirac. Qed.

Lemma mapm_p_comp (S R U : ordType) (g : R -> {prob U}) (f : S -> R) m :
  mapm_p g (mapm f m) = mapm_p (g \o f) m.
Proof. rewrite /mapm_p !unlock; exact/mapim_p_comp. Qed.

Lemma supp_mapm_p (S R : ordType) (f : S -> {prob R}) m1 m2 :
  m2 \in supp (mapm_p f m1) =
  (domm m1 == domm m2) &&
  all (fun x => match m1 x, m2 x with
                | Some y, Some z => z \in supp (f y)
                | _, _ => true
                end) (domm m1).
Proof.
rewrite /mapm_p unlock /mapim_p.
apply/(sameP supp_sampleP)/(iffP andP).
- case=> /eqP edomm ecodomm; exists (val m2); last first.
    by apply/supp_diracP; rewrite fmvalK.
  have /= esize: size (val m1) = size (val m2).
    move/(congr1 (size \o val)): edomm.
    by rewrite /= !val_domm !size_map => ->.
  rewrite supp_map_p all2E esize eqxx /=.
  apply/allP=> /= - xyz /(nthP xyz) [i].
  rewrite size_zip -esize minnn=> isize.
  case: xyz=> [[x1 y] [x2 z]]; rewrite nth_zip //=; case=> e1 e2.
  have e1' : nth x1 (domm m1) i = x1.
    by rewrite val_domm (nth_map (x1, y)) // e1.
  have e2' : nth x2 (domm m2) i = x2.
    by rewrite val_domm (nth_map (x2, z)) -?esize // e2.
  have ex: x1 = x2.
    rewrite -e1' -e2' -edomm; apply/set_nth_default.
    by rewrite val_domm size_map.
  move: x1 ex e1 e2 {e1' e2'}=> {x2} x <- e1 e2.
  have {}e1: m1 x = Some y.
    by apply/getmP; rewrite -e1; apply/mem_nth.
  have {}e2: m2 x = Some z.
    by apply/getmP; rewrite -e2; apply/mem_nth; rewrite -esize.
  have xP: x \in domm m1 by rewrite mem_domm e1.
  move/allP/(_ _ xP): ecodomm; rewrite e1 e2=> yz.
  by apply/supp_sampleP; exists z; rewrite // supp_dirac in_fset1 eqxx.
- case=> {}m2 m2P /supp_diracP ->.
  move: m2P; rewrite supp_map_p all2E; case/andP=> /eqP esize ecodomm.
  have edomm : unzip1 m1 = unzip1 m2.
    move: m2 esize ecodomm; rewrite /=.
    elim: {m1} (val m1)=> [|[x1 y] m1 IH] [|[x2 z] m2] //= [esize].
    case/andP=> /supp_sampleP [{}z _ /supp_diracP [-> _]] ecodomm.
    congr cons; exact: IH.
  have m2_sorted: mkfmap m2 = m2 :> seq _.
    apply/mkfmapK; rewrite -edomm; exact: (valP m1).
  have {}edomm : domm m1 = domm (mkfmap m2).
    by apply/val_inj; rewrite /= !val_domm m2_sorted.
  split; first by rewrite edomm.
  move: esize ecodomm; rewrite -m2_sorted.
  move: (mkfmap m2) edomm=> {m2_sorted} {}m2 /= edomm esize ecodomm.
  apply/allP=> x /dommP [y yP]; rewrite yP fmvalK.
  have /dommP [z zP]: x \in domm m2 by rewrite -edomm mem_domm yP.
  rewrite zP.
  case/getmP/(nthP (x, y)): (yP)=> /= i isize ei1.
  have ei1domm : nth x (domm m1) i = x.
    by rewrite val_domm (nth_map (x, y)) // ei1.
  have ei2: nth (x, z) m2 i = (x, z).
    move: zP; rewrite -{1} ei1domm edomm.
    rewrite (getm_nth (x, z)) -?esize //; case=> {2}<-.
    move: ei1domm; rewrite edomm val_domm (nth_map (x, z)) -?esize //.
    by move=> e; rewrite -[X in _ = (X, _)]e [LHS]surjective_pairing.
  have inzip : nth ((x, y), (x, z)) (zip m1 m2) i = ((x, y), (x, z)).
    by rewrite nth_zip // ei1 ei2.
  have {}inzip: ((x, y), (x, z)) \in zip m1 m2.
    by rewrite -inzip; apply/mem_nth; rewrite size_zip -esize minnn.
  move/allP/(_ _ inzip): ecodomm; rewrite inE /=.
  by case/supp_sampleP=> ?? /supp_diracP [->].
Qed.

Lemma supp_mapm_pP (S R : ordType) (f : S -> {prob R}) m1 m2 :
  reflect (domm m1 = domm m2 /\
           forall x y z, m1 x = Some y -> m2 x = Some z -> z \in supp (f y))
          (m2 \in supp (mapm_p f m1)).
Proof.
rewrite supp_mapm_p; apply/(iffP andP).
- case=> /eqP edomm ecodomm; split=> // x y z yP zP.
  have xP: x \in domm m1 by rewrite mem_domm yP.
  by move/allP/(_ _ xP): ecodomm; rewrite yP zP.
- case=> edomm ecodomm; split; first by rewrite edomm.
  apply/allP=> x xP; case/dommP: (xP)=> y yP.
  move: (xP); rewrite edomm; case/dommP=> z zP.
  rewrite yP zP; apply: ecodomm yP zP.
Qed.

Lemma mapm_p0 (S R : ordType) f : @mapm_p S R f emptym = dirac emptym.
Proof. by rewrite mapm_pE /=. Qed.

Lemma mapm_p_setm (S R : ordType) f m x y :
  @mapm_p S R f (setm m x y) =
  sample: m <- mapm_p f m;
  sample: z  <- f y;
  dirac (setm m x z).
Proof.
rewrite /mapm_p unlock !mapim_pE [setm m x y]/setm /=.
case: m=> /=; elim=> [|[x' y'] m IH] //=.
case: Ord.ltgtP=> [//|xx'|{x'} <-] //= mP.
  rewrite IH ?(path_sorted mP) //.
  rewrite !sampleA; apply/eq_sample=> /= m'.
  rewrite sampleC sampleA; apply/eq_sample=> /= z.
  rewrite sampleA sample_diracL; apply/eq_sample=> /= z'.
  rewrite sample_diracL setmC // eq_sym.
  by case: Ord.ltgtP xx'.
rewrite sampleA; apply/eq_sample=> m'.
rewrite sampleC; apply/eq_sample=> z.
rewrite sampleA.
under eq_sample=> ? do rewrite sample_diracL setmxx.
by rewrite sample_const.
Qed.

(*
Lemma mapm_pE' (S R : ordType) f my mz :
  @mapm_p S R f my mz =
  if domm my == domm mz then
    \prod_(x <- domm my)
      (match my x, mz x with
       | Some y, Some z => f y z
       | _, _ => 0
       end)
  else 0.
Proof.
case: eqP=>[e|ne]; last first.
  apply/suppPn/supp_mapm_pP; case=> ??; congruence.
move ez: (domm mz) e=> X ey.
elim/fset_ind: X my mz ey ez=> [|x X x_X IH] my mz ey ez.
  move: ey ez=> /eqP/emptymP -> /eqP/emptymP ->.
  by rewrite domm0 big_nil mapm_pE /= diracE.
rewrite ey big_fsetU1 //=.
have /dommP [y yP]: x \in domm my by rewrite ey in_fsetU1 eqxx.
have /dommP [z zP]: x \in domm mz by rewrite ez in_fsetU1 eqxx.
set my' := remm my x; set mz' := remm mz x.
have e: (x |: X) :\ x = X.
  apply/eq_fset=> x'; rewrite in_fsetD1 in_fsetU1.
  by case: eqP=> [{x'}->|] //=; rewrite (negbTE x_X).
have eyD: domm my' = X by rewrite domm_rem ey.
have ezD: domm mz' = X by rewrite domm_rem ez.
have {}ey: my = setm my' x y.
  apply/eq_fmap=> x'; rewrite setmE remmE.
  by case: eqP yP=> [->|].
have {}ez: mz = setm mz' x z.
  apply/eq_fmap=> x'; rewrite setmE remmE.
  by case: eqP zP=> [->|].
rewrite yP zP ey ez mapm_p_setm sampleE big_seq.
under eq_big=> [m|m mP]; first over.
  rewrite sampleE.
  under eq_big=> [z'|z' z'P]; first over.
    rewrite diracE eq_setm remmI ?ezD // remmI; last first.
      by case/supp_mapm_pP: mP=> <- _; rewrite eyD.
    over.
  rewrite /=.
rewrite (IH _ _ ey ez).

have e: perm_eq (x |: X) (x :: X).
  apply: uniq_perm.
  - exact: uniq_fset.
  - by rewrite /= x_X uniq_fset.
  by move=> x'; rewrite inE in_fsetU1.
rewrite ey (perm_big _ e) /= big_cons.

move e: (size my)=> n; elim: n my mz e=> [|n IH] /= my mz.
  rewrite sizeES mapm_pE [domm my]dommES.
  case e: (splitm my)=> [[[x y] my']|] //= _.
  by rewrite diracE -eq_domm0 eq_sym big_nil.
rewrite sizeES mapm_pE [domm my]dommES.
case ey: (splitm my)=> [[[x y] my']|] //= [e].
rewrite sampleE [domm mz]dommES.
case ez: (splitm mz)=> [[[x' z] mz']|]; last first.
  rewrite eq_domm0.

rewrite mapm_pE /= -[domm _ == domm _]val_eqE /= val_domm.
(* FIXME: why doesn't under work here? *)
have e: forall x, true -> match my x, mz x with
                  | Some y, Some z => f y z
                  | _ , _ => 0
                  end =
                  match getm_def (val my) x, mz x with
                  | Some y, Some z => f y z
                  | _, _ => 0
                  end.
  by move=> x; rewrite -mkfmapE fmvalK.
rewrite (eq_big xpredT _ _ e) //= {e}.
elim: (val my) mz=> /= {my} [|[x y] my IH] /= mz.
  by rewrite big_nil diracE -eq_domm0 -val_eqE /= eq_sym.
rewrite sampleE.
*)

Lemma mapm_p_sample (S R U : ordType) (g : R -> {prob U}) (f : S -> {prob R}) (m : {fmap T -> S}) :
  mapm_p (fun x => sample: y <- f x; g y) m =
  sample: m' <- mapm_p f m; mapm_p g m'.
Proof.
elim/fmap_ind: m=> [|m IH x y x_m]; first by rewrite !(mapm_p0, sample_diracL).
rewrite !mapm_p_setm IH // !sampleA; apply/eq_sample=> m'.
rewrite sampleA; under [in LHS]eq_sample=> ? do rewrite sampleA.
rewrite sampleC; apply/eq_sample=> z.
by rewrite sample_diracL mapm_p_setm sampleC.
Qed.

End MapMapProb.

Section SampleFun.

Variables (T : ordType) (S : ordType) (def : T -> S).

Definition samplef (f : ffun (dirac \o def)) : {prob ffun def} :=
  sample: m <- mapm_p id (val f);
  dirac (mkffun (fun x => odflt (def x) (m x)) (domm m)).

Lemma samplef0 : samplef emptyf = dirac emptyf.
Proof.
rewrite /samplef /= mapm_p0 sample_diracL; congr dirac.
by apply/eq_ffun=> x; rewrite mkffunE /= domm0 /= emptyfE.
Qed.

Lemma samplef_map f : samplef (mapf dirac f) = dirac f.
Proof.
rewrite /samplef /= val_mapf; last exact: dirac_inj.
rewrite mapm_p_comp mapm_p_dirac sample_diracL; congr dirac.
apply/eq_ffun=> x; rewrite mkffunE /appf /= mem_domm.
by case: (ffval f x).
Qed.

End SampleFun.

Section MapFunProb.

Variables (T S R : ordType) (def : T -> S).

Definition mapf_p (g : S -> R) (f : ffun (dirac \o def)) : ffun (dirac \o g \o def) :=
  mkffun (fun x => sample: y <- f x; dirac (g y)) (supp f).

Lemma mapf_pE g f x :
  mapf_p g f x = sample: y <- f x; dirac (g y).
Proof.
rewrite /mapf_p mkffunE mem_supp /=.
by case: eqP=> //= ->; rewrite sample_diracL.
Qed.

End MapFunProb.

Lemma coupling_mapm_p (T S11 S12 S21 S22 : ordType)
    (R1 : T -> S11 -> S12 -> Prop) (R2 : T -> S21 -> S22 -> Prop)
    (f1 : S11 -> {prob S21}) (f2 : S12 -> {prob S22}) :
  (forall x y1 y2, R1 x y1 y2 -> coupling (R2 x) (f1 y1) (f2 y2)) ->
  forall (m1 : {fmap _}) (m2 : {fmap _}),
    (forall x y1 y2, m1 x = Some y1 -> m2 x = Some y2 -> R1 x y1 y2) ->
    coupling (fun (m1' : {fmap _}) (m2' : {fmap _}) =>
                forall x y1 y2, m1' x = Some y1 -> m2' x = Some y2 -> R2 x y1 y2)
             (mapm_p f1 m1) (mapm_p f2 m2).
Proof.
move=> RP; elim/fmap_rect.
  rewrite mapm_p0; move=> m2 em; apply: coupling_trivial.
  move=> _ m2' /supp_diracP -> //.
move=> m1 IH x y1 fresh m2 em.
have y1P y2 : m2 x = Some y2 -> R1 x y1 y2.
  by apply: em; rewrite setmE eqxx.
have /IH {}em x' y1' y2 : m1 x' = Some y1' -> m2 x' = Some y2 -> R1 x' y1' y2.
  move=> m1x' m2x'; apply: em; rewrite // setmE.
  by case: eqP m1x'=> // ->; rewrite (dommPn fresh).
case m2x: (m2 x)=> [y2|].
  rewrite -(setmI m2x) !mapm_p_setm; apply: coupling_sample em _.
  move=> /= m1' m2' _ _ em'; apply: coupling_sample.
    apply: RP; exact: y1P.
  move=> y1' y2' _ _ yP; apply: coupling_dirac.
  move=> x' y1'' y2''; rewrite !setmE; case: eqP=> [->|_].
    by move=> [<-] [<-].
  exact: em'.
rewrite mapm_p_setm -[mapm_p _ m2]sample_diracR.
apply: coupling_sample em _ => m1' m2' _ /supp_mapm_pP [ed2 _].
move=> {}em; apply: coupling_sampleL=> y1' _; apply: coupling_dirac.
move=> x' y1'' y2''; rewrite setmE; case: eqP=> [->|_]; last exact: em.
by have /dommPn -> : x \notin domm m2' by rewrite -ed2 mem_domm m2x.
Qed.

Lemma coupling_mapm_p_eq (T S1 S2 : ordType) (f : S1 -> {prob S2}) (m1 m2 : {fmap T -> S1}) (P : T -> bool) :
  {in P, m1 =1 m2} ->
  coupling (fun m1' m2' : {fmap T -> S2} => {in P, m1' =1 m2'}) (mapm_p f m1) (mapm_p f m2).
Proof.
move=> em.
pose R1 x (y1 y2 : S1) := P x -> y1 = y2.
pose R2 x (y1' y2' : S2) := P x -> y1' = y2'.
have emW x y1 y2 : m1 x = Some y1 -> m2 x = Some y2 -> R1 x y1 y2.
  move=> m1x m2x Px; apply: Some_inj; rewrite -m1x -m2x; exact: em.
have RP x y1 y2 : R1 x y1 y2 -> coupling (R2 x) (f y1) (f y2).
  move=> xP; case: (boolP (P x))=> [{}/xP ->|nPx].
    apply: couplingW; last exact: coupling_same.
    by move=> ?? ->.
  apply: coupling_trivial=> ?? _ _ contra; by rewrite contra in nPx.
apply: in_couplingW; last by apply: coupling_mapm_p emW; eauto.
move=> /= m1' m2' /supp_mapm_pP [ed1 _] /supp_mapm_pP [ed2 _].
move=> em' x Px.
case m1'x: (m1' x)=> [y1'|].
  suff /dommP [y2' m2'x] : x \in domm m2'.
    by rewrite m2'x (em' _ _ _ m1'x m2'x) // mem_domm.
  by rewrite -ed2 mem_domm -em // -mem_domm ed1 mem_domm m1'x.
by apply/esym/dommPn; rewrite -ed2 mem_domm -em // -mem_domm ed1 mem_domm m1'x.
Qed.


Module Base.
Inductive t :=
| A
| B.

Definition f0 x :=
  match x with
  | A => A
  | B => B
  end.

Print f0.

Definition f0bis x y :=
  match x, y with
  | A, A => A
  | _, _ => B
  end.

Print f0bis.
End Base.

Definition f1 n :=
  match n with
  | O => Base.A
  | S _ => Base.B
  end.

Print f1.

Definition f2 n :=
  match n with
  | O => Base.A
  | S (S n) => Base.B
  | _ => Base.A
  end.

Print f2.

Definition pred n :=
  match n with
  | O => O
  | S O => O
  | S (S n) => S n
  end.

Print pred.

Fixpoint add n m :=
  match n, m with
  | O, _ => m
  | _, O => n
  | S p, S q => S (S (add p q))
  end.

Print add.

Theorem add_0_l n : add O n = n.
Proof.
reflexivity.
Qed.

Inductive vector (A : Type) : nat -> Type :=
| Nil : vector A O
| Cons n (_ : A) (_ : vector A n) : vector A (S n).

Definition rectS {A} (P:forall n, vector A (S n) -> Type)
 (bas: forall a: A, P _ (Cons _ _ a (Nil _)))
 (rect: forall a {n} (v: vector A (S n)), P _ v -> P _ (Cons _ _ a v)) :=
 fix rectS_fix {n} (v: vector A (S n)) : P _ v :=
 match v with
 |Cons _ 0 a v =>
   match v with
     |Nil _ => bas a
     |_ => fun devil => False_ind (@IDProp) devil (* subterm !!! *)
   end
 |Cons _ (S nn') a v => rect a v (rectS_fix v)
 |_ => fun devil => False_ind (@IDProp) devil (* subterm !!! *)
 end.

Definition head' (A : Type) (n : nat) (v : vector A (S n)) :=
  match v in vector _ n0 return match n0 return Type with
  | S m => A
  | O => IDProp
  end with
  | Cons _ _ hd _tl => hd
  | Nil _ => idProp
  end.

Print head'.

Definition head (A : Type) (n : nat) (v : vector A (S n)) :=
  match v with
  | Cons _ _ hd _tl => hd
  end.

Print head.

Definition tail' (A : Type) (n : nat) (v : vector A (S n)) : vector A n :=
  match v in vector _ n0 return match n0 return Type with
  | S m => vector A m
  | O => IDProp
  end with
  | Cons _ _ _hd tl => tl
  | Nil _ => idProp
  end.

Print tail'.

Definition tail (A : Type) (n : nat) (v : vector A (S n)) : vector A n (*bug!*) :=
  match v with
  | Cons _ _ _hd tl => tl
  end.

Print tail.

Definition tail2' (A : Type) (n : nat) (v : vector A (S (S n))) :=
   match
     v as v0 in (vector _ H)
     return
       (match H as H0 return Type with
        | O => IDProp
        | S x =>
            match x as H1 return Type with
            | O => IDProp
            | S x0 => vector A (S x0)
            end
        end)
   with
   | Nil _ => idProp
   | Cons _ x _hd tl =>
       match x as H, tl return match H return Type with
         | O => IDProp
         | S n => vector A (S n)
       end with
       | O, _ => idProp
       | S x0, tl => tl
       end
   end.

Print tail2'.

Definition tail2 (A : Type) (n : nat) (v : vector A (S (S n))) :=
  match v with
  | Cons _ _ _hd tl => tl
  end.

Print tail2.

Definition snd (A : Type) (n : nat) (v : vector A (S (S n))) :=
  match v with
  | Cons _ _ _ (Cons _ _ hd _) => hd
  end.

Print snd.

Definition third (A : Type) (n : nat) (v : vector A (S (S (S n)))) :=
  match v with
  | Cons _ _ _ (Cons _ _ _ (Cons _ _ hd _)) => hd
  end.

Print third.

Definition forth (A : Type) (n : nat) (v : vector A (S (S (S (S n))))) :=
  match v with
  | Cons _ _ _ (Cons _ _ _ (Cons _ _ _ (Cons _ _ hd _))) => hd
  end.

Print forth.

Definition fifth (A : Type) (n : nat) (v : vector A (S (S (S (S (S n)))))) :=
  match v with
  | Cons _ _ _ (Cons _ _ _ (Cons _ _ _ (Cons _ _ _ (Cons _ _ hd _)))) => hd
  end.

Print fifth.

Fixpoint map A B (n : nat) (f : A -> B) (v : vector A n) : vector B n :=
  match v with
  | Nil _ => Nil B
  | Cons _ _ hd tl => Cons _ _ (f hd) (map _ _ _ f tl)
  end.

Print map.

(*
Fixpoint mapn A B (n : nat) (f : A -> B) (v : vector A n) : vector B n :=
  match n, v with
  | O, Nil _ => Nil B
  | S m, Cons _ _ hd tl => Cons _ _ (f hd) (mapn _ _ m f tl)
  end.

Print mapn.
*)

Fixpoint map2' A0 A1 B (n : nat) (f : A0 -> A1 -> B) (v0 : vector A0 n) (v1 : vector A1 n) :=
  match v0 in vector _ n0 return
    vector A1 n0 -> vector B n0
 with
  | Nil _ => fun v1 : vector A1 O =>
     match v1 in vector _ n0 return
       match n0 return Type with
       | O => vector B O
       | S _ => IDProp
       end with
     | Nil _ => Nil B
     | Cons _ _ _ _ => idProp
     end
  | Cons _ n0 hd0 tl0 => fun v1 : vector A1 (S n0) =>
     match v1 in vector _ n1 return
       match n1 return Type with
       | O => IDProp
       | S n2 => vector A0 n2 -> vector B (S n2)
     end with
     | Nil _ => idProp
     | Cons _ n1 hd1 tl1 => fun tl0 : vector A0 n1 =>
         Cons _ n1 (*bug!*) (f hd0 hd1) (map2' A0 A1 B _ f tl0 tl1)
     end tl0
  end v1.

Print map2'.

Fixpoint map2 A0 A1 B (n : nat) (f : A0 -> A1 -> B) (v0 : vector A0 n) (v1 : vector A1 n) : vector B n :=
  match v0, v1 with
  | Nil _, Nil _ => Nil _
  | Cons _ _ hd0 tl0, Cons _ _ hd1 tl1 =>
    Cons _ _ (f hd0 hd1) (map2 _ _ _ _ f tl0 tl1)
  end.

Print map2.

Module M0.
Inductive t : nat -> Type :=
| A : t (S O).

Definition test := fun n (v : t n) =>
match n, v with
| S n, A => tt
end.
End M0.

Module M1.
Inductive t : nat -> Type := C : t 0 | D n : t n -> t n.

Definition test (v : t 0) : t 0 :=
match v with
| C => C
| D n x => D 0 x
end.

Print test.
End M1.

(*
Fixpoint map3 A0 A1 A2 B (n : nat) (f : A0 -> A1 -> A2 -> B) (v0 : vector A0 n) (v1 : vector A1 n) (v2 : vector A2 n) : vector B n :=
  match v0, v1, v2 with
  | Nil _, Nil _, Nil _ => Nil _
  | Cons _ _ hd0 tl0, Cons _ _ hd1 tl1, Cons _ _ hd2 tl2 =>
    Cons _ _ (f hd0 hd1 hd2) (map3 _ _ _ _ _ f tl0 tl1 tl2)
  end.

Print map3.
*)
(* from Logic.v *)

Inductive ex (A:Type) (P:A -> Prop) : Prop :=
  ex_intro : forall x:A, P x -> ex A P.

Section Projections.

  Variables (A:Prop) (P:A->Prop).

  Definition ex_proj1 (x:ex A P) : A :=
    match x with ex_intro _ _ a _ => a end.

  Print ex_proj1.

  Definition ex_proj2 (x:ex A P) : P (ex_proj1 x) :=
    match x with ex_intro _ _ _ b => b end.

End Projections.

Inductive bool : Set :=
  | true : bool
  | false : bool.

Definition andb (b1 b2:bool) : bool := if b1 then b2 else false.

Require Import Lists.List.

Import ListNotations.

  Fixpoint last (A: Type) (l:list A) (d:A) : A :=
  match l with
    | [] => d
    | [a] => a
    | a :: l => last A l d
  end.

Section Equivalences.

  Variable U:Type.

  Definition UIP_refl_on_ (x : U) :=
    forall (p : x = x), p = eq_refl x.

End Equivalences.

Theorem UIP_shift_on (X : Type) (x : X) :
  UIP_refl_on_ X x -> forall y : x = x, UIP_refl_on_ (x = x) y.
Proof.
  intros UIP_refl y.
  rewrite (UIP_refl y).
  intros z.
  assert (UIP:forall y' y'' : x = x, y' = y'').
  { intros. apply eq_trans_r with (eq_refl x); apply UIP_refl. }
  transitivity (eq_trans (eq_trans (UIP (eq_refl x) (eq_refl x)) z)
                         (eq_sym (UIP (eq_refl x) (eq_refl x)))).
  - destruct z. destruct (UIP _ _). reflexivity.
  - change
      (match eq_refl x as y' in _ = x' return y' = y' -> Prop with
       | eq_refl => fun z => z = (eq_refl (eq_refl x))
       end (eq_trans (eq_trans (UIP (eq_refl x) (eq_refl x)) z)
                     (eq_sym (UIP (eq_refl x) (eq_refl x))))).
    destruct z. destruct (UIP _ _). reflexivity.
Qed.

Lemma UIP_refl_bool (b:bool) (x : b = b) : x = eq_refl.
Proof.
  destruct b.
  - change (match true as b return true=b -> Prop with
            | true => fun x => x = eq_refl
            | _ => fun _ => True
            end x).
    destruct x; reflexivity.
  - change (match false as b return false=b -> Prop with
            | false => fun x => x = eq_refl
            | _ => fun _ => True
            end x).
    destruct x; reflexivity.
Defined.

Module Fin.
Inductive t : nat -> Set :=
|F1 : forall {n}, t (S n)
|FS : forall {n}, t n -> t (S n).

Definition case0 P (p: t O): P p :=
  match p with | F1 | FS  _ => fun devil => False_rect (@IDProp) devil (* subterm !!! *) end.

Definition caseS' {n : nat} (p : t (S n)) : forall (P : t (S n) -> Type)
  (P1 : P F1) (PS : forall (p : t n), P (FS p)), P p :=
  match p with
  | @F1 k => fun P P1 PS => P1
  | FS pp => fun P P1 PS => PS pp
  end.

Fixpoint weak {m}{n} p (f : t m -> t n) :
  t (p + m) -> t (p + n) :=
match p as p' return t (p' + m) -> t (p' + n) with
  |0 => f
  |S p' => fun x => match x with
     |@F1 n' => fun eq : n' = p' + m => F1
     |@FS n' y => fun eq : n' = p' + m => FS (weak p' f (eq_rect _ t y _ eq))
  end (eq_refl _)
end.

End Fin.


Require Import Coq.Init.Prelude.

Fixpoint fnd (n : nat) : Prop :=
  match n with
  | O => True
  | S n => True /\ fnd n
  end.

Fixpoint big_and (xs : list Prop) : Prop :=
  match xs with
  | nil => True
  | cons x xs => and x (big_and xs)
  end.
Fixpoint typeof_big_conj (xs : list Prop) (P : Prop) : Prop :=
  match xs with
  | nil => P
  | cons x xs => x -> typeof_big_conj xs P
  end.
Lemma apply_rconj_and xs (P Q : Prop) : P -> typeof_big_conj xs Q -> typeof_big_conj xs (P /\ Q).
Proof. revert Q; revert P; induction xs; intros; cbn in *; eauto. Qed.
Lemma big_conj xs : typeof_big_conj xs (big_and xs).
Proof. induction xs. exact I. cbn. intros. apply apply_rconj_and; eauto. Qed.

Require Import Coq.Lists.List.

Require Import Ltac2.Ltac2.
Require Import Ltac2.Control.
Require Import Ltac2.Std.
Require Import Ltac2.Notations.

Ltac2 get_constant (c : constr)
  := match Constr.Unsafe.kind c with
     | Constr.Unsafe.Constant c _ => c
     | _ => throw Match_failure
     end.

Ltac2 Notation "eval" "cbv" s(strategy) "in" v(constr) :=
  Std.eval_cbv s v.

Goal fnd 1000.
  Time
    let n := match! goal with [ |- fnd ?n ] => n end in
    let ls := constr:(repeat True $n) in
    let ls := (eval cbv in $ls) in
    let pf := constr:(big_conj $ls) in
    let t := Constr.type pf in
    let t := (eval cbv [typeof_big_conj big_and] in $t) in
    (assert (H : $t) by refine pf); apply H; exact I.
Time Qed.
(*
Finished transaction in 0.281 secs (0.23u,0.049s) (successful)
Finished transaction in 0.044 secs (0.044u,0.s) (successful)
*)


(* generated by Ott 0.31, locally-nameless from: src/ott/DKDeclarative.ott *)
Require Import Bool.
Require Import Metalib.Metatheory.
Require Import List.
(** syntax *)
Definition tpvar : Set := var. (*r variables *)

Inductive type : Set :=  (*r types *)
 | t_var_b (_:nat) (*r variable *)
 | t_var_f (x:tpvar) (*r variable *)
 | t_int : type (*r integer type *)
 | t_all (A:type) (*r polymorphic types *)
 | t_arrow (A:type) (B:type) (*r function types *).

Definition dkcontext : Set := (list tpvar).

(* EXPERIMENTAL *)
(** auxiliary functions on the new list types *)
(** library functions *)
(** subrules *)
(** arities *)
(** opening up abstractions *)
Fixpoint open_type_wrt_type_rec (k:nat) (A5:type) (A_6:type) {struct A_6}: type :=
  match A_6 with
  | (t_var_b nat) => if (k === nat) then A5 else (t_var_b nat)
  | (t_var_f x) => t_var_f x
  | t_int => t_int
  | (t_all A) => t_all (open_type_wrt_type_rec (S k) A5 A)
  | (t_arrow A B) => t_arrow (open_type_wrt_type_rec k A5 A) (open_type_wrt_type_rec k A5 B)
end.

Definition open_type_wrt_type A5 A_6 := open_type_wrt_type_rec 0 A_6 A5.

(** terms are locally-closed pre-terms *)
(** definitions *)

(* defns LC_type *)
Inductive lc_type : type -> Prop :=    (* defn lc_type *)
 | lc_t_var_f : forall (x:tpvar),
     (lc_type (t_var_f x))
 | lc_t_int :
     (lc_type t_int)
 | lc_t_all : forall (L:vars) (A:type),
      ( forall x , x \notin  L  -> lc_type  ( open_type_wrt_type A (t_var_f x) )  )  ->
     (lc_type (t_all A))
 | lc_t_arrow : forall (A B:type),
     (lc_type A) ->
     (lc_type B) ->
     (lc_type (t_arrow A B)).
(** free variables *)
Fixpoint fv_type (A5:type) : vars :=
  match A5 with
  | (t_var_b nat) => {}
  | (t_var_f x) => {{x}}
  | t_int => {}
  | (t_all A) => (fv_type A)
  | (t_arrow A B) => (fv_type A) \u (fv_type B)
end.

(** substitutions *)
Fixpoint subst_type (A5:type) (x5:tpvar) (A_6:type) {struct A_6} : type :=
  match A_6 with
  | (t_var_b nat) => t_var_b nat
  | (t_var_f x) => (if eq_var x x5 then A5 else (t_var_f x))
  | t_int => t_int
  | (t_all A) => t_all (subst_type A5 x5 A)
  | (t_arrow A B) => t_arrow (subst_type A5 x5 A) (subst_type A5 x5 B)
end.

Fixpoint ctx_dom (c : list tpvar) :=
  match c with
  | nil => empty
  | x :: c' => add x (ctx_dom c')
  end
.

(** definitions *)

(* defns Monotype *)
Inductive dk_mono : type -> Prop :=    (* defn dk_mono *)
 | dkm_var : forall (x:tpvar),
     dk_mono (t_var_f x)
 | dkm_int :
     dk_mono t_int
 | dkm_arrow : forall (A B:type),
     dk_mono A ->
     dk_mono B ->
     dk_mono (t_arrow A B).

(* defns WellformedContext *)
Inductive dk_wf_ctx : dkcontext -> Prop :=    (* defn dk_wf_ctx *)
 | wf_ctx_nil :
     dk_wf_ctx  nil
 | wf_ctx_cons : forall (G:dkcontext) (x:tpvar),
      dk_wf_ctx G  ->
      (not (In  x   G ))  ->
     dk_wf_ctx  ( x  ::  G ) .

(* defns WellformedType *)
Inductive dk_wf_tp : dkcontext -> type -> Prop :=    (* defn dk_wf_tp *)
 | wf_tp_var : forall (G:dkcontext) (x:tpvar),
      dk_wf_ctx G  ->
      (In  x   G )  ->
     dk_wf_tp G (t_var_f x)
 | wf_tp_int : forall (G:dkcontext),
     dk_wf_ctx G ->
     dk_wf_tp G t_int
 | wf_tp_arrow : forall (G:dkcontext) (A B:type),
      dk_wf_tp G A  ->
     dk_wf_tp G B ->
     dk_wf_tp G (t_arrow A B)
 | wf_tp_all : forall (L:vars) (G:dkcontext) (A:type),
      ( forall x , x \notin  L  -> dk_wf_tp  ( x  ::  G )   ( open_type_wrt_type A (t_var_f x) )  )  ->
     dk_wf_tp G (t_all A).

(* defns Subtyping *)
Inductive dk_sub : dkcontext -> type -> type -> Prop :=    (* defn dk_sub *)
 | dks_var : forall (G:dkcontext) (x:tpvar),
      dk_wf_ctx G  ->
      (In  x   G )  ->
     dk_sub G (t_var_f x) (t_var_f x)
 | dks_int : forall (G:dkcontext),
     dk_wf_ctx G ->
     dk_sub G t_int t_int
 | dks_arrow : forall (G:dkcontext) (A1 A2 B1 B2:type),
      dk_sub G B1 A1  ->
     dk_sub G A2 B2 ->
     dk_sub G (t_arrow A1 A2) (t_arrow B1 B2)
 | dks_forall_l : forall (G:dkcontext) (A B C:type),
     dk_wf_tp G C ->
     dk_mono C ->
     dk_sub G  (open_type_wrt_type  A   C )  B ->
     dk_sub G (t_all A) B
 | dks_forall_r : forall (L:vars) (G:dkcontext) (A B:type),
      ( forall x , x \notin  L  -> dk_sub  ( x  ::  G )  A  ( open_type_wrt_type B (t_var_f x) )  )  ->
     dk_sub G A (t_all B).

(** infrastructure *)
Hint Constructors dk_mono dk_wf_ctx dk_wf_tp dk_sub lc_type : core.


(** Adapted from "Elements of Set Theory" Chapter 7 **)
(** Coq coding by choukh, Dec 2020 **)

Global Set Warnings "-unrecognized-unicode".

Require Import Relation_Definitions.
Require Import RelationClasses.
Require Import ZFC.Lib.FuncFacts.
Require Import ZFC.Lib.WosetMin.
Require Export ZFC.Elements.EST7_2.

(*** EST第七章3：序结构，序同构，良序结构，伊普西隆像 ***)

(* 序结构 *)
Module OrderedStruct.
Declare Scope OrderedStruct_scope.
Delimit Scope OrderedStruct_scope with os.
Open Scope OrderedStruct_scope.

Record OrderedStruct : Type := constr {
  A : set;
  R : set;
  ordered_struct : is_binRel R A;
}.

Lemma eq_intro : ∀ S T, A S = A T → R S = R T → S = T.
Proof.
  intros S T HA HR. destruct S. destruct T.
  simpl in *. subst. f_equal. apply proof_irrelevance.
Qed.

Definition po := λ S, poset (A S) (R S).
Definition lo := λ S, loset (A S) (R S).
Definition wo := λ S, woset (A S) (R S).

(* 序嵌入双射 *)
Definition order_embedding_bijection := λ f S T, f: A S ⟺ A T ∧
  ∀ x y ∈ A S, (x <ᵣ y) (R S) ↔ (f[x] <ᵣ f[y]) (R T).
Notation "f :ₒₑ S ⟺ T" := (order_embedding_bijection f S T)
  (at level 70) : OrderedStruct_scope.

(* 序同构 *)
Definition isomorphic : relation OrderedStruct :=
  λ S T, ∃ f, f :ₒₑ S ⟺ T.

Notation "S ≅ T" := ( isomorphic S T) (at level 60) : OrderedStruct_scope.
Notation "S ≇ T" := (¬isomorphic S T) (at level 60) : OrderedStruct_scope.

Theorem iso_refl : Reflexive isomorphic.
Proof with auto.
  intros S. exists (Ident (A S)).
  split. apply ident_bijection.
  intros x Hx y Hy. rewrite ident_ap, ident_ap... reflexivity.
Qed.

Theorem iso_symm : Symmetric isomorphic.
Proof with eauto.
  intros S T [f [Hf H]]. exists (f⁻¹).
  pose proof (inv_bijection f (A S) (A T) Hf) as Hf'.
  destruct Hf as [Hi [_ Hr]].
  split... intros x Hx y Hy. symmetry.
  replace x with (f[f⁻¹[x]]) at 2.
  replace y with (f[f⁻¹[y]]) at 2. apply H.
  eapply ap_ran... apply bijection_is_func...
  eapply ap_ran... apply bijection_is_func...
  rewrite inv_ran_reduction... congruence.
  rewrite inv_ran_reduction... congruence.
Qed.

Theorem iso_trans : Transitive isomorphic.
Proof with eauto; try congruence.
  intros S T U [f [Hf H1]] [g [Hg H2]]. exists (g ∘ f).
  pose proof (compo_bijection f g (A S) (A T) (A U) Hf Hg) as Hcom.
  apply bijection_is_func in Hf as [Hf _].
  apply bijection_is_func in Hg as [Hg _].
  assert (H := Hf). destruct H as [Hff [Hdf _]].
  assert (H := Hg). destruct H as [Hfg [Hdg _]].
  split... intros x Hx y Hy.
  rewrite (H1 x Hx y Hy).
  rewrite compo_correct, compo_correct...
  - apply H2. eapply ap_ran... eapply ap_ran...
  - rewrite compo_dom... apply SepI...
    rewrite Hdg. eapply ap_ran...
  - rewrite compo_dom... apply SepI...
    rewrite Hdg. eapply ap_ran...
Qed.

(* 序同构是等价关系 *)
Theorem iso_equiv : Equivalence isomorphic.
Proof.
  split. apply iso_refl. apply iso_symm. apply iso_trans.
Qed.
Existing Instance iso_equiv.

(* 像关系 *)
Module ImageRel.

Definition constr := λ f B R,
  BinRel B (λ x y, (f[x] <ᵣ f[y]) R).

Lemma po : ∀ f B A R, poset A R → poset B (constr f B R).
Proof with eauto.
  intros * [_ [_ [Htr Hir]]]. repeat split...
  - apply binRel_is_binRel.
  - eapply binRel_is_rel... apply binRel_is_binRel.
  - intros x y z Hxy Hyz.
    apply binRelE2 in Hxy as [Hx [Hy Hxy]].
    apply binRelE2 in Hyz as [_  [Hz Hyz]].
    apply binRelI... eapply Htr...
  - intros x Hxx.
    apply binRelE3 in Hxx. eapply Hir...
Qed.

Lemma lo : ∀ f B A R, f: B ⇔ A →
  loset A R → loset B (constr f B R).
Proof with eauto; try congruence.
  intros * Hf Hlo.
  apply injection_is_func in Hf as [Hf Hi].
  apply loset_iff_connected_poset in Hlo as [Hcon Hpo].
  apply loset_iff_connected_poset. split; [|eapply po]...
  intros x Hx y Hy Hnq.
  assert (Hnqf: f[x] ≠ f[y]). {
    destruct Hf as [_ [Hd _]].
    intros H. apply Hnq. eapply injectiveE...
  }
  edestruct Hcon... eapply ap_ran... eapply ap_ran...
  - left. apply binRelI...
  - right. apply binRelI...
Qed.

Lemma wo : ∀ f B A R, f: B ⇔ A →
  woset A R → woset B (constr f B R).
Proof with eauto; try congruence.
  intros * Hf [Hlo Hmin]. split. eapply lo...
  intros A' [a Ha'] Hsub. apply Hsub in Ha' as Ha.
  apply injection_is_func in Hf as [[Hf [Hd Hr]] Hi].
  specialize Hmin with (f⟦A'⟧) as [b [Hb Hmin]]. {
    exists (f[a]). eapply imgI... apply func_correct...
  } {
    intros y Hy. apply imgE in Hy as [x [Hx Hp]].
    eapply ranI in Hp. apply Hr...
  }
  apply imgE in Hb as [m [Hm' Hp]].
  apply Hsub in Hm' as Hm.
  apply func_ap in Hp... subst b.
  exists m. split... intros x Hx'.
  destruct (classic (m = x)). right... left.
  apply Hsub in Hx' as Hx.
  apply binRelI... edestruct Hmin...
  eapply imgI... apply func_correct...
  exfalso. apply H. eapply injectiveE...
Qed.

Fact eq_subRel : ∀ R B, R ⥏ B = constr (Ident B) B R.
Proof with auto.
  intros. ext Hx.
  - apply SepE in Hx as [Hlt Hp].
    apply CPrdE1 in Hp as [a [Ha [b [Hb Heq]]]]. subst x.
    apply SepI... apply CPrdI... zfc_simple.
    rewrite ident_ap, ident_ap...
  - apply SepE in Hx as [Hp Hlt].
    apply CPrdE1 in Hp as [a [Ha [b [Hb Heq]]]]. subst x.
    zfc_simple. rewrite ident_ap, ident_ap in Hlt...
    apply SepI... apply CPrdI...
Qed.

End ImageRel.

(* 与偏序结构同构的结构也是偏序结构 *)
Theorem iso_po : ∀ S T, S ≅ T → po S → po T.
Proof with auto.
  intros * Hiso Hpo. symmetry in Hiso.
  destruct Hiso as [f [_ H]].
  eapply (ImageRel.po f (A T)) in Hpo.
  replace (ImageRel.constr f (A T) (R S)) with (R T) in Hpo...
  ext p Hp.
  - apply ordered_struct in Hp as Hcp.
    apply CPrdE1 in Hcp as [a [Ha [b [Hb Heq]]]].
    subst p. apply binRelI... apply H...
  - apply binRelE1 in Hp as [a [Ha [b [Hb [Heq Hlt]]]]].
    subst p. apply H...
Qed.

(* 与线序结构同构的结构也是线序结构 *)
Theorem iso_lo : ∀ S T, S ≅ T → lo S → lo T.
Proof with auto.
  intros * Hiso Hlo. symmetry in Hiso.
  destruct Hiso as [f [Hf H]].
  apply bijection_is_injection in Hf as [Hf _].
  apply (ImageRel.lo f (A T)) in Hlo...
  replace (ImageRel.constr f (A T) (R S)) with (R T) in Hlo...
  ext p Hp.
  - apply ordered_struct in Hp as Hcp.
    apply CPrdE1 in Hcp as [a [Ha [b [Hb Heq]]]].
    subst p. apply binRelI... apply H...
  - apply binRelE1 in Hp as [a [Ha [b [Hb [Heq Hlt]]]]].
    subst p. apply H...
Qed.

(* 与良序结构同构的结构也是良序结构 *)
Theorem iso_wo : ∀ S T, S ≅ T → wo S → wo T.
Proof with auto.
  intros * Hiso Hwo. symmetry in Hiso.
  destruct Hiso as [f [Hf H]].
  apply bijection_is_injection in Hf as [Hf _].
  apply (ImageRel.wo f (A T)) in Hwo...
  replace (ImageRel.constr f (A T) (R S)) with (R T) in Hwo...
  ext p Hp.
  - apply ordered_struct in Hp as Hcp.
    apply CPrdE1 in Hcp as [a [Ha [b [Hb Heq]]]].
    subst p. apply binRelI... apply H...
  - apply binRelE1 in Hp as [a [Ha [b [Hb [Heq Hlt]]]]].
    subst p. apply H...
Qed.

End OrderedStruct.

(* 良序结构 *)
Module WoStruct.
Declare Scope WoStruct_scope.
Delimit Scope WoStruct_scope with wo.
Open Scope WoStruct_scope.

Record WoStruct : Type := constr {
  A : set;
  R : set;
  wo : woset A R;
}.
Global Hint Immediate wo : core.

Lemma empty_loset : woset ∅ ∅.
Proof with auto.
  repeat split.
  - intros x H. exfalso0.
  - intros x y z H. exfalso0.
  - intros x H. exfalso0.
  - intros B Hne Hsub. apply sub_empty in Hsub.
    exfalso. eapply EmptyNI; eauto.
Qed.

Notation "ø̃" := (constr ∅ ∅ empty_loset) : WoStruct_scope.
Notation "ℕ̃" := (constr ω Lt Lt_wellOrder) : WoStruct_scope.

Lemma eq_intro : ∀ S T, A S = A T → R S = R T → S = T.
Proof.
  intros S T HA HR. destruct S. destruct T.
  simpl in *. subst. f_equal. apply proof_irrelevance.
Qed.

(* 序嵌入双射 *)
Definition order_embedding_bijection := λ f S T, f: A S ⟺ A T ∧
  ∀ x y ∈ A S, (x <ᵣ y) (R S) ↔ (f[x] <ᵣ f[y]) (R T).
Notation "f :ₒₑ S ⟺ T" := (order_embedding_bijection f S T)
  (at level 70) : WoStruct_scope.

Definition isomorphic : relation WoStruct :=
  λ S T, ∃ f, f :ₒₑ S ⟺ T.

Notation "S ≅ T" := ( isomorphic S T) (at level 60) : WoStruct_scope.
Notation "S ≇ T" := (¬isomorphic S T) (at level 60) : WoStruct_scope.

Definition SubStruct := λ S T, A S ⊆ A T ∧ R S = R T ⥏ A S.
Notation "S ⊑ T" := (SubStruct S T) (at level 70) : WoStruct_scope.

Local Lemma woset_is_binRel : ∀ A R, woset A R → is_binRel R A.
Proof. intros. apply H. Qed.

Local Definition br := λ S, woset_is_binRel (A S) (R S) (wo S).
Definition parent := λ S, OrderedStruct.constr (A S) (R S) (br S).

Lemma parent_iso : ∀ S T,
  S ≅ T ↔ OrderedStruct.isomorphic (parent S) (parent T).
Proof.
  split; intros [f [Hf Hoe]]; exists f; split; auto.
Qed.

Lemma parent_wo : ∀ S, OrderedStruct.wo (parent S).
Proof. intros. apply wo. Qed.

Theorem iso_refl : Reflexive isomorphic.
Proof.
  intros S. apply parent_iso. apply OrderedStruct.iso_refl.
Qed.

Theorem iso_symm : Symmetric isomorphic.
Proof.
  intros S T H. rewrite parent_iso in H. apply parent_iso.
  apply OrderedStruct.iso_symm; auto.
Qed.

Theorem iso_trans : Transitive isomorphic.
Proof.
  intros S T U H1 H2. rewrite parent_iso in H1, H2.
  apply parent_iso. eapply OrderedStruct.iso_trans; eauto.
Qed.

Theorem iso_equiv : Equivalence isomorphic.
Proof.
  split. apply iso_refl. apply iso_symm. apply iso_trans.
Qed.
Existing Instance iso_equiv.

(* 良序结构间的态射唯一 *)
Theorem wo_iso_unique : ∀ S T, S ≅ T → ∃! f, f :ₒₑ S ⟺ T.
Proof with eauto; try congruence.
  intros S T Hiso. rewrite <- unique_existence. split...
  intros f g [Hf H1] [Hg H2].
  cut (∀ f g S T, f :ₒₑ S ⟺ T → g :ₒₑ S ⟺ T → f ⊆ g). {
    intros H. apply sub_antisym; eapply H; split...
  }
  clear Hiso Hf H1 Hg H2 f g S T.
  intros f g S T [Hf Hoe] [Hg Hopg] p Hp.
  destruct (wo S) as [Hlo Hmin].
  apply inv_bijection in Hf as Hf'.
  apply bijection_is_func in Hf' as [Hf' _].
  apply bijection_is_func in Hf as [Hf [Hif Hrf]].
  assert (H := Hf). destruct H as [Hff [Hdf _]].
  apply inv_bijection in Hg as Hg'.
  apply bijection_is_func in Hg' as [Hg' _].
  apply bijection_is_func in Hg as [Hg [Hig Hrg]].
  assert (H := Hg). destruct H as [Hfg [Hdg _]].
  apply func_pair' in Hp as [x [y [Hp Heqp]]]... subst p.
  apply domI in Hp as Hx. rewrite Hdf in Hx.
  set {x ∊ A S | f[x] = g[x]} as AS'.
  replace (A S) with AS' in Hx. {
    apply SepE in Hx as [Hx Heq].
    apply func_ap in Hp... apply func_point...
  }
  eapply transfinite_induction...
  split. intros a Ha. apply SepE1 in Ha...
  intros t Ht Hsub. apply SepI...
  contra.
  apply (lo_connected (R T) (A T)) in H as []; revgoals...
  eapply ap_ran... eapply ap_ran... apply (wo T).
  - assert (Hgt: g[t] ∈ ran f). {
      rewrite Hrf, <- Hrg. eapply ranI. apply func_correct...
    }
    rewrite <- (inv_ran_reduction f Hif (g[t])) in H...
    apply Hoe in H; revgoals... { eapply ap_ran... }
    remember (f⁻¹[g[t]]) as t'.
    assert (Ht': t' ∈ seg t (R S)) by (apply segI; auto).
    apply Hsub in Ht'. apply SepE in Ht' as [Ht' Heq].
    rewrite Heqt' in Heq at 1.
    rewrite inv_ran_reduction in Heq...
    eapply injectiveE in Heq... rewrite Heq in H.
    eapply (lo_irrefl (R S) (A S))...
  - assert (Hft: f[t] ∈ ran g). {
      rewrite Hrg, <- Hrf. eapply ranI. apply func_correct...
    }
    rewrite <- (inv_ran_reduction g Hig (f[t])) in H...
    apply Hopg in H; revgoals... { eapply ap_ran... }
    remember (g⁻¹[f[t]]) as t'.
    assert (Ht': t' ∈ seg t (R S)) by (apply segI; auto).
    apply Hsub in Ht'. apply SepE in Ht' as [Ht' Heq].
    rewrite Heqt' in Heq at 2.
    rewrite inv_ran_reduction in Heq...
    eapply injectiveE in Heq... rewrite Heq in H.
    eapply (lo_irrefl (R S) (A S))...
Qed.

Section Seg.

Let seg := λ t A R, {x ∊ A | (x <ᵣ t) R}.

Lemma seg_woset : ∀ t A R, woset A R →
  let B := seg t A R in woset B (R ⥏ B).
Proof.
  intros t A R H B. eapply subRel_woset; eauto.
  intros x Hx. apply SepE1 in Hx. apply Hx.
Qed.

Definition Seg := λ t S,
  let B := seg t (A S) (R S) in
  constr B ((R S) ⥏ B) (seg_woset _ _ _ (wo S)).

End Seg.

Lemma seg_a_eq : ∀ S, ∀t ∈ A S, A (Seg t S) = seg t (R S).
Proof with eauto.
  intros S t Ha.
  ext Hx.
  - apply SepE in Hx as [Hx Hxa].
    apply SepI... eapply domI...
  - apply SepE in Hx as [Hx Hxa].
    apply SepI... eapply dom_binRel... apply wo.
Qed.

Lemma seg_lt : ∀ S r s t, (s <ᵣ t) (R S) →
  (r <ᵣ s) (R (Seg t S)) ↔ (r <ᵣ s) (R S).
Proof with eauto.
  intros S r s t Hst. split; intros Hlt.
  - apply SepE1 in Hlt...
  - apply SepI... apply CPrdI; apply SepI...
    eapply dom_binRel. apply wo. eapply domI...
    destruct (wo S) as [[_ [Htr _]] _]. eapply Htr...
    eapply dom_binRel. apply wo. eapply domI...
Qed.

Lemma seg_of_seg : ∀ S r s t, (s <ᵣ t) (R S) →
  r ∈ seg s (R (Seg t S)) ↔ r ∈ seg s (R S).
Proof with auto.
  intros S r s t Hst. split; intros Hlt.
  - apply SepE2 in Hlt. apply seg_lt in Hlt... apply segI...
  - apply segI. apply seg_lt... apply SepE2 in Hlt...
Qed.

Section RecursionOnWoStruct.
Import TransfiniteRecursion.

Definition Recursion := λ S γ, Recursion (A S) (R S) γ.

Definition recursion_spec := λ S (γ : set → set → Prop) F,
  is_function F ∧ dom F = A S ∧ ∀t ∈ A S, γ (F ↾ seg t (R S)) F[t].

Lemma recursion_spec_intro : ∀ S γ, (∀ x, ∃! y, γ x y) →
  recursion_spec S γ (Recursion S γ).
Proof. intros. apply recursion_spec_intro; auto. Qed.

End RecursionOnWoStruct.

(*  像 *)
Module Import EpsilonImage.

Definition γ := λ x y, y = ran x.
Definition E := λ S, Recursion S γ.
Definition α := λ S, ran (E S).
Definition ε := λ S, MemberRel (α S).

Lemma e_spec : ∀ S, recursion_spec S γ (E S).
Proof.
  intros. apply recursion_spec_intro. intros f.
  rewrite <- unique_existence. split.
  exists (ran f). congruence. congruence.
Qed.

Lemma e_empty : ∀ S, A S = ∅ → E S = ∅.
Proof.
  intros. destruct (e_spec S) as [Hf [Hd _]].
  apply empty_dom; congruence.
Qed.

Lemma e_ap : ∀ S,
  ∀t ∈ A S, (E S)[t] = {(E S)[x] | x ∊ seg t (R S)}.
Proof with eauto.
  intros S t Ht. destruct (e_spec S) as [Hf [Hd Hγ]].
  apply Hγ in Ht. rewrite Ht.
  ext y Hy.
  - apply ranE in Hy as [x Hp]. apply restrE2 in Hp as [Hp Hx].
    apply ReplAx. exists x. split... apply func_ap in Hp...
  - apply ReplAx in Hy as [x [Hx Hap]].
    apply (ranI _ x). apply restrI... apply func_point...
    rewrite Hd. apply SepE in Hx as [Hx _].
    eapply dom_binRel... destruct (wo S) as [[] _]...
Qed.

Lemma e_ap_order : ∀ S,
  ∀ s t ∈ A S, (s <ᵣ t) (R S) → (E S)[s] ∈ (E S)[t].
Proof with auto.
  intros S s Hs t Ht Hst. rewrite (e_ap S t Ht).
  apply ReplI... apply segI...
Qed.

Lemma e_intro : ∀ S, ∀ x, ∀ s t ∈ A S,
  (s <ᵣ t) (R S) → x = (E S)[s] → x ∈ (E S)[t].
Proof.
  intros S x s Hs t Ht Hst Hest. subst.
  apply e_ap_order; auto.
Qed.

Lemma e_elim : ∀ S, ∀t ∈ A S, ∀x ∈ (E S)[t],
  ∃s ∈ A S, (s <ᵣ t) (R S) ∧ x = (E S)[s] ∧ x ∈ (E S)[t].
Proof with eauto.
  intros S t Ht x Hx.
  assert (H := Hx). rewrite e_ap in H...
  apply ReplAx in H as [s [Hs Heq]].
  apply SepE in Hs as [Hs Hst].
  eapply dom_binRel in Hs... apply wo.
Qed.

Theorem e_irrefl : ∀ S, ∀t ∈ A S, (E S)[t] ∉ (E S)[t].
Proof with eauto.
  intros S t Ht Hin.
  destruct (wo S) as [Hlo Hmin].
  set {t ∊ A S | (E S)[t] ∈ (E S)[t]} as B.
  specialize Hmin with B as [t₀ [Ht₀ Hmin]].
  - exists t. apply SepI...
  - intros x Hx. apply SepE1 in Hx...
  - apply SepE in Ht₀ as [Ht₀ Hin₀]. assert (H := Hin₀).
    apply e_elim in H as [s [Hs [Hst₀ [Heq H]]]]...
    assert (Hsb: s ∈ B). { apply SepI... rewrite Heq in H... }
    apply Hmin in Hsb. eapply lo_not_le_gt...
Qed.

Lemma e_injective : ∀ S, injective (E S).
Proof with eauto; try congruence.
  intros. destruct (wo S) as [Hlo Hmin].
  destruct (e_spec S) as [Hf [Hd _]]...
  split... intros y Hy. rewrite <- unique_existence.
  split. apply ranE in Hy...
  intros s t Hs Ht.
  apply func_ap in Hs as H1... apply domI in Hs.
  apply func_ap in Ht as H2... apply domI in Ht.
  rewrite <- H2 in H1.
  contra.
  eapply lo_connected in H as []...
  - apply e_ap_order in H...
    rewrite H1 in H. eapply e_irrefl...
  - apply e_ap_order in H...
    rewrite H1 in H. eapply e_irrefl...
Qed.

Theorem e_bijective : ∀ S, (E S): A S ⟺ α S.
Proof with auto.
  intros. destruct (e_spec S) as [Hf [Hd _]]...
  apply bijection_is_surjection. split. split...
  apply e_injective...
Qed.

Lemma e_ap_in_α : ∀ S, ∀x ∈ A S, (E S)[x] ∈ α S.
Proof with eauto.
  intros S x Hx. eapply ap_ran...
  apply bijection_is_func. apply e_bijective...
Qed.

Theorem e_preserve_order : ∀ S, ∀ x y ∈ A S,
  (x <ᵣ y) (R S) ↔ ((E S)[x] <ᵣ (E S)[y]) (ε S).
Proof with eauto; try congruence.
  intros S x Hx y Hy.
  destruct (e_spec S) as [Hf [Hd _]]...
  split; intros.
  - apply binRelI. apply e_ap_in_α... apply e_ap_in_α...
    apply e_ap_order...
  - apply binRelE3 in H.
    apply (e_elim S) in H as [u [Hu [Hut [Heq He]]]]...
    apply injectiveE in Heq... apply e_injective...
Qed.

Section ImportOrderedStruct.
Import OrderedStruct.

Let Epsilon := λ S,
  constr (α S) (ε S) (memberRel_is_binRel _).

(* 任意良序结构与其伊普西隆结构同构 *)
Lemma wo_iso_epsilon : ∀ S, parent S ≅ Epsilon S.
Proof with auto.
  intros S. exists (E S). split.
  - apply e_bijective...
  - apply e_preserve_order...
Qed.

(* 任意良序结构的伊普西隆结构也是良序结构 *)
Lemma epsilon_wo : ∀ S, wo (Epsilon S).
Proof.
  intros. eapply iso_wo. apply wo_iso_epsilon. apply parent_wo.
Qed.

End ImportOrderedStruct.

Definition Epsilon := λ S, constr (α S) (ε S) (epsilon_wo S).

(* 任意良序结构与其伊普西隆结构同构 *)
Theorem iso_epsilon : ∀ S, S ≅ Epsilon S.
Proof with auto.
  intros S. exists (E S). split.
  - apply e_bijective...
  - apply e_preserve_order...
Qed.

Theorem α_trans : ∀ S, trans (α S).
Proof with auto; try congruence.
  intros S x y Hxy Hy.
  destruct (e_spec S) as [Hf [Hd _]]...
  apply ranE in Hy as [t Hp]. apply domI in Hp as Ht.
  apply func_ap in Hp... subst y.
  apply e_elim in Hxy as [s [Hs [_ [Heq _]]]]...
  subst x. eapply ranI. apply func_correct...
Qed.

Lemma α_intro : ∀ S t, ∀s ∈ A S, (E S)[s] = t → t ∈ α S.
Proof with auto.
  intros S t s Hs Heqt.
  pose proof (e_spec S) as [Hf [Hd _]].
  apply (ranI _ s). apply func_point... rewrite Hd...
Qed.

Lemma α_elim : ∀ S, ∀t ∈ α S, ∃s ∈ A S, (E S)[s] = t.
Proof with auto.
  intros S t Ht.
  pose proof (e_spec S) as [Hf [Hd _]].
  apply ranE in Ht as [s Hp]. apply domI in Hp as Hs.
  apply func_ap in Hp... exists s. split. congruence. unfold E...
Qed.

Fact ε_iff : ∀ S, ∀ x y,
  (x <ᵣ y) (ε S) ↔ x ∈ α S ∧ y ∈ α S ∧ x ∈ y.
Proof with auto.
  intros S x y. split.
  - intros Hxy. apply binRelE2...
  - intros [Hx [Hy Hxy]]. apply binRelI...
Qed.

Lemma seg_e_ap : ∀ S, ∀t ∈ A S, ∀ s, (s <ᵣ t) (R S) →
  (E (Seg t S))[s] = (E S)[s].
Proof with eauto.
  intros S t Ht s Hst.
  destruct (wo S) as [[_ [Htr _]] _].
  assert (Hs: s ∈ A S). {
    eapply dom_binRel. apply wo. eapply domI...
  }
  generalize dependent Hst.
  set {x ∊ A S | (x <ᵣ t) (R S) → (E (Seg t S))[x] = (E S)[x]} as AS'.
  replace (A S) with AS' in Hs. apply SepE2 in Hs... clear Hs s.
  eapply transfinite_induction... split.
  intros s Hs. apply SepE1 in Hs...
  intros s Hs Hsub. apply SepI... intros Hst.
  rewrite e_ap; revgoals... apply SepI...
  ext Hx.
  - apply ReplAx in Hx as [r [Hr Heqx]]. subst x.
    assert (Hrs: (r <ᵣ s) (R S)). {
      apply SepE in Hr as [_ Hrs]. apply seg_lt in Hrs...
    }
    assert (Hrt: (r <ᵣ t) (R S)). eapply Htr...
    apply seg_of_seg in Hr... apply Hsub in Hr.
    apply SepE in Hr as [Hr IH]. apply IH in Hrt.
    eapply e_intro; revgoals...
  - apply e_elim in Hx as [r [Hr [Hrs [Heqx _]]]]...
    assert (Hr': r ∈ seg s (R S)) by (apply segI; auto).
    assert (Hrt: (r <ᵣ t) (R S)). eapply Htr...
    apply ReplAx. exists r. split. apply seg_of_seg...
    apply Hsub in Hr'. apply SepE in Hr' as [_ IH].
    apply IH in Hrt. congruence.
Qed.

Lemma seg_α : ∀ S, ∀t ∈ A S, α (Seg t S) = (E S)[t].
Proof with eauto.
  intros S t Ht.
  pose proof (e_spec (Seg t S)) as [Hf [Hd _]].
  ext Hx.
  - apply α_elim in Hx as [s [Hs Hp]].
    apply SepE2 in Hs. subst x. rewrite seg_e_ap...
    apply e_ap_order... eapply dom_binRel. apply wo. eapply domI...
  - apply e_elim in Hx as [s [Hs [Hst [Heqx Hx]]]]...
    apply (α_intro _ _ s). apply SepI... rewrite seg_e_ap...
Qed.

End EpsilonImage.

(* Examples *)
Module Export EpsilonImageOfNats.

Example e_ω_nat : ∀n ∈ ω, (E ℕ̃)[n] = n.
Proof with neauto.
  intros n Hn.
  ω_induction n.
  - ext Hx.
    + apply e_elim in Hx as [k [_ [Hk _]]]...
      apply binRelE3 in Hk. exfalso0.
    + exfalso0.
  - assert (Hmp: m⁺ ∈ ω) by (apply ω_inductive; auto).
    ext Hx.
    + apply e_elim in Hx as [k [Hk [Hkmp [Heqx Hx]]]]...
      apply binRelE3 in Hkmp.
      apply le_iff_lt_suc in Hkmp as []...
      * apply BUnionI1. rewrite <- IH.
        eapply e_intro... apply binRelI...
      * apply BUnionI2. subst. rewrite <- IH at 2...
    + apply le_iff_lt_suc in Hx as []; [| |eapply ω_trans|]...
      * rewrite <- IH in H.
        apply e_elim in H as [k [Hk [Hkm [Heqx Hx]]]]...
        apply binRelE3 in Hkm.
        eapply e_intro... apply binRelI... apply BUnionI1...
      * apply (e_intro ℕ̃ x m)... apply binRelI... congruence.
Qed.

Example e_nat_nat : ∀ n m ∈ ω, n ∈ m → (E (Seg m ℕ̃))[n] = n.
Proof with neauto.
  intros n Hn p Hp.
  ω_induction n; intros Hnp.
  - ext Hx; [|exfalso0].
    apply (e_elim (Seg p ℕ̃)) in Hx as [k [_ [Hk _]]].
    + apply SepE in Hk as [Hk _].
      apply binRelE3 in Hk. exfalso0.
    + apply SepI... apply binRelI...
  - assert (Hp': p⁺ ∈ ω) by (apply ω_inductive; auto).
    assert (Hm': m⁺ ∈ ω) by (apply ω_inductive; auto).
    assert (Hmp: m ∈ p). { apply (nat_trans p Hp m m⁺)... }
    assert (Hmseg: m ∈ A (Seg p ℕ̃)). {
      apply SepI... apply binRelI...
    }
    assert (Hm'seg: m⁺ ∈ A (Seg p ℕ̃)). {
      apply SepI... apply binRelI...
    }
    ext Hx.
    + apply (e_elim (Seg p ℕ̃)) in Hx as [k [Hk [Hkm [Heqx Hx]]]]...
      apply SepE in Hkm as [Hkm _].
      apply binRelE2 in Hkm as [Hkw [_ Hkm]].
      apply le_iff_lt_suc in Hkm as []...
      * apply BUnionI1. rewrite <- IH...
        eapply e_intro... apply seg_lt; apply binRelI...
      * apply BUnionI2. subst. rewrite <- IH at 2...
    + apply le_iff_lt_suc in Hx as []; [| |eapply ω_trans|]...
      * rewrite <- IH in H...
        apply (e_elim (Seg p ℕ̃)) in H as [k [Hk [Hkm [Heqx Hx]]]]...
        apply SepE in Hkm as [Hkm _].
        apply binRelE2 in Hkm as [Hkw [_ Hkm]].
        eapply (e_intro (Seg p ℕ̃))...
        apply seg_lt; apply binRelI... apply BUnionI1...
      * apply (e_intro (Seg p ℕ̃) x m)...
        apply seg_lt; apply binRelI... rewrite IH...
Qed.

Example α_nat : ∀n ∈ ω, α (Seg n ℕ̃) = n.
Proof with neauto; try congruence.
  intros n Hn.
  ω_induction n.
  - unfold α. replace (E (Seg ∅ ℕ̃)) with ∅.
    apply ran_of_empty. symmetry. apply e_empty.
    ext Hx; [|exfalso0].
    apply SepE2 in Hx. apply binRelE3 in Hx...
  - assert (Hm': m⁺ ∈ ω) by (apply ω_inductive; auto).
    destruct (e_spec (Seg m⁺ ℕ̃)) as [Hf [Hd _]]...
    ext y Hy.
    + apply ranE in Hy as [x Hp]. apply domI in Hp as Hx.
      rewrite Hd in Hx. apply SepE2 in Hx.
      apply binRelE2 in Hx as [Hx [_ Hx']].
      apply func_ap in Hp... rewrite e_nat_nat in Hp...
    + assert (Hyw: y ∈ ω) by (eapply ω_trans; eauto).
      apply (ranI _ y). apply func_point...
      * rewrite Hd. apply SepI... apply binRelI...
      * rewrite seg_e_ap... apply e_ω_nat... apply binRelI...
Qed.

Example α_ω : α ℕ̃ = ω.
Proof with auto.
  destruct (e_spec ℕ̃) as [Hf [Hd _]].
  ext m Hm.
  - apply ranE in Hm as [n Hp]. apply domI in Hp as Hn.
    rewrite Hd in Hn. apply func_ap in Hp...
    rewrite e_ω_nat in Hp... subst...
  - apply (ranI _ m). apply func_point...
    rewrite Hd... apply e_ω_nat...
Qed.

End EpsilonImageOfNats.

Import WosetMin.FullVer.

Definition Min := λ S, Min (A S) (R S).

Lemma min_correct : ∀ S B,
  ⦿ B → B ⊆ A S → minimum (Min S)[B] B (R S).
Proof. intros; apply min_correct; auto. Qed.

(* 良序结构的同构关系具有至少三歧性 *)
Theorem wo_iso_at_least_trich : ∀ S T,
  S ≅ T ∨
  (∃t ∈ A T, S ≅ Seg t T) ∨
  (∃t ∈ A S, Seg t S ≅ T).
Proof with eauto; try congruence.
  intros.
  destruct (wo S) as [Hlo HminS]. assert (Hlo' := Hlo).
  destruct Hlo' as [_ [Htr _]].
  destruct (wo T) as [[_ [Htr' _]] _].
  set (Extraneous (A S ∪ A T)) as e.
  set (λ f y, y = match (ixm (⦿ (A T - ran f))) with
    | inl _ => (Min T)[A T - ran f]
    | inr _ => e
  end) as γ.
  set (Recursion S γ) as F.
  pose proof (recursion_spec_intro S γ) as [HfF [HdF Hγ]]. {
    intros f. rewrite <- unique_existence. split... unfold γ.
    destruct (ixm (⦿ (A T - ran f))).
    - exists ((Min T)[A T - ran f])...
    - exists e...
  }
  fold F in HfF, HdF, Hγ. unfold γ in Hγ.
  assert (HL1: ∀ x y, (x <ᵣ y) (R S) →
    A T - F⟦seg y (R S)⟧ ⊆ A T - F⟦seg x (R S)⟧
  ). {
    intros x y Hxy. intros b Hb.
    apply SepE in Hb as [Hb Hb'].
    apply SepI... intros Hb''. apply Hb'.
    apply imgE in Hb'' as [t [Htx Hp]].
    eapply imgI... apply SepE in Htx as [_ Htx].
    apply segI. eapply Htr...
  }
  (* Case I *)
  destruct (classic (e ∈ ran F)) as [He|He]. {
    right; right.
    set {x ∊ A S | F[x] = e} as B.
    specialize HminS with B as [a [Ha HminS]]. {
      apply ranE in He as [a Hp]. apply domI in Hp as Ha.
      apply func_ap in Hp... exists a. apply SepI...
    } {
      intros b Hb. apply SepE1 in Hb...
    }
    assert (HL2: ∀ x y ∈ A S, (x <ᵣ y) (R S) → (y <ᵣ a) (R S) →
      (F[x] <ᵣ F[y]) (R T)
    ). {
      intros x Hx y Hy Hxy Hya. rewrite Hγ, Hγ...
      set (A T - ran (F ↾ seg x (R S))) as C.
      set (A T - ran (F ↾ seg y (R S))) as D.
      apply Hγ in Hx as HFx. fold C in HFx.
      apply Hγ in Hy as HFy. fold D in HFy.
      destruct (ixm (⦿ C)); destruct (ixm (⦿ D)).
      + pose proof (min_correct T C) as [_ Hmin]...
          { intros c Hc. apply SepE1 in Hc... }
        pose proof (min_correct T D) as [Hmd _]...
          { intros d Hd. apply SepE1 in Hd... }
        assert (Hmc: (Min T)[D] ∈ C). { eapply HL1... }
        apply Hmin in Hmc as []... exfalso.
        rewrite <- H, <- HFx in Hmd. apply SepE2 in Hmd.
        apply Hmd. eapply ranI. apply restrI.
        apply segI... apply func_correct...
      + exfalso. assert (y ∈ B). apply SepI...
        apply HminS in H. eapply lo_not_le_gt...
      + exfalso. assert (x ∈ B). apply SepI...
        apply HminS in H as []; subst; eapply lo_irrefl...
      + exfalso. assert (x ∈ B). apply SepI...
        apply HminS in H as []; subst; eapply lo_irrefl...
    }
    set (F ↾ seg a (R S)) as Fᵒ.
    apply SepE1 in Ha as Haa.
    exists a. split... exists Fᵒ. split; [split; split|].
    - apply restr_func...
    - intros b Hb. rewrite <- unique_existence.
      split. apply ranE in Hb...
      intros x y Hpx Hpy.
      apply restrE2 in Hpx as [Hpx Hx]. apply func_ap in Hpx...
      apply restrE2 in Hpy as [Hpy Hy]. apply func_ap in Hpy...
      apply SepE in Hx as [Hx Hxa].
      apply SepE in Hy as [Hy Hya].
      eapply dom_binRel in Hx; [|apply wo].
      eapply dom_binRel in Hy; [|apply wo].
      contra.
      eapply lo_connected in H as []; eauto;
      (eapply lo_irrefl; [apply (wo T)|]).
      + apply (HL2 _ Hx _ Hy) in Hya... rewrite Hpx, Hpy in Hya...
      + apply (HL2 _ Hy _ Hx) in Hxa... rewrite Hpx, Hpy in Hxa...
    - unfold Fᵒ. rewrite <- seg_a_eq... apply restr_dom...
      intros x Hx. apply SepE1 in Hx...
    - apply sub_antisym.
      + intros y Hy. apply ranE in Hy as [x Hp].
        apply restrE2 in Hp as [Hp Hx].
        apply SepE in Hx as [Hx Hxa].
        eapply dom_binRel in Hx; [|apply wo].
        apply func_ap in Hp... subst y.
        apply Hγ in Hx as Hap. rewrite Hap.
        set (A T - ran (F ↾ seg x (R S))) as C. fold C in Hap.
        destruct (ixm (⦿ C)).
        * pose proof (min_correct T C) as [Hm _]...
          intros c Hc. apply SepE1 in Hc... apply SepE1 in Hm...
        * exfalso. assert (x ∈ B). apply SepI...
          apply HminS in H. eapply lo_not_le_gt...
      + apply sub_iff_no_comp.
        destruct (classic (A T - ran Fᵒ = ∅))...
        exfalso. apply EmptyNE in H.
        apply Hγ in Haa. fold Fᵒ in Haa.
        destruct (ixm (⦿ (A T - ran Fᵒ)))...
        assert (Hm: (Min T)[A T - ran Fᵒ] ∈ A T - ran Fᵒ). apply min_correct...
        apply SepE1 in Hm. apply SepE2 in Ha.
        apply (extraneous (A S ∪ A T)). fold e. apply BUnionI2...
    - intros x Hx y Hy.
      apply SepE in Hx as [Hx Hxa].
      apply SepE in Hy as [Hy Hya].
      replace (Fᵒ[x]) with (F[x]).
      replace (Fᵒ[y]) with (F[y]).
      + split; intros Hxy.
        apply HL2... apply SepE1 in Hxy...
        destruct (classic (x = y)).
        * exfalso. subst.
          eapply lo_irrefl. apply (wo T). apply Hxy.
        * apply SepI; [|apply CPrdI; apply SepI]...
          eapply lo_connected in H as []... exfalso.
          eapply lo_irrefl. apply (wo T).
          eapply Htr'... apply HL2...
      + symmetry. apply func_ap. apply restr_func...
        apply restrI. apply segI... apply func_correct...
      + symmetry. apply func_ap. apply restr_func...
        apply restrI. apply segI... apply func_correct...
  }
  (* End of Case I *)
  (* Lemma for Case II & III *)
  assert (HL2: ∀ x y ∈ A S, (x <ᵣ y) (R S) →
    (F[x] <ᵣ F[y]) (R T)
  ). {
    intros x Hx y Hy Hxy. rewrite Hγ, Hγ...
    set (A T - ran (F ↾ seg x (R S))) as C.
    set (A T - ran (F ↾ seg y (R S))) as D.
    apply Hγ in Hx as HFx. fold C in HFx.
    apply Hγ in Hy as HFy. fold D in HFy.
    destruct (ixm (⦿ C)); destruct (ixm (⦿ D)).
    + pose proof (min_correct T C) as [_ Hmin]...
        { intros c Hc. apply SepE1 in Hc... }
      pose proof (min_correct T D) as [Hmd _]...
        { intros d Hd. apply SepE1 in Hd... }
      assert (Hmc: (Min T)[D] ∈ C). { eapply HL1... }
      apply Hmin in Hmc as []... exfalso.
      rewrite <- H in Hmd. apply SepE2 in Hmd.
      apply Hmd. eapply ranI. apply restrI.
      apply segI... apply func_point...
    + exfalso. apply He. rewrite <- HFy.
      eapply ranI. apply func_correct...
    + exfalso. apply He. rewrite <- HFx.
      eapply ranI. apply func_correct...
    + exfalso. apply He. rewrite <- HFx.
      eapply ranI. apply func_correct...
  }
  (* Case II *)
  destruct (classic (ran F = A T)) as [Hr|Hnq]. {
    left.
    exists F. split; [split; split|]...
    - intros b Hb. rewrite <- unique_existence.
      split. apply ranE in Hb...
      intros x y Hpx Hpy.
      apply domI in Hpx as Hx. apply func_ap in Hpx...
      apply domI in Hpy as Hy. apply func_ap in Hpy...
      rewrite HdF in Hx, Hy.
      contra.
      eapply lo_connected in H as []; eauto;
      (eapply lo_irrefl; [apply (wo T)|]).
      + apply (HL2 _ Hx _ Hy) in H... rewrite Hpx, Hpy in H...
      + apply (HL2 _ Hy _ Hx) in H... rewrite Hpx, Hpy in H...
    - intros x Hx y Hy. split; intros Hxy. apply HL2...
      destruct (classic (x = y)).
      + exfalso. subst.
        eapply lo_irrefl. apply (wo T). apply Hxy.
      + eapply lo_connected in H as []... exfalso.
        eapply lo_irrefl. apply (wo T).
        eapply Htr'... apply HL2...
  }
  (* Case III *)
  right; left.
  assert (Hsub: ran F ⊆ A T). {
    intros y Hy.
    apply ranE in Hy as [x Hp]. apply domI in Hp as Hx.
    apply func_ap in Hp as Hap... rewrite Hγ in Hap...
    set (A T - ran (F ↾ seg x (R S))) as B. fold B in Hap.
    destruct (ixm (⦿ B)).
    - assert ((Min T)[B] ∈ B). {
        apply min_correct...
        intros b Hb. apply SepE1 in Hb...
      }
      apply SepE1 in H...
    - exfalso. apply He. apply (ranI _ x)...
  }
  assert (Hpsub: ran F ⊂ A T) by (split; auto).
  apply comp_nonempty in Hpsub.
  set ((Min T)[A T - ran F]) as b.
  pose proof (min_correct T (A T -ran F)) as [Hm Hminb]...
  fold b in Hm, Hminb.
  assert (Hr: ran F = A (Seg b T)). {
    apply sub_antisym; intros y Hy.
    - destruct (wo T) as [Hlo' _].
      apply SepE in Hm as [Hb1 Hb2]. apply Hsub in Hy as Hyt.
      apply SepI... destruct (classic (y = b)). subst...
      apply (lo_connected (R T) (A T)) in H as []...
      exfalso. apply ranE in Hy as [x Hp]. apply domI in Hp as Hx.
      apply func_ap in Hp... rewrite Hγ in Hp...
      set (A T - ran (F ↾ seg x (R S))) as C. fold C in Hp.
      destruct (ixm (⦿ C)); subst y.
      + pose proof (min_correct T C) as [Hm Hminc]...
        intros c Hc. apply SepE1 in Hc...
        assert (Hb: b ∈ C). {
          apply SepI... intros Hb'. apply restr_ran_included in Hb'...
        }
        apply Hminc in Hb as [];
        (eapply lo_irrefl; [apply (wo T)|])...
        rewrite H0 in H...
      + assert (e ∈ A S ∪ A T) by (apply BUnionI2; auto).
        eapply extraneous...
    - apply SepE in Hy as [Hy Hyb].
      contra as Hy'.
      assert (y ∈ A T - ran F) by (apply SepI; auto).
      apply Hminb in H as []; subst;
      (eapply lo_irrefl; [apply (wo T)|])...
  }
  exists b. split. apply SepE1 in Hm...
  exists F. split; [split; split|]...
  - intros c Hc. rewrite <- unique_existence.
    split. apply ranE in Hc...
    intros x y Hpx Hpy.
    apply domI in Hpx as Hx. apply func_ap in Hpx...
    apply domI in Hpy as Hy. apply func_ap in Hpy...
    rewrite HdF in Hx, Hy.
    contra.
    eapply lo_connected in H as []; eauto;
    (eapply lo_irrefl; [apply (wo T)|]).
    + apply (HL2 _ Hx _ Hy) in H... rewrite Hpx, Hpy in H...
    + apply (HL2 _ Hy _ Hx) in H... rewrite Hpx, Hpy in H...
  - intros x Hx y Hy. split; intros Hxy.
    + apply SepI. apply HL2...
      rewrite <- HdF in Hx, Hy.
      apply func_correct in Hx... apply ranI in Hx.
      apply func_correct in Hy... apply ranI in Hy.
      apply CPrdI; (apply SepI; [apply Hsub|])...
      * rewrite Hr in Hx. apply SepE2 in Hx...
      * rewrite Hr in Hy. apply SepE2 in Hy...
    + apply SepE1 in Hxy.
      destruct (classic (x = y)).
      * exfalso. subst.
        eapply lo_irrefl. apply (wo T). apply Hxy.
      * eapply lo_connected in H as []... exfalso.
        eapply lo_irrefl. apply (wo T).
        eapply Htr'... apply HL2...
Qed.

End WoStruct.


Set Implicit Arguments.
From TLC Require Import LibTactics LibLogic LibProd LibEpsilon LibContainer
     LibSet LibRelation LibPer LibMonoid.
Local Notation path := rtclosure.
From iris_time.union_find.math Require Import LibNatExtra LibIter LibRewrite
     InverseNatNat Ackermann InverseAckermann MiscArith TLCBuffer
     UnionFind01Data UnionFind11Rank.

(* We now define the parent [p] of a vertex and prove several properties of
   this function. *)

(* -------------------------------------------------------------------------- *)

Section Parent.

(* In the following, we again assume that [F] is a ranked disjoint set forest
   with domain [D] and ranks [K]. *)

Variable V : Type.
Variable D : set V.
Variable F : binary V.
Variable K : V -> nat.
Hypothesis is_rdsf_F:
  is_rdsf D F K.

(* -------------------------------------------------------------------------- *)

(* Remark: in the special case where [R] has type [A->A->Prop],
   one may get away without providing a proof of [Inhab A].
  --TEMPORARY from TLC*)

  Definition rel_to_fun' A (R : A -> A -> Prop) (a : A) : A :=
    @rel_to_fun A A (Inhab_of_val a) R a.

(* For convenience, following Tarjan's notation, let [p x] refer to the parent
   of [x], if it exists. *)

Definition p :=
  rel_to_fun' F.

Section NonRoot.

(* Let [x] be a non-root. *)

Variable x : V.
Hypothesis x_non_root: ~ is_root F x.

(* There is an edge from [x] to its parent. *)

Lemma parent_spec:
  F x (p x).
Proof using x_non_root.
  unfold p, rel_to_fun'. eapply rel_rel_to_fun_of_not_forall. eauto.
Qed.

(* Every edge out of [x] is to its parent. *)

Lemma parent_unique:
  forall y,
  F x y ->
  p x = y.
Proof using is_rdsf_F.
  unfold p, rel_to_fun'. intros. eapply rel_to_fun_eq_of_functional.
  eapply is_rdsf_F. (* miracle! *) assumption.
Qed.

(* The parent of [x] is in [D]. *)

Lemma parent_in_D:
  p x \in D.
Proof using is_rdsf_F x_non_root.
  (* We use the fact that [F] is confined to [D]. *)
  eapply is_rdsf_F. eapply parent_spec.
Qed.

(* The rank of [x] is less than the rank of its parent. *)

Lemma parent_has_greater_rank:
  K x < K (p x).
Proof using is_rdsf_F x_non_root.
  (* Rank increases along edges. *)
  eapply is_rdsf_F. eapply parent_spec.
Qed.

(* There is a path from [x] to its parent. *)

Lemma path_to_parent:
  path F x (p x).
Proof using is_rdsf_F x_non_root.
  eauto using parent_spec with rtclosure.
Qed.

(* [x] is equivalent to its parent. *)

Lemma is_equiv_parent:
  is_equiv F x (p x).
Proof using is_rdsf_F x_non_root.
  eauto using path_is_equiv, path_to_parent with is_dsf.
Qed.

(* If [z] is the representative of [x], then [z] is the
   representative of [x]'s parent, too. *)

Lemma is_repr_parent:
  forall z,
  is_repr F x z ->
  is_repr F (p x) z.
Proof using is_rdsf_F x_non_root.
  eauto 11 using is_equiv_trans, is_equiv_parent, path_is_equiv with is_repr is_dsf.
Qed.

(* Thus, there is path from the parent of [x] to the representative of [x]. *)

Lemma path_from_parent_to_repr:
  forall z,
  is_repr F x z ->
  path F (p x) z.
Proof using is_rdsf_F x_non_root.
  intros. forwards [ ? ? ]: is_repr_parent; eauto.
Qed.

End NonRoot.

End Parent.

(* -------------------------------------------------------------------------- *)

(* [intro_parent F v] introduces a fresh name, say [parent], for the parent
   of [v] in the graph [F]. The assertion [F v parent] is automatically
   introduced as well. *)

Ltac intro_parent F v :=
  (* Prove that there is an edge from [v] to [p F v]. *)
  forwards: parent_spec F v; [
    (* We must prove that [v] has a parent in [F]. *)
    unfold not in *; eauto with is_root
  | (* Now, introduce a fresh name for the parent of [v]. *)
    let parent := fresh "parent" in
    generalize dependent (p F v); intro parent; intros
  ].


(**

Syntax of the calculus of constructions as in Streicher
"Semantics of Type Theory" formalized as a multisorted
binding signature.

Written by: Anders Mörtberg, 2017

*)

Require Import UniMath.Foundations.PartD.
Require Import UniMath.Foundations.Sets.

Require Import UniMath.MoreFoundations.Tactics.

Require Import UniMath.Combinatorics.StandardFiniteSets.
Require Import UniMath.Combinatorics.Lists.

Require Import UniMath.CategoryTheory.Core.Categories.
Require Import UniMath.CategoryTheory.Core.Functors.
Require Import UniMath.CategoryTheory.FunctorCategory.
Require Import UniMath.CategoryTheory.categories.HSET.Core.
Require Import UniMath.CategoryTheory.categories.HSET.Colimits.
Require Import UniMath.CategoryTheory.categories.HSET.Limits.
Require Import UniMath.CategoryTheory.categories.HSET.Slice.
Require Import UniMath.CategoryTheory.Chains.Chains.
Require Import UniMath.CategoryTheory.Chains.OmegaCocontFunctors.
Require Import UniMath.CategoryTheory.limits.graphs.limits.
Require Import UniMath.CategoryTheory.limits.graphs.colimits.
Require Import UniMath.CategoryTheory.limits.initial.
Require Import UniMath.CategoryTheory.limits.binproducts.
Require Import UniMath.CategoryTheory.limits.products.
Require Import UniMath.CategoryTheory.limits.bincoproducts.
Require Import UniMath.CategoryTheory.limits.coproducts.
Require Import UniMath.CategoryTheory.limits.terminal.
Require Import UniMath.CategoryTheory.FunctorAlgebras.
Require Import UniMath.CategoryTheory.exponentials.
Require Import UniMath.CategoryTheory.whiskering.
Require Import UniMath.CategoryTheory.Monads.Monads.
Require Import UniMath.CategoryTheory.slicecat.

Require Import UniMath.SubstitutionSystems.Signatures.
Require Import UniMath.SubstitutionSystems.SumOfSignatures.
Require Import UniMath.SubstitutionSystems.BinProductOfSignatures.
Require Import UniMath.SubstitutionSystems.SubstitutionSystems.
Require Import UniMath.SubstitutionSystems.LiftingInitial_alt.
Require Import UniMath.SubstitutionSystems.MonadsFromSubstitutionSystems.
Require Import UniMath.SubstitutionSystems.Notation.
Local Open Scope subsys.
Require Import UniMath.SubstitutionSystems.SignatureExamples.
Require Import UniMath.SubstitutionSystems.BindingSigToMonad.
Require Import UniMath.SubstitutionSystems.MonadsMultiSorted.
Require Import UniMath.SubstitutionSystems.MultiSorted.
Require Import UniMath.SubstitutionSystems.MultiSortedMonadConstruction.

Local Open Scope cat.

Section ccs.

(* Preliminary stuff, upstream? *)
Local Infix "::" := (@cons _).
Local Notation "[]" := (@nil _) (at level 0, format "[]").
Local Notation "C / X" := (slice_cat C X).
Local Notation "a + b" := (setcoprod a b) : set.

(* Was there a general version of this somewhere? *)
Definition six_rec {A : UU} (a b c d e f : A) : stn 6 -> A.
Proof.
induction 1 as [n p].
induction n as [|n _]; [apply a|].
induction n as [|n _]; [apply b|].
induction n as [|n _]; [apply c|].
induction n as [|n _]; [apply d|].
induction n as [|n _]; [apply e|].
induction n as [|n _]; [apply f|].
induction (nopathsfalsetotrue p).
Defined.

(** We assume a two element set of sorts *)
Definition sort : hSet := @tpair _ (λ X, isaset X) bool isasetbool.

Definition ty : sort := true.
Definition el : sort := false.

Local Definition HSET_over_sort : category.
Proof.
exists (HSET / sort).
now apply has_homsets_slice_precat.
Defined.

Let HSET_over_sort2 := [HSET/sort,HSET_over_sort].

(** The grammar of expressions and objects from page 157:
<<
E ::= (Πx:E) E                product of types
    | Prop                    type of propositions
    | Proof(t)                type of proofs of proposition t

t ::= x                       variable
    | (λx:E) t                function abstraction
    | App([x:E] E, t, t)      function application
    | (∀x:E) t                universal quantification
>>

We refer to the first syntactic class as ty and the second as el. We first reformulate the rules as
follows:
<<
A,B ::= Π(A,x.B)              product of types
      | Prop                  type of propositions
      | Proof(t)              type of proofs of proposition t

t,u ::= x                     variable
      | λ(A,x.t)              function abstraction
      | App(A,x.B,t,u)        function application
      | ∀(A,x.t)              universal quantification
>>

This grammar then gives 6 operations, to the left as Vladimir's restricted 2-sorted signature (where
el is 0 and ty is 1) and to the right as a multisorted signature:

((0, 1), (1, 1), 1)                 = (([],ty), ([el], ty), ty)
(1)                                 = ([],ty)
((0, 0), 1)                         = (([], el), ty)
((0, 1), (1, 0), 0)                 = (([], ty), ([el], el), el)
((0, 1), (1, 1), (0, 0), (0, 0), 0) = (([], ty), ([el], ty), ([], el), ([], el), el)
((0, 1), (1, 0), 0)                 = (([], ty), ([el], el), el)

*)

(** The multisorted signature of CC-S *)
Definition CCS_Sig : MultiSortedSig sort.
Proof.
use make_MultiSortedSig.
- exact (stn 6,,isasetstn 6).
- apply six_rec.
  + exact ((([],,ty) :: (cons el [],,ty) :: nil),,ty).
  + exact ([],,ty).
  + exact ((([],,el) :: nil),,ty).
  + exact ((([],,ty) :: (cons el [],,el) :: nil),,el).
  + exact ((([],,ty) :: (cons el [],,ty) :: ([],,el) :: ([],,el) :: nil),,el).
  + exact ((([],,ty) :: (cons el [],,el) :: nil),,el).
Defined.

Definition CCS_Signature : Signature (HSET / sort) _ _ :=
  MultiSortedSigToSignature sort CCS_Sig.

Let Id_H := Id_H _ (BinCoproducts_HSET_slice sort).

Definition CCS_Functor : functor HSET_over_sort2 HSET_over_sort2 :=
  Id_H CCS_Signature.

Lemma CCS_Functor_Initial : Initial (FunctorAlg CCS_Functor).
Proof.
apply SignatureInitialAlgebraSetSort.
apply is_omega_cocont_MultiSortedSigToSignature.
apply slice_precat_colims_of_shape, ColimsHSET_of_shape.
Defined.

Definition CCS_Monad : Monad (HSET / sort) :=
  MultiSortedSigToMonad sort CCS_Sig.

(** Extract the constructors from the initial algebra *)
Definition CCS : HSET_over_sort2 :=
  alg_carrier _ (InitialObject CCS_Functor_Initial).

Let CCS_mor : HSET_over_sort2⟦CCS_Functor CCS,CCS⟧ :=
  alg_map _ (InitialObject CCS_Functor_Initial).

Let CCS_alg : algebra_ob CCS_Functor :=
  InitialObject CCS_Functor_Initial.

Local Lemma BP : BinProducts [HSET_over_sort,HSET].
Proof.
apply BinProducts_functor_precat, BinProductsHSET.
Defined.

Local Notation "'Id'" := (functor_identity HSET_over_sort).
Local Notation "x ⊗ y" := (BinProductObject _ (BP x y)).

(** The variables *)
Definition var_map : HSET_over_sort2⟦Id,CCS⟧ :=
  BinCoproductIn1 (BinCoproducts_functor_precat _ _ _ _ _) · CCS_mor.

Definition Pi_source (X : HSET_over_sort2) : HSET_over_sort2 :=
  ((X ∙ proj_functor sort ty) ⊗ (sorted_option_functor sort el ∙ X ∙ proj_functor sort ty)) ∙ hat_functor sort ty.

(** The Pi constructor *)
Definition Pi_map : HSET_over_sort2⟦Pi_source CCS,CCS⟧ :=
  (CoproductIn _ _ (Coproducts_functor_precat _ _ _ _ _) (● 0)%stn)
    · (BinCoproductIn2 (BinCoproducts_functor_precat _ _ _ _ _))
    · CCS_mor.

Definition Prop_source (X : HSET_over_sort2) : HSET_over_sort2 :=
  constant_functor (slice_cat HSET _) HSET 1%CS  ∙ hat_functor sort ty.

Definition Prop_map : HSET_over_sort2⟦Prop_source CCS,CCS⟧.
Proof.
use ((CoproductIn _ _ (Coproducts_functor_precat _ _ _ _ _) (● 1)%stn) · (BinCoproductIn2 (BinCoproducts_functor_precat _ _ _ _ _)) · CCS_mor).
Defined.

Definition Proof_source (X : HSET_over_sort2) : HSET_over_sort2 :=
  (X ∙ proj_functor sort el) ∙ hat_functor sort ty.

(** The Proof constructor *)
Definition Proof_map : HSET_over_sort2⟦Proof_source CCS,CCS⟧ :=
  (CoproductIn _ _ (Coproducts_functor_precat _ _ _ _ _) (● 2)%stn)
    · (BinCoproductIn2 (BinCoproducts_functor_precat _ _ _ _ _))
    · CCS_mor.

Definition lam_source (X : HSET_over_sort2) : HSET_over_sort2 :=
  ((X ∙ proj_functor sort ty) ⊗ (sorted_option_functor sort el ∙ X ∙ proj_functor sort el)) ∙ hat_functor sort el.

(** The lambda constructor *)
Definition lam_map : HSET_over_sort2⟦lam_source CCS,CCS⟧ :=
  (CoproductIn _ _ (Coproducts_functor_precat _ _ _ _ _) (● 3)%stn)
    · (BinCoproductIn2 (BinCoproducts_functor_precat _ _ _ _ _))
    · CCS_mor.

Definition app_source (X : HSET_over_sort2) : HSET_over_sort2 :=
  ((X ∙ proj_functor sort ty) ⊗
  ((sorted_option_functor sort el ∙ X ∙ proj_functor sort ty) ⊗
  ((X ∙ proj_functor sort el) ⊗
   (X ∙ proj_functor sort el)))) ∙ hat_functor sort el.

(** The app constructor *)
Definition app_map : HSET_over_sort2⟦app_source CCS,CCS⟧ :=
  (CoproductIn _ _ (Coproducts_functor_precat _ _ _ _ _) (● 4)%stn)
    · (BinCoproductIn2 (BinCoproducts_functor_precat _ _ _ _ _))
    · CCS_mor.

Definition forall_source (X : HSET_over_sort2) : HSET_over_sort2 :=
  ((X ∙ proj_functor sort ty) ⊗ (sorted_option_functor sort el ∙ X ∙ proj_functor sort el)) ∙ hat_functor sort el.

(** The ∀ constructor *)
Definition forall_map : HSET_over_sort2⟦forall_source CCS,CCS⟧ :=
  (CoproductIn _ _ (Coproducts_functor_precat _ _ _ _ _) (● 5)%stn)
    · (BinCoproductIn2 (BinCoproducts_functor_precat _ _ _ _ _))
    · CCS_mor.

Definition make_CCS_Algebra X
  (fvar    : HSET_over_sort2⟦Id,X⟧)
  (fPi     : HSET_over_sort2⟦Pi_source X,X⟧)
  (fProp   : HSET_over_sort2⟦Prop_source X,X⟧)
  (fProof  : HSET_over_sort2⟦Proof_source X,X⟧)
  (flam    : HSET_over_sort2⟦lam_source X,X⟧)
  (fapp    : HSET_over_sort2⟦app_source X,X⟧)
  (fforall : HSET_over_sort2⟦forall_source X,X⟧) : algebra_ob CCS_Functor.
Proof.
apply (tpair _ X).
use (BinCoproductArrow _ fvar).
use CoproductArrow.
intros i.
induction i as [n p].
repeat (induction n as [|n _]; try induction (nopathsfalsetotrue p)).
- exact fPi.
- exact fProp.
- exact fProof.
- exact flam.
- simpl in fapp.
  exact fapp.
- exact fforall.
Defined. (* This is slow *)

(** The recursor for ccs *)
Definition foldr_map X
  (fvar    : HSET_over_sort2⟦Id,X⟧)
  (fPi     : HSET_over_sort2⟦Pi_source X,X⟧)
  (fProp   : HSET_over_sort2⟦Prop_source X,X⟧)
  (fProof  : HSET_over_sort2⟦Proof_source X,X⟧)
  (flam    : HSET_over_sort2⟦lam_source X,X⟧)
  (fapp    : HSET_over_sort2⟦app_source X,X⟧)
  (fforall : HSET_over_sort2⟦forall_source X,X⟧) :
  algebra_mor _ CCS_alg (make_CCS_Algebra X fvar fPi fProp fProof flam fapp fforall).
Proof.
apply (InitialArrow CCS_Functor_Initial (make_CCS_Algebra X fvar fPi fProp fProof flam fapp fforall)).
Defined.

End ccs.