import Mathlib
import Aesop

set_option maxHeartbeats 0

open Topology Filter

theorem putnam_2019_a1
: {n : ℤ | ∃ A B C : ℤ, A ≥ 0 ∧ B ≥ 0 ∧ C ≥ 0 ∧ A^3 + B^3 + C^3 - 3*A*B*C = n} = (({n : ℤ | n ≥ 0 ∧ ¬Int.ModEq 9 n 3 ∧ ¬Int.ModEq 9 n 6}) : Set ℤ ) := by
  have h_subset₁ : {n : ℤ | ∃ A B C : ℤ, A ≥ 0 ∧ B ≥ 0 ∧ C ≥ 0 ∧ A^3 + B^3 + C^3 - 3*A*B*C = n} ⊆ ({n : ℤ | n ≥ 0 ∧ ¬Int.ModEq 9 n 3 ∧ ¬Int.ModEq 9 n 6} : Set ℤ) := by
    intro n hn
    rcases hn with ⟨A, B, C, hA, hB, hC, hn⟩
    have h₁ : n ≥ 0 := by
      have h₁ : (A : ℤ) ≥ 0 := by exact_mod_cast hA
      have h₂ : (B : ℤ) ≥ 0 := by exact_mod_cast hB
      have h₃ : (C : ℤ) ≥ 0 := by exact_mod_cast hC
      have h₄ : (A : ℤ)^3 + (B : ℤ)^3 + (C : ℤ)^3 - 3 * (A : ℤ) * (B : ℤ) * (C : ℤ) ≥ 0 := by
        nlinarith [sq_nonneg (A - B), sq_nonneg (B - C), sq_nonneg (C - A),
          mul_nonneg h₁ h₂, mul_nonneg h₂ h₃, mul_nonneg h₃ h₁,
          sq_nonneg (A + B + C), sq_nonneg (A + B - C), sq_nonneg (A + C - B),
          sq_nonneg (B + C - A)]
      linarith
    have h₂ : ¬Int.ModEq 9 n 3 := by
      have h₂ : ∀ (A B C : ℤ), A ≥ 0 → B ≥ 0 → C ≥ 0 → (A^3 + B^3 + C^3 - 3*A*B*C) % 9 ≠ 3 := by
        intro A B C hA hB hC
        have h₃ : A % 9 = 0 ∨ A % 9 = 1 ∨ A % 9 = 2 ∨ A % 9 = 3 ∨ A % 9 = 4 ∨ A % 9 = 5 ∨ A % 9 = 6 ∨ A % 9 = 7 ∨ A % 9 = 8 := by
          omega
        have h₄ : B % 9 = 0 ∨ B % 9 = 1 ∨ B % 9 = 2 ∨ B % 9 = 3 ∨ B % 9 = 4 ∨ B % 9 = 5 ∨ B % 9 = 6 ∨ B % 9 = 7 ∨ B % 9 = 8 := by
          omega
        have h₅ : C % 9 = 0 ∨ C % 9 = 1 ∨ C % 9 = 2 ∨ C % 9 = 3 ∨ C % 9 = 4 ∨ C % 9 = 5 ∨ C % 9 = 6 ∨ C % 9 = 7 ∨ C % 9 = 8 := by
          omega
        rcases h₃ with (h₃ | h₃ | h₃ | h₃ | h₃ | h₃ | h₃ | h₃ | h₃) <;>
          rcases h₄ with (h₄ | h₄ | h₄ | h₄ | h₄ | h₄ | h₄ | h₄ | h₄) <;>
            rcases h₅ with (h₅ | h₅ | h₅ | h₅ | h₅ | h₅ | h₅ | h₅ | h₅) <;>
              (try omega) <;>
                (try {
                  simp [h₃, h₄, h₅, pow_three, Int.mul_emod, Int.add_emod, Int.sub_emod]
                  <;> norm_num <;> omega
                })
      have h₃ := h₂ A B C hA hB hC
      intro h₄
      have h₅ : (n : ℤ) % 9 = 3 := by
        rw [Int.ModEq] at h₄
        omega
      have h₆ : (A^3 + B^3 + C^3 - 3*A*B*C) % 9 = 3 := by
        omega
      contradiction
    have h₃ : ¬Int.ModEq 9 n 6 := by
      have h₂ : ∀ (A B C : ℤ), A ≥ 0 → B ≥ 0 → C ≥ 0 → (A^3 + B^3 + C^3 - 3*A*B*C) % 9 ≠ 6 := by
        intro A B C hA hB hC
        have h₃ : A % 9 = 0 ∨ A % 9 = 1 ∨ A % 9 = 2 ∨ A % 9 = 3 ∨ A % 9 = 4 ∨ A % 9 = 5 ∨ A % 9 = 6 ∨ A % 9 = 7 ∨ A % 9 = 8 := by
          omega
        have h₄ : B % 9 = 0 ∨ B % 9 = 1 ∨ B % 9 = 2 ∨ B % 9 = 3 ∨ B % 9 = 4 ∨ B % 9 = 5 ∨ B % 9 = 6 ∨ B % 9 = 7 ∨ B % 9 = 8 := by
          omega
        have h₅ : C % 9 = 0 ∨ C % 9 = 1 ∨ C % 9 = 2 ∨ C % 9 = 3 ∨ C % 9 = 4 ∨ C % 9 = 5 ∨ C % 9 = 6 ∨ C % 9 = 7 ∨ C % 9 = 8 := by
          omega
        rcases h₃ with (h₃ | h₃ | h₃ | h₃ | h₃ | h₃ | h₃ | h₃ | h₃) <;>
          rcases h₄ with (h₄ | h₄ | h₄ | h₄ | h₄ | h₄ | h₄ | h₄ | h₄) <;>
            rcases h₅ with (h₅ | h₅ | h₅ | h₅ | h₅ | h₅ | h₅ | h₅ | h₅) <;>
              (try omega) <;>
                (try {
                  simp [h₃, h₄, h₅, pow_three, Int.mul_emod, Int.add_emod, Int.sub_emod]
                  <;> norm_num <;> omega
                })
      have h₃ := h₂ A B C hA hB hC
      intro h₄
      have h₅ : (n : ℤ) % 9 = 6 := by
        rw [Int.ModEq] at h₄
        omega
      have h₆ : (A^3 + B^3 + C^3 - 3*A*B*C) % 9 = 6 := by
        omega
      contradiction
    exact ⟨h₁, h₂, h₃⟩

  have h_subset₂ : ({n : ℤ | n ≥ 0 ∧ ¬Int.ModEq 9 n 3 ∧ ¬Int.ModEq 9 n 6} : Set ℤ) ⊆ {n : ℤ | ∃ A B C : ℤ, A ≥ 0 ∧ B ≥ 0 ∧ C ≥ 0 ∧ A^3 + B^3 + C^3 - 3*A*B*C = n} := by
    intro n hn
    have h₁ : n ≥ 0 := hn.1
    have h₂ : ¬Int.ModEq 9 n 3 := hn.2.1
    have h₃ : ¬Int.ModEq 9 n 6 := hn.2.2

    have h₄ : n % 3 = 0 ∨ n % 3 = 1 ∨ n % 3 = 2 := by
      have h₄ : n % 3 = 0 ∨ n % 3 = 1 ∨ n % 3 = 2 := by omega
      exact h₄
    rcases h₄ with (h₄ | h₄ | h₄)
    ·
      have h₅ : n % 9 ≠ 3 := by
        intro h₅
        have h₆ : Int.ModEq 9 n 3 := by
          rw [Int.ModEq]
          omega
        contradiction
      have h₆ : n % 9 ≠ 6 := by
        intro h₆
        have h₇ : Int.ModEq 9 n 6 := by
          rw [Int.ModEq]
          omega
        contradiction

      have h₇ : n % 9 = 0 := by
        have h₇ : n % 9 = 0 ∨ n % 9 = 3 ∨ n % 9 = 6 := by
          omega
        rcases h₇ with (h₇ | h₇ | h₇)
        · exact h₇
        · exfalso
          exact h₅ h₇
        · exfalso
          exact h₆ h₇

      have h₈ : ∃ (k : ℤ), n = 9 * k := by
        use n / 9
        have h₈ : n % 9 = 0 := h₇
        have h₉ : n = 9 * (n / 9) := by
          omega
        exact h₉
      rcases h₈ with ⟨k, hk⟩
      by_cases h₉ : k = 0
      ·
        use 0, 0, 0
        simp [h₉, hk]
        <;> norm_num
      ·
        have h₁₀ : k > 0 := by
          by_contra h₁₀
          have h₁₁ : k ≤ 0 := by linarith
          have h₁₂ : n ≤ 0 := by
            have h₁₃ : (9 : ℤ) * k ≤ 0 := by nlinarith
            linarith
          have h₁₃ : n ≥ 0 := h₁
          have h₁₄ : n = 0 := by linarith
          have h₁₅ : k = 0 := by
            omega
          contradiction

        use (k + 1 : ℤ), (k : ℤ), (k - 1 : ℤ)
        have h₁₁ : (k + 1 : ℤ) ≥ 0 := by linarith
        have h₁₂ : (k : ℤ) ≥ 0 := by linarith
        have h₁₃ : (k - 1 : ℤ) ≥ 0 := by
          have h₁₄ : k ≥ 1 := by
            by_contra h₁₄
            have h₁₅ : k ≤ 0 := by linarith
            have h₁₆ : n ≤ 0 := by
              have h₁₇ : (9 : ℤ) * k ≤ 0 := by nlinarith
              linarith
            have h₁₇ : n ≥ 0 := h₁
            have h₁₈ : n = 0 := by linarith
            have h₁₉ : k = 0 := by
              omega
            contradiction
          linarith
        constructor
        · exact_mod_cast h₁₁
        constructor
        · exact_mod_cast h₁₂
        constructor
        · exact_mod_cast h₁₃

        have h₁₄ : (k + 1 : ℤ)^3 + (k : ℤ)^3 + (k - 1 : ℤ)^3 - 3 * (k + 1 : ℤ) * (k : ℤ) * (k - 1 : ℤ) = 9 * k := by
          ring_nf
          <;> nlinarith
        rw [hk]
        linarith
    ·

      have h₅ : (n - 1) % 3 = 0 := by
        omega
      have h₆ : ∃ (m : ℤ), n = 3 * m + 1 := by
        use (n - 1) / 3
        have h₆ : (n - 1) % 3 = 0 := h₅
        have h₇ : n - 1 = 3 * ((n - 1) / 3) := by
          omega
        linarith
      rcases h₆ with ⟨m, hm⟩
      have h₇ : m ≥ 0 := by
        by_contra h₇
        have h₈ : m < 0 := by linarith
        have h₉ : n < 1 := by
          nlinarith
        have h₁₀ : n ≥ 0 := h₁
        omega
      use (m + 1 : ℤ), (m : ℤ), (m : ℤ)
      have h₈ : (m + 1 : ℤ) ≥ 0 := by linarith
      have h₉ : (m : ℤ) ≥ 0 := by linarith
      have h₁₀ : (m : ℤ) ≥ 0 := by linarith
      constructor
      · exact_mod_cast h₈
      constructor
      · exact_mod_cast h₉
      constructor
      · exact_mod_cast h₁₀

      have h₁₁ : (m + 1 : ℤ)^3 + (m : ℤ)^3 + (m : ℤ)^3 - 3 * (m + 1 : ℤ) * (m : ℤ) * (m : ℤ) = 3 * m + 1 := by
        ring_nf
        <;> nlinarith
      rw [hm]
      linarith
    ·

      have h₅ : (n - 2) % 3 = 0 := by
        omega
      have h₆ : ∃ (m : ℤ), n = 3 * m + 2 := by
        use (n - 2) / 3
        have h₆ : (n - 2) % 3 = 0 := h₅
        have h₇ : n - 2 = 3 * ((n - 2) / 3) := by
          omega
        linarith
      rcases h₆ with ⟨m, hm⟩
      have h₇ : m ≥ 0 := by
        by_contra h₇
        have h₈ : m < 0 := by linarith
        have h₉ : n < 2 := by
          nlinarith
        have h₁₀ : n ≥ 0 := h₁
        have h₁₁ : m ≤ -1 := by
          omega
        have h₁₂ : n < 0 := by
          nlinarith
        linarith
      use (m + 1 : ℤ), (m + 1 : ℤ), (m : ℤ)
      have h₈ : (m + 1 : ℤ) ≥ 0 := by linarith
      have h₉ : (m + 1 : ℤ) ≥ 0 := by linarith
      have h₁₀ : (m : ℤ) ≥ 0 := by linarith
      constructor
      · exact_mod_cast h₈
      constructor
      · exact_mod_cast h₉
      constructor
      · exact_mod_cast h₁₀

      have h₁₁ : (m + 1 : ℤ)^3 + (m + 1 : ℤ)^3 + (m : ℤ)^3 - 3 * (m + 1 : ℤ) * (m + 1 : ℤ) * (m : ℤ) = 3 * m + 2 := by
        ring_nf
        <;> nlinarith
      rw [hm]
      linarith

  have h_main : {n : ℤ | ∃ A B C : ℤ, A ≥ 0 ∧ B ≥ 0 ∧ C ≥ 0 ∧ A^3 + B^3 + C^3 - 3*A*B*C = n} = (({n : ℤ | n ≥ 0 ∧ ¬Int.ModEq 9 n 3 ∧ ¬Int.ModEq 9 n 6}) : Set ℤ ) := by
    apply Set.Subset.antisymm
    · exact h_subset₁
    · exact h_subset₂

  exact h_main