-- Prefixes and suffixes of a list: take and drop, their lengths, how they
-- recombine, and how they read through concatenation, mapping, and indexing
-- Part of Cumulus, the Grace17 mathematics library.
-- Copyright (c) 2026 Grant Passmore.
-- Released under the Apache License 2.0. See LICENSE in the repository root.
-- Authors: Grant Passmore

module Grace.Cumulus.List.TakeDrop
import Replicate

def listTake {A : Set} (n : Nat) (xs : List(A)) → Set
  | n, List.nil => []
  | n, List.cons(a, tail) =>
    boolChoice(isZero(n), [], a :: listTake(predecessor(n), tail))
  descends xs

def listDrop {A : Set} (n : Nat) (xs : List(A)) → Set
  | n, List.nil => []
  | n, List.cons(a, tail) =>
    boolChoice(isZero(n), a :: tail, listDrop(predecessor(n), tail))
  descends xs

theorem listTake.nil (A : Set) (n : Nat)
  shows listTake{A}(n, []) = []
proof
  simp [listTake]

theorem listTake.zero (A : Set) (xs : List(A))
  shows listTake{A}(0, xs) = []
proof
  induction xs
  | List.nil => simp [listTake]
  | List.cons(a, tail) using ih => simp [listTake, isZero, boolChoice]

theorem listTake.succ_cons (A a : Set) (xs : List(A)) (n : Nat)
  requires element : a ∈ A
  shows listTake{A}(succ(n), a :: xs) = a :: listTake{A}(n, xs)
proof
  simp [listTake, isZero, predecessor, boolChoice]

theorem listDrop.nil (A : Set) (n : Nat)
  shows listDrop{A}(n, []) = []
proof
  simp [listDrop]

theorem listDrop.zero (A : Set) (xs : List(A))
  shows listDrop{A}(0, xs) = xs
proof
  induction xs
  | List.nil => simp [listDrop]
  | List.cons(a, tail) using ih => simp [listDrop, isZero, boolChoice]

theorem listDrop.succ_cons (A a : Set) (xs : List(A)) (n : Nat)
  requires element : a ∈ A
  shows listDrop{A}(succ(n), a :: xs) = listDrop{A}(n, xs)
proof
  simp [listDrop, isZero, predecessor, boolChoice]

theorem listTake.closed (A : Set) (xs : List(A)) (n : Nat)
  shows List(A, listTake{A}(n, xs))
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat
    rw listTake.nil
    typecheck
  | List.cons(a, tail) using ih =>
    intros n nNat
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listTake.zero]
      typecheck
    . obtain ⟨k, kNat, equal⟩ := successor
      have inner := ih with k, kNat
      rw [equal, listTake.succ_cons]
      typecheck

theorem listDrop.closed (A : Set) (xs : List(A)) (n : Nat)
  shows List(A, listDrop{A}(n, xs))
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat
    rw listDrop.nil
    typecheck
  | List.cons(a, tail) using ih =>
    intros n nNat
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listDrop.zero]
      typecheck
    . obtain ⟨k, kNat, equal⟩ := successor
      have inner := ih with k, kNat
      rw [equal, listDrop.succ_cons]
      exact inner

-- The prefix followed by the suffix is the list.
theorem listTake.append_drop (A : Set) (xs : List(A)) (n : Nat)
  shows append(listDrop{A}(n, xs), listTake{A}(n, xs)) = xs
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat
    rw [listTake.nil, listDrop.nil]
    simp [append]
  | List.cons(a, tail) using ih =>
    intros n nNat
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listTake.zero, listDrop.zero]
      simp [append]
    . obtain ⟨k, kNat, equal⟩ := successor
      have inner := ih with k, kNat
      have taken := listTake.closed at A, tail, k with *
      have dropped := listDrop.closed at A, tail, k with *
      rw [equal, listTake.succ_cons, listDrop.succ_cons, append.cons, inner]
      rfl

