Require Export NuTree.

Fixpoint NuNumber nt : nat :=
  match nt with
  | Tr _ l => list_sum (List.map NuNumber l)
  | Nu _ x => S ( NuNumber x)
  end.

Fact no_nus_nu_number_z nt :
  NoNus nt -> NuNumber nt = 0.
Proof.
  induction 1 as [a l H1 H2]. simpl. apply list_sum_zero.
  intros x [s [<- H3]] % in_map_iff. apply H2, H3.
Qed.

Fact no_nus_nu_number_z' nt :
  Ntwr nt -> NoNus nt <-> NuNumber nt = 0.
Proof.
  intros HW1. split.
  - apply no_nus_nu_number_z.
  - intros H1. induction HW1.
    + constructor. intros x H3. apply H0; trivial.
      simpl in H1. rewrite list_sum_zero in H1. apply H1, in_map_iff. eauto.
    + simpl in H1. discriminate H1.
Qed.

(*
  Embedding Tree in NuTree
 *)
Fixpoint t_to_nt t : NuTree :=
  match t with TrA a l => Tr a (map t_to_nt l) end.

Notation "'%' t" := (t_to_nt t) (at level 65).

Fact pure_no_nus t :
  Twr t -> NoNus (% t).
Proof.
  induction 1 as [a l H1 H2 _]. simpl. constructor.
  intros x [s [<- H3]] % in_map_iff. apply H2, H3.
Qed.

(* The number of nu nodes in a pure wellranked tree is zero *)
Fact pure_nu_number_z t :
  Twr t -> NuNumber (% t) = 0.
Proof.
  now intros H1 % pure_no_nus % no_nus_nu_number_z.
Qed.

Lemma t_to_nt_in l n :
  n el map t_to_nt l -> exists t, t el l /\ n = % t.
Proof.
  rewrite in_map_iff. intros [t [H2 H3]]. now exists t.
Qed.

(*
  The translation of a wellranked tree to a nu-tree is also wellranked
 *)
Fact twr_ntwr t :
  Twr t -> Ntwr (% t).
Proof.
  induction 1 as [a l H1 H2 H3]. simpl. constructor; try now rewrite map_length.
  intros s [s' [H4 ->]] % t_to_nt_in. apply H2, H4.
Qed.

(*
  When translating a tree t into a nu-tree, t is in the language of the translation
 *)
Fact ntl_refl t A:
  Twr t -> t El [% t] A.
Proof.
  intros H1. revert A. induction H1 as [a l _ H3 _]. intros A.
  constructor; fold t_to_nt; auto using map_length.
  intros k H5. rewrite map_length in H5.
  assert(H6: nth k l (DTR) el l) by apply nth_In, H5.
  rewrite nth_indep with (d':= % DTR). rewrite map_nth.
  apply (H3 (nth k l (DTR)) H6). now rewrite map_length.
Qed.

Section Equality.

  Variable n : NuTree.
  Variable n' : NuTree.
  Hypothesis HW1 : Ntwr n.
  Hypothesis HW2 : Ntwr n'.

  (* Deciding equality on wellranked nu-trees *)

  Inductive Nteq : NuTree -> NuTree -> Prop :=
  | NteqA a l l' : |l| = |l'| -> |l| = rk a -> (forall k, k < |l| -> Nteq (nth k l DNT) (nth k l' DNT)) -> Nteq (Tr a l) (Tr a l')
  | NteqB a n n' : Nteq n n' -> Nteq (Nu a n) (Nu a n').

  Fixpoint nt_eq n n' : bool :=
    match n, n' with
    | Tr a l, Tr b l' =>
      if Dec (a=b)
      then all_true (map_pw (map nt_eq l) l')
      else false
    | Nu a m, Nu b m' =>
      if Dec (a=b)
      then nt_eq m m'
      else false
    | _, _ => false
    end.

  Fact nt_eq_correct_1 :
    nt_eq n n' = true -> n = n'.
  Proof.
    revert HW2. revert n'.
    induction HW1 as [a l H1 H2| a n H1 H2]; intros n' HW2;
      induction HW2 as [b l' H4 H5| b n' H4 H5]; simpl; try discriminate 1;
        decide (a=b) as [E|]; try discriminate 1; subst b.
    - rewrite all_true_map_pw_map with (x:=DNT)(y:=DNT); try congruence.
      intros H6. f_equal. apply plist_eq with (x:=DNT); try congruence. intros k H7. apply H2; auto. apply H4, nth_In. congruence.
    - intros H6. f_equal. apply H2; auto.
  Qed.

  Fact nt_eq_correct_2 :
    n = n' -> nt_eq n n' = true.
  Proof.
    intros. subst n'. induction HW1; simpl; decide (a=a) as [_|]; auto.
    rewrite all_true_map_pw_map with (x:=DNT)(y:=DNT); auto.
  Qed.

  Fact nt_eq_correct :
    nt_eq n n' = true <-> n = n'.
  Proof.
    split; eauto using nt_eq_correct_1, nt_eq_correct_2.
  Qed.

  Instance nt_eq_dec :
    dec(n = n').
  Proof.
    destruct (nt_eq n n') eqn:E.
    - left. apply nt_eq_correct; assumption.
    - right. intros H3. apply nt_eq_correct in H3; try assumption. congruence.
  Qed.

End Equality.