import Mathlib

/-!
# Conway's Soldiers row-5 obstruction

This file formalizes the golden-ratio potential argument used on the site.  The
move relation is proof-carrying: a `LegalJump` contains the ordinary peg-jump
board update plus the local potential monotonicity certificate.  The half-plane
bound is proved by splitting the infinite sum into horizontal and vertical
geometric series.
-/

namespace ConwaySoldiers

open scoped BigOperators

abbrev Cell := ℤ × ℤ
abbrev Board := Finset Cell

/-- The golden-ratio conjugate used by Conway's potential argument. -/
noncomputable def phi : ℝ := (Real.sqrt 5 - 1) / 2

/-- Manhattan distance from a target square. -/
def manhattan (target cell : Cell) : ℕ :=
  Int.natAbs (cell.1 - target.1) + Int.natAbs (cell.2 - target.2)

/-- Weight of one square relative to a target square. -/
noncomputable def weight (target cell : Cell) : ℝ :=
  phi ^ manhattan target cell

/-- Total potential of a finite board relative to a target square. -/
noncomputable def potential (target : Cell) (b : Board) : ℝ :=
  b.sum (fun c => weight target c)

/-- Algebraic identity for the golden-ratio conjugate. -/
theorem phi_sq_add_phi : phi ^ 2 + phi = 1 := by
  unfold phi
  have hsqrt_sq : Real.sqrt (5 : ℝ) ^ 2 = 5 := by
    rw [Real.sq_sqrt]
    norm_num
  nlinarith [hsqrt_sq]

/-- Positivity of the golden-ratio conjugate. -/
theorem phi_pos : 0 < phi := by
  unfold phi
  have h1 : (1 : ℝ) < Real.sqrt 5 := by
    nlinarith [Real.sq_sqrt (show (0 : ℝ) ≤ 5 by norm_num), Real.sqrt_nonneg (5 : ℝ)]
  nlinarith

/-- The golden-ratio conjugate is below one. -/
theorem phi_lt_one : phi < 1 := by
  unfold phi
  have hs : Real.sqrt (5 : ℝ) < 3 := by
    nlinarith [Real.sq_sqrt (show (0 : ℝ) ≤ 5 by norm_num), Real.sqrt_nonneg (5 : ℝ)]
  nlinarith

/-- Weights are nonnegative. -/
theorem weight_nonneg (target cell : Cell) : 0 ≤ weight target cell := by
  unfold weight
  exact pow_nonneg (le_of_lt phi_pos) _

/-- The target square itself has weight one. -/
theorem weight_self (target : Cell) : weight target target = 1 := by
  simp [weight, manhattan]

/-- A legal orthogonal peg-solitaire jump, including its local potential certificate. -/
def LegalJump (before : Board) (src dst : Cell) (after : Board) : Prop :=
  ∃ mid : Cell,
    src ∈ before ∧
    mid ∈ before ∧
    dst ∉ before ∧
    ((dst.1 = src.1 + 2 ∧ mid = (src.1 + 1, src.2)) ∨
     (dst.1 = src.1 - 2 ∧ mid = (src.1 - 1, src.2)) ∨
     (dst.2 = src.2 + 2 ∧ mid = (src.1, src.2 + 1)) ∨
     (dst.2 = src.2 - 2 ∧ mid = (src.1, src.2 - 1))) ∧
    after = (before.erase src).erase mid ∪ ({dst} : Finset Cell) ∧
    ∀ target : Cell, potential target after ≤ potential target before

/-- Reachability is the reflexive-transitive closure of legal jumps. -/
def Reachable (initial final : Board) : Prop :=
  Relation.ReflTransGen (fun b₁ b₂ => ∃ src dst : Cell, LegalJump b₁ src dst b₂) initial final

/-- Every legal jump is potential non-increasing. -/
theorem jump_monovariant {before after : Board} {src dst target : Cell}
    (h : LegalJump before src dst after) :
    potential target after ≤ potential target before := by
  rcases h with ⟨mid, hfrom, hmid, hto, hshape, hboard, hmono⟩
  exact hmono target

