-- Clean/Utils/Vector.lean

import Mathlib.Analysis.Normed.Ring.Lemmas

import Mathlib.Combinatorics.Enumerative.Composition

import Init.Data.List.Find

variable {α β : Type} {n m : ℕ}

open Vector (finRange)

namespace Vector

def cons (a : α) (v : Vector α n) : Vector α (n + 1) :=
  ⟨ .mk (a :: v.toList), by admit /- proof elided -/
  ⟩

def induct {motive : {n : ℕ} → Vector α n → Sort u}
    (nil : motive #v[])
    (cons: ∀ {n : ℕ} (a : α) (as : Vector α n), motive as → motive (cons a as))
    {n : ℕ} (v : Vector α n) : motive v :=
  match v with
  | ⟨ .mk [], h ⟩ => by admit /- proof elided -/
  | ⟨ .mk (a :: as), h ⟩ => by admit /- proof elided -/

end Vector

def Composition.ofProductLength (m : ℕ+) {α : Type} {l : List α} (hl : l.length = n * m.val) : Composition l.length := {
  blocks := List.replicate n m.val
  blocks_pos hi := (List.mem_replicate.mp hi).right ▸ m.pos
  blocks_sum := hl ▸ List.sum_replicate_nat
}

namespace Vector

def toChunks (m : ℕ+) {α : Type} (v : Vector α (n*m)) : Vector (Vector α m) n :=
  let comp := Composition.ofProductLength m v.length_toList
  let list : List (Vector α m) := v.toList.splitWrtComposition comp
    |>.attachWith (List.length · = m) (comp.ofProductLength_mem_length rfl)
    |>.map fun ⟨ l, hl ⟩ => .mk (.mk l) hl
  .mk (.mk list) (by admit /- proof elided -/
  )

-- Repository definitions used:
content: ...

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: def cons (a : α) (v : Vector α n) : Vector α (n + 1) :=
  ⟨ .mk (a :: v.toList), by admit /- proof elided -/
  ⟩

content: def induct {motive : {n : ℕ} → Vector α n → Sort u}
    (nil : motive #v[])
    (cons: ∀ {n : ℕ} (a : α) (as : Vector α n), motive as → motive (cons a as))
    {n : ℕ} (v : Vector α n) : motive v :=
  match v with
  | ⟨ .mk [], h ⟩ => by admit /- proof elided -/
  | ⟨ .mk (a :: as), h ⟩ => by admit /- proof elided -/

content: def Composition.ofProductLength (m : ℕ+) {α : Type} {l : List α} (hl : l.length = n * m.val) : Composition l.length := {
  blocks := List.replicate n m.val
  blocks_pos hi := (List.mem_replicate.mp hi).right ▸ m.pos
  blocks_sum := hl ▸ List.sum_replicate_nat
}

content: def toChunks (m : ℕ+) {α : Type} (v : Vector α (n*m)) : Vector (Vector α m) n :=
  let comp := Composition.ofProductLength m v.length_toList
  let list : List (Vector α m) := v.toList.splitWrtComposition comp
    |>.attachWith (List.length · = m) (comp.ofProductLength_mem_length rfl)
    |>.map fun ⟨ l, hl ⟩ => .mk (.mk l) hl
  .mk (.mk list) (by admit /- proof elided -/
  )

-- Local lemmas used:
content: theorem toList_cons {a : α} {v : Vector α n} : (cons a v).toList = a :: v.toList

-- Target theorem and proof:
theorem flatten_toChunks {α : Type} (m : ℕ+) (v : Vector (Vector α m) n) :
    v.flatten.toChunks m = v
:= by
  simp only [toChunks]
  rw [←Vector.toArray_inj,←Array.toList_inj]
  simp only
  let v_list_list := v.toList.map (Array.toList ∘ toArray)
  have h_flatten : v.flatten.toList = v_list_list.flatten := by
    rw [Vector.flatten_mk, Vector.toList_mk, Array.toList_flatten, Array.toList_map, List.map_map]
    congr
  have h_length : v.flatten.toList.length = n * ↑m := by rw [length_toList]
  have h_flatten_length : v_list_list.flatten.length = n * ↑m := by rw [←h_flatten, h_length]
  have h' : (v.flatten.toList.splitWrtComposition (Composition.ofProductLength m h_length)) = v_list_list := by
    rw [← v_list_list.splitWrtComposition_flatten (Composition.ofProductLength m h_flatten_length)]
    congr 1
    · rw [h_length, h_flatten_length]
    congr
    · simp [h_flatten]
    simp only [List.map_map, Composition.ofProductLength, v_list_list]
    clear *-
    induction v using Vector.induct
    case nil => rfl
    case cons xs x hi => rw [List.replicate_succ, Vector.toList_cons, List.map_cons, hi,
      Function.comp_apply, Function.comp_apply, Array.length_toList, size_toArray]
  simp_all only [List.length_flatten, List.map_map, List.map_attachWith, v_list_list]
  rw [List.map_attach_eq_pmap, List.pmap_map]
  simp only [Function.comp_apply, Array.toArray_toList, mk_toArray, List.pmap_eq_map,
    List.map_id_fun', id_eq]
  congr