-- The infinite Ramsey theorem: for every n, every infinite set S, and every
-- colouring of the n-subsets of S with finitely many colours, some infinite
-- H ⊆ S has all its n-subsets of one colour. The chain of stages supplies
-- points a₀, a₁, … with colours c₀, c₁, …; one colour occurs
-- infinitely often, and the points of that colour form H: an
-- (n + 1)-subset of H with least index j is insert(aⱼ, T) for an
-- n-subset T of Sⱼ₊₁.
-- 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.Combinatorics.RamseyInfinite
import RamseyChain

-- The colours and the points of the chain, as graphs on ℕ.

defgraph ramseyColourMap (C S K : Set, n : Nat) (i ∈ ℕ) : K :=
  ramseyColour(C, S, K, n, i)
requires ramsey : RamseyProperty(n, K)
requires colouring : C : subsetsOfSize(S, n + 1) → K
requires infinite : Infinite(S)
proof
  have iNat := nat.omega.backward with member
  exact ramseyColour.member at C, S, K, n, i with *

defgraph ramseyPointMap (C S K : Set, n : Nat) (i ∈ ℕ) : S :=
  ramseyPoint(C, S, K, n, i)
requires ramsey : RamseyProperty(n, K)
requires colouring : C : subsetsOfSize(S, n + 1) → K
requires infinite : Infinite(S)
proof
  have iNat := nat.omega.backward with member
  exact ramseyPoint.in_set at C, S, K, n, i with *

-- The stages whose colour is c, and the points of those stages.

def ramseyIndices (C S K c : Set, n : Nat) : Set :=
  functionPreimage(ramseyColourMap(C, S, K, n), ℕ, {c})

theorem ramseyIndices.value (C S K c : Set, n : Nat)
  shows ramseyIndices(C, S, K, c, n) =
    functionPreimage(ramseyColourMap(C, S, K, n), ℕ, {c})
proof
  simp [ramseyIndices]

def ramseyHomogeneous (C S K c : Set, n : Nat) : Set :=
  functionImage(ramseyPointMap(C, S, K, n), ramseyIndices(C, S, K, c, n), S)

theorem ramseyHomogeneous.value (C S K c : Set, n : Nat)
  shows ramseyHomogeneous(C, S, K, c, n) =
    functionImage(ramseyPointMap(C, S, K, n), ramseyIndices(C, S, K, c, n), S)
proof
  simp [ramseyHomogeneous]

-- The indices of colour c whose points lie in U.
def ramseyIndicesOf (C S K c U : Set, n : Nat) : Set :=
  { i ∈ ramseyIndices(C, S, K, c, n) | ramseyPoint(C, S, K, n, i) ∈ U }

section RamseyInfinite
variable (C S K : Set) (n : Nat)
hypothesis ramsey : RamseyProperty(n, K)
hypothesis colouring : FunctionBetween(C, subsetsOfSize(S, n + 1), K)
hypothesis infinite : Infinite(S)

theorem ramseyIndices.subset (c : Set)
  shows ramseyIndices(C, S, K, c, n) ⊆ ℕ
proof
  rw ramseyIndices.value
  exact functionPreimage.subset at ramseyColourMap(C, S, K, n), ℕ, {c}

theorem ramseyIndices.membership (c i : Set)
  shows i ∈ ramseyIndices(C, S, K, c, n) ↔
    i ∈ ℕ ∧ ramseyColour(C, S, K, n, i) = c
proof
  rw ramseyIndices.value
  constructor
  . intro h
    obtain ⟨hi, valueC⟩ := functionPreimage.point_membership.forward with h
    have value := ramseyColourMap.value at C, S, K, n, i with *
    rw value at valueC
    exact ⟨hi, valueC⟩
  . intro h
    obtain ⟨hi, valueC⟩ := h
    have value := ramseyColourMap.value at C, S, K, n, i with *
    apply functionPreimage.point_membership.backward
    rw value
    exact ⟨hi, valueC⟩

theorem ramseyHomogeneous.subset (c : Set)
  shows ramseyHomogeneous(C, S, K, c, n) ⊆ S
proof
  rw ramseyHomogeneous.value
  exact functionImage.subset at
    ramseyPointMap(C, S, K, n), ramseyIndices(C, S, K, c, n), S

theorem ramseyHomogeneous.membership (c x : Set)
  shows x ∈ ramseyHomogeneous(C, S, K, c, n) ↔
    ∃ i ∈ ramseyIndices(C, S, K, c, n), x = ramseyPoint(C, S, K, n, i)
proof
  have indices := ramseyIndices.subset at C, S, K, n, c with *
  rw ramseyHomogeneous.value
  constructor
  . intro h
    obtain ⟨_, i, hi, shape⟩ := functionImage.membership.forward with h
    have hiN := indices with i, hi
    have value := ramseyPointMap.value at C, S, K, n, i with *
    rw value at shape
    exact ⟨i, hi, shape⟩
  . intro h
    obtain ⟨i, hi, shape⟩ := h
    have hiN := indices with i, hi
    have iNat := nat.omega.backward with hiN
    have value := ramseyPointMap.value at C, S, K, n, i with *
    have inS := ramseyPoint.in_set at C, S, K, n, i with *
    apply functionImage.membership.backward
    constructor
    . rw shape
      exact inS
    . refine ⟨i, hi, ?_⟩
      rw [shape, value]
      rfl

