/-
Copyright (c) 2020 Anne Baanen. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Anne Baanen, Kexing Ying, Eric Wieser
-/

import algebra.invertible
import linear_algebra.matrix.determinant
import linear_algebra.matrix.bilinear_form
import linear_algebra.matrix.symmetric

/-!
# Quadratic forms

This file defines quadratic forms over a `R`-module `M`.
A quadratic form on a ring `R` is a map `Q : M → R` such that:
* `quadratic_form.map_smul`: `Q (a • x) = a * a * Q x`
* `quadratic_form.polar_add_left`, `quadratic_form.polar_add_right`,
  `quadratic_form.polar_smul_left`, `quadratic_form.polar_smul_right`:
  the map `quadratic_form.polar Q := λ x y, Q (x + y) - Q x - Q y` is bilinear.

This notion generalizes to semirings using the approach in [izhakian2016][] which requires that
there be a (possibly non-unique) companion bilinear form `B` such that
`∀ x y, Q (x + y) = Q x + Q y + B x y`. Over a ring, this `B` is precisely `quadratic_form.polar Q`.

To build a `quadratic_form` from the `polar` axioms, use `quadratic_form.of_polar`.

Quadratic forms come with a scalar multiplication, `(a • Q) x = Q (a • x) = a * a * Q x`,
and composition with linear maps `f`, `Q.comp f x = Q (f x)`.

## Main definitions

 * `quadratic_form.of_polar`: a more familiar constructor that works on rings
 * `quadratic_form.associated`: associated bilinear form
 * `quadratic_form.pos_def`: positive definite quadratic forms
 * `quadratic_form.anisotropic`: anisotropic quadratic forms
 * `quadratic_form.discr`: discriminant of a quadratic form

## Main statements

 * `quadratic_form.associated_left_inverse`,
 * `quadratic_form.associated_right_inverse`: in a commutative ring where 2 has
  an inverse, there is a correspondence between quadratic forms and symmetric
  bilinear forms
 * `bilin_form.exists_orthogonal_basis`: There exists an orthogonal basis with
  respect to any nondegenerate, symmetric bilinear form `B`.

## Notation

In this file, the variable `R` is used when a `ring` structure is sufficient and
`R₁` is used when specifically a `comm_ring` is required. This allows us to keep
`[module R M]` and `[module R₁ M]` assumptions in the variables without
confusion between `*` from `ring` and `*` from `comm_ring`.

The variable `S` is used when `R` itself has a `•` action.

## References

 * https://en.wikipedia.org/wiki/Quadratic_form
 * https://en.wikipedia.org/wiki/Discriminant#Quadratic_forms

## Tags

quadratic form, homogeneous polynomial, quadratic polynomial
-/

universes u v w
variables {S : Type*}
variables {R R₁: Type*} {M : Type*}

open_locale big_operators

section polar
variables [ring R] [comm_ring R₁] [add_comm_group M]

namespace quadratic_form

/-- Up to a factor 2, `Q.polar` is the associated bilinear form for a quadratic form `Q`.

Source of this name: https://en.wikipedia.org/wiki/Quadratic_form#Generalization
-/
def polar (f : M → R) (x y : M) :=
f (x + y) - f x - f y

lemma polar_add (f g : M → R) (x y : M) :
  polar (f + g) x y = polar f x y + polar g x y :=
by { simp only [polar, pi.add_apply], abel }

lemma polar_neg (f : M → R) (x y : M) :
  polar (-f) x y = - polar f x y :=
by { simp only [polar, pi.neg_apply, sub_eq_add_neg, neg_add] }

lemma polar_smul [monoid S] [distrib_mul_action S R] (f : M → R) (s : S) (x y : M) :
  polar (s • f) x y = s • polar f x y :=
by { simp only [polar, pi.smul_apply, smul_sub] }

lemma polar_comm (f : M → R) (x y : M) : polar f x y = polar f y x :=
by rw [polar, polar, add_comm, sub_sub, sub_sub, add_comm (f x) (f y)]

