(** This greatly seconds "TakeDrop" module, but for streams instead of lists *)

Set Implicit Arguments.
Require Coq.Arith.Le.
Require Import Coq.Lists.List.
Require Import Coq.Lists.Streams.
Require Beroal.Init.
Require Import Beroal.TakeDrop.

Section element.
Variable e : Type.

Definition drop_pres_eq : forall n (xs ys : Stream e),
	EqSt xs ys -> EqSt (Str_nth_tl n xs) (Str_nth_tl n ys).
Proof.
refine (fix rec (n : nat) xs ys eq0 := _).
destruct n as [|pn].
simpl.
	refine (eq0).
simpl.
	destruct eq0 as [p0 p1].
	refine (rec _ _ _ p1).
Qed.

Definition In (x : e) := Exists (fun xs => x = hd xs).
Definition non_dup := ForAll (fun xs => ~ In (hd xs) (tl xs)).

Definition filtered_finite (p : e -> Prop)
	:= Exists (ForAll (fun ys => ~ p (hd ys))).
Definition filtered_infinite (p : e -> Prop)
	:= ForAll (Exists (fun ys => p (hd ys))).

Definition app xs0 (ys : Stream e) := (cofix rec (xs : list e)
	:= match xs with
	| nil => ys
	| List.cons x xss => Streams.Cons x (rec xss)
	end) xs0.

Definition take := fix rec (n : nat) (xs : Stream e)
	:= match n with
	| O => nil
	| S pn => List.cons (hd xs) (rec pn (tl xs)) 
	end.

Definition take_pres_eq : forall xs ys,
	EqSt xs ys -> forall n, take n xs = take n ys.
Proof.
refine (fix rec xs ys eq0 (n : nat) {struct n} := _).
destruct n as [|pn].
simpl.
	refine (refl_equal _).
simpl.
	destruct eq0 as [p0 p1].
	refine (eq_ind _ (fun dep => _ = dep :: _) _ _ p0).
	refine (f_equal _ (rec _ _ p1 _)).
Qed.

(** permuting take and tail *)
Check fun n xs => (refl_equal _) : take n (tl xs) = List.tail (take (S n) xs).

(** obvious simplifications *)
Check fun x xs ys => conj (refl_equal _)
		(conj (refl_equal _)
		(conj (refl_equal _)
		(refl_equal _)))
	: hd ys = hd (app nil ys)
	/\ tl ys = tl (app nil ys)
	/\ x = hd (app (x :: xs) ys)
	/\ app xs ys = tl (app (x :: xs) ys).

Definition app_nil_id : forall xs, EqSt xs (app nil xs).
Proof.
refine (fun xs => eqst _ _ _ _).
refine (refl_equal _).
simpl.
refine (EqSt_reflex _).
Qed.

Definition app_comm_cons : forall x xs ys,
	EqSt (Streams.Cons x (app xs ys)) (app (x :: xs) ys).
Proof.
refine (fun x xs ys => eqst _ _ _ _).
refine (refl_equal _).
refine (EqSt_reflex _).
Qed.

Definition app_single_cons : forall x xs,
	EqSt (Streams.Cons x xs) (app (x :: nil) xs).
Proof.
refine (fun x xs => eqst _ _ _ _).
refine (refl_equal _).
refine (eqst _ _ _ _).
refine (refl_equal _).
refine (EqSt_reflex _).
Qed.

Definition app_pres_eq : forall xs ys, EqSt xs ys
	-> forall zs, EqSt (app zs xs) (app zs ys).
Proof.
refine (fun xs ys eq0 => fix rec (zs : list e) := _).
destruct zs as [|z zss].
destruct eq0 as [eq1 eq2].
	refine (eqst (app nil xs) (app nil ys) eq1 eq2).
refine (eqst _ _ _ _).
	refine (refl_equal _).
	refine (rec zss).
Qed.

Definition take_app_degen : forall n xs ys, n <= length xs
	-> TakeDrop.take n xs = take n (app xs ys).
