Require Import ZArith List Bool Lia.

Open Scope Z_scope.

Inductive bin :=
  L (x : Z) | N (t1 t2 : bin).

Definition leaf_value (t : bin) :=
match t with
  L x => Some x
| N t1 t2 => None
end.

Fixpoint Zb (t : bin) : Z :=
match t with
| L _ => 1
| N t1 t2 => Zb t1 + Zb t2
end.

Fixpoint insert_one_element (x : Z) (t : bin) :=
match t with
| L y => N (L y) (L x)
| N t1 t2 => N (insert_one_element x t2) t1
end.

Fixpoint list_to_bin (x : Z) (l : list Z) :=
match l with
| nil => L x
| y :: tl => insert_one_element x (list_to_bin y tl)
end.

Fixpoint bin_to_list (l : list Z) (t : bin) : list Z :=
match t with
| L x => x :: l
| N t1 t2 => bin_to_list  (bin_to_list  l t2) t1
end.

Lemma to_list_size l t : Zlength (bin_to_list l t) = Zlength l + Zb t.
Proof.
revert l; induction t as [ x | t1 Ih1 t2 Ih2].
  intros l.
  cbn [Zb bin_to_list].
  rewrite Zlength_cons.
  ring.
intros l.
cbn [Zb bin_to_list].
rewrite Ih1.
rewrite Ih2.
ring.
Qed.

Fixpoint count_in_list (x : Z) (l : list Z) : Z :=
match l with
| nil => 0
| y :: tl => (if x =? y then 1 else 0) + count_in_list x tl
end.

Fixpoint count_in_bin (x : Z) (t : bin) : Z :=
match t with
| L y => if x =? y then 1 else 0
| N t1 t2 => count_in_bin x t1 + count_in_bin x t2
end.

Lemma bin_to_list_preserves_counts (l : list Z) (t : bin) (x : Z) :
  count_in_list x (bin_to_list l t) = count_in_list x l + count_in_bin x t.
Proof.
revert l.
induction t as [y | t1 Ih1 t2 Ih2].
  intros l.
  cbn [count_in_list bin_to_list count_in_bin].
  ring.
intros l.
cbn [count_in_list bin_to_list count_in_bin].
rewrite Ih1.
rewrite Ih2.
ring.
Qed.

Definition to_bin_and_back (l : list Z) :=
  match l with nil => l | a :: tl => bin_to_list nil (list_to_bin a tl) end.

Compute (to_bin_and_back (1 :: 2 :: 3 :: 4 :: 5 :: 6 :: nil)).

Compute Z.iter 5 to_bin_and_back (1 :: 2 :: 3 :: 4 :: 5 :: 6 :: nil).

Check bin_ind.
Check list_ind.

Compute list_to_bin 0 (1 :: 2 :: 3 :: nil).

Fixpoint insert (a : Z) (l : list Z) :=
  match l with
    nil => a::nil
  | b::tl => if Zle_bool a b then a::b::tl else b::insert a tl
  end.

Fixpoint sort (l : list Z) :=
  match l with nil => nil | a::tl => insert a (sort tl) end.