theorem listTake.length_le_count (A : Set) (xs : List(A)) (n : Nat)
  shows length{A}(listTake{A}(n, xs)) ≤ n
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat
    rw [listTake.nil, length.nil]
    arith
  | List.cons(a, tail) using ih =>
    intros n nNat
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listTake.zero, length.nil]
      arith
    . obtain ⟨k, kNat, equal⟩ := successor
      have inner := ih with k, kNat
      have taken := listTake.closed at A, tail, k with *
      rw [equal, listTake.succ_cons, length.cons]
      arith

theorem listTake.length_le (A : Set) (xs : List(A)) (n : Nat)
  shows length{A}(listTake{A}(n, xs)) ≤ length{A}(xs)
proof
  have split := listTake.append_drop at A, xs, n with *
  have taken := listTake.closed at A, xs, n with *
  have dropped := listDrop.closed at A, xs, n with *
  have bound := append.length_left_le at A, listTake{A}(n, xs), listDrop{A}(n, xs) with *
  rwa split at bound

theorem listDrop.length_le (A : Set) (xs : List(A)) (n : Nat)
  shows length{A}(listDrop{A}(n, xs)) ≤ length{A}(xs)
proof
  have split := listTake.append_drop at A, xs, n with *
  have taken := listTake.closed at A, xs, n with *
  have dropped := listDrop.closed at A, xs, n with *
  have bound := append.length_right_le at A, listTake{A}(n, xs), listDrop{A}(n, xs) with *
  rwa split at bound

theorem listTake.length (A : Set) (xs : List(A)) (n : Nat)
  requires bound : n ≤ length{A}(xs)
  shows length{A}(listTake{A}(n, xs)) = n
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat bound
    rw length.nil at bound
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listTake.nil, length.nil]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      rw equal at bound
      exfalso
      arith
  | List.cons(a, tail) using ih =>
    intros n nNat bound
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    rw length.cons at bound
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listTake.zero, length.nil]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      rw equal at bound
      have smaller : k ≤ length{A}(tail) by arith
      have inner := ih with k, kNat, smaller
      have taken := listTake.closed at A, tail, k with *
      rw [equal, listTake.succ_cons, length.cons, inner]
      rfl

theorem listTake.all (A : Set) (xs : List(A)) (n : Nat)
  requires bound : length{A}(xs) ≤ n
  shows listTake{A}(n, xs) = xs
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat bound
    exact listTake.nil at A, n
  | List.cons(a, tail) using ih =>
    intros n nNat bound
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    rw length.cons at bound
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw zero at bound
      exfalso
      arith
    . obtain ⟨k, kNat, equal⟩ := successor
      rw equal at bound
      have smaller : length{A}(tail) ≤ k by arith
      have inner := ih with k, kNat, smaller
      rw [equal, listTake.succ_cons, inner]
      rfl

theorem listTake.self (A : Set) (xs : List(A))
  shows listTake{A}(length{A}(xs), xs) = xs
proof
  have reflexive := LessEqual.reflexive at length{A}(xs) with *
  exact listTake.all with reflexive

theorem listDrop.all (A : Set) (xs : List(A)) (n : Nat)
  requires bound : length{A}(xs) ≤ n
  shows listDrop{A}(n, xs) = []
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat bound
    exact listDrop.nil at A, n
  | List.cons(a, tail) using ih =>
    intros n nNat bound
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    rw length.cons at bound
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw zero at bound
      exfalso
      arith
    . obtain ⟨k, kNat, equal⟩ := successor
      rw equal at bound
      have smaller : length{A}(tail) ≤ k by arith
      have inner := ih with k, kNat, smaller
      rw [equal, listDrop.succ_cons, inner]
      rfl

theorem listDrop.self (A : Set) (xs : List(A))
  shows listDrop{A}(length{A}(xs), xs) = []
proof
  have reflexive := LessEqual.reflexive at length{A}(xs) with *
  exact listDrop.all with reflexive

theorem listDrop.length (A : Set) (xs : List(A)) (n : Nat)
  requires bound : n ≤ length{A}(xs)
  shows length{A}(listDrop{A}(n, xs)) + n = length{A}(xs)
