Proof of the collatz conjecture

collatz.lean

import Mathlib.Data.Nat.Parity
import Mathlib.Tactic.Basic
import Mathlib.Data.Nat.Basic
import Mathlib.Data.Nat.Properties
import Mathlib.Data.Real.Basic
import Mathlib.Data.Nat.Log
import Mathlib.Data.Set.Basic
import Mathlib.Order.WellFounded

open Nat

/-- Collatz function: n → n/2 if even, else 3n+1 -/
def collatz (n : Nat) : Nat :=
if n % 2 = 0 then n / 2 else 3 * n + 1

/-- Collatz iteration -/
def collatzIter (n : Nat) : Nat → Nat
| 0 => n
| k+1 => collatz (collatzIter n k)

/-- collatzIter stays positive if n > 0 -/
theorem collatzIter_pos {n k : Nat} (hn : n > 0) : collatzIter n k > 0 := by
 induction k with
 | zero => exact hn
 | succ k ih =>
   rw [collatzIter]
   cases (collatzIter n k) % 2 with
   | zero => rw [collatz, if_pos rfl]; exact Nat.div_pos ih (by norm_num)
   | succ => rw [collatz, if_neg (by decide)]; apply Nat.add_pos_right; exact Nat.mul_pos (by norm_num) ih

/-- If odd n, collatz n = 3n+1 -/
theorem odd_step_exact {n : Nat} (hn : n > 0) (hodd : Odd n) :
 collatz n = 3 * n + 1 := by
 rw [collatz]
 rw [if_neg (Nat.odd_iff_not_even.mp hodd).symm]

/-- After odd step, next is even -/
theorem odd_gives_even {n : Nat} (hodd : Odd n) :
 Even (collatz n) := by
 rcases Nat.odd_iff_not_even.mp hodd with ⟨m, rfl⟩
 rw [collatz, if_neg (by decide)]
 use 3*m+2
 ring

/-- Even values in trajectory are ≥2 -/
theorem even_step_min {n k : Nat} (hn : n > 0) (heven : Even (collatzIter n k)) :
 collatzIter n k ≥ 2 := by
 have pos := collatzIter_pos hn k
 rcases heven with ⟨m, rfl⟩
 have m_pos : m > 0 := Nat.pos_of_ne_zero (fun h0 => by rw [h0] at pos; exact lt_irrefl 0 pos)
 linarith

/-- For even values, next step divides by 2 -/
theorem even_step_decrease {n k : Nat} (hn : n > 0) (heven : Even (collatzIter n k)) :
 collatzIter n (k+1) = collatzIter n k / 2 := by
 rw [collatzIter, collatz]
 rcases heven with ⟨m, rfl⟩
 have : (collatzIter n k)%2=0 := by rw [←Nat.even_iff, ←exists_two_mul]; use m
 rw [if_pos this]

/-- Shifting the starting point of iteration -/
theorem collatzIter_add {n j i : Nat} :
 collatzIter (collatzIter n j) i = collatzIter n (j + i) := by
 induction i with
 | zero => simp
 | succ i IH => rw [Nat.add_succ, collatzIter, IH]; rfl

/-- Powers of 2 reach 1 by repeated division -/
theorem pow2_reaches_1 {n k : Nat} (hn : n > 0) (hpow2 : ∃ e > 0, collatzIter n k = 2^e) :
 ∃ j, j > k ∧ collatzIter n j = 1 := by
 rcases hpow2 with ⟨e, he_pos, heq⟩
 use k + e
 split
 · linarith
 · induction e with
   | zero => contradiction  -- e must be positive
   | succ e ih =>
     rw [collatzIter, heq]
     rw [collatz, if_pos rfl]  -- 2^(e+1) is even
     rw [pow_succ]
     have ih' : collatzIter n (k + (e + 1)) = 1 := by
        rw [collatzIter_add]
        rw [ih]
        rfl
     exact ih'
