-- Permutations of lists: equality of all counts, its equivalence laws,
-- splitting a list at an occurrence, and invariance of length, support,
-- and filtering
-- 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.Permutation
import Count

-- Two lists are permutations of each other when every element occurs
-- equally often in both.
predicate ListPermutation (A xs ys : Set) :=
  ∀ x, listCount{A}(x, xs) = listCount{A}(x, ys)

theorem ListPermutation.count (A x : Set) (xs ys : List(A))
  requires permutation : ListPermutation(A, xs, ys)
  shows listCount{A}(x, xs) = listCount{A}(x, ys)
proof
  exact permutation with x

theorem ListPermutation.of_eq (A : Set) (xs ys : List(A))
  requires equal : xs = ys
  shows ListPermutation(A, xs, ys)
proof
  unfold ListPermutation
  intro x
  rw equal
  rfl

theorem ListPermutation.reflexive (A : Set) (xs : List(A))
  shows ListPermutation(A, xs, xs)
proof
  unfold ListPermutation
  intro x
  rfl

theorem ListPermutation.symmetric (A : Set) (xs ys : List(A))
  requires permutation : ListPermutation(A, xs, ys)
  shows ListPermutation(A, ys, xs)
proof
  unfold ListPermutation
  intro x
  have count := permutation with x
  symm at count
  exact count

theorem ListPermutation.transitive (A : Set) (xs ys zs : List(A))
  requires first : ListPermutation(A, xs, ys)
  requires second : ListPermutation(A, ys, zs)
  shows ListPermutation(A, xs, zs)
proof
  unfold ListPermutation
  intro x
  have left := first with x
  have right := second with x
  rw left
  exact right