-- The points are pairwise distinct, so H is as large as its index set.
theorem ramseyPointMap.injective
  shows InjectiveBetween(ramseyPointMap(C, S, K, n), ℕ, S)
proof
  have graph := ramseyPointMap.function at C, S, K, n with *
  apply Injective.of_pointwise with graph
  intros i hi j hj same
  have iNat := nat.omega.backward with hi
  have jNat := nat.omega.backward with hj
  have valueI := ramseyPointMap.value at C, S, K, n, i with *
  have valueJ := ramseyPointMap.value at C, S, K, n, j with *
  rw [valueI, valueJ] at same
  exact ramseyPoint.injective_on at C, S, K, n, i, j with *

theorem ramseyHomogeneous.infinite (c : Set)
  requires indicesInfinite : Infinite(ramseyIndices(C, S, K, c, n))
  shows Infinite(ramseyHomogeneous(C, S, K, c, n))
proof
  have injective := ramseyPointMap.injective at C, S, K, n with *
  have graph := ramseyPointMap.function at C, S, K, n with *
  have indices := ramseyIndices.subset at C, S, K, n, c with *
  have restricted := restrict.injective with injective, indices
  have imageInfinite := Infinite.injective_image with
    restricted, indicesInfinite
  have sameImage := functionImage.restrict with graph, indices
  rw sameImage at imageInfinite
  rw ramseyHomogeneous.value
  exact imageInfinite

-- An (n + 1)-subset U of H whose least index is j is insert(aⱼ, T) for the
-- n-subset T = U ∖ {aⱼ} of Sⱼ₊₁, so its colour is cⱼ = c.
theorem ramseyHomogeneous.colour_of_least (c U : Set) (j : Nat)
  requires hU : U ∈ subsetsOfSize(ramseyHomogeneous(C, S, K, c, n), n + 1)
  requires jIndex : j ∈ ramseyIndices(C, S, K, c, n)
  requires jIn : ramseyPoint(C, S, K, n, j) ∈ U
  requires least : ∀ i ∈ ramseyIndices(C, S, K, c, n),
    ramseyPoint(C, S, K, n, i) ∈ U → j ≤ i
  shows C(U) = c
proof
  have subsetU := subsetsOfSize.subset with hU
  have indices := ramseyIndices.subset at C, S, K, n, c with *
  have jParts := ramseyIndices.membership.forward with jIndex
  have restSubset : U ∖ {ramseyPoint(C, S, K, n, j)} ⊆ ramseyStage(C, S, K, n, j + 1)
    intros x hx
    obtain ⟨xIn, xNe⟩ := difference.singleton_membership.forward with hx
    have xH := subsetU with x, xIn
    obtain ⟨i, hi, shape⟩ := ramseyHomogeneous.membership.forward with xH
    have hiN := indices with i, hi
    have iNat := nat.omega.backward with hiN
    have pointIn : ramseyPoint(C, S, K, n, i) ∈ U
      rw ← shape
      exact xIn
    have jLe := least with i, hi, pointIn
    have jNe : j ≠ i
      intro same
      apply xNe
      cc
    have jLt := Less.of_LessEqual_ne with jLe, jNe
    have later := ramseyPoint.later_member at C, S, K, n, j, i with *
    rw shape
    exact later
  have counted := subsetsOfSize.counted with hU
  have one : n + 1 = succ(n) := add.right_one with *
  rw one at counted
  have restCounted := HasCardinality.remove with counted, jIn
  have restIn := subsetsOfSize.of_subset with restSubset, restCounted
  have law := ramseyColour.law at
    C, S, K, n, U ∖ {ramseyPoint(C, S, K, n, j)}, j with *
  have restored := insert.restore with jIn
  rw restored at law
  rw law
  exact jParts.right

theorem ramseyHomogeneous.homogeneous (c U : Set)
  requires hU : U ∈ subsetsOfSize(ramseyHomogeneous(C, S, K, c, n), n + 1)
  shows C(U) = c