proof
  have split := listTake.append_drop at A, xs, n with *
  have taken := listTake.closed at A, xs, n with *
  have dropped := listDrop.closed at A, xs, n with *
  have lengths := length.append at A, listDrop{A}(n, xs), listTake{A}(n, xs) with *
  rw [split, listTake.length] at lengths
  symm at lengths
  rwa add.commutative at lengths

theorem listTake.eq_nil (A : Set) (xs : List(A)) (n : Nat)
  requires empty : listTake{A}(n, xs) = []
  shows n = 0 ∨ xs = []
proof
  have shape := natural.zero_or_successor with *
  cases shape with zero | successor
  . left
    exact zero
  . right
    obtain ⟨k, kNat, equal⟩ := successor
    have listShape := List.cases with *
    cases listShape with nil | cons
    . exact nil
    . obtain ⟨a, element, tail, tailList, equal2⟩ := cons
      rw [equal, equal2, listTake.succ_cons] at empty
      have taken := listTake.closed at A, tail, k with *
      exfalso
      have distinct := List.cons_ne_nil at A, a, listTake{A}(k, tail) with element
      exact distinct with empty

theorem listTake.take (A : Set) (xs : List(A)) (m n : Nat)
  requires bound : m ≤ n
  shows listTake{A}(m, listTake{A}(n, xs)) = listTake{A}(m, xs)
proof
  induction xs generalizing m n
  | List.nil =>
    intros m n mNat nNat bound
    rw [listTake.nil, listTake.nil]
    rfl
  | List.cons(a, tail) using ih =>
    intros m n mNat nNat bound
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have outer := listTake.closed at A, a :: tail, n with *
    have shapeM := natural.zero_or_successor at m with mNat
    cases shapeM with zeroM | successorM
    . rw [zeroM, listTake.zero, listTake.zero]
      rfl
    . obtain ⟨j, jNat, equalM⟩ := successorM
      have shapeN := natural.zero_or_successor at n with nNat
      cases shapeN with zeroN | successorN
      . rw [equalM, zeroN] at bound
        exfalso
        arith
      . obtain ⟨k, kNat, equalN⟩ := successorN
        rw [equalM, equalN] at bound
        have smaller : j ≤ k by arith
        have inner := ih with j, k, jNat, kNat, smaller
        have taken := listTake.closed at A, tail, k with *
        rw [equalM, equalN, listTake.succ_cons, listTake.succ_cons, inner]
        rw listTake.succ_cons
        rfl

theorem listTake.idempotent (A : Set) (xs : List(A)) (n : Nat)
  shows listTake{A}(n, listTake{A}(n, xs)) = listTake{A}(n, xs)
proof
  have reflexive := LessEqual.reflexive at n with *
  exact listTake.take with reflexive

theorem listDrop.drop (A : Set) (xs : List(A)) (m n : Nat)
  shows listDrop{A}(m, listDrop{A}(n, xs)) = listDrop{A}(n + m, xs)
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat
    rw [listDrop.nil, listDrop.nil, listDrop.nil]
    rfl
  | List.cons(a, tail) using ih =>
    intros n nNat
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have shape := natural.zero_or_successor at n with nNat
    cases shape with zero | successor
    . rw [zero, listDrop.zero, add.left_unit]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      have inner := ih with k, kNat
      have shift : succ(k) + m = succ(k + m) by simp [add, add.succ_right]
      rw [equal, listDrop.succ_cons, inner, shift, listDrop.succ_cons]
      rfl

theorem listDrop.succ (A : Set) (xs : List(A)) (n : Nat)
  shows listDrop{A}(succ(n), xs) = listTail{A}(listDrop{A}(n, xs))
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat
    rw [listDrop.nil, listDrop.nil, listTail.nil]
    rfl
  | List.cons(a, tail) using ih =>
    intros n nNat
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have shape := natural.zero_or_successor at n with nNat
    cases shape with zero | successor
    . rw [zero, listDrop.zero, listDrop.succ_cons, listDrop.zero, listTail.cons]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      have inner := ih with k, kNat
      rw [equal, listDrop.succ_cons, listDrop.succ_cons, inner]
      rfl