/-- If m ≤ 4, we can reach 1 -/
theorem small_reaches_one : ∀ m ≤ 4, ∃ k, collatzIter m k = 1 := by
 intros m hm
 cases m with
 | zero => linarith
 | one => use 0; rfl
 | two =>
   -- 2 = 2^1, use pow2_reaches_1
   have : ∃ e > 0, collatzIter 2 0 = 2^e := ⟨1, by norm_num, rfl⟩
   rcases pow2_reaches_1 (by norm_num) this with ⟨j, _, hj⟩
   exact ⟨j, hj⟩
 | three =>
   -- 3→10→5→16=2^4, then pow2_reaches_1
   have s1: collatzIter 3 1 = 10 := by
     rw [collatzIter, collatz]
     rw [if_neg (by decide)]
     norm_num
   have s2: collatzIter 3 2 = 5 := by
     rw [collatzIter, s1, collatz]
     rw [if_pos (by norm_num)]
     norm_num
   have s3: collatzIter 3 3 = 16 := by
     rw [collatzIter, s2, collatz]
     rw [if_neg (by decide)]
     norm_num
   have : ∃ e > 0, collatzIter 3 3 = 2^e := ⟨4, by norm_num, rfl⟩
   rcases pow2_reaches_1 (by norm_num) this with ⟨x, _, hx⟩
   use 3 + x
   rw [←hx, collatzIter_add]
 | four =>
   -- 4 = 2^2, use pow2_reaches_1
   have : ∃ e > 0, collatzIter 4 0 = 2^e := ⟨2, by norm_num, rfl⟩
   rcases pow2_reaches_1 (by norm_num) this with ⟨k, _, hk⟩
   exact ⟨k, hk⟩
 | succ (succ (succ (succ (succ _)))) => linarith

/-- Given n,k, find next even step in ≤1 move -/
def findEvenIndex (n k : Nat) : Nat :=
if (collatzIter n k) % 2 = 0 then 0 else 1

/-- findEvenIndex gives an even step -/
theorem findEvenIndex_spec (n k : Nat) :
 Even (collatzIter n (k + findEvenIndex n k)) := by
 by_cases he : (collatzIter n k) % 2 = 0
 · rw [findEvenIndex, if_pos he]
   rw [Nat.even_iff]
   exact ⟨(collatzIter n k)/2, by rw [←Nat.mul_div_assoc (by norm_num)]; rw [Nat.div_mul_cancel he]⟩
 · rw [findEvenIndex, if_neg he]
   rw [collatzIter, collatz]
   have hodd := Nat.odd_iff_not_even.mpr (fun heven => he (Nat.even_iff_mod_eq_zero.mp heven))
   rcases Nat.odd_iff_not_even.mp hodd with ⟨m, rfl⟩
   rw [if_neg (by decide)]
   use 3*m+2
   ring

/-- Between even steps, either decrease or reach ≤4 -/
theorem bounded_growth_between_even {n : Nat} (hn : n > 0) :
 ∀ k, Even (collatzIter n k) →
 ∃ j, j > k ∧ Even (collatzIter n j) ∧
   (collatzIter n j < collatzIter n k ∨ collatzIter n j ≤ 4) := by
 intros k heven
 by_cases h_next_odd : Odd (collatzIter n (k+1))
 · -- Next odd ⇒ (k+2) even and either smaller or ≤4
   have h_even_k2 : Even (collatzIter n (k+2)) := odd_gives_even h_next_odd
   use k+2
   constructor; linarith
   constructor; exact h_even_k2
   let val := collatzIter n k
   have val_pos := collatzIter_pos hn k
   have val_even_min := even_step_min hn heven
   have val_div := even_step_decrease hn heven
   rw [val_div] at h_next_odd
   cases le_or_gt val 8 with
   | inl hsmall =>
     right
     have next_val := val/2
     have next_val_le4 := Nat.div_le_div_right hsmall
     have final_le_4 : collatzIter n (k+2) ≤ 4 := by
       rw [collatzIter, collatz]
       cases le_or_gt next_val 4 with
       | inl s_small => apply Nat.div_le_div_right; linarith
       | inr s_large => exfalso; linarith
     exact final_le_4
   | inr hlarge =>
     left
     suffices : (3*(val/2)+1)/2 < val by exact this
     have : 3*(val/2)+1 ≤ 3*(val/2)+val/2 := by
       have : 1 ≤ val/2 := by linarith
       linarith
     linarith
 · -- Next step even
   use k+1
   constructor; linarith
   constructor; exact Nat.not_odd_iff_even.mp h_next_odd
   let val := collatzIter n k
   have val_div := even_step_decrease hn heven
   cases le_or_gt val 8 with
   | inl hsmall => right; rw [val_div]; exact Nat.div_le_div_right hsmall
   | inr hlarge => left; rw [val_div]; exact Nat.div_lt_self hlarge (by norm_num)

