-- Constant lists: n copies of one element, their length, support, and
-- interaction with concatenation, mapping, reversal, 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.Replicate
import Basic

def listReplicate (n : Nat, x : Set) → Set
  | zero, x => []
  | succ(k), x => x :: listReplicate(k, x)

theorem listReplicate.zero (x : Set)
  shows listReplicate(0, x) = []
proof
  simp [listReplicate]

theorem listReplicate.succ (n : Nat) (x : Set)
  shows listReplicate(succ(n), x) = x :: listReplicate(n, x)
proof
  simp [listReplicate]

theorem listReplicate.closed (A x : Set) (n : Nat)
  requires element : x ∈ A
  shows List(A, listReplicate(n, x))
proof
  induction n
  | zero =>
    rw listReplicate.zero
    typecheck
  | succ(k) using ih =>
    rw listReplicate.succ
    typecheck

theorem listReplicate.length (A x : Set) (n : Nat)
  requires element : x ∈ A
  shows length{A}(listReplicate(n, x)) = n
proof
  induction n
  | zero =>
    rw [listReplicate.zero, length.nil]
    rfl
  | succ(k) using ih =>
    have list := listReplicate.closed at A, x, k with element
    rw [listReplicate.succ, length.cons, ih]
    rfl

theorem listReplicate.one (x : Set)
  shows listReplicate(1, x) = x :: []
proof
  simp [listReplicate]

theorem listReplicate.support_subset (A x : Set) (n : Nat)
  requires element : x ∈ A
  shows listSupport{A}(listReplicate(n, x)) ⊆ {x}
proof
  induction n
  | zero =>
    rw [listReplicate.zero, listSupport.nil]
    exact empty.subset
  | succ(k) using ih =>
    have list := listReplicate.closed at A, x, k with element
    rw [listReplicate.succ, listSupport.cons]
    intros y hy
    have alternatives := insert.membership.forward with hy
    cases alternatives with head | rest
    . rw head
      exact singleton.self at x
    . exact ih with y, rest

theorem listReplicate.support_succ (A x : Set) (n : Nat)
  requires element : x ∈ A
  shows listSupport{A}(listReplicate(succ(n), x)) = {x}
proof
  have list := listReplicate.closed at A, x, n with element
  have subset := listReplicate.support_subset at A, x, succ(n) with element
  apply Subset.antisymmetric with subset
  intros y hy
  have equal := singleton.membership.forward with hy
  rw [equal, listReplicate.succ]
  exact listSupport.mem_cons_self with element

theorem listReplicate.mem_support (A x y : Set) (n : Nat)
  requires element : x ∈ A
  requires member : y ∈ listSupport{A}(listReplicate(n, x))
  shows y = x
proof
  have subset := listReplicate.support_subset at A, x, n with element
  have single := subset with y, member
  exact singleton.membership.forward with single

theorem listReplicate.mem_support_iff (A x y : Set) (n : Nat)
  requires element : x ∈ A
  shows y ∈ listSupport{A}(listReplicate(n, x)) ↔ 0 < n ∧ y = x
proof
  constructor
  . intro member
    have equal := listReplicate.mem_support with element, member
    have nonempty : listReplicate(n, x) ≠ []
      intro empty
      rw [empty, listSupport.nil] at member
      exact set.empty with member
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listReplicate.zero] at nonempty
      exfalso
      have same : [] = [] by rfl
      exact nonempty with same
    . obtain ⟨k, kNat, equal2⟩ := successor
      rw equal2
      constructor
      . arith
      . exact equal
  . intro both
    obtain ⟨positive, equal⟩ := both
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw zero at positive
      exfalso
      exact Less.irreflexive with Naturals.constructor.zero, positive
    . obtain ⟨k, kNat, equal2⟩ := successor
      rw [equal2, listReplicate.support_succ, equal]
      exact singleton.self at x

theorem listReplicate.append (A x : Set) (m n : Nat)
  requires element : x ∈ A
  shows append{A}(listReplicate(m, x), listReplicate(n, x)) = listReplicate(n + m, x)