Proof.
refine (fix rec n (xs : list e) ys {struct xs} := _).
destruct xs as [|x xss].
simpl.
	refine (fun le0 => _).
	rewrite <- (Coq.Arith.Le.le_n_O_eq _ le0).
	simpl.
	refine (refl_equal _).
simpl.
destruct n as [|pn].
	simpl.
		refine (fun _ => refl_equal _).
	refine (fun le0 => _).
		refine (f_equal (List.cons x)
			(rec _ _ _ (Coq.Arith.Le.le_S_n _ _ le0))).
Qed.

Definition drop_app_degen : forall n xs ys, n <= length xs
	-> EqSt (app (TakeDrop.drop n xs) ys)
		(Str_nth_tl n (app xs ys)).
Proof.
refine (fix rec n (xs : list e) ys {struct xs} := _).
destruct xs as [|x xss].
simpl.
	refine (fun le0 => _).
	rewrite <- (Coq.Arith.Le.le_n_O_eq _ le0).
	simpl.
	refine (EqSt_reflex _).
simpl.
	destruct n as [|pn].
	simpl.
		refine (fun _ => EqSt_reflex _).
	refine (fun le0 => _).
		refine (rec _ _ _ (Coq.Arith.Le.le_S_n _ _ le0)).
Qed.

Definition take_app_part : forall n ys xs,
	xs ++ (take n ys) = take (length xs + n) (app xs ys).
Proof.
refine (fun n ys => fix rec (xs : list e) := _).
destruct xs as [|x xss].
refine (take_pres_eq _ _).
	refine (app_nil_id _).
refine (f_equal (List.cons x) (rec xss)).
Qed.

Definition drop_app_part : forall n ys xs,
	EqSt (Str_nth_tl n ys)
		(Str_nth_tl (length xs + n) (app xs ys)).
Proof.
refine (fun n ys => fix rec (xs : list e) := _).
destruct xs as [|x xss].
refine (drop_pres_eq _ _).
	refine (app_nil_id _).
refine (rec xss).
Qed.

Definition take_app_exact : forall xs ys,
	xs = take (length xs) (app xs ys).
Proof.
refine (fun xs ys => _).
refine (eq_ind _ (fun dep => _ = take dep _)
	_ _ (Coq.Arith.Plus.plus_0_r _)).
refine (eq_ind _ (fun dep => dep = _)
	_ _ (sym_eq (app_nil_end _))).
refine (take_app_part 0 _ _).
Qed.

Definition drop_app_exact : forall xs ys,
	EqSt ys (Str_nth_tl (length xs) (app xs ys)).
Proof.
refine (fun xs ys => _).
refine (eq_ind _ (fun dep => EqSt _ (Str_nth_tl dep _))
	_ _ (Coq.Arith.Plus.plus_0_r _)).
refine (drop_app_part 0 _ _).
Qed.

Definition take_app_drop : forall n xs,
	EqSt xs (app (take n xs) (Str_nth_tl n xs)).
Proof.
refine (fix rec (n : nat) xs := _).
destruct n as [|pn].
simpl.
	refine (app_nil_id _).
refine (eqst _ _ _ _).
	simpl.
		refine (refl_equal _).
	refine (rec pn (tl xs)).
Qed.

Definition drop_plus : forall n m xs,
	EqSt (Str_nth_tl n xs)
		(app (take m (Str_nth_tl n xs)) (Str_nth_tl (n + m) xs)).
Proof.
refine (fix rec (n : nat) m xs {struct n} := _).
destruct n as [|pn].
simpl.
	refine (take_app_drop _ _).
simpl.
	refine (rec pn m (tl xs)).
Qed.

Definition take_plus : forall n m xs,
	take (n + m) xs = take n xs ++ take m (Str_nth_tl n xs).
Proof.
refine (fix rec (n : nat) m xs {struct n} := _).
destruct n as [|pn].
simpl.
	refine (refl_equal _).
