import Std
import Batteries
import Lean.Elab.Tactic.Omega

namespace Hunch.ChunkedSequence
def contents : ((A : Type) → List (List A) → List A) :=
  fun _ chunks => chunks.flatten

def splitAt : ((A : Type) → Nat → List (List A) → List A × List A) :=
  fun A n chunks => ((contents A chunks).take n, (contents A chunks).drop n)

def splice : ((A : Type) → Nat → Nat → List (List A) → List (List A) → List A) :=
  fun A n deleted inserted chunks => (contents A chunks).take n ++ contents A inserted ++ (contents A chunks).drop (n + deleted)

theorem concat_contents : (∀ (A : Type) (xs ys : List (List A)), contents A (xs ++ ys) = contents A xs ++ contents A ys) :=
  by
    intro A xs ys
    exact List.flatten_append

theorem split_roundtrip : (∀ (A : Type) (n : Nat) (chunks : List (List A)), (splitAt A n chunks).1 ++ (splitAt A n chunks).2 = contents A chunks) :=
  by
    intro A n chunks
    exact List.take_append_drop n (contents A chunks)

theorem measure_correct : (∀ (A : Type) (chunks : List (List A)), (contents A chunks).length = (chunks.map List.length).sum) :=
  by
    intro A chunks
    induction chunks with
    | nil => rfl
    | cons chunk rest ih => simp [contents] at ih ⊢

theorem splice_length : (∀ (A : Type) (n deleted : Nat) (inserted chunks : List (List A)), n ≤ (contents A chunks).length → deleted ≤ (contents A chunks).length - n → (splice A n deleted inserted chunks).length = (contents A chunks).length - deleted + (contents A inserted).length) :=
  by
    intro A n deleted inserted chunks hn hd
    simp only [splice, List.length_append, List.length_take, List.length_drop]
    omega

end Hunch.ChunkedSequence

def OA_statement : Prop := (∀ (A : Type) (n : Nat) (xs ys : List (List A)), Hunch.ChunkedSequence.contents A (xs ++ ys) = Hunch.ChunkedSequence.contents A xs ++ Hunch.ChunkedSequence.contents A ys ∧ (Hunch.ChunkedSequence.splitAt A n xs).1 ++ (Hunch.ChunkedSequence.splitAt A n xs).2 = Hunch.ChunkedSequence.contents A xs ∧ (Hunch.ChunkedSequence.contents A xs).length = (xs.map List.length).sum)

theorem OA_target : OA_statement :=
  by
    intro A n xs ys
    exact ⟨Hunch.ChunkedSequence.concat_contents A xs ys, Hunch.ChunkedSequence.split_roundtrip A n xs, Hunch.ChunkedSequence.measure_correct A xs⟩

#print axioms OA_target
#check OA_target