My solution to IMO 1 2024

imo1_2024.lean

import Mathlib.Data.Real.Basic
import Mathlib.Data.Real.Archimedean
import Mathlib.Algebra.BigOperators.Group.Finset.Basic
import Mathlib.Algebra.BigOperators.Ring.Finset
import Mathlib.Algebra.BigOperators.Intervals
open Finset (range)

-- STATEMENT

def IsSpecial (x : ℝ) :=
    ∀ (n : ℕ), ↑n ∣ ∑ i ∈ range (n+1), ⌊ i * x ⌋

def Problem := DecidablePred IsSpecial

-- SOLUTION

theorem cast_toNat {z : ℤ} :
    0 ≤ z → ↑z.toNat = (↑z : ℝ) := by norm_cast; simp

theorem arithDiv (k : ℤ) : 2 ∣ (k+1) * k :=
    Even.two_dvd $ match Int.even_or_odd k with
    | .inl h => h.mul_left _
    | .inr h => h.add_one.mul_right _

theorem le_self_mul_lt1 {x y : ℝ} :
    y ≤ 1 → 0 ≤ x → x * y ≤ x := by
    convert IsOrderedRing.mul_le_mul_of_nonneg_left y 1 x
    simp

theorem IntSpecial.simplify_sum {n} {k : ℤ} :
    (∑ i ∈ range (n+1), ⌊ ↑i * (↑k : ℝ) ⌋) = ((n+1) * n / 2) * k := by
    conv => enter [1, 2, i]; norm_cast; rw [Int.floor_intCast]
    rw [←Finset.sum_mul, ←Nat.cast_sum, Finset.sum_range_id]
    simp

theorem IntSpecial (k : ℤ) : Even k ↔ IsSpecial ↑k := by
    constructor <;> intro
    case mp h =>
        intro n
        rw [IntSpecial.simplify_sum]
        rw [←Int.mul_ediv_assoc' _ (arithDiv n),
            Int.mul_ediv_assoc _ h.two_dvd]
        apply Int.dvd_trans _ (Int.dvd_mul_right _ _)
        apply Int.dvd_mul_left
    case mpr h =>
        replace h := even_iff_two_dvd.mpr (h 2)
        contrapose h
        rw [IntSpecial.simplify_sum]
        apply Int.not_even_iff_odd.mpr
        apply Odd.mul
        . rw [Int.mul_ediv_assoc]
          apply Odd.mul
          repeat simp [Int.odd_iff.mpr]
        . apply Int.not_even_iff_odd.mp h

theorem SpecialPlusEven.fwd {x : ℝ} {k : ℤ} (k_even : Even k) :
    IsSpecial x → IsSpecial (x + k) := by
    intros x_special n
    conv =>
        enter [2, 2, i]
        rw [left_distrib]
        norm_cast; rw [Int.floor_add_intCast]
    rw [Finset.sum_add_distrib]
    apply Int.dvd_add
    . exact x_special n
    . convert (IntSpecial k).mp k_even n
      norm_cast; rw [Int.floor_intCast]

theorem SpecialPlusEven {x : ℝ} {k : ℤ} (k_even : Even k) :
    IsSpecial x ↔ IsSpecial (x + k) := by
    constructor
    exact SpecialPlusEven.fwd k_even
    convert SpecialPlusEven.fwd (Even.neg k_even)
    simp

theorem PosLtOneNotSpecial x (x_ioo : x ∈ Set.Ioo 0 1) : ¬(IsSpecial x) := by
    apply Set.mem_setOf.mp at x_ioo
    have ceil_pos : 0 ≤ ⌈x⁻¹⌉ := by simp [Int.le_of_lt, *]
    have x_cancel : 1 = x⁻¹ * x := by
      rw [mul_comm, Field.mul_inv_cancel x]
      exact ne_of_gt x_ioo.left
    have term_pos : ⌊⌈x⁻¹⌉ * x⌋ > 0 := by
      simp [Int.lt_floor_iff.mpr, *, Int.le_ceil x⁻¹]

    apply mt (· ⌈x⁻¹⌉.toNat)
    rw [Finset.sum_range_succ, cast_toNat ceil_pos]
    rw [Finset.sum_eq_zero, Int.zero_add]
    . apply mt Int.eq_zero_of_dvd_of_natAbs_lt_natAbs
      simp [x_ioo, x_cancel, Int.le_ceil]
      apply Int.natAbs_lt_natAbs_of_nonneg_of_lt
      exact le_of_lt term_pos
      apply (@Int.cast_lt ℝ).mp
      apply lt_of_le_of_lt (Int.floor_le _)
      simp [x_ioo]
    . intro i i_range
      simp [Finset.mem_range.mp] at i_range
      apply Int.lt_ceil.mp at i_range
      simp at i_range
      simp [x_ioo, x_cancel, i_range]