simpl.
	refine (f_equal _ (rec pn m (tl xs))).
Qed.



Definition iterate (f : e -> e) := cofix rec (x : e)
	:= (Streams.Cons x (rec (f x))).
Definition cycle (x : e) (xs : list e)
	:= cofix rec := Streams.Cons x (app xs rec).
Definition cyclic (period : nat) (xs : Stream e)
	:= EqSt xs (Str_nth_tl period xs).

Definition repeat (x : e)
	:= cofix rec := Streams.Cons x rec.
Definition replicate n (x : e) := take n (repeat x).

Definition iterate_id_cyclic : forall x n,
	cyclic n (iterate (@Beroal.Init.id _) x).
Proof.
refine (fun x => _).
refine (fix rec (n : nat) := _).
destruct n as [|pn].
refine (EqSt_reflex _).
unfold cyclic.
	refine (rec pn).
Qed.

Definition iterate_id_nth : forall x n,
	x = Str_nth n (iterate (@Beroal.Init.id _) x).
Proof.
refine (fun x n => match iterate_id_cyclic x n with eqst p _ => p end).
Qed.

Definition cycle_single_cyclic : forall x n, cyclic n (cycle x nil).
Proof.
refine (fun x => _).
refine (fix rec (n : nat) := _).
destruct n as [|pn].
refine (EqSt_reflex _).
unfold cyclic.
	simpl.
	refine (trans_EqSt (rec pn) (drop_pres_eq _ (app_nil_id _))).
Qed.

Definition cycle_single_nth : forall x n,
	x = Str_nth n (cycle x nil).
Proof.
refine (fun x n => match cycle_single_cyclic x n with eqst p _ => p end).
Qed.

Definition iterate_id_eq_cycle_single : forall x,
	EqSt (iterate (@Beroal.Init.id _) x) (cycle x nil).
Proof.
refine (fun x => _).
refine (ntheq_eqst _ _ _).
refine (fun n => _).
unfold Str_nth.
refine (trans_eq (sym_eq (iterate_id_nth _ _)) _).
refine (trans_eq _ (cycle_single_nth _ _)).
simpl.
refine (refl_equal _).
Qed.

Definition unfold_cycle : forall x xs,
	EqSt ((app (x :: xs) (cycle x xs))) (cycle x xs).
Proof.
refine (fun x xs => _).
refine (eqst _ _ _ _).
simpl.
refine (refl_equal _).
simpl.
refine (EqSt_reflex _).
Qed.

Definition take_cycle_id : forall x xs,
	x :: xs = take (S (length xs)) (cycle x xs).
Proof.
refine (fun x xs => take_app_exact _ _).
Qed.

Definition cycle_cyclic : forall x xs, cyclic (S (length xs)) (cycle x xs).
Proof.
refine (fun x xs => _).
unfold cyclic.
refine (match @drop_app_exact xs (cycle x xs) with
	| eqst p0 p1 => _ end).
refine (eqst _ _ p0 p1).
Qed.

Definition cyclic_tl : forall period xs,
	cyclic period xs -> cyclic period (tl xs).
Proof.
unfold cyclic.
refine (fun period xs eq0 => _).
destruct eq0 as [p0 p1].
refine (eq_ind _ (fun dep => EqSt _ dep) p1 _ (tl_nth_tl _ _)).
Qed.

Definition cyclic_drop : forall period xs,
	cyclic period xs -> forall n, cyclic period (Str_nth_tl n xs).
Proof.
refine (fun period xs eq0 => fix rec (n : nat) := _).
destruct n as [|pn].
refine (eq0).
simpl.
refine (eq_ind _ (cyclic period) (cyclic_tl (rec _)) _ (tl_nth_tl _ _)).
Qed.

Definition cyclic_hd_eq : forall xs ys,
	cyclic 1 xs -> cyclic 1 ys -> hd xs = hd ys -> EqSt xs ys.
