/-
  WP34 — SHA-256 message schedule (FIPS 180-4 §6.2.2 step 1):
    W_t = σ1(W_{t−2}) + W_{t−7} + σ0(W_{t−15}) + W_{t−16}   for 16 ≤ t < 64.
  Each schedule step is invertible in its oldest input: W_{t−16} is determined by W_{t−15..t}.
-/
import Wp34.Sha256Words

namespace Wp34

/-- One schedule step, argument order `(W_{t−16}, W_{t−15}, W_{t−7}, W_{t−2})`
(`message_schedule` in the Python reference). -/
def scheduleStep (w16 w15 w7 w2 : Word) : Word := sigma1 w2 + w7 + sigma0 w15 + w16

/-- `W_{t−16} = W_t − σ1(W_{t−2}) − W_{t−7} − σ0(W_{t−15})`. -/
theorem scheduleStep_inv (w16 w15 w7 w2 : Word) :
    scheduleStep w16 w15 w7 w2 - sigma1 w2 - w7 - sigma0 w15 = w16 := by
  unfold scheduleStep sigma0 sigma1 rotr; bv_decide

/-- For fixed `(W_{t−15}, W_{t−7}, W_{t−2})`, the map `W_{t−16} ↦ W_t` is a bijection. -/
theorem scheduleStep_bijective (w15 w7 w2 : Word) :
    Bijective (fun w16 => scheduleStep w16 w15 w7 w2) :=
  ⟨fun x y h => by
      have hx := scheduleStep_inv x w15 w7 w2
      have hy := scheduleStep_inv y w15 w7 w2
      simp only at h
      rw [h] at hx; exact hx.symm.trans hy,
   fun wt => ⟨wt - sigma1 w2 - w7 - sigma0 w15, by
      unfold scheduleStep sigma0 sigma1 rotr; bv_decide⟩⟩

/-- The expanded message schedule as a function of the 16 message words `M`
(well-founded recursion on `t`). -/
def schedule (M : Fin 16 → Word) : Nat → Word
  | t =>
    if h : t < 16 then M ⟨t, h⟩
    else scheduleStep (schedule M (t - 16)) (schedule M (t - 15)) (schedule M (t - 7))
           (schedule M (t - 2))
  decreasing_by all_goals omega

theorem schedule_lt (M : Fin 16 → Word) (t : Nat) (h : t < 16) : schedule M t = M ⟨t, h⟩ := by
  rw [schedule]; simp [h]

theorem schedule_ge (M : Fin 16 → Word) (t : Nat) (h : 16 ≤ t) :
    schedule M t = scheduleStep (schedule M (t - 16)) (schedule M (t - 15))
      (schedule M (t - 7)) (schedule M (t - 2)) := by
  rw [schedule]; simp [Nat.not_lt.mpr h]

/-- Backward schedule: `W_{t−16}` is determined by `(W_{t−15}, W_{t−7}, W_{t−2}, W_t)`.
Consequently the 64-word expanded schedule is determined by ANY 16 consecutive words. -/
theorem schedule_backward (M : Fin 16 → Word) (t : Nat) (h : 16 ≤ t) :
    schedule M (t - 16) =
      schedule M t - sigma1 (schedule M (t - 2)) - schedule M (t - 7)
        - sigma0 (schedule M (t - 15)) := by
  rw [schedule_ge M t h, scheduleStep_inv]

end Wp34