import Mathlib.Tactic
import Mathlib
/-!
# USA Mathematical Olympiad 2023, Problem 5
Let n be an integer greater than 2. We will be arranging the numbers
1, 2, ... n² into an n × n grid. Such an arrangement is called *row-valid*
if the numbers in each row can be permuted to make an arithmetic progression.
Similarly, such an arrangement is called *column-valid* if the numbers
in each column can be permuted to make an arithmetic progression.
Determine the values of n for which it possible to transform
any row-valid arrangement into a column-valid arrangement by permuting
the numbers in each row.
-/
namespace Usa2023P5
def PermutedArithSeq {n : ℕ} (hn : 0 < n) (a : Fin n ↪ Fin (n^2)) : Prop :=
∃ p : Fin n → Fin n, p.Bijective ∧
∃ k : ℕ, ∀ m : Fin n, (a (p m)).val = a (p ⟨0, hn⟩) + m.val * k
def row_valid {n : ℕ} (hn : 0 < n) (a : Fin n → Fin n → Fin (n^2)) (ha : a.Injective2) : Prop :=
∀ r : Fin n, PermutedArithSeq hn ⟨(a r ·), Function.Injective2.right ha r⟩
def col_valid {n : ℕ} (hn : 0 < n) (a : Fin n → Fin n → Fin (n^2)) (ha : a.Injective2) : Prop :=
∀ c : Fin n, PermutedArithSeq hn ⟨(a · c), Function.Injective2.left ha c⟩
theorem injective_of_permuted_rows {α β γ : Type}
{f : α → β → γ} (hf : f.Injective2) {p : α → β → β} (hp : ∀ a, (p a).Injective) :
Function.Injective2 (fun r c ↦ f r (p r c)) := by
intro a1 a2 b1 b2 hab
obtain ⟨ha1, hp1⟩ := hf hab
rw [ha1] at *
rw [hp a2 hp1]
simp only [and_self]
/- determine -/ abbrev solution_set : Set ℕ := sorry
theorem usa2023_p5 (n : ℕ) (hn : 2 < n) :
n ∈ solution_set ↔
(∀ a : (Fin n → Fin n → Fin (n^2)),
(ha : a.Injective2) → row_valid (Nat.zero_lt_of_lt hn) a ha →
∃ p : Fin n → Fin n → Fin n, ∃ hp : (∀ r, (p r).Injective),
col_valid (Nat.zero_lt_of_lt hn) (fun r c ↦ a r (p r c))
(injective_of_permuted_rows ha hp)) := sorry
end Usa2023P5
This problem does not yet have a complete formalized solution.