import AntiDynamics.Quantified

/-!
# Lemma 1 (Section 5.2): unembedded quantificational clauses

Both parts are proved *with the constructions repaired* where the paper's Case 2 (Zone B)
has a gap (see RESULTS.md, issue Q2).
-/

namespace AntiDyn
namespace Quantified

variable {W ι ρ κ : Type} (M : QModel W ι ρ κ)

@[simp] theorem eval_quant (q : κ) (P R : Pred ρ) (w : W) :
    (sys M).eval (.leaf (.quant q P R)) w = M.qsem q P R w := rfl

/-- Transparency of the restrictor trigger `P̲_k P'_j`: for all `X = P_m`, `Y`,
`C ⊨ (Q (P_k and P_m) . Y) ⇔ (Q P_m . Y)`. -/
def RT (C : WSet W) (q : κ) (k : ρ) : Prop :=
  ∀ (m : ρ) (Y : Pred ρ) w, C w → M.qsem q (.conj k m) Y w = M.qsem q (.plain m) Y w

/-- Transparency of the nuclear-scope trigger `R̲_k R'_j` after restrictor `P`: for all `Y = P_m`,
`C ⊨ (Q P . (R_k and P_m)) ⇔ (Q P . P_m)`. -/
def ST (C : WSet W) (q : κ) (P : Pred ρ) (k : ρ) : Prop :=
  ∀ (m : ρ) w, C w → M.qsem q P (.conj k m) w = M.qsem q P (.plain m) w

/-- Unpacking Transparency for a quantificational clause. -/
theorem transp_quant (C : WSet W) (q : κ) (P R : Pred ρ) :
    (sys M).Transp C (.leaf (.quant q P R)) ↔
      (∀ k j, P = .trig k j → RT M C q k) ∧ (∀ k j, R = .trig k j → ST M C q P k) := by
  rw [Sys.transp_leaf]
  constructor
  · intro h
    refine ⟨?_, ?_⟩
    · rintro k j rfl m Y w hw
      exact h _ _ (Or.inl ⟨k, j, m, Y, rfl, rfl, rfl⟩) w hw
    · rintro k j rfl m w hw
      exact h _ _ (Or.inr ⟨k, j, m, rfl, rfl, rfl⟩) w hw
  · rintro ⟨h1, h2⟩ φ₁ φ₂ hv w hw
    rcases hv with ⟨k, j, m, Y, hP, rfl, rfl⟩ | ⟨k, j, m, hR, rfl, rfl⟩
    · exact h1 k j hP m Y w hw
    · exact h2 k j hR m w hw

@[simp] theorem pstat_conj (k m : ρ) (w : W) (d : Nat) :
    M.pstat (.conj k m) w d = (M.I k w d && M.I m w d) := rfl
@[simp] theorem pstat_plain (k : ρ) (w : W) (d : Nat) :
    M.pstat (.plain k) w d = M.I k w d := rfl

theorem cnt_and_comm (n : Nat) (A B : Nat → Bool) :
    cnt n (fun x => A x && B x) = cnt n (fun x => B x && A x) :=
  cnt_congr n (fun x _ => by cases A x <;> cases B x <;> rfl)

/-- World-wise "if" half of Lemma 1(i): if every individual satisfies `P_k` at `w`,
`w ⊨ (Q (P_k and P_m) . Y) ⇔ (Q P_m . Y)`. -/
theorem rt_at (q : κ) (k m : ρ) (Y : Pred ρ) (w : W)
    (h : ∀ d, d < M.n → M.I k w d = true) :
    M.qsem q (.conj k m) Y w = M.qsem q (.plain m) Y w := by
  simp only [QModel.qsem, pstat_conj, pstat_plain]
  have e1 : cnt M.n (fun d => (M.I k w d && M.I m w d) && !(M.pstat Y w d)) =
      cnt M.n (fun d => M.I m w d && !(M.pstat Y w d)) :=
    cnt_congr _ (fun d hd => by simp [h d hd])
  have e2 : cnt M.n (fun d => (M.I k w d && M.I m w d) && M.pstat Y w d) =
      cnt M.n (fun d => M.I m w d && M.pstat Y w d) :=
    cnt_congr _ (fun d hd => by simp [h d hd])
  rw [e1, e2]

