-- Reading a list by position: the default is invisible inside the list,
-- positions mirror under reversal, the last entry of a snoc, and
-- extensionality by positions in both directions
-- 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.Index
import Sum

theorem listCoefficient.congr_default (A z w : Set) (xs : List(A)) (n : Nat)
  requires bound : n < length{A}(xs)
  shows listCoefficient{A}(z, xs, n) = listCoefficient{A}(w, xs, n)
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat bound
    rw length.nil at bound
    exfalso
    arith
  | List.cons(a, tail) using ih =>
    intros n nNat bound
    have element : a ∈ A by assumption
    rw length.cons at bound
    have shape := natural.zero_or_successor with *
    cases shape with zero | successor
    . rw [zero, listCoefficient.zero, listCoefficient.zero]
      rfl
    . obtain ⟨k, kNat, equal⟩ := successor
      rw equal at bound
      have smaller : k < length{A}(tail) by arith
      have inner := ih with k, kNat, smaller
      rw [equal, listCoefficient.succ, listCoefficient.succ]
      exact inner

theorem listCoefficient.singleton (A z a : Set)
  requires element : a ∈ A
  shows listCoefficient{A}(z, a :: [], 0) = a
proof
  simp [listCoefficient, listHeadOr]

theorem listCoefficient.lt_length_of_ne_default (A z : Set) (xs : List(A)) (n : Nat)
  requires different : listCoefficient{A}(z, xs, n) ≠ z
  shows n < length{A}(xs)
proof
  have alternatives : n < length{A}(xs) ∨ ¬(n < length{A}(xs)) by tauto
  cases alternatives with inside | outside
  . exact inside
  . have bound := LessEqual.of_not_Less at n, length{A}(xs) with outside
    have equal := listCoefficient.outside at A, z, xs, n with bound
    exfalso
    exact different with equal

theorem listCoefficient.count_positive (A z : Set) (xs : List(A)) (n : Nat)
  requires bound : n < length{A}(xs)
  shows 0 < listCount{A}(listCoefficient{A}(z, xs, n), xs)
proof
  have member := listCoefficient.mem_support at A, z, xs, n with bound
  exact listCount.positive_of_mem with member

theorem listSupport.mem_iff_index (A z y : Set) (xs : List(A))
  shows y ∈ listSupport{A}(xs) ↔
    ∃ n, Nat(n) ∧ n < length{A}(xs) ∧ listCoefficient{A}(z, xs, n) = y
proof
  constructor
  . intro member
    have cover := listSupport.exists_index at A, z, xs with *
    exact cover with y, member
  . intro witness
    obtain ⟨n, nNat, bound, value⟩ := witness
    have member := listCoefficient.mem_support at A, z, xs, n with *
    rwa value at member

-- The last entry of a list extended by a, and reading the reversal.
theorem listCoefficient.snoc_last (A z a : Set) (xs : List(A))
  requires element : a ∈ A
  shows listCoefficient{A}(z, append(a :: [], xs), length{A}(xs)) = a
proof
  have single : List(A, a :: []) by typecheck
  have step := listCoefficient.append_right at A, z, xs, a :: [], 0 with *
  rw [add.left_unit, listCoefficient.singleton] at step
  exact step

theorem listCoefficient.snoc_below (A z a : Set) (xs : List(A)) (n : Nat)
  requires element : a ∈ A
  requires bound : n < length{A}(xs)
  shows listCoefficient{A}(z, append(a :: [], xs), n) = listCoefficient{A}(z, xs, n)
proof
  have single : List(A, a :: []) by typecheck
  exact listCoefficient.append_left with bound

theorem listCoefficient.reverse (A z : Set) (xs : List(A)) (n m : Nat)
  requires split : n + m + 1 = length{A}(xs)
  shows listCoefficient{A}(z, reverse(xs), n) = listCoefficient{A}(z, xs, m)