proof
  have subsetU := subsetsOfSize.subset with hU
  have indices := ramseyIndices.subset at C, S, K, n, c with *
  obtain ⟨u, hu⟩ := subsetsOfSize.exists_member with hU
  have uH := subsetU with u, hu
  obtain ⟨i0, hi0, shape0⟩ := ramseyHomogeneous.membership.forward with uH
  have pointIn0 : ramseyPoint(C, S, K, n, i0) ∈ U
    rw ← shape0
    exact hu
  have candidate : i0 ∈ ramseyIndicesOf(C, S, K, c, U, n) :=
    ramseyIndicesOf.membership.backward with ⟨hi0, pointIn0⟩
  have inhabited : Nonempty(ramseyIndicesOf(C, S, K, c, U, n)) :=
    ⟨i0, candidate⟩
  have naturals : ramseyIndicesOf(C, S, K, c, U, n) ⊆ ℕ
    intros i hi
    have parts := ramseyIndicesOf.membership.forward with hi
    exact indices with i, parts.left
  have jMember := leastNatural.member with inhabited, naturals
  obtain ⟨jIndex, jIn⟩ := ramseyIndicesOf.membership.forward with jMember
  have jN := indices with
    leastNatural(ramseyIndicesOf(C, S, K, c, U, n)), jIndex
  have jNat := nat.omega.backward with jN
  have least : ∀ i ∈ ramseyIndices(C, S, K, c, n),
      ramseyPoint(C, S, K, n, i) ∈ U →
        leastNatural(ramseyIndicesOf(C, S, K, c, U, n)) ≤ i
    intros i hi pointIn
    have iNat := nat.omega.backward with (indices with i, hi)
    have member : i ∈ ramseyIndicesOf(C, S, K, c, U, n) :=
      ramseyIndicesOf.membership.backward with ⟨hi, pointIn⟩
    exact leastNatural.least with inhabited, naturals, member
  exact ramseyHomogeneous.colour_of_least at
    C, S, K, n, c, U, leastNatural(ramseyIndicesOf(C, S, K, c, U, n)) with *

end RamseyInfinite

-- The induction on n.

theorem RamseyProperty.zero (K : Set)
  shows RamseyProperty(0, K)
proof
  intros S C infinite colouring
  have emptyIn : ∅ ∈ subsetsOfSize(S, 0)
    rw subsetsOfSize.zero
    exact singleton.self at ∅
  have value := FunctionBetween.apply_closed with colouring, emptyIn
  refine ⟨S, Subset.reflexive at S, infinite, C(∅), value, ?_⟩
  intros U hU
  rw subsetsOfSize.zero at hU
  have equal := singleton.membership.forward with hU
  rw equal
  rfl

theorem RamseyProperty.succ (K : Set) (n : Nat)
  requires finite : Finite(K)
  requires ramsey : RamseyProperty(n, K)
  shows RamseyProperty(n + 1, K)
proof
  intros S C infinite colouring
  have colourGraph := ramseyColourMap.function at C, S, K, n with *
  have omegaInfinite : Infinite(ℕ) := omega.infinite
  obtain ⟨c, hc, indicesInfinite⟩ :=
    pigeonhole.infinite_fiber with colourGraph, omegaInfinite, finite
  rw ← ramseyIndices.value at indicesInfinite
  have subset := ramseyHomogeneous.subset at C, S, K, n, c with *
  have hInfinite := ramseyHomogeneous.infinite at C, S, K, n, c with *
  refine ⟨ramseyHomogeneous(C, S, K, c, n), subset, hInfinite, c, hc, ?_⟩
  intros U hU
  exact ramseyHomogeneous.homogeneous at C, S, K, n, c, U with *

theorem RamseyProperty.all (K : Set) (n : Nat)
  requires finite : Finite(K)
  shows RamseyProperty(n, K)
proof
  induction n
  | zero => exact RamseyProperty.zero at K
  | k + 1 using ih => exact RamseyProperty.succ at K, k with finite, ih

-- The infinite Ramsey theorem, for any finite set of colours.
theorem ramsey.infinite_colours (n : Nat) (S C K : Set)
  requires infinite : Infinite(S)
  requires finite : Finite(K)
  requires colouring : C : subsetsOfSize(S, n) → K
  shows ∃ H, H ⊆ S ∧ Infinite(H) ∧
    ∃ c ∈ K, ∀ U ∈ subsetsOfSize(H, n), C(U) = c
proof
  have ramsey := RamseyProperty.all at K, n with finite
  exact ramsey with S, C, infinite, colouring

-- The infinite Ramsey theorem for two colours.
theorem ramsey.infinite (n : Nat) (S C : Set)
  requires infinite : Infinite(S)
  requires colouring : C : subsetsOfSize(S, n) → 2
  shows ∃ H, H ⊆ S ∧ Infinite(H) ∧
    ∃ c ∈ 2, ∀ U ∈ subsetsOfSize(H, n), C(U) = c
proof
  have finite := Finite.natural at 2
  exact ramsey.infinite_colours with infinite, finite, colouring

-- The infinite Ramsey theorem for k colours.
theorem ramsey.infinite_numeral (n k : Nat) (S C : Set)
  requires infinite : Infinite(S)
  requires colouring : C : subsetsOfSize(S, n) → k
  shows ∃ H, H ⊆ S ∧ Infinite(H) ∧
    ∃ c ∈ k, ∀ U ∈ subsetsOfSize(H, n), C(U) = c
proof
  have finite := Finite.natural at k
  exact ramsey.infinite_colours with infinite, finite, colouring