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 deleted : Nat) (inserted chunks : List (List A)), n ≤ (Hunch.ChunkedSequence.contents A chunks).length → deleted ≤ (Hunch.ChunkedSequence.contents A chunks).length - n → (Hunch.ChunkedSequence.splice A n deleted inserted chunks).length = (Hunch.ChunkedSequence.contents A chunks).length - deleted + (Hunch.ChunkedSequence.contents A inserted).length)

theorem OA_target : OA_statement :=
  Hunch.ChunkedSequence.splice_length

#print axioms OA_target
#check OA_target