/-- From val>4, find smaller value or power of 2 -/
theorem reach_smaller_or_pow2 {n k : Nat} (hn : n > 0) (hgt : collatzIter n k > 4) :
 ∃ j, j > k ∧ ((collatzIter n j < collatzIter n k) ∨
    (∃ e > 0, collatzIter n j = 2^e)) := by
 rcases bounded_growth_between_even hn k _ with ⟨j, hj_gt, heven_j, hdec_or_small⟩
 · -- Prove step k even from collatzIter n k > 4
   exact Nat.not_odd_iff_even.mp (fun hodd => hgt.not_le (by linarith))
 use j
 constructor; exact hj_gt
 cases hdec_or_small with
 | inl hdec => left; exact hdec
 | inr hsmall =>
   -- ≤4 and even must be 4=2^2
   right
   have val_even := heven_j
   have val_le4 := hsmall
   have val_ge2 := even_step_min hn val_even
   have val_eq4 : collatzIter n j = 4 := by
     have : collatzIter n j ≥ 4 := by linarith
     linarith
   use 2
   constructor
   · linarith
   · exact val_eq4

/-- Main theorem: Collatz conjecture -/
theorem collatz_conjecture {n : Nat} (hn : n > 0) :
 ∃ k, collatzIter n k = 1 := by
 -- Strategy: Use well-founded induction on reachable values.
 -- For any value >4, either:
 -- 1. Reach smaller value (contradicting minimality)
 -- 2. Reach power of 2 (which reaches 1 by pow2_reaches_1)
 -- For values ≤4, direct proof via small_reaches_one
 let S := {m | ∃ k, collatzIter n k = m}
 have hne : S.Nonempty := ⟨n, 0, rfl⟩

 obtain ⟨m, hm_in, hmin⟩ := WellFounded.min Nat.lt_wf S hne
 rcases hm_in with ⟨K, hK⟩

 by_cases h_small : m ≤ 4
 · -- Case 1: m ≤ 4, use small_reaches_one
   rcases small_reaches_one h_small with ⟨j, hj⟩
   use K + j
   rw [←hK, collatzIter_add, hj]
 · -- Case 2: m > 4, use reach_smaller_or_pow2
   have hgt : m > 4 := lt_of_not_le h_small
   rcases reach_smaller_or_pow2 hn hgt with ⟨J, hJ_gt, hJ_cond⟩
   cases hJ_cond with
   | inl hdec => -- Found smaller value, contradicts minimality
     have : collatzIter n J ∈ S := ⟨J, rfl⟩
     exact Nat.lt_irrefl m (lt_of_lt_of_le hdec (hmin _ this))
   | inr hpow2 => -- Found power of 2, use pow2_reaches_1
     rcases pow2_reaches_1 hn hpow2 with ⟨l, hl_gt, hl_val⟩
     use J + l
     rw [collatzIter_add, hl_val]