-- Tagged disjoint unions and cardinal addition
-- 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.Sum
import Finite

-- Tags are Kuratowski pairs. Distinct von Neumann tags make the two
-- images disjoint even when the original carriers overlap.
def inl (x : Set) : Set := (x, 0)
def inr (y : Set) : Set := (y, 1)
def disjointSum (A B : Set) : Set := (A × {0}) ∪ (B × {1})
notation infixr 55 "⊎" => disjointSum
def sumMap (f g A B C D : Set) : Set := { p ∈ (A ⊎ B) × (C ⊎ D) | (∃ x ∈ A, ∃ y ∈ C, relates(f, x, y) ∧ p = (inl(x) ↦ inl(y))) ∨ (∃ x ∈ B, ∃ y ∈ D, relates(g, x, y) ∧ p = (inr(x) ↦ inr(y))) }
def rightProjection (n : Set) : Set := { p ∈ (0 ⊎ n) × n | ∃ y ∈ n, p = (inr(y) ↦ y) }

def natSumEquiv (m n : Nat) → Set
  | zero, n => rightProjection(n)
  | succ(m), n => adjoinMap(natSumEquiv(m, n), inl(m), add(m, n))

theorem disjointSum.left (x A B : Set)
  requires hx : x ∈ A
  shows inl(x) ∈ A ⊎ B
proof
  unfold inl
  unfold disjointSum
  apply union2.left
  apply product.pair
  exact hx
  exact singleton.self

theorem disjointSum.right (y A B : Set)
  requires hy : y ∈ B
  shows inr(y) ∈ A ⊎ B
proof
  unfold inr
  unfold disjointSum
  apply union2.right
  apply product.pair
  exact hy
  exact singleton.self

theorem disjointSum.membership (p A B : Set)
  shows p ∈ A ⊎ B ↔ ((∃ x ∈ A, p = inl(x)) ∨ (∃ y ∈ B, p = inr(y)))
proof
  constructor
  intro hp
  simp [disjointSum, union2.membership, product.membership, singleton.membership] at hp
  cases hp with hleft | hright
  simp [product.membership, singleton.membership] at hleft
  rcases hleft with (hbound, witnesses)
  obtain (x, hx) := witnesses
  rcases hx with (hxA, rest)
  obtain ⟨tag, tagZero, pairShape⟩ := rest
  simp [singleton.membership] at tagZero
  left
  witness x
  constructor
  exact hxA
  unfold inl
  cc
  simp [product.membership, singleton.membership] at hright
  rcases hright with ⟨hbound, y, hyB, tagRight, tagOne, pairShapeRight⟩
  simp [singleton.membership] at tagOne
  right
  witness y
  constructor
  exact hyB
  unfold inr
  cc
  intro hp
  cases hp with hleft | hright
  obtain ⟨xBack, hxBackA, pairShapeBack⟩ := hleft
  rw pairShapeBack
  exact disjointSum.left with hxBackA
  obtain ⟨yBack, hyBackB, pairShapeRightBack⟩ := hright
  rw pairShapeRightBack
  exact disjointSum.right with hyBackB

theorem zero.ne_one
  shows 0 ≠ 1
proof
  intro equal
  have member := next.contains at 0
  have successorZero := next.natural at 0
  rw [← successorZero, ← equal, nat.zero] at member
  simp

theorem tags.disjoint (x y : Set)
  shows inl(x) ≠ inr(y)
proof
  intro equal
  unfold inl at equal
  unfold inr at equal
  obtain (first, tagsEqual) := orderedPair.injective with equal
  exact zero.ne_one with tagsEqual

theorem inl.injective (x y : Set)
  requires h : inl(x) = inl(y)
  shows x = y
proof
  simp [inl] at h
  obtain (first, second) := orderedPair.injective with h
  exact first

theorem inr.injective (x y : Set)
  requires h : inr(x) = inr(y)
  shows x = y
proof
  simp [inr] at h
  obtain (first, second) := orderedPair.injective with h
  exact first

theorem sumMap.relates_left (f g A B C D x y : Set)
  requires hx : x ∈ A
  requires hy : y ∈ C
  requires hxy : relates(f, x, y)
  shows relates(sumMap(f, g, A, B, C, D), inl(x), inl(y))
proof
  unfold relates
  apply sumMap.membership.backward
  constructor
  apply product.pair
  exact disjointSum.left with hx
  exact disjointSum.left with hy
  left
  refine ⟨x, hx, y, hy, hxy, rfl⟩

