-- Countability and Cantor's diagonal theorem
-- 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.Cardinal.Countable
import Finite

-- A set is countable when it injects into the natural numbers. This includes
-- finite sets and gives the usual "at most countable" convention.
predicate Countable (A : Set) := ∃ f, InjectiveBetween(f, A, ℕ)
predicate Uncountable (A : Set) := ¬Countable(A)

-- The diagonal subset associated with a proposed enumeration of P(A).
def cantorDiagonal (f A : Set) : Set :=
  { x ∈ A | ¬(x ∈ f(x)) }

theorem cantorDiagonal.subset (f A : Set)
  shows cantorDiagonal(f, A) ⊆ A
proof
  simp

theorem cantorDiagonal.in_powerset (f A : Set)
  shows cantorDiagonal(f, A) ∈ 𝒫(A)
proof
  simp [Subset, set.powerset, cantorDiagonal.subset]

-- Cantor's theorem in its graph-theoretic form: no set-valued function from A
-- onto its powerset exists. The proof is the literal diagonal argument.
theorem Cantor.no_surjection_onto_powerset (f A : Set)
  shows ¬SurjectiveBetween(f, A, 𝒫(A))
proof
  intro surjective
  rcases surjective with ⟨function, onto⟩
  have diagonalInPower := cantorDiagonal.in_powerset at f, A
  specialize onto with cantorDiagonal(f, A), diagonalInPower
  obtain ⟨d, dInA, mapsToDiagonal⟩ := onto
  have value := FunctionBetween.apply_value with function, dInA, mapsToDiagonal
  have diagonalShape := cantorDiagonal.membership at f, A, d
  have membershipCases : d ∈ cantorDiagonal(f, A) ∨ ¬(d ∈ cantorDiagonal(f, A)) by tauto
  cases membershipCases with diagonal | notDiagonal
  . have shape := diagonalShape.forward with diagonal
    rcases shape with ⟨dInAAgain, notInValue⟩
    rw value at notInValue
    exact notInValue with diagonal
  . apply notDiagonal
    apply diagonalShape.backward
    constructor
    . exact dInA
    . rw value
      exact notDiagonal

theorem powerset.not_equipotent (A : Set)
  shows ¬(A ≃ 𝒫(A))
proof
  intro equivalent
  obtain ⟨f, bijective⟩ := equivalent
  have surjective := Bijective.surjective with bijective
  exact Cantor.no_surjection_onto_powerset with surjective

-- The injection form of Cantor's theorem for the naturals. Given a proposed
-- code f : P(N) → N, diagonalize at the code assigned to the diagonal set.
def countableDiagonal (f : Set) : Set :=
  { n ∈ ℕ | ∃ X ∈ 𝒫(ℕ), relates(f, X, n) ∧ ¬(n ∈ X) }

theorem countableDiagonal.subset (f : Set)
  shows countableDiagonal(f) ⊆ ℕ
proof
  simp

theorem powerset.naturals_not_injects_into_naturals (f : Set)
  shows ¬InjectiveBetween(f, 𝒫(ℕ), ℕ)
proof
  intro injection
  rcases injection with ⟨function, oneToOne⟩
  have diagonalSubset := countableDiagonal.subset at f
  have diagonalInPower := set.powerset.backward at countableDiagonal(f), ℕ with
    diagonalSubset
  obtain ⟨d, dInN, diagonalMaps⟩ := FunctionBetween.exists_value with function, diagonalInPower
  have cases : d ∈ countableDiagonal(f) ∨ ¬(d ∈ countableDiagonal(f)) by tauto
  cases cases with member | missing
  . have data := countableDiagonal.membership.forward at
      f, d with member
    rcases data with ⟨valueNatural, X, xInPower, xMaps, valueMissingX⟩
    have sameSet := oneToOne with
      X, xInPower, countableDiagonal(f), diagonalInPower,
      d, dInN,
      xMaps, diagonalMaps
    apply valueMissingX
    rwa sameSet
  . apply missing
    apply countableDiagonal.membership.backward
    constructor
    . exact dInN
    . witness countableDiagonal(f)
      constructor
      . exact diagonalInPower
      . constructor
        . exact diagonalMaps
        . exact missing

theorem powerset.naturals_uncountable
  shows Uncountable(𝒫(ℕ))
proof
  unfold Uncountable
  intro countable
  obtain ⟨f, injection⟩ := countable
  exact powerset.naturals_not_injects_into_naturals with injection