-- A set with n elements has 2^n subsets
-- 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.PowersetCount
import Counting
import Sum
import FiniteSubsets

-- Subsets of insert(x, A) are sorted by whether they contain x: those that do
-- not are subsets of A, and those that do become subsets of A once x is
-- removed. The graph sends S to inl(S) or inr(S ∖ {x}) accordingly.
def powersetSplit (A x : Set) : Set :=
  { p ∈ 𝒫(insert(x, A)) × (𝒫(A) ⊎ 𝒫(A)) |
    ∃ S ∈ 𝒫(insert(x, A)), ((¬(x ∈ S) ∧ p = (S, inl(S))) ∨ (x ∈ S ∧ p = (S, inr(S ∖ {x})))) }

theorem subset.without (A x S : Set)
  requires subset : S ⊆ insert(x, A)
  shows S ∖ {x} ⊆ A
proof
  simp

theorem subset.absent (A x S : Set)
  requires subset : S ⊆ insert(x, A)
  requires absent : ¬(x ∈ S)
  shows S ⊆ A
proof
  simp

theorem insert.remove_fresh (T x : Set)
  requires fresh : ¬(x ∈ T)
  shows insert(x, T) ∖ {x} = T
proof
  ext y
  simp

theorem powersetSplit.relates_absent (A x S : Set)
  requires member : S ∈ 𝒫(insert(x, A))
  requires absent : ¬(x ∈ S)
  shows relates(powersetSplit(A, x), S, inl(S))
proof
  have subset := powerset.membership_iff.forward with member
  have small := subset.absent with *
  have inPower := powerset.membership_iff.backward with small
  unfold relates
  apply powersetSplit.membership.backward
  constructor
  . apply product.pair at S, inl(S), 𝒫(insert(x, A)), 𝒫(A) ⊎ 𝒫(A)
    . exact member
    . exact disjointSum.left with inPower
  . witness S
    constructor
    . exact member
    . left
      refine ⟨absent, rfl⟩

theorem powersetSplit.relates_present (A x S : Set)
  requires member : S ∈ 𝒫(insert(x, A))
  requires present : x ∈ S
  shows relates(powersetSplit(A, x), S, inr(S ∖ {x}))
proof
  have subset := powerset.membership_iff.forward with member
  have rest := subset.without with *
  have inPower := powerset.membership_iff.backward with rest
  unfold relates
  apply powersetSplit.membership.backward
  constructor
  . apply product.pair at S, inr(S ∖ {x}), 𝒫(insert(x, A)), 𝒫(A) ⊎ 𝒫(A)
    . exact member
    . exact disjointSum.right with inPower
  . witness S
    constructor
    . exact member
    . right
      refine ⟨present, rfl⟩

theorem powersetSplit.factors (A x u v : Set)
  requires h : relates(powersetSplit(A, x), u, v)
  shows u ∈ 𝒫(insert(x, A)) ∧ ((¬(x ∈ u) ∧ v = inl(u)) ∨ (x ∈ u ∧ v = inr(u ∖ {x})))
proof
  unfold relates at h
  simp [powersetSplit.membership] at h
  rcases h with ⟨bounded, S, inPower, shape⟩
  cases shape with absent | present
  . rcases absent with ⟨notIn, pair⟩
    obtain ⟨source, target⟩ := orderedPair.injective with pair
    rw ← source at inPower
    rw ← source at notIn
    rw ← source at target
    constructor
    . exact inPower
    . left
      refine ⟨notIn, target⟩
  . rcases present with ⟨isIn, pair⟩
    obtain ⟨source, target⟩ := orderedPair.injective with pair
    rw ← source at inPower
    rw ← source at isIn
    rw ← source at target
    constructor
    . exact inPower
    . right
      refine ⟨isIn, target⟩

theorem powersetSplit.relation (A x : Set)
  shows RelationBetween(powersetSplit(A, x), 𝒫(insert(x, A)), 𝒫(A) ⊎ 𝒫(A))
proof
  simp [RelationBetween, Subset]

theorem powersetSplit.total (A x : Set)
  shows TotalBetween(powersetSplit(A, x), 𝒫(insert(x, A)), 𝒫(A) ⊎ 𝒫(A))
