/-
Copyright (c) 2020 Johan Commelin. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Johan Commelin
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.algebra.group_power.default
import Mathlib.PostPort

universes u_1 u_2

namespace Mathlib

/-!
# Powers of elements of groups with an adjoined zero element

In this file we define integer power functions for groups with an adjoined zero element.
This generalises the integer power function on a division ring.
-/

@[simp] theorem zero_pow' {M : Type u_1} [monoid_with_zero M] (n : ℕ) : n ≠ 0 → 0 ^ n = 0 := sorry

theorem ne_zero_pow {M : Type u_1} [monoid_with_zero M] {a : M} {n : ℕ} (hn : n ≠ 0) :
    a ^ n ≠ 0 → a ≠ 0 :=
  imp_of_not_imp_not (a ^ n ≠ 0) (a ≠ 0)
    (eq.mpr (id (imp_congr_eq (push_neg.not_not_eq (a = 0)) (push_neg.not_not_eq (a ^ n = 0))))
      fun (ᾰ : a = 0) => Eq._oldrec (zero_pow' n hn) (Eq.symm ᾰ))

@[simp] theorem zero_pow_eq_zero {M : Type u_1} [monoid_with_zero M] [nontrivial M] {n : ℕ} :
    0 ^ n = 0 ↔ 0 < n :=
  sorry

@[simp] theorem inv_pow' {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℕ) :
    a⁻¹ ^ n = (a ^ n⁻¹) :=
  sorry

theorem pow_sub' {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) {m : ℕ} {n : ℕ} (ha : a ≠ 0)
    (h : n ≤ m) : a ^ (m - n) = a ^ m * (a ^ n⁻¹) :=
  sorry

theorem pow_inv_comm' {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (m : ℕ) (n : ℕ) :
    a⁻¹ ^ m * a ^ n = a ^ n * a⁻¹ ^ m :=
  commute.pow_pow (commute.inv_left' (commute.refl a)) m n

/--
The power operation in a group with zero.
This extends `monoid.pow` to negative integers
with the definition `a ^ (-n) = (a ^ n)⁻¹`.
-/
def fpow {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) : ℤ → G₀ := sorry

protected instance int.has_pow {G₀ : Type u_1} [group_with_zero G₀] : has_pow G₀ ℤ :=
  has_pow.mk fpow

@[simp] theorem fpow_coe_nat {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℕ) :
    a ^ ↑n = a ^ n :=
  rfl

theorem fpow_of_nat {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℕ) :
    a ^ Int.ofNat n = a ^ n :=
  rfl

@[simp] theorem fpow_neg_succ_of_nat {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℕ) :
    a ^ Int.negSucc n = (a ^ Nat.succ n⁻¹) :=
  rfl

@[simp] theorem fpow_zero {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) : a ^ 0 = 1 := rfl

@[simp] theorem fpow_one {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) : a ^ 1 = a := mul_one a

@[simp] theorem one_fpow {G₀ : Type u_1} [group_with_zero G₀] (n : ℤ) : 1 ^ n = 1 := sorry

theorem zero_fpow {G₀ : Type u_1} [group_with_zero G₀] (z : ℤ) : z ≠ 0 → 0 ^ z = 0 := sorry

@[simp] theorem fpow_neg {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℤ) :
    a ^ (-n) = (a ^ n⁻¹) :=
  sorry

theorem fpow_neg_one {G₀ : Type u_1} [group_with_zero G₀] (x : G₀) : x ^ (-1) = (x⁻¹) :=
  congr_arg has_inv.inv (pow_one x)

theorem inv_fpow {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℤ) : a⁻¹ ^ n = (a ^ n⁻¹) :=
  int.cases_on n (fun (n : ℕ) => idRhs (a⁻¹ ^ n = (a ^ n⁻¹)) (inv_pow' a n))
    fun (n : ℕ) =>
      idRhs (a⁻¹ ^ (n + 1)⁻¹ = (a ^ (n + 1)⁻¹⁻¹)) (congr_arg has_inv.inv (inv_pow' a (n + 1)))

theorem fpow_add_one {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} (ha : a ≠ 0) (n : ℤ) :
    a ^ (n + 1) = a ^ n * a :=
  sorry

theorem fpow_sub_one {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} (ha : a ≠ 0) (n : ℤ) :
    a ^ (n - 1) = a ^ n * (a⁻¹) :=
  sorry

theorem fpow_add {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} (ha : a ≠ 0) (m : ℤ) (n : ℤ) :
    a ^ (m + n) = a ^ m * a ^ n :=
  sorry

theorem fpow_add' {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} {m : ℤ} {n : ℤ}
    (h : a ≠ 0 ∨ m + n ≠ 0 ∨ m = 0 ∧ n = 0) : a ^ (m + n) = a ^ m * a ^ n :=
  sorry

theorem fpow_one_add {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} (h : a ≠ 0) (i : ℤ) :
    a ^ (1 + i) = a * a ^ i :=
  eq.mpr (id (Eq._oldrec (Eq.refl (a ^ (1 + i) = a * a ^ i)) (fpow_add h 1 i)))
    (eq.mpr (id (Eq._oldrec (Eq.refl (a ^ 1 * a ^ i = a * a ^ i)) (fpow_one a)))
      (Eq.refl (a * a ^ i)))

theorem semiconj_by.fpow_right {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} {x : G₀} {y : G₀}
    (h : semiconj_by a x y) (m : ℤ) : semiconj_by a (x ^ m) (y ^ m) :=
  int.cases_on m (fun (m : ℕ) => idRhs (semiconj_by a (x ^ m) (y ^ m)) (semiconj_by.pow_right h m))
    fun (m : ℕ) =>
      idRhs (semiconj_by a (x ^ (m + 1)⁻¹) (y ^ (m + 1)⁻¹))
        (semiconj_by.inv_right' (semiconj_by.pow_right h (m + 1)))

theorem commute.fpow_right {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} {b : G₀} (h : commute a b)
    (m : ℤ) : commute a (b ^ m) :=
  semiconj_by.fpow_right h

theorem commute.fpow_left {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} {b : G₀} (h : commute a b)
    (m : ℤ) : commute (a ^ m) b :=
  commute.symm (commute.fpow_right (commute.symm h) m)

theorem commute.fpow_fpow {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} {b : G₀} (h : commute a b)
    (m : ℤ) (n : ℤ) : commute (a ^ m) (b ^ n) :=
  commute.fpow_right (commute.fpow_left h m) n

theorem commute.fpow_self {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℤ) :
    commute (a ^ n) a :=
  commute.fpow_left (commute.refl a) n

theorem commute.self_fpow {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℤ) :
    commute a (a ^ n) :=
  commute.fpow_right (commute.refl a) n

theorem commute.fpow_fpow_self {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (m : ℤ) (n : ℤ) :
    commute (a ^ m) (a ^ n) :=
  commute.fpow_fpow (commute.refl a) m n

theorem fpow_bit0 {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℤ) :
    a ^ bit0 n = a ^ n * a ^ n :=
  sorry

theorem fpow_bit1 {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℤ) :
    a ^ bit1 n = a ^ n * a ^ n * a :=
  sorry

theorem fpow_mul {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (m : ℤ) (n : ℤ) :
    a ^ (m * n) = (a ^ m) ^ n :=
  sorry

theorem fpow_mul' {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (m : ℤ) (n : ℤ) :
    a ^ (m * n) = (a ^ n) ^ m :=
  eq.mpr (id (Eq._oldrec (Eq.refl (a ^ (m * n) = (a ^ n) ^ m)) (mul_comm m n)))
    (eq.mpr (id (Eq._oldrec (Eq.refl (a ^ (n * m) = (a ^ n) ^ m)) (fpow_mul a n m)))
      (Eq.refl ((a ^ n) ^ m)))

@[simp] theorem units.coe_gpow' {G₀ : Type u_1} [group_with_zero G₀] (u : units G₀) (n : ℤ) :
    ↑(u ^ n) = ↑u ^ n :=
  sorry

theorem fpow_ne_zero_of_ne_zero {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} (ha : a ≠ 0) (z : ℤ) :
    a ^ z ≠ 0 :=
  int.cases_on z (fun (z : ℕ) => idRhs (a ^ z ≠ 0) (pow_ne_zero z ha))
    fun (z : ℕ) => idRhs (a ^ Nat.succ z⁻¹ ≠ 0) (inv_ne_zero (pow_ne_zero (Nat.succ z) ha))

theorem fpow_sub {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} (ha : a ≠ 0) (z1 : ℤ) (z2 : ℤ) :
    a ^ (z1 - z2) = a ^ z1 / a ^ z2 :=
  sorry

theorem commute.mul_fpow {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} {b : G₀} (h : commute a b)
    (i : ℤ) : (a * b) ^ i = a ^ i * b ^ i :=
  sorry

theorem mul_fpow {G₀ : Type u_1} [comm_group_with_zero G₀] (a : G₀) (b : G₀) (m : ℤ) :
    (a * b) ^ m = a ^ m * b ^ m :=
  commute.mul_fpow (commute.all a b) m

theorem fpow_bit0' {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℤ) :
    a ^ bit0 n = (a * a) ^ n :=
  Eq.trans (fpow_bit0 a n) (Eq.symm (commute.mul_fpow (commute.refl a) n))

theorem fpow_bit1' {G₀ : Type u_1} [group_with_zero G₀] (a : G₀) (n : ℤ) :
    a ^ bit1 n = (a * a) ^ n * a :=
  eq.mpr (id (Eq._oldrec (Eq.refl (a ^ bit1 n = (a * a) ^ n * a)) (fpow_bit1 a n)))
    (eq.mpr
      (id
        (Eq._oldrec (Eq.refl (a ^ n * a ^ n * a = (a * a) ^ n * a))
          (commute.mul_fpow (commute.refl a) n)))
      (Eq.refl (a ^ n * a ^ n * a)))

theorem fpow_eq_zero {G₀ : Type u_1} [group_with_zero G₀] {x : G₀} {n : ℤ} (h : x ^ n = 0) :
    x = 0 :=
  classical.by_contradiction fun (hx : ¬x = 0) => fpow_ne_zero_of_ne_zero hx n h

theorem fpow_ne_zero {G₀ : Type u_1} [group_with_zero G₀] {x : G₀} (n : ℤ) : x ≠ 0 → x ^ n ≠ 0 :=
  mt fpow_eq_zero

theorem fpow_neg_mul_fpow_self {G₀ : Type u_1} [group_with_zero G₀] (n : ℤ) {x : G₀} (h : x ≠ 0) :
    x ^ (-n) * x ^ n = 1 :=
  eq.mpr (id (Eq._oldrec (Eq.refl (x ^ (-n) * x ^ n = 1)) (fpow_neg x n)))
    (inv_mul_cancel (fpow_ne_zero n h))

theorem one_div_pow {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} (n : ℕ) :
    (1 / a) ^ n = 1 / a ^ n :=
  sorry

theorem one_div_fpow {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} (n : ℤ) :
    (1 / a) ^ n = 1 / a ^ n :=
  sorry

@[simp] theorem inv_fpow' {G₀ : Type u_1} [group_with_zero G₀] {a : G₀} (n : ℤ) :
    a⁻¹ ^ n = a ^ (-n) :=
  sorry

@[simp] theorem div_pow {G₀ : Type u_1} [comm_group_with_zero G₀] (a : G₀) (b : G₀) (n : ℕ) :
    (a / b) ^ n = a ^ n / b ^ n :=
  sorry

@[simp] theorem div_fpow {G₀ : Type u_1} [comm_group_with_zero G₀] (a : G₀) {b : G₀} (n : ℤ) :
    (a / b) ^ n = a ^ n / b ^ n :=
  sorry

theorem div_sq_cancel {G₀ : Type u_1} [comm_group_with_zero G₀] {a : G₀} (ha : a ≠ 0) (b : G₀) :
    a ^ bit0 1 * b / a = a * b :=
  eq.mpr (id (Eq._oldrec (Eq.refl (a ^ bit0 1 * b / a = a * b)) (pow_two a)))
    (eq.mpr (id (Eq._oldrec (Eq.refl (a * a * b / a = a * b)) (mul_assoc a a b)))
      (eq.mpr (id (Eq._oldrec (Eq.refl (a * (a * b) / a = a * b)) (mul_div_cancel_left (a * b) ha)))
        (Eq.refl (a * b))))

/-- If a monoid homomorphism `f` between two `group_with_zero`s maps `0` to `0`, then it maps `x^n`,
`n : ℤ`, to `(f x)^n`. -/
theorem monoid_with_zero_hom.map_fpow {G₀ : Type u_1} {G₀' : Type u_2} [group_with_zero G₀]
    [group_with_zero G₀'] (f : monoid_with_zero_hom G₀ G₀') (x : G₀) (n : ℤ) :
    coe_fn f (x ^ n) = coe_fn f x ^ n :=
  sorry

end Mathlib


/-
Copyright (c) 2020 Ruben Van de Velde. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
-/

import Flt.Primes
import Flt.Spts
import Flt.OddPrimeOrFour

theorem Zsqrtd.associated_norm_of_associated {d : ℤ} {f g : ℤ√d} (h : Associated f g) :
    Associated f.norm g.norm := by
  obtain ⟨u, rfl⟩ := h
  have := (Zsqrtd.isUnit_iff_norm_isUnit u).mp u.isUnit
  rw [Zsqrtd.norm_mul]
  exact associated_mul_unit_right (Zsqrtd.norm f) _ this
#align zsqrtd.associated_norm_of_associated Zsqrtd.associated_norm_of_associated

theorem OddPrimeOrFour.im_ne_zero {p : ℤ√(-3)} (h : OddPrimeOrFour p.norm)
    (hcoprime : IsCoprime p.re p.im) : p.im ≠ 0 :=
  by
  intro H
  rw [Zsqrtd.norm_def, H, MulZeroClass.mul_zero, sub_zero, ← pow_two] at h
  obtain h | ⟨hp, -⟩ := h
  · rw [H, isCoprime_zero_right, Int.isUnit_iff_abs_eq] at hcoprime
    rw [← sq_abs, hcoprime] at h
    norm_num at h
  · exact pow_not_prime one_lt_two.ne' hp
#align odd_prime_or_four.im_ne_zero OddPrimeOrFour.im_ne_zero

theorem Zsqrt3.isUnit_iff {z : ℤ√(-3)} : IsUnit z ↔ abs z.re = 1 ∧ z.im = 0 :=
  by
  rw [← Zsqrtd.norm_eq_one_iff, ← Int.coe_nat_inj', Int.ofNat_one,
    Int.natAbs_of_nonneg (Zsqrtd.norm_nonneg (by norm_num) z)]
  refine' ⟨Spts.eq_one, fun h => _⟩
  rw [Zsqrtd.norm_def, h.2, MulZeroClass.mul_zero, sub_zero, ← sq, ← sq_abs, h.1, one_pow]
#align zsqrt3.is_unit_iff Zsqrt3.isUnit_iff

theorem Zsqrt3.coe_of_isUnit {z : ℤ√(-3)} (h : IsUnit z) : ∃ u : Units ℤ, z = u :=
  by
  obtain ⟨u, hu⟩ := Zsqrt3.isUnit_iff.mp h
  obtain ⟨v, hv⟩ : IsUnit z.re := by rwa [Int.isUnit_iff_abs_eq]
  use v
  rw [Zsqrtd.ext, hu, ← hv]
  simp only [Zsqrtd.coe_int_re, and_true_iff, eq_self_iff_true, Zsqrtd.coe_int_im]
#align zsqrt3.coe_of_is_unit Zsqrt3.coe_of_isUnit

theorem Zsqrt3.coe_of_isUnit' {z : ℤ√(-3)} (h : IsUnit z) : ∃ u : ℤ, z = u ∧ abs u = 1 :=
  by
  obtain ⟨u, hu⟩ := Zsqrt3.coe_of_isUnit h
  refine' ⟨u, _, _⟩
  · rw [hu]
  · exact Int.isUnit_iff_abs_eq.mp u.isUnit
#align zsqrt3.coe_of_is_unit' Zsqrt3.coe_of_isUnit'

theorem OddPrimeOrFour.factors (a : ℤ√(-3)) (x : ℤ) (hcoprime : IsCoprime a.re a.im)
    (hx : OddPrimeOrFour x) (hfactor : x ∣ a.norm) :
    ∃ c : ℤ√(-3), |x| = c.norm ∧ 0 ≤ c.re ∧ c.im ≠ 0 :=
  by
  obtain rfl | ⟨hprime, hodd⟩ := hx
  · refine' ⟨⟨1, 1⟩, _, zero_le_one, one_ne_zero⟩
    simp only [Zsqrtd.norm_def, mul_one, abs_of_pos, zero_lt_four, sub_neg_eq_add]
  · obtain ⟨c, hc⟩ := _root_.factors a x hcoprime hodd hfactor
    rw [hc]
    apply Zsqrtd.exists c
    intro H
    rw [Zsqrtd.norm_def, H, MulZeroClass.mul_zero, sub_zero, ← pow_two, eq_comm] at hc
    have := hprime.abs
    rw [←hc] at this
    exact pow_not_prime one_lt_two.ne' this
#align odd_prime_or_four.factors OddPrimeOrFour.factors

theorem step1a {a : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) (heven : Even a.norm) :
    Odd a.re ∧ Odd a.im :=
  by
  rw [Int.odd_iff_not_even, Int.odd_iff_not_even]
  have : Even a.re ↔ Even a.im := by simpa [Zsqrtd.norm_def, parity_simps] using heven
  apply (iff_iff_and_or_not_and_not.mp this).resolve_left
  rw [even_iff_two_dvd, even_iff_two_dvd]
  rintro ⟨hre, him⟩
  have := hcoprime.isUnit_of_dvd' hre him
  rw [isUnit_iff_dvd_one] at this
  norm_num at this
#align step1a step1a

theorem step1 {a : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) (heven : Even a.norm) :
    ∃ u : ℤ√(-3), a = ⟨1, 1⟩ * u ∨ a = ⟨1, -1⟩ * u :=
  by
  obtain ⟨ha, hb⟩ := step1a hcoprime heven
  obtain hdvd | hdvd := Int.four_dvd_add_or_sub_of_odd ha hb
  · obtain ⟨v, hv⟩ := hdvd
    rw [← eq_sub_iff_add_eq] at hv
    use ⟨v - a.im, v⟩
    right
    rw [Zsqrtd.ext, hv, Zsqrtd.mul_re, Zsqrtd.mul_im]
    dsimp only
    constructor <;> ring
  · obtain ⟨v, hv⟩ := hdvd
    rw [sub_eq_iff_eq_add] at hv
    use ⟨a.im + v, -v⟩
    left
    rw [Zsqrtd.ext, hv, Zsqrtd.mul_re, Zsqrtd.mul_im]
    dsimp only
    constructor <;> ring
#align step1 step1

theorem step1' {a : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) (heven : Even a.norm) :
    ∃ u : ℤ√(-3), IsCoprime u.re u.im ∧ a.norm = 4 * u.norm :=
  by
  obtain ⟨u', hu'⟩ := step1 hcoprime heven
  refine' ⟨u', _, _⟩
  · apply Zsqrtd.coprime_of_dvd_coprime hcoprime
    obtain rfl | rfl := hu' <;> apply dvd_mul_left
  · cases' hu' with hu' hu' <;>
      · rw [hu', Zsqrtd.norm_mul]
        congr
#align step1' step1'

theorem step1'' {a p : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) (hp : p.norm = 4) (hq : p.im ≠ 0)
    (heven : Even a.norm) :
    ∃ u : ℤ√(-3), IsCoprime u.re u.im ∧ (a = p * u ∨ a = star p * u) ∧ a.norm = 4 * u.norm :=
  by
  obtain ⟨u, h2⟩ := step1 hcoprime heven
  set q : ℤ√(-3) := ⟨1, 1⟩
  replace h2 : a = q * u ∨ a = star q * u := h2
  obtain ⟨hp', hq'⟩ := Spts.four hp hq
  refine' ⟨p.re * u, _, _, _⟩
  · rw [← Int.isUnit_iff_abs_eq] at hp'
    rw [Zsqrtd.smul_re, Zsqrtd.smul_im, isCoprime_mul_unit_left hp']
    apply Zsqrtd.coprime_of_dvd_coprime hcoprime
    obtain rfl | rfl := h2 <;> apply dvd_mul_left
  · rw [abs_eq <| zero_le_one' ℤ] at hp' hq'
    cases hp' with
    | inl hp' =>
      have h4 : p = q ∨ p = star q := by
        apply Or.imp _ _ hq' <;>
          · intro h5
            simp [Zsqrtd.ext, hp', h5]
      simp only [hp', one_mul, Int.cast_one]
      cases h4 <;> cases h2 <;> simp [*, or_comm]
    | inr hp' =>
      have h4 : p = -q ∨ p = -star q := by
        apply Or.imp _ _ hq'.symm
        · intro h5
          simp [Zsqrtd.ext, hp', h5]
        · intro h5
          simp [Zsqrtd.ext, hp', h5]
      simp only [hp', one_mul, Zsqrtd.norm_neg, Int.cast_one, Int.cast_neg, neg_one_mul]
      cases h4 <;> cases h2 <;> simp [*]
  · rw [Zsqrtd.norm_mul, Zsqrtd.norm_int_cast, ← sq, ← sq_abs, hp', one_pow, one_mul]
    cases' h2 with h2 h2 <;>
      · rw [h2, Zsqrtd.norm_mul]
        congr
#align step1'' step1''

theorem step2 {a p : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) (hdvd : p.norm ∣ a.norm)
    (hpprime : Prime p.norm) :
    ∃ u : ℤ√(-3), IsCoprime u.re u.im ∧ (a = p * u ∨ a = star p * u) ∧ a.norm = p.norm * u.norm :=
  by
  obtain ⟨u', h, h'⟩ := Spts.mul_of_dvd'' hdvd hpprime
  refine' ⟨u', _, h, h'⟩
  apply Zsqrtd.coprime_of_dvd_coprime hcoprime
  obtain rfl | rfl := h <;> apply dvd_mul_left
#align step2 step2

theorem step1_2 {a p : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) (hdvd : p.norm ∣ a.norm)
    (hp : OddPrimeOrFour p.norm) (hq : p.im ≠ 0) :
    ∃ u : ℤ√(-3), IsCoprime u.re u.im ∧ (a = p * u ∨ a = star p * u) ∧ a.norm = p.norm * u.norm :=
  by
  obtain hp | ⟨hpprime, -⟩ := hp
  · rw [hp] at hdvd⊢
    have heven : Even a.norm :=
      by
      apply even_iff_two_dvd.mpr (dvd_trans _ hdvd)
      norm_num
    exact step1'' hcoprime hp hq heven
  · apply step2 hcoprime hdvd hpprime
#align step1_2 step1_2

theorem OddPrimeOrFour.factors' {a : ℤ√(-3)} (h2 : a.norm ≠ 1) (hcoprime : IsCoprime a.re a.im) :
    ∃ u q : ℤ√(-3),
      0 ≤ q.re ∧
        q.im ≠ 0 ∧ OddPrimeOrFour q.norm ∧ a = q * u ∧ IsCoprime u.re u.im ∧ u.norm < a.norm :=
  by
  have h2 : 2 < a.norm := by
    apply lt_of_le_of_ne _ (Spts.not_two _).symm
    rw [← one_mul (2 : ℤ), mul_two, Int.add_one_le_iff]
    apply lt_of_le_of_ne _ h2.symm
    rw [← Int.sub_one_lt_iff, sub_self]
    exact Spts.pos_of_coprime' hcoprime
  obtain ⟨p, hpdvd, hp⟩ := OddPrimeOrFour.exists_and_dvd h2
  obtain ⟨q', hcd, hc, hd⟩ := OddPrimeOrFour.factors a p hcoprime hp hpdvd
  rw [← abs_dvd, hcd] at hpdvd
  have hp' := hp.abs
  rw [hcd] at hp'
  obtain ⟨u, huvcoprime, huv, huvdvd⟩ := step1_2 hcoprime hpdvd hp' hd
  use u
  cases' huv with huv huv <;> [(use q'), (use star q')] <;>
    simp only [IsCoprime, Zsqrtd.star_re, Zsqrtd.star_im, Ne.def, neg_eq_zero, Zsqrtd.norm_conj] <;>
    use hc, hd, hp', huv, huvcoprime <;>
    · rw [huvdvd, lt_mul_iff_one_lt_left (Spts.pos_of_coprime' huvcoprime), ← hcd]
      exact hp.one_lt_abs
#align odd_prime_or_four.factors' OddPrimeOrFour.factors'

theorem step3 {a : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) :
    ∃ f : Multiset (ℤ√(-3)),
      (a = f.prod ∨ a = -f.prod) ∧
        ∀ pq : ℤ√(-3), pq ∈ f → 0 ≤ pq.re ∧ pq.im ≠ 0 ∧ OddPrimeOrFour pq.norm :=
  by
  suffices
    ∀ n : ℕ,
      a.norm.natAbs = n →
        ∃ f : Multiset (ℤ√(-3)),
          (a = f.prod ∨ a = -f.prod) ∧
            ∀ pq : ℤ√(-3), pq ∈ f → 0 ≤ pq.re ∧ pq.im ≠ 0 ∧ OddPrimeOrFour pq.norm
    by exact this a.norm.natAbs rfl
  intro n hn
  induction' n using Nat.strong_induction_on with n ih generalizing a
  dsimp only at ih
  cases' eq_or_ne a.norm 1 with h h
  · use 0
    constructor
    · simp only [Multiset.prod_zero, Spts.eq_one' h]
    · intro pq hpq
      exact absurd hpq (Multiset.not_mem_zero _)
  · obtain ⟨U, q, hc, hd, hp, huv, huvcoprime, descent⟩ := OddPrimeOrFour.factors' h hcoprime
    replace descent :=
      Int.natAbs_lt_natAbs_of_nonneg_of_lt (Zsqrtd.norm_nonneg (by norm_num) _) descent
    rw [hn] at descent
    obtain ⟨g, hgprod, hgfactors⟩ := ih U.norm.natAbs descent huvcoprime rfl
    refine' ⟨q ::ₘ g, _, _⟩
    · simp only [huv, Multiset.prod_cons, ← mul_neg]
      exact Or.imp (congr_arg _) (congr_arg _) hgprod
    · rintro pq hpq
      rw [Multiset.mem_cons] at hpq
      obtain rfl | ind := hpq
      · exact ⟨hc, hd, hp⟩
      · exact hgfactors pq ind
#align step3 step3

theorem step4_3 {p p' : ℤ√(-3)} (hcoprime : IsCoprime p.re p.im) (hcoprime' : IsCoprime p'.re p'.im)
    (h : OddPrimeOrFour p.norm) (heq : p.norm = p'.norm) :
    abs p.re = abs p'.re ∧ abs p.im = abs p'.im :=
  by
  have hdvd : p'.norm ∣ p.norm := by rw [heq]
  have him : p'.im ≠ 0 := by
    apply OddPrimeOrFour.im_ne_zero _ hcoprime'
    rwa [← heq]
  obtain ⟨u, -, H | H, h1⟩ := step1_2 hcoprime hdvd (by rwa [← heq]) him <;>
    · rw [heq] at h1
      have := Int.eq_one_of_mul_eq_self_right (Spts.ne_zero_of_coprime' _ hcoprime') h1.symm
      obtain ⟨ha, hb⟩ := Spts.eq_one this
      simp only [hb, Zsqrtd.ext, Zsqrtd.mul_re, Zsqrtd.mul_im, add_zero, zero_add,
        MulZeroClass.mul_zero] at H
      rw [H.1, H.2, abs_mul, abs_mul, ha, mul_one, mul_one]
      try rw [zsqrtd.conj_re, zsqrtd.conj_im, abs_neg]
      simp
#align step4_3 step4_3

theorem prod_map_norm {d : ℤ} {s : Multiset (ℤ√d)} :
    (s.map Zsqrtd.norm).prod = Zsqrtd.norm s.prod :=
  Multiset.prod_hom s (Zsqrtd.normMonoidHom : ℤ√d →* ℤ)
#align prod_map_norm prod_map_norm

theorem factors_unique_prod :
    ∀ {f g : Multiset (ℤ√(-3))},
      (∀ x ∈ f, OddPrimeOrFour (Zsqrtd.norm x)) →
        (∀ x ∈ g, OddPrimeOrFour (Zsqrtd.norm x)) →
          Associated f.prod.norm g.prod.norm →
            Multiset.Rel Associated (f.map Zsqrtd.norm) (g.map Zsqrtd.norm) :=
  by
  intro f g hf hg hassoc
  apply factors_unique_prod'
  · intro x hx
    rw [Multiset.mem_map] at hx
    obtain ⟨y, hy, rfl⟩ := hx
    exact hf y hy
  · intro x hx
    rw [Multiset.mem_map] at hx
    obtain ⟨y, hy, rfl⟩ := hx
    exact hg y hy
  · rwa [prod_map_norm, prod_map_norm]
#align factors_unique_prod factors_unique_prod

/-- The multiset of factors. -/
noncomputable def factorization' {a : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) :
    Multiset (ℤ√(-3)) :=
  Classical.choose (step3 hcoprime)
#align factorization' factorization'

theorem factorization'_prop {a : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) :
    (a = (factorization' hcoprime).prod ∨ a = -(factorization' hcoprime).prod) ∧
      ∀ pq : ℤ√(-3),
        pq ∈ factorization' hcoprime → 0 ≤ pq.re ∧ pq.im ≠ 0 ∧ OddPrimeOrFour pq.norm :=
  Classical.choose_spec (step3 hcoprime)
#align factorization'_prop factorization'_prop

theorem factorization'.coprime_of_mem {a b : ℤ√(-3)} (h : IsCoprime a.re a.im)
    (hmem : b ∈ factorization' h) : IsCoprime b.re b.im :=
  by
  obtain ⟨h1, -⟩ := factorization'_prop h
  set f := factorization' h
  apply Zsqrtd.coprime_of_dvd_coprime h
  apply (Multiset.dvd_prod hmem).trans
  cases' h1 with h1 h1 <;> rw [h1]
  simp only [dvd_neg]
  rfl
#align factorization'.coprime_of_mem factorization'.coprime_of_mem

theorem no_conj (s : Multiset (ℤ√(-3))) {p : ℤ√(-3)} (hp : ¬IsUnit p)
    (hcoprime : IsCoprime s.prod.re s.prod.im) : ¬(p ∈ s ∧ star p ∈ s) :=
  by
  contrapose! hp
  obtain ⟨h1, h2⟩ := hp
  by_cases him : p.im = 0
  · obtain ⟨t, rfl⟩ := Multiset.exists_cons_of_mem h1
    rw [Multiset.prod_cons] at hcoprime
    simp only [him, add_zero, MulZeroClass.zero_mul, Zsqrtd.mul_im, Zsqrtd.mul_re,
      MulZeroClass.mul_zero] at hcoprime
    rw [Zsqrt3.isUnit_iff]
    simp only [him, and_true_iff, eq_self_iff_true]
    rw [← Int.isUnit_iff_abs_eq]
    apply IsCoprime.isUnit_of_dvd' hcoprime <;> apply dvd_mul_right
  · have : star p ≠ p := by
      rw [Ne.def, Zsqrtd.ext]
      rintro ⟨-, H⟩
      apply him
      apply eq_zero_of_neg_eq
      simpa using H
    obtain ⟨t1, rfl⟩ := Multiset.exists_cons_of_mem h1
    have : star p ∈ t1 := by
      rw [Multiset.mem_cons] at h2
      exact h2.resolve_left this
    obtain ⟨t2, rfl⟩ := Multiset.exists_cons_of_mem this
    rw [Multiset.prod_cons, Multiset.prod_cons, ← mul_assoc, ← Zsqrtd.norm_eq_mul_conj] at hcoprime
    rw [Zsqrtd.isUnit_iff_norm_isUnit]
    apply IsCoprime.isUnit_of_dvd' hcoprime <;>
      simp only [add_zero, Zsqrtd.coe_int_re, MulZeroClass.zero_mul, dvd_mul_right, Zsqrtd.mul_re,
        MulZeroClass.mul_zero, Zsqrtd.mul_im, Zsqrtd.coe_int_im]
#align no_conj no_conj

/-- Associated elements in `ℤ√-3`. -/
def Associated' (x y : ℤ√(-3)) : Prop :=
  Associated x y ∨ Associated x (star y)
#align associated' Associated'

@[refl]
theorem Associated'.refl (x : ℤ√(-3)) : Associated' x x :=
  Or.inl (by rfl)
#align associated'.refl Associated'.refl

theorem Associated'.norm_eq {x y : ℤ√(-3)} (h : Associated' x y) : x.norm = y.norm := by
  cases' h with h h <;> simp only [Zsqrtd.norm_eq_of_associated (by norm_num) h, Zsqrtd.norm_conj]
#align associated'.norm_eq Associated'.norm_eq

theorem associated'_of_abs_eq {x y : ℤ√(-3)} (hre : abs x.re = abs y.re)
    (him : abs x.im = abs y.im) : Associated' x y :=
  by
  rw [abs_eq_abs] at hre him
  cases' hre with h1 h1 <;> cases' him with h2 h2 <;>
    [
      (left; use 1),
      (right; use 1),
      (right; use -1),
      (left; use -1)
    ] <;>
    simp [Zsqrtd.ext, h1, h2]
#align associated'_of_abs_eq associated'_of_abs_eq

theorem associated'_of_associated_norm {x y : ℤ√(-3)}
    (h : Associated (Zsqrtd.norm x) (Zsqrtd.norm y)) (hx : IsCoprime x.re x.im)
    (hy : IsCoprime y.re y.im) (h' : OddPrimeOrFour x.norm) : Associated' x y :=
  by
  have heq : x.norm = y.norm :=
    haveI hd : (-3 : ℤ) ≤ 0 := by norm_num
    Int.eq_of_associated_of_nonneg h (Zsqrtd.norm_nonneg hd _) (Zsqrtd.norm_nonneg hd _)
  obtain ⟨hre, him⟩ := step4_3 hx hy h' heq
  exact associated'_of_abs_eq hre him
#align associated'_of_associated_norm associated'_of_associated_norm

theorem factorization'.associated'_of_norm_eq {a b c : ℤ√(-3)} (h : IsCoprime a.re a.im)
    (hbmem : b ∈ factorization' h) (hcmem : c ∈ factorization' h) (hnorm : b.norm = c.norm) :
    Associated' b c := by
  apply associated'_of_associated_norm
  · rw [hnorm]
  · exact factorization'.coprime_of_mem h hbmem
  · exact factorization'.coprime_of_mem h hcmem
  · exact ((factorization'_prop h).2 b hbmem).2.2
#align factorization'.associated'_of_norm_eq factorization'.associated'_of_norm_eq

theorem factors_unique {f g : Multiset (ℤ√(-3))} (hf : ∀ x ∈ f, OddPrimeOrFour (Zsqrtd.norm x))
    (hf' : ∀ x ∈ f, IsCoprime (Zsqrtd.re x) (Zsqrtd.im x))
    (hg : ∀ x ∈ g, OddPrimeOrFour (Zsqrtd.norm x))
    (hg' : ∀ x ∈ g, IsCoprime (Zsqrtd.re x) (Zsqrtd.im x)) (h : Associated f.prod g.prod) :
    Multiset.Rel Associated' f g :=
  by
  have p :
    ∀ (x : ℤ√(-3)) (_ : x ∈ f) (y : ℤ√(-3)) (_ : y ∈ g),
      Associated x.norm y.norm → Associated' x y :=
    by
    intro x hx y hy hxy
    apply associated'_of_associated_norm hxy
    · exact hf' x hx
    · exact hg' y hy
    · exact hf x hx
  refine' Multiset.Rel.mono _ p
  rw [← Multiset.rel_map]
  apply factors_unique_prod hf hg
  exact Zsqrtd.associated_norm_of_associated h
#align factors_unique factors_unique

theorem factors_2_even' {a : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) :
    Even (evenFactorExp a.norm) :=
  by
  induction' hn : a.norm.natAbs using Nat.strong_induction_on with n ih generalizing a
  dsimp only at ih
  by_cases hparity : Even a.norm
  · obtain ⟨u, huvcoprime, huvprod⟩ := step1' hcoprime hparity
    have huv := Spts.ne_zero_of_coprime' _ huvcoprime
    rw [huvprod, factors_2_even huv, Nat.even_add]
    apply iff_of_true even_two
    apply ih _ _ huvcoprime rfl
    rw [← hn, huvprod, Int.natAbs_mul, lt_mul_iff_one_lt_left (Int.natAbs_pos.mpr huv)]
    norm_num
  · convert even_zero (α := ℕ)
    simp only [evenFactorExp, Multiset.count_eq_zero, hn]
    contrapose! hparity with hfactor
    rw [even_iff_two_dvd]
    exact UniqueFactorizationMonoid.dvd_of_mem_normalizedFactors hfactor
#align factors_2_even' factors_2_even'

theorem factorsOddPrimeOrFour.unique {a : ℤ√(-3)} (hcoprime : IsCoprime a.re a.im) {f : Multiset ℤ}
    (hf : ∀ x ∈ f, OddPrimeOrFour x) (hf' : ∀ x ∈ f, (0 : ℤ) ≤ x)
    (hassoc : Associated f.prod a.norm) : f = factorsOddPrimeOrFour a.norm :=
  factorsOddPrimeOrFour.unique' hf hf' (Spts.pos_of_coprime' hcoprime) (factors_2_even' hcoprime)
    hassoc
#align factors_odd_prime_or_four.unique factorsOddPrimeOrFour.unique

theorem eq_or_eq_conj_of_associated_of_re_zero {x A : ℤ√(-3)} (hx : x.re = 0) (h : Associated x A) :
    x = A ∨ x = star A := by
  obtain ⟨u, hu⟩ := h
  obtain ⟨v, hv1, hv2⟩ := Zsqrt3.coe_of_isUnit' u.isUnit
  have hA : A.re = 0 := by
    simp only [← hu, hv1, hx, add_zero, MulZeroClass.zero_mul, Zsqrtd.mul_re, MulZeroClass.mul_zero,
      Zsqrtd.coe_int_im]
  rw [abs_eq <| zero_le_one' ℤ] at hv2
  cases' hv2 with hv2 hv2
  · left
    simpa only [hv1, hv2, mul_one, Int.cast_one] using hu
  · right
    simp only [hv1, hv2, mul_one, Int.cast_one, mul_neg, Int.cast_neg] at hu
    simp [← hu, hx, Zsqrtd.ext]
#align eq_or_eq_conj_of_associated_of_re_zero eq_or_eq_conj_of_associated_of_re_zero

theorem eq_or_eq_conj_iff_associated'_of_nonneg {x A : ℤ√(-3)} (hx : 0 ≤ x.re) (hA : 0 ≤ A.re) :
    Associated' x A ↔ x = A ∨ x = star A := by
  constructor
  · rintro (⟨u, hu⟩ | ⟨u, hu⟩) <;> obtain ⟨v, hv1, hv2⟩ := Zsqrt3.coe_of_isUnit' u.isUnit
    -- associated x A
    · by_cases hxre : x.re = 0
      · apply eq_or_eq_conj_of_associated_of_re_zero hxre ⟨u, hu⟩
      · rw [hv1] at hu
        rw [abs_eq <| zero_le_one' ℤ] at hv2
        cases' hv2 with habsv habsv
        · left
          rw [← hu, habsv, Int.cast_one, mul_one]
        · exfalso
          simp only [habsv, mul_one, Int.cast_one, mul_neg, Int.cast_neg] at hu
          apply lt_irrefl (0 : ℤ)
          calc
            0 ≤ A.re := hA
            _ = -x.re := ?_
            _ < 0 := ?_

          · simp only [← hu, Zsqrtd.neg_re]
          · simp only [neg_neg_iff_pos]
            exact lt_of_le_of_ne hx (Ne.symm hxre)
    -- associated x A.conj
    · by_cases hxre : x.re = 0
      · convert(eq_or_eq_conj_of_associated_of_re_zero hxre ⟨u, hu⟩).symm
        rw [star_star]
      · rw [hv1] at hu
        rw [abs_eq <| zero_le_one' ℤ] at hv2
        cases' hv2 with habsv habsv
        · right
          rw [← hu, habsv, Int.cast_one, mul_one]
        · exfalso
          simp only [habsv, mul_one, Int.cast_one, mul_neg, Int.cast_neg] at hu
          apply lt_irrefl (0 : ℤ)
          calc
            0 ≤ A.re := hA
            _ = -x.re := ?_
            _ < 0 := ?_

          · rw [← star_star A, ← hu]
            simp only [Zsqrtd.neg_re, Zsqrtd.star_re]
          · simp only [neg_neg_iff_pos]
            apply lt_of_le_of_ne hx (Ne.symm hxre)
  · rintro (rfl | rfl)
    · rfl
    · right
      rfl
#align eq_or_eq_conj_iff_associated'_of_nonneg eq_or_eq_conj_iff_associated'_of_nonneg

theorem step5'
    -- lemma page 54
    (a : ℤ√(-3))
    (r : ℤ) (hcoprime : IsCoprime a.re a.im) (hcube : r ^ 3 = a.norm) :
    ∃ g : Multiset (ℤ√(-3)), factorization' hcoprime = 3 • g := by
  classical
    obtain ⟨h1, h2⟩ := factorization'_prop hcoprime
    set f := factorization' hcoprime with hf
    apply Multiset.exists_smul_of_dvd_count
    intro x hx
    convert dvd_mul_right 3 (Multiset.count (Zsqrtd.norm x) (factorsOddPrimeOrFour r))
    have : Even (evenFactorExp r) :=
      by
      suffices Even (3 * evenFactorExp r)
        by
        rw [Nat.even_mul] at this
        apply this.resolve_left
        norm_num
      rw [← evenFactorExp.pow r 3, hcube]
      exact factors_2_even' hcoprime
    calc
      Multiset.count x f = Multiset.card (Multiset.filter (x = ·) f) := by
        rw [Multiset.count, Multiset.countp_eq_card_filter]
      _ = Multiset.card (Multiset.filter (fun a : ℤ√(-3) => x.norm = a.norm) f) :=
        (congr_arg _ (Multiset.filter_congr fun A HA => ?_))
      _ = Multiset.countp (x.norm = ·) (Multiset.map Zsqrtd.norm f) :=
        (Multiset.countp_map Zsqrtd.norm f (x.norm = ·)).symm
      _ = Multiset.countp (x.norm = ·) (factorsOddPrimeOrFour a.norm) := ?_
      _ = Multiset.count x.norm (factorsOddPrimeOrFour a.norm) := by rw [Multiset.count]
      _ = Multiset.count x.norm (factorsOddPrimeOrFour (r ^ 3)) := by rw [hcube]
      _ = Multiset.count x.norm (3 • _) := (congr_arg _ <| factorsOddPrimeOrFour.pow _ _ this)
      _ = 3 * _ := Multiset.count_nsmul x.norm _ _

    swap
    show
      Multiset.countp (Eq x.norm) (Multiset.map Zsqrtd.norm f) =
        Multiset.countp (Eq x.norm) (factorsOddPrimeOrFour a.norm)
    congr
    · apply factorsOddPrimeOrFour.unique hcoprime
      · intro x hx
        obtain ⟨y, hy, rfl⟩ := Multiset.mem_map.mp hx
        exact (h2 y hy).2.2
      · intro x hx
        obtain ⟨y, -, rfl⟩ := Multiset.mem_map.mp hx
        apply Zsqrtd.norm_nonneg
        norm_num
      · rw [prod_map_norm]
        cases' h1 with h1 h1 <;> rw [h1]
        rw [Zsqrtd.norm_neg]
    have h2x := h2 x hx
    refine' ⟨congr_arg _, fun h => _⟩
    have hassoc := factorization'.associated'_of_norm_eq hcoprime hx HA h
    have eq_or_eq_conj := (eq_or_eq_conj_iff_associated'_of_nonneg h2x.1 (h2 A HA).1).mp hassoc
    refine' eq_or_eq_conj.resolve_right fun H => _
    apply no_conj f
    · intro HH
      apply h2x.2.1
      rw [Zsqrt3.isUnit_iff] at HH
      exact HH.2
    · cases' h1 with h1 h1 <;> rw [h1] at hcoprime
      · exact hcoprime
      · rwa [← IsCoprime.neg_neg_iff, ← Zsqrtd.neg_im, ← Zsqrtd.neg_re]
    · refine' ⟨hx, _⟩
      rwa [H, star_star]
#align step5' step5'

theorem step5
    -- lemma page 54
    (a : ℤ√(-3))
    (r : ℤ) (hcoprime : IsCoprime a.re a.im) (hcube : r ^ 3 = a.norm) : ∃ p : ℤ√(-3), a = p ^ 3 :=
  by
  obtain ⟨f, hf⟩ := step5' a r hcoprime hcube
  obtain ⟨h1, -⟩ := factorization'_prop hcoprime
  cases h1 with
  | inl h1 =>
    use f.prod
    rw [h1, hf, Multiset.prod_nsmul]
  | inr h1 =>
    use -f.prod
    rw [h1, hf, Multiset.prod_nsmul, Odd.neg_pow]
    norm_num
#align step5 step5

theorem step6 (a b r : ℤ) (hcoprime : IsCoprime a b) (hcube : r ^ 3 = a ^ 2 + 3 * b ^ 2) :
    ∃ p q, a = p ^ 3 - 9 * p * q ^ 2 ∧ b = 3 * p ^ 2 * q - 3 * q ^ 3 :=
  by
  set A : ℤ√(-3) := ⟨a, b⟩ with hA
  have hcube' : r ^ 3 = A.norm :=
    by
    simp only [hcube, Zsqrtd.norm_def, hA]
    ring
  obtain ⟨p, hp⟩ := step5 A r hcoprime hcube'
  use p.re, p.im
  simp only [Zsqrtd.ext, pow_succ', pow_two, Zsqrtd.mul_re, Zsqrtd.mul_im] at hp
  convert hp using 2 <;> simp <;> ring
#align step6 step6


/-
Copyright (c) 2021 Eric Wieser. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Eric Wieser
-/
import algebra.group.inj_surj
import data.list.big_operators.basic
import data.list.fin_range
import group_theory.group_action.defs
import group_theory.submonoid.basic
import data.set_like.basic
import data.sigma.basic

/-!
# Additively-graded multiplicative structures

> THIS FILE IS SYNCHRONIZED WITH MATHLIB4.
> Any changes to this file require a corresponding PR to mathlib4.

This module provides a set of heterogeneous typeclasses for defining a multiplicative structure
over the sigma type `graded_monoid A` such that `(*) : A i → A j → A (i + j)`; that is to say, `A`
forms an additively-graded monoid. The typeclasses are:

* `graded_monoid.ghas_one A`
* `graded_monoid.ghas_mul A`
* `graded_monoid.gmonoid A`
* `graded_monoid.gcomm_monoid A`

With the `sigma_graded` locale open, these respectively imbue:

* `has_one (graded_monoid A)`
* `has_mul (graded_monoid A)`
* `monoid (graded_monoid A)`
* `comm_monoid (graded_monoid A)`

the base type `A 0` with:

* `graded_monoid.grade_zero.has_one`
* `graded_monoid.grade_zero.has_mul`
* `graded_monoid.grade_zero.monoid`
* `graded_monoid.grade_zero.comm_monoid`

and the `i`th grade `A i` with `A 0`-actions (`•`) defined as left-multiplication:

* (nothing)
* `graded_monoid.grade_zero.has_smul (A 0)`
* `graded_monoid.grade_zero.mul_action (A 0)`
* (nothing)

For now, these typeclasses are primarily used in the construction of `direct_sum.ring` and the rest
of that file.

## Dependent graded products

This also introduces `list.dprod`, which takes the (possibly non-commutative) product of a list
of graded elements of type `A i`. This definition primarily exist to allow `graded_monoid.mk`
and `direct_sum.of` to be pulled outside a product, such as in `graded_monoid.mk_list_dprod` and
`direct_sum.of_list_dprod`.

## Internally graded monoids

In addition to the above typeclasses, in the most frequent case when `A` is an indexed collection of
`set_like` subobjects (such as `add_submonoid`s, `add_subgroup`s, or `submodule`s), this file
provides the `Prop` typeclasses:

* `set_like.has_graded_one A` (which provides the obvious `graded_monoid.ghas_one A` instance)
* `set_like.has_graded_mul A` (which provides the obvious `graded_monoid.ghas_mul A` instance)
* `set_like.graded_monoid A` (which provides the obvious `graded_monoid.gmonoid A` and
  `graded_monoid.gcomm_monoid A` instances)

which respectively provide the API lemmas

* `set_like.one_mem_graded`
* `set_like.mul_mem_graded`
* `set_like.pow_mem_graded`, `set_like.list_prod_map_mem_graded`

Strictly this last class is unecessary as it has no fields not present in its parents, but it is
included for convenience. Note that there is no need for `set_like.graded_ring` or similar, as all
the information it would contain is already supplied by `graded_monoid` when `A` is a collection
of objects satisfying `add_submonoid_class` such as `submodule`s. These constructions are explored
in `algebra.direct_sum.internal`.

This file also defines:

* `set_like.is_homogeneous A` (which says that `a` is homogeneous iff `a ∈ A i` for some `i : ι`)
* `set_like.homogeneous_submonoid A`, which is, as the name suggests, the submonoid consisting of
  all the homogeneous elements.

## tags

graded monoid
-/

set_option old_structure_cmd true

variables {ι : Type*}

/-- A type alias of sigma types for graded monoids. -/
def graded_monoid (A : ι → Type*) := sigma A

namespace graded_monoid

instance {A : ι → Type*} [inhabited ι] [inhabited (A default)]: inhabited (graded_monoid A) :=
sigma.inhabited

/-- Construct an element of a graded monoid. -/
def mk {A : ι → Type*} : Π i, A i → graded_monoid A := sigma.mk

/-! ### Typeclasses -/
section defs

variables (A : ι → Type*)

/-- A graded version of `has_one`, which must be of grade 0. -/
class ghas_one [has_zero ι] :=
(one : A 0)

/-- `ghas_one` implies `has_one (graded_monoid A)` -/
instance ghas_one.to_has_one [has_zero ι] [ghas_one A] : has_one (graded_monoid A) :=
⟨⟨_, ghas_one.one⟩⟩

/-- A graded version of `has_mul`. Multiplication combines grades additively, like
`add_monoid_algebra`. -/
class ghas_mul [has_add ι] :=
(mul {i j} : A i → A j → A (i + j))

/-- `ghas_mul` implies `has_mul (graded_monoid A)`. -/
instance ghas_mul.to_has_mul [has_add ι] [ghas_mul A] :
  has_mul (graded_monoid A) :=
⟨λ (x y : graded_monoid A), ⟨_, ghas_mul.mul x.snd y.snd⟩⟩

lemma mk_mul_mk [has_add ι] [ghas_mul A] {i j} (a : A i) (b : A j) :
  mk i a * mk j b = mk (i + j) (ghas_mul.mul a b) :=
rfl

namespace gmonoid

variables {A} [add_monoid ι] [ghas_mul A] [ghas_one A]

/-- A default implementation of power on a graded monoid, like `npow_rec`.
`gmonoid.gnpow` should be used instead. -/
def gnpow_rec : Π (n : ℕ) {i}, A i → A (n • i)
| 0 i a := cast (congr_arg A (zero_nsmul i).symm) ghas_one.one
| (n + 1) i a := cast (congr_arg A (succ_nsmul i n).symm) (ghas_mul.mul a $ gnpow_rec _ a)

@[simp] lemma gnpow_rec_zero (a : graded_monoid A) : graded_monoid.mk _ (gnpow_rec 0 a.snd) = 1 :=
sigma.ext (zero_nsmul _) (heq_of_cast_eq _ rfl).symm

/-- Tactic used to autofill `graded_monoid.gmonoid.gnpow_zero'` when the default
`graded_monoid.gmonoid.gnpow_rec` is used. -/
meta def apply_gnpow_rec_zero_tac : tactic unit := `[apply graded_monoid.gmonoid.gnpow_rec_zero]

@[simp] lemma gnpow_rec_succ (n : ℕ) (a : graded_monoid A) :
  (graded_monoid.mk _ $ gnpow_rec n.succ a.snd) = a * ⟨_, gnpow_rec n a.snd⟩ :=
sigma.ext (succ_nsmul _ _) (heq_of_cast_eq _ rfl).symm

/-- Tactic used to autofill `graded_monoid.gmonoid.gnpow_succ'` when the default
`graded_monoid.gmonoid.gnpow_rec` is used. -/
meta def apply_gnpow_rec_succ_tac : tactic unit := `[apply graded_monoid.gmonoid.gnpow_rec_succ]

end gmonoid

/-- A graded version of `monoid`.

Like `monoid.npow`, this has an optional `gmonoid.gnpow` field to allow definitional control of
natural powers of a graded monoid. -/
class gmonoid [add_monoid ι]  extends ghas_mul A, ghas_one A :=
(one_mul (a : graded_monoid A) : 1 * a = a)
(mul_one (a : graded_monoid A) : a * 1 = a)
(mul_assoc (a b c : graded_monoid A) : a * b * c = a * (b * c))
(gnpow : Π (n : ℕ) {i}, A i → A (n • i) := gmonoid.gnpow_rec)
(gnpow_zero' : Π (a : graded_monoid A), graded_monoid.mk _ (gnpow 0 a.snd) = 1
  . gmonoid.apply_gnpow_rec_zero_tac)
(gnpow_succ' : Π (n : ℕ) (a : graded_monoid A),
  (graded_monoid.mk _ $ gnpow n.succ a.snd) = a * ⟨_, gnpow n a.snd⟩
  . gmonoid.apply_gnpow_rec_succ_tac)

/-- `gmonoid` implies a `monoid (graded_monoid A)`. -/
instance gmonoid.to_monoid [add_monoid ι] [gmonoid A] :
  monoid (graded_monoid A) :=
{ one := (1), mul := (*),
  npow := λ n a, graded_monoid.mk _ (gmonoid.gnpow n a.snd),
  npow_zero' := λ a, gmonoid.gnpow_zero' a,
  npow_succ' := λ n a, gmonoid.gnpow_succ' n a,
  one_mul := gmonoid.one_mul, mul_one := gmonoid.mul_one, mul_assoc := gmonoid.mul_assoc }

lemma mk_pow [add_monoid ι] [gmonoid A] {i} (a : A i) (n : ℕ) :
  mk i a ^ n = mk (n • i) (gmonoid.gnpow _ a) :=
begin
  induction n with n,
  { rw [pow_zero],
    exact (gmonoid.gnpow_zero' ⟨_, a⟩).symm, },
  { rw [pow_succ, n_ih, mk_mul_mk],
    exact (gmonoid.gnpow_succ' n ⟨_, a⟩).symm, },
end

/-- A graded version of `comm_monoid`. -/
class gcomm_monoid [add_comm_monoid ι] extends gmonoid A :=
(mul_comm (a : graded_monoid A) (b : graded_monoid A) : a * b = b * a)

/-- `gcomm_monoid` implies a `comm_monoid (graded_monoid A)`, although this is only used as an
instance locally to define notation in `gmonoid` and similar typeclasses. -/
instance gcomm_monoid.to_comm_monoid [add_comm_monoid ι] [gcomm_monoid A] :
  comm_monoid (graded_monoid A) :=
{ mul_comm := gcomm_monoid.mul_comm, ..gmonoid.to_monoid A }

end defs

/-! ### Instances for `A 0`

The various `g*` instances are enough to promote the `add_comm_monoid (A 0)` structure to various
types of multiplicative structure.
-/

section grade_zero

variables (A : ι → Type*)

section one
variables [has_zero ι] [ghas_one A]

/-- `1 : A 0` is the value provided in `ghas_one.one`. -/
@[nolint unused_arguments]
instance grade_zero.has_one : has_one (A 0) :=
⟨ghas_one.one⟩

end one

section mul
variables [add_zero_class ι] [ghas_mul A]

/-- `(•) : A 0 → A i → A i` is the value provided in `graded_monoid.ghas_mul.mul`, composed with
an `eq.rec` to turn `A (0 + i)` into `A i`.
-/
instance grade_zero.has_smul (i : ι) : has_smul (A 0) (A i) :=
{ smul := λ x y, (zero_add i).rec (ghas_mul.mul x y) }

/-- `(*) : A 0 → A 0 → A 0` is the value provided in `graded_monoid.ghas_mul.mul`, composed with
an `eq.rec` to turn `A (0 + 0)` into `A 0`.
-/
instance grade_zero.has_mul : has_mul (A 0) :=
{ mul := (•) }

variables {A}

@[simp] lemma mk_zero_smul {i} (a : A 0) (b : A i) : mk _ (a • b) = mk _ a * mk _ b :=
sigma.ext (zero_add _).symm $ eq_rec_heq _ _

@[simp] lemma grade_zero.smul_eq_mul (a b : A 0) : a • b = a * b := rfl

end mul

section monoid
variables [add_monoid ι] [gmonoid A]

instance : has_pow (A 0) ℕ :=
{ pow := λ x n, (nsmul_zero n).rec (gmonoid.gnpow n x : A (n • 0)) }

variables {A}

@[simp] lemma mk_zero_pow (a : A 0) (n : ℕ) : mk _ (a ^ n) = mk _ a ^ n :=
sigma.ext (nsmul_zero n).symm $ eq_rec_heq _ _

variables (A)

/-- The `monoid` structure derived from `gmonoid A`. -/
instance grade_zero.monoid : monoid (A 0) :=
function.injective.monoid (mk 0) sigma_mk_injective rfl mk_zero_smul mk_zero_pow

end monoid

section monoid
variables [add_comm_monoid ι] [gcomm_monoid A]

/-- The `comm_monoid` structure derived from `gcomm_monoid A`. -/
instance grade_zero.comm_monoid : comm_monoid (A 0) :=
function.injective.comm_monoid (mk 0) sigma_mk_injective rfl mk_zero_smul mk_zero_pow

end monoid

section mul_action
variables [add_monoid ι] [gmonoid A]

/-- `graded_monoid.mk 0` is a `monoid_hom`, using the `graded_monoid.grade_zero.monoid` structure.
-/
def mk_zero_monoid_hom : A 0 →* (graded_monoid A) :=
{ to_fun := mk 0, map_one' := rfl, map_mul' := mk_zero_smul }

/-- Each grade `A i` derives a `A 0`-action structure from `gmonoid A`. -/
instance grade_zero.mul_action {i} : mul_action (A 0) (A i) :=
begin
  letI := mul_action.comp_hom (graded_monoid A) (mk_zero_monoid_hom A),
  exact function.injective.mul_action (mk i) sigma_mk_injective mk_zero_smul,
end

end mul_action

end grade_zero

end graded_monoid

/-! ### Dependent products of graded elements -/

section dprod

variables {α : Type*} {A : ι → Type*} [add_monoid ι] [graded_monoid.gmonoid A]

/-- The index used by `list.dprod`. Propositionally this is equal to `(l.map fι).sum`, but
definitionally it needs to have a different form to avoid introducing `eq.rec`s in `list.dprod`. -/
def list.dprod_index (l : list α) (fι : α → ι) : ι :=
l.foldr (λ i b, fι i + b) 0

@[simp] lemma list.dprod_index_nil (fι : α → ι) : ([] : list α).dprod_index fι = 0 := rfl
@[simp] lemma list.dprod_index_cons (a : α) (l : list α) (fι : α → ι) :
  (a :: l).dprod_index fι = fι a + l.dprod_index fι := rfl

lemma list.dprod_index_eq_map_sum (l : list α) (fι : α → ι) :
  l.dprod_index fι = (l.map fι).sum :=
begin
  dunfold list.dprod_index,
  induction l,
  { simp, },
  { simp [l_ih], },
end

/-- A dependent product for graded monoids represented by the indexed family of types `A i`.
This is a dependent version of `(l.map fA).prod`.

For a list `l : list α`, this computes the product of `fA a` over `a`, where each `fA` is of type
`A (fι a)`. -/
def list.dprod (l : list α) (fι : α → ι) (fA : Π a, A (fι a)) :
  A (l.dprod_index fι) :=
l.foldr_rec_on _ _ graded_monoid.ghas_one.one (λ i x a ha, graded_monoid.ghas_mul.mul (fA a) x)

@[simp] lemma list.dprod_nil (fι : α → ι) (fA : Π a, A (fι a)) :
  (list.nil : list α).dprod fι fA = graded_monoid.ghas_one.one := rfl

-- the `( : _)` in this lemma statement results in the type on the RHS not being unfolded, which
-- is nicer in the goal view.
@[simp] lemma list.dprod_cons (fι : α → ι) (fA : Π a, A (fι a)) (a : α) (l : list α) :
  (a :: l).dprod fι fA = (graded_monoid.ghas_mul.mul (fA a) (l.dprod fι fA) : _) := rfl

lemma graded_monoid.mk_list_dprod (l : list α) (fι : α → ι) (fA : Π a, A (fι a)) :
  graded_monoid.mk _ (l.dprod fι fA) = (l.map (λ a, graded_monoid.mk (fι a) (fA a))).prod :=
begin
  induction l,
  { simp, refl  },
  { simp [←l_ih, graded_monoid.mk_mul_mk, list.prod_cons],
    refl, },
end

/-- A variant of `graded_monoid.mk_list_dprod` for rewriting in the other direction. -/
lemma graded_monoid.list_prod_map_eq_dprod (l : list α) (f : α → graded_monoid A) :
  (l.map f).prod = graded_monoid.mk _ (l.dprod (λ i, (f i).1) (λ i, (f i).2)) :=
begin
  rw [graded_monoid.mk_list_dprod, graded_monoid.mk],
  simp_rw sigma.eta,
end

lemma graded_monoid.list_prod_of_fn_eq_dprod {n : ℕ} (f : fin n → graded_monoid A) :
  (list.of_fn f).prod =
    graded_monoid.mk _ ((list.fin_range n).dprod (λ i, (f i).1) (λ i, (f i).2)) :=
by rw [list.of_fn_eq_map, graded_monoid.list_prod_map_eq_dprod]

end dprod

/-! ### Concrete instances -/
section

variables (ι) {R : Type*}

@[simps one]
instance has_one.ghas_one [has_zero ι] [has_one R] : graded_monoid.ghas_one (λ i : ι, R) :=
{ one := 1 }

@[simps mul]
instance has_mul.ghas_mul [has_add ι] [has_mul R] : graded_monoid.ghas_mul (λ i : ι, R) :=
{ mul := λ i j, (*) }

/-- If all grades are the same type and themselves form a monoid, then there is a trivial grading
structure. -/
@[simps gnpow]
instance monoid.gmonoid [add_monoid ι] [monoid R] : graded_monoid.gmonoid (λ i : ι, R) :=
{ one_mul := λ a, sigma.ext (zero_add _) (heq_of_eq (one_mul _)),
  mul_one := λ a, sigma.ext (add_zero _) (heq_of_eq (mul_one _)),
  mul_assoc := λ a b c, sigma.ext (add_assoc _ _ _) (heq_of_eq (mul_assoc _ _ _)),
  gnpow := λ n i a, a ^ n,
  gnpow_zero' := λ a, sigma.ext (zero_nsmul _) (heq_of_eq (monoid.npow_zero' _)),
  gnpow_succ' := λ n ⟨i, a⟩, sigma.ext (succ_nsmul _ _) (heq_of_eq (monoid.npow_succ' _ _)),
  ..has_one.ghas_one ι,
  ..has_mul.ghas_mul ι }

/-- If all grades are the same type and themselves form a commutative monoid, then there is a
trivial grading structure. -/
instance comm_monoid.gcomm_monoid [add_comm_monoid ι] [comm_monoid R] :
  graded_monoid.gcomm_monoid (λ i : ι, R) :=
{ mul_comm := λ a b, sigma.ext (add_comm _ _) (heq_of_eq (mul_comm _ _)),
  ..monoid.gmonoid ι }

/-- When all the indexed types are the same, the dependent product is just the regular product. -/
@[simp] lemma list.dprod_monoid {α} [add_monoid ι] [monoid R] (l : list α) (fι : α → ι)
  (fA : α → R) :
  (l.dprod fι fA : (λ i : ι, R) _) = ((l.map fA).prod : _) :=
begin
  induction l,
  { rw [list.dprod_nil, list.map_nil, list.prod_nil], refl },
  { rw [list.dprod_cons, list.map_cons, list.prod_cons, l_ih], refl },
end

end

/-! ### Shorthands for creating instance of the above typeclasses for collections of subobjects -/

section subobjects

variables {R : Type*}

/-- A version of `graded_monoid.ghas_one` for internally graded objects. -/
class set_like.has_graded_one {S : Type*} [set_like S R] [has_one R] [has_zero ι]
  (A : ι → S) : Prop :=
(one_mem : (1 : R) ∈ A 0)

lemma set_like.one_mem_graded {S : Type*} [set_like S R] [has_one R] [has_zero ι] (A : ι → S)
  [set_like.has_graded_one A] : (1 : R) ∈ A 0 := set_like.has_graded_one.one_mem

instance set_like.ghas_one {S : Type*} [set_like S R] [has_one R] [has_zero ι] (A : ι → S)
  [set_like.has_graded_one A] : graded_monoid.ghas_one (λ i, A i) :=
{ one := ⟨1, set_like.one_mem_graded _⟩ }

@[simp] lemma set_like.coe_ghas_one {S : Type*} [set_like S R] [has_one R] [has_zero ι] (A : ι → S)
  [set_like.has_graded_one A] : ↑(@graded_monoid.ghas_one.one _ (λ i, A i) _ _) = (1 : R) := rfl

/-- A version of `graded_monoid.ghas_one` for internally graded objects. -/
class set_like.has_graded_mul {S : Type*} [set_like S R] [has_mul R] [has_add ι]
  (A : ι → S) : Prop :=
(mul_mem : ∀ ⦃i j⦄ {gi gj}, gi ∈ A i → gj ∈ A j → gi * gj ∈ A (i + j))

lemma set_like.mul_mem_graded {S : Type*} [set_like S R] [has_mul R] [has_add ι] {A : ι → S}
  [set_like.has_graded_mul A] ⦃i j⦄ {gi gj} (hi : gi ∈ A i) (hj : gj ∈ A j) :
  gi * gj ∈ A (i + j) :=
set_like.has_graded_mul.mul_mem hi hj

instance set_like.ghas_mul {S : Type*} [set_like S R] [has_mul R] [has_add ι] (A : ι → S)
  [set_like.has_graded_mul A] :
  graded_monoid.ghas_mul (λ i, A i) :=
{ mul := λ i j a b, ⟨(a * b : R), set_like.mul_mem_graded a.prop b.prop⟩ }

@[simp] lemma set_like.coe_ghas_mul {S : Type*} [set_like S R] [has_mul R] [has_add ι] (A : ι → S)
  [set_like.has_graded_mul A] {i j : ι} (x : A i) (y : A j) :
    ↑(@graded_monoid.ghas_mul.mul _ (λ i, A i) _ _ _ _ x y) = (x * y : R) := rfl

/-- A version of `graded_monoid.gmonoid` for internally graded objects. -/
class set_like.graded_monoid {S : Type*} [set_like S R] [monoid R] [add_monoid ι]
  (A : ι → S) extends set_like.has_graded_one A, set_like.has_graded_mul A : Prop

namespace set_like
variables {S : Type*} [set_like S R] [monoid R] [add_monoid ι]
variables {A : ι → S} [set_like.graded_monoid A]

lemma pow_mem_graded (n : ℕ) {r : R} {i : ι} (h : r ∈ A i) : r ^ n ∈ A (n • i) :=
begin
  induction n,
  { rw [pow_zero, zero_nsmul], exact one_mem_graded _ },
  { rw [pow_succ', succ_nsmul'], exact mul_mem_graded n_ih h },
end

lemma list_prod_map_mem_graded {ι'} (l : list ι') (i : ι' → ι) (r : ι' → R)
  (h : ∀ j ∈ l, r j ∈ A (i j)) :
  (l.map r).prod ∈ A (l.map i).sum :=
begin
  induction l,
  { rw [list.map_nil, list.map_nil, list.prod_nil, list.sum_nil],
    exact one_mem_graded _ },
  { rw [list.map_cons, list.map_cons, list.prod_cons, list.sum_cons],
    exact mul_mem_graded
      (h _ $ list.mem_cons_self _ _) (l_ih $ λ j hj, h _ $ list.mem_cons_of_mem _ hj) },
end

lemma list_prod_of_fn_mem_graded {n} (i : fin n → ι) (r : fin n → R) (h : ∀ j, r j ∈ A (i j)) :
  (list.of_fn r).prod ∈ A (list.of_fn i).sum :=
begin
  rw [list.of_fn_eq_map, list.of_fn_eq_map],
  exact list_prod_map_mem_graded _ _ _ (λ _ _, h _),
end

end set_like

/-- Build a `gmonoid` instance for a collection of subobjects. -/
instance set_like.gmonoid {S : Type*} [set_like S R] [monoid R] [add_monoid ι] (A : ι → S)
  [set_like.graded_monoid A] :
  graded_monoid.gmonoid (λ i, A i) :=
{ one_mul := λ ⟨i, a, h⟩, sigma.subtype_ext (zero_add _) (one_mul _),
  mul_one := λ ⟨i, a, h⟩, sigma.subtype_ext (add_zero _) (mul_one _),
  mul_assoc := λ ⟨i, a, ha⟩ ⟨j, b, hb⟩ ⟨k, c, hc⟩,
    sigma.subtype_ext (add_assoc _ _ _) (mul_assoc _ _ _),
  gnpow := λ n i a, ⟨a ^ n, set_like.pow_mem_graded n a.prop⟩,
  gnpow_zero' := λ n, sigma.subtype_ext (zero_nsmul _) (pow_zero _),
  gnpow_succ' := λ n a, sigma.subtype_ext (succ_nsmul _ _) (pow_succ _ _),
  ..set_like.ghas_one A,
  ..set_like.ghas_mul A }

@[simp] lemma set_like.coe_gnpow {S : Type*} [set_like S R] [monoid R] [add_monoid ι] (A : ι → S)
  [set_like.graded_monoid A] {i : ι} (x : A i) (n : ℕ) :
    ↑(@graded_monoid.gmonoid.gnpow _ (λ i, A i) _ _ n _ x) = (x ^ n : R) := rfl

/-- Build a `gcomm_monoid` instance for a collection of subobjects. -/
instance set_like.gcomm_monoid {S : Type*} [set_like S R] [comm_monoid R] [add_comm_monoid ι]
  (A : ι → S) [set_like.graded_monoid A] :
  graded_monoid.gcomm_monoid (λ i, A i) :=
{ mul_comm := λ ⟨i, a, ha⟩ ⟨j, b, hb⟩, sigma.subtype_ext (add_comm _ _) (mul_comm _ _),
  ..set_like.gmonoid A}

section dprod
open set_like set_like.graded_monoid
variables {α S : Type*} [set_like S R] [monoid R] [add_monoid ι]

/-- Coercing a dependent product of subtypes is the same as taking the regular product of the
coercions. -/
@[simp] lemma set_like.coe_list_dprod (A : ι → S) [set_like.graded_monoid A]
  (fι : α → ι) (fA : Π a, A (fι a)) (l : list α) :
  ↑(l.dprod fι fA : (λ i, ↥(A i)) _) = (list.prod (l.map (λ a, fA a)) : R) :=
begin
  induction l,
  { rw [list.dprod_nil, coe_ghas_one, list.map_nil, list.prod_nil] },
  { rw [list.dprod_cons, coe_ghas_mul, list.map_cons, list.prod_cons, l_ih], },
end

include R

/-- A version of `list.coe_dprod_set_like` with `subtype.mk`. -/
lemma set_like.list_dprod_eq (A : ι → S) [set_like.graded_monoid A]
  (fι : α → ι) (fA : Π a, A (fι a)) (l : list α) :
  (l.dprod fι fA : (λ i, ↥(A i)) _) =
    ⟨list.prod (l.map (λ a, fA a)), (l.dprod_index_eq_map_sum fι).symm ▸
      list_prod_map_mem_graded l _ _ (λ i hi, (fA i).prop)⟩ :=
subtype.ext $ set_like.coe_list_dprod _ _ _ _

end dprod

end subobjects

section homogeneous_elements

variables {R S : Type*} [set_like S R]

/-- An element `a : R` is said to be homogeneous if there is some `i : ι` such that `a ∈ A i`. -/
def set_like.is_homogeneous (A : ι → S) (a : R) : Prop := ∃ i, a ∈ A i

@[simp] lemma set_like.is_homogeneous_coe {A : ι → S} {i} (x : A i) :
  set_like.is_homogeneous A (x : R) :=
⟨i, x.prop⟩

lemma set_like.is_homogeneous_one [has_zero ι] [has_one R]
  (A : ι → S) [set_like.has_graded_one A] : set_like.is_homogeneous A (1 : R) :=
⟨0, set_like.one_mem_graded _⟩

lemma set_like.is_homogeneous.mul [has_add ι] [has_mul R] {A : ι → S}
  [set_like.has_graded_mul A] {a b : R} :
  set_like.is_homogeneous A a → set_like.is_homogeneous A b → set_like.is_homogeneous A (a * b)
| ⟨i, hi⟩ ⟨j, hj⟩ := ⟨i + j, set_like.mul_mem_graded hi hj⟩

/-- When `A` is a `set_like.graded_monoid A`, then the homogeneous elements forms a submonoid. -/
def set_like.homogeneous_submonoid [add_monoid ι] [monoid R]
  (A : ι → S) [set_like.graded_monoid A] : submonoid R :=
{ carrier := { a | set_like.is_homogeneous A a },
  one_mem' := set_like.is_homogeneous_one A,
  mul_mem' := λ a b, set_like.is_homogeneous.mul }

end homogeneous_elements


/-
Copyright (c) 2015 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro
-/
import data.multiset.powerset
import data.multiset.range

/-!
# The `nodup` predicate for multisets without duplicate elements.
-/

namespace multiset
open list

variables {α β γ : Type*}

/- nodup -/

/-- `nodup s` means that `s` has no duplicates, i.e. the multiplicity of
  any element is at most 1. -/
def nodup (s : multiset α) : Prop :=
quot.lift_on s nodup (λ s t p, propext p.nodup_iff)

@[simp] theorem coe_nodup {l : list α} : @nodup α l ↔ l.nodup := iff.rfl

@[simp] theorem nodup_zero : @nodup α 0 := pairwise.nil

@[simp] theorem nodup_cons {a : α} {s : multiset α} : nodup (a ::ₘ s) ↔ a ∉ s ∧ nodup s :=
quot.induction_on s $ λ l, nodup_cons

theorem nodup_cons_of_nodup {a : α} {s : multiset α} (m : a ∉ s) (n : nodup s) : nodup (a ::ₘ s) :=
nodup_cons.2 ⟨m, n⟩

theorem nodup_singleton : ∀ a : α, nodup ({a} : multiset α) := nodup_singleton

theorem nodup_of_nodup_cons {a : α} {s : multiset α} (h : nodup (a ::ₘ s)) : nodup s :=
(nodup_cons.1 h).2

theorem not_mem_of_nodup_cons {a : α} {s : multiset α} (h : nodup (a ::ₘ s)) : a ∉ s :=
(nodup_cons.1 h).1

theorem nodup_of_le {s t : multiset α} (h : s ≤ t) : nodup t → nodup s :=
le_induction_on h $ λ l₁ l₂, nodup_of_sublist

theorem not_nodup_pair : ∀ a : α, ¬ nodup (a ::ₘ a ::ₘ 0) := not_nodup_pair

theorem nodup_iff_le {s : multiset α} : nodup s ↔ ∀ a : α, ¬ a ::ₘ a ::ₘ 0 ≤ s :=
quot.induction_on s $ λ l, nodup_iff_sublist.trans $ forall_congr $ λ a,
not_congr (@repeat_le_coe _ a 2 _).symm

lemma nodup_iff_ne_cons_cons {s : multiset α} : s.nodup ↔ ∀ a t, s ≠ a ::ₘ a ::ₘ t :=
nodup_iff_le.trans
  ⟨λ h a t s_eq, h a (s_eq.symm ▸ cons_le_cons a (cons_le_cons a (zero_le _))),
   λ h a le, let ⟨t, s_eq⟩ := le_iff_exists_add.mp le in
     h a t (by rwa [cons_add, cons_add, zero_add] at s_eq )⟩

theorem nodup_iff_count_le_one [decidable_eq α] {s : multiset α} : nodup s ↔ ∀ a, count a s ≤ 1 :=
quot.induction_on s $ λ l, nodup_iff_count_le_one

@[simp] theorem count_eq_one_of_mem [decidable_eq α] {a : α} {s : multiset α}
  (d : nodup s) (h : a ∈ s) : count a s = 1 :=
le_antisymm (nodup_iff_count_le_one.1 d a) (count_pos.2 h)

lemma nodup_iff_pairwise {α} {s : multiset α} : nodup s ↔ pairwise (≠) s :=
quotient.induction_on s $ λ l, (pairwise_coe_iff_pairwise (by exact λ a b, ne.symm)).symm

lemma pairwise_of_nodup {r : α → α → Prop} {s : multiset α} :
  (∀a∈s, ∀b∈s, a ≠ b → r a b) → nodup s → pairwise r s :=
quotient.induction_on s $ assume l h hl, ⟨l, rfl, hl.imp_of_mem $ assume a b ha hb, h a ha b hb⟩

lemma forall_of_pairwise {r : α → α → Prop} (H : symmetric r) {s : multiset α}
   (hs : pairwise r s) : (∀a∈s, ∀b∈s, a ≠ b → r a b) :=
let ⟨l, hl₁, hl₂⟩ := hs in hl₁.symm ▸ list.forall_of_pairwise H hl₂

theorem nodup_add {s t : multiset α} : nodup (s + t) ↔ nodup s ∧ nodup t ∧ disjoint s t :=
quotient.induction_on₂ s t $ λ l₁ l₂, nodup_append

theorem disjoint_of_nodup_add {s t : multiset α} (d : nodup (s + t)) : disjoint s t :=
(nodup_add.1 d).2.2

theorem nodup_add_of_nodup {s t : multiset α} (d₁ : nodup s) (d₂ : nodup t) :
  nodup (s + t) ↔ disjoint s t :=
by simp [nodup_add, d₁, d₂]

theorem nodup_of_nodup_map (f : α → β) {s : multiset α} : nodup (map f s) → nodup s :=
quot.induction_on s $ λ l, nodup_of_nodup_map f

theorem nodup_map_on {f : α → β} {s : multiset α} : (∀x∈s, ∀y∈s, f x = f y → x = y) →
  nodup s → nodup (map f s) :=
quot.induction_on s $ λ l, nodup_map_on

theorem nodup_map {f : α → β} {s : multiset α} (hf : function.injective f) :
  nodup s → nodup (map f s) :=
nodup_map_on (λ x _ y _ h, hf h)

theorem inj_on_of_nodup_map {f : α → β} {s : multiset α} :
  nodup (map f s) → ∀ (x ∈ s) (y ∈ s), f x = f y → x = y :=
quot.induction_on s $ λ l, inj_on_of_nodup_map

theorem nodup_map_iff_inj_on {f : α → β} {s : multiset α} (d : nodup s) :
  nodup (map f s) ↔ (∀ (x ∈ s) (y ∈ s), f x = f y → x = y) :=
⟨inj_on_of_nodup_map, λ h, nodup_map_on h d⟩

theorem nodup_filter (p : α → Prop) [decidable_pred p] {s} : nodup s → nodup (filter p s) :=
quot.induction_on s $ λ l, nodup_filter p

@[simp] theorem nodup_attach {s : multiset α} : nodup (attach s) ↔ nodup s :=
quot.induction_on s $ λ l, nodup_attach

theorem nodup_pmap {p : α → Prop} {f : Π a, p a → β} {s : multiset α} {H}
  (hf : ∀ a ha b hb, f a ha = f b hb → a = b) : nodup s → nodup (pmap f s H) :=
quot.induction_on s (λ l H, nodup_pmap hf) H

instance nodup_decidable [decidable_eq α] (s : multiset α) : decidable (nodup s) :=
quotient.rec_on_subsingleton s $ λ l, l.nodup_decidable

theorem nodup_erase_eq_filter [decidable_eq α] (a : α) {s} : nodup s → s.erase a = filter (≠ a) s :=
quot.induction_on s $ λ l d, congr_arg coe $ nodup_erase_eq_filter a d

theorem nodup_erase_of_nodup [decidable_eq α] (a : α) {l} : nodup l → nodup (l.erase a) :=
nodup_of_le (erase_le _ _)

theorem mem_erase_iff_of_nodup [decidable_eq α] {a b : α} {l} (d : nodup l) :
  a ∈ l.erase b ↔ a ≠ b ∧ a ∈ l :=
by rw nodup_erase_eq_filter b d; simp [and_comm]

theorem mem_erase_of_nodup [decidable_eq α] {a : α} {l} (h : nodup l) : a ∉ l.erase a :=
by rw mem_erase_iff_of_nodup h; simp

theorem nodup_product {s : multiset α} {t : multiset β} : nodup s → nodup t → nodup (product s t) :=
quotient.induction_on₂ s t $ λ l₁ l₂ d₁ d₂, by simp [nodup_product d₁ d₂]

theorem nodup_sigma {σ : α → Type*} {s : multiset α} {t : Π a, multiset (σ a)} :
  nodup s → (∀ a, nodup (t a)) → nodup (s.sigma t) :=
quot.induction_on s $ assume l₁,
begin
  choose f hf using assume a, quotient.exists_rep (t a),
  rw show t = λ a, f a, from (eq.symm $ funext $ λ a, hf a),
  simpa using nodup_sigma
end

theorem nodup_filter_map (f : α → option β) {s : multiset α}
  (H : ∀ (a a' : α) (b : β), b ∈ f a → b ∈ f a' → a = a') :
  nodup s → nodup (filter_map f s) :=
quot.induction_on s $ λ l, nodup_filter_map H

theorem nodup_inter_left [decidable_eq α] {s : multiset α} (t) : nodup s → nodup (s ∩ t) :=
nodup_of_le $ inter_le_left _ _

theorem nodup_inter_right [decidable_eq α] (s) {t : multiset α} : nodup t → nodup (s ∩ t) :=
nodup_of_le $ inter_le_right _ _

@[simp] theorem nodup_union [decidable_eq α] {s t : multiset α} :
  nodup (s ∪ t) ↔ nodup s ∧ nodup t :=
⟨λ h, ⟨nodup_of_le (le_union_left _ _) h, nodup_of_le (le_union_right _ _) h⟩,
 λ ⟨h₁, h₂⟩, nodup_iff_count_le_one.2 $ λ a, by rw [count_union]; exact
   max_le (nodup_iff_count_le_one.1 h₁ a) (nodup_iff_count_le_one.1 h₂ a)⟩

@[simp] theorem nodup_powerset {s : multiset α} : nodup (powerset s) ↔ nodup s :=
⟨λ h, nodup_of_nodup_map _ (nodup_of_le (map_single_le_powerset _) h),
  quotient.induction_on s $ λ l h,
  by simp; refine list.nodup_map_on _ (nodup_sublists'.2 h); exact
  λ x sx y sy e,
    (h.sublist_ext (mem_sublists'.1 sx) (mem_sublists'.1 sy)).1
      (quotient.exact e)⟩

theorem nodup_powerset_len {n : ℕ} {s : multiset α}
  (h : nodup s) : nodup (powerset_len n s) :=
nodup_of_le (powerset_len_le_powerset _ _) (nodup_powerset.2 h)

@[simp] lemma nodup_bind {s : multiset α} {t : α → multiset β} :
  nodup (bind s t) ↔ ((∀a∈s, nodup (t a)) ∧ (s.pairwise (λa b, disjoint (t a) (t b)))) :=
have h₁ : ∀a, ∃l:list β, t a = l, from
  assume a, quot.induction_on (t a) $ assume l, ⟨l, rfl⟩,
let ⟨t', h'⟩ := classical.axiom_of_choice h₁ in
have t = λa, t' a, from funext h',
have hd : symmetric (λa b, list.disjoint (t' a) (t' b)), from assume a b h, h.symm,
quot.induction_on s $ by simp [this, list.nodup_bind, pairwise_coe_iff_pairwise hd]

theorem nodup_ext {s t : multiset α} : nodup s → nodup t → (s = t ↔ ∀ a, a ∈ s ↔ a ∈ t) :=
quotient.induction_on₂ s t $ λ l₁ l₂ d₁ d₂, quotient.eq.trans $ perm_ext d₁ d₂

theorem le_iff_subset {s t : multiset α} : nodup s → (s ≤ t ↔ s ⊆ t) :=
quotient.induction_on₂ s t $ λ l₁ l₂ d, ⟨subset_of_le, subperm_of_subset_nodup d⟩

theorem range_le {m n : ℕ} : range m ≤ range n ↔ m ≤ n :=
(le_iff_subset (nodup_range _)).trans range_subset

theorem mem_sub_of_nodup [decidable_eq α] {a : α} {s t : multiset α} (d : nodup s) :
  a ∈ s - t ↔ a ∈ s ∧ a ∉ t :=
⟨λ h, ⟨mem_of_le tsub_le_self h, λ h',
  by refine count_eq_zero.1 _ h; rw [count_sub a s t, tsub_eq_zero_iff_le];
     exact le_trans (nodup_iff_count_le_one.1 d _) (count_pos.2 h')⟩,
 λ ⟨h₁, h₂⟩, or.resolve_right (mem_add.1 $ mem_of_le le_tsub_add h₁) h₂⟩

lemma map_eq_map_of_bij_of_nodup (f : α → γ) (g : β → γ) {s : multiset α} {t : multiset β}
  (hs : s.nodup) (ht : t.nodup) (i : Πa∈s, β)
  (hi : ∀a ha, i a ha ∈ t) (h : ∀a ha, f a = g (i a ha))
  (i_inj : ∀a₁ a₂ ha₁ ha₂, i a₁ ha₁ = i a₂ ha₂ → a₁ = a₂)
  (i_surj : ∀b∈t, ∃a ha, b = i a ha) :
  s.map f = t.map g :=
have t = s.attach.map (λ x, i x.1 x.2),
  from (nodup_ext ht (nodup_map
      (show function.injective (λ x : {x // x ∈ s}, i x.1 x.2), from λ x y hxy,
        subtype.eq (i_inj x.1 y.1 x.2 y.2 hxy))
      (nodup_attach.2 hs))).2
    (λ x, by simp only [mem_map, true_and, subtype.exists, eq_comm, mem_attach];
      exact ⟨i_surj _, λ ⟨y, hy⟩, hy.snd.symm ▸ hi _ _⟩),
calc s.map f = s.pmap  (λ x _, f x) (λ _, id) : by rw [pmap_eq_map]
... = s.attach.map (λ x, f x.1) : by rw [pmap_eq_map_attach]
... = t.map g : by rw [this, multiset.map_map]; exact map_congr (λ x _, h _ _)

end multiset


/-
Copyright (c) 2022 Jun Yoshida. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
-/
import Std.Logic

import Algdata.Init.UInt

set_option autoImplicit false

/-!
# Quotient ring ℤ/256ℤ
-/

/--
The 8-bit unsigned integer type as the quotient ring ℤ/256ℤ.
-/
@[unbox]
structure ℤ256 : Type :=
  mk :: (val : UInt8)
  deriving DecidableEq, Repr

namespace ℤ256

/-!
## Basic instances

@remark Since we will give `div`/`mod` later in terms of the 2-adic pseudo-norm, they do not appear in this section.
-/

@[inline]
protected
def toNat (x : ℤ256) : Nat := x.val.val.val

@[inline]
protected
def ofNat (n : @& Nat) : ℤ256 := mk (OfNat.ofNat n)

@[inline, match_pattern]
protected
abbrev zero : ℤ256 := mk UInt8.zero

@[inline, match_pattern]
protected
abbrev one : ℤ256 := mk UInt8.one

@[inline]
protected
def add (x y : ℤ256) : ℤ256 := mk (x.val + y.val)

@[inline]
protected
def sub (x y : ℤ256) : ℤ256 := mk (x.val - y.val)

@[inline]
protected
def neg (x : ℤ256) : ℤ256 := mk (0 - x.val)

@[inline]
protected
def mul (x y : ℤ256) : ℤ256 := mk (x.val * y.val)

/-- @remark `instOfNatNat` has priority 100. -/
instance (priority := 50) (n : Nat) : OfNat ℤ256 n := ⟨ℤ256.ofNat n⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 0) := ⟨ℤ256.zero⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 1) := ⟨ℤ256.one⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 2) := ⟨mk (.raw 2 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 3) := ⟨mk (.raw 3 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 4) := ⟨mk (.raw 4 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 5) := ⟨mk (.raw 5 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 6) := ⟨mk (.raw 6 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 7) := ⟨mk (.raw 7 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 8) := ⟨mk (.raw 8 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 0x10) := ⟨mk (.raw 0x10 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 0x20) := ⟨mk (.raw 0x20 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 0x40) := ⟨mk (.raw 0x40 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 0x80) := ⟨mk (.raw 0x80 (by decide))⟩
instance (priority := 80) : OfNat ℤ256 (nat_lit 0xFF) := ⟨mk (.raw 0xFF 0xFF.lt_succ_self)⟩

instance : Add ℤ256 := ⟨ℤ256.add⟩
instance : Sub ℤ256 := ⟨ℤ256.sub⟩
instance : Neg ℤ256 := ⟨ℤ256.neg⟩
instance : Mul ℤ256 := ⟨ℤ256.mul⟩

@[inline]
protected
def toString (x : ℤ256) := x.val.binaryDigits

instance : ToString ℤ256 := ⟨ℤ256.toString⟩

/-!
## Bit operations

@remark For some of the operations, we provide two versions of implementations; one uses the corresponding operations on `UInt8` while the other avoid well-founded recursions. The equivalences of the two versions are proved in [Lemmas.lean](Algdata/Data/Z128/Lemmas.lean).
-/

@[inline]
protected
def complement (x : ℤ256) : ℤ256 := mk (~~~x.val)

/-- The (homogeneous) left shift operator. -/
@[inline]
protected
def shiftLeft (x y : ℤ256) : ℤ256 := mk (x.val <<< y.val)

/-- The (homogeneous) right shift operator. -/
@[inline]
protected
def shiftRight (x y : ℤ256) : ℤ256 := mk (x.val >>> y.val)

instance : Complement ℤ256 := ⟨ℤ256.complement⟩
instance : ShiftLeft ℤ256 := ⟨ℤ256.shiftLeft⟩
instance : ShiftRight ℤ256 := ⟨ℤ256.shiftRight⟩

/--
OR-operation: defined by `UInt8.lor`.
@warning As `UInt8.lor` uses well-founded recursion, do not try to reduce the function directly.
 -/
@[inline]
protected
def lor_wf (x y : ℤ256) : ℤ256 := mk (x.val ||| y.val)

/-- OR-operation: implemented without well-founded recursion. -/
@[implemented_by ℤ256.lor_wf]
protected
def lor (x y : ℤ256) : ℤ256 := aux (nat_lit 8) x y
  where
    aux : Nat → ℤ256 → ℤ256 → ℤ256
    | 0, _, _ => 0
    | n+1, x, y =>
      let z := aux n (x >>> 1) (y >>> 1)
      (z <<< 1) + mk (if x.val % 2 = 0 then y.val % 2 else 1)

/--
AND-operation: defined by `UInt8.land`.
@warning As `Uint8.land` uses well-founded recursion, do not try to reduce the function directly.
-/
@[inline]
protected
def land_wf (x y : ℤ256) : ℤ256 := mk (x.val &&& y.val)

/-- AND-operation: implemented without well-founded recursion. -/
@[implemented_by ℤ256.land_wf]
protected
def land (x y : ℤ256) : ℤ256 := aux (nat_lit 8) x y
  where
    aux : Nat → ℤ256 → ℤ256 → ℤ256
    | 0, _, _ => 0
    | n+1, x, y =>
      let z := aux n (x >>> 1) (y >>> 1)
      (z <<< 1) + mk (if x.val % 2 = 0 then 0 else y.val % 2)

/--
XOR-operation: defined by `UInt8.xor`.
@warning As `UInt8.xor` uses well-founded recursion, do not try to reduce the function directly.
-/
@[inline]
protected
def xor_wf (x y : ℤ256) : ℤ256 := mk (x.val ^^^ y.val)

/-- XOR-operation: implemented without well-founded recursion. -/
@[implemented_by ℤ256.xor_wf]
protected
def xor (x y : ℤ256) : ℤ256 := aux (nat_lit 8) x y
  where
    aux : Nat → ℤ256 → ℤ256 → ℤ256
    | 0, _, _ => 0
    | n+1, x, y =>
      let z := aux n (x >>> 1) (y >>> 1)
      (z <<< 1) + mk (if x.val % 2 = y.val % 2 then 0 else 1)

instance : AndOp ℤ256 := ⟨ℤ256.land⟩
instance : OrOp ℤ256 := ⟨ℤ256.lor⟩
instance : Xor ℤ256 := ⟨ℤ256.xor⟩

/--
Counting non-zero bits.
-/
@[inline]
def bitCount (x : ℤ256) : ℤ256 :=
  Id.run <| do
    let mut a := x
    a := (a &&& 0b01010101 : ℤ256) + (a &&& 0b10101010) >>> 1
    a := (a &&& 0b00110011 : ℤ256) + (a &&& 0b11001100) >>> 2
    a := (a &&& 0b00001111 : ℤ256) + (a &&& 0b11110000) >>> 4
    return a

/--
Given `x : ℤ256`, `x.tzmask` is the bit mask for trailing zeros; e.g.
```lean
tzmask 0b10110100 = 0b00000011
```
-/
@[inline]
def tzmask (x : ℤ256) : ℤ256 := ~~~x &&& (x-1)

/--
Count trailing zeros; e.g. `ℤ256.ctz 0b10110100 = 2`.
This function is exactly the *2-adic pseudo-norm*; for example, `x : ℤ256` is invertible if and only if `x.ctz = 0`.
-/
@[inline]
def ctz (x : ℤ256) : ℤ256 := x.tzmask.bitCount

@[inline]
def norm₂ (x : ℤ256) : Nat := x.ctz.toNat

/--
The position of the most significant bit; note that `n.toNat.log2 = n.msb.pred`.
```lean
-- []
#eval Nat.fold (λ n s => if (ℤ256.ofNat n).msb.pred = n.log2 then s else (n::s)) 256 []
```
-/
@[inline]
def msb (x : ℤ256) : ℤ256 :=
  Id.run <| do
    let mut a := x
    a := a ||| (a >>> 1)
    a := a ||| (a >>> 2)
    a := a ||| (a >>> 4)
    return a.bitCount

/--
Think of a term `a : ℤ256` as a bit-array of length 8.
-/
instance : GetElem ℤ256 Nat ℤ256 (λ _ i => i < 8) where
  getElem x i h := (x >>> mk (.raw i (trans h (by decide : 8 < UInt8.size)))) &&& 1

/-!
## Div/Mod

The quotient ring ℤ/256ℤ is an Euclidean ring in the sense that, given pair `x y : ℤ256`, there are unique elements `q r : ℤ256` such that

  - `x = q * y + r`
  - `r = 0 ∨ r.ctz < y.ctz`

Indeed, every element `x : ℤ256` is written in the form `x = u * 2^n` for `u : ℤ256` and `n : Nat` such that `u` is invertible and `n = x.norm₂`.
Moreover, such `r` is unique up to unit multiples of `y`.

Test code:
```lean
-- true
#eval let a := Array.ofFn (n:=UInt8.size) (λ n => mk ⟨n⟩); a.foldl (λ b x => b && a.foldl (λ b y => b && (x = (x / y) * y + x % y)) True) True
```
-/

/--
2-adic mod: `x % y` is the lower bits of `x` up to the `y.norm₂`-th bit.
Specifically, we have `x % 0 = x`.
-/
@[inline]
protected
def mod (x y : ℤ256) : ℤ256 := x &&& y.tzmask

/--
The 2-adic division (with 2-adic remainders).
@remark `x / 0 = 0`.
-/
@[inline]
protected
def div (x y : ℤ256) : ℤ256 :=
  if y = 0
  then 0
  else Id.run <| do
    let mut x := x >>> y.ctz
    let y := y >>> y.ctz
    let mut q := 0
    q := q ||| (x &&& 0x01); x := (x - ((x &&& 0x01) * y)) >>> 1
    q := q ||| (x &&& 0x01) <<< 1; x := (x - ((x &&& 0x01) * y)) >>> 1
    q := q ||| ((x &&& 0x01) <<< 2); x := (x - ((x &&& 0x01) * y)) >>> 1
    q := q ||| ((x &&& 0x01) <<< 3); x := (x - ((x &&& 0x01) * y)) >>> 1
    q := q ||| ((x &&& 0x01) <<< 4); x := (x - ((x &&& 0x01) * y)) >>> 1
    q := q ||| ((x &&& 0x01) <<< 5); x := (x - ((x &&& 0x01) * y)) >>> 1
    q := q ||| ((x &&& 0x01) <<< 6); x := (x - ((x &&& 0x01) * y)) >>> 1
    q := q ||| ((x &&& 0x01) <<< 7); x := (x - ((x &&& 0x01) * y)) >>> 1
    return q

instance : Mod ℤ256 := ⟨ℤ256.mod⟩
instance : Div ℤ256 := ⟨ℤ256.div⟩

/-!
## Recursors

As a truncated 2-adic integer, `ℤ256` has a recursor based on the 2-adic pseudo-norm.
-/

section recursor

private
theorem shiftRec_terminate_aux : ∀ (n : UInt8), n = 0 ∨ ((n >>> 1) < n)
| .zero => Or.inl rfl
| .raw (n+1) h => Or.inr $ by
  conv => change (n+1) / 2 % UInt8.size < n+1
  calc
    (n+1) / 2 % UInt8.size
      = (n+1) / 2 := Nat.mod_eq_of_lt (Trans.trans (Nat.div_le_self _ 2) h)
    _ < n+1 := Nat.div_lt_self n.zero_lt_succ (Nat.lt_succ_self 1)

protected
def shiftRec.{u} {motive : ℤ256 → Sort u} (zero : motive 0) (shift : (x : ℤ256) → x ≠ 0 → motive (x >>> 1) → motive x) : (x : ℤ256) → motive x
| 0 => zero
| mk (.raw (n+1) hn) =>
  have hneq : mk (UInt8.mk (Fin.mk (n+1) hn)) ≠ 0 := by intro h; cases h
  have : UInt8.mk (Fin.mk (n+1) hn) >>> 1 < UInt8.mk (Fin.mk (n+1) hn) :=
    Or.resolve_left (shiftRec_terminate_aux _) (by intro h; cases h)
  shift _ hneq (ℤ256.shiftRec zero shift ( mk (UInt8.mk (Fin.mk (n+1) hn) >>> 1)))
termination_by _ _ _ _ x => x.val.val.val

end recursor

end ℤ256


/-
Copyright (c) 2022 Andrew Yang. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Andrew Yang
-/
import algebraic_geometry.Gamma_Spec_adjunction
import algebraic_geometry.open_immersion
import category_theory.limits.opposites

/-!
# Affine schemes

We define the category of `AffineScheme`s as the essential image of `Spec`.
We also define predicates about affine schemes and affine open sets.

## Main definitions

* `algebraic_geometry.AffineScheme`: The category of affine schemes.
* `algebraic_geometry.is_affine`: A scheme is affine if the canonical map `X ⟶ Spec Γ(X)` is an
  isomorphism.
* `algebraic_geometry.Scheme.iso_Spec`: The canonical isomorphism `X ≅ Spec Γ(X)` for an affine
  scheme.
* `algebraic_geometry.AffineScheme.equiv_CommRing`: The equivalence of categories
  `AffineScheme ≌ CommRingᵒᵖ` given by `AffineScheme.Spec : CommRingᵒᵖ ⥤ AffineScheme` and
  `AffineScheme.Γ : AffineSchemeᵒᵖ ⥤ CommRing`.
* `algebraic_geometry.is_affine_open`: An open subset of a scheme is affine if the open subscheme is
  affine.
* `algebraic_geometry.is_affine_open.from_Spec`: The immersion `Spec 𝒪ₓ(U) ⟶ X` for an affine `U`.

-/

noncomputable theory

open category_theory category_theory.limits opposite topological_space

universe u

namespace algebraic_geometry

open Spec (structure_sheaf)

/-- The category of affine schemes -/
def AffineScheme := Scheme.Spec.ess_image

/-- A Scheme is affine if the canonical map `X ⟶ Spec Γ(X)` is an isomorphism. -/
class is_affine (X : Scheme) : Prop :=
(affine : is_iso (Γ_Spec.adjunction.unit.app X))

attribute [instance] is_affine.affine

/-- The canonical isomorphism `X ≅ Spec Γ(X)` for an affine scheme. -/
def Scheme.iso_Spec (X : Scheme) [is_affine X] :
  X ≅ Scheme.Spec.obj (op $ Scheme.Γ.obj $ op X) :=
as_iso (Γ_Spec.adjunction.unit.app X)

lemma mem_AffineScheme (X : Scheme) : X ∈ AffineScheme ↔ is_affine X :=
⟨λ h, ⟨functor.ess_image.unit_is_iso h⟩, λ h, @@mem_ess_image_of_unit_is_iso _ _ _ X h.1⟩

instance is_affine_AffineScheme (X : AffineScheme.{u}) : is_affine (X : Scheme.{u}) :=
(mem_AffineScheme _).mp X.prop

instance Spec_is_affine (R : CommRingᵒᵖ) : is_affine (Scheme.Spec.obj R) :=
(mem_AffineScheme _).mp (Scheme.Spec.obj_mem_ess_image R)

lemma is_affine_of_iso {X Y : Scheme} (f : X ⟶ Y) [is_iso f] [h : is_affine Y] :
  is_affine X :=
by { rw [← mem_AffineScheme] at h ⊢, exact functor.ess_image.of_iso (as_iso f).symm h }

namespace AffineScheme

/-- The `Spec` functor into the category of affine schemes. -/
@[derive [full, faithful, ess_surj], simps]
def Spec : CommRingᵒᵖ ⥤ AffineScheme := Scheme.Spec.to_ess_image

/-- The forgetful functor `AffineScheme ⥤ Scheme`. -/
@[derive [full, faithful], simps]
def forget_to_Scheme : AffineScheme ⥤ Scheme := Scheme.Spec.ess_image_inclusion

/-- The global section functor of an affine scheme. -/
def Γ : AffineSchemeᵒᵖ ⥤ CommRing := forget_to_Scheme.op ⋙ Scheme.Γ

/-- The category of affine schemes is equivalent to the category of commutative rings. -/
def equiv_CommRing : AffineScheme ≌ CommRingᵒᵖ :=
equiv_ess_image_of_reflective.symm

instance Γ_is_equiv : is_equivalence Γ.{u} :=
begin
  haveI : is_equivalence Γ.{u}.right_op.op := is_equivalence.of_equivalence equiv_CommRing.op,
  exact (functor.is_equivalence_trans Γ.{u}.right_op.op (op_op_equivalence _).functor : _),
end

instance : has_colimits AffineScheme.{u} :=
begin
  haveI := adjunction.has_limits_of_equivalence.{u} Γ.{u},
  haveI : has_colimits AffineScheme.{u} ᵒᵖᵒᵖ := has_colimits_op_of_has_limits,
  exactI adjunction.has_colimits_of_equivalence.{u} (op_op_equivalence AffineScheme.{u}).inverse
end

instance : has_limits AffineScheme.{u} :=
begin
  haveI := adjunction.has_colimits_of_equivalence Γ.{u},
  haveI : has_limits AffineScheme.{u} ᵒᵖᵒᵖ := limits.has_limits_op_of_has_colimits,
  exactI adjunction.has_limits_of_equivalence (op_op_equivalence AffineScheme.{u}).inverse
end

end AffineScheme

/-- An open subset of a scheme is affine if the open subscheme is affine. -/
def is_affine_open {X : Scheme} (U : opens X.carrier) : Prop :=
is_affine (X.restrict U.open_embedding)

lemma range_is_affine_open_of_open_immersion {X Y : Scheme} [is_affine X] (f : X ⟶ Y)
  [H : is_open_immersion f] : is_affine_open ⟨set.range f.1.base, H.base_open.open_range⟩ :=
begin
  refine is_affine_of_iso (is_open_immersion.iso_of_range_eq f (Y.of_restrict _) _).inv,
  exact subtype.range_coe.symm,
  apply_instance
end

lemma top_is_affine_open (X : Scheme) [is_affine X] : is_affine_open (⊤ : opens X.carrier) :=
begin
  convert range_is_affine_open_of_open_immersion (𝟙 X),
  ext1,
  exact set.range_id.symm
end

instance Scheme.affine_basis_cover_is_affine (X : Scheme) (i : X.affine_basis_cover.J) :
  is_affine (X.affine_basis_cover.obj i) :=
algebraic_geometry.Spec_is_affine _

lemma is_basis_affine_open (X : Scheme) :
  opens.is_basis { U : opens X.carrier | is_affine_open U } :=
begin
  rw opens.is_basis_iff_nbhd,
  rintros U x (hU : x ∈ (U : set X.carrier)),
  obtain ⟨S, hS, hxS, hSU⟩ := X.affine_basis_cover_is_basis.exists_subset_of_mem_open hU U.prop,
  refine ⟨⟨S, X.affine_basis_cover_is_basis.is_open hS⟩, _, hxS, hSU⟩,
  rcases hS with ⟨i, rfl⟩,
  exact range_is_affine_open_of_open_immersion _,
end

/-- The open immersion `Spec 𝒪ₓ(U) ⟶ X` for an affine `U`. -/
def is_affine_open.from_Spec {X : Scheme} {U : opens X.carrier} (hU : is_affine_open U) :
  Scheme.Spec.obj (op $ X.presheaf.obj $ op U) ⟶ X :=
begin
  haveI : is_affine (X.restrict U.open_embedding) := hU,
  have : U.open_embedding.is_open_map.functor.obj ⊤ = U,
  { ext1, exact set.image_univ.trans subtype.range_coe },
  exact Scheme.Spec.map (X.presheaf.map (eq_to_hom this.symm).op).op ≫
    (X.restrict U.open_embedding).iso_Spec.inv ≫ X.of_restrict _
end

instance is_affine_open.is_open_immersion_from_Spec {X : Scheme} {U : opens X.carrier}
  (hU : is_affine_open U) :
  is_open_immersion hU.from_Spec :=
by { delta is_affine_open.from_Spec, apply_instance }

lemma is_affine_open.from_Spec_range {X : Scheme} {U : opens X.carrier} (hU : is_affine_open U) :
  set.range hU.from_Spec.1.base = (U : set X.carrier) :=
begin
  delta is_affine_open.from_Spec,
  erw [← category.assoc, Scheme.comp_val_base],
  rw [coe_comp, set.range_comp, set.range_iff_surjective.mpr, set.image_univ],
  exact subtype.range_coe,
  rw ← Top.epi_iff_surjective,
  apply_instance
end

lemma is_affine_open.from_Spec_image_top {X : Scheme} {U : opens X.carrier}
  (hU : is_affine_open U) :
  hU.is_open_immersion_from_Spec.base_open.is_open_map.functor.obj ⊤ = U :=
by { ext1, exact set.image_univ.trans hU.from_Spec_range }

lemma is_affine_open.is_compact {X : Scheme} {U : opens X.carrier} (hU : is_affine_open U) :
  is_compact (U : set X.carrier) :=
begin
  convert @is_compact.image _ _ _ _ set.univ hU.from_Spec.1.base
    prime_spectrum.compact_space.1 (by continuity),
  convert hU.from_Spec_range.symm,
  exact set.image_univ
end

instance Scheme.quasi_compact_of_affine (X : Scheme) [is_affine X] : compact_space X.carrier :=
⟨(top_is_affine_open X).is_compact⟩

lemma is_affine_open.from_Spec_base_preimage
  {X : Scheme} {U : opens X.carrier} (hU : is_affine_open U) :
    (opens.map hU.from_Spec.val.base).obj U = ⊤ :=
begin
  ext1,
  change hU.from_Spec.1.base ⁻¹' (U : set X.carrier) = set.univ,
  rw [← hU.from_Spec_range, ← set.image_univ],
  exact set.preimage_image_eq _ PresheafedSpace.is_open_immersion.base_open.inj
end

lemma Scheme.Spec_map_presheaf_map_eq_to_hom {X : Scheme} {U V : opens X.carrier} (h : U = V) (W) :
  (Scheme.Spec.map (X.presheaf.map (eq_to_hom h).op).op).val.c.app W =
    eq_to_hom (by { cases h, dsimp, induction W using opposite.rec, congr, ext1, simpa }) :=
begin
  have : Scheme.Spec.map (X.presheaf.map (𝟙 (op U))).op = 𝟙 _,
  { rw [X.presheaf.map_id, op_id, Scheme.Spec.map_id]  },
  cases h,
  refine (Scheme.congr_app this _).trans _,
  erw category.id_comp,
  simpa
end

lemma is_affine_open.Spec_Γ_identity_hom_app_from_Spec {X : Scheme} {U : opens X.carrier}
  (hU : is_affine_open U) :
  (Spec_Γ_identity.hom.app (X.presheaf.obj $ op U)) ≫ hU.from_Spec.1.c.app (op U) =
    (Scheme.Spec.obj _).presheaf.map (eq_to_hom hU.from_Spec_base_preimage).op :=
begin
  haveI : is_affine _ := hU,
  have e₁ :=
    Spec_Γ_identity.hom.naturality (X.presheaf.map (eq_to_hom U.open_embedding_obj_top).op),
  rw ← is_iso.comp_inv_eq at e₁,
  have e₂ := Γ_Spec.adjunction_unit_app_app_top (X.restrict U.open_embedding),
  erw ← e₂ at e₁,
  simp only [functor.id_map, quiver.hom.unop_op, functor.comp_map, ← functor.map_inv, ← op_inv,
    LocallyRingedSpace.Γ_map, category.assoc, functor.right_op_map, inv_eq_to_hom] at e₁,
  delta is_affine_open.from_Spec Scheme.iso_Spec,
  rw [Scheme.comp_val_c_app, Scheme.comp_val_c_app, ← e₁],
  simp_rw category.assoc,
  erw ← X.presheaf.map_comp_assoc,
  rw ← op_comp,
  have e₃ : U.open_embedding.is_open_map.adjunction.counit.app U ≫
    eq_to_hom U.open_embedding_obj_top.symm =
    U.open_embedding.is_open_map.functor.map (eq_to_hom U.inclusion_map_eq_top) :=
    subsingleton.elim _ _,
  have e₄ : X.presheaf.map _ ≫ _ = _ :=
    (as_iso (Γ_Spec.adjunction.unit.app (X.restrict U.open_embedding)))
    .inv.1.c.naturality_assoc (eq_to_hom U.inclusion_map_eq_top).op _,
  erw [e₃, e₄, ← Scheme.comp_val_c_app_assoc, iso.inv_hom_id],
  simp only [eq_to_hom_map, eq_to_hom_op, Scheme.Spec_map_presheaf_map_eq_to_hom],
  erw [Scheme.Spec_map_presheaf_map_eq_to_hom, category.id_comp],
  simpa only [eq_to_hom_trans]
end

@[elementwise]
lemma is_affine_open.from_Spec_app_eq {X : Scheme} {U : opens X.carrier}
  (hU : is_affine_open U) :
  hU.from_Spec.1.c.app (op U) = Spec_Γ_identity.inv.app (X.presheaf.obj $ op U) ≫
    (Scheme.Spec.obj _).presheaf.map (eq_to_hom hU.from_Spec_base_preimage).op :=
by rw [← hU.Spec_Γ_identity_hom_app_from_Spec, iso.inv_hom_id_app_assoc]

lemma is_affine_open.basic_open_is_affine {X : Scheme} {U : opens X.carrier}
  (hU : is_affine_open U) (f : X.presheaf.obj (op U)) : is_affine_open (X.basic_open f) :=
begin
  convert range_is_affine_open_of_open_immersion (Scheme.Spec.map (CommRing.of_hom
    (algebra_map (X.presheaf.obj (op U)) (localization.away f))).op ≫ hU.from_Spec),
  ext1,
  have : hU.from_Spec.val.base '' (hU.from_Spec.val.base ⁻¹' (X.basic_open f : set X.carrier)) =
    (X.basic_open f : set X.carrier),
  { rw [set.image_preimage_eq_inter_range, set.inter_eq_left_iff_subset, hU.from_Spec_range],
    exact Scheme.basic_open_subset _ _ },
  rw [subtype.coe_mk, Scheme.comp_val_base, ← this, coe_comp, set.range_comp],
  congr' 1,
  refine (congr_arg coe $ Scheme.preimage_basic_open hU.from_Spec f).trans _,
  refine eq.trans _ (prime_spectrum.localization_away_comap_range (localization.away f) f).symm,
  congr' 1,
  have : (opens.map hU.from_Spec.val.base).obj U = ⊤,
  { ext1,
    change hU.from_Spec.1.base ⁻¹' (U : set X.carrier) = set.univ,
    rw [← hU.from_Spec_range, ← set.image_univ],
    exact set.preimage_image_eq _ PresheafedSpace.is_open_immersion.base_open.inj },
  refine eq.trans _ (basic_open_eq_of_affine f),
  have lm : ∀ s, (opens.map hU.from_Spec.val.base).obj U ⊓ s = s := λ s, this.symm ▸ top_inf_eq,
  refine eq.trans _ (lm _),
  refine eq.trans _
    ((Scheme.Spec.obj $ op $ X.presheaf.obj $ op U).basic_open_res _ (eq_to_hom this).op),
  rw ← comp_apply,
  congr' 2,
  rw iso.eq_inv_comp,
  erw hU.Spec_Γ_identity_hom_app_from_Spec,
end

lemma Scheme.map_prime_spectrum_basic_open_of_affine (X : Scheme) [is_affine X]
  (f : Scheme.Γ.obj (op X)) :
  (opens.map X.iso_Spec.hom.1.base).obj (prime_spectrum.basic_open f) = X.basic_open f :=
begin
  rw ← basic_open_eq_of_affine,
  transitivity (opens.map X.iso_Spec.hom.1.base).obj ((Scheme.Spec.obj
    (op (Scheme.Γ.obj (op X)))).basic_open ((inv (X.iso_Spec.hom.1.c.app
      (op ((opens.map (inv X.iso_Spec.hom).val.base).obj ⊤)))) ((X.presheaf.map (eq_to_hom _)) f))),
  congr,
  { rw [← is_iso.inv_eq_inv, is_iso.inv_inv, is_iso.iso.inv_inv, nat_iso.app_hom],
    erw ← Γ_Spec.adjunction_unit_app_app_top,
    refl },
  { rw eq_to_hom_map, refl },
  { dsimp, congr },
  { refine (Scheme.preimage_basic_open _ _).trans _,
    rw [is_iso.inv_hom_id_apply, Scheme.basic_open_res_eq] }
end

lemma is_basis_basic_open (X : Scheme) [is_affine X] :
  opens.is_basis (set.range (X.basic_open : X.presheaf.obj (op ⊤) → opens X.carrier)) :=
begin
  delta opens.is_basis,
  convert prime_spectrum.is_basis_basic_opens.inducing
    (Top.homeo_of_iso (Scheme.forget_to_Top.map_iso X.iso_Spec)).inducing using 1,
  ext,
  simp only [set.mem_image, exists_exists_eq_and],
  split,
  { rintro ⟨_, ⟨x, rfl⟩, rfl⟩,
    refine ⟨_, ⟨_, ⟨x, rfl⟩, rfl⟩, _⟩,
    exact congr_arg subtype.val (X.map_prime_spectrum_basic_open_of_affine x) },
  { rintro ⟨_, ⟨_, ⟨x, rfl⟩, rfl⟩, rfl⟩,
    refine ⟨_, ⟨x, rfl⟩, _⟩,
    exact congr_arg subtype.val (X.map_prime_spectrum_basic_open_of_affine x).symm }
end

/-- The prime ideal of `𝒪ₓ(U)` corresponding to a point `x : U`. -/
noncomputable
def is_affine_open.prime_ideal_of {X : Scheme} {U : opens X.carrier}
  (hU : is_affine_open U) (x : U) :
  prime_spectrum (X.presheaf.obj $ op U) :=
((Scheme.Spec.map (X.presheaf.map (eq_to_hom $
  show U.open_embedding.is_open_map.functor.obj ⊤ = U, from
    opens.ext (set.image_univ.trans subtype.range_coe)).op).op).1.base
  ((@@Scheme.iso_Spec (X.restrict U.open_embedding) hU).hom.1.base x))

lemma is_affine_open.from_Spec_prime_ideal_of {X : Scheme} {U : opens X.carrier}
  (hU : is_affine_open U) (x : U) :
  hU.from_Spec.val.base (hU.prime_ideal_of x) = x.1 :=
begin
  dsimp only [is_affine_open.from_Spec, subtype.coe_mk],
  erw [← Scheme.comp_val_base_apply, ← Scheme.comp_val_base_apply],
  simpa only [← functor.map_comp_assoc, ← functor.map_comp, ← op_comp, eq_to_hom_trans, op_id,
    eq_to_hom_refl, category_theory.functor.map_id, category.id_comp, iso.hom_inv_id_assoc]
end

lemma is_affine_open.is_localization_stalk_aux {X : Scheme} (U : opens X.carrier)
  [is_affine (X.restrict U.open_embedding)] :
  (inv (Γ_Spec.adjunction.unit.app (X.restrict U.open_embedding))).1.c.app
    (op ((opens.map U.inclusion).obj U)) =
      X.presheaf.map (eq_to_hom $ by rw opens.inclusion_map_eq_top :
        U.open_embedding.is_open_map.functor.obj ⊤ ⟶
          (U.open_embedding.is_open_map.functor.obj ((opens.map U.inclusion).obj U))).op ≫
      to_Spec_Γ (X.presheaf.obj $ op (U.open_embedding.is_open_map.functor.obj ⊤)) ≫
      (Scheme.Spec.obj $ op $ X.presheaf.obj $ _).presheaf.map
        (eq_to_hom (by { rw [opens.inclusion_map_eq_top], refl }) : unop _ ⟶ ⊤).op :=
begin
  have e : (opens.map (inv (Γ_Spec.adjunction.unit.app (X.restrict U.open_embedding))).1.base).obj
    ((opens.map U.inclusion).obj U) = ⊤,
  by { rw [opens.inclusion_map_eq_top], refl },
  rw [Scheme.inv_val_c_app, is_iso.comp_inv_eq, Scheme.app_eq _ e,
    Γ_Spec.adjunction_unit_app_app_top],
  simp only [category.assoc, eq_to_hom_op],
  erw ← functor.map_comp_assoc,
  rw [eq_to_hom_trans, eq_to_hom_refl, category_theory.functor.map_id,
    category.id_comp],
  erw Spec_Γ_identity.inv_hom_id_app_assoc,
  simp only [eq_to_hom_map, eq_to_hom_trans],
end

lemma is_affine_open.is_localization_stalk {X : Scheme} {U : opens X.carrier}
  (hU : is_affine_open U) (x : U) :
  is_localization.at_prime (X.presheaf.stalk x) (hU.prime_ideal_of x).as_ideal :=
begin
  haveI : is_affine _ := hU,
  haveI : nonempty U := ⟨x⟩,
  rcases x with ⟨x, hx⟩,
  let y := hU.prime_ideal_of ⟨x, hx⟩,
  have : hU.from_Spec.val.base y = x := hU.from_Spec_prime_ideal_of ⟨x, hx⟩,
  change is_localization y.as_ideal.prime_compl _,
  clear_value y,
  subst this,
  apply (is_localization.is_localization_iff_of_ring_equiv _
    (as_iso $ PresheafedSpace.stalk_map hU.from_Spec.1 y).CommRing_iso_to_ring_equiv).mpr,
  convert structure_sheaf.is_localization.to_stalk _ _ using 1,
  delta structure_sheaf.stalk_algebra,
  congr' 1,
  rw ring_hom.algebra_map_to_algebra,
  refine (PresheafedSpace.stalk_map_germ hU.from_Spec.1 _ ⟨_, _⟩).trans _,
  delta is_affine_open.from_Spec Scheme.iso_Spec structure_sheaf.to_stalk,
  simp only [Scheme.comp_val_c_app, category.assoc],
  dsimp only [functor.op, as_iso_inv, unop_op],
  erw is_affine_open.is_localization_stalk_aux,
  simp only [category.assoc],
  conv_lhs { rw ← category.assoc },
  erw [← X.presheaf.map_comp, Spec_Γ_naturality_assoc],
  congr' 1,
  simp only [← category.assoc],
  transitivity _ ≫ (structure_sheaf (X.presheaf.obj $ op U)).1.germ ⟨_, _⟩,
  { refl },
  convert ((structure_sheaf (X.presheaf.obj $ op U)).1.germ_res (hom_of_le le_top) ⟨_, _⟩) using 2,
  rw category.assoc,
  erw nat_trans.naturality,
  rw [← LocallyRingedSpace.Γ_map_op, ← LocallyRingedSpace.Γ.map_comp_assoc, ← op_comp],
  erw ← Scheme.Spec.map_comp,
  rw [← op_comp, ← X.presheaf.map_comp],
  transitivity LocallyRingedSpace.Γ.map (quiver.hom.op $ Scheme.Spec.map
    (X.presheaf.map (𝟙 (op U))).op) ≫ _,
  { congr },
  simp only [category_theory.functor.map_id, op_id],
  erw category_theory.functor.map_id,
  rw category.id_comp,
  refl
end

end algebraic_geometry