theorem listDrop.one (A : Set) (xs : List(A))
  shows listDrop{A}(1, xs) = listTail{A}(xs)
proof
  have step := listDrop.succ at A, xs, 0 with *
  rw [listDrop.zero] at step
  exact step

theorem listDrop.tail (A : Set) (xs : List(A)) (n : Nat)
  shows listDrop{A}(succ(n), xs) = listDrop{A}(n, listTail{A}(xs))
proof
  have shape := List.cases with *
  cases shape with nil | cons
  . rw [nil, listTail.nil, listDrop.nil, listDrop.nil]
    rfl
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    rw [equal, listTail.cons, listDrop.succ_cons]
    rfl

-- Take and drop across a concatenation.
theorem listTake.append_left (A : Set) (xs ys : List(A)) (n : Nat)
  requires bound : n ≤ length{A}(xs)
  shows listTake{A}(n, append(ys, xs)) = listTake{A}(n, xs)
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat bound
    rw length.nil at bound
    have ysList : List(A, ys) by typecheck
    rw append.nil
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listTake.zero, listTake.nil]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      rw equal at bound
      exfalso
      arith
  | List.cons(a, tail) using ih =>
    intros n nNat bound
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    rw length.cons at bound
    have appended := append.closed at A, tail, ys with *
    have appendedWhole : List(A, append(ys, a :: tail)) by typecheck
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listTake.zero, listTake.zero]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      rw equal at bound
      have smaller : k ≤ length{A}(tail) by arith
      have inner := ih with k, kNat, smaller
      rw [equal, append.cons, listTake.succ_cons, listTake.succ_cons, inner]
      rfl

theorem listDrop.append_left (A : Set) (xs ys : List(A)) (n : Nat)
  requires bound : n ≤ length{A}(xs)
  shows listDrop{A}(n, append(ys, xs)) = append(ys, listDrop{A}(n, xs))
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat bound
    rw length.nil at bound
    have ysList : List(A, ys) by typecheck
    rw append.nil
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listDrop.zero, listDrop.nil, append.nil]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      rw equal at bound
      exfalso
      arith
  | List.cons(a, tail) using ih =>
    intros n nNat bound
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    rw length.cons at bound
    have appended := append.closed at A, tail, ys with *
    have appendedWhole : List(A, append(ys, a :: tail)) by typecheck
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listDrop.zero, listDrop.zero]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      rw equal at bound
      have smaller : k ≤ length{A}(tail) by arith
      have inner := ih with k, kNat, smaller
      rw [equal, append.cons, listDrop.succ_cons, listDrop.succ_cons, inner]
      rfl

theorem listTake.append_right (A : Set) (xs ys : List(A)) (n : Nat)
  shows listTake{A}(length{A}(xs) + n, append(ys, xs)) =
    append(listTake{A}(n, ys), xs)
proof
  induction xs
  | List.nil =>
    have taken := listTake.closed at A, ys, n with *
    rw [length.nil, add.left_unit]
    simp [append]
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have appended := append.closed at A, tail, ys with *
    have taken := listTake.closed at A, ys, n with *
    have shift : succ(length{A}(tail)) + n = succ(length{A}(tail) + n)
      simp [add, add.succ_right]
    rw [length.cons, shift, append.cons, listTake.succ_cons, ih, append.cons]
    rfl

theorem listDrop.append_right (A : Set) (xs ys : List(A)) (n : Nat)
  shows listDrop{A}(length{A}(xs) + n, append(ys, xs)) = listDrop{A}(n, ys)
proof
  induction xs
  | List.nil =>
    rw [length.nil, add.left_unit]
    simp [append]
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have appended := append.closed at A, tail, ys with *
    have shift : succ(length{A}(tail)) + n = succ(length{A}(tail) + n)
      simp [add, add.succ_right]
    rw [length.cons, shift, append.cons, listDrop.succ_cons, ih]
    rfl

