-- The exchange step: an edge e of least finite level in an even set of the
-- last colour is traded for an edge of lower level on the even set it closes
-- in the colour refined at that level, and the first k colours stay trees.
-- 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.GraphTheory.TreePacking.Exchange
import Order

-- The exchange: e of least finite level m in an even set of T_k, P its class
-- at level m, c_m the colour refined there, e' an edge of least level on the
-- even set closed by e in T_{c_m}. Exchanging e and e' between the two
-- colours keeps the first k colours trees.
def exchange (E T e f : Set, c k : Nat) : Set :=
  { p ∈ succ(k) × 𝒫(E) | ∃ i ∈ succ(k),
    (i = c ∧ p = (i, (T(c) ∪ {e}) ∖ {f})) ∨
    (i = k ∧ p = (i, (T(k) ∪ {f}) ∖ {e})) ∨
    (i ≠ c ∧ i ≠ k ∧ p = (i, T(i))) }

-- An edge has one colour.
theorem Decomposition.colour_unique (E T g : Set, k i j : Nat)
  requires decomposition : Decomposition(E, T, k)
  requires indexI : i ∈ k
  requires indexJ : j ∈ k
  requires inI : g ∈ T(i)
  requires inJ : g ∈ T(j)
  shows i = j
proof
  by_contra different
  exact decomposition.right.right with
    i, indexI, j, indexJ, different, g, inI, inJ

-- The facts the exchange rests on: c and k are distinct colours, e and f
-- are distinct edges of E, and the two swapped classes lie in E.
theorem exchange.setting (E T e f : Set, c k : Nat)
  requires decomposition : Decomposition(E, T, succ(k))
  requires colour : c ∈ k
  requires first : e ∈ T(k)
  requires second : f ∈ T(c)
  shows c ≠ k ∧ e ≠ f ∧
    (T(c) ∪ {e}) ∖ {f} ∈ 𝒫(E) ∧ (T(k) ∪ {f}) ∖ {e} ∈ 𝒫(E)
proof
  have less := Less.of_membership at c, k with colour
  have distinct : c ≠ k by arith
  have cIn : c ∈ succ(k) by simp [colour]
  have kIn : k ∈ succ(k) by simp
  have classC := Decomposition.class_subset with decomposition, cIn
  have classK := Decomposition.class_subset with decomposition, kIn
  have eInE := classK with e, first
  have fInE := classC with f, second
  have unionC : T(c) ∪ {e} ⊆ E :=
    Subset.union2_iff.backward with
      ⟨classC, (singleton.subset_iff.backward with eInE)⟩
  have unionK : T(k) ∪ {f} ⊆ E :=
    Subset.union2_iff.backward with
      ⟨classK, (singleton.subset_iff.backward with fInE)⟩
  refine ⟨distinct, ?_, ?_, ?_⟩
  . intro same
    rw same at first
    exact decomposition.right.right with
      c, cIn, k, kIn, distinct, f, second, first
  . exact powerset.membership_iff.backward with
      (Subset.transitive with difference.subset_left, unionC)
  . exact powerset.membership_iff.backward with
      (Subset.transitive with difference.subset_left, unionK)

-- An entry of the exchange is one of the three kinds of class.
theorem exchange.shape (E T e f x y : Set, c k : Nat)
  requires member : (x, y) ∈ exchange(E, T, e, f, c, k)
  shows (x = c ∧ y = (T(c) ∪ {e}) ∖ {f}) ∨
    (x = k ∧ y = (T(k) ∪ {f}) ∖ {e}) ∨
    (x ≠ c ∧ x ≠ k ∧ y = T(x))
proof
  obtain ⟨_, i, _, cases⟩ := member
  rcases cases with
    ⟨same, shape⟩ | ⟨same, shape⟩ | ⟨notC, notK, shape⟩
  . obtain ⟨xi, yA⟩ := orderedPair.injective with shape
    left
    rw xi
    exact ⟨same, yA⟩
  . obtain ⟨xi, yB⟩ := orderedPair.injective with shape
    right
    left
    rw xi
    exact ⟨same, yB⟩
  . obtain ⟨xi, yC⟩ := orderedPair.injective with shape
    right
    right
    rw xi
    exact ⟨notC, notK, yC⟩

theorem exchange.function (E T e f : Set, c k : Nat)
  requires decomposition : Decomposition(E, T, succ(k))
  requires colour : c ∈ k
  requires first : e ∈ T(k)
  requires second : f ∈ T(c)
  shows exchange(E, T, e, f, c, k) : succ(k) → 𝒫(E)
