Partial proof of https://math.stackexchange.com/a/4894379/259262

imo-shortlist.lean

-- https://math.stackexchange.com/a/4894379/259262

import Mathlib.Data.Real.Sqrt
import Mathlib.Tactic.Hint
import Mathlib.Tactic.Cases
import Mathlib.Algebra.Order.Floor

open Real Set

namespace Problem

def Property (f : ℝ → ℝ) : Prop := ∀ (x : ℝ), ∀ (y : ℝ), (f (x^2 + y^2)) + (f (x * f y + f x * f y)) = (f (x+y))^2

theorem abs_ne_zero_of_ne_zero {x : ℝ} (ne_zero : x ≠ 0) : |x| ≠ 0 := by
  by_cases h : 0 ≤ x
  . rw [abs_of_nonneg h]
    exact ne_zero
  . rw [abs_of_neg (by linarith)]
    linarith

theorem ne_zero_of_abs_ne_zero {x : ℝ} (abs_ne_zero : |x| ≠ 0) : x ≠ 0 := by
  by_cases h : 0 ≤ x
  . rw [abs_of_nonneg h] at abs_ne_zero
    exact abs_ne_zero
  . rw [abs_of_neg (by linarith)] at abs_ne_zero
    linarith

theorem lt_of_le {x y : ℝ} (le : x ≤ y) (ne : x ≠ y) : x < y := by
  rcases lt_trichotomy x y with (x_lt | (eq | x_gt))
  . linarith
  . tauto
  . linarith

theorem ZeroRootImpliesSqCommutes {f} (pr : Property f) (f_zero : f 0 = 0) (y : ℝ) : f (y ^ 2) = (f y) ^ 2 := by
    let x := pr 0
    rw [f_zero] at x
    simp at x
    rw [f_zero] at x
    simp at x
    exact x y

theorem ZeroRootImpliesFOne {f} (pr : Property f) (f_zero : f 0 = 0) : f 1 = 1 ∨ f 1 = 0 := by
    have x := ZeroRootImpliesSqCommutes pr f_zero 1
    simp only [one_pow] at x
    have : f 1 * (f 1 - 1) = 0 := by linarith
    have nearly : f 1 = 0 ∨ f 1 - 1 = 0 := by
      rw [mul_eq_zero] at this
      exact this
    cases' nearly with p1 p2
    . tauto
    . left; linarith

theorem Commutation {f} (pr : Property f) (x : ℝ) (y : ℝ) : f (x * f y + f x * f y) = f (y * f x + f x * f y) := by
    have r := pr x y
    have s := pr y x
    rw [add_comm y x, add_comm (y ^ 2) (x ^ 2), mul_comm (f y) (f x)] at s
    linear_combination r - s

theorem ClassifyIfInjective {f} (pr : Property f) (inj : ∀ x, ∀ y, f x = f y → x = y) : ∀ x, f x = x := by
  have: ∀ x, ∀ y, x * f y = y * f x := by 
    intros x y
    linarith [inj _ _ (Commutation pr x y)]

  have is_linear: ∀ y, y * f 1 = f y := by
    intros y
    linarith [this 1 y]

  have f_zero: f 0 = 0 := by linarith [is_linear 0]

  have: f 1 = 1 ∨ f 1 = 0 := ZeroRootImpliesFOne pr f_zero
  have: f 1 ≠ 0 := by
    intros bad
    have: 1 = 0 := inj 1 0 (by linarith)
    linarith

  have coeff: f 1 = 1 := by tauto

  intros y
  have := is_linear y
  rw [coeff] at this
  ring_nf at this
  tauto