theorem listTake.append_self (A : Set) (xs ys : List(A))
  shows listTake{A}(length{A}(xs), append(ys, xs)) = xs
proof
  have reflexive := LessEqual.reflexive at length{A}(xs) with *
  have step := listTake.append_left at A, xs, ys, length{A}(xs) with reflexive
  rw [step, listTake.self]
  rfl

theorem listDrop.append_self (A : Set) (xs ys : List(A))
  shows listDrop{A}(length{A}(xs), append(ys, xs)) = ys
proof
  have step := listDrop.append_right at A, xs, ys, 0 with *
  rw [add.right_unit, listDrop.zero] at step
  exact step

-- Mapping commutes with take and drop.
theorem listTake.map (A B f : Set) (xs : List(A)) (n : Nat)
  requires function : f : A → B
  shows listMap{A}(f, listTake{A}(n, xs)) = listTake{B}(n, listMap{A}(f, xs))
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat
    rw [listTake.nil, listMap.nil, listTake.nil]
    rfl
  | List.cons(a, tail) using ih =>
    intros n nNat
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have image := FunctionBetween.apply_closed with function, element
    have mapped : List(B, listMap{A}(f, tail)) by exact listMap.closed with function
    have mappedWhole : List(B, f(a) :: listMap{A}(f, tail)) by typecheck
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listTake.zero, listMap.nil, listMap.cons, listTake.zero]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      have inner := ih with k, kNat
      have taken := listTake.closed at A, tail, k with *
      rw [equal, listTake.succ_cons, listMap.cons, inner, listMap.cons]
      rw listTake.succ_cons
      rfl

theorem listDrop.map (A B f : Set) (xs : List(A)) (n : Nat)
  requires function : f : A → B
  shows listMap{A}(f, listDrop{A}(n, xs)) = listDrop{B}(n, listMap{A}(f, xs))
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat
    rw [listDrop.nil, listMap.nil, listDrop.nil]
    rfl
  | List.cons(a, tail) using ih =>
    intros n nNat
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have image := FunctionBetween.apply_closed with function, element
    have mapped : List(B, listMap{A}(f, tail)) by exact listMap.closed with function
    have mappedWhole : List(B, f(a) :: listMap{A}(f, tail)) by typecheck
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listDrop.zero, listMap.cons, listDrop.zero]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      have inner := ih with k, kNat
      rw [equal, listDrop.succ_cons, inner, listMap.cons, listDrop.succ_cons]
      rfl

-- Take and drop of a constant list.
theorem listTake.replicate (A x : Set) (m k : Nat)
  requires element : x ∈ A
  shows listTake{A}(m, listReplicate(k + m, x)) = listReplicate(m, x)
proof
  induction m
  | zero =>
    have list := listReplicate.closed at A, x, k + 0 with element
    rw [listTake.zero, listReplicate.zero]
    rfl
  | succ(j) using ih =>
    have list := listReplicate.closed at A, x, k + j with element
    rw [add.succ_right, listReplicate.succ, listTake.succ_cons, ih]
    rw listReplicate.succ
    rfl

theorem listDrop.replicate (A x : Set) (m k : Nat)
  requires element : x ∈ A
  shows listDrop{A}(m, listReplicate(k + m, x)) = listReplicate(k, x)
proof
  induction m
  | zero =>
    have list := listReplicate.closed at A, x, k with element
    rw [add.right_unit, listDrop.zero]
    rfl
  | succ(j) using ih =>
    have list := listReplicate.closed at A, x, k + j with element
    rw [add.succ_right, listReplicate.succ, listDrop.succ_cons, ih]
    rfl

-- Supports of prefixes and suffixes.
theorem listTake.support_subset (A : Set) (xs : List(A)) (n : Nat)
  shows listSupport{A}(listTake{A}(n, xs)) ⊆ listSupport{A}(xs)