theorem sumMap.relates_right (f g A B C D x y : Set)
  requires hx : x ∈ B
  requires hy : y ∈ D
  requires hxy : relates(g, x, y)
  shows relates(sumMap(f, g, A, B, C, D), inr(x), inr(y))
proof
  unfold relates
  apply sumMap.membership.backward
  constructor
  apply product.pair
  exact disjointSum.right with hx
  exact disjointSum.right with hy
  right
  refine ⟨x, hx, y, hy, hxy, rfl⟩

theorem sumMap.factors (f g A B C D u v : Set)
  requires h : relates(sumMap(f, g, A, B, C, D), u, v)
  shows (∃ x ∈ A, ∃ y ∈ C, u = inl(x) ∧ v = inl(y) ∧ relates(f, x, y)) ∨ (∃ x ∈ B, ∃ y ∈ D, u = inr(x) ∧ v = inr(y) ∧ relates(g, x, y))
proof
  unfold relates at h
  simp [sumMap.membership] at h
  rcases h with (hProduct, hShape)
  cases hShape with hleft | hright
  obtain ⟨x, hxA, y, hyC, hxy, hPair⟩ := hleft
  obtain (source, target) := orderedPair.injective with hPair
  left
  refine ⟨x, hxA, y, hyC, source, target, hxy⟩
  obtain ⟨xRight, hxB, yRight, hyD, hgxy, hPairRight⟩ := hright
  obtain (sourceRight, targetRight) := orderedPair.injective with hPairRight
  right
  refine ⟨xRight, hxB, yRight, hyD, sourceRight, targetRight, hgxy⟩

theorem sumMap.relation (f g A B C D : Set)
  shows RelationBetween(sumMap(f, g, A, B, C, D), A ⊎ B, C ⊎ D)
proof
  simp [RelationBetween, Subset]

theorem sumMap.total (f g A B C D : Set)
  requires hf : f : A → C
  requires hg : g : B → D
  shows TotalBetween(sumMap(f, g, A, B, C, D), A ⊎ B, C ⊎ D)
proof
  intros u hu
  have sourceShape := disjointSum.membership.forward with hu
  cases sourceShape with leftShape | rightShape
  obtain ⟨x, hxA, uShape⟩ := leftShape
  obtain ⟨y, hyC, hxy⟩ := FunctionBetween.exists_value with hf, hxA
  witness inl(y)
  constructor
  exact disjointSum.left with hyC
  rw uShape
  exact sumMap.relates_left with hxA, hyC, hxy
  obtain ⟨xRight, hxB, uShapeRight⟩ := rightShape
  obtain ⟨yRight, hyD, hgxy⟩ := FunctionBetween.exists_value with hg, hxB
  witness inr(yRight)
  constructor
  exact disjointSum.right with hyD
  rw uShapeRight
  exact sumMap.relates_right with hxB, hyD, hgxy

theorem sumMap.functional (f g A B C D : Set)
  requires hf : f : A → C
  requires hg : g : B → D
  shows Functional(sumMap(f, g, A, B, C, D))
proof
  intros u v huv w huw
  have firstFactor := sumMap.factors with huv
  have secondFactor := sumMap.factors with huw
  cases firstFactor with firstLeft | firstRight
  cases secondFactor with secondLeft | secondRight
  obtain ⟨x1, hx1A, y1, hy1C, uLeft1, vLeft, hfxy1⟩ := firstLeft
  obtain ⟨x2, hx2A, y2, hy2C, uLeft2, wLeft, hfxy2⟩ := secondLeft
  rw uLeft1 at uLeft2
  have sameInput := inl.injective with uLeft2
  rw ← sameInput at hfxy2
  have unique := FunctionBetween.functional with hf, x1, y1, hfxy1, y2, hfxy2
  rw [vLeft, wLeft, unique]
  rfl
  obtain ⟨xLeftCross, hxLeftCrossA, yLeftCross, hyLeftCrossC, uLeftCross, hyLeftCrossLast⟩ := firstLeft
  obtain ⟨xr, hxrB, yr, hyrD, uRight2, hyrLast⟩ := secondRight
  rw uLeftCross at uRight2
  have impossible := tags.disjoint with uRight2
  contradiction
  cases secondFactor with secondLeftAfterRight | secondRightAfterRight
  obtain ⟨xl, hxlA, yl, hylC, uLeftAfterRight, hylLast⟩ := secondLeftAfterRight
  obtain ⟨xr1, hxr1B, yr1, hyr1D, uRight1, hyr1Last⟩ := firstRight
  rw uLeftAfterRight at uRight1
  have impossibleReverse := tags.disjoint with uRight1
  contradiction
  obtain ⟨xr2, hxr2B, yr2, hyr2D, uRight2b, wRight, hgxy2⟩ := secondRightAfterRight
  obtain ⟨xrFirst, hxrFirstB, yrFirst, hyrFirstD, uRightFirst, vRight, hgxyFirst⟩ := firstRight
  rw uRightFirst at uRight2b
  have sameInputRight := inr.injective with uRight2b
  rw ← sameInputRight at hgxy2
  have uniqueRight := FunctionBetween.functional with hg, xrFirst, yrFirst, hgxyFirst, yr2, hgxy2
  rw [vRight, wRight, uniqueRight]
  rfl