proof
  obtain ⟨distinct, _, swappedC, swappedK⟩ :=
    exchange.setting with decomposition, colour, first, second
  constructor
  . intros p member
    exact member.left
  constructor
  . intros x y relatedY z relatedZ
    have shapeY := exchange.shape with relatedY
    have shapeZ := exchange.shape with relatedZ
    -- In each of the nine cases the two values agree or the indices clash.
    rcases shapeY with ⟨xc, yA⟩ | ⟨xk, yB⟩ | ⟨xnc, xnk, yC⟩
    . rcases shapeZ with ⟨zc, zA⟩ | ⟨zk, zB⟩ | ⟨znc, znk, zC⟩; . cc; . cc; . cc
    . rcases shapeZ with ⟨zc, zA⟩ | ⟨zk, zB⟩ | ⟨znc, znk, zC⟩; . cc; . cc; . cc
    . rcases shapeZ with ⟨zc, zA⟩ | ⟨zk, zB⟩ | ⟨znc, znk, zC⟩; . cc; . cc; . cc
  . intros x input
    by_cases xc : x = c
    . refine ⟨(T(c) ∪ {e}) ∖ {f}, swappedC, ?_⟩
      unfold relates
      refine ⟨(product.pair with input, swappedC), x, input, ?_⟩
      left
      exact ⟨xc, rfl⟩
    by_cases xk : x = k
    . refine ⟨(T(k) ∪ {f}) ∖ {e}, swappedK, ?_⟩
      unfold relates
      refine ⟨(product.pair with input, swappedK), x, input, ?_⟩
      right
      left
      exact ⟨xk, rfl⟩
    have class := FunctionBetween.apply_closed with decomposition.left, input
    refine ⟨T(x), class, ?_⟩
    unfold relates
    refine ⟨(product.pair with input, class), x, input, ?_⟩
    right
    right
    exact ⟨xc, xk, rfl⟩

theorem exchange.value_colour (E T e f : Set, c k : Nat)
  requires decomposition : Decomposition(E, T, succ(k))
  requires colour : c ∈ k
  requires first : e ∈ T(k)
  requires second : f ∈ T(c)
  shows exchange(E, T, e, f, c, k)(c) = (T(c) ∪ {e}) ∖ {f}
proof
  obtain ⟨_, _, swappedC, _⟩ :=
    exchange.setting with decomposition, colour, first, second
  have cIn : c ∈ succ(k) by simp [colour]
  have function := exchange.function with decomposition, colour, first, second
  have related : relates(exchange(E, T, e, f, c, k), c, (T(c) ∪ {e}) ∖ {f})
    unfold relates
    refine ⟨(product.pair with cIn, swappedC), c, cIn, ?_⟩
    left
    exact ⟨rfl, rfl⟩
  exact FunctionBetween.apply_value with function, cIn, related

theorem exchange.value_last (E T e f : Set, c k : Nat)
  requires decomposition : Decomposition(E, T, succ(k))
  requires colour : c ∈ k
  requires first : e ∈ T(k)
  requires second : f ∈ T(c)
  shows exchange(E, T, e, f, c, k)(k) = (T(k) ∪ {f}) ∖ {e}
proof
  obtain ⟨_, _, _, swappedK⟩ :=
    exchange.setting with decomposition, colour, first, second
  have kIn : k ∈ succ(k) by simp
  have function := exchange.function with decomposition, colour, first, second
  have related : relates(exchange(E, T, e, f, c, k), k, (T(k) ∪ {f}) ∖ {e})
    unfold relates
    refine ⟨(product.pair with kIn, swappedK), k, kIn, ?_⟩
    right
    left
    exact ⟨rfl, rfl⟩
  exact FunctionBetween.apply_value with function, kIn, related

theorem exchange.value_other (E T e f : Set, c k i : Nat)
  requires decomposition : Decomposition(E, T, succ(k))
  requires colour : c ∈ k
  requires first : e ∈ T(k)
  requires second : f ∈ T(c)
  requires index : i ∈ succ(k)
  requires notC : i ≠ c
  requires notK : i ≠ k
  shows exchange(E, T, e, f, c, k)(i) = T(i)
proof
  have function := exchange.function with decomposition, colour, first, second
  have class := FunctionBetween.apply_closed with decomposition.left, index
  have related : relates(exchange(E, T, e, f, c, k), i, T(i))
    unfold relates
    refine ⟨(product.pair with index, class), i, index, ?_⟩
    right
    right
    exact ⟨notC, notK, rfl⟩
  exact FunctionBetween.apply_value with function, index, related

-- An edge lies in the exchanged class of a colour when it is e and the
-- colour is c, f and the colour is k, or neither edge and of that colour.
theorem exchange.class_membership (E T e f g : Set, c k i : Nat)
  requires decomposition : Decomposition(E, T, succ(k))
  requires colour : c ∈ k
  requires first : e ∈ T(k)
  requires second : f ∈ T(c)
  requires index : i ∈ succ(k)
  shows g ∈ exchange(E, T, e, f, c, k)(i) ↔
    (g = e ∧ i = c) ∨ (g = f ∧ i = k) ∨
    (g ≠ e ∧ g ≠ f ∧ g ∈ T(i))