proof
  have split := listTake.append_drop at A, xs, n with *
  have taken := listTake.closed at A, xs, n with *
  have dropped := listDrop.closed at A, xs, n with *
  intros y hy
  have member := listSupport.mem_append_left at A, y, listTake{A}(n, xs), listDrop{A}(n, xs) with hy
  rwa split at member

theorem listDrop.support_subset (A : Set) (xs : List(A)) (n : Nat)
  shows listSupport{A}(listDrop{A}(n, xs)) ⊆ listSupport{A}(xs)
proof
  have split := listTake.append_drop at A, xs, n with *
  have taken := listTake.closed at A, xs, n with *
  have dropped := listDrop.closed at A, xs, n with *
  intros y hy
  have member := listSupport.mem_append_right at A, y, listTake{A}(n, xs), listDrop{A}(n, xs) with hy
  rwa split at member

theorem listSupport.take_union_drop (A : Set) (xs : List(A)) (n : Nat)
  shows listSupport{A}(xs) =
    listSupport{A}(listTake{A}(n, xs)) ∪ listSupport{A}(listDrop{A}(n, xs))
proof
  have split := listTake.append_drop at A, xs, n with *
  have taken := listTake.closed at A, xs, n with *
  have dropped := listDrop.closed at A, xs, n with *
  have union := listSupport.append at A, listDrop{A}(n, xs), listTake{A}(n, xs) with *
  rwa split at union

-- Indexing into prefixes and suffixes.
theorem listTake.coefficient (A z : Set) (xs : List(A)) (n k : Nat)
  requires bound : k < n
  shows listCoefficient{A}(z, listTake{A}(n, xs), k) = listCoefficient{A}(z, xs, k)
proof
  induction xs generalizing n k
  | List.nil =>
    intros n k nNat kNat bound
    rw listTake.nil
    rfl
  | List.cons(a, tail) using ih =>
    intros n k nNat kNat bound
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have shapeN := natural.zero_or_successor at n with nNat
    cases shapeN with zeroN | successorN
    . rw zeroN at bound
      exfalso
      arith
    . obtain ⟨j, jNat, equalN⟩ := successorN
      have taken := listTake.closed at A, tail, j with *
      rw [equalN, listTake.succ_cons]
      have shapeK := natural.zero_or_successor at k with kNat
      cases shapeK with zeroK | successorK
      . rw [zeroK, listCoefficient.zero, listCoefficient.zero]
        rfl
      . obtain ⟨r, rNat, equalK⟩ := successorK
        rw [equalN, equalK] at bound
        have smaller : r < j by arith
        have inner := ih with j, r, jNat, rNat, smaller
        rw [equalK, listCoefficient.succ, listCoefficient.succ]
        exact inner

theorem listDrop.coefficient (A z : Set) (xs : List(A)) (n k : Nat)
  shows listCoefficient{A}(z, listDrop{A}(n, xs), k) = listCoefficient{A}(z, xs, n + k)
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat
    rw [listDrop.nil, listCoefficient.nil, listCoefficient.nil]
    rfl
  | List.cons(a, tail) using ih =>
    intros n nNat
    have element : a ∈ A by assumption
    have whole : List(A, a :: tail) by typecheck
    have shape := natural.zero_or_successor at n with nNat
    cases shape with zero | successor
    . rw [zero, listDrop.zero, add.left_unit]
      rfl
    . obtain ⟨j, jNat, equal⟩ := successor
      have inner := ih with j, jNat
      have shift : succ(j) + k = succ(j + k) by simp [add, add.succ_right]
      rw [equal, listDrop.succ_cons, inner, shift, listCoefficient.succ]
      rfl

theorem listCoefficient.head (A z : Set) (xs : List(A))
  shows listCoefficient{A}(z, xs, 0) = listHeadOr{A}(z, xs)
proof
  simp [listCoefficient]

theorem listCoefficient.tail (A z : Set) (xs : List(A)) (n : Nat)
  shows listCoefficient{A}(z, listTail{A}(xs), n) = listCoefficient{A}(z, xs, succ(n))
proof
  simp [listCoefficient]