proof
  induction n
  | zero =>
    have list := listReplicate.closed at A, x, m with element
    rw [listReplicate.zero, add.left_unit]
    simp [append]
  | succ(k) using ih =>
    have list := listReplicate.closed at A, x, k with element
    have other := listReplicate.closed at A, x, m with element
    rw [listReplicate.succ, append.cons, ih]
    have shift : succ(k) + m = succ(k + m) by simp [add, add.succ_right]
    rw [shift, listReplicate.succ]
    rfl

theorem listReplicate.map (A B f x : Set) (n : Nat)
  requires function : f : A → B
  requires element : x ∈ A
  shows listMap{A}(f, listReplicate(n, x)) = listReplicate(n, f(x))
proof
  induction n
  | zero =>
    rw [listReplicate.zero, listMap.nil, listReplicate.zero]
    rfl
  | succ(k) using ih =>
    have list := listReplicate.closed at A, x, k with element
    rw [listReplicate.succ, listMap.cons, ih, listReplicate.succ]
    rfl

theorem listReplicate.reverse (A x : Set) (n : Nat)
  requires element : x ∈ A
  shows reverse{A}(listReplicate(n, x)) = listReplicate(n, x)
proof
  induction n
  | zero =>
    rw listReplicate.zero
    simp [reverse]
  | succ(k) using ih =>
    have list := listReplicate.closed at A, x, k with element
    have reversed : List(A, reverse(listReplicate(k, x))) by typecheck
    have single : List(A, x :: []) by typecheck
    rw [listReplicate.succ, reverse.cons, ih]
    have one := listReplicate.one at x
    rw [← one, listReplicate.append, add.right_one]
    exact listReplicate.succ at k, x

theorem listReplicate.tail (A x : Set) (n : Nat)
  requires element : x ∈ A
  shows listTail{A}(listReplicate(succ(n), x)) = listReplicate(n, x)
proof
  have list := listReplicate.closed at A, x, n with element
  rw [listReplicate.succ, listTail.cons]
  rfl

theorem listReplicate.coefficient (A z x : Set) (n k : Nat)
  requires element : x ∈ A
  requires bound : k < n
  shows listCoefficient{A}(z, listReplicate(n, x), k) = x
proof
  induction n generalizing k
  | zero =>
    intros k kNat bound
    exfalso
    arith
  | succ(j) using ih =>
    intros k kNat bound
    have list := listReplicate.closed at A, x, j with element
    have shape := natural.zero_or_successor at k with *
    rw listReplicate.succ
    cases shape with zero | successor
    . rw [zero, listCoefficient.zero]
      rfl
    . obtain ⟨r, rNat, equal⟩ := successor
      rw equal at bound
      have smaller : r < j by arith
      rw [equal, listCoefficient.succ]
      exact ih with r, rNat, smaller

-- A list whose support lies in a singleton is constant.
theorem List.constant_of_support (A x : Set) (xs : List(A))
  requires element : x ∈ A
  requires support : listSupport{A}(xs) ⊆ {x}
  shows xs = listReplicate(length{A}(xs), x)
proof
  induction xs
  | List.nil =>
    intro support
    rw [length.nil, listReplicate.zero]
    rfl
  | List.cons(a, tail) using ih =>
    intro support
    have headA : a ∈ A by assumption
    have headIn := listSupport.mem_cons_self with headA
    have headSingle := support with a, headIn
    have head := singleton.membership.forward with headSingle
    have rest : listSupport{A}(tail) ⊆ {x}
      intros y hy
      have member := listSupport.mem_cons_of_mem with headA, hy
      exact support with y, member
    have tails := ih with rest
    rw [length.cons, listReplicate.succ, head]
    rw ← tails
    rfl

theorem listReplicate.injective_length (A x : Set) (m n : Nat)
  requires element : x ∈ A
  requires equal : listReplicate(m, x) = listReplicate(n, x)
  shows m = n
proof
  have lengths : length{A}(listReplicate(m, x)) = length{A}(listReplicate(n, x))
    rw equal
    rfl
  rw [listReplicate.length, listReplicate.length] at lengths
  exact lengths