proof
  obtain ⟨distinct, edges, _, _⟩ :=
    exchange.setting with decomposition, colour, first, second
  have cIn : c ∈ succ(k) by simp [colour]
  have kIn : k ∈ succ(k) by simp
  have eNotC : ¬e ∈ T(c)
    intro inC
    exact distinct with
      (Decomposition.colour_unique with decomposition, cIn, kIn, inC, first)
  have fNotK : ¬f ∈ T(k)
    intro inK
    exact distinct with
      (Decomposition.colour_unique with decomposition, cIn, kIn, second, inK)
  by_cases ic : i = c
  . rw [ic, (exchange.value_colour with decomposition, colour, first, second)]
    constructor
    . intro member
      simp [union2.membership, singleton.membership] at member
      rcases member with ⟨inC | same, notF⟩
      . right
        right
        refine ⟨?_, notF, inC⟩
        intro same
        rw same at inC
        exact eNotC with inC
      . left
        exact ⟨same, rfl⟩
    . intro cases
      rcases cases with ⟨same, _⟩ | ⟨_, ck⟩ | ⟨_, notF, inC⟩
      . rw same
        simp [edges]
      . exfalso
        exact distinct with ck
      . simp [inC, notF]
  by_cases ik : i = k
  . rw [ik, (exchange.value_last with decomposition, colour, first, second)]
    constructor
    . intro member
      simp [union2.membership, singleton.membership] at member
      rcases member with ⟨inK | same, notE⟩
      . right
        right
        refine ⟨notE, ?_, inK⟩
        intro same
        rw same at inK
        exact fNotK with inK
      . right
        left
        exact ⟨same, rfl⟩
    . intro cases
      rcases cases with ⟨_, kc⟩ | ⟨same, _⟩ | ⟨notE, _, inK⟩
      . exfalso
        cc
      . rw same
        have different : f ≠ e by cc
        simp [different]
      . simp [inK, notE]
  rw (exchange.value_other with
    decomposition, colour, first, second, index, ic, ik)
  constructor
  . intro inT
    right
    right
    refine ⟨?_, ?_, inT⟩
    . intro same
      rw same at inT
      exact ik with (Decomposition.colour_unique with
        decomposition, index, kIn, inT, first)
    . intro same
      rw same at inT
      exact ic with (Decomposition.colour_unique with
        decomposition, index, cIn, inT, second)
  . intro cases
    rcases cases with ⟨_, same⟩ | ⟨_, same⟩ | ⟨_, _, inT⟩
    . exfalso
      exact ic with same
    . exfalso
      exact ik with same
    . exact inT

theorem exchange.decomposition (V E leftEnd rightEnd T e f : Set, c k : Nat)
  requires decomposition : Decomposition(E, T, succ(k))
  requires colour : c ∈ k
  requires first : e ∈ T(k)
  requires second : f ∈ T(c)
  shows Decomposition(E, exchange(E, T, e, f, c, k), succ(k))
proof
  obtain ⟨_, edges, _, _⟩ :=
    exchange.setting with decomposition, colour, first, second
  have cIn : c ∈ succ(k) by simp [colour]
  have kIn : k ∈ succ(k) by simp
  refine ⟨(exchange.function with decomposition, colour, first, second),
    ?_, ?_⟩
  -- Every edge keeps a colour: e moves to c, f to k, the rest stay.
  . intros g inE
    by_cases ge : g = e
    . refine ⟨c, cIn, ?_⟩
      rw exchange.class_membership with
        decomposition, colour, first, second, cIn
      left
      exact ⟨ge, rfl⟩
    by_cases gf : g = f
    . refine ⟨k, kIn, ?_⟩
      rw exchange.class_membership with
        decomposition, colour, first, second, kIn
      right
      left
      exact ⟨gf, rfl⟩
    obtain ⟨i, index, inT⟩ := decomposition.right.left with g, inE
    have natural := natural.member at succ(k) with i, index
    refine ⟨i, index, ?_⟩
    rw exchange.class_membership with
      decomposition, colour, first, second, index
    right
    right
    exact ⟨ge, gf, inT⟩
  -- And only one colour, as under T.
  . intros i indexI j indexJ different
    have naturalI := natural.member at succ(k) with i, indexI
    have naturalJ := natural.member at succ(k) with j, indexJ
    apply Disjoint.iff.backward
    intros g inI inJ
    rw exchange.class_membership with
      decomposition, colour, first, second, indexI at inI
    rw exchange.class_membership with
      decomposition, colour, first, second, indexJ at inJ
    rcases inI with ⟨ge, ic⟩ | ⟨gf, ik⟩ | ⟨notE, notF, inTi⟩
    . rcases inJ with ⟨_, jc⟩ | ⟨gf, _⟩ | ⟨notE, _, _⟩
      . cc
      . cc
      . exact notE with ge
    . rcases inJ with ⟨ge, _⟩ | ⟨_, jk⟩ | ⟨_, notF, _⟩
      . cc
      . cc
      . exact notF with gf
    . rcases inJ with ⟨ge, _⟩ | ⟨gf, _⟩ | ⟨_, _, inTj⟩
      . exact notE with ge
      . exact notF with gf
      . exact different with
          (Decomposition.colour_unique with
            decomposition, indexI, indexJ, inTi, inTj)