/-- World-wise "if" half of Lemma 1(ii): if every restrictor-individual satisfies `R_k` at
`w`, `w ⊨ (Q P . (R_k and P_m)) ⇔ (Q P . P_m)`. -/
theorem st_at (q : κ) (P : Pred ρ) (k m : ρ) (w : W)
    (h : ∀ d, d < M.n → M.pstat P w d = true → M.I k w d = true) :
    M.qsem q P (.conj k m) w = M.qsem q P (.plain m) w := by
  simp only [QModel.qsem, pstat_conj, pstat_plain]
  have e1 : cnt M.n (fun d => M.pstat P w d && !(M.I k w d && M.I m w d)) =
      cnt M.n (fun d => M.pstat P w d && !M.I m w d) :=
    cnt_congr _ (fun d hd => by
      by_cases hd' : M.pstat P w d = true
      · simp [hd', h d hd hd']
      · simp at hd'; simp [hd'])
  have e2 : cnt M.n (fun d => M.pstat P w d && (M.I k w d && M.I m w d)) =
      cnt M.n (fun d => M.pstat P w d && M.I m w d) :=
    cnt_congr _ (fun d hd => by
      by_cases hd' : M.pstat P w d = true
      · simp [hd', h d hd hd']
      · simp at hd'; simp [hd'])
  rw [e1, e2]

/-- The "if" half of Lemma 1(i): if every individual satisfies `P_k` throughout `C`, the
restrictor trigger is transparent.  (No hypothesis on `f` or on expressibility needed.) -/
theorem rt_of_ent (C : WSet W) (q : κ) (k : ρ)
    (h : ∀ w, C w → ∀ d, d < M.n → M.I k w d = true) : RT M C q k :=
  fun m Y w hw => rt_at M q k m Y w (h w hw)

/-- The "if" half of Lemma 1(ii). -/
theorem st_of_ent (C : WSet W) (q : κ) (P : Pred ρ) (k : ρ)
    (h : ∀ w, C w → ∀ d, d < M.n → M.pstat P w d = true → M.I k w d = true) :
    ST M C q P k :=
  fun m w hw => st_at M q P k m w (h w hw)

/-! ### Lemma 1 (i): trigger in the restrictor -/

theorem lemma1_i (hexp : M.Expressive) (C : WSet W) (q : κ) (k : ρ)
    (htri : ∃ a b a' b', a + b ≤ M.n ∧ a' + b' ≤ M.n ∧ M.f q a b ≠ M.f q a' b') :
    RT M C q k ↔ ∀ w, C w → ∀ d, d < M.n → M.I k w d = true := by
  constructor
  · intro h
    apply Classical.byContradiction
    intro hne
    -- a world `w ∈ C` and an individual `d0 < n` with `¬ P(d0)`
    have : ∃ w, C w ∧ ∃ d, d < M.n ∧ M.I k w d ≠ true := by
      apply Classical.byContradiction
      intro hh
      apply hne
      intro w hw d hd
      apply Classical.byContradiction
      intro hd'
      exact hh ⟨w, hw, d, hd, hd'⟩
    obtain ⟨w, hw, d0, hd0, hP0⟩ := this
    have hkn : cnt M.n (M.I k w) < M.n := by
      have h1 := cnt_compl M.n (M.I k w)
      have h2 : 0 < cnt M.n (fun x => !M.I k w x) :=
        cnt_pos (fun x => !M.I k w x) hd0 (by simpa using hP0)
      omega
    obtain ⟨a, b, i, j, hia, hjb, hab, hij, hfne⟩ :=
      find_step (M.f q) M.n (cnt M.n (M.I k w)) hkn htri
    obtain ⟨X, Y, e1, e2, e3, e4⟩ := realize_four M.n (M.I k w) (a - i) (b - j) i j (by omega) hij
    obtain ⟨mX, hmX⟩ := hexp (fun _ => X)
    obtain ⟨mY, hmY⟩ := hexp (fun _ => Y)
    have hmX' : ∀ d, M.I mX w d = X d := fun d => hmX w d
    have hmY' : ∀ d, M.I mY w d = Y d := fun d => hmY w d
    have key := h mX (.plain mY) w hw
    simp only [QModel.qsem, pstat_conj, pstat_plain, hmX', hmY'] at key
    have hR1 : cnt M.n (fun d => X d && !Y d) = a := by
      rw [cnt_split M.n (fun d => X d && !Y d) (M.I k w)]
      have h1 := cnt_congr M.n (g := fun d => X d && !Y d && M.I k w d)
        (h := fun d => M.I k w d && X d && !Y d)
        (fun x _ => by cases M.I k w x <;> cases X x <;> cases Y x <;> rfl)
      have h2 := cnt_congr M.n (g := fun d => X d && !Y d && !M.I k w d)
        (h := fun d => !M.I k w d && X d && !Y d)
        (fun x _ => by cases M.I k w x <;> cases X x <;> cases Y x <;> rfl)
      omega
    have hR2 : cnt M.n (fun d => X d && Y d) = b := by
      rw [cnt_split M.n (fun d => X d && Y d) (M.I k w)]
      have h1 := cnt_congr M.n (g := fun d => X d && Y d && M.I k w d)
        (h := fun d => M.I k w d && X d && Y d)
        (fun x _ => by cases M.I k w x <;> cases X x <;> cases Y x <;> rfl)
      have h2 := cnt_congr M.n (g := fun d => X d && Y d && !M.I k w d)
        (h := fun d => !M.I k w d && X d && Y d)
        (fun x _ => by cases M.I k w x <;> cases X x <;> cases Y x <;> rfl)
      omega
    rw [e1, e2, hR1, hR2] at key
    exact hfne key
  · exact fun h => rt_of_ent M C q k h