theorem ListPermutation.cons (A a : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires permutation : ListPermutation(A, xs, ys)
  shows ListPermutation(A, a :: xs, a :: ys)
proof
  unfold ListPermutation
  intro x
  have count := permutation with x
  rw [listCount.cons, listCount.cons, count]
  rfl

theorem ListPermutation.cons_cancel (A a : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires permutation : ListPermutation(A, a :: xs, a :: ys)
  shows ListPermutation(A, xs, ys)
proof
  unfold ListPermutation
  intro x
  have count := permutation with x
  have mark := indicator.natural at {x}, a
  rw [listCount.cons, listCount.cons] at count
  exact add.left_cancel with count

theorem ListPermutation.append_right (A : Set) (xs ys zs : List(A))
  requires permutation : ListPermutation(A, xs, ys)
  shows ListPermutation(A, append(zs, xs), append(zs, ys))
proof
  unfold ListPermutation
  intro x
  have count := permutation with x
  rw [listCount.append, listCount.append, count]
  rfl

theorem ListPermutation.append_left (A : Set) (xs ys zs : List(A))
  requires permutation : ListPermutation(A, xs, ys)
  shows ListPermutation(A, append(xs, zs), append(ys, zs))
proof
  unfold ListPermutation
  intro x
  have count := permutation with x
  rw [listCount.append, listCount.append, count]
  rfl

theorem ListPermutation.append (A : Set) (xs ys us vs : List(A))
  requires first : ListPermutation(A, xs, ys)
  requires second : ListPermutation(A, us, vs)
  shows ListPermutation(A, append(us, xs), append(vs, ys))
proof
  unfold ListPermutation
  intro x
  have left := first with x
  have right := second with x
  rw [listCount.append, listCount.append, left, right]
  rfl

theorem ListPermutation.append_commutative (A : Set) (xs ys : List(A))
  shows ListPermutation(A, append(ys, xs), append(xs, ys))
proof
  unfold ListPermutation
  intro x
  have xsNat := listCount.natural at A, x, xs with *
  have ysNat := listCount.natural at A, x, ys with *
  rw [listCount.append, listCount.append, add.commutative]
  rfl

theorem ListPermutation.reverse (A : Set) (xs : List(A))
  shows ListPermutation(A, reverse(xs), xs)
proof
  unfold ListPermutation
  intro x
  exact listCount.reverse at A, x, xs with *

theorem ListPermutation.swap (A a b : Set) (xs : List(A))
  requires first : a ∈ A
  requires second : b ∈ A
  shows ListPermutation(A, a :: b :: xs, b :: a :: xs)
proof
  unfold ListPermutation
  intro x
  have inner : List(A, b :: xs) by typecheck
  have other : List(A, a :: xs) by typecheck
  have markA := indicator.natural at {x}, a
  have markB := indicator.natural at {x}, b
  have count := listCount.natural at A, x, xs with *
  rw [listCount.cons, listCount.cons, listCount.cons, listCount.cons]
  exact add.swap at indicator({x}, b), listCount{A}(x, xs), indicator({x}, a) with *

theorem ListPermutation.middle (A a : Set) (xs ys : List(A))
  requires element : a ∈ A
  shows ListPermutation(A, append(ys, a :: xs), a :: append(ys, xs))
proof
  unfold ListPermutation
  intro x
  have whole : List(A, a :: xs) by typecheck
  have appended := append.closed at A, xs, ys with *
  have mark := indicator.natural at {x}, a
  have xsNat := listCount.natural at A, x, xs with *
  have ysNat := listCount.natural at A, x, ys with *
  rw [listCount.append, listCount.cons, listCount.cons, listCount.append]
  exact add.associative at listCount{A}(x, xs), listCount{A}(x, ys), indicator({x}, a) with *

theorem ListPermutation.rotate (A : Set) (xs : List(A)) (n : Nat)
  shows ListPermutation(A, xs, append(listTake{A}(n, xs), 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 swapped := ListPermutation.append_commutative at A, listTake{A}(n, xs), listDrop{A}(n, xs) with *
  rwa split at swapped

theorem ListPermutation.mem_support (A y : Set) (xs ys : List(A))
  requires permutation : ListPermutation(A, xs, ys)
  requires member : y ∈ listSupport{A}(xs)
  shows y ∈ listSupport{A}(ys)
proof
  have positive := listCount.positive_of_mem with member
  have count := permutation with y
  rw count at positive
  exact listCount.mem_of_positive with positive

theorem ListPermutation.support (A : Set) (xs ys : List(A))
  requires permutation : ListPermutation(A, xs, ys)
  shows listSupport{A}(xs) = listSupport{A}(ys)
proof
  have back := ListPermutation.symmetric with permutation
  ext y
  constructor
  . intro member
    exact ListPermutation.mem_support with permutation, member
  . intro member
    exact ListPermutation.mem_support with back, member

theorem ListPermutation.nil (A : Set) (xs : List(A))
  requires permutation : ListPermutation(A, xs, [])
  shows xs = []
proof
  have support := ListPermutation.support with permutation
  rw listSupport.nil at support
  exact listSupport.eq_empty with support

theorem ListPermutation.filter (A S : Set) (xs ys : List(A))
  requires permutation : ListPermutation(A, xs, ys)
  shows ListPermutation(A, listFilter{A}(S, xs), listFilter{A}(S, ys))
proof
  unfold ListPermutation
  intro x
  have alternatives : x ∈ S ∨ ¬(x ∈ S) by tauto
  cases alternatives with member | absent
  . have left := listCount.filter_member at A, S, x, xs with member
    have right := listCount.filter_member at A, S, x, ys with member
    rw [left, right]
    exact permutation with x
  . have left := listCount.filter_other at A, S, x, xs with absent
    have right := listCount.filter_other at A, S, x, ys with absent
    rw [left, right]
    rfl

-- A list is split at any element of its support.
theorem listSupport.split (A a : Set) (xs : List(A))
  requires member : a ∈ listSupport{A}(xs)
  shows ∃ front, List(A, front) ∧ ∃ back, List(A, back) ∧ xs = append(a :: back, front)
proof
  induction xs
  | List.nil =>
    intro member
    rw listSupport.nil at member
    exfalso
    exact set.empty with member
  | List.cons(b, tail) using ih =>
    intro member
    have element : b ∈ A by assumption
    have split := listSupport.mem_cons_iff.forward with element, member
    cases split with head | rest
    . have nilList : List(A, []) by typecheck
      refine ⟨[], nilList, tail, assumption, ?_⟩
      rw head
      simp [append]
    . obtain ⟨front, frontList, back, backList, equal⟩ := ih with rest
      have subset := listSupport.subset at A, tail with *
      have aA := subset with a, rest
      have frontCons : List(A, b :: front) by typecheck
      have backCons : List(A, a :: back) by typecheck
      refine ⟨b :: front, frontCons, back, backList, ?_⟩
      rw [append.cons, ← equal]
      rfl

-- Removing the head of a permuted list from the other side.
theorem ListPermutation.cons_split (A a : Set) (xs ys : List(A))
  requires element : a ∈ A
  requires permutation : ListPermutation(A, a :: xs, ys)
  shows ∃ front, List(A, front) ∧ ∃ back, List(A, back) ∧
    ys = append(a :: back, front) ∧ ListPermutation(A, xs, append(back, front))
proof
  have headIn := listSupport.mem_cons_self at A, a, xs with element
  have member := ListPermutation.mem_support with permutation, headIn
  obtain ⟨front, frontList, back, backList, equal⟩ := listSupport.split with member
  refine ⟨front, frontList, back, backList, equal, ?_⟩
  unfold ListPermutation
  intro x
  have count := permutation with x
  have backCons : List(A, a :: back) by typecheck
  have mark := indicator.natural at {x}, a
  have xsNat := listCount.natural at A, x, xs with *
  have frontNat := listCount.natural at A, x, front with *
  have backNat := listCount.natural at A, x, back with *
  rw [equal, listCount.cons, listCount.append, listCount.cons] at count
  rw listCount.append
  have shuffled : listCount{A}(x, front) + (indicator({x}, a) + listCount{A}(x, back)) =
    indicator({x}, a) + (listCount{A}(x, front) + listCount{A}(x, back))
    exact add.swap at indicator({x}, a), listCount{A}(x, back), listCount{A}(x, front) with *
  rw shuffled at count
  exact add.left_cancel with count

theorem ListPermutation.length (A : Set) (xs ys : List(A))
  requires permutation : ListPermutation(A, xs, ys)
  shows length{A}(xs) = length{A}(ys)
proof
  induction xs generalizing ys
  | List.nil =>
    intros ys ysList permutation
    have back := ListPermutation.symmetric with permutation
    have empty := ListPermutation.nil with back
    rw empty
    rfl
  | List.cons(a, tail) using ih =>
    intros ys ysList permutation
    have element : a ∈ A by assumption
    obtain ⟨front, frontList, back, backList, equal, rest⟩ :=
      ListPermutation.cons_split with element, permutation
    have appended := append.closed at A, front, back with *
    have inner := ih with append(back, front), appended, rest
    have backCons : List(A, a :: back) by typecheck
    rw [equal, length.cons, inner, length.append, length.append, length.cons]
    have shift : succ(length{A}(front) + length{A}(back)) = length{A}(front) + succ(length{A}(back))
      rw add.succ_right
      rfl
    exact shift

theorem ListPermutation.singleton (A a : Set) (xs : List(A))
  requires element : a ∈ A
  requires permutation : ListPermutation(A, a :: [], xs)
  shows xs = a :: []
proof
  obtain ⟨front, frontList, back, backList, equal, rest⟩ :=
    ListPermutation.cons_split with element, permutation
  have appended := append.closed at A, front, back with *
  have back2 := ListPermutation.symmetric with rest
  have empty := ListPermutation.nil with back2
  have both := append.eq_nil with empty
  obtain ⟨frontNil, backNil⟩ := both
  rw [equal, frontNil, backNil]
  simp [append]

theorem ListPermutation.replicate (A x : Set) (xs : List(A)) (n : Nat)
  requires element : x ∈ A
  requires permutation : ListPermutation(A, xs, listReplicate(n, x))
  shows xs = listReplicate(n, x)
proof
  have list := listReplicate.closed at A, x, n with element
  have support := ListPermutation.support with permutation
  have lengths := ListPermutation.length with permutation
  have covered := listReplicate.support_subset at A, x, n with element
  rw ← support at covered
  have constant := List.constant_of_support with element, covered
  rw [listReplicate.length] at lengths
  rw lengths at constant
  exact constant

theorem ListPermutation.length_filter (A S : Set) (xs ys : List(A))
  requires permutation : ListPermutation(A, xs, ys)
  shows length{A}(listFilter{A}(S, xs)) = length{A}(listFilter{A}(S, ys))
proof
  have filtered := ListPermutation.filter at A, S, xs, ys with permutation
  have leftList := listFilter.closed at A, S, xs with *
  have rightList := listFilter.closed at A, S, ys with *
  exact ListPermutation.length with filtered