-- Occurrences of an element in a list: the count, its interaction with
-- concatenation, reversal, filtering, and constant lists, and positivity
-- as membership in the support
-- 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.Count
import Filter

def listCount {A : Set} (x : Set) (xs : List(A)) → Set
  | x, List.nil => 0
  | x, List.cons(a, tail) => add(indicator({x}, a), listCount(x, tail))
  descends xs

theorem listCount.nil (A x : Set)
  shows listCount{A}(x, []) = 0
proof
  simp [listCount]

theorem listCount.cons (A x a : Set) (xs : List(A))
  requires element : a ∈ A
  shows listCount{A}(x, a :: xs) = indicator({x}, a) + listCount{A}(x, xs)
proof
  simp [listCount]

theorem listCount.natural (A x : Set) (xs : List(A))
  shows Nat(listCount{A}(x, xs))
proof
  induction xs
  | List.nil =>
    rw listCount.nil
    typecheck
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have mark := indicator.natural at {x}, a
    rw listCount.cons
    typecheck

theorem listCount.cons_self (A x : Set) (xs : List(A))
  requires element : x ∈ A
  shows listCount{A}(x, x :: xs) = 1 + listCount{A}(x, xs)
proof
  have mark := indicator.singleton_self at x
  rw [listCount.cons, mark]
  rfl

theorem listCount.cons_other (A x a : Set) (xs : List(A))
  requires element : a ∈ A
  requires different : x ≠ a
  shows listCount{A}(x, a :: xs) = listCount{A}(x, xs)
proof
  have mark := indicator.singleton_other with different
  have natural := listCount.natural at A, x, xs with *
  rw [listCount.cons, mark, add.left_unit]
  rfl

theorem listCount.singleton_self (A x : Set)
  requires element : x ∈ A
  shows listCount{A}(x, x :: []) = 1
proof
  have mark := indicator.singleton_self at x
  simp [listCount, mark, add]

theorem listCount.singleton_other (A x a : Set)
  requires element : a ∈ A
  requires different : x ≠ a
  shows listCount{A}(x, a :: []) = 0
proof
  have mark := indicator.singleton_other with different
  simp [listCount, mark, add]

theorem listCount.append (A x : Set) (xs ys : List(A))
  shows listCount{A}(x, append(ys, xs)) = listCount{A}(x, xs) + listCount{A}(x, ys)
proof
  induction xs
  | List.nil =>
    have natural := listCount.natural at A, x, ys with *
    rw [append.nil, listCount.nil, add.left_unit]
    rfl
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have appended := append.closed at A, tail, ys with *
    have mark := indicator.natural at {x}, a
    have tailNat := listCount.natural at A, x, tail with *
    have ysNat := listCount.natural at A, x, ys with *
    rw [append.cons, listCount.cons, ih, listCount.cons, add.associative]
    rfl

theorem listCount.reverse (A x : Set) (xs : List(A))
  shows listCount{A}(x, reverse(xs)) = listCount{A}(x, xs)
proof
  induction xs
  | List.nil =>
    rw reverse.nil
    rfl
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have single : List(A, a :: []) by typecheck
    have reversed : List(A, reverse(tail)) by typecheck
    have mark := indicator.natural at {x}, a
    have tailNat := listCount.natural at A, x, tail with *
    have nilList : List(A, []) by typecheck
    rw [reverse.cons, listCount.append, ih, listCount.cons, listCount.cons, listCount.nil]
    rw [add.right_unit, add.commutative]
    rfl

theorem listCount.le_length (A x : Set) (xs : List(A))
  shows listCount{A}(x, xs) ≤ length{A}(xs)
proof
  induction xs
  | List.nil =>
    rw [listCount.nil, length.nil]
    arith
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have tailNat := listCount.natural at A, x, tail with *
    rw [listCount.cons, length.cons]
    have alternatives : x = a ∨ x ≠ a by tauto
    cases alternatives with same | different
    . have mark : indicator({x}, a) = 1
        rw same
        exact indicator.singleton_self at a
      rw mark
      arith
    . have mark := indicator.singleton_other with different
      rw [mark, add.left_unit]
      arith

