-- Prefixes and suffixes as relations between lists: order laws, the
-- characterizations by take and drop, and what they preserve
-- 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.Prefix
import Index

-- xs is an initial segment of ys.
predicate ListPrefix (A xs ys : Set) := ∃ zs, List(A, zs) ∧ ys = append(zs, xs)

-- xs is a final segment of ys.
predicate ListSuffix (A xs ys : Set) := ∃ zs, List(A, zs) ∧ ys = append(xs, zs)

theorem ListPrefix.reflexive (A : Set) (xs : List(A))
  shows ListPrefix(A, xs, xs)
proof
  have nilList : List(A, []) by typecheck
  refine ⟨[], nilList, ?_⟩
  rw append.right_unit
  rfl

theorem ListPrefix.nil (A : Set) (ys : List(A))
  shows ListPrefix(A, [], ys)
proof
  refine ⟨ys, assumption, ?_⟩
  simp [append]

theorem ListPrefix.append (A : Set) (xs zs : List(A))
  shows ListPrefix(A, xs, append(zs, xs))
proof
  refine ⟨zs, assumption, rfl⟩

theorem ListPrefix.transitive (A : Set) (xs ys zs : List(A))
  requires first : ListPrefix(A, xs, ys)
  requires second : ListPrefix(A, ys, zs)
  shows ListPrefix(A, xs, zs)
proof
  obtain ⟨us, usList, equalYs⟩ := first
  obtain ⟨vs, vsList, equalZs⟩ := second
  have appended : List(A, append(vs, us)) by typecheck
  refine ⟨append(vs, us), appended, ?_⟩
  rw [equalZs, equalYs, append.associative]
  rfl

theorem ListPrefix.length (A : Set) (xs ys : List(A))
  requires prefix : ListPrefix(A, xs, ys)
  shows length{A}(xs) ≤ length{A}(ys)
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  rw equal
  exact append.length_left_le at A, xs, zs with *

theorem ListPrefix.eq_take (A : Set) (xs ys : List(A))
  requires prefix : ListPrefix(A, xs, ys)
  shows xs = listTake{A}(length{A}(xs), ys)
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  rw [equal, listTake.append_self]
  rfl

theorem ListPrefix.take (A : Set) (ys : List(A)) (n : Nat)
  shows ListPrefix(A, listTake{A}(n, ys), ys)
proof
  have dropped := listDrop.closed at A, ys, n with *
  refine ⟨listDrop{A}(n, ys), dropped, ?_⟩
  symm
  exact listTake.append_drop at A, ys, n with *

theorem ListPrefix.iff_take (A : Set) (xs ys : List(A))
  shows ListPrefix(A, xs, ys) ↔ xs = listTake{A}(length{A}(xs), ys)
proof
  constructor
  . intro prefix
    exact ListPrefix.eq_take with prefix
  . intro equal
    have taken := ListPrefix.take at A, ys, length{A}(xs) with *
    rwa ← equal at taken

theorem ListPrefix.antisymmetric (A : Set) (xs ys : List(A))
  requires first : ListPrefix(A, xs, ys)
  requires second : ListPrefix(A, ys, xs)
  shows xs = ys
proof
  have forward := ListPrefix.length with first
  have backward := ListPrefix.length with second
  have same := LessEqual.antisymmetric with forward, backward
  have taken := ListPrefix.eq_take with first
  rw [taken, same, listTake.self]
  rfl

theorem ListPrefix.of_eq_length (A : Set) (xs ys : List(A))
  requires prefix : ListPrefix(A, xs, ys)
  requires same : length{A}(xs) = length{A}(ys)
  shows xs = ys
proof
  have taken := ListPrefix.eq_take with prefix
  rw [taken, same, listTake.self]
  rfl