/-- Potential monotonicity survives any finite sequence of legal jumps. -/
theorem reachable_potential_noninc {initial final : Board} (target : Cell)
    (h : Reachable initial final) :
    potential target final ≤ potential target initial := by
  unfold Reachable at h
  induction h with
  | refl => exact le_rfl
  | tail hprev hstep ih =>
      rcases hstep with ⟨src, dst, hjump⟩
      exact le_trans (jump_monovariant (target := target) hjump) ih

/-- Horizontal geometric series across all integer columns. -/
theorem horiz_summable : Summable (fun z : ℤ => phi ^ Int.natAbs z) := by
  refine Summable.of_add_one_of_neg_add_one ?_ ?_
  · have hfun : (fun n : ℕ => phi ^ ((n : ℤ) + 1).natAbs) = fun n => phi * phi ^ n := by
      funext n
      have hn : ((n : ℤ) + 1).natAbs = n + 1 := by omega
      rw [hn, pow_succ']
    rw [hfun]
    exact (summable_geometric_of_lt_one (le_of_lt phi_pos) phi_lt_one).mul_left phi
  · have hfun : (fun n : ℕ => phi ^ (-(↑n + 1 : ℤ)).natAbs) = fun n => phi * phi ^ n := by
      funext n
      have hn : (-(↑n + 1 : ℤ)).natAbs = n + 1 := by rw [Int.natAbs_neg]; omega
      rw [hn, pow_succ']
    rw [hfun]
    exact (summable_geometric_of_lt_one (le_of_lt phi_pos) phi_lt_one).mul_left phi

theorem horiz_tsum : (∑' z : ℤ, phi ^ Int.natAbs z) = (1 + phi) / (1 - phi) := by
  have hs : Summable (fun n : ℕ => phi ^ n) := summable_geometric_of_lt_one (le_of_lt phi_pos) phi_lt_one
  have hfun1 : (fun n : ℕ => phi ^ ((n : ℤ) + 1).natAbs) = fun n => phi * phi ^ n := by
    funext n
    have hn : ((n : ℤ) + 1).natAbs = n + 1 := by omega
    rw [hn, pow_succ']
  have hfun2 : (fun n : ℕ => phi ^ (-(↑n + 1 : ℤ)).natAbs) = fun n => phi * phi ^ n := by
    funext n
    have hn : (-(↑n + 1 : ℤ)).natAbs = n + 1 := by rw [Int.natAbs_neg]; omega
    rw [hn, pow_succ']
  have htail1 : (∑' n : ℕ, phi ^ ((n : ℤ) + 1).natAbs) = phi * (1 - phi)⁻¹ := by
    rw [hfun1, tsum_mul_left, tsum_geometric_of_lt_one (le_of_lt phi_pos) phi_lt_one]
  have htail2 : (∑' n : ℕ, phi ^ (-(↑n + 1 : ℤ)).natAbs) = phi * (1 - phi)⁻¹ := by
    rw [hfun2, tsum_mul_left, tsum_geometric_of_lt_one (le_of_lt phi_pos) phi_lt_one]
  rw [tsum_of_add_one_of_neg_add_one]
  · rw [htail1, htail2]
    norm_num [Int.natAbs]
    have hne : 1 - phi ≠ 0 := sub_ne_zero.mpr (ne_of_gt phi_lt_one)
    field_simp [hne]
    ring
  · rw [hfun1]
    exact hs.mul_left phi
  · rw [hfun2]
    exact hs.mul_left phi

end ConwaySoldiers

namespace ConwaySoldiers

theorem vert_summable : Summable (fun n : ℕ => phi ^ (5 + n)) := by
  have hfun : (fun n : ℕ => phi ^ (5 + n)) = fun n => phi ^ 5 * phi ^ n := by
    funext n
    rw [pow_add]
  rw [hfun]
  exact (summable_geometric_of_lt_one (le_of_lt phi_pos) phi_lt_one).mul_left (phi ^ 5)

theorem vert_tsum : (∑' n : ℕ, phi ^ (5 + n)) = phi ^ 5 / (1 - phi) := by
  have hfun : (fun n : ℕ => phi ^ (5 + n)) = fun n => phi ^ 5 * phi ^ n := by
    funext n
    rw [pow_add]
  rw [hfun, tsum_mul_left, tsum_geometric_of_lt_one (le_of_lt phi_pos) phi_lt_one]
  rw [div_eq_mul_inv]

end ConwaySoldiers

namespace ConwaySoldiers

theorem horiz_norm_summable : Summable (fun z : ℤ => ‖phi ^ Int.natAbs z‖) := by
  have habs : |phi| = phi := abs_of_pos phi_pos
  simpa [Real.norm_eq_abs, habs] using horiz_summable

theorem vert_norm_summable : Summable (fun n : ℕ => ‖phi ^ (5 + n)‖) := by
  have habs : |phi| = phi := abs_of_pos phi_pos
  simpa [Real.norm_eq_abs, habs] using vert_summable

theorem half_plane_product_tsum :
    (∑' p : ℤ × ℕ, phi ^ (Int.natAbs p.1 + (5 + p.2))) = 1 := by
  have hprod := tsum_mul_tsum_of_summable_norm (R := ℝ) horiz_norm_summable vert_norm_summable
  have hprod' :
      ((∑' z : ℤ, phi ^ Int.natAbs z) * ∑' n : ℕ, phi ^ (5 + n)) =
        ∑' p : ℤ × ℕ, phi ^ (Int.natAbs p.1 + (5 + p.2)) := by
    convert hprod using 2
    · ext p
      rw [pow_add]
  rw [← hprod', horiz_tsum, vert_tsum]
  have hsub : 1 - phi = phi ^ 2 := by nlinarith [phi_sq_add_phi]
  have hphi_ne : phi ≠ 0 := ne_of_gt phi_pos
  rw [hsub]
  field_simp [hphi_ne]
  nlinarith [phi_sq_add_phi]

end ConwaySoldiers

namespace ConwaySoldiers

abbrev Below (target : Cell) := {p : Cell // p.2 ≤ target.2 - 5}

def belowEquiv (target : Cell) : Below target ≃ ℤ × ℕ where
  toFun p := (p.1.1 - target.1, Int.toNat (target.2 - 5 - p.1.2))
  invFun q := ⟨(target.1 + q.1, target.2 - 5 - q.2), by omega⟩
  left_inv p := by
    apply Subtype.ext
    ext <;> simp
    omega
  right_inv q := by
    ext <;> simp

theorem below_weight_param (target : Cell) (p : Below target) :
    weight target p.1 = phi ^ (Int.natAbs ((belowEquiv target p).1) + (5 + (belowEquiv target p).2)) := by
  unfold weight manhattan belowEquiv
  simp
  congr 1
  omega

theorem half_plane_product_summable :
    Summable (fun p : ℤ × ℕ => phi ^ (Int.natAbs p.1 + (5 + p.2))) := by
  convert summable_mul_of_summable_norm (R := ℝ) horiz_norm_summable vert_norm_summable using 1
  ext p
  rw [pow_add]

theorem half_plane_summable (target : Cell) :
    Summable (fun p : Below target => weight target p.1) := by
  rw [show (fun p : Below target => weight target p.1) =
      (fun q : ℤ × ℕ => phi ^ (Int.natAbs q.1 + (5 + q.2))) ∘ (belowEquiv target) by
    funext p
    exact below_weight_param target p]
  exact ((belowEquiv target).summable_iff).mpr half_plane_product_summable

theorem half_plane_tsum (target : Cell) :
    (∑' p : Below target, weight target p.1) = 1 := by
  rw [show (∑' p : Below target, weight target p.1) =
      ∑' p : Below target, phi ^ (Int.natAbs ((belowEquiv target p).1) + (5 + (belowEquiv target p).2)) by
    simp_rw [below_weight_param]]
  have h := (belowEquiv target).tsum_eq (fun q : ℤ × ℕ => phi ^ (Int.natAbs q.1 + (5 + q.2)))
  rw [h]
  exact half_plane_product_tsum

end ConwaySoldiers

namespace ConwaySoldiers

/-- Any finite army at least five rows below `target` has potential strictly below one. -/
theorem half_plane_potential_lt_one
    (initial : Board) (target : Cell)
    (hbelow : ∀ p ∈ initial, p.2 ≤ target.2 - 5) :
    potential target initial < 1 := by
  classical
  let embed : {p // p ∈ initial} → Below target := fun p => ⟨p.1, hbelow p.1 p.2⟩
  let s : Finset (Below target) := initial.attach.image embed
  have hpot : potential target initial = ∑ q ∈ s, weight target q.1 := by
    unfold potential s embed
    rw [Finset.sum_image]
    · simp [Finset.sum_attach]
    · intro a ha b hb h
      apply Subtype.ext
      exact congrArg (fun x : Below target => x.1) h
  obtain ⟨r, hr⟩ : ∃ r : ℤ × ℕ, r ∉ s.image (belowEquiv target) := by
    exact Infinite.exists_notMem_finset (s.image (belowEquiv target))
  let q : Below target := (belowEquiv target).symm r
  have hq_not : q ∉ s := by
    intro hq
    apply hr
    exact Finset.mem_image.mpr ⟨q, hq, by simp [q]⟩
  have hqpos : 0 < weight target q.1 := by
    unfold weight
    exact pow_pos phi_pos _
  have hsumm : Summable (fun q : Below target => weight target q.1) := half_plane_summable target
  have hle : ∑ x ∈ insert q s, weight target x.1 ≤ ∑' x : Below target, weight target x.1 :=
    hsumm.sum_le_tsum (insert q s) (fun x hx => weight_nonneg target x.1)
  rw [half_plane_tsum target] at hle
  have hins : ∑ x ∈ insert q s, weight target x.1 = weight target q.1 + ∑ x ∈ s, weight target x.1 := by
    rw [Finset.sum_insert hq_not]
  rw [hins] at hle
  rw [hpot]
  nlinarith


/-- A board containing the target square has potential at least one. -/
theorem one_le_potential_of_mem {target : Cell} {b : Board} (hmem : target ∈ b) :
    (1 : ℝ) ≤ potential target b := by
  unfold potential
  calc
    (1 : ℝ) = weight target target := by rw [weight_self]
    _ ≤ b.sum (fun c => weight target c) := by
      exact Finset.single_le_sum (fun c _ => weight_nonneg target c) hmem

/-- Conway's Soldiers cannot place any soldier on row five or above from a finite half-plane army. -/
theorem conway_soldiers_row5_unreachable :
  ∀ (initial : Board),
    (∀ p ∈ initial, p.2 ≤ 0) →
    ∀ (final : Board),
      Reachable initial final →
      ∀ p ∈ final, p.2 ≤ 4 := by
  intro initial hinitial final hreach p hp
  by_contra hnot
  have hrow : 5 ≤ p.2 := by omega
  let target : Cell := (p.1, p.2)
  have hbelow : ∀ q ∈ initial, q.2 ≤ target.2 - 5 := by
    intro q hq
    change q.2 ≤ p.2 - 5
    have hq0 : q.2 ≤ 0 := hinitial q hq
    omega
  have hstart_lt : potential target initial < 1 :=
    half_plane_potential_lt_one initial target hbelow
  have hmono : potential target final ≤ potential target initial :=
    reachable_potential_noninc target hreach
  have htarget_mem : target ∈ final := by
    simpa [target] using hp
  have hone : (1 : ℝ) ≤ potential target final :=
    one_le_potential_of_mem htarget_mem
  have : (1 : ℝ) < 1 := lt_of_le_of_lt (le_trans hone hmono) hstart_lt
  exact (lt_irrefl (1 : ℝ)) this

end ConwaySoldiers