/-!
# Finite counting on `{0, …, n-1}` (core Lean only)

Individuals are the natural numbers `< n`; a *set of individuals* is a `Nat → Bool`.
`cnt n g` is the number of `x < n` with `g x`.  The main tools are the "rank pieces"
(used to build the sets `X`, `Y` in the proof of Lemma 1) and the two realization lemmas.
-/

namespace AntiDyn

def cnt (n : Nat) (g : Nat → Bool) : Nat := (List.range n).countP g

theorem cnt_zero (g : Nat → Bool) : cnt 0 g = 0 := by simp [cnt]

theorem cnt_succ (n : Nat) (g : Nat → Bool) :
    cnt (n+1) g = cnt n g + (if g n then 1 else 0) := by
  simp [cnt, List.range_succ, List.countP_append]

theorem cnt_congr (n : Nat) {g h : Nat → Bool} (e : ∀ x, x < n → g x = h x) :
    cnt n g = cnt n h := by
  induction n with
  | zero => simp [cnt_zero]
  | succ n ih =>
    rw [cnt_succ, cnt_succ, ih (fun x hx => e x (by omega)), e n (by omega)]

theorem cnt_le (n : Nat) (g : Nat → Bool) : cnt n g ≤ n := by
  induction n with
  | zero => simp [cnt_zero]
  | succ n ih => rw [cnt_succ]; split <;> omega

theorem cnt_mono (n : Nat) {g h : Nat → Bool} (e : ∀ x, g x = true → h x = true) :
    cnt n g ≤ cnt n h := by
  induction n with
  | zero => simp [cnt_zero]
  | succ n ih =>
    rw [cnt_succ, cnt_succ]
    have := e n
    split <;> split <;> simp_all <;> omega

/-- `cnt g = cnt (g ∧ P) + cnt (g ∧ ¬P)`. -/
theorem cnt_split (n : Nat) (g P : Nat → Bool) :
    cnt n g = cnt n (fun x => g x && P x) + cnt n (fun x => g x && !P x) := by
  induction n with
  | zero => simp [cnt_zero]
  | succ n ih =>
    rw [cnt_succ, cnt_succ, cnt_succ, ih]
    cases g n <;> cases P n <;> simp <;> omega

theorem cnt_compl (n : Nat) (g : Nat → Bool) : cnt n g + cnt n (fun x => !g x) = n := by
  induction n with
  | zero => simp [cnt_zero]
  | succ n ih =>
    rw [cnt_succ, cnt_succ]
    cases g n <;> simp <;> omega

theorem cnt_pos (g : Nat → Bool) {n d : Nat} (hd : d < n) (hg : g d = true) : 0 < cnt n g := by
  induction n with
  | zero => omega
  | succ n ih =>
    rw [cnt_succ]
    by_cases h : d < n
    · have := ih h; omega
    · have : d = n := by omega
      subst this; simp [hg]

theorem cnt_or_disj (n : Nat) (g h : Nat → Bool) (e : ∀ x, ¬ (g x = true ∧ h x = true)) :
    cnt n (fun x => g x || h x) = cnt n g + cnt n h := by
  induction n with
  | zero => simp [cnt_zero]
  | succ n ih =>
    rw [cnt_succ, cnt_succ, cnt_succ, ih]
    have := e n
    cases hg : g n <;> cases hh : h n <;> simp_all <;> omega

/-- The elements of `S` whose rank (number of elements of `S` below them) lies in `[lo, hi)`. -/
def piece (S : Nat → Bool) (lo hi : Nat) (x : Nat) : Bool :=
  S x && decide (lo ≤ cnt x S) && decide (cnt x S < hi)

theorem piece_sub (S : Nat → Bool) (lo hi x : Nat) (h : piece S lo hi x = true) : S x = true := by
  simp [piece] at h; exact h.1.1

theorem cnt_piece (S : Nat → Bool) (lo hi : Nat) (hlh : lo ≤ hi) (n : Nat) :
    cnt n (piece S lo hi) = min hi (cnt n S) - min lo (cnt n S) := by
  induction n with
  | zero => simp [cnt_zero]
  | succ n ih =>
    rw [cnt_succ, cnt_succ, ih]
    by_cases hS : S n = true
    · simp only [piece, hS, Bool.true_and, Bool.and_eq_true, decide_eq_true_eq]
      by_cases h1 : lo ≤ cnt n S <;> by_cases h2 : cnt n S < hi <;> simp [h1, h2] <;> omega
    · simp [piece, hS]

/-- Choose a subset of `S` (a set of individuals) of any size `k ≤ |S|`. -/
theorem exists_subset (n : Nat) (S : Nat → Bool) (k : Nat) (hk : k ≤ cnt n S) :
    ∃ T : Nat → Bool, (∀ x, T x = true → S x = true) ∧ cnt n T = k := by
  refine ⟨piece S 0 k, piece_sub S 0 k, ?_⟩
  rw [cnt_piece S 0 k (by omega)]
  omega