proof
  intros S member
  have subset := powerset.membership_iff.forward with member
  have alternatives : x ∈ S ∨ ¬(x ∈ S) by tauto
  cases alternatives with present | absent
  . have rest := subset.without with *
    have inPower := powerset.membership_iff.backward with rest
    witness inr(S ∖ {x})
    constructor
    . exact disjointSum.right with inPower
    . exact powersetSplit.relates_present with *
  . have small := subset.absent with *
    have inPower := powerset.membership_iff.backward with small
    witness inl(S)
    constructor
    . exact disjointSum.left with inPower
    . exact powersetSplit.relates_absent with *

theorem powersetSplit.functional (A x : Set)
  shows Functional(powersetSplit(A, x))
proof
  intros u v huv w huw
  have first := powersetSplit.factors with huv
  have second := powersetSplit.factors with huw
  rcases first with ⟨inPower, firstShape⟩
  rcases second with ⟨inPowerAgain, secondShape⟩
  cases firstShape with firstAbsent | firstPresent
  . rcases firstAbsent with ⟨absent, vShape⟩
    cases secondShape with secondAbsent | secondPresent
    . rcases secondAbsent with ⟨absentAgain, wShape⟩
      rw [vShape, wShape]
      rfl
    . rcases secondPresent with ⟨present, wShape⟩
      exfalso
      exact absent with present
  . rcases firstPresent with ⟨present, vShape⟩
    cases secondShape with secondAbsent | secondPresent
    . rcases secondAbsent with ⟨absent, wShape⟩
      exfalso
      exact absent with present
    . rcases secondPresent with ⟨presentAgain, wShape⟩
      rw [vShape, wShape]
      rfl

theorem powersetSplit.function (A x : Set)
  shows powersetSplit(A, x) : 𝒫(insert(x, A)) → 𝒫(A) ⊎ 𝒫(A)
proof
  constructor
  . exact powersetSplit.relation at A, x
  . constructor
    . exact powersetSplit.functional at A, x
    . exact powersetSplit.total at A, x

theorem powersetSplit.injective (A x : Set)
  shows InjectiveBetween(powersetSplit(A, x), 𝒫(insert(x, A)), 𝒫(A) ⊎ 𝒫(A))
proof
  constructor
  . exact powersetSplit.function at A, x
  . intros u hu v hv z hz huz hvz
    have first := powersetSplit.factors with huz
    have second := powersetSplit.factors with hvz
    rcases first with ⟨uPower, firstShape⟩
    rcases second with ⟨vPower, secondShape⟩
    cases firstShape with firstAbsent | firstPresent
    . rcases firstAbsent with ⟨uAbsent, zLeft⟩
      cases secondShape with secondAbsent | secondPresent
      . rcases secondAbsent with ⟨vAbsent, zLeftAgain⟩
        rw zLeft at zLeftAgain
        exact inl.injective with zLeftAgain
      . rcases secondPresent with ⟨vPresent, zRight⟩
        rw zLeft at zRight
        exfalso
        exact tags.disjoint with zRight
    . rcases firstPresent with ⟨uPresent, zRight⟩
      cases secondShape with secondAbsent | secondPresent
      . rcases secondAbsent with ⟨vAbsent, zLeft⟩
        rw zLeft at zRight
        exfalso
        exact tags.disjoint with zRight
      . rcases secondPresent with ⟨vPresent, zRightAgain⟩
        rw zRight at zRightAgain
        have sameRest := inr.injective with zRightAgain
        have restoreU := insert.restore with uPresent
        have restoreV := insert.restore with vPresent
        rw [← restoreU, ← restoreV, sameRest]
        rfl

theorem powersetSplit.surjective (A x : Set)
  requires fresh : ¬(x ∈ A)
  shows SurjectiveBetween(powersetSplit(A, x), 𝒫(insert(x, A)), 𝒫(A) ⊎ 𝒫(A))
proof
  constructor
  . exact powersetSplit.function at A, x
  . intros v hv
    have shape := disjointSum.membership.forward with hv
    cases shape with leftShape | rightShape
    . obtain ⟨T, inPower, vShape⟩ := leftShape
      have subset := powerset.membership_iff.forward with inPower
      have larger : T ⊆ insert(x, A) by simp
      have member := powerset.membership_iff.backward with larger
      have absent : ¬(x ∈ T) by simp
      witness T
      constructor
      . exact member
      . rw vShape
        exact powersetSplit.relates_absent with *
    . obtain ⟨T, inPower, vShape⟩ := rightShape
      have subset := powerset.membership_iff.forward with inPower
      have larger : insert(x, T) ⊆ insert(x, A) by simp
      have member := powerset.membership_iff.backward with larger
      have present := insert.self at x, T
      have absent : ¬(x ∈ T) by simp
      have removed := insert.remove_fresh with absent
      have related := powersetSplit.relates_present with *
      rw removed at related
      witness insert(x, T)
      constructor
      . exact member
      . rw vShape
        exact related