theorem listCount.positive_of_mem (A x : Set) (xs : List(A))
  requires member : x ∈ listSupport{A}(xs)
  shows 0 < listCount{A}(x, 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 tailNat := listCount.natural at A, x, tail with *
    have split := listSupport.mem_cons_iff.forward with element, member
    rw listCount.cons
    cases split with head | rest
    . have mark : indicator({x}, a) = 1
        rw head
        exact indicator.singleton_self at a
      rw mark
      arith
    . have inner := ih with rest
      have mark := indicator.natural at {x}, a
      arith

theorem listCount.mem_of_positive (A x : Set) (xs : List(A))
  requires positive : 0 < listCount{A}(x, xs)
  shows x ∈ listSupport{A}(xs)
proof
  induction xs
  | List.nil =>
    intro positive
    rw listCount.nil at positive
    exfalso
    exact Less.irreflexive with Naturals.constructor.zero, positive
  | List.cons(a, tail) using ih =>
    intro positive
    have element : a ∈ A by assumption
    have tailNat := listCount.natural at A, x, tail with *
    have alternatives : x = a ∨ x ≠ a by tauto
    cases alternatives with same | different
    . rw same
      exact listSupport.mem_cons_self with element
    . rw listCount.cons_other at positive
      have inner := ih with positive
      exact listSupport.mem_cons_of_mem with element, inner

theorem listCount.positive_iff (A x : Set) (xs : List(A))
  shows 0 < listCount{A}(x, xs) ↔ x ∈ listSupport{A}(xs)
proof
  constructor
  . intro positive
    exact listCount.mem_of_positive with positive
  . intro member
    exact listCount.positive_of_mem with member

theorem listCount.zero_of_absent (A x : Set) (xs : List(A))
  requires absent : ¬(x ∈ listSupport{A}(xs))
  shows listCount{A}(x, xs) = 0
proof
  have natural := listCount.natural at A, x, xs with *
  have shape := natural.zero_or_successor at listCount{A}(x, xs) with natural
  cases shape with zero | successor
  . exact zero
  . obtain ⟨k, kNat, equal⟩ := successor
    have positive : 0 < listCount{A}(x, xs)
      rw equal
      arith
    have member := listCount.mem_of_positive with positive
    exfalso
    exact absent with member

theorem listCount.absent_of_zero (A x : Set) (xs : List(A))
  requires zero : listCount{A}(x, xs) = 0
  shows ¬(x ∈ listSupport{A}(xs))
proof
  intro member
  have positive := listCount.positive_of_mem with member
  rw zero at positive
  exact Less.irreflexive with Naturals.constructor.zero, positive

theorem listCount.zero_iff (A x : Set) (xs : List(A))
  shows listCount{A}(x, xs) = 0 ↔ ¬(x ∈ listSupport{A}(xs))
proof
  constructor
  . intro zero
    exact listCount.absent_of_zero with zero
  . intro absent
    exact listCount.zero_of_absent with absent

theorem listCount.outside_carrier (A x : Set) (xs : List(A))
  requires outside : ¬(x ∈ A)
  shows listCount{A}(x, xs) = 0
proof
  have absent : ¬(x ∈ listSupport{A}(xs))
    intro member
    have subset := listSupport.subset at A, xs with *
    have inA := subset with x, member
    exact outside with inA
  exact listCount.zero_of_absent with absent

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

theorem listCount.replicate_other (A x y : Set) (n : Nat)
  requires element : x ∈ A
  requires different : y ≠ x
  shows listCount{A}(y, listReplicate(n, x)) = 0
proof
  have absent : ¬(y ∈ listSupport{A}(listReplicate(n, x)))
    intro member
    have equal := listReplicate.mem_support with element, member
    exact different with equal
  exact listCount.zero_of_absent with absent

theorem listCount.filter_member (A S x : Set) (xs : List(A))
  requires member : x ∈ S
  shows listCount{A}(x, listFilter{A}(S, xs)) = listCount{A}(x, xs)
proof
  induction xs
  | List.nil =>
    rw listFilter.nil
    rfl
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have filtered := listFilter.closed at A, S, tail with *
    have alternatives : a ∈ S ∨ ¬(a ∈ S) by tauto
    cases alternatives with inS | absent
    . rw [listFilter.cons_member, listCount.cons, listCount.cons, ih]
      rfl
    . have different : x ≠ a
        intro same
        rw same at member
        exact absent with member
      rw [listFilter.cons_other, listCount.cons_other, ih]
      rfl

theorem listCount.filter_other (A S x : Set) (xs : List(A))
  requires absent : ¬(x ∈ S)
  shows listCount{A}(x, listFilter{A}(S, xs)) = 0
proof
  have filtered := listFilter.closed at A, S, xs with *
  have outside : ¬(x ∈ listSupport{A}(listFilter{A}(S, xs)))
    intro member
    have both := listFilter.mem_support with member
    exact absent with both.right
  exact listCount.zero_of_absent with outside

-- The count is the length of the entries equal to x.
theorem listCount.filter_length (A x : Set) (xs : List(A))
  shows listCount{A}(x, xs) = length{A}(listFilter{A}({x}, xs))
proof
  induction xs
  | List.nil =>
    rw [listCount.nil, listFilter.nil, length.nil]
    rfl
  | List.cons(a, tail) using ih =>
    have element : a ∈ A by assumption
    have filtered := listFilter.closed at A, {x}, tail with *
    have alternatives : x = a ∨ x ≠ a by tauto
    cases alternatives with same | different
    . have elementX : x ∈ A
        rw same
        exact element
      have inSingle : x ∈ {x} := singleton.self at x
      rw [← same, listCount.cons_self, listFilter.cons_member, length.cons, ih]
      rw [add.commutative, add.right_one]
      rfl
    . have notSingle : ¬(a ∈ {x})
        intro inSingle
        have equal := singleton.membership.forward with inSingle
        symm at equal
        exact different with equal
      rw [listCount.cons_other, listFilter.cons_other, ih]
      rfl

theorem listCount.take_add_drop (A x : Set) (xs : List(A)) (n : Nat)
  shows listCount{A}(x, listTake{A}(n, xs)) + listCount{A}(x, listDrop{A}(n, xs)) =
    listCount{A}(x, 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 counts := listCount.append at A, x, listTake{A}(n, xs), listDrop{A}(n, xs) with *
  rw split at counts
  symm
  exact counts

theorem listCount.cons_le (A x a : Set) (xs : List(A))
  requires element : a ∈ A
  shows listCount{A}(x, xs) ≤ listCount{A}(x, a :: xs)
proof
  have mark := indicator.natural at {x}, a
  have tailNat := listCount.natural at A, x, xs with *
  rw listCount.cons
  arith

theorem listCount.eq_length (A x : Set) (xs : List(A))
  requires element : x ∈ A
  requires full : listCount{A}(x, xs) = length{A}(xs)
  shows xs = listReplicate(length{A}(xs), x)
proof
  have filtered := listFilter.closed at A, {x}, xs with *
  rw listCount.filter_length at full
  have same := listFilter.length_eq_of_all with full
  have covered := listFilter.support_in at A, {x}, xs with *
  rw same at covered
  exact List.constant_of_support with element, covered

theorem listCount.le_of_support_subset (A x : Set) (xs : List(A))
  requires element : x ∈ A
  requires covered : listSupport{A}(xs) ⊆ {x}
  shows listCount{A}(x, xs) = length{A}(xs)
proof
  have constant := List.constant_of_support with element, covered
  have count := listCount.replicate_self at A, x, length{A}(xs) with element
  rw ← constant at count
  exact count