theorem sumMap.function (f g A B C D : Set)
  requires hf : f : A → C
  requires hg : g : B → D
  shows sumMap(f, g, A, B, C, D) : A ⊎ B → C ⊎ D
proof
  constructor
  exact sumMap.relation
  constructor
  exact sumMap.functional with hf, hg
  exact sumMap.total with hf, hg

theorem sumMap.surjective (f g A B C D : Set)
  requires hf : SurjectiveBetween(f, A, C)
  requires hg : SurjectiveBetween(g, B, D)
  shows SurjectiveBetween(sumMap(f, g, A, B, C, D), A ⊎ B, C ⊎ D)
proof
  constructor
  have fFunction := Surjective.function with hf
  have gFunction := Surjective.function with hg
  exact sumMap.function with fFunction, gFunction
  intros v hv
  have targetShape := disjointSum.membership.forward with hv
  cases targetShape with targetLeft | targetRight
  obtain ⟨y, hyC, vShape⟩ := targetLeft
  rcases hf with (fFunctionAgain, fOnto)
  specialize fOnto with y, hyC
  obtain ⟨x, hxA, hxy⟩ := fOnto
  witness inl(x)
  constructor
  exact disjointSum.left with hxA
  rw vShape
  exact sumMap.relates_left with hxA, hyC, hxy
  obtain ⟨yRight, hyD, vShapeRight⟩ := targetRight
  rcases hg with (gFunctionAgain, gOnto)
  specialize gOnto with yRight, hyD
  obtain ⟨xRight, hxB, hgxy⟩ := gOnto
  witness inr(xRight)
  constructor
  exact disjointSum.right with hxB
  rw vShapeRight
  exact sumMap.relates_right with hxB, hyD, hgxy

theorem sumMap.injective (f g A B C D : Set)
  requires hf : InjectiveBetween(f, A, C)
  requires hg : InjectiveBetween(g, B, D)
  shows InjectiveBetween(sumMap(f, g, A, B, C, D), A ⊎ B, C ⊎ D)
proof
  constructor
  have fFunction := Injective.function with hf
  have gFunction := Injective.function with hg
  exact sumMap.function with fFunction, gFunction
  intros u hu v hv z hz huz hvz
  have firstFactor := sumMap.factors with huz
  have secondFactor := sumMap.factors with hvz
  cases firstFactor with firstLeft | firstRight
  cases secondFactor with secondLeft | secondRight
  obtain ⟨x1, hx1A, y1, hy1C, uLeft, zLeft1, hfxy1⟩ := firstLeft
  obtain ⟨x2, hx2A, y2, hy2C, vLeft, zLeft2, hfxy2⟩ := secondLeft
  rw zLeft1 at zLeft2
  have sameOutput := inl.injective with zLeft2
  rw ← sameOutput at hfxy2
  have sameInput := Injective.cancel with hf, hx1A, hx2A, hy1C, hfxy1, hfxy2
  rw [uLeft, vLeft, sameInput]
  rfl
  obtain ⟨xLeftCross, hxLeftCrossA, yLeftCross, hyLeftCrossC, uLeftCross, zLeftCross, hfLeftCross⟩ := firstLeft
  obtain ⟨xRightCross, hxRightCrossB, yRightCross, hyRightCrossD, vRightCross, zRightCross, hgRightCross⟩ := secondRight
  rw zLeftCross at zRightCross
  have impossible := tags.disjoint with zRightCross
  contradiction
  cases secondFactor with secondLeftAfterRight | secondRightAfterRight
  obtain ⟨xLeftAfterRight, hxLeftAfterRightA, yLeftAfterRight, hyLeftAfterRightC, vLeftAfterRight, zLeftAfterRight, hfLeftAfterRight⟩ := secondLeftAfterRight
  obtain ⟨xRightFirst, hxRightFirstB, yRightFirst, hyRightFirstD, uRightFirst, zRightFirst, hgRightFirst⟩ := firstRight
  rw zLeftAfterRight at zRightFirst
  have impossibleReverse := tags.disjoint with zRightFirst
  contradiction
  obtain ⟨xRight2, hxRight2B, yRight2, hyRight2D, vRight2, zRight2, hgxy2⟩ := secondRightAfterRight
  obtain ⟨xRight1, hxRight1B, yRight1, hyRight1D, uRight1, zRight1, hgxy1⟩ := firstRight
  rw zRight1 at zRight2
  have sameOutputRight := inr.injective with zRight2
  rw ← sameOutputRight at hgxy2
  have sameInputRight := Injective.cancel with hg, hxRight1B, hxRight2B, hyRight1D, hgxy1, hgxy2
  rw [uRight1, vRight2, sameInputRight]
  rfl