theorem NegLtOneNotSpecial (x : ℝ) (x_ioo : x ∈ Set.Ioo 0 1) : ¬(IsSpecial (-x)) := by
    apply Set.mem_setOf.mp at x_ioo
    have floor_pos : 0 ≤ ⌊x⁻¹⌋ := by
      simp [Int.le_floor.mpr, le_of_lt, *]
    have floor_succ_pos : (0 : ℝ) < ↑⌊x⁻¹⌋ + 1 := by
      norm_cast
      apply lt_of_lt_of_le (by simp) (add_right_mono floor_pos)
    have x_cancel : 1 = x⁻¹ * x := by
      rw [mul_comm, DivisionRing.mul_inv_cancel x]
      exact ne_of_gt x_ioo.left

    apply mt (· (⌊x⁻¹⌋.toNat + 1))

    -- cut off the i = 0 term
    rw [Finset.sum_range_succ']
    simp [-Int.ofNat_toNat]

    -- introduce a +1 on every term, and the other -1 is taken
    -- out of the sum and disappears since it's a multiple
    conv => enter [1, 2, 2, i]; apply (add_neg_cancel_right _ 1).symm
    rw [Finset.sum_add_distrib, Finset.sum_const, Finset.card_range]
    apply Not.imp _ Int.dvd_add_self_mul.mp

    rw [Finset.sum_range_succ, cast_toNat floor_pos]
    rw [Finset.sum_eq_zero, Int.zero_add]
    . rw [Int.floor_neg]
      apply mt Int.dvd_neg.mpr
      apply mt Int.eq_zero_of_dvd_of_natAbs_lt_natAbs
      simp [x_ioo, floor_pos]
      constructor
      . apply Int.natAbs_lt_natAbs_of_nonneg_of_lt <;> simp [x_ioo, floor_succ_pos]
        apply Int.ceil_lt_iff.mpr
        simp
        apply le_self_mul_lt1 (le_of_lt x_ioo.right) (le_of_lt floor_succ_pos)
      . apply ne_of_gt
        simp
        apply Int.lt_ceil.mpr
        norm_cast
        rw [x_cancel]
        simp [x_ioo]
    . intro i i_range
      simp [Finset.mem_range.mp] at i_range
      apply Int.lt_floor_iff.mp at i_range
      norm_cast at *
      generalize h : i + 1 = i at *
      rw [(Int.floor_eq_iff (z := -1)).mpr] <;> simp
      simp [x_ioo, x_cancel, i_range]
      apply (Nat.zero_lt_succ _).trans_eq h

theorem Ico02Special.NegIoo10 x (x_ioo : x ∈ Set.Ioo (-1) 0) : ¬IsSpecial x := by
    convert NegLtOneNotSpecial (x := -x) _ <;> simp [x_ioo.out, neg_lt.mp]

noncomputable def Ico02Special x (x_range : x ∈ Set.Ico 0 2) :
    Decidable (IsSpecial x) := by
    by_cases 0 = x
    case pos h =>
        convert Decidable.isTrue ((IntSpecial 0).mp _) <;> simp [h]
    replace x_range : 0 < x ∧ x < 2 := by
        simp [lt_of_le_of_ne, x_range.out, *]
    apply Decidable.isFalse

    by_cases 1 < x
    case pos x_gt1 =>
        apply mt (SpecialPlusEven (k := -2) _).mp <;> simp
        apply Ico02Special.NegIoo10
        have : 2 < 1 + x := by
            convert add_left_strictMono (a := 1) x_gt1; ring
        simp [x_range, this]
    replace x_range : 0 < x ∧ x ≤ 1 := by
        simp [le_of_not_gt, *]

    by_cases x = 1
    case pos h =>
        convert (IntSpecial 1).not.mp _ <;> simp [h]
    replace x_range : 0 < x ∧ x < 1 := by
        simp [lt_of_le_of_ne, *]

    convert PosLtOneNotSpecial _ _; simp [x_range]

noncomputable def RealSpecial : Problem := by
    intro x
    let k := 2 * -⌊x / 2⌋
    have k_even : Even k := Even.two_nsmul _
    suffices h : Decidable (IsSpecial (x + ↑k)) from
        decidable_of_iff _ (SpecialPlusEven k_even).symm
    generalize x_lt2 : x + ↑k = x'
    replace x_lt2 : x' = 2 * Int.fract (x/2) := by
        simp [←x_lt2, k, -Int.self_sub_floor, Int.fract]
        ring
    apply Ico02Special
    simp [x_lt2, Int.fract_lt_one]

-- many of the convs above use enter and not 'in ⌊_⌋' because that
-- spits out a weird error, seems like it has trouble finding things
-- inside of the sums, perhaps due to the macro?