theorem ClassifyIfZeroIsRoot {f} (pr : Property f) (f_zero_eq_zero : f 0 = 0) : (∀ x, f x = 0) ∨ ∀ x, f x = x := by
  have eqn_1 : ∀ (y : ℝ), f (y ^ 2) = (f (-y)) ^ 2 := by
    intros y
    have s := ZeroRootImpliesSqCommutes pr f_zero_eq_zero (-y)
    simp only [even_two, Even.neg_pow] at s
    tauto

  have eqn_1' : ∀ (y: ℝ), f y = f (-y) ∨ f y = - f (-y) := by
    intros y
    have : (f (-y)) ^ 2 = (f y) ^ 2 :=
      calc (f (-y)) ^ 2 = f (y ^ 2) := by rw [eqn_1 y]
      _ = f ((-y) ^ 2) := by ring_nf
      _ = f (-(-y)) ^ 2 := by rw [eqn_1 (-y)]
      _ = (f y) ^ 2 := by ring_nf
    exact sq_eq_sq_iff_eq_or_eq_neg.mp (id this.symm)

  have pos_for_pos : ∀ (y : ℝ), 0 ≤ y → 0 ≤ f y := by
    intros y y_pos
    let x := ZeroRootImpliesSqCommutes pr f_zero_eq_zero (Real.sqrt y)
    rw [sq_sqrt y_pos] at x
    rw [x]
    exact sq_nonneg (f (sqrt y))

  have pre_2 : ∀ (x : ℝ), ∀ (y : ℝ), 0 ≤ x -> 0 ≤ y -> 0 ≤ f (x * f y + f x * f y) := by
    intros x y x_pos y_pos
    apply pos_for_pos
    have : 0 ≤ f x := pos_for_pos x x_pos
    have : 0 ≤ f y := pos_for_pos y y_pos
    positivity

  have eqn_2 : ∀ (x : ℝ), ∀ (y : ℝ), f (x * f y + f x * f y) - f (x * f (-y) + f x * f (-y)) = (f (x + y)) ^ 2 - (f (x + (-y))) ^ 2 := by
    intros x y
    let e1 := pr x y
    let e2 := pr x (-y)
    simp at e1
    simp at e2
    linarith

  have increasing_lemma : ∀ (x : ℝ), ∀ (y : ℝ), 0 ≤ x → 0 ≤ y → 0 ≤ (f (x + y)) ^ 2 - (f (x + (-y))) ^ 2 := by 
    intros x y x_pos y_pos
    rw [<-eqn_2 x y]
    cases' (eqn_1' y) with is_id switch
    . rw [←is_id]
      simp
    . have switch : f (-y) = - (f y) := by linarith
      rw [switch]
      have thing : -(x * f y + f x * f y) = x * -f y + f x * -(f y) := by ring
      suffices hyp : 0 ≤ f (x * f y + f x * f y) - f (-(x * f y + f x * f y))
        by rw [<-thing]; tauto
      cases' (eqn_1' (x * f y + f x * f y)) with is_id switch
      . rw [is_id]; simp
      . have thing : f (x * f y + f x * f y) - f (-(x * f y + f x * f y)) = f (x * f y + f x * f y) + (-f (-(x * f y + f x * f y))) := by ring
        rw [thing]
        rw [<-switch]
        simp
        tauto

  have increasing : ∀ (u : ℝ), ∀ (v : ℝ), 0 ≤ u → u ≤ v → f u ≤ f v := by
    intros u v u_pos u_le_v
    have v_pos : 0 ≤ v := by linarith
    have thing : (f u) ^ 2 ≤ (f v) ^ 2 := by
      have thing := increasing_lemma ((v + u) / 2) ((v - u) / 2) (by positivity) (by linarith)
      simp at thing
      calc f u ^ 2 = f ((v + u) / 2 + -((v - u) / 2)) ^ 2 := by ring_nf
      _ ≤ f ((v + u) / 2 + (v - u) / 2) ^ 2 := thing
      _ = f v ^ 2 := by ring_nf
    have pos_u : 0 ≤ f u := pos_for_pos u u_pos
    have pos_v : 0 ≤ f v := pos_for_pos v v_pos
    rw [sq_le_sq] at thing
    rw [abs_of_nonneg pos_u, abs_of_nonneg pos_v] at thing
    exact thing

  by_cases has_root : ∃ a, a ≠ 0 ∧ f a = 0 -- with f_one_eq_one f_one_eq_zero
  . left
    rcases has_root with ⟨a, ⟨a_ne_zero, a_root⟩⟩
    have mod_a_root : f |a| = 0 := by
      by_cases a_sign : 0 ≤ a
      . rw [abs_of_nonneg a_sign]
        exact a_root
      . rw [abs_of_neg (by linarith)]
        cases' eqn_1' |a| with p1 p2
        . ring_nf at p1
          rw [abs_of_neg (by linarith)] at p1
          ring_nf at p1
          rw [p1, a_root]
        . rw [abs_of_neg (by linarith)] at p2
          ring_nf at p2
          rw [p2, a_root]
          ring_nf

    have : ∀ (n : ℕ), f (2 * n * |a|) = 0 := by
      intros n
      induction' n with n
      . simp [f_zero_eq_zero]
      . have ind_step : ∀ x, f (x + |a|) ^ 2 = f (x + (-|a|)) ^ 2 := by
          intros x
          have whatnot := eqn_2 x |a|
          have f_minus_a : f (-|a|) = 0 := by
            cases' eqn_1' |a| with hyp hyp
            . linarith
            . linarith
          simp [mod_a_root, f_zero_eq_zero, f_minus_a] at whatnot
          linarith
        have ind_step' : ∀ (n : ℕ), f (2 * n * |a|) ^ 2 = f ((2 * (n - 1)) * |a|) ^ 2 := by
          intros n
          have answer := ind_step ((2 * n - 1) * |a|)
          ring_nf at answer
          ring_nf
          exact answer
        have : ∀ (n : ℕ), f (2 * n * |a|) ^ 2 = 0 := by
          intros n
          induction' n with n ind_hyp
          . simp; tauto
          . calc f (2 * (Nat.succ n) * |a|) ^ 2 = f ((2 * ((Nat.succ n) - 1)) * |a|) ^ 2 := ind_step' (Nat.succ n)
            _ = f (2 * n * |a|) ^ 2 := by simp
            _ = 0 := ind_hyp
        have : ∀ (n : ℕ), f (2 * n * |a|) = 0 := by
          intros n
          exact sq_eq_zero_iff.1 (this n)
        exact this (Nat.succ n)

    have zero_for_pos : ∀ (x : ℝ), 0 ≤ x → f x = 0 := by
      intros x x_pos
      let ceil := ⌈x / (2 * |a|)⌉
      have mod_a_nonzero : |a| ≠ 0 := abs_ne_zero_of_ne_zero a_ne_zero
      have mod_a_pos : 0 < |a| := by linarith [abs_pos.2 a_ne_zero]
      have two_a_big : 0 < 2 * |a| := by linarith
      have ceil_big : x / (2 * |a|) ≤ ceil := Int.le_ceil _
      have ceil_pos : 0 ≤ ceil := by positivity
      have r : x ≤ 2 * (Int.toNat ceil) * |a| := by
        calc x = (x / (2 * |a|)) * (2 * |a|) := by field_simp
          _ ≤ ceil * (2 * |a|) := (mul_le_mul_right two_a_big).2 ceil_big
          _ = 2 * ceil * |a| := by ring_nf
          _ = 2 * (Int.toNat ceil) * |a| := by apply congrFun (congrArg HMul.hMul (congrArg (HMul.hMul _) _)) |a|; norm_cast; simp_all only [sub_nonneg, ne_eq, abs_eq_zero, not_false_eq_true, abs_pos, gt_iff_lt, Nat.ofNat_pos, mul_pos_iff_of_pos_left, Int.toNat_of_nonneg, ceil]
      have t : f x ≤ f (2 * (Int.toNat ceil) * |a|) := increasing _ _ x_pos r
      rw [this (Int.toNat ceil)] at t
      have u : 0 ≤ f x := by
        have answer := increasing 0 x (by linarith) x_pos
        rw [f_zero_eq_zero] at answer
        exact answer
      linarith
    intros x
    have f_abs_eq_zero : f |x| = 0 := zero_for_pos |x| (abs_nonneg _)
    by_cases h : 0 ≤ x 
    . simp [abs_of_nonneg h] at f_abs_eq_zero; exact f_abs_eq_zero
    . have abs_x_is_neg : |x| = -x := abs_of_neg (by linarith)
      cases' eqn_1' |x| with p1 p2
      . rw [abs_x_is_neg] at p1
        ring_nf at p1
        rw [abs_x_is_neg] at f_abs_eq_zero
        rw [f_abs_eq_zero] at p1
        tauto
      . rw [f_abs_eq_zero, abs_x_is_neg] at p2
        simp at p2
        tauto

  . right
    have eqn_1'' : ∀ x, f x = - f (-x) := by
      intros x
      by_cases x_zero : x = 0
      . rw [x_zero]; simpa
      .
        cases' eqn_1' x with even odd
        . have eq2 := eqn_2 |x| x
          by_cases h : 0 ≤ x
          . rw [abs_of_nonneg h, <-even] at eq2
            simp only [sub_self, add_right_neg] at eq2
            rw [f_zero_eq_zero] at eq2
            simp only [ne_eq, OfNat.ofNat_ne_zero, not_false_eq_true, zero_pow, sub_zero] at eq2
            have: f (x + x) = 0 := sq_eq_zero_iff.1 eq2.symm
            have: x + x ≠ 0 := by ring_nf; exact mul_ne_zero x_zero (by norm_num)
            have : False := has_root ⟨x + x, by tauto⟩
            tauto
          . have h_neg : x < 0 := by linarith
            rw [abs_of_neg h_neg, <-even] at eq2
            simp only [neg_mul, sub_self, add_left_neg] at eq2
            rw [f_zero_eq_zero] at eq2
            simp only [ne_eq, OfNat.ofNat_ne_zero, not_false_eq_true, zero_pow, zero_sub, zero_eq_neg, pow_eq_zero_iff] at eq2
            have: -x + -x ≠ 0 := by ring_nf; linarith
            have : False := has_root ⟨-x + -x, by tauto⟩
            tauto
        . exact odd

    have eqn_2' : ∀ (x : ℝ), ∀ (y : ℝ), 2 * f (x * f y + f x * f y) = (f (x + y)) ^ 2 - (f (x + (-y))) ^ 2 := by
      intros x y
      calc
        2 * f (x * f y + f x * f y) = f (x * f y + f x * f y) - (-(f (x * f y + f x * f y))) := by ring_nf
        _ = f (x * f y + f x * f y) - (f (-(x * f y + f x * f y))) := by rw [eqn_1'' (x * f y + f x * f y)]; ring_nf
        _ = f (x * f y + f x * f y) - (f (x * (-f y) + f x * (-f y))) := by ring_nf
        _ = f (x * f y + f x * f y) - (f (x * f (-y) + f x * f (-y))) := by rw [eqn_1'' y]; ring_nf
        _ = _ := eqn_2 x y

    have pos_for_pos' : ∀ y, 0 < y → 0 < f y := by
      intros y y_pos
      have f_y_nonneg : 0 ≤ f y := pos_for_pos y (by linarith)
      have: f y ≠ 0 := fun eq => has_root ⟨y, ⟨by linarith, eq⟩⟩
      exact lt_of_le f_y_nonneg this.symm

    have increasing_on_pos : ∀ x, ∀ y, 0 ≤ x → x < y → (f y) ^ 2 - (f x) ^ 2 > 0 := by
      intros x y x_nonneg y_big
      let h := (y - x) / 2
      have concrete := eqn_2' (x + h) h
      simp only [add_neg_cancel_right] at concrete
      have: 0 < h := by calc
        0 = 0 / 2 := by norm_num
        _ < (y - x) / 2 := by linarith
      have: 0 < x + h := by positivity
      have: 0 < f (x + h) := pos_for_pos' (x + h) (by tauto)
      have: 0 < f h := pos_for_pos' h (by tauto)
      calc (f y) ^ 2 - (f x) ^ 2 = f (x + h + h) ^ 2 - f x ^ 2 := by ring_nf
        _ = 2 * _ := concrete.symm
        _ > 0 := mul_pos (by norm_num) (pos_for_pos' _ (by positivity))

    have increasing_on_pos' : ∀ x, ∀ y, 0 ≤ x → x < y → f y - f x > 0 := by
      intros x y x_pos x_le_y
      have: _ := increasing_on_pos x y x_pos x_le_y
      have: 0 ≤ f x := pos_for_pos x x_pos
      have: 0 < f y := pos_for_pos' y (by linarith)
      have factored: (f y + f x) * (f y - f x) > 0 := by linarith
      exact pos_of_mul_pos_right factored (by positivity)

    have increasing : ∀ x, ∀ y, x < y → f x < f y := by
      intros x y x_le_y
      by_cases x_pos : 0 ≤ x
      . linarith [increasing_on_pos' x y x_pos x_le_y]
      . simp only [not_le] at x_pos
        have: 0 < f (-x) := pos_for_pos' (-x) (by linarith)
        have rewrite_x: f x = - f (-x) := eqn_1'' x
        have: f x < 0 := by linarith
        by_cases y_pos : 0 ≤ y
        . have: 0 ≤ f y := pos_for_pos y y_pos
          linarith
        . simp only [not_le] at y_pos
          have rewrite_y: f y = - f (-y) := eqn_1'' y
          rw [rewrite_y, rewrite_x]
          simp only [neg_lt_neg_iff, gt_iff_lt]
          have: -y < -x := by linarith
          have: _ := increasing_on_pos' (-y) (-x) (by linarith) (by linarith)
          linarith

    have injective : ∀ x, ∀ y, f x = f y → x = y := by
      intros x y f_eq
      rcases lt_trichotomy x y with (x_lt_y | ( eq | y_lt_x ) )
      . have: f x < f y := increasing x y x_lt_y
        linarith
      . exact eq
      . have: f y < f x := increasing y x y_lt_x
        linarith

    exact ClassifyIfInjective pr injective

theorem Classify' {f} (pr : Property f) : ((∀ (x : ℝ), f x = 0) ∨ (∀ x, f x = x) ∨ (∀ x, f x = 2)) := by
  by_cases f 0 = 0
  . cases' ClassifyIfZeroIsRoot pr (by assumption) with p q
    . tauto
    . tauto
  . right
    right
    sorry

theorem Classify {f} : (Property f) ↔ ((∀ (x : ℝ), f x = 0) ∨ (∀ x, f x = x) ∨ (∀ x, f x = 2)) := by
  apply Iff.intro _ _
  . intros property
    exact Classify' property
  . unfold Property
    rintro (zero | (id | two))
    . intros x y
      simp only [zero]
      norm_num
    . intros x y
      simp only [id]
      ring_nf
    . intros x y
      simp only [two]
      norm_num