/-- Auxiliary lemma to express bilinearity of `quadratic_form.polar` without subtraction. -/
lemma polar_add_left_iff {f : M → R} {x x' y : M} :
  polar f (x + x') y = polar f x y + polar f x' y ↔
    f (x + x' + y) + (f x + f x' + f y) = f (x + x') + f (x' + y) + f (y + x) :=
begin
  simp only [←add_assoc],
  simp only [polar, sub_eq_iff_eq_add, eq_sub_iff_add_eq, sub_add_eq_add_sub, add_sub],
  simp only [add_right_comm _ (f y) _, add_right_comm _ (f x') (f x)],
  rw [add_comm y x, add_right_comm _ _ (f (x + y)), add_comm _ (f (x + y)),
    add_right_comm (f (x + y)), add_left_inj],
end

lemma polar_comp {F : Type*} [ring S] [add_monoid_hom_class F R S] (f : M → R) (g : F) (x y : M) :
  polar (g ∘ f) x y = g (polar f x y) :=
by simp only [polar, pi.smul_apply, function.comp_apply, map_sub]

end quadratic_form

end polar

/-- A quadratic form over a module.

For a more familiar constructor when `R` is a ring, see `quadratic_form.of_polar`. -/
structure quadratic_form (R : Type u) (M : Type v) [semiring R] [add_comm_monoid M] [module R M] :=
(to_fun : M → R)
(to_fun_smul : ∀ (a : R) (x : M), to_fun (a • x) = a * a * to_fun x)
(exists_companion' : ∃ B : bilin_form R M, ∀ x y, to_fun (x + y) = to_fun x + to_fun y + B x y)

namespace quadratic_form

section fun_like
variables [semiring R] [add_comm_monoid M] [module R M]
variables {Q Q' : quadratic_form R M}

instance fun_like : fun_like (quadratic_form R M) M (λ _, R) :=
{ coe := to_fun,
  coe_injective' := λ x y h, by cases x; cases y; congr' }

/-- Helper instance for when there's too many metavariables to apply
`fun_like.has_coe_to_fun` directly. -/
instance : has_coe_to_fun (quadratic_form R M) (λ _, M → R) := ⟨to_fun⟩

variables (Q)

/-- The `simp` normal form for a quadratic form is `coe_fn`, not `to_fun`. -/
@[simp] lemma to_fun_eq_coe : Q.to_fun = ⇑Q := rfl

-- this must come after the coe_to_fun definition
initialize_simps_projections quadratic_form (to_fun → apply)

variables {Q}

@[ext] lemma ext (H : ∀ (x : M), Q x = Q' x) : Q = Q' := fun_like.ext _ _ H

lemma congr_fun (h : Q = Q') (x : M) : Q x = Q' x := fun_like.congr_fun h _

lemma ext_iff : Q = Q' ↔ (∀ x, Q x = Q' x) := fun_like.ext_iff

/-- Copy of a `quadratic_form` with a new `to_fun` equal to the old one. Useful to fix definitional
equalities. -/
protected def copy (Q : quadratic_form R M) (Q' : M → R) (h : Q' = ⇑Q) : quadratic_form R M :=
{ to_fun := Q',
  to_fun_smul := h.symm ▸ Q.to_fun_smul,
  exists_companion' := h.symm ▸ Q.exists_companion' }

@[simp]
lemma coe_copy (Q : quadratic_form R M) (Q' : M → R) (h : Q' = ⇑Q) : ⇑(Q.copy Q' h) = Q' := rfl

lemma copy_eq (Q : quadratic_form R M) (Q' : M → R) (h : Q' = ⇑Q) : Q.copy Q' h = Q :=
fun_like.ext' h

end fun_like

section semiring
variables [semiring R] [add_comm_monoid M] [module R M]
variables (Q : quadratic_form R M)

lemma map_smul (a : R) (x : M) : Q (a • x) = a * a * Q x := Q.to_fun_smul a x

lemma exists_companion : ∃ B : bilin_form R M, ∀ x y, Q (x + y) = Q x + Q y + B x y :=
Q.exists_companion'

lemma map_add_add_add_map (x y z : M) :
  Q (x + y + z) + (Q x + Q y + Q z) = Q (x + y) + Q (y + z) + Q (z + x) :=
begin
  obtain ⟨B, h⟩ := Q.exists_companion,
  rw [add_comm z x],
  simp [h],
  abel,
end

lemma map_add_self (x : M) : Q (x + x) = 4 * Q x :=
by { rw [←one_smul R x, ←add_smul, map_smul], norm_num }

@[simp] lemma map_zero : Q 0 = 0 :=
by rw [←@zero_smul R _ _ _ _ (0 : M), map_smul, zero_mul, zero_mul]

instance zero_hom_class : zero_hom_class (quadratic_form R M) M R :=
{ map_zero := map_zero,
  ..quadratic_form.fun_like }

lemma map_smul_of_tower [comm_semiring S] [algebra S R] [module S M] [is_scalar_tower S R M]
  (a : S) (x : M) :
  Q (a • x) = (a * a) • Q x :=
by rw [←is_scalar_tower.algebra_map_smul R a x, map_smul, ←ring_hom.map_mul, algebra.smul_def]

end semiring

section ring
variables [ring R] [comm_ring R₁] [add_comm_group M]
variables [module R M] (Q : quadratic_form R M)

@[simp] lemma map_neg (x : M) : Q (-x) = Q x :=
by rw [←@neg_one_smul R _ _ _ _ x, map_smul, neg_one_mul, neg_neg, one_mul]

lemma map_sub (x y : M) : Q (x - y) = Q (y - x) :=
by rw [←neg_sub, map_neg]

@[simp]
lemma polar_zero_left (y : M) : polar Q 0 y = 0 :=
by simp only [polar, zero_add, quadratic_form.map_zero, sub_zero, sub_self]

@[simp]
lemma polar_add_left (x x' y : M) :
  polar Q (x + x') y = polar Q x y + polar Q x' y :=
polar_add_left_iff.mpr $ Q.map_add_add_add_map x x' y

@[simp]
lemma polar_smul_left (a : R) (x y : M) :
  polar Q (a • x) y = a * polar Q x y :=
begin
  obtain ⟨B, h⟩ := Q.exists_companion,
  simp_rw [polar, h, Q.map_smul, bilin_form.smul_left, sub_sub, add_sub_cancel'],
end

@[simp]
lemma polar_neg_left (x y : M) :
  polar Q (-x) y = -polar Q x y :=
by rw [←neg_one_smul R x, polar_smul_left, neg_one_mul]

@[simp]
lemma polar_sub_left (x x' y : M) :
  polar Q (x - x') y = polar Q x y - polar Q x' y :=
by rw [sub_eq_add_neg, sub_eq_add_neg, polar_add_left, polar_neg_left]

@[simp]
lemma polar_zero_right (y : M) : polar Q y 0 = 0 :=
by simp only [add_zero, polar, quadratic_form.map_zero, sub_self]

@[simp]
lemma polar_add_right (x y y' : M) :
  polar Q x (y + y') = polar Q x y + polar Q x y' :=
by rw [polar_comm Q x, polar_comm Q x, polar_comm Q x, polar_add_left]

@[simp]
lemma polar_smul_right (a : R) (x y : M) :
  polar Q x (a • y) = a * polar Q x y :=
by rw [polar_comm Q x, polar_comm Q x, polar_smul_left]

@[simp]
lemma polar_neg_right (x y : M) :
  polar Q x (-y) = -polar Q x y :=
by rw [←neg_one_smul R y, polar_smul_right, neg_one_mul]

@[simp]
lemma polar_sub_right (x y y' : M) :
  polar Q x (y - y') = polar Q x y - polar Q x y' :=
by rw [sub_eq_add_neg, sub_eq_add_neg, polar_add_right, polar_neg_right]

@[simp]
lemma polar_self (x : M) : polar Q x x = 2 * Q x :=
begin
  rw [polar, map_add_self, sub_sub, sub_eq_iff_eq_add, ←two_mul, ←two_mul, ←mul_assoc],
  norm_num
end

/-- `quadratic_form.polar` as a bilinear form -/
@[simps]
def polar_bilin : bilin_form R M :=
{ bilin := polar Q,
  bilin_add_left := polar_add_left Q,
  bilin_smul_left := polar_smul_left Q,
  bilin_add_right := λ x y z, by simp_rw [polar_comm _ x, polar_add_left Q],
  bilin_smul_right := λ r x y, by simp_rw [polar_comm _ x, polar_smul_left Q] }

variables [comm_semiring S] [algebra S R] [module S M] [is_scalar_tower S R M]

@[simp]
lemma polar_smul_left_of_tower (a : S) (x y : M) :
  polar Q (a • x) y = a • polar Q x y :=
by rw [←is_scalar_tower.algebra_map_smul R a x, polar_smul_left, algebra.smul_def]

@[simp]
lemma polar_smul_right_of_tower (a : S) (x y : M) :
  polar Q x (a • y) = a • polar Q x y :=
by rw [←is_scalar_tower.algebra_map_smul R a y, polar_smul_right, algebra.smul_def]

/-- An alternative constructor to `quadratic_form.mk`, for rings where `polar` can be used. -/
@[simps]
def of_polar (to_fun : M → R) (to_fun_smul : ∀ (a : R) (x : M), to_fun (a • x) = a * a * to_fun x)
  (polar_add_left : ∀ (x x' y : M), polar to_fun (x + x') y = polar to_fun x y + polar to_fun x' y)
  (polar_smul_left : ∀ (a : R) (x y : M), polar to_fun (a • x) y = a • polar to_fun x y) :
  quadratic_form R M :=
{ to_fun := to_fun,
  to_fun_smul := to_fun_smul,
  exists_companion' := ⟨
    { bilin := polar to_fun,
      bilin_add_left := polar_add_left,
      bilin_smul_left := polar_smul_left,
      bilin_add_right := λ x y z, by simp_rw [polar_comm _ x, polar_add_left],
      bilin_smul_right := λ r x y, by simp_rw [polar_comm _ x, polar_smul_left, smul_eq_mul] },
    λ x y, by rw [bilin_form.coe_fn_mk, polar, sub_sub, add_sub_cancel'_right]⟩ }

/-- In a ring the companion bilinear form is unique and equal to `quadratic_form.polar`. -/
lemma some_exists_companion : Q.exists_companion.some = polar_bilin Q :=
bilin_form.ext $ λ x y,
  by rw [polar_bilin_apply, polar, Q.exists_companion.some_spec, sub_sub, add_sub_cancel']

end ring

section semiring_operators
variables [semiring R] [add_comm_monoid M] [module R M]

section has_smul

variables [monoid S] [distrib_mul_action S R] [smul_comm_class S R R]

/-- `quadratic_form R M` inherits the scalar action from any algebra over `R`.

When `R` is commutative, this provides an `R`-action via `algebra.id`. -/
instance : has_smul S (quadratic_form R M) :=
⟨ λ a Q,
  { to_fun := a • Q,
    to_fun_smul := λ b x, by rw [pi.smul_apply, map_smul, pi.smul_apply, mul_smul_comm],
    exists_companion' := let ⟨B, h⟩ := Q.exists_companion in ⟨a • B,
      by simp [h]⟩ } ⟩

@[simp] lemma coe_fn_smul (a : S) (Q : quadratic_form R M) : ⇑(a • Q) = a • Q := rfl

@[simp] lemma smul_apply (a : S) (Q : quadratic_form R M) (x : M) :
  (a • Q) x = a • Q x := rfl

end has_smul

instance : has_zero (quadratic_form R M) :=
⟨ { to_fun := λ x, 0,
    to_fun_smul := λ a x, by simp only [mul_zero],
    exists_companion' := ⟨0, λ x y, by simp only [add_zero, bilin_form.zero_apply]⟩ } ⟩

@[simp] lemma coe_fn_zero : ⇑(0 : quadratic_form R M) = 0 := rfl

@[simp] lemma zero_apply (x : M) : (0 : quadratic_form R M) x = 0 := rfl

instance : inhabited (quadratic_form R M) := ⟨0⟩

instance : has_add (quadratic_form R M) :=
⟨ λ Q Q',
  { to_fun := Q + Q',
    to_fun_smul := λ a x,
      by simp only [pi.add_apply, map_smul, mul_add],
    exists_companion' :=
      let ⟨B, h⟩ := Q.exists_companion, ⟨B', h'⟩ := Q'.exists_companion in
      ⟨B + B', λ x y, by simp_rw [pi.add_apply, h, h', bilin_form.add_apply, add_add_add_comm] ⟩ } ⟩

@[simp] lemma coe_fn_add (Q Q' : quadratic_form R M) : ⇑(Q + Q') = Q + Q' := rfl

@[simp] lemma add_apply (Q Q' : quadratic_form R M) (x : M) : (Q + Q') x = Q x + Q' x := rfl

instance : add_comm_monoid (quadratic_form R M) :=
fun_like.coe_injective.add_comm_monoid _ coe_fn_zero coe_fn_add (λ _ _, coe_fn_smul _ _)

/-- `@coe_fn (quadratic_form R M)` as an `add_monoid_hom`.

This API mirrors `add_monoid_hom.coe_fn`. -/
@[simps apply]
def coe_fn_add_monoid_hom : quadratic_form R M →+ (M → R) :=
{ to_fun := coe_fn, map_zero' := coe_fn_zero, map_add' := coe_fn_add }

/-- Evaluation on a particular element of the module `M` is an additive map over quadratic forms. -/
@[simps apply]
def eval_add_monoid_hom (m : M) : quadratic_form R M →+ R :=
(pi.eval_add_monoid_hom _ m).comp coe_fn_add_monoid_hom

section sum

@[simp] lemma coe_fn_sum {ι : Type*} (Q : ι → quadratic_form R M) (s : finset ι) :
  ⇑(∑ i in s, Q i) = ∑ i in s, Q i :=
(coe_fn_add_monoid_hom : _ →+ (M → R)).map_sum Q s

@[simp] lemma sum_apply {ι : Type*} (Q : ι → quadratic_form R M) (s : finset ι) (x : M) :
  (∑ i in s, Q i) x = ∑ i in s, Q i x :=
(eval_add_monoid_hom x : _ →+ R).map_sum Q s

end sum

instance [monoid S] [distrib_mul_action S R] [smul_comm_class S R R] :
  distrib_mul_action S (quadratic_form R M) :=
{ mul_smul := λ a b Q, ext (λ x, by simp only [smul_apply, mul_smul]),
  one_smul := λ Q, ext (λ x, by simp only [quadratic_form.smul_apply, one_smul]),
  smul_add := λ a Q Q', by { ext, simp only [add_apply, smul_apply, smul_add] },
  smul_zero := λ a, by { ext, simp only [zero_apply, smul_apply, smul_zero] }, }

instance [semiring S] [module S R] [smul_comm_class S R R] : module S (quadratic_form R M) :=
{ zero_smul := λ Q, by { ext, simp only [zero_apply, smul_apply, zero_smul] },
  add_smul := λ a b Q, by { ext, simp only [add_apply, smul_apply, add_smul] } }

end semiring_operators

section ring_operators
variables [ring R] [add_comm_group M] [module R M]

instance : has_neg (quadratic_form R M) :=
⟨ λ Q,
  { to_fun := -Q,
    to_fun_smul := λ a x,
      by simp only [pi.neg_apply, map_smul, mul_neg],
    exists_companion' :=
      let ⟨B, h⟩ := Q.exists_companion in
      ⟨-B, λ x y, by simp_rw [pi.neg_apply, h, bilin_form.neg_apply, neg_add] ⟩ } ⟩

@[simp] lemma coe_fn_neg (Q : quadratic_form R M) : ⇑(-Q) = -Q := rfl

@[simp] lemma neg_apply (Q : quadratic_form R M) (x : M) : (-Q) x = -Q x := rfl

instance : has_sub (quadratic_form R M) :=
⟨ λ Q Q', (Q + -Q').copy (Q - Q') (sub_eq_add_neg _ _) ⟩

@[simp] lemma coe_fn_sub (Q Q' : quadratic_form R M) : ⇑(Q - Q') = Q - Q' := rfl

@[simp] lemma sub_apply (Q Q' : quadratic_form R M) (x : M) : (Q - Q') x = Q x - Q' x := rfl

instance : add_comm_group (quadratic_form R M) :=
fun_like.coe_injective.add_comm_group _
  coe_fn_zero coe_fn_add coe_fn_neg coe_fn_sub (λ _ _, coe_fn_smul _ _) (λ _ _, coe_fn_smul _ _)

end ring_operators

section comp

variables [semiring R] [add_comm_monoid M] [module R M]
variables {N : Type v} [add_comm_monoid N] [module R N]

/-- Compose the quadratic form with a linear function. -/
def comp (Q : quadratic_form R N) (f : M →ₗ[R] N) :
  quadratic_form R M :=
{ to_fun := λ x, Q (f x),
  to_fun_smul := λ a x, by simp only [map_smul, f.map_smul],
  exists_companion' :=
    let ⟨B, h⟩ := Q.exists_companion in
    ⟨B.comp f f, λ x y, by simp_rw [f.map_add, h, bilin_form.comp_apply]⟩ }

@[simp] lemma comp_apply (Q : quadratic_form R N) (f : M →ₗ[R] N) (x : M) :
  (Q.comp f) x = Q (f x) := rfl

/-- Compose a quadratic form with a linear function on the left. -/
@[simps {simp_rhs := tt}]
def _root_.linear_map.comp_quadratic_form {S : Type*}
  [comm_semiring S] [algebra S R] [module S M] [is_scalar_tower S R M]
  (f : R →ₗ[S] S) (Q : quadratic_form R M) :
  quadratic_form S M :=
{ to_fun := λ x, f (Q x),
  to_fun_smul := λ b x, by rw [Q.map_smul_of_tower b x, f.map_smul, smul_eq_mul],
  exists_companion' :=
    let ⟨B, h⟩ := Q.exists_companion in
    ⟨f.comp_bilin_form B, λ x y, by simp_rw [h, f.map_add, linear_map.comp_bilin_form_apply]⟩ }

end comp

section comm_ring
variables [comm_semiring R] [add_comm_monoid M] [module R M]

/-- The product of linear forms is a quadratic form. -/
def lin_mul_lin (f g : M →ₗ[R] R) : quadratic_form R M :=
{ to_fun := f * g,
  to_fun_smul := λ a x,
    by { simp only [smul_eq_mul, ring_hom.id_apply, pi.mul_apply, linear_map.map_smulₛₗ], ring },
  exists_companion' := ⟨
    bilin_form.lin_mul_lin f g + bilin_form.lin_mul_lin g f, λ x y, by { simp, ring }⟩ }

@[simp]
lemma lin_mul_lin_apply (f g : M →ₗ[R] R) (x) : lin_mul_lin f g x = f x * g x := rfl

@[simp]
lemma add_lin_mul_lin (f g h : M →ₗ[R] R) :
  lin_mul_lin (f + g) h = lin_mul_lin f h + lin_mul_lin g h :=
ext (λ x, add_mul _ _ _)

@[simp]
lemma lin_mul_lin_add (f g h : M →ₗ[R] R) :
  lin_mul_lin f (g + h) = lin_mul_lin f g + lin_mul_lin f h :=
ext (λ x, mul_add _ _ _)

variables {N : Type v} [add_comm_monoid N] [module R N]

@[simp]
lemma lin_mul_lin_comp (f g : M →ₗ[R] R) (h : N →ₗ[R] M) :
  (lin_mul_lin f g).comp h = lin_mul_lin (f.comp h) (g.comp h) :=
rfl

variables {n : Type*}

/-- `sq` is the quadratic form mapping the vector `x : R₁` to `x * x` -/
@[simps]
def sq : quadratic_form R R :=
lin_mul_lin linear_map.id linear_map.id

/-- `proj i j` is the quadratic form mapping the vector `x : n → R₁` to `x i * x j` -/
def proj (i j : n) : quadratic_form R (n → R) :=
lin_mul_lin (@linear_map.proj _ _ _ (λ _, R) _ _ i) (@linear_map.proj _ _ _ (λ _, R) _ _ j)

@[simp]
lemma proj_apply (i j : n) (x : n → R) : proj i j x = x i * x j := rfl

end comm_ring

end quadratic_form

/-!
### Associated bilinear forms

Over a commutative ring with an inverse of 2, the theory of quadratic forms is
basically identical to that of symmetric bilinear forms. The map from quadratic
forms to bilinear forms giving this identification is called the `associated`
quadratic form.
-/

namespace bilin_form
open quadratic_form

section semiring

variables [semiring R] [add_comm_monoid M] [module R M]
variables {B : bilin_form R M}

/-- A bilinear form gives a quadratic form by applying the argument twice. -/
def to_quadratic_form (B : bilin_form R M) : quadratic_form R M :=
{ to_fun := λ x, B x x,
  to_fun_smul := λ a x, by simp only [mul_assoc, smul_right, smul_left],
  exists_companion' :=
    ⟨B + bilin_form.flip_hom ℕ B, λ x y, by { simp [add_add_add_comm, add_comm] }⟩ }

@[simp] lemma to_quadratic_form_apply (B : bilin_form R M) (x : M) :
  B.to_quadratic_form x = B x x :=
rfl

section
variables (R M)
@[simp] lemma to_quadratic_form_zero : (0 : bilin_form R M).to_quadratic_form = 0 := rfl
end

@[simp] lemma to_quadratic_form_add (B₁ B₂ : bilin_form R M) :
  (B₁ + B₂).to_quadratic_form = B₁.to_quadratic_form + B₂.to_quadratic_form := rfl

@[simp] lemma to_quadratic_form_smul [monoid S] [distrib_mul_action S R] [smul_comm_class S R R]
  (a : S) (B : bilin_form R M) :
  (a • B).to_quadratic_form = a • B.to_quadratic_form := rfl

section
variables (R M)
/-- `bilin_form.to_quadratic_form` as an additive homomorphism -/
@[simps] def to_quadratic_form_add_monoid_hom : bilin_form R M →+ quadratic_form R M :=
{ to_fun := to_quadratic_form,
  map_zero' := to_quadratic_form_zero _ _,
  map_add' := to_quadratic_form_add }
end

@[simp] lemma to_quadratic_form_list_sum (B : list (bilin_form R M)) :
  B.sum.to_quadratic_form = (B.map to_quadratic_form).sum :=
map_list_sum (to_quadratic_form_add_monoid_hom R M) B

@[simp] lemma to_quadratic_form_multiset_sum (B : multiset (bilin_form R M)) :
  B.sum.to_quadratic_form = (B.map to_quadratic_form).sum :=
map_multiset_sum (to_quadratic_form_add_monoid_hom R M) B

@[simp] lemma to_quadratic_form_sum {ι : Type*} (s : finset ι) (B : ι → bilin_form R M) :
  (∑ i in s, B i).to_quadratic_form = ∑ i in s, (B i).to_quadratic_form :=
map_sum (to_quadratic_form_add_monoid_hom R M) B s

@[simp] lemma to_quadratic_form_eq_zero {B : bilin_form R M} :
  B.to_quadratic_form = 0 ↔ B.is_alt :=
quadratic_form.ext_iff

end semiring

section ring
variables [ring R] [add_comm_group M] [module R M]
variables {B : bilin_form R M}

lemma polar_to_quadratic_form (x y : M) : polar (λ x, B x x) x y = B x y + B y x :=
by { simp only [add_assoc, add_sub_cancel', add_right, polar, add_left_inj, add_neg_cancel_left,
  add_left, sub_eq_add_neg _ (B y y), add_comm (B y x) _] }

@[simp] lemma to_quadratic_form_neg (B : bilin_form R M) :
  (-B).to_quadratic_form = -B.to_quadratic_form := rfl

@[simp] lemma to_quadratic_form_sub (B₁ B₂ : bilin_form R M) :
  (B₁ - B₂).to_quadratic_form = B₁.to_quadratic_form - B₂.to_quadratic_form := rfl

end ring

end bilin_form

namespace quadratic_form
open bilin_form

section associated_hom
variables [ring R] [comm_ring R₁] [add_comm_group M] [module R M] [module R₁ M]
variables (S) [comm_semiring S] [algebra S R]
variables [invertible (2 : R)] {B₁ : bilin_form R M}

/-- `associated_hom` is the map that sends a quadratic form on a module `M` over `R` to its
associated symmetric bilinear form.  As provided here, this has the structure of an `S`-linear map
where `S` is a commutative subring of `R`.

Over a commutative ring, use `associated`, which gives an `R`-linear map.  Over a general ring with
no nontrivial distinguished commutative subring, use `associated'`, which gives an additive
homomorphism (or more precisely a `ℤ`-linear map.) -/
def associated_hom : quadratic_form R M →ₗ[S] bilin_form R M :=
{ to_fun := λ Q,
  ((•) : submonoid.center R → bilin_form R M → bilin_form R M)
    (⟨⅟2, λ x, (commute.one_right x).bit0_right.inv_of_right⟩) Q.polar_bilin,
  map_add' := λ Q Q', by { ext, simp only [bilin_form.add_apply, bilin_form.smul_apply, coe_fn_mk,
    polar_bilin_apply, polar_add, coe_fn_add, smul_add] },
  map_smul' := λ s Q, by { ext, simp only [ring_hom.id_apply, polar_smul, smul_comm s,
    polar_bilin_apply, coe_fn_mk, coe_fn_smul, bilin_form.smul_apply] } }

variables (Q : quadratic_form R M) (S)

@[simp] lemma associated_apply (x y : M) :
  associated_hom S Q x y = ⅟2 * (Q (x + y) - Q x - Q y) := rfl

lemma associated_is_symm : (associated_hom S Q).is_symm :=
λ x y, by simp only [associated_apply, add_comm, add_left_comm, sub_eq_add_neg]

@[simp] lemma associated_comp {N : Type v} [add_comm_group N] [module R N] (f : N →ₗ[R] M) :
  associated_hom S (Q.comp f) = (associated_hom S Q).comp f f :=
by { ext, simp only [quadratic_form.comp_apply, bilin_form.comp_apply, associated_apply,
  linear_map.map_add] }

lemma associated_to_quadratic_form (B : bilin_form R M) (x y : M) :
  associated_hom S B.to_quadratic_form x y = ⅟2 * (B x y + B y x) :=
by simp only [associated_apply, ← polar_to_quadratic_form, polar, to_quadratic_form_apply]

lemma associated_left_inverse (h : B₁.is_symm) :
  associated_hom S (B₁.to_quadratic_form) = B₁ :=
bilin_form.ext $ λ x y,
by rw [associated_to_quadratic_form, is_symm.eq h x y, ←two_mul, ←mul_assoc, inv_of_mul_self,
  one_mul]

lemma to_quadratic_form_associated : (associated_hom S Q).to_quadratic_form = Q :=
quadratic_form.ext $ λ x,
  calc (associated_hom S Q).to_quadratic_form x
      = ⅟2 * (Q x + Q x) : by simp only [add_assoc, add_sub_cancel', one_mul,
    to_quadratic_form_apply, add_mul, associated_apply, map_add_self, bit0]
  ... = Q x : by rw [← two_mul (Q x), ←mul_assoc, inv_of_mul_self, one_mul]

-- note: usually `right_inverse` lemmas are named the other way around, but this is consistent
-- with historical naming in this file.
lemma associated_right_inverse :
  function.right_inverse (associated_hom S)
    (bilin_form.to_quadratic_form : _ → quadratic_form R M) :=
λ Q, to_quadratic_form_associated S Q

lemma associated_eq_self_apply (x : M) : associated_hom S Q x x = Q x :=
begin
  rw [associated_apply, map_add_self],
  suffices : (⅟2) * (2 * Q x) = Q x,
  { convert this,
    simp only [bit0, add_mul, one_mul],
    abel },
  simp only [← mul_assoc, one_mul, inv_of_mul_self],
end

/-- `associated'` is the `ℤ`-linear map that sends a quadratic form on a module `M` over `R` to its
associated symmetric bilinear form. -/
abbreviation associated' : quadratic_form R M →ₗ[ℤ] bilin_form R M :=
associated_hom ℤ

/-- Symmetric bilinear forms can be lifted to quadratic forms -/
instance can_lift :
  can_lift (bilin_form R M) (quadratic_form R M) (associated_hom ℕ) bilin_form.is_symm :=
{ prf := λ B hB, ⟨B.to_quadratic_form, associated_left_inverse _ hB⟩ }

/-- There exists a non-null vector with respect to any quadratic form `Q` whose associated
bilinear form is non-zero, i.e. there exists `x` such that `Q x ≠ 0`. -/
lemma exists_quadratic_form_ne_zero {Q : quadratic_form R M} (hB₁ : Q.associated' ≠ 0) :
  ∃ x, Q x ≠ 0 :=
begin
  rw ←not_forall,
  intro h,
  apply hB₁,
  rw [(quadratic_form.ext h : Q = 0), linear_map.map_zero],
end

end associated_hom

section associated
variables [comm_ring R₁] [add_comm_group M] [module R₁ M]
variables [invertible (2 : R₁)]

-- Note:  When possible, rather than writing lemmas about `associated`, write a lemma applying to
-- the more general `associated_hom` and place it in the previous section.

/-- `associated` is the linear map that sends a quadratic form over a commutative ring to its
associated symmetric bilinear form. -/
abbreviation associated : quadratic_form R₁ M →ₗ[R₁] bilin_form R₁ M :=
associated_hom R₁

@[simp] lemma associated_lin_mul_lin (f g : M →ₗ[R₁] R₁) :
  (lin_mul_lin f g).associated =
    ⅟(2 : R₁) • (bilin_form.lin_mul_lin f g + bilin_form.lin_mul_lin g f) :=
by { ext, simp only [smul_add, algebra.id.smul_eq_mul, bilin_form.lin_mul_lin_apply,
  quadratic_form.lin_mul_lin_apply, bilin_form.smul_apply, associated_apply, bilin_form.add_apply,
  linear_map.map_add], ring }

end associated

section anisotropic
section semiring
variables [semiring R] [add_comm_monoid M] [module R M]

/-- An anisotropic quadratic form is zero only on zero vectors. -/
def anisotropic (Q : quadratic_form R M) : Prop := ∀ x, Q x = 0 → x = 0

lemma not_anisotropic_iff_exists (Q : quadratic_form R M) :
  ¬anisotropic Q ↔ ∃ x ≠ 0, Q x = 0 :=
by simp only [anisotropic, not_forall, exists_prop, and_comm]

lemma anisotropic.eq_zero_iff {Q : quadratic_form R M} (h : anisotropic Q) {x : M} :
  Q x = 0 ↔ x = 0 :=
⟨h x, λ h, h.symm ▸ map_zero Q⟩

end semiring
section ring
variables [ring R] [add_comm_group M] [module R M]

/-- The associated bilinear form of an anisotropic quadratic form is nondegenerate. -/
lemma nondegenerate_of_anisotropic [invertible (2 : R)] (Q : quadratic_form R M)
  (hB : Q.anisotropic) : Q.associated'.nondegenerate :=
begin
  intros x hx,
  refine hB _ _,
  rw ← hx x,
  exact (associated_eq_self_apply _ _ x).symm,
end

end ring

end anisotropic

section pos_def

variables {R₂ : Type u} [ordered_ring R₂] [add_comm_monoid M] [module R₂ M]
variables {Q₂ : quadratic_form R₂ M}

/-- A positive definite quadratic form is positive on nonzero vectors. -/
def pos_def (Q₂ : quadratic_form R₂ M) : Prop := ∀ x ≠ 0, 0 < Q₂ x

lemma pos_def.smul {R} [linear_ordered_comm_ring R] [module R M]
  {Q : quadratic_form R M} (h : pos_def Q) {a : R} (a_pos : 0 < a) : pos_def (a • Q) :=
λ x hx, mul_pos a_pos (h x hx)

variables {n : Type*}

lemma pos_def.nonneg {Q : quadratic_form R₂ M} (hQ : pos_def Q) (x : M) :
  0 ≤ Q x :=
(eq_or_ne x 0).elim (λ h, h.symm ▸ (map_zero Q).symm.le) (λ h, (hQ _ h).le)

lemma pos_def.anisotropic {Q : quadratic_form R₂ M} (hQ : Q.pos_def) : Q.anisotropic :=
λ x hQx, classical.by_contradiction $ λ hx, lt_irrefl (0 : R₂) $ begin
  have := hQ _ hx,
  rw hQx at this,
  exact this,
end

lemma pos_def_of_nonneg {Q : quadratic_form R₂ M} (h : ∀ x, 0 ≤ Q x) (h0 : Q.anisotropic) :
  pos_def Q :=
λ x hx, lt_of_le_of_ne (h x) (ne.symm $ λ hQx, hx $ h0 _ hQx)

lemma pos_def_iff_nonneg {Q : quadratic_form R₂ M} :
  pos_def Q ↔ (∀ x, 0 ≤ Q x) ∧ Q.anisotropic  :=
⟨λ h, ⟨h.nonneg, h.anisotropic⟩, λ ⟨n, a⟩, pos_def_of_nonneg n a⟩

lemma pos_def.add (Q Q' : quadratic_form R₂ M) (hQ : pos_def Q) (hQ' : pos_def Q') :
  pos_def (Q + Q') :=
λ x hx, add_pos (hQ x hx) (hQ' x hx)

lemma lin_mul_lin_self_pos_def {R} [linear_ordered_comm_ring R] [module R M]
  (f : M →ₗ[R] R) (hf : linear_map.ker f = ⊥) :
  pos_def (lin_mul_lin f f) :=
λ x hx, mul_self_pos.2 (λ h, hx $ linear_map.ker_eq_bot'.mp hf _ h)

end pos_def
end quadratic_form

section
/-!
### Quadratic forms and matrices

Connect quadratic forms and matrices, in order to explicitly compute with them.
The convention is twos out, so there might be a factor 2⁻¹ in the entries of the
matrix.
The determinant of the matrix is the discriminant of the quadratic form.
-/

variables {n : Type w} [fintype n] [decidable_eq n]
variables [comm_ring R₁] [add_comm_monoid M] [module R₁ M]

/-- `M.to_quadratic_form` is the map `λ x, col x ⬝ M ⬝ row x` as a quadratic form. -/
def matrix.to_quadratic_form' (M : matrix n n R₁) :
  quadratic_form R₁ (n → R₁) :=
M.to_bilin'.to_quadratic_form

variables [invertible (2 : R₁)]

/-- A matrix representation of the quadratic form. -/
def quadratic_form.to_matrix' (Q : quadratic_form R₁ (n → R₁)) : matrix n n R₁ :=
Q.associated.to_matrix'

open quadratic_form
lemma quadratic_form.to_matrix'_smul (a : R₁) (Q : quadratic_form R₁ (n → R₁)) :
  (a • Q).to_matrix' = a • Q.to_matrix' :=
by simp only [to_matrix', linear_equiv.map_smul, linear_map.map_smul]

lemma quadratic_form.is_symm_to_matrix' (Q : quadratic_form R₁ (n → R₁)) :
  Q.to_matrix'.is_symm :=
begin
  ext i j,
  rw [to_matrix', bilin_form.to_matrix'_apply, bilin_form.to_matrix'_apply, associated_is_symm]
end

end

namespace quadratic_form

variables {n : Type w} [fintype n]
variables [comm_ring R₁] [decidable_eq n] [invertible (2 : R₁)]
variables {m : Type w} [decidable_eq m] [fintype m]
open_locale matrix

@[simp]
lemma to_matrix'_comp (Q : quadratic_form R₁ (m → R₁)) (f : (n → R₁) →ₗ[R₁] (m → R₁)) :
  (Q.comp f).to_matrix' = f.to_matrix'ᵀ ⬝ Q.to_matrix' ⬝ f.to_matrix' :=
by { ext, simp only [quadratic_form.associated_comp, bilin_form.to_matrix'_comp, to_matrix'] }

section discriminant
variables {Q : quadratic_form R₁ (n → R₁)}

/-- The discriminant of a quadratic form generalizes the discriminant of a quadratic polynomial. -/
def discr (Q : quadratic_form R₁ (n → R₁)) : R₁ := Q.to_matrix'.det

lemma discr_smul (a : R₁) : (a • Q).discr = a ^ fintype.card n * Q.discr :=
by simp only [discr, to_matrix'_smul, matrix.det_smul]

lemma discr_comp (f : (n → R₁) →ₗ[R₁] (n → R₁)) :
  (Q.comp f).discr = f.to_matrix'.det * f.to_matrix'.det * Q.discr :=
by simp only [matrix.det_transpose, mul_left_comm, quadratic_form.to_matrix'_comp, mul_comm,
  matrix.det_mul, discr]

end discriminant

end quadratic_form

namespace quadratic_form

end quadratic_form

namespace bilin_form
section semiring
variables [semiring R] [add_comm_monoid M] [module R M]

/-- A bilinear form is nondegenerate if the quadratic form it is associated with is anisotropic. -/
lemma nondegenerate_of_anisotropic
  {B : bilin_form R M} (hB : B.to_quadratic_form.anisotropic) : B.nondegenerate :=
λ x hx, hB _ (hx x)
end semiring

variables [ring R] [add_comm_group M] [module R M]

/-- There exists a non-null vector with respect to any symmetric, nonzero bilinear form `B`
on a module `M` over a ring `R` with invertible `2`, i.e. there exists some
`x : M` such that `B x x ≠ 0`. -/

open finite_dimensional

variables {V : Type u} {K : Type v} [field K] [add_comm_group V] [module K V]
variable [finite_dimensional K V]

/-- Given a symmetric bilinear form `B` on some vector space `V` over a field `K`
in which `2` is invertible, there exists an orthogonal basis with respect to `B`. -/
lemma exists_orthogonal_basis [hK : invertible (2 : K)]
  {B : bilin_form K V} (hB₂ : B.is_symm) :
  ∃ (v : basis (fin (finrank K V)) K V), B.is_Ortho v :=
begin
  unfreezingI { induction hd : finrank K V with d ih generalizing V },
  { exact ⟨basis_of_finrank_zero hd, λ _ _ _, zero_left _⟩ },
  haveI := finrank_pos_iff.1 (hd.symm ▸ nat.succ_pos d : 0 < finrank K V),
  -- either the bilinear form is trivial or we can pick a non-null `x`
  obtain rfl | hB₁ := eq_or_ne B 0,
  { let b := finite_dimensional.fin_basis K V,
    rw hd at b,
    refine ⟨b, λ i j hij, rfl⟩, },
  obtain ⟨x, hx⟩ := exists_bilin_form_self_ne_zero hB₁ hB₂,
  rw [← submodule.finrank_add_eq_of_is_compl (is_compl_span_singleton_orthogonal hx).symm,
      finrank_span_singleton (ne_zero_of_not_is_ortho_self x hx)] at hd,
  let B' := B.restrict (B.orthogonal $ K ∙ x),
  obtain ⟨v', hv₁⟩ := ih (B.restrict_symm hB₂ _ : B'.is_symm) (nat.succ.inj hd),
  -- concatenate `x` with the basis obtained by induction
  let b := basis.mk_fin_cons x v'
    (begin
      rintros c y hy hc,
      rw add_eq_zero_iff_neg_eq at hc,
      rw [← hc, submodule.neg_mem_iff] at hy,
      have := (is_compl_span_singleton_orthogonal hx).disjoint,
      rw submodule.disjoint_def at this,
      have := this (c • x) (submodule.smul_mem _ _ $ submodule.mem_span_singleton_self _) hy,
      exact (smul_eq_zero.1 this).resolve_right (λ h, hx $ h.symm ▸ zero_left _),
    end)
    (begin
      intro y,
      refine ⟨-B x y/B x x, λ z hz, _⟩,
      obtain ⟨c, rfl⟩ := submodule.mem_span_singleton.1 hz,
      rw [is_ortho, smul_left, add_right, smul_right, div_mul_cancel _ hx, add_neg_self, mul_zero],
    end),
  refine ⟨b, _⟩,
  { rw basis.coe_mk_fin_cons,
    intros j i,
    refine fin.cases _ (λ i, _) i; refine fin.cases _ (λ j, _) j; intro hij;
      simp only [function.on_fun, fin.cons_zero, fin.cons_succ, function.comp_apply],
    { exact (hij rfl).elim },
    { rw [is_ortho, hB₂],
      exact (v' j).prop _ (submodule.mem_span_singleton_self x) },
    { exact (v' i).prop _ (submodule.mem_span_singleton_self x) },
    { exact hv₁ (ne_of_apply_ne _ hij), }, }
end

end bilin_form

namespace quadratic_form

open finset bilin_form

variables {M₁ : Type*} [semiring R] [comm_semiring R₁] [add_comm_monoid M] [add_comm_monoid M₁]
variables [module R M] [module R M₁]
variables {ι : Type*} [fintype ι] {v : basis ι R M}

/-- Given a quadratic form `Q` and a basis, `basis_repr` is the basis representation of `Q`. -/
noncomputable def basis_repr (Q : quadratic_form R M) (v : basis ι R M) :
  quadratic_form R (ι → R) :=
Q.comp v.equiv_fun.symm

@[simp]
lemma basis_repr_apply (Q : quadratic_form R M) (w : ι → R) :
  Q.basis_repr v w = Q (∑ i : ι, w i • v i) :=
by { rw ← v.equiv_fun_symm_apply, refl }

section

variables (R₁)

/-- The weighted sum of squares with respect to some weight as a quadratic form.

The weights are applied using `•`; typically this definition is used either with `S = R₁` or
`[algebra S R₁]`, although this is stated more generally. -/
def weighted_sum_squares [monoid S] [distrib_mul_action S R₁]
  [smul_comm_class S R₁ R₁]
  (w : ι → S) : quadratic_form R₁ (ι → R₁) :=
∑ i : ι, w i • proj i i

end

@[simp]
lemma weighted_sum_squares_apply [monoid S] [distrib_mul_action S R₁] [smul_comm_class S R₁ R₁]
  (w : ι → S) (v : ι → R₁) :
  weighted_sum_squares R₁ w v = ∑ i : ι, w i • (v i * v i) :=
quadratic_form.sum_apply _ _ _

/-- On an orthogonal basis, the basis representation of `Q` is just a sum of squares. -/
lemma basis_repr_eq_of_is_Ortho
  {R₁ M} [comm_ring R₁] [add_comm_group M] [module R₁ M] [invertible (2 : R₁)]
  (Q : quadratic_form R₁ M) (v : basis ι R₁ M) (hv₂ : (associated Q).is_Ortho v) :
  Q.basis_repr v = weighted_sum_squares _ (λ i, Q (v i)) :=
begin
  ext w,
  rw [basis_repr_apply, ←@associated_eq_self_apply R₁, sum_left, weighted_sum_squares_apply],
  refine sum_congr rfl (λ j hj, _),
  rw [←@associated_eq_self_apply R₁, sum_right, sum_eq_single_of_mem j hj],
  { rw [smul_left, smul_right, smul_eq_mul], ring },
  { intros i _ hij,
    rw [smul_left, smul_right,
        show associated_hom R₁ Q (v j) (v i) = 0, from hv₂ hij.symm,
        mul_zero, mul_zero] },
end

end quadratic_form


/-
Copyright (c) 2022 Sébastien Gouëzel. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Sébastien Gouëzel

! This file was ported from Lean 3 source module topology.algebra.order.left_right_lim
! leanprover-community/mathlib commit 0a0ec35061ed9960bf0e7ffb0335f44447b58977
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
import Mathlib.Tactic.WLOG
import Mathlib.Topology.Order.Basic
import Mathlib.Topology.Algebra.Order.LeftRight

/-!
# Left and right limits

We define the (strict) left and right limits of a function.

* `leftLim f x` is the strict left limit of `f` at `x` (using `f x` as a garbage value if `x`
  is isolated to its left).
* `rightLim f x` is the strict right limit of `f` at `x` (using `f x` as a garbage value if `x`
  is isolated to its right).

We develop a comprehensive API for monotone functions. Notably,

* `Monotone.continuousAt_iff_leftLim_eq_rightLim` states that a monotone function is continuous
  at a point if and only if its left and right limits coincide.
* `Monotone.countable_not_continuousAt` asserts that a monotone function taking values in a
  second-countable space has at most countably many discontinuity points.

We also port the API to antitone functions.

## TODO

Prove corresponding stronger results for `StrictMono` and `StrictAnti` functions.
-/

open Set Filter

open Topology

section

variable {α β : Type _} [LinearOrder α] [TopologicalSpace β]

/-- Let `f : α → β` be a function from a linear order `α` to a topological space `β`, and
let `a : α`. The limit strictly to the left of `f` at `a`, denoted with `leftLim f a`, is defined
by using the order topology on `α`. If `a` is isolated to its left or the function has no left
limit, we use `f a` instead to guarantee a good behavior in most cases. -/
noncomputable def Function.leftLim (f : α → β) (a : α) : β := by
  classical
  haveI : Nonempty β := ⟨f a⟩
  letI : TopologicalSpace α := Preorder.topology α
  exact if 𝓝[<] a = ⊥ ∨ ¬∃ y, Tendsto f (𝓝[<] a) (𝓝 y) then f a else limUnder (𝓝[<] a) f
#align function.left_lim Function.leftLim

/-- Let `f : α → β` be a function from a linear order `α` to a topological space `β`, and
let `a : α`. The limit strictly to the right of `f` at `a`, denoted with `rightLim f a`, is defined
by using the order topology on `α`. If `a` is isolated to its right or the function has no right
limit, , we use `f a` instead to guarantee a good behavior in most cases. -/
noncomputable def Function.rightLim (f : α → β) (a : α) : β :=
  @Function.leftLim αᵒᵈ β _ _ f a
#align function.right_lim Function.rightLim

open Function

theorem leftLim_eq_of_tendsto [hα : TopologicalSpace α] [h'α : OrderTopology α] [T2Space β]
    {f : α → β} {a : α} {y : β} (h : 𝓝[<] a ≠ ⊥) (h' : Tendsto f (𝓝[<] a) (𝓝 y)) :
    leftLim f a = y := by
  have h'' : ∃ y, Tendsto f (𝓝[<] a) (𝓝 y) := ⟨y, h'⟩
  rw [h'α.topology_eq_generate_intervals] at h h' h''
  simp only [leftLim, h, h'', not_true, or_self_iff, if_false]
  haveI := neBot_iff.2 h
  exact lim_eq h'
#align left_lim_eq_of_tendsto leftLim_eq_of_tendsto

theorem leftLim_eq_of_eq_bot [hα : TopologicalSpace α] [h'α : OrderTopology α] (f : α → β) {a : α}
    (h : 𝓝[<] a = ⊥) : leftLim f a = f a := by
  rw [h'α.topology_eq_generate_intervals] at h
  simp [leftLim, ite_eq_left_iff, h]
#align left_lim_eq_of_eq_bot leftLim_eq_of_eq_bot

end

open Function

namespace Monotone

variable {α β : Type _} [LinearOrder α] [ConditionallyCompleteLinearOrder β] [TopologicalSpace β]
  [OrderTopology β] {f : α → β} (hf : Monotone f) {x y : α}

theorem leftLim_eq_supₛ [TopologicalSpace α] [OrderTopology α] (h : 𝓝[<] x ≠ ⊥) :
    leftLim f x = supₛ (f '' Iio x) :=
  leftLim_eq_of_tendsto h (hf.tendsto_nhdsWithin_Iio x)
#align monotone.left_lim_eq_Sup Monotone.leftLim_eq_supₛ

theorem leftLim_le (h : x ≤ y) : leftLim f x ≤ f y := by
  letI : TopologicalSpace α := Preorder.topology α
  haveI : OrderTopology α := ⟨rfl⟩
  rcases eq_or_ne (𝓝[<] x) ⊥ with (h' | h')
  · simpa [leftLim, h'] using hf h
  haveI A : NeBot (𝓝[<] x) := neBot_iff.2 h'
  rw [leftLim_eq_supₛ hf h']
  refine' csupₛ_le _ _
  · simp only [nonempty_image_iff]
    exact (forall_mem_nonempty_iff_neBot.2 A) _ self_mem_nhdsWithin
  · simp only [mem_image, mem_Iio, forall_exists_index, and_imp, forall_apply_eq_imp_iff₂]
    intro z hz
    exact hf (hz.le.trans h)
#align monotone.left_lim_le Monotone.leftLim_le

theorem le_leftLim (h : x < y) : f x ≤ leftLim f y := by
  letI : TopologicalSpace α := Preorder.topology α
  haveI : OrderTopology α := ⟨rfl⟩
  rcases eq_or_ne (𝓝[<] y) ⊥ with (h' | h')
  · rw [leftLim_eq_of_eq_bot _ h']
    exact hf h.le
  rw [leftLim_eq_supₛ hf h']
  refine' le_csupₛ ⟨f y, _⟩ (mem_image_of_mem _ h)
  simp only [upperBounds, mem_image, mem_Iio, forall_exists_index, and_imp,
    forall_apply_eq_imp_iff₂, mem_setOf_eq]
  intro z hz
  exact hf hz.le
#align monotone.le_left_lim Monotone.le_leftLim

@[mono]
protected theorem leftLim : Monotone (leftLim f) := by
  intro x y h
  rcases eq_or_lt_of_le h with (rfl | hxy)
  · exact le_rfl
  · exact (hf.leftLim_le le_rfl).trans (hf.le_leftLim hxy)
#align monotone.left_lim Monotone.leftLim

theorem le_rightLim (h : x ≤ y) : f x ≤ rightLim f y :=
  hf.dual.leftLim_le h
#align monotone.le_right_lim Monotone.le_rightLim

theorem rightLim_le (h : x < y) : rightLim f x ≤ f y :=
  hf.dual.le_leftLim h
#align monotone.right_lim_le Monotone.rightLim_le

@[mono]
protected theorem rightLim : Monotone (rightLim f) := fun _ _ h => hf.dual.leftLim h
#align monotone.right_lim Monotone.rightLim

theorem leftLim_le_rightLim (h : x ≤ y) : leftLim f x ≤ rightLim f y :=
  (hf.leftLim_le le_rfl).trans (hf.le_rightLim h)
#align monotone.left_lim_le_right_lim Monotone.leftLim_le_rightLim

theorem rightLim_le_leftLim (h : x < y) : rightLim f x ≤ leftLim f y := by
  letI : TopologicalSpace α := Preorder.topology α
  haveI : OrderTopology α := ⟨rfl⟩
  rcases eq_or_ne (𝓝[<] y) ⊥ with (h' | h')
  · simp [leftLim, h']
    exact rightLim_le hf h
  obtain ⟨a, ⟨xa, ay⟩⟩ : (Ioo x y).Nonempty :=
    forall_mem_nonempty_iff_neBot.2 (neBot_iff.2 h') (Ioo x y)
      (Ioo_mem_nhdsWithin_Iio ⟨h, le_refl _⟩)
  calc
    rightLim f x ≤ f a := hf.rightLim_le xa
    _ ≤ leftLim f y := hf.le_leftLim ay

#align monotone.right_lim_le_left_lim Monotone.rightLim_le_leftLim

variable [TopologicalSpace α] [OrderTopology α]

theorem tendsto_leftLim (x : α) : Tendsto f (𝓝[<] x) (𝓝 (leftLim f x)) := by
  rcases eq_or_ne (𝓝[<] x) ⊥ with (h' | h')
  · simp [h']
  rw [leftLim_eq_supₛ hf h']
  exact hf.tendsto_nhdsWithin_Iio x
#align monotone.tendsto_left_lim Monotone.tendsto_leftLim

theorem tendsto_leftLim_within (x : α) : Tendsto f (𝓝[<] x) (𝓝[≤] leftLim f x) := by
  apply tendsto_nhdsWithin_of_tendsto_nhds_of_eventually_within f (hf.tendsto_leftLim x)
  filter_upwards [@self_mem_nhdsWithin _ _ x (Iio x)] with y hy using hf.le_leftLim hy
#align monotone.tendsto_left_lim_within Monotone.tendsto_leftLim_within

theorem tendsto_rightLim (x : α) : Tendsto f (𝓝[>] x) (𝓝 (rightLim f x)) :=
  hf.dual.tendsto_leftLim x
#align monotone.tendsto_right_lim Monotone.tendsto_rightLim

theorem tendsto_rightLim_within (x : α) : Tendsto f (𝓝[>] x) (𝓝[≥] rightLim f x) :=
  hf.dual.tendsto_leftLim_within x
#align monotone.tendsto_right_lim_within Monotone.tendsto_rightLim_within

/-- A monotone function is continuous to the left at a point if and only if its left limit
coincides with the value of the function. -/
theorem continuousWithinAt_Iio_iff_leftLim_eq :
    ContinuousWithinAt f (Iio x) x ↔ leftLim f x = f x := by
  rcases eq_or_ne (𝓝[<] x) ⊥ with (h' | h')
  · simp [leftLim_eq_of_eq_bot f h', ContinuousWithinAt, h']
  haveI : (𝓝[Iio x] x).NeBot := neBot_iff.2 h'
  refine' ⟨fun h => tendsto_nhds_unique (hf.tendsto_leftLim x) h.tendsto, fun h => _⟩
  have := hf.tendsto_leftLim x
  rwa [h] at this
#align monotone.continuous_within_at_Iio_iff_left_lim_eq Monotone.continuousWithinAt_Iio_iff_leftLim_eq

/-- A monotone function is continuous to the right at a point if and only if its right limit
coincides with the value of the function. -/
theorem continuousWithinAt_Ioi_iff_rightLim_eq :
    ContinuousWithinAt f (Ioi x) x ↔ rightLim f x = f x :=
  hf.dual.continuousWithinAt_Iio_iff_leftLim_eq
#align monotone.continuous_within_at_Ioi_iff_right_lim_eq Monotone.continuousWithinAt_Ioi_iff_rightLim_eq

/-- A monotone function is continuous at a point if and only if its left and right limits
coincide. -/
theorem continuousAt_iff_leftLim_eq_rightLim : ContinuousAt f x ↔ leftLim f x = rightLim f x := by
  refine' ⟨fun h => _, fun h => _⟩
  · have A : leftLim f x = f x :=
      hf.continuousWithinAt_Iio_iff_leftLim_eq.1 h.continuousWithinAt
    have B : rightLim f x = f x :=
      hf.continuousWithinAt_Ioi_iff_rightLim_eq.1 h.continuousWithinAt
    exact A.trans B.symm
  · have h' : leftLim f x = f x :=
      by
      apply le_antisymm (leftLim_le hf (le_refl _))
      rw [h]
      exact le_rightLim hf (le_refl _)
    refine' continuousAt_iff_continuous_left'_right'.2 ⟨_, _⟩
    · exact hf.continuousWithinAt_Iio_iff_leftLim_eq.2 h'
    · rw [h] at h'
      exact hf.continuousWithinAt_Ioi_iff_rightLim_eq.2 h'
#align monotone.continuous_at_iff_left_lim_eq_right_lim Monotone.continuousAt_iff_leftLim_eq_rightLim

/-- In a second countable space, the set of points where a monotone function is not right-continuous
is at most countable. Superseded by `countable_not_continuousAt` which gives the two-sided
version. -/
theorem countable_not_continuousWithinAt_Ioi [TopologicalSpace.SecondCountableTopology β] :
    Set.Countable { x | ¬ContinuousWithinAt f (Ioi x) x } := by
  /- If `f` is not continuous on the right at `x`, there is an interval `(f x, z x)` which is not
    reached by `f`. This gives a family of disjoint open intervals in `β`. Such a family can only
    be countable as `β` is second-countable. -/
  nontriviality α
  let s := { x | ¬ContinuousWithinAt f (Ioi x) x }
  have : ∀ x, x ∈ s → ∃ z, f x < z ∧ ∀ y, x < y → z ≤ f y :=
    by
    rintro x (hx : ¬ContinuousWithinAt f (Ioi x) x)
    contrapose! hx
    refine' tendsto_order.2 ⟨fun m hm => _, fun u hu => _⟩
    · filter_upwards [@self_mem_nhdsWithin _ _ x (Ioi x)] with y hy using hm.trans_le
        (hf (le_of_lt hy))
    rcases hx u hu with ⟨v, xv, fvu⟩
    have : Ioo x v ∈ 𝓝[>] x := Ioo_mem_nhdsWithin_Ioi ⟨le_refl _, xv⟩
    filter_upwards [this]with y hy
    apply (hf hy.2.le).trans_lt fvu
  -- choose `z x` such that `f` does not take the values in `(f x, z x)`.
  choose! z hz using this
  have I : InjOn f s := by
    apply StrictMonoOn.injOn
    intro x hx y _ hxy
    calc
      f x < z x := (hz x hx).1
      _ ≤ f y := (hz x hx).2 y hxy

  -- show that `f s` is countable by arguing that a disjoint family of disjoint open intervals
  -- (the intervals `(f x, z x)`) is at most countable.
  have fs_count : (f '' s).Countable :=
    by
    have A : (f '' s).PairwiseDisjoint fun x => Ioo x (z (invFunOn f s x)) := by
      rintro _ ⟨u, us, rfl⟩ _ ⟨v, vs, rfl⟩ huv
      wlog hle : u ≤ v generalizing u v
      · exact (this v vs u us huv.symm (le_of_not_le hle)).symm
      have hlt : u < v := hle.lt_of_ne (ne_of_apply_ne _ huv)
      apply disjoint_iff_forall_ne.2
      rintro a ha b hb rfl
      simp only [I.leftInvOn_invFunOn us, I.leftInvOn_invFunOn vs] at ha hb
      exact lt_irrefl _ ((ha.2.trans_le ((hz u us).2 v hlt)).trans hb.1)
    apply Set.PairwiseDisjoint.countable_of_Ioo A
    rintro _ ⟨y, ys, rfl⟩
    simpa only [I.leftInvOn_invFunOn ys] using (hz y ys).1
  exact MapsTo.countable_of_injOn (mapsTo_image f s) I fs_count
#align monotone.countable_not_continuous_within_at_Ioi Monotone.countable_not_continuousWithinAt_Ioi

/-- In a second countable space, the set of points where a monotone function is not left-continuous
is at most countable. Superseded by `countable_not_continuousAt` which gives the two-sided
version. -/
theorem countable_not_continuousWithinAt_Iio [TopologicalSpace.SecondCountableTopology β] :
    Set.Countable { x | ¬ContinuousWithinAt f (Iio x) x } :=
  hf.dual.countable_not_continuousWithinAt_Ioi
#align monotone.countable_not_continuous_within_at_Iio Monotone.countable_not_continuousWithinAt_Iio

/-- In a second countable space, the set of points where a monotone function is not continuous
is at most countable. -/
theorem countable_not_continuousAt [TopologicalSpace.SecondCountableTopology β] :
    Set.Countable { x | ¬ContinuousAt f x } := by
  apply
    (hf.countable_not_continuousWithinAt_Ioi.union hf.countable_not_continuousWithinAt_Iio).mono
      _
  refine' compl_subset_compl.1 _
  simp only [compl_union]
  rintro x ⟨hx, h'x⟩
  simp only [mem_setOf_eq, Classical.not_not, mem_compl_iff] at hx h'x⊢
  exact continuousAt_iff_continuous_left'_right'.2 ⟨h'x, hx⟩
#align monotone.countable_not_continuous_at Monotone.countable_not_continuousAt

end Monotone

namespace Antitone

variable {α β : Type _} [LinearOrder α] [ConditionallyCompleteLinearOrder β] [TopologicalSpace β]
  [OrderTopology β] {f : α → β} (hf : Antitone f) {x y : α}

theorem le_leftLim (h : x ≤ y) : f y ≤ leftLim f x :=
  hf.dual_right.leftLim_le h
#align antitone.le_left_lim Antitone.le_leftLim

theorem leftLim_le (h : x < y) : leftLim f y ≤ f x :=
  hf.dual_right.le_leftLim h
#align antitone.left_lim_le Antitone.leftLim_le

@[mono]
protected theorem leftLim : Antitone (leftLim f) :=
  hf.dual_right.leftLim
#align antitone.left_lim Antitone.leftLim

theorem rightLim_le (h : x ≤ y) : rightLim f y ≤ f x :=
  hf.dual_right.le_rightLim h
#align antitone.right_lim_le Antitone.rightLim_le

theorem le_rightLim (h : x < y) : f y ≤ rightLim f x :=
  hf.dual_right.rightLim_le h
#align antitone.le_right_lim Antitone.le_rightLim

@[mono]
protected theorem rightLim : Antitone (rightLim f) :=
  hf.dual_right.rightLim
#align antitone.right_lim Antitone.rightLim

theorem rightLim_le_leftLim (h : x ≤ y) : rightLim f y ≤ leftLim f x :=
  hf.dual_right.leftLim_le_rightLim h
#align antitone.right_lim_le_left_lim Antitone.rightLim_le_leftLim

theorem leftLim_le_rightLim (h : x < y) : leftLim f y ≤ rightLim f x :=
  hf.dual_right.rightLim_le_leftLim h
#align antitone.left_lim_le_right_lim Antitone.leftLim_le_rightLim

variable [TopologicalSpace α] [OrderTopology α]

theorem tendsto_leftLim (x : α) : Tendsto f (𝓝[<] x) (𝓝 (leftLim f x)) :=
  hf.dual_right.tendsto_leftLim x
#align antitone.tendsto_left_lim Antitone.tendsto_leftLim

theorem tendsto_leftLim_within (x : α) : Tendsto f (𝓝[<] x) (𝓝[≥] leftLim f x) :=
  hf.dual_right.tendsto_leftLim_within x
#align antitone.tendsto_left_lim_within Antitone.tendsto_leftLim_within

theorem tendsto_rightLim (x : α) : Tendsto f (𝓝[>] x) (𝓝 (rightLim f x)) :=
  hf.dual_right.tendsto_rightLim x
#align antitone.tendsto_right_lim Antitone.tendsto_rightLim

theorem tendsto_rightLim_within (x : α) : Tendsto f (𝓝[>] x) (𝓝[≤] rightLim f x) :=
  hf.dual_right.tendsto_rightLim_within x
#align antitone.tendsto_right_lim_within Antitone.tendsto_rightLim_within

/-- An antitone function is continuous to the left at a point if and only if its left limit
coincides with the value of the function. -/
theorem continuousWithinAt_Iio_iff_leftLim_eq :
    ContinuousWithinAt f (Iio x) x ↔ leftLim f x = f x :=
  hf.dual_right.continuousWithinAt_Iio_iff_leftLim_eq
#align antitone.continuous_within_at_Iio_iff_left_lim_eq Antitone.continuousWithinAt_Iio_iff_leftLim_eq

/-- An antitone function is continuous to the right at a point if and only if its right limit
coincides with the value of the function. -/
theorem continuousWithinAt_Ioi_iff_rightLim_eq :
    ContinuousWithinAt f (Ioi x) x ↔ rightLim f x = f x :=
  hf.dual_right.continuousWithinAt_Ioi_iff_rightLim_eq
#align antitone.continuous_within_at_Ioi_iff_right_lim_eq Antitone.continuousWithinAt_Ioi_iff_rightLim_eq

/-- An antitone function is continuous at a point if and only if its left and right limits
coincide. -/
theorem continuousAt_iff_leftLim_eq_rightLim : ContinuousAt f x ↔ leftLim f x = rightLim f x :=
  hf.dual_right.continuousAt_iff_leftLim_eq_rightLim
#align antitone.continuous_at_iff_left_lim_eq_right_lim Antitone.continuousAt_iff_leftLim_eq_rightLim

/-- In a second countable space, the set of points where an antitone function is not continuous
is at most countable. -/
theorem countable_not_continuousAt [TopologicalSpace.SecondCountableTopology β] :
    Set.Countable { x | ¬ContinuousAt f x } :=
  hf.dual_right.countable_not_continuousAt
#align antitone.countable_not_continuous_at Antitone.countable_not_continuousAt

end Antitone


import tactic
import data.nat.basic
import data.nat.prime

noncomputable theory
open_locale classical

open nat

theorem sqrt2_irrational' :
  ¬ ∃ (m n : ℕ),
  2 * m^2 = n^2 ∧
  m.coprime n
:=
begin
  sorry,
end

lemma ge_zero_sq_ge_zero {n : ℕ} (hne : 0 < n) : (0 < n^2)
:=
begin
  sorry,
end

lemma cancellation_lemma {k m n : ℕ}
(hk_pos : 0 < k^2)
(hmn : 2 * (m * k) ^ 2 = (n * k) ^ 2)
: 2 * m ^ 2 = n ^ 2
:=
begin
  sorry,
end

-- Assume that everything above this line is true.

--BEGIN--
/-
gcd_pos_of_pos_left : ∀ {m : ℕ} (n : ℕ), 0 < m → 0 < m.gcd n
gcd_pos_of_pos_right : ∀ (m : ℕ) {n : ℕ}, 0 < n → 0 < m.gcd n
exists_coprime : ∀ {m n : ℕ}, 0 < m.gcd n → (∃ (m' n' : ℕ), m'.coprime n' ∧ m = m' * m.gcd n ∧ n = n' * m.gcd n)
nat.pos_of_ne_zero : ∀ {n : ℕ}, n ≠ 0 → 0 < n

-/
theorem wlog_coprime :
  (∃ (m n : ℕ),
  2 * m^2 = n^2 ∧
  m ≠ 0 )
  → (∃ (m' n' : ℕ),
    2 * m'^2 = n'^2 ∧
    m'.coprime n' )
:=
begin
  rintro ⟨m, n, hmn, hm0⟩,
  set k := m.gcd n with hk,
  -- this abbreviation reduces clutter
  -- ``set`` is similar to ``have``
  -- you can replace all the ``m.gcd n`` with ``k`` using ``rw ←hk,`` if needed
  sorry,
end

theorem sqrt2_irrational'' :
  ¬ ∃ (m n : ℕ),
  2 * m^2 = n^2 ∧
  m ≠ 0
:=
begin
  sorry,
end

--END--


/-
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.big_operators.basic
import algebra.star.basic

/-! # Big-operators lemmas about `star` algebraic operations

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

These results are kept separate from `algebra.star.basic` to avoid it needing to import `finset`.
-/

variables {R : Type*}

open_locale big_operators

@[simp] lemma star_prod [comm_monoid R] [star_semigroup R] {α : Type*}
  (s : finset α) (f : α → R):
  star (∏ x in s, f x) = ∏ x in s, star (f x) :=
map_prod (star_mul_aut : R ≃* R) _ _

@[simp] lemma star_sum [add_comm_monoid R] [star_add_monoid R] {α : Type*}
  (s : finset α) (f : α → R):
  star (∑ x in s, f x) = ∑ x in s, star (f x) :=
(star_add_equiv : R ≃+ R).map_sum _ _


import
  data.real.basic
  extra.periodic.big_operators
  data.nat.periodic
  algebra.big_operators.order
  tactic.fin_cases

/-! # IMO 2018 A2 (P2) -/

namespace IMOSL
namespace IMO2018A2

open finset function extra
open_locale classical

def good (a : ℕ → ℝ) := ∀ i : ℕ, a i * a (i + 1) + 1 = a (i + 2)

section extra

lemma sum_nonneg_eq_zero {M : Type*} [linear_ordered_add_comm_group M] {a : ℕ → M}
  (h : ∀ i : ℕ, 0 ≤ a i) {n : ℕ} (h0 : (range n).sum a = 0) {j : ℕ} (h1 : j < n) : a j = 0 :=
begin
  revert h1 j; induction n with n n_ih,
  intros j h1; exfalso; exact not_le_of_lt h1 j.zero_le,
  rw [sum_range_succ, add_eq_zero_iff_eq_neg] at h0,
  have h1 : 0 ≤ (range n).sum (λ (x : ℕ), a x) := sum_nonneg (λ i _, h i),
  rw [h0, neg_nonneg] at h1,
  replace h1 := le_antisymm h1 (h n),
  rw [h1, neg_zero] at h0,
  intros j h2,
  rw [nat.lt_succ_iff, le_iff_eq_or_lt] at h2,
  rcases h2 with rfl | h2,
  exacts [h1, n_ih h0 h2]
end

lemma periodic_shift {α M : Type*} [add_comm_monoid M] {f : M → α} {c : M}
    (h : periodic f c) (d : M) : periodic (λ m, f (m + d)) c :=
  λ m, by simp only []; rw [add_right_comm, h]

lemma sq_add_one_ne_self (x : ℝ) : x * x + 1 ≠ x :=
begin
  intros h,
  apply not_lt.mpr (sq_nonneg (2 * x - 1)),
  rw ← eq_sub_iff_add_eq at h,
  rw [sub_sq, mul_pow, sq, sq, h, mul_one, one_pow, mul_sub_one, mul_assoc, sub_sub_cancel_left],
  norm_num
end

end extra

/-- A good periodic function must be 3-periodic. -/
private lemma periodic_to_periodic3 {a : ℕ → ℝ} (h : good a)
  {n : ℕ} (h0 : 0 < n) (h1 : periodic a n) : periodic a 3 :=
begin
  have h2 : (range n).sum (λ i, 2 * a (i + 2) ^ 2 + (a (i + 3) - a i) ^ 2) =
    (range n).sum (λ i, a (i + 3) ^ 2 + a i ^ 2 + 2 * (a (i + 2) - a i)) :=
  begin
    apply congr_arg; funext i,
    rw [add_comm, sub_sq', ← add_sub_right_comm, add_sub_assoc, add_right_inj,
        mul_assoc, ← mul_sub, mul_right_inj' (two_ne_zero : (2 : ℝ) ≠ 0),
        sub_eq_sub_iff_sub_eq_sub, ← sub_one_mul, mul_comm, sq, ← sub_one_mul],
    dsimp only [good] at h,
    conv at h { find (_ = _) { rw ← eq_sub_iff_add_eq } },
    rw [← h, mul_assoc, bit0, ← add_assoc, h]
  end,
  repeat { rw sum_add_distrib at h2 },
  rw [← mul_sum, ← mul_sum] at h2,
  repeat { rw periodic_sum_const (periodic.comp h1 (λ x, x ^ 2)) at h2 },
  dsimp only [comp] at h2,
  rw [← two_mul, add_right_inj, sum_sub_distrib, periodic_sum_const h1, sub_self, mul_zero] at h2,
  intros i,
  replace h2 := sum_nonneg_eq_zero (λ i, sq_nonneg (a (i + 3) - a i)) h2 (i.mod_lt h0),
  rwa [pow_eq_zero_iff, periodic.map_mod_nat h1, ← periodic.map_mod_nat h1,
       nat.mod_add_mod, periodic.map_mod_nat h1, sub_eq_zero] at h2,
  exact two_pos
end

/-- One example of a good function. -/
def good_ex (i : ℕ) : ℝ := ite (i % 3 = 0) 2 (-1)

private lemma good_ex_good : good good_ex :=
begin
  intros i; dsimp only [good_ex],
  rw [← nat.mod_add_mod, ← nat.mod_add_mod i],
  let p : fin 3 := ⟨i % 3, i.mod_lt three_pos⟩,
  fin_cases p using hp,
  all_goals {
    rw fin.ext_iff at hp,
    simp [p] at hp,
    rw hp; norm_num }
end

private lemma good_shift {a : ℕ → ℝ} (h : good a) (k : ℕ) : good (λ i, a (i + k)) :=
  λ i, by simp only []; rw [add_right_comm, h, add_right_comm]

/-- Solution of the underlying equation for a good 3-periodic function. -/
private lemma good_solution {a b c : ℝ}
    (h : a * b + 1 = c) (h0 : b * c + 1 = a) (h1 : c * a + 1 = b) :
  (a, b, c) = (2, -1, -1) ∨ (a, b, c) = (-1, 2, -1) ∨ (a, b, c) = (-1, -1, 2) :=
begin
  have h2 : ∀ x y z : ℝ, x * y + 1 = z → y * z + 1 = x → y = -1 ∨ x = z :=
  begin
    intros x y z h h0,
    replace h0 := congr_arg2 (λ x y, x - y) h h0; simp only [] at h0,
    rwa [add_sub_add_right_eq_sub, mul_comm, ← mul_sub, ← neg_sub x,
         ← neg_one_mul, mul_eq_mul_right_iff, sub_eq_zero] at h0
  end,
  rcases h2 a b c h h0 with rfl | rfl,
  { rcases h2 _ _ _ h0 h1 with rfl | rfl,
    rw [← h0, mul_neg_one, neg_neg]; left; refl,
    rw [← h, mul_neg_one, neg_neg]; right; right; refl },
  { rcases h2 _ _ _ h0 h1 with rfl | rfl,
    rw [← h1, mul_neg_one, neg_neg]; right; left; refl,
    exfalso; exact sq_add_one_ne_self b h }
end

/-- Characterization of all good 3-periodic functions. -/
private lemma good_period3_iff (a : ℕ → ℝ) : (good a ∧ periodic a 3) ↔
  a = good_ex ∨ a = (λ i, good_ex (i + 1)) ∨ a = λ i, good_ex (i + 2) :=
begin
  split,
  /- The → direction is quite costly, but with a relatively shorter code -/
  { rintros ⟨h, h0⟩,
    obtain (h1 | h1 | h1) : (a 0, a 1, a 2) = (2, -1, -1) ∨
      (a 0, a 1, a 2) = (-1, 2, -1) ∨ (a 0, a 1, a 2) = (-1, -1, 2) :=
    begin
      apply good_solution,
      replace h := h 0,
      rwa [zero_add, zero_add] at h,
      replace h := h 1,
      rwa [bit0, ← add_assoc, ← periodic.map_mod_nat h0 3] at h,
      replace h := h 2,
      rwa [← periodic.map_mod_nat h0 3, ← periodic.map_mod_nat h0 4] at h
    end,
    left,
    work_on_goal 2 { right; right },
    work_on_goal 3 { right; left },
    all_goals {
      rw [prod.mk.inj_iff, prod.mk.inj_iff] at h1,
      rcases h1 with ⟨h1, h2, h3⟩,
      funext i; dsimp only [good_ex],
      rw ← periodic.map_mod_nat h0,
      try { rw ← nat.mod_add_mod },
      let p : fin 3 := ⟨i % 3, i.mod_lt three_pos⟩,
      fin_cases p using hp,
      all_goals {
        rw fin.ext_iff at hp,
        simp [p] at hp,
        rw hp; assumption } } },
  { have h : periodic good_ex 3 := λ i, by simp only [good_ex, nat.add_mod_right],
    rintros (rfl | rfl | rfl); split,
    exacts [good_ex_good, h, good_shift good_ex_good 1, periodic_shift h 1,
            good_shift good_ex_good 2, periodic_shift h 2] }
end

/-- Final solution, extra version -/
theorem final_solution_extra {n : ℕ} (h : 0 < n) {a : ℕ → ℝ} (h0 : periodic a n) :
  good a ↔ a = good_ex ∨ a = (λ i, good_ex (i + 1)) ∨ a = λ i, good_ex (i + 2) :=
begin
  rw ← good_period3_iff a,
  symmetry; apply and_iff_left_of_imp,
  intros h1; exact periodic_to_periodic3 h1 h h0
end

/-- Final solution, extra version, case 3 ∤ n -/
theorem final_solution_extra_not3dvd {n : ℕ} (h : 0 < n) (h0 : ¬ 3 ∣ n)
  {a : ℕ → ℝ} (h1 : periodic a n) : ¬good a :=
begin
  replace h0 := (or_iff_left h0).mp (nat.coprime_or_dvd_of_prime nat.prime_three n),
  rw nat.coprime_comm at h0,
  intros h2,
  replace h1 : periodic a 1 :=
  begin
    cases nat.exists_mul_mod_eq_one_of_coprime h0 (by norm_num) with q h3,
    have h4 := periodic_to_periodic3 h2 h h1,
    intros i,
    rw [← h3, ← periodic.map_mod_nat h4, nat.add_mod_mod, periodic.map_mod_nat h4,
        ← periodic.map_mod_nat h1, nat.add_mul_mod_self_left, periodic.map_mod_nat h1]
  end,
  replace h2 := h2 0,
  rw [bit0, ← add_assoc, h1 (0 + 1), h1] at h2,
  exact sq_add_one_ne_self (a 0) h2
end

/-- Final solution, original version -/
theorem final_solution (n : ℕ) (h : 0 < n) : (∃ a : ℕ → ℝ, good a ∧ periodic a n) ↔ 3 ∣ n :=
begin
  split,
  { rintros ⟨a, h0, h1⟩,
    by_contra h2,
    exact final_solution_extra_not3dvd h h2 h1 h0 },
  { rintros ⟨n, rfl⟩,
    use good_ex; split,
    exact good_ex_good,
    intros i,
    simp only [good_ex, nat.add_mul_mod_self_left] }
end

end IMO2018A2
end IMOSL