/-- Choose `Y` meeting two *disjoint* sets `U`, `V` in prescribed numbers of elements. -/
theorem realize_two (n : Nat) (U V : Nat → Bool) (hd : ∀ x, ¬ (U x = true ∧ V x = true))
    (u v : Nat) (hu : u ≤ cnt n U) (hv : v ≤ cnt n V) :
    ∃ Y : Nat → Bool, cnt n (fun x => U x && Y x) = u ∧ cnt n (fun x => V x && Y x) = v := by
  refine ⟨fun x => piece U 0 u x || piece V 0 v x, ?_, ?_⟩
  · have : cnt n (fun x => U x && (piece U 0 u x || piece V 0 v x)) = cnt n (piece U 0 u) := by
      apply cnt_congr; intro x _
      have := hd x
      by_cases hU : U x = true
      · have hV : V x = false := by simpa [hU] using this
        have : piece V 0 v x = false := by
          cases h : piece V 0 v x
          · rfl
          · have := piece_sub V 0 v x h; simp_all
        simp [this, hU]
      · simp [hU]
        have := piece_sub U 0 u x
        cases h : piece U 0 u x
        · rfl
        · simp_all
    rw [this, cnt_piece U 0 u (by omega)]
    omega
  · have : cnt n (fun x => V x && (piece U 0 u x || piece V 0 v x)) = cnt n (piece V 0 v) := by
      apply cnt_congr; intro x _
      have := hd x
      by_cases hV : V x = true
      · have hU : U x = false := by simpa [hV] using this
        have : piece U 0 u x = false := by
          cases h : piece U 0 u x
          · rfl
          · have := piece_sub U 0 u x h; simp_all
        simp [this, hV]
      · simp [hV]
        have := piece_sub V 0 v x
        cases h : piece V 0 v x
        · rfl
        · simp_all
    rw [this, cnt_piece V 0 v (by omega)]
    omega

/-- **Realization lemma** (the constructions of `X`, `Y` in the proof of Lemma 1(i)).
Let `P` have `k` elements among the `n` individuals.  For any `s₁ + s₂ ≤ k` and
`c₁ + c₂ ≤ n − k` there are `X, Y` with
`|P ∩ X ∖ Y| = s₁`, `|P ∩ X ∩ Y| = s₂`, `|X ∖ P ∖ Y| = c₁`, `|(X ∖ P) ∩ Y| = c₂`. -/
theorem realize_four (n : Nat) (P : Nat → Bool) (s₁ s₂ c₁ c₂ : Nat)
    (hs : s₁ + s₂ ≤ cnt n P) (hc : c₁ + c₂ ≤ n - cnt n P) :
    ∃ X Y : Nat → Bool,
      cnt n (fun x => P x && X x && !Y x) = s₁ ∧ cnt n (fun x => P x && X x && Y x) = s₂ ∧
      cnt n (fun x => !P x && X x && !Y x) = c₁ ∧ cnt n (fun x => !P x && X x && Y x) = c₂ := by
  let N : Nat → Bool := fun x => !P x
  have hcomp : cnt n P + cnt n N = n := cnt_compl n P
  let X : Nat → Bool := fun x => piece P 0 (s₁ + s₂) x || piece N 0 (c₁ + c₂) x
  have hXP : ∀ x, (P x && X x) = piece P 0 (s₁ + s₂) x := by
    intro x
    by_cases hp : P x = true
    · have : piece N 0 (c₁ + c₂) x = false := by
        cases h : piece N 0 (c₁ + c₂) x
        · rfl
        · have := piece_sub N 0 _ x h; simp [N, hp] at this
      simp [X, hp, this]
    · have h1 : piece P 0 (s₁ + s₂) x = false := by
        cases h : piece P 0 (s₁ + s₂) x
        · rfl
        · have := piece_sub P 0 _ x h; simp_all
      simp [hp, h1]
  have hXN : ∀ x, (N x && X x) = piece N 0 (c₁ + c₂) x := by
    intro x
    by_cases hp : P x = true
    · have h1 : piece N 0 (c₁ + c₂) x = false := by
        cases h : piece N 0 (c₁ + c₂) x
        · rfl
        · have := piece_sub N 0 _ x h; simp [N, hp] at this
      simp [N, hp, h1]
    · have h1 : piece P 0 (s₁ + s₂) x = false := by
        cases h : piece P 0 (s₁ + s₂) x
        · rfl
        · have := piece_sub P 0 _ x h; simp_all
      have hp' : P x = false := by simpa using hp
      simp [N, X, hp', h1]
  have cP : cnt n (fun x => P x && X x) = s₁ + s₂ := by
    rw [cnt_congr n (fun x _ => hXP x), cnt_piece P 0 _ (by omega)]; omega
  have cN : cnt n (fun x => N x && X x) = c₁ + c₂ := by
    rw [cnt_congr n (fun x _ => hXN x), cnt_piece N 0 _ (by omega)]; omega
  obtain ⟨Y, hY1, hY2⟩ := realize_two n (fun x => P x && X x) (fun x => N x && X x)
    (by intro x ⟨h1, h2⟩; simp [N] at h1 h2; simp_all) s₂ c₂
    (by rw [cP]; omega) (by rw [cN]; omega)
  refine ⟨X, Y, ?_, hY1, ?_, hY2⟩
  · have := cnt_split n (fun x => P x && X x) Y
    rw [cP, hY1] at this
    have e : cnt n (fun x => P x && X x && !Y x) = s₁ := by omega
    exact e
  · have := cnt_split n (fun x => N x && X x) Y
    rw [cN, hY2] at this
    have e : cnt n (fun x => N x && X x && !Y x) = c₁ := by omega
    exact e

end AntiDyn