import Mathlib.Tactic

theorem quadratic_residue_mod_4 (n : Nat) : n ^ 2 % 4 = 0 ∨ n ^ 2 % 4 = 1 := by
  have h : ∀ k : Fin 4, k.val ^ 2 % 4 = 0 ∨ k.val ^ 2 % 4 = 1 := by
    decide
  -- Use the fact that n % 4 is one of 0, 1, 2, 3 to reduce the problem
  have h_mod : n % 4 = 0 ∨ n % 4 = 1 ∨ n % 4 = 2 ∨ n % 4 = 3 := by omega
  -- Case analysis on n % 4
  rcases h_mod with (h0 | h1 | h2 | h3)
  · -- Case n % 4 = 0
    have : n ^ 2 % 4 = 0 := by
      have : n % 4 = 0 := h0
      have : n ^ 2 % 4 = (n % 4) ^ 2 % 4 := by simp [Nat.pow_mod]
      rw [this, h0]
      norm_num
    exact Or.inl this
  · -- Case n % 4 = 1
    have : n ^ 2 % 4 = 1 := by
      have : n % 4 = 1 := h1
      have : n ^ 2 % 4 = (n % 4) ^ 2 % 4 := by simp [Nat.pow_mod]
      rw [this, h1]
      norm_num
    exact Or.inr this
  · -- Case n % 4 = 2
    have : n ^ 2 % 4 = 0 := by
      have : n % 4 = 2 := h2
      have : n ^ 2 % 4 = (n % 4) ^ 2 % 4 := by simp [Nat.pow_mod]
      rw [this, h2]
      norm_num
    exact Or.inl this
  · -- Case n % 4 = 3
    have : n ^ 2 % 4 = 1 := by
      have : n % 4 = 3 := h3
      have : n ^ 2 % 4 = (n % 4) ^ 2 % 4 := by simp [Nat.pow_mod]
      rw [this, h3]
      norm_num
    exact Or.inr this