theorem ListPrefix.cons (A a : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires prefix : ListPrefix(A, xs, ys)
  shows ListPrefix(A, a :: xs, a :: ys)
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  refine ⟨zs, zsList, ?_⟩
  rw [equal, append.cons]
  rfl

theorem ListPrefix.cons_cancel (A a : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires prefix : ListPrefix(A, a :: xs, a :: ys)
  shows ListPrefix(A, xs, ys)
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  refine ⟨zs, zsList, ?_⟩
  rw append.cons at equal
  have appended : List(A, append(zs, xs)) by typecheck
  exact List.cons_tail with equal

theorem ListPrefix.coefficient (A z : Set) (xs ys : List(A)) (n : Nat)
  requires prefix : ListPrefix(A, xs, ys)
  requires bound : n < length{A}(xs)
  shows listCoefficient{A}(z, xs, n) = listCoefficient{A}(z, ys, n)
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  rw equal
  symm
  exact listCoefficient.append_left with bound

theorem ListPrefix.support_subset (A : Set) (xs ys : List(A))
  requires prefix : ListPrefix(A, xs, ys)
  shows listSupport{A}(xs) ⊆ listSupport{A}(ys)
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  rw equal
  intros y hy
  exact listSupport.mem_append_left with hy

theorem ListPrefix.count_le (A x : Set) (xs ys : List(A))
  requires prefix : ListPrefix(A, xs, ys)
  shows listCount{A}(x, xs) ≤ listCount{A}(x, ys)
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  have xsNat := listCount.natural at A, x, xs with *
  have zsNat := listCount.natural at A, x, zs with *
  rw [equal, listCount.append]
  arith

theorem ListPrefix.of_cons (A a : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires prefix : ListPrefix(A, a :: xs, ys)
  shows ∃ tail, List(A, tail) ∧ ys = a :: tail ∧ ListPrefix(A, xs, tail)
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  have appended : List(A, append(zs, xs)) by typecheck
  refine ⟨append(zs, xs), appended, ?_, ?_⟩
  . rw [equal, append.cons]
    rfl
  . exact ⟨zs, zsList, rfl⟩

theorem ListPrefix.map (A B f : Set) (xs ys : List(A))
  requires function : f : A → B
  requires prefix : ListPrefix(A, xs, ys)
  shows ListPrefix(B, listMap{A}(f, xs), listMap{A}(f, ys))
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  have mapped : List(B, listMap{A}(f, zs)) by exact listMap.closed with function
  refine ⟨listMap{A}(f, zs), mapped, ?_⟩
  rw equal
  exact listMap.append with function

-- Suffixes.
theorem ListSuffix.reflexive (A : Set) (xs : List(A))
  shows ListSuffix(A, xs, xs)
proof
  have nilList : List(A, []) by typecheck
  refine ⟨[], nilList, ?_⟩
  rw append.nil
  rfl

theorem ListSuffix.nil (A : Set) (ys : List(A))
  shows ListSuffix(A, [], ys)
proof
  refine ⟨ys, assumption, ?_⟩
  rw append.right_unit
  rfl

theorem ListSuffix.append (A : Set) (xs zs : List(A))
  shows ListSuffix(A, xs, append(xs, zs))
proof
  refine ⟨zs, assumption, rfl⟩

theorem ListSuffix.cons (A a : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires suffix : ListSuffix(A, xs, ys)
  shows ListSuffix(A, xs, a :: ys)
proof
  obtain ⟨zs, zsList, equal⟩ := suffix
  have consList : List(A, a :: zs) by typecheck
  refine ⟨a :: zs, consList, ?_⟩
  rw [equal, append.cons]
  rfl

theorem ListSuffix.tail (A a : Set) (xs : List(A))
  requires element : a ∈ A
  shows ListSuffix(A, xs, a :: xs)
proof
  have reflexive := ListSuffix.reflexive at A, xs with *
  exact ListSuffix.cons with element, reflexive

theorem ListSuffix.transitive (A : Set) (xs ys zs : List(A))
  requires first : ListSuffix(A, xs, ys)
  requires second : ListSuffix(A, ys, zs)
  shows ListSuffix(A, xs, zs)
proof
  obtain ⟨us, usList, equalYs⟩ := first
  obtain ⟨vs, vsList, equalZs⟩ := second
  have appended : List(A, append(us, vs)) by typecheck
  refine ⟨append(us, vs), appended, ?_⟩
  rw [equalZs, equalYs, ← append.associative]
  rfl

theorem ListSuffix.length (A : Set) (xs ys : List(A))
  requires suffix : ListSuffix(A, xs, ys)
  shows length{A}(xs) ≤ length{A}(ys)
proof
  obtain ⟨zs, zsList, equal⟩ := suffix
  rw equal
  exact append.length_right_le at A, zs, xs with *

theorem ListSuffix.drop (A : Set) (ys : List(A)) (n : Nat)
  shows ListSuffix(A, listDrop{A}(n, ys), ys)
proof
  have taken := listTake.closed at A, ys, n with *
  refine ⟨listTake{A}(n, ys), taken, ?_⟩
  symm
  exact listTake.append_drop at A, ys, n with *

theorem ListSuffix.eq_drop (A : Set) (xs ys : List(A))
  requires suffix : ListSuffix(A, xs, ys)
  shows ∃ n, Nat(n) ∧ xs = listDrop{A}(n, ys)
proof
  obtain ⟨zs, zsList, equal⟩ := suffix
  have zsNat : Nat(length{A}(zs)) by typecheck
  refine ⟨length{A}(zs), zsNat, ?_⟩
  rw [equal, listDrop.append_self]
  rfl

theorem ListSuffix.antisymmetric (A : Set) (xs ys : List(A))
  requires first : ListSuffix(A, xs, ys)
  requires second : ListSuffix(A, ys, xs)
  shows xs = ys
proof
  obtain ⟨zs, zsList, equal⟩ := first
  have forward := ListSuffix.length with first
  have backward := ListSuffix.length with second
  have same := LessEqual.antisymmetric with forward, backward
  have lengths : length{A}(append(xs, zs)) = length{A}(xs)
    rw ← equal
    symm
    exact same
  rw length.append at lengths
  have padded : length{A}(zs) + length{A}(xs) = 0 + length{A}(xs)
    rw add.left_unit
    exact lengths
  have zero := add.right_cancel with padded
  have empty := length.eq_zero with zero
  rw [equal, empty, append.nil]
  rfl

theorem ListSuffix.of_eq_length (A : Set) (xs ys : List(A))
  requires suffix : ListSuffix(A, xs, ys)
  requires same : length{A}(xs) = length{A}(ys)
  shows xs = ys
proof
  obtain ⟨zs, zsList, equal⟩ := suffix
  have lengths : length{A}(append(xs, zs)) = length{A}(xs)
    rw ← equal
    symm
    exact same
  rw length.append at lengths
  have padded : length{A}(zs) + length{A}(xs) = 0 + length{A}(xs)
    rw add.left_unit
    exact lengths
  have zero := add.right_cancel with padded
  have empty := length.eq_zero with zero
  rw [equal, empty, append.nil]
  rfl

theorem ListSuffix.support_subset (A : Set) (xs ys : List(A))
  requires suffix : ListSuffix(A, xs, ys)
  shows listSupport{A}(xs) ⊆ listSupport{A}(ys)
proof
  obtain ⟨zs, zsList, equal⟩ := suffix
  rw equal
  intros y hy
  exact listSupport.mem_append_right with hy

theorem ListSuffix.count_le (A x : Set) (xs ys : List(A))
  requires suffix : ListSuffix(A, xs, ys)
  shows listCount{A}(x, xs) ≤ listCount{A}(x, ys)
proof
  obtain ⟨zs, zsList, equal⟩ := suffix
  have xsNat := listCount.natural at A, x, xs with *
  have zsNat := listCount.natural at A, x, zs with *
  rw [equal, listCount.append]
  arith

-- Reversal exchanges prefixes and suffixes.
theorem ListPrefix.reverse (A : Set) (xs ys : List(A))
  requires prefix : ListPrefix(A, xs, ys)
  shows ListSuffix(A, reverse(xs), reverse(ys))
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  have reversed : List(A, reverse(zs)) by typecheck
  refine ⟨reverse(zs), reversed, ?_⟩
  rw [equal, reverse.append]
  rfl

theorem ListSuffix.reverse (A : Set) (xs ys : List(A))
  requires suffix : ListSuffix(A, xs, ys)
  shows ListPrefix(A, reverse(xs), reverse(ys))
proof
  obtain ⟨zs, zsList, equal⟩ := suffix
  have reversed : List(A, reverse(zs)) by typecheck
  refine ⟨reverse(zs), reversed, ?_⟩
  rw [equal, reverse.append]
  rfl

theorem ListPrefix.iff_reverse (A : Set) (xs ys : List(A))
  shows ListPrefix(A, xs, ys) ↔ ListSuffix(A, reverse(xs), reverse(ys))
proof
  constructor
  . intro prefix
    exact ListPrefix.reverse with prefix
  . intro suffix
    have reversedXs : List(A, reverse(xs)) by typecheck
    have reversedYs : List(A, reverse(ys)) by typecheck
    have back := ListSuffix.reverse at A, reverse(xs), reverse(ys) with suffix
    rw [reverse.involutive, reverse.involutive] at back
    exact back

theorem ListSuffix.map (A B f : Set) (xs ys : List(A))
  requires function : f : A → B
  requires suffix : ListSuffix(A, xs, ys)
  shows ListSuffix(B, listMap{A}(f, xs), listMap{A}(f, ys))
proof
  obtain ⟨zs, zsList, equal⟩ := suffix
  have mapped : List(B, listMap{A}(f, zs)) by exact listMap.closed with function
  refine ⟨listMap{A}(f, zs), mapped, ?_⟩
  rw equal
  exact listMap.append with function

theorem ListPrefix.filter (A S : Set) (xs ys : List(A))
  requires prefix : ListPrefix(A, xs, ys)
  shows ListPrefix(A, listFilter{A}(S, xs), listFilter{A}(S, ys))
proof
  obtain ⟨zs, zsList, equal⟩ := prefix
  have filtered := listFilter.closed at A, S, zs with *
  refine ⟨listFilter{A}(S, zs), filtered, ?_⟩
  rw equal
  exact listFilter.append at A, S, xs, zs with *

theorem ListSuffix.filter (A S : Set) (xs ys : List(A))
  requires suffix : ListSuffix(A, xs, ys)
  shows ListSuffix(A, listFilter{A}(S, xs), listFilter{A}(S, ys))
proof
  obtain ⟨zs, zsList, equal⟩ := suffix
  have filtered := listFilter.closed at A, S, zs with *
  refine ⟨listFilter{A}(S, zs), filtered, ?_⟩
  rw equal
  exact listFilter.append at A, S, zs, xs with *