proof
  induction xs generalizing n m
  | List.nil =>
    intros n m nNat mNat split
    rw [length.nil] at split
    have sum : Nat(n + m) by typecheck
    have distinct := Naturals.disjoint.zero.succ at n + m
    rw add.right_one at split
    exfalso
    symm at split
    exact distinct with split
  | List.cons(a, tail) using ih =>
    intros n m nNat mNat split
    have element : a ∈ A by assumption
    have single : List(A, a :: []) by typecheck
    have reversed : List(A, reverse(tail)) by typecheck
    have reversedLength := reverse.length at A, tail with *
    rw [length.cons, add.right_one] at split
    have sum := Naturals.injective.succ with split
    rw reverse.cons
    have shape := natural.zero_or_successor at m with mNat
    cases shape with zero | successor
    . rw [zero, add.right_unit] at sum
      rw [zero, listCoefficient.zero, sum, ← reversedLength]
      exact listCoefficient.snoc_last with element
    . obtain ⟨k, kNat, equal⟩ := successor
      rw [equal, add.succ_right] at sum
      have smaller : n < length{A}(reverse(tail))
        rw reversedLength
        rw ← sum
        arith
      have inner : n + k + 1 = length{A}(tail)
        rw add.right_one
        exact sum
      have recursive := ih with n, k, nNat, kNat, inner
      rw [equal, listCoefficient.succ, listCoefficient.snoc_below, recursive]
      rfl

theorem listCoefficient.reverse_head (A z : Set) (xs : List(A)) (n : Nat)
  requires split : n + 1 = length{A}(xs)
  shows listCoefficient{A}(z, reverse(xs), 0) = listCoefficient{A}(z, xs, n)
proof
  have shifted : 0 + n + 1 = length{A}(xs)
    rw add.left_unit
    exact split
  exact listCoefficient.reverse with shifted

theorem listCoefficient.reverse_last (A z : Set) (xs : List(A)) (n : Nat)
  requires split : n + 1 = length{A}(xs)
  shows listCoefficient{A}(z, reverse(xs), n) = listCoefficient{A}(z, xs, 0)
proof
  have shifted : n + 0 + 1 = length{A}(xs)
    rw add.right_unit
    exact split
  exact listCoefficient.reverse with shifted

-- Extensionality by positions, in both directions and with every position.
theorem listCoefficient.ext_iff (A z : Set) (xs ys : List(A))
  shows xs = ys ↔ (length{A}(xs) = length{A}(ys) ∧ (∀ n, Nat(n) → n < length{A}(xs) →
    listCoefficient{A}(z, xs, n) = listCoefficient{A}(z, ys, n)))
proof
  constructor
  . intro equal
    constructor
    . rw equal
      rfl
    . intros n nNat bound
      rw equal
      rfl
  . intro both
    obtain ⟨sameLength, entries⟩ := both
    exact listCoefficient.ext with sameLength, entries

theorem listCoefficient.ext_all (A z : Set) (xs ys : List(A))
  requires sameLength : length{A}(xs) = length{A}(ys)
  requires entries : ∀ n, Nat(n) → listCoefficient{A}(z, xs, n) = listCoefficient{A}(z, ys, n)
  shows xs = ys
proof
  have bounded : ∀ n, Nat(n) → n < length{A}(xs) →
      listCoefficient{A}(z, xs, n) = listCoefficient{A}(z, ys, n)
    intros n nNat bound
    exact entries with n, nNat
  exact listCoefficient.ext with sameLength, bounded

theorem listCoefficient.eq_of_entries (A z : Set) (xs ys : List(A))
  requires entries : ∀ n, Nat(n) → listCoefficient{A}(z, xs, n) = listCoefficient{A}(z, ys, n)
  requires default : ¬(z ∈ A)
  shows xs = ys