Proof.
unfold cyclic.
simpl.
refine (cofix rec xs ys cx cy eq0 : EqSt _ _ := _).
refine (eqst _ _ _ _).
refine (eq0).
destruct cx as [cx0 cx1].
destruct cy as [cy0 cy1].
refine (rec _ _ cx1 cy1 (trans_eq (sym_eq cx0) (trans_eq eq0 cy0))).
Qed.

Definition eqst_many : forall n xs ys, take n xs = take n ys
	-> EqSt (Str_nth_tl n xs) (Str_nth_tl n ys) -> EqSt xs ys.
Proof.
refine (cofix rec (n : nat) (xs ys : Stream e) : _ -> _ -> EqSt _ _ := _).
destruct n as [|pn].
refine (fun _ p => p).
simpl.
refine (fun eq0 eq1 => _).
refine (eqst _ _
	(f_equal (@List.hd _ (hd xs)) eq0)
	(rec _ _ _ (f_equal (@List.tail _) eq0) eq1)).
Qed.


(* cycle_uncurry roth_ex cycle_rotate_tl cycle_rotate_drop; rep (length xs) roth = id *)
Definition cycle_rotate_tl_app_nil : forall x0 x1 xs,
	EqSt (tl (cycle x0 (x1 :: xs)))
		(app nil (cycle x1 (xs ++ (x0 :: nil)))).
Proof.
refine (fun x0 x1 xs => _).
refine (cofix recx : EqSt _ _ := _).
refine (eqst _ _ _ _).
simpl.
	refine (refl_equal _).
refine ((_ : forall ys, EqSt (app ys (cycle x0 (x1 :: xs)))
	((app (ys ++ x0 :: nil) (cycle x1 (xs ++ x0 :: nil))))
	) xs).
refine (cofix recy (ys : list e)
	: EqSt (app ys (cycle x0 (x1 :: xs)))
		((app (ys ++ x0 :: nil) (cycle x1 (xs ++ x0 :: nil))))
	:= _).
destruct ys as [| y0 yss].
refine (eqst _ _ _ _).
	simpl.
		refine (refl_equal _).
	refine (recx).
refine (eqst _ _ _ _).
	simpl.
		refine (refl_equal _).
	refine (recy yss).
Qed.

End element.

Require Coq.Arith.Plus.
Section Stream_nat.

Definition enum := iterate S.

Definition enum_nth n s : Str_nth n (enum s) = n + s.
Proof.
refine (fix rec (n : nat) {struct n} := fun s => _).
destruct n as [|pn].
unfold Str_nth.
simpl.
refine (refl_equal _).
unfold Str_nth.
simpl.
change (Str_nth pn (enum (S s)) = S pn + s).
refine (trans_eq (rec _ _) (sym_eq (Coq.Arith.Plus.plus_Snm_nSm _ _))).
Qed.

Definition enum_step step := iterate (plus step).

Definition enum_step_nth step n start0
: Str_nth n (enum_step step start0) = n * step + start0.
Proof.
refine (fun step => fix rec (n : nat) {struct n} := fun start0 => _).
destruct n as [|pn].
unfold Str_nth.
simpl.
refine (refl_equal _).
change (Str_nth pn (enum_step step (step + start0))
= step + pn * step + start0).
refine (trans_eq (rec _ _)
(trans_eq (Coq.Arith.Plus.plus_permute _ _ _)
(Coq.Arith.Plus.plus_assoc _ _ _))).
Qed.

End Stream_nat.

Definition zip_with (e0 e1 e2 : Type) (f : e0 -> e1 -> e2)
:= fix rec (xs : list e0) (ys : Stream e1) := match xs with
	| nil => nil
	| List.cons x xs => match ys with
		| Streams.Cons y ys => List.cons (f x y) (rec xs ys) end
	end.

Definition zip (e0 e1 : Type) := @zip_with e0 e1 _
	(fun x0 x1 => (x0, x1)).