theorem sumMap.bijective (f g A B C D : Set)
  requires hf : BijectiveBetween(f, A, C)
  requires hg : BijectiveBetween(g, B, D)
  shows BijectiveBetween(sumMap(f, g, A, B, C, D), A ⊎ B, C ⊎ D)
proof
  constructor
  have fInjective := Bijective.injective with hf
  have gInjective := Bijective.injective with hg
  exact sumMap.injective with fInjective, gInjective
  have fSurjective := Bijective.surjective with hf
  have gSurjective := Bijective.surjective with hg
  exact sumMap.surjective with fSurjective, gSurjective

theorem Equipotent.disjointSum (A B C D : Set)
  requires hAC : Equipotent(A, C)
  requires hBD : Equipotent(B, D)
  shows Equipotent(A ⊎ B, C ⊎ D)
proof
  obtain (f, hf) := hAC
  obtain (g, hg) := hBD
  witness sumMap(f, g, A, B, C, D)
  exact sumMap.bijective with hf, hg

theorem rightProjection.relates (y n : Set)
  requires hy : y ∈ n
  shows relates(rightProjection(n), inr(y), y)
proof
  unfold relates
  apply rightProjection.membership.backward
  constructor
  apply product.pair
  exact disjointSum.right with hy
  exact hy
  refine ⟨y, hy, rfl⟩

theorem rightProjection.factors (u v n : Set)
  requires h : relates(rightProjection(n), u, v)
  shows v ∈ n ∧ u = inr(v)
proof
  simp [relates, rightProjection.membership, orderedPair.injective]

theorem rightProjection.relation (n : Set)
  shows RelationBetween(rightProjection(n), 0 ⊎ n, n)
proof
  simp [RelationBetween, Subset]

theorem rightProjection.total (n : Set)
  shows TotalBetween(rightProjection(n), 0 ⊎ n, n)
proof
  intros u hu
  have sourceShape := disjointSum.membership.forward with hu
  cases sourceShape with impossibleLeft | rightShape
  obtain ⟨x, hxZero, uLeft⟩ := impossibleLeft
  rw nat.zero at hxZero
  simp
  obtain ⟨y, hyN, uRight⟩ := rightShape
  witness y
  constructor
  exact hyN
  rw uRight
  exact rightProjection.relates with hyN

theorem rightProjection.functional (n : Set)
  shows Functional(rightProjection(n))
proof
  simp [Functional, rightProjection.factors, inr.injective]

theorem rightProjection.function (n : Set)
  shows rightProjection(n) : 0 ⊎ n → n
proof
  constructor
  exact rightProjection.relation
  constructor
  exact rightProjection.functional
  exact rightProjection.total at n

theorem rightProjection.injective (n : Set)
  shows InjectiveBetween(rightProjection(n), 0 ⊎ n, n)
proof
  constructor
  exact rightProjection.function
  intros u hu v hv y hy huy hvy
  obtain (hyN, uShape) := rightProjection.factors with huy
  obtain (hyNAgain, vShape) := rightProjection.factors with hvy
  rw [uShape, vShape]
  rfl

theorem rightProjection.surjective (n : Set)
  shows SurjectiveBetween(rightProjection(n), 0 ⊎ n, n)
proof
  constructor
  exact rightProjection.function
  intros y hy
  witness inr(y)
  constructor
  exact disjointSum.right with hy
  exact rightProjection.relates with hy

theorem rightProjection.bijective (n : Set)
  shows BijectiveBetween(rightProjection(n), 0 ⊎ n, n)