proof
  have sameLength : length{A}(xs) = length{A}(ys)
    have xsNat : Nat(length{A}(xs)) by typecheck
    have ysNat : Nat(length{A}(ys)) by typecheck
    have trichotomy := Less.trichotomy at length{A}(xs), length{A}(ys) with *
    rcases trichotomy with less | equal | greater
    . have bound : length{A}(xs) ≤ length{A}(xs) by arith
      have outside := listCoefficient.outside at A, z, xs, length{A}(xs) with bound
      have inside := listCoefficient.mem_support at A, z, ys, length{A}(xs) with less
      have value := entries with length{A}(xs), xsNat
      rw [← value, outside] at inside
      have subset := listSupport.subset at A, ys with *
      have inA := subset with z, inside
      exfalso
      exact default with inA
    . exact equal
    . have bound : length{A}(ys) ≤ length{A}(ys) by arith
      have outside := listCoefficient.outside at A, z, ys, length{A}(ys) with bound
      have inside := listCoefficient.mem_support at A, z, xs, length{A}(ys) with greater
      have value := entries with length{A}(ys), ysNat
      rw [value, outside] at inside
      have subset := listSupport.subset at A, xs with *
      have inA := subset with z, inside
      exfalso
      exact default with inA
  exact listCoefficient.ext_all with sameLength, entries

theorem listCoefficient.take_ext (A z : Set) (xs ys : List(A)) (n : Nat)
  requires boundXs : n ≤ length{A}(xs)
  requires boundYs : n ≤ length{A}(ys)
  requires entries : ∀ k, Nat(k) → k < n →
    listCoefficient{A}(z, xs, k) = listCoefficient{A}(z, ys, k)
  shows listTake{A}(n, xs) = listTake{A}(n, ys)
proof
  have takenXs := listTake.closed at A, xs, n with *
  have takenYs := listTake.closed at A, ys, n with *
  have lengthXs := listTake.length with boundXs
  have lengthYs := listTake.length with boundYs
  have sameLength : length{A}(listTake{A}(n, xs)) = length{A}(listTake{A}(n, ys))
    rw [lengthXs, lengthYs]
    rfl
  have positions : ∀ k, Nat(k) → k < length{A}(listTake{A}(n, xs)) →
      listCoefficient{A}(z, listTake{A}(n, xs), k) = listCoefficient{A}(z, listTake{A}(n, ys), k)
    intros k kNat bound
    rw lengthXs at bound
    have left := listTake.coefficient at A, z, xs, n, k with bound
    have right := listTake.coefficient at A, z, ys, n, k with bound
    rw [left, right]
    exact entries with k, kNat, bound
  exact listCoefficient.ext with sameLength, positions

theorem listCoefficient.drop_ext (A z : Set) (xs ys : List(A)) (n : Nat)
  requires sameLength : length{A}(xs) = length{A}(ys)
  requires entries : ∀ k, Nat(k) → n ≤ k →
    listCoefficient{A}(z, xs, k) = listCoefficient{A}(z, ys, k)
  shows listDrop{A}(n, xs) = listDrop{A}(n, ys)
proof
  have droppedXs := listDrop.closed at A, xs, n with *
  have droppedYs := listDrop.closed at A, ys, n with *
  have positions : ∀ k, Nat(k) →
      listCoefficient{A}(z, listDrop{A}(n, xs), k) = listCoefficient{A}(z, listDrop{A}(n, ys), k)
    intros k kNat
    have left := listDrop.coefficient at A, z, xs, n, k with *
    have right := listDrop.coefficient at A, z, ys, n, k with *
    have shifted : Nat(n + k) by typecheck
    have above : n ≤ n + k by arith
    rw [left, right]
    exact entries with n + k, shifted, above
  have lengths : length{A}(listDrop{A}(n, xs)) = length{A}(listDrop{A}(n, ys))
    have alternatives : n ≤ length{A}(xs) ∨ ¬(n ≤ length{A}(xs)) by tauto
    cases alternatives with inside | outside
    . have insideYs : n ≤ length{A}(ys) by rwa ← sameLength
      have left := listDrop.length with inside
      have right := listDrop.length with insideYs
      rw sameLength at left
      rw ← right at left
      exact add.right_cancel with left
    . have beyond := Less.of_not_LessEqual at n, length{A}(xs) with outside
      have boundXs : length{A}(xs) ≤ n by arith
      have boundYs : length{A}(ys) ≤ n by rwa ← sameLength
      have left := listDrop.all with boundXs
      have right := listDrop.all with boundYs
      rw [left, right]
      rfl
  exact listCoefficient.ext_all with lengths, positions