/-! ### Lemma 1 (ii): trigger in the nuclear scope -/

theorem lemma1_ii (hexp : M.Expressive) (C : WSet W) (q : κ) (k : ρ) (P : Pred ρ) (p : Nat)
    (hp : ∀ w, C w → cnt M.n (M.pstat P w) = p)
    (hline : ∃ b₁ b₂, b₁ ≤ p ∧ b₂ ≤ p ∧ M.f q (p - b₁) b₁ ≠ M.f q (p - b₂) b₂) :
    ST M C q P k ↔ ∀ w, C w → ∀ d, d < M.n → M.pstat P w d = true → M.I k w d = true := by
  constructor
  · intro h
    apply Classical.byContradiction
    intro hne
    have : ∃ w, C w ∧ ∃ d, d < M.n ∧ M.pstat P w d = true ∧ M.I k w d ≠ true := by
      apply Classical.byContradiction
      intro hh
      apply hne
      intro w hw d hd hPd
      apply Classical.byContradiction
      intro hd'
      exact hh ⟨w, hw, d, hd, hPd, hd'⟩
    obtain ⟨w, hw, d0, hd0, hPd0, hRd0⟩ := this
    have hpw := hp w hw
    -- b = |P ∩ R| < p
    have hsplit : cnt M.n (M.pstat P w) =
        cnt M.n (fun d => M.pstat P w d && M.I k w d) +
        cnt M.n (fun d => M.pstat P w d && !M.I k w d) := cnt_split M.n (M.pstat P w) (M.I k w)
    have hpos : 0 < cnt M.n (fun d => M.pstat P w d && !M.I k w d) :=
      cnt_pos (fun d => M.pstat P w d && !M.I k w d) hd0 (by simp [hPd0]; simpa using hRd0)
    generalize hb : cnt M.n (fun d => M.pstat P w d && M.I k w d) = b at hsplit
    have hb_lt : b < p := by omega
    -- some `b*` with a different value
    have : ∃ bs, bs ≤ p ∧ M.f q (p - bs) bs ≠ M.f q (p - b) b := by
      obtain ⟨b₁, b₂, h1, h2, hne'⟩ := hline
      by_cases e : M.f q (p - b₁) b₁ = M.f q (p - b) b
      · exact ⟨b₂, h2, fun e2 => hne' (e.trans e2.symm)⟩
      · exact ⟨b₁, h1, e⟩
    obtain ⟨bs, hbs, hbsne⟩ := this
    -- pick β₁ ≤ β₂ with f(p-β₁, β₁) ≠ f(p-β₂, β₂), β₁ ≤ b, β₂ - β₁ ≤ p - b
    have : ∃ β₁ β₂, β₁ ≤ β₂ ∧ β₁ ≤ b ∧ β₂ - β₁ ≤ p - b ∧
        M.f q (p - β₁) β₁ ≠ M.f q (p - β₂) β₂ := by
      by_cases hlt : bs < b
      · have hne2 : (fun x => M.f q (p - x) x) bs ≠ (fun x => M.f q (p - x) x) (bs + (b - bs)) := by
          have : bs + (b - bs) = b := by omega
          simp only [this]; exact hbsne
        obtain ⟨x, hx1, hx2, hx3⟩ := ivt (fun x => M.f q (p - x) x) (b - bs) bs (by omega) hne2
        exact ⟨x, x + 1, by omega, by omega, by omega, hx3⟩
      · have hgt : b < bs := by
          have : bs ≠ b := fun e => hbsne (by subst e; rfl)
          omega
        exact ⟨b, bs, by omega, Nat.le_refl _, by omega, fun e => hbsne e.symm⟩
    obtain ⟨β₁, β₂, hβ12, hβ1, hβ2, hfne⟩ := this
    obtain ⟨Y, hY1, hY2⟩ := realize_two M.n (fun d => M.pstat P w d && M.I k w d)
      (fun d => M.pstat P w d && !M.I k w d)
      (by intro x ⟨h1, h2⟩; simp at h1 h2; simp_all) β₁ (β₂ - β₁)
      (by omega) (by omega)
    obtain ⟨mY, hmY⟩ := hexp (fun _ => Y)
    have hmY' : ∀ d, M.I mY w d = Y d := fun d => hmY w d
    have key := h mY w hw
    simp only [QModel.qsem, pstat_conj, pstat_plain, hmY'] at key
    -- counts
    have c1 : cnt M.n (fun d => M.pstat P w d && (M.I k w d && Y d)) = β₁ := by
      rw [← hY1]
      exact cnt_congr _ (fun x _ => by cases M.pstat P w x <;> cases M.I k w x <;> cases Y x <;> rfl)
    have c2 : cnt M.n (fun d => M.pstat P w d && !(M.I k w d && Y d)) = p - β₁ := by
      have h1 : cnt M.n (M.pstat P w) =
          cnt M.n (fun d => M.pstat P w d && (M.I k w d && Y d)) +
          cnt M.n (fun d => M.pstat P w d && !(M.I k w d && Y d)) :=
        cnt_split M.n (M.pstat P w) (fun d => M.I k w d && Y d)
      omega
    have c3 : cnt M.n (fun d => M.pstat P w d && Y d) = β₂ := by
      have h1 : cnt M.n (fun d => M.pstat P w d && Y d) =
          cnt M.n (fun d => (M.pstat P w d && Y d) && M.I k w d) +
          cnt M.n (fun d => (M.pstat P w d && Y d) && !M.I k w d) :=
        cnt_split M.n (fun d => M.pstat P w d && Y d) (M.I k w)
      have h2 : cnt M.n (fun d => (M.pstat P w d && Y d) && M.I k w d) = β₁ := by
        rw [← hY1]
        exact cnt_congr _ (fun x _ => by cases M.pstat P w x <;> cases M.I k w x <;> cases Y x <;> rfl)
      have h3 : cnt M.n (fun d => (M.pstat P w d && Y d) && !M.I k w d) = β₂ - β₁ := by
        rw [← hY2]
        exact cnt_congr _ (fun x _ => by cases M.pstat P w x <;> cases M.I k w x <;> cases Y x <;> rfl)
      omega
    have c4 : cnt M.n (fun d => M.pstat P w d && !Y d) = p - β₂ := by
      have h1 : cnt M.n (M.pstat P w) =
          cnt M.n (fun d => M.pstat P w d && Y d) +
          cnt M.n (fun d => M.pstat P w d && !Y d) := cnt_split M.n (M.pstat P w) Y
      omega
    rw [c1, c2, c3, c4] at key
    exact hfne key
  · exact fun h => st_of_ent M C q P k h

end Quantified
end AntiDyn