-- Heads, tails, constructor laws, the equation lemmas of mapping,
-- reversal, and support that the remaining list modules rewrite with, and
-- the image of a set of lists under a fixed cons
-- 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.Basic
import Append
import Reverse

-- A Boolean choice between two sets; list recursions branch through it.
def boolChoice (test : Bool, yes no : Set) → Set
  | Bool.false, yes, no => no
  | Bool.true, yes, no => yes

theorem boolChoice.false (yes no : Set)
  shows boolChoice(Bool.false, yes, no) = no
proof
  simp [boolChoice]

theorem boolChoice.true (yes no : Set)
  shows boolChoice(Bool.true, yes, no) = yes
proof
  simp [boolChoice]

-- Constructor laws.
theorem List.cons_ne_nil (A a : Set) (xs : List(A))
  requires element : a ∈ A
  shows a :: xs ≠ []
proof
  intro equal
  symm at equal
  have distinct := List.inductive.disjoint.nil.cons at a, xs
  exact distinct with equal

theorem List.nil_ne_cons (A a : Set) (xs : List(A))
  requires element : a ∈ A
  shows [] ≠ a :: xs
proof
  exact List.inductive.disjoint.nil.cons at a, xs

theorem List.cons_injective (A a b : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires other : b ∈ A
  requires equal : a :: xs = b :: ys
  shows a = b ∧ xs = ys
proof
  exact List.inductive.injective.cons with equal

theorem List.cons_head (A a b : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires other : b ∈ A
  requires equal : a :: xs = b :: ys
  shows a = b
proof
  have both := List.cons_injective with equal
  exact both.left

theorem List.cons_tail (A a b : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires other : b ∈ A
  requires equal : a :: xs = b :: ys
  shows xs = ys
proof
  have both := List.cons_injective with equal
  exact both.right

theorem List.closed_cons (A a : Set) (xs : List(A))
  requires element : a ∈ A
  shows List(A, a :: xs)
proof
  typecheck

-- Length.
theorem length.eq_zero (A : Set) (xs : List(A))
  requires zero : length{A}(xs) = 0
  shows xs = []
proof
  have shape := List.cases with *
  cases shape with empty | nonempty
  . exact empty
  . obtain ⟨a, element, tail, tailList, equal⟩ := nonempty
    rw [equal, length.cons] at zero
    have distinct := Naturals.disjoint.zero.succ at length{A}(tail)
    exfalso
    symm at zero
    exact distinct with zero

theorem length.nil_iff (A : Set) (xs : List(A))
  shows length{A}(xs) = 0 ↔ xs = []
proof
  constructor
  . intro zero
    exact length.eq_zero with zero
  . intro empty
    rw empty
    exact length.nil at A

theorem length.cons_positive (A a : Set) (xs : List(A))
  requires element : a ∈ A
  shows 0 < length{A}(a :: xs)
proof
  rw length.cons
  arith

theorem length.positive_of_ne_nil (A : Set) (xs : List(A))
  requires nonempty : xs ≠ []
  shows 0 < length{A}(xs)
proof
  have shape := List.cases with *
  cases shape with empty | cons
  . exfalso
    exact nonempty with empty
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    rw equal
    exact length.cons_positive with element

theorem List.ne_nil_of_length_positive (A : Set) (xs : List(A))
  requires positive : 0 < length{A}(xs)
  shows xs ≠ []
proof
  intro empty
  rw [empty, length.nil] at positive
  exact Less.irreflexive with Naturals.constructor.zero, positive

theorem length.positive_iff (A : Set) (xs : List(A))
  shows 0 < length{A}(xs) ↔ xs ≠ []
proof
  constructor
  . intro positive
    exact List.ne_nil_of_length_positive with positive
  . intro nonempty
    exact length.positive_of_ne_nil with nonempty

theorem length.singleton (A a : Set)
  requires element : a ∈ A
  shows length{A}(a :: []) = 1
proof
  simp [length]

-- Heads and tails.
theorem listTail.nil (A : Set)
  shows listTail{A}([]) = []
proof
  simp [listTail]

theorem listTail.cons (A a : Set) (xs : List(A))
  requires element : a ∈ A
  shows listTail{A}(a :: xs) = xs
proof
  simp [listTail]

theorem listHeadOr.nil (A z : Set)
  shows listHeadOr{A}(z, []) = z
proof
  simp [listHeadOr]

theorem listHeadOr.cons (A z a : Set) (xs : List(A))
  requires element : a ∈ A
  shows listHeadOr{A}(z, a :: xs) = a
proof
  simp [listHeadOr]

theorem listHeadOr.mem_carrier (A z : Set) (xs : List(A))
  requires default : z ∈ A
  shows listHeadOr{A}(z, xs) ∈ A
proof
  have shape := List.cases with *
  cases shape with empty | cons
  . rw [empty, listHeadOr.nil]
    exact default
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    rw [equal, listHeadOr.cons]
    exact element

theorem listHeadOr.mem_support (A z : Set) (xs : List(A))
  requires nonempty : xs ≠ []
  shows listHeadOr{A}(z, xs) ∈ listSupport{A}(xs)
proof
  have shape := List.cases with *
  cases shape with empty | cons
  . exfalso
    exact nonempty with empty
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    rw [equal, listHeadOr.cons, listSupport.cons]
    exact insert.self

theorem List.head_tail (A z : Set) (xs : List(A))
  requires nonempty : xs ≠ []
  shows xs = listHeadOr{A}(z, xs) :: listTail{A}(xs)
proof
  have shape := List.cases with *
  cases shape with empty | cons
  . exfalso
    exact nonempty with empty
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    rw [equal, listHeadOr.cons, listTail.cons]
    rfl

theorem listTail.length (A : Set) (xs : List(A))
  requires nonempty : xs ≠ []
  shows succ(length{A}(listTail{A}(xs))) = length{A}(xs)
proof
  have shape := List.cases with *
  cases shape with empty | cons
  . exfalso
    exact nonempty with empty
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    rw [equal, listTail.cons, length.cons]
    rfl

theorem listTail.support_subset (A : Set) (xs : List(A))
  shows listSupport{A}(listTail{A}(xs)) ⊆ listSupport{A}(xs)
proof
  have shape := List.cases with *
  cases shape with empty | cons
  . rw [empty, listTail.nil]
    exact Subset.reflexive at listSupport{A}([])
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    rw [equal, listTail.cons, listSupport.cons]
    intros y hy
    exact insert.contains with hy

-- Append.
theorem append.closed (A : Set) (xs ys : List(A))
  shows List(A, append(ys, xs))
proof
  typecheck

theorem append.nil (A : Set) (ys : List(A))
  shows append(ys, []) = ys
proof
  simp [append]

theorem append.cons (A a : Set) (xs ys : List(A))
  requires element : a ∈ A
  shows append(ys, a :: xs) = a :: append(ys, xs)
proof
  simp [append]

theorem append.singleton (A a : Set) (ys : List(A))
  requires element : a ∈ A
  shows append(ys, a :: []) = a :: ys
proof
  simp [append]

theorem append.eq_nil (A : Set) (xs ys : List(A))
  requires empty : append(ys, xs) = []
  shows xs = [] ∧ ys = []
proof
  have lengths : length{A}(append(ys, xs)) = 0
    rw empty
    exact length.nil at A
  rw length.append at lengths
  have both := add.eq_zero with lengths
  obtain ⟨xsZero, ysZero⟩ := both
  constructor
  . exact length.eq_zero with xsZero
  . exact length.eq_zero with ysZero

theorem append.nil_iff (A : Set) (xs ys : List(A))
  shows append(ys, xs) = [] ↔ xs = [] ∧ ys = []
proof
  constructor
  . intro empty
    exact append.eq_nil with empty
  . intro both
    obtain ⟨xsNil, ysNil⟩ := both
    rw [xsNil, ysNil]
    simp [append]

theorem append.length_left_le (A : Set) (xs ys : List(A))
  shows length{A}(xs) ≤ length{A}(append(ys, xs))
proof
  rw length.append
  arith

theorem append.length_right_le (A : Set) (xs ys : List(A))
  shows length{A}(ys) ≤ length{A}(append(ys, xs))
proof
  rw length.append
  arith

theorem append.left_nil_of_eq (A : Set) (xs ys : List(A))
  requires equal : append(ys, xs) = xs
  shows ys = []
proof
  have lengths : length{A}(append(ys, xs)) = length{A}(xs)
    rw equal
    rfl
  rw length.append at lengths
  have padded : length{A}(xs) + length{A}(ys) = length{A}(xs) + 0
    rw add.right_unit
    exact lengths
  have zero := add.left_cancel with padded
  exact length.eq_zero with zero

theorem append.right_nil_of_eq (A : Set) (xs ys : List(A))
  requires equal : append(ys, xs) = ys
  shows xs = []
proof
  have lengths : length{A}(append(ys, xs)) = length{A}(ys)
    rw equal
    rfl
  rw length.append at lengths
  have padded : length{A}(xs) + length{A}(ys) = 0 + length{A}(ys)
    rw add.left_unit
    exact lengths
  have zero := add.right_cancel with padded
  exact length.eq_zero with zero

-- Support.
theorem listSupport.mem_cons_self (A a : Set) (xs : List(A))
  requires element : a ∈ A
  shows a ∈ listSupport{A}(a :: xs)
proof
  rw listSupport.cons
  exact insert.self

theorem listSupport.mem_cons_of_mem (A a y : Set) (xs : List(A))
  requires element : a ∈ A
  requires member : y ∈ listSupport{A}(xs)
  shows y ∈ listSupport{A}(a :: xs)
proof
  rw listSupport.cons
  exact insert.contains with member

theorem listSupport.mem_cons_iff (A a y : Set) (xs : List(A))
  requires element : a ∈ A
  shows y ∈ listSupport{A}(a :: xs) ↔ y = a ∨ y ∈ listSupport{A}(xs)
proof
  rw listSupport.cons
  exact insert.membership at y, a, listSupport{A}(xs)

theorem listSupport.mem_append_iff (A y : Set) (xs ys : List(A))
  shows y ∈ listSupport{A}(append(ys, xs)) ↔
    y ∈ listSupport{A}(xs) ∨ y ∈ listSupport{A}(ys)
proof
  rw listSupport.append
  exact union2.membership at y, listSupport{A}(xs), listSupport{A}(ys)

theorem listSupport.mem_append_left (A y : Set) (xs ys : List(A))
  requires member : y ∈ listSupport{A}(xs)
  shows y ∈ listSupport{A}(append(ys, xs))
proof
  rw listSupport.append
  simp

theorem listSupport.mem_append_right (A y : Set) (xs ys : List(A))
  requires member : y ∈ listSupport{A}(ys)
  shows y ∈ listSupport{A}(append(ys, xs))
proof
  rw listSupport.append
  simp

theorem listSupport.singleton (A a : Set)
  requires element : a ∈ A
  shows listSupport{A}(a :: []) = {a}
proof
  simp [listSupport]
  ext y
  simp

theorem listSupport.eq_empty (A : Set) (xs : List(A))
  requires empty : listSupport{A}(xs) = ∅
  shows xs = []
proof
  have shape := List.cases with *
  cases shape with nil | cons
  . exact nil
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    have member : a ∈ listSupport{A}(xs)
      rw equal
      exact listSupport.mem_cons_self with element
    rw empty at member
    exfalso
    exact set.empty with member

theorem listSupport.empty_iff (A : Set) (xs : List(A))
  shows listSupport{A}(xs) = ∅ ↔ xs = []
proof
  constructor
  . intro empty
    exact listSupport.eq_empty with empty
  . intro nil
    rw nil
    exact listSupport.nil at A

theorem listSupport.nonempty_of_ne_nil (A : Set) (xs : List(A))
  requires nonempty : xs ≠ []
  shows ∃ y, y ∈ listSupport{A}(xs)
proof
  have shape := List.cases with *
  cases shape with nil | cons
  . exfalso
    exact nonempty with nil
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    witness a
    rw equal
    exact listSupport.mem_cons_self with element

-- Reversal.
theorem reverse.nil (A : Set)
  shows reverse{A}([]) = []
proof
  simp [reverse]

theorem reverse.cons (A a : Set) (xs : List(A))
  requires element : a ∈ A
  shows reverse(a :: xs) = append(a :: [], reverse(xs))
proof
  simp [reverse]

theorem reverse.closed (A : Set) (xs : List(A))
  shows List(A, reverse(xs))
proof
  typecheck

theorem reverse.snoc (A a : Set) (xs : List(A))
  requires element : a ∈ A
  shows reverse(append(a :: [], xs)) = a :: reverse(xs)
proof
  have single : List(A, a :: []) by typecheck
  have reversed : List(A, reverse(xs)) by typecheck
  rw [reverse.append, reverse.singleton, append.singleton]
  rfl

theorem reverse.injective (A : Set) (xs ys : List(A))
  requires equal : reverse(xs) = reverse(ys)
  shows xs = ys
proof
  have doubled : reverse(reverse(xs)) = reverse(reverse(ys))
    rw equal
    rfl
  rw [reverse.involutive, reverse.involutive] at doubled
  exact doubled

theorem reverse.eq_nil (A : Set) (xs : List(A))
  requires empty : reverse(xs) = []
  shows xs = []
proof
  have nilReversed : reverse{A}([]) = [] by simp [reverse]
  rw ← nilReversed at empty
  exact reverse.injective with empty

theorem reverse.nil_iff (A : Set) (xs : List(A))
  shows reverse(xs) = [] ↔ xs = []
proof
  constructor
  . intro empty
    exact reverse.eq_nil with empty
  . intro nil
    rw nil
    simp [reverse]

-- Every nonempty list ends with an element: the mirror image of List.cases.
theorem List.snoc_cases (A : Set) (xs : List(A))
  shows xs = [] ∨ ∃ a ∈ A, ∃ front, List(A, front) ∧ xs = append(a :: [], front)
proof
  have reversed : List(A, reverse(xs)) by typecheck
  have shape := List.cases with reversed
  cases shape with nil | cons
  . left
    exact reverse.eq_nil with nil
  . right
    obtain ⟨a, element, tail, tailList, equal⟩ := cons
    refine ⟨a, element, reverse(tail), ?_, ?_⟩
    . typecheck
    . have back : reverse(reverse(xs)) = reverse(a :: tail)
        rw equal
        rfl
      have single : List(A, a :: []) by typecheck
      rw [reverse.involutive, reverse.cons] at back
      rw back
      rfl

-- Mapping.
theorem listMap.nil (A f : Set)
  shows listMap{A}(f, []) = []
proof
  simp [listMap]

theorem listMap.cons (A f a : Set) (xs : List(A))
  requires element : a ∈ A
  shows listMap{A}(f, a :: xs) = f(a) :: listMap{A}(f, xs)
proof
  simp [listMap]

theorem listMap.singleton (A f a : Set)
  requires element : a ∈ A
  shows listMap{A}(f, a :: []) = f(a) :: []
proof
  simp [listMap]

theorem listMap.reverse (A B f : Set) (xs : List(A))
  requires function : f : A → B
  shows listMap{A}(f, reverse(xs)) = reverse{B}(listMap{A}(f, xs))
proof
  induction xs
  | List.nil =>
    rw [reverse.nil, listMap.nil, reverse.nil]
    rfl
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have image := FunctionBetween.apply_closed with function, element
    have single : List(A, a :: []) by typecheck
    have reversed : List(A, reverse(tail)) by typecheck
    have mapped : List(B, listMap{A}(f, tail)) by exact listMap.closed with function
    have step := listMap.append at A, B, f, reverse(tail), a :: [] with function
    rw [reverse.cons, step, ih, listMap.singleton, listMap.cons, reverse.cons]
    rfl

theorem listMap.identity (A : Set) (xs : List(A))
  shows listMap{A}(identity(A), xs) = xs
proof
  induction xs
  | List.nil => simp [listMap]
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have value := identity.apply with element
    rw [listMap.cons, ih, value]
    rfl

theorem listMap.congr (A f g : Set) (xs : List(A))
  requires agreement : ∀ x ∈ listSupport{A}(xs), f(x) = g(x)
  shows listMap{A}(f, xs) = listMap{A}(g, xs)
proof
  induction xs
  | List.nil =>
    intro agreement
    simp [listMap]
  | List.cons(a, tail) using ih =>
    intro agreement
    have element : a ∈ A by assumption
    have headIn := listSupport.mem_cons_self with element
    have head := agreement with a, headIn
    have rest : ∀ x ∈ listSupport{A}(tail), f(x) = g(x)
      intros x hx
      have member := listSupport.mem_cons_of_mem with element, hx
      exact agreement with x, member
    have tails := ih with rest
    rw [listMap.cons, listMap.cons, head, tails]
    rfl

theorem listMap.eq_nil (A f : Set) (xs : List(A))
  requires empty : listMap{A}(f, xs) = []
  shows xs = []
proof
  have shape := List.cases with *
  cases shape with nil | cons
  . exact nil
  . obtain ⟨a, element, tail, tailList, equal⟩ := cons
    rw [equal, listMap.cons] at empty
    have distinct := List.inductive.disjoint.nil.cons at f(a), listMap{A}(f, tail)
    exfalso
    symm at empty
    exact distinct with empty

theorem listMap.mem_support (A B f x : Set) (xs : List(A))
  requires function : f : A → B
  requires member : x ∈ listSupport{A}(xs)
  shows f(x) ∈ listSupport{B}(listMap{A}(f, xs))
proof
  induction xs
  | List.nil =>
    intro member
    rw listSupport.nil at member
    exfalso
    exact set.empty with member
  | List.cons(a, tail) using ih =>
    intro member
    have element : a ∈ A by assumption
    have image := FunctionBetween.apply_closed with function, element
    have mapped : List(B, listMap{A}(f, tail)) by exact listMap.closed with function
    have alternatives := listSupport.mem_cons_iff.forward with element, member
    rw [listMap.cons, listSupport.cons]
    cases alternatives with head | rest
    . rw head
      exact insert.self
    . have inner := ih with rest
      exact insert.contains with inner

theorem listMap.support_subset (A B f : Set) (xs : List(A))
  requires function : f : A → B
  shows listSupport{B}(listMap{A}(f, xs)) ⊆ functionImage(f, listSupport{A}(xs), B)
proof
  induction xs
  | List.nil =>
    rw [listMap.nil, listSupport.nil]
    exact empty.subset
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have image := FunctionBetween.apply_closed with function, element
    have mapped : List(B, listMap{A}(f, tail)) by exact listMap.closed with function
    rw [listMap.cons, listSupport.cons]
    intros y hy
    have alternatives := insert.membership.forward with hy
    cases alternatives with head | rest
    . rw head
      have headIn := listSupport.mem_cons_self with element
      exact functionImage.membership.backward with ⟨image, a, headIn, rfl⟩
    . have inner := ih with y, rest
      obtain ⟨yIn, x, hx, equal⟩ := functionImage.membership.forward with inner
      have member := listSupport.mem_cons_of_mem with element, hx
      exact functionImage.membership.backward with ⟨yIn, x, member, equal⟩

theorem listMap.support (A B f : Set) (xs : List(A))
  requires function : f : A → B
  shows listSupport{B}(listMap{A}(f, xs)) = functionImage(f, listSupport{A}(xs), B)
proof
  have forward := listMap.support_subset with function
  apply Subset.antisymmetric with forward
  intros y hy
  obtain ⟨yIn, x, hx, equal⟩ := functionImage.membership.forward with hy
  rw equal
  exact listMap.mem_support with function, hx

-- The head and tail of a list over A.
theorem List.cons_parts (A x xs : Set)
  requires list : List(A, x :: xs)
  shows x ∈ A ∧ List(A, xs)
proof
  have shape := List.cases at A, x :: xs with list
  cases shape with isNil | isCons
  . exfalso
    have distinct := List.inductive.disjoint.nil.cons at x, xs
    symm at isNil
    exact distinct with isNil
  . obtain ⟨y, hy, tail, tailList, equal⟩ := isCons
    obtain ⟨headEqual, tailEqual⟩ :=
      List.inductive.injective.cons with equal
    rw [headEqual, tailEqual]
    exact ⟨hy, tailList⟩

-- The lists of a set S with a fixed entry prepended.
def consImage (x S : Set) : Set := { x :: ys | ys ∈ S }