⟨ Index | Module Fin-Pigeonhole-Principle 
{-# OPTIONS --safe  #-}

open import Level using (Level; 0ℓ)
import Level as Level
open import Data.Nat
open import Data.Sum
open import Data.Fin as Fin using (Fin)
open import Data.Empty
open import Data.List
open import Data.Vec as V using (Vec)
import Data.Vec.Properties as V
open import Data.List as L
import Data.List.Properties as L
import Data.Nat.Properties as ℕ
open import Data.Product
open import Function.Base
open import Relation.Binary.PropositionalEquality
open import Data.List.Relation.Unary.Any using (Any)
open import Relation.Nullary.Decidable

open import Data.List.Membership.Propositional
open import Data.List.Membership.Propositional.Properties as L

module Fin-Pigeonhole-Principle  where

private
  variable
    ℓ : Level
    A : Set ℓ

∷⁻¹ : {k : ℕ} → Vec A (suc k) → A × Vec A k
∷⁻¹ (x V.∷ xs) = x , xs

VΣ : { k : ℕ } → (A → ℕ) →  Vec A k → ℕ
VΣ A→ℕ vec = V.sum (V.map A→ℕ vec)

-- Pigeonhole principle for sums: if the sum of k natural numbers is smaller than
-- k * n, then at least one of the numbers must have been smaller than n.
VΣ-php : { k : ℕ } → {f : A → ℕ} → (vec : Vec A k) → ∀ n → VΣ f vec < k * n → ∃[ i ] f (V.lookup vec i) < n
VΣ-php {k = ℕ.suc k} {f = f} (x V.∷ vec) n Σx∷vec<suc[k]*n with (ℕ.<-≤-connex (f x) n)
... | inj₁ f[x]<n = Fin.zero , f[x]<n
... | inj₂ n≤f[x] = (Fin.suc (VΣ-php' .proj₁)) , VΣ-php' .proj₂
  where
    Σvec<k*n : VΣ f vec < k * n
    Σvec<k*n = ℕ.+-cancelˡ-< (f x) _ (k * n) (ℕ.<-≤-trans Σx∷vec<suc[k]*n (ℕ.+-monoˡ-≤ (k * n) n≤f[x]))
    VΣ-php' : ∃[ j ] f (V.lookup vec j) < n
    VΣ-php' = VΣ-php vec n Σvec<k*n

VΣ-suc : { k : ℕ } → (f : A → ℕ) → (cons : A → A) → (cons+1 : ∀ {a} → f (cons a) ≡ suc (f a)) → (list : Vec A k) (i : Fin k) → VΣ f (list V.[ i ]%= cons) ≡ suc (VΣ f list)
VΣ-suc f cons cons+1 (x V.∷ list) (Fin.zero) = cong₂ _+_ cons+1 refl
VΣ-suc f cons cons+1 (x V.∷ list) (Fin.suc i) =
  begin
  VΣ f ((x V.∷ list) V.[ Fin.suc i ]%= cons) ≡⟨⟩
  f x + (VΣ f (list V.[ i ]%= cons)) ≡⟨ cong (_ +_) (VΣ-suc f cons cons+1 (list) _) ⟩
  f x + suc (VΣ f list) ≡⟨ ℕ.+-suc (f x) _ ⟩
  suc (VΣ f (x V.∷ list))
  ∎
  where
    open ≡-Reasoning

split-buckets : {k : ℕ} → List ((Fin k) × A) → Vec (List A) k
split-buckets [] = V.replicate _ []
split-buckets ((hd , tl) ∷ xs) = V.updateAt (split-buckets xs) hd (tl L.∷_)

split-buckets-∈ : {k : ℕ} → (list : List ((Fin k) × A)) → {h : Fin k} {t : A} → (h , t) ∈ list → t ∈ V.lookup (split-buckets list) h
split-buckets-∈ [] ()
split-buckets-∈ ((h , t) ∷ xs) {_} {_} (Any.here refl) = subst (t ∈_) (sym (V.lookup∘updateAt h {t ∷_} (split-buckets xs)))
                                                          (Any.here refl)
split-buckets-∈ ((x , tl) ∷ xs) {h} {t} (Any.there h,t∈xs) with (h Fin.≟ x)
... | no h≢x = subst (t ∈_) (sym (V.lookup∘updateAt′ h x h≢x (split-buckets xs))) (split-buckets-∈ xs h,t∈xs)
... | yes refl = subst (t ∈_) (sym (V.lookup∘updateAt h (split-buckets xs))) (Any.there (split-buckets-∈ xs h,t∈xs))

split-buckets-length : {k : ℕ} → {list : List ((Fin k) × A)} → VΣ L.length (split-buckets list) ≡ L.length list
split-buckets-length {k = zero} {list = []} = refl
split-buckets-length {k = suc k} {list = []} = split-buckets-length {k = k} {list = []}
split-buckets-length {k = k} {list = (hd , tl) ∷ xs} =
  begin
  VΣ length (split-buckets ((hd , tl) ∷ xs)) ≡⟨⟩
  VΣ length ((split-buckets xs) V.[ hd ]%= (tl L.∷_)) ≡⟨ VΣ-suc L.length ((tl ∷_)) refl (split-buckets xs) hd ⟩
  1 + (VΣ length (split-buckets xs)) ≡⟨ cong suc (split-buckets-length {list = xs}) ⟩
  1 + length xs
  ∎
  where
    open ≡-Reasoning

-- for an alphabet of size k, a list of fewer than k ^ n words can't contain all words of length n:
find-missing : {k n : ℕ} → (list : List (Vec (Fin k) n)) → length list < k ^ n → ∃[ w ] w ∉ list
find-missing {n = zero} [] len<k^n = V.[] , λ ()
find-missing {n = zero} (x ∷ xs) (s≤s ())
find-missing {k} {n = suc n} list len<k*k^n =
  (i V.∷ w) , λ iw∈list →
    proj₂ small-bucket-lacks ((w ∈ bucket[ i ]) ∋
          split-buckets-∈ _ (L.∈-map⁺ V.uncons iw∈list))
  where
    unconsed : List (Fin k × Vec (Fin k) n)
    unconsed = L.map V.uncons list 
    buckets : Vec (List (Vec (Fin k) n)) k
    buckets = split-buckets unconsed

    bucket[_] : Fin k → List (Vec (Fin k) n)
    bucket[_] = V.lookup buckets

    Σlen-buckets≡len-list : VΣ L.length buckets ≡ L.length list
    Σlen-buckets≡len-list = begin
      VΣ length buckets ≡⟨ split-buckets-length {k = k} {unconsed} ⟩
      length unconsed ≡⟨ L.length-map V.uncons list ⟩
      length list
      ∎
      where open ≡-Reasoning
    ∃small-bucket : ∃[ i ] L.length bucket[ i ] < k ^ n
    ∃small-bucket = VΣ-php {f = length} buckets (k ^ n)
      (subst (_< k * k ^ n) (sym Σlen-buckets≡len-list) len<k*k^n)
    i = proj₁ ∃small-bucket
    small-bucket-lacks : ∃[ w ] w ∉ bucket[ i ]
    small-bucket-lacks = find-missing bucket[ i ] (proj₂ ∃small-bucket)
    w : Vec (Fin k) n
    w = proj₁ small-bucket-lacks