theorem powersetSplit.bijective (A x : Set)
  requires fresh : ¬(x ∈ A)
  shows BijectiveBetween(powersetSplit(A, x), 𝒫(insert(x, A)), 𝒫(A) ⊎ 𝒫(A))
proof
  simp [BijectiveBetween, powersetSplit.injective, powersetSplit.surjective]

theorem Equipotent.powerset_insert (A x : Set)
  requires fresh : ¬(x ∈ A)
  shows 𝒫(insert(x, A)) ≃ 𝒫(A) ⊎ 𝒫(A)
proof
  simp [Equipotent, powersetSplit.bijective]

-- Adding a fresh element doubles the number of subsets.
theorem HasCardinality.powerset_insert (A x m : Set)
  requires counted : HasCardinality(𝒫(A), m)
  requires fresh : ¬(x ∈ A)
  shows HasCardinality(𝒫(insert(x, A)), add(m, m))
proof
  have doubled := HasCardinality.disjointSum with *
  rcases doubled with ⟨natural, equivalent⟩
  have split := Equipotent.powerset_insert with *
  constructor
  . exact natural
  . exact Equipotent.transitive with *

theorem HasCardinality.exists_member (A : Set, n : Nat)
  requires counted : HasCardinality(A, succ(n))
  shows ∃ a, a ∈ A
proof
  rcases counted with ⟨natural, f, bijective⟩
  obtain ⟨function, onto⟩ := Bijective.surjective with bijective
  have last := membership.self_successor at n with *
  obtain ⟨a, member, related⟩ := onto with n, last
  witness a
  exact member

theorem powerset.empty
  shows 𝒫(∅) = {∅}
proof
  ext S
  constructor
  . intro member
    have subset := powerset.membership_iff.forward with member
    have reverse := empty.subset at S
    have equal := Subset.antisymmetric at S, ∅ with *
    exact singleton.membership.backward with equal
  . intro member
    have equal := singleton.membership.forward with member
    rw equal
    have reflexive := empty.subset at ∅
    exact powerset.membership_iff.backward with reflexive

theorem HasCardinality.powerset_empty
  shows HasCardinality(𝒫(∅), 1)
proof
  constructor
  . exact Naturals.constructor.succ at 0 with Naturals.constructor.zero
  . rw powerset.empty
    exact Equipotent.singleton_one

-- A set with n elements has 2^n subsets. Induction on n: remove an element,
-- count the subsets of the rest, and double.
theorem powerset.cardinality (A : Set) (n : Nat)
  requires counted : HasCardinality(A, n)
  shows HasCardinality(𝒫(A), 2 ^ n)
proof
  induction n generalizing A
  | zero =>
    intros A counted
    have empty := HasCardinality.empty_of_zero at A with counted
    rw [empty, pow.zero]
    exact HasCardinality.powerset_empty
  | succ(k) using ih =>
    intros A counted
    obtain ⟨a, member⟩ := HasCardinality.exists_member with counted
    have finite := Finite.from_cardinality with counted
    have count := cardinality.eq_of_has with counted
    have removal := cardinality.remove at A, a with *
    rw count at removal
    have subset := difference.subset_left at A, {a}
    have finiteRest := Finite.subset at A ∖ {a}, A with *
    obtain ⟨restNatural, restEquivalent⟩ := cardinality.correct with finiteRest
    have equal := Naturals.injective.succ with removal
    have restHas : HasCardinality(A ∖ {a}, k)
      unfold HasCardinality
      rw equal
      refine ⟨restNatural, restEquivalent⟩
    have powerRest := ih with A ∖ {a}, restHas
    have fresh : ¬(a ∈ A ∖ {a}) by simp [difference.membership, singleton.membership]
    have doubled := HasCardinality.powerset_insert at A ∖ {a}, a, 2 ^ k with *
    have restored := insert.restore with member
    rw restored at doubled
    have power : 2 ^ succ(k) = add(2 ^ k, 2 ^ k) by simp [pow.succ, mul.two_left]
    rwa power