proof
  refine ⟨rightProjection.injective at n, rightProjection.surjective at n⟩

theorem inl.fresh (m n : Nat)
  shows ¬inl(m) ∈ m ⊎ n
proof
  intro member
  have shape := disjointSum.membership.forward with member
  cases shape with leftShape | rightShape
  obtain ⟨x, hxM, tagsEqual⟩ := leftShape
  have same := inl.injective with tagsEqual
  rw ← same at hxM
  have irreflexive := membership.irreflexive with hxM
  contradiction
  obtain ⟨y, hyN, tagsEqualRight⟩ := rightShape
  have disjoint := tags.disjoint with tagsEqualRight
  contradiction

theorem disjointSum.successor_left (m n : Nat)
  shows succ(m) ⊎ n = insert(inl(m), m ⊎ n)
proof
  ext p
  constructor
  . intro hp
    have shape := disjointSum.membership.forward with hp
    cases shape with leftShape | rightShape
    . obtain (x, hx) := leftShape
      rcases hx with (hxSucc, pairShape)
      have successorShape := membership.successor_set{n := m}.forward with hxSucc
      cases successorShape with oldMember | newest
      . apply insert.contains at p, inl(m), m ⊎ n
        rw pairShape
        exact disjointSum.left with oldMember
      . rw [pairShape, newest]
        exact insert.self
    . obtain (y, hy) := rightShape
      rcases hy with (hyN, pairShape)
      apply insert.contains
      rw pairShape
      exact disjointSum.right with hyN
  . intro hp
    have shape := insert.membership.forward with hp
    cases shape with newest | oldMember
    . rw newest
      simp [disjointSum.membership, membership.successor_set]
    . have oldShape := disjointSum.membership.forward with oldMember
      cases oldShape with leftShape | rightShape
      . obtain (xOld, hxOld) := leftShape
        rcases hxOld with (hxM, pairShapeOld)
        rw pairShapeOld
        simp [disjointSum.membership, membership.successor_set]
      . obtain (yOld, hyOld) := rightShape
        rcases hyOld with (hyN, pairShapeRight)
        rw pairShapeRight
        exact disjointSum.right with hyN

theorem natural.successor_as_insert (n : Nat)
  shows succ(n) = insert(n, n)
proof
  simp [next.natural, next, insert, union2.commutative]

theorem natSumEquiv.bijective (n m : Nat)
  shows BijectiveBetween(natSumEquiv(m, n), m ⊎ n, add(m, n))
proof
  induction m
  | zero => settle with natSumEquiv, add, rightProjection.bijective
  | succ(k) using ih =>
    rw [natSumEquiv, add, disjointSum.successor_left, natural.successor_as_insert]
    apply adjoinMap.bijective
    . exact ih
    . intro member
      have shape := disjointSum.membership.forward with member
      cases shape with leftShape | rightShape
      . obtain (xFresh, hxFresh) := leftShape
        rcases hxFresh with (hxK, tagsEqual)
        have same := inl.injective with tagsEqual
        rw ← same at hxK
        have irreflexive := membership.irreflexive with hxK
        contradiction
      . obtain (yFresh, hyFresh) := rightShape
        rcases hyFresh with (hyN, tagsEqualRight)
        have disjoint := tags.disjoint with tagsEqualRight
        contradiction
    . exact membership.irreflexive at add(k, n)

theorem Equipotent.nat_disjointSum (m n : Set)
  requires hm : Nat(m)
  requires hn : Nat(n)
  shows (m ⊎ n) ≃ add(m, n)
proof
  witness natSumEquiv(m, n)
  exact natSumEquiv.bijective with hn, hm

theorem HasCardinality.disjointSum (A B m n : Set)
  requires hA : HasCardinality(A, m)
  requires hB : HasCardinality(B, n)
  shows HasCardinality(A ⊎ B, add(m, n))
proof
  rcases hA with (hm, equipotentA)
  rcases hB with (hn, equipotentB)
  constructor
  . exact Grace.Core.Nat.Algebra.add.closed at n, m
  . have tagged := Equipotent.disjointSum with equipotentA, equipotentB
    have numeric := Equipotent.nat_disjointSum with hm, hn
    exact Equipotent.transitive with tagged, numeric

theorem Finite.disjointSum (A B : Set)
  requires finiteA : Finite(A)
  requires finiteB : Finite(B)
  shows Finite(A ⊎ B)
proof
  obtain (m, hA) := finiteA
  obtain (n, hB) := finiteB
  witness add(m, n)
  exact HasCardinality.disjointSum with hA, hB