(* 1/10 of a second for this computation. *)
Time Compute firstn 10
  (sort (snd (Z.iter 1000 (fun '(n, l) => (n + 1, n :: l)) (0, nil)))).

(* 4/10 of a second for this computation. 2000 elements *)
Time Compute firstn 10
  (sort (snd (Z.iter 2000 (fun '(n, l) => (n + 1, n :: l)) (0, nil)))).

Fixpoint merge_aux (a : Z) (l' : list Z)
  (f : list Z -> list Z) (l : list Z) :=
  match l with
  | b :: tl =>
    if Zle_bool a b then a :: f l else b :: merge_aux a l' f tl
  | nil => a :: l'
  end.

Fixpoint merge (l1 : list Z) : list Z -> list Z:=
  match l1 with
  | a :: l1' =>
    merge_aux a l1' (merge l1')
  | nil => fun l2 => l2
  end.

Fixpoint bin_sort (t : bin) :=
match t with
| L a => a :: nil
| N t1 t2 => merge (bin_sort t1) (bin_sort t2)
end.

Definition merge_sort (l : list Z) :=
match l with
| nil => nil
| a :: tl => bin_sort (list_to_bin a tl)
end.

Compute let l :=
   snd (Z.iter 1000 (fun '(n, l) => (n + 1, n :: l)) (0, nil))
  in
 (forallb (fun '(x1, x2) => x1 =? x2)
   (combine (merge_sort l) (sort l))).

Definition head_constraint (l1 l2 l : list Z) :=
  (exists l1', exists l2', exists l', exists a, exists b, exists c,
    l1 = a::l1' /\ l2 = b::l2' /\ l = c::l' /\
    (c = a \/ c = b)) \/
  (l1 = nil /\ exists l2', exists l', exists b, l2 = b::l2' /\ l = b::l') \/
  (l2 = nil /\ exists l1', exists l', exists a, l1 = a::l1' /\ l = a::l') \/
  (l1 = nil /\ l2 = nil /\ l = nil).

Fixpoint sorted (l : list Z) :=
match l with
| a :: (b :: _) as tl => (a <=? b) && sorted tl
| _ => true
end.

Lemma merge_sorted :
  forall l1 l2, sorted l1 = true -> sorted l2 = true ->
    sorted (merge l1 l2) = true /\
    head_constraint l1 l2 (merge l1 l2).
induction l1; intros l2.
 intros sn sl2;split;[exact sl2 | ].
 destruct l2 as [ | b l2'].
  right; right; right; repeat split; reflexivity.
 right; left;split;[reflexivity | exists l2'; exists l2'; exists b].
 split; reflexivity.
induction l2 as [ | b l2 IHl2].
 intros sal1 _; split.
  assumption.
 right; right; left; split;[reflexivity | ].
 exists l1; exists l1; exists a; split; reflexivity.
intros sal1 sbl2.
change (merge (a::l1) (b::l2)) with
  (if Zle_bool a b then a::merge l1 (b::l2) else b::merge (a::l1) l2).
case_eq (a <=? b)%Z.
 intros cab; split.
  destruct l1 as [ | a' l1].
   simpl merge.
   change ((if (a <=? b)%Z then sorted (b::l2) else false) = true).
   rewrite cab; assumption.
  assert (int: sorted (a'::l1) = true).
   simpl in sal1; destruct (a <=? a')%Z;[exact sal1 | discriminate].
  apply (IHl1 (b::l2)) in int.
   destruct int as [int1 int2].
   case_eq (merge (a'::l1) (b::l2)).
    intros q; rewrite q in int2.
    destruct int2 as [t2 | [t2 | [t2 | t2]]].
       destruct t2 as [abs1 [abs2 [abs3 [abs4 [abs5 [abs6 [_ [_ [abs9 _]]]]]]]]].
       discriminate.
      destruct t2 as [abs _]; discriminate.
     destruct t2 as [abs _]; discriminate.
    destruct t2 as [abs _]; discriminate.
   intros r l' qm.
   change ((if (a <=? r)%Z then sorted (r :: l') else false) = true).
   destruct int2 as [t2 | [t2 | [t2 | t2]]];
      try (destruct t2 as [abs _]; discriminate).
   destruct t2 as [l1' [l2' [l'' [u [v [w [q1 [q2 [qm' [H | H]]]]]]]]]].
    rewrite <- qm, int1.
    rewrite qm in qm'; injection qm';
     injection q1; injection q2; intros; subst.
    simpl in sal1; destruct (a <=? u)%Z; reflexivity || discriminate.
   rewrite qm in qm'; injection qm'; injection q1; injection q2.
   intros; subst; rewrite <- qm, int1, cab; reflexivity.
  assumption.
 left.
 exists l1; exists l2; exists (merge l1 (b :: l2)); exists a; exists b.
 exists a.
 repeat split; left; reflexivity.
intros cab.
assert (cab' : (b <=? a = true)%Z).
 apply Zle_imp_le_bool; rewrite Z.leb_nle in cab; lia.
assert (sl2 : sorted l2 = true).
 destruct l2 as [ | b' l2].
  reflexivity.
 simpl in sbl2; destruct (b <=? b')%Z; assumption || discriminate.
destruct (IHl2 sal1 sl2) as [int1 int2].
destruct l2 as [ | b' l2].
 split;[ | left; exists l1; exists nil; exists (merge (a::l1) nil); exists a;
          exists b; exists b; repeat split; right; reflexivity].
 destruct int2 as [t2 | [t2 | [ t2 | t2]]];
  try (destruct t2 as [abs _]; discriminate).
  destruct t2 as [ab1 [ab2 [ab3 [ab4 [ab5 [ab6 [_ [ab _]]]]]]]]; discriminate.
 destruct t2 as [_ [l1' [l'' [u [q qm]]]]].
 rewrite qm in int1 |- *.
 change ((if (b <=? u)%Z then sorted (u::l'') else false) = true).
 injection q; intros; subst.
 rewrite cab', int1; reflexivity.
destruct int2 as [t2 | [t2 | [t2 | t2]]];
  try (destruct t2 as [abs _]; discriminate).
destruct t2 as [l1' [l2' [l' [u [v [w [q [q' [qm [H | H]]]]]]]]]].
 injection q; injection q'; intros; subst.
 rewrite qm in *.
 split;[ | left; exists l1'; exists (v::l2'); exists (u::l');
              exists u; exists b; exists b; repeat split; right; reflexivity].
 change ((if (b <=? u)%Z then sorted (u::l') else false) = true).
 rewrite cab'; assumption.
injection q; injection q'; intros; subst.
rewrite qm in *.
split;[ | left; exists l1'; exists (v::l2'); exists (v::l');
              exists u; exists b; exists b; repeat split; right; reflexivity].
change ((if (b <=?v)%Z then sorted (v :: l') else false) = true).
simpl in sbl2; destruct (b <=? v)%Z; assumption || discriminate.
Qed.


(* Sorting this list takes 12 seconds with insertion sort, and 0.05 seconds
  with merge_sort*)
Time Compute firstn 10 (merge_sort
   (snd (Z.iter 10000 (fun '(n, l) => (n + 1, n :: l)) (0, nil)))).

Lemma bin_sort_sorted : forall t, sorted (bin_sort t) = true.
Proof.
induction t.
  reflexivity.
cbn [sorted merge bin_sort].
assert (tmp := merge_sorted (bin_sort t1) (bin_sort t2) IHt1 IHt2).
destruct tmp as [tmp' _]; assumption.
Qed.

Lemma merge_sort_sorted : forall l, sorted (merge_sort l) = true.
Proof.
intros [ | a l].
 reflexivity.
unfold merge_sort; apply bin_sort_sorted.
Qed.