End Zip.
Prenex Implicits zip unzip1 unzip2.
Section Flatten.
Variable T : Type.
Implicit Types (s : seq T) (ss : seq (seq T)).
Definition flatten := foldr cat (Nil T).
Definition shape := map (@size T).
Fixpoint reshape sh s :=
if sh is n :: sh' then take n s :: reshape sh' (drop n s) else [::].
Definition flatten_index sh r c := sumn (take r sh) + c.
Definition reshape_index sh i := find (pred1 0) (scanl subn i.+1 sh).
Definition reshape_offset sh i := i - sumn (take (reshape_index sh i) sh).

(* Specifications: *)
Lemma size_flatten ss : size (flatten ss) = sumn (shape ss).
Lemma flattenK ss : reshape (shape ss) (flatten ss) = ss.
Lemma reshapeKr sh s : size s <= sumn sh -> flatten (reshape sh s) = s.
Lemma reshapeKl sh s : size s >= sumn sh -> shape (reshape sh s) = sh.
Lemma flatten_rcons ss s : flatten (rcons ss s) = flatten ss ++ s.
Lemma nth_shape ss i : nth 0 (shape ss) i = size (nth [::] ss i).
Lemma shape_rev ss : shape (rev ss) = rev (shape ss).
Lemma reshape_rcons s sh n (m := sumn sh) :
m + n = size s ->
reshape (rcons sh n) s = rcons (reshape sh (take m s)) (drop m s).
Lemma flatten_indexP sh r c :
c < nth 0 sh r -> flatten_index sh r c < sumn sh.
Lemma reshape_indexP sh i : i < sumn sh -> reshape_index sh i < size sh.
Lemma reshape_offsetP sh i :
i < sumn sh -> reshape_offset sh i < nth 0 sh (reshape_index sh i).
Lemma reshape_indexK sh i :
flatten_index sh (reshape_index sh i) (reshape_offset sh i) = i.
Lemma flatten_indexKl sh r c :
c < nth 0 sh r -> reshape_index sh (flatten_index sh r c) = r.
Lemma flatten_indexKr sh r c :
c < nth 0 sh r -> reshape_offset sh (flatten_index sh r c) = c.
Lemma nth_flatten x0 ss i (r := reshape_index (shape ss) i) :
nth x0 (flatten ss) i = nth x0 (nth [::] ss r) (reshape_offset (shape ss) i).
Lemma reshape_leq sh i1 i2
(r1 := reshape_index sh i1) (c1 := reshape_offset sh i1)
(r2 := reshape_index sh i2) (c2 := reshape_offset sh i2) :
(i1 <= i2) = ((r1 < r2) || ((r1 == r2) && (c1 <= c2))).