Library Containers.W_From_Nat

Require Export HoTT.
Require Export Containers.Initial_Algebra.
Require Export Containers.Container.

Axiom propresize :
   (A : Type) {H : IsHProp A}, Type.
Axiom equiv_propresize :
   (A : Type) {H : IsHProp A}, propresize A <~> A.
Axiom ishprop_propresize :
   (X : Type) (H : IsHProp X), IsHProp (propresize X).
Existing Instance ishprop_propresize.

Variable A : Type.
Variable B : A Type.
Context {funext : Funext}.

Class is_tree_type (T : Type) := {
  sup : label : A, (B label T) T;
  label : T A;
  arg : tree : T, B (tree.(label)) T
}.

Section Tree_Type.

Context {T : Type}.
Context {is_tt : is_tree_type T}.

Fixpoint Addr' (t : T) (depth : nat) :=
match depth with
  | 0 ⇒ Unit
  | depth.+1{ b : B t.(label) & Addr' (t.(arg) b) depth }
end.

Definition Addr (t : T) :=
{ depth : nat & Addr' t depth }.

Definition root_addr {t : T} : Addr t :=
(0; tt).

Definition subtree_addr {t : T}
    (b : B t.(label)) (addr' : Addr (t.(arg) b)) :
  Addr t :=
let (depth, addr') := addr' in
(depth.+1; (b; addr')).

Definition Addr_rect (P : t, Addr t Type)
    (root_case : t, P t root_addr)
    (subtree_case : t b addr',
      P (t.(arg) b) addr' P t (subtree_addr b addr'))
    (t : T) (addr : Addr t) :
  P t addr.
Proof.
  destruct addr as [depth addr'].
  revert t addr'; induction depth; intros t addr'.
  - destruct addr'.
    apply root_case.
  - destruct addr' as [b addr'].
    apply (subtree_case t b (depth; addr')).
    apply IHdepth.
Defined.

Definition subtree_at (t : T) (addr : Addr t) : T.
Proof.
  revert t addr; apply Addr_rect.
  - exact idmap.
  - intros _ _ _.
    exact idmap.
Defined.

Definition label_at (t : T) (addr : Addr t) : A :=
(t.(subtree_at) addr).(label).

Definition extend_addr {t : T}
    (addr : Addr t) (b : B (t.(label_at) addr)) : Addr t.
Proof.
  revert t addr b.
  erapply Addr_rect.
  - exact (fun t bsubtree_addr b root_addr).
  - exact (fun t b addr' IHaddr b'subtree_addr b (IHaddr b')).
Defined.

Lemma subtree_at_extend_addr {t : T}
    (addr : Addr t) :
  t.(subtree_at) o extend_addr addr = arg (t.(subtree_at) addr).
Proof.
  revert t addr.
  erapply Addr_rect.
  - reflexivity.
  - intros t b addr' IH.
    exact IH.
Defined.

Lemma equiv_addr_match t P :
  (P root_addr × (b : B t.(label)) addr, P (subtree_addr b addr)) <~>
  ( addr, P addr).
Proof.
  serapply equiv_adjointify.
  - intros [root_case subtree_case].
    intros addr.
    revert t addr P root_case subtree_case.
    erapply Addr_rect.
    + intros t P root_case subtree_case.
      exact root_case.
    + intros t b addr' _ P root_case subtree_case.
      exact (subtree_case b addr').
  - exact (fun f(f root_addr, fun b addr'f (subtree_addr b addr'))).
  - intros f.
    apply path_forall; intros addr.
    revert t addr P f.
    erapply Addr_rect.
    + intros.
      reflexivity.
    + intros.
      reflexivity.
  - intros [root_x arg_f].
    reflexivity.
Defined.

End Tree_Type.

Arguments root_addr {_ _ _}.
Arguments subtree_addr {_ _ _} _ _.
Arguments extend_addr {_ _ _} _ _.

Axiom M : Type.
Axiom m_out : M <~> {a : A & B a M}.
Definition m_sup a f := m_out^-1 (a; f).
Definition m_label := pr1 o m_out.
Definition m_arg := pr2 o m_out.
Lemma m_label_sup : a f, m_label (m_sup a f) = a.
Proof.
  intros a f.
  unfold m_label, m_sup.
  exact (eisretr m_out (a; f))..1.
Defined.
Lemma m_arg_sup : a f b,
  (m_arg (m_sup a f)) b =
  f (m_label_sup a f # b).
Proof.
  intros a f b.
  unfold m_arg, m_sup.
  rewrite <- (inv_V (m_label_sup a f)).
  rewrite <- transport_arrow_toconst.
  revert b; apply ap10.
  apply (moveL_transport_V (fun aB a M)).
  unfold m_label_sup.
  exact (eisretr m_out (a; f))..2.
Qed.
Instance is_tree_type_M : is_tree_type M :=
  Build_is_tree_type M m_sup m_label m_arg.

Definition is_wf (m : M) : Type :=
P : M hProp, ( a f, ( b, P (f b)) P (sup a f)) P m.

Definition W : Type :=
{ m : M & is_wf m }.

Definition w_sup (a : A) (f : B a W) : W.
Proof.
   (m_sup a (pr1 o f)).
  intros P step.
  apply step.
  intros b.
  apply (f b).2.
  exact step.
Defined.

Definition w_label (w : W) : A :=
(w.1).(label).

Definition is_wf_m_arg {m : M} (wf : is_wf m) (b : B m.(label)) :
  is_wf (m.(arg) b).
Proof.
  unfold is_wf in ×.
  intros P step.
  revert b.
  apply (wf (fun mBuildhProp _)).
  intros a f IH b.
  simpl rewrite m_arg_sup.
  rewrite <- (eissect m_out (f (transport B (m_label_sup a f) b))).
  apply step.
  apply IH.
Defined.

Definition w_arg (w : W) (b : B w.(w_label)) : W.
Proof.
  destruct w as [m wf].
   (m_arg m b).
  apply is_wf_m_arg.
  exact wf.
Defined.

Instance is_tree_type_W : is_tree_type W :=
Build_is_tree_type W w_sup w_label w_arg.

Section Recursor.

Variable (C : Algebra (| A ||> B |)).
Definition X := alg_carrier C.
Definition step a f := alg_fun C (a; f).

Definition Local_Recursor (w : W) :=
{f : Addr w X &
   addr, f addr = step (w.(label_at) addr) (f o extend_addr addr)}.

Lemma equiv_arg_recursor (w : W) :
  Local_Recursor w <~> b, Local_Recursor (w.(arg) b).
Proof.
  unfold Local_Recursor.
  transitivity {
      f : Addr w X &
         addr, f addr =
          step (label_at w addr) (fun xf (extend_addr addr x))
    }. {
    apply reflexivity.
  }
  transitivity {
      xf : X × b, Addr (w.(arg) b) X &
        (xf.(fst) = step (w.(label)) (fun bxf.(snd) b root_addr)) ×
        ( b addr, xf.(snd) b addr =
          step (label_at (w.(arg) b) addr)
               (fun b'xf.(snd) b (extend_addr addr b')))
    }. {
    serapply equiv_functor_sigma'.
    - symmetry.
      rapply equiv_addr_match.
    - simpl.
      intros f.
      symmetry.
      exact (equiv_addr_match w _).
  }
  transitivity {
      xf : {x : X & b, Addr (w.(arg) b) X} &
        (xf.1 = step (w.(label)) (fun bxf.2 b root_addr)) ×
        ( b addr, xf.2 b addr =
          step (label_at (w.(arg) b) addr)
               (fun b'xf.2 b (extend_addr addr b')))
    }. {
    serapply equiv_functor_sigma'.
    - symmetry.
      apply equiv_sigma_prod0.
    - simpl.
      intros xf.
      apply reflexivity.
  }
  transitivity {
      x : X & {
        f : b, Addr (w.(arg) b) X &
          (x = step (w.(label)) (fun bf b root_addr)) ×
          ( b addr, f b addr =
            step (label_at (w.(arg) b) addr)
                 (fun b'f b (extend_addr addr b')))
      }
    }. {
    symmetry.
    apply (
        equiv_sigma_assoc (fun _ b, Addr (w.(arg) b) X)
        (fun xf(xf.1 = step (w.(label)) (fun bxf.2 b root_addr)) ×
          ( b addr, xf.2 b addr =
            step (label_at (w.(arg) b) addr)
               (fun b'xf.2 b (extend_addr addr b'))))).
  }
  transitivity {
      x : X & {
        f : b, Addr (w.(arg) b) X & {
          _ : x = step (w.(label)) (fun bf b root_addr) &
             b addr, f b addr =
              step (label_at (w.(arg) b) addr)
                   (fun b'f b (extend_addr addr b'))
        }
      }
    }. {
    apply equiv_functor_sigma_id; intros x.
    apply equiv_functor_sigma_id; intros f.
    symmetry.
    apply equiv_sigma_prod0.
  }
  apply (transitivity (equiv_sigma_symm _)).
  transitivity {
      f : b, Addr (w.(arg) b) X & {
        _ : {x : X & x = step (w.(label)) (fun bf b root_addr)} &
           b addr, f b addr =
            step (label_at (w.(arg) b) addr)
                 (fun b'f b (extend_addr addr b'))
      }
    }. {
    apply equiv_functor_sigma_id; intros f.
    erapply equiv_sigma_assoc.
  }
  transitivity {
      f : b, Addr (w.(arg) b) X & {
        _ : Unit &
           b addr, f b addr =
            step (label_at (w.(arg) b) addr)
                 (fun b'f b (extend_addr addr b'))
      }
    }. {
    apply equiv_functor_sigma_id; intros f.
    apply equiv_functor_sigma'.
    - apply equiv_contr_contr.
    - intros _.
      exact equiv_idmap.
  }
  transitivity {
      f : b, Addr (w.(arg) b) X &
         b addr, f b addr =
          step (label_at (w.(arg) b) addr)
               (fun b'f b (extend_addr addr b'))
    }. {
    apply (equiv_functor_sigma' equiv_idmap); intros f.
    erapply equiv_contr_sigma.
  }
  erapply equiv_sigT_coind.
Qed.

Instance contr_local_recursor (w : W) :
  Contr (Local_Recursor w).
Proof.
  destruct w as [m wf].
  generalize wf as wf'.
  erapply equiv_propresize.
  apply (wf (fun mBuildhProp _)). clear.
  intros a f IH.
  apply equiv_propresize.
  intros wf.
  apply (contr_equiv' ( b, Local_Recursor ((sup a f; wf).(arg) b))).
  - exact (equiv_inverse (equiv_arg_recursor (sup a f; wf))).
  - serapply contr_forall; intros b; simpl.
    unfold w_arg.
    rewrite (path_sigma' is_wf (m_arg_sup a f b) idpath).
    change (
        (fun wf'Contr (Local_Recursor
          (f (transport B (m_label_sup a f) b); wf')))
        (transport is_wf (m_arg_sup a f b) (is_wf_m_arg wf b))).
    generalize (transport is_wf (m_arg_sup a f b) (is_wf_m_arg wf b))
        as wf'.
    erapply equiv_propresize.
    apply IH.
Defined.

Definition local_recursor (w : W) :=
center (Local_Recursor w).

Definition Recursor :=
{f : W X & w, f w = step w.(label) (f o w.(arg))}.

Definition restrict (h : Recursor) (w : W) :
  Local_Recursor w.
Proof.
   (h.1 o w.(subtree_at)).
  intros addr.
  apply (transitivity (h.2 (w.(subtree_at) addr))).
  apply (ap (fun fstep (w.(label_at) addr) (h.1 o f))).
  symmetry.
  apply subtree_at_extend_addr.
Defined.

Goal h w, (restrict h w).2 root_addr = h.2 w.
Proof.
  intros.
  simpl.
  apply concat_p1.
Qed.

Definition arg_recursor {w : W}
    (h : Local_Recursor w) (b : B w.(label)) :
  Local_Recursor (w.(arg) b) :=
(h.1 o subtree_addr b; h.2 o subtree_addr b).

Definition lift (h : w, Local_Recursor w) :
  Recursor.
Proof.
   (fun w(h w).1 root_addr).
  intros w.
  apply (transitivity ((h w).2 root_addr)).
  apply (ap (step w.(label))).
  apply apD10^-1; intros b.
  change ((arg_recursor (h w) b).1 root_addr =
          (h (arg w b)).1 root_addr).
  apply (ap (fun h' : Local_Recursor _h'.1 root_addr)).
  apply path_contr.
Defined.

Lemma restrict_is_sect :
  Sect restrict lift.
Proof.
  intros h.
  serapply path_sigma'.
  - reflexivity.
  - apply path_forall; intros w'.
    unfold transport, transitivity, transitive_paths.
    simpl.
    rewrite concat_p1.
    apply moveR_Mp.
    rewrite concat_Vp.
    transitivity (ap (step (w_label w'))
    (apD10^-1
       (fun b
        ap (fun h'h'.1 root_addr)
          (@idpath _ (restrict h (w_arg w' b)))))). {
      repeat f_ap. apply path_forall; intros b. f_ap.
      apply path_contr.
    }
    simpl.
    change (ap (step w'.(label)) (apD10^-1 (apD10 (@idpath _
        (h.1 o subtree_at w' o extend_addr root_addr)))) = idpath).
    rewrite eissect.
    reflexivity.
Qed.

Instance contr_recursor :
  Contr Recursor.
Proof.
   (lift local_recursor).
  intros r.
  rewrite <- restrict_is_sect.
  f_ap.
  apply path_contr.
Defined.

Definition recursor :=
center Recursor.

End Recursor.