/-
Copyright (c) 2020 Sébastien Gouëzel. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Sébastien Gouëzel, Johan Commelin
-/
import analysis.analytic.basic
import combinatorics.composition

/-!
# Composition of analytic functions

in this file we prove that the composition of analytic functions is analytic.

The argument is the following. Assume `g z = ∑' qₙ (z, ..., z)` and `f y = ∑' pₖ (y, ..., y)`. Then

`g (f y) = ∑' qₙ (∑' pₖ (y, ..., y), ..., ∑' pₖ (y, ..., y))
= ∑' qₙ (p_{i₁} (y, ..., y), ..., p_{iₙ} (y, ..., y))`.

For each `n` and `i₁, ..., iₙ`, define a `i₁ + ... + iₙ` multilinear function mapping
`(y₀, ..., y_{i₁ + ... + iₙ - 1})` to
`qₙ (p_{i₁} (y₀, ..., y_{i₁-1}), p_{i₂} (y_{i₁}, ..., y_{i₁ + i₂ - 1}), ..., p_{iₙ} (....)))`.
Then `g ∘ f` is obtained by summing all these multilinear functions.

To formalize this, we use compositions of an integer `N`, i.e., its decompositions into
a sum `i₁ + ... + iₙ` of positive integers. Given such a composition `c` and two formal
multilinear series `q` and `p`, let `q.comp_along_composition p c` be the above multilinear
function. Then the `N`-th coefficient in the power series expansion of `g ∘ f` is the sum of these
terms over all `c : composition N`.

To complete the proof, we need to show that this power series has a positive radius of convergence.
This follows from the fact that `composition N` has cardinality `2^(N-1)` and estimates on
the norm of `qₙ` and `pₖ`, which give summability. We also need to show that it indeed converges to
`g ∘ f`. For this, we note that the composition of partial sums converges to `g ∘ f`, and that it
corresponds to a part of the whole sum, on a subset that increases to the whole space. By
summability of the norms, this implies the overall convergence.

## Main results

* `q.comp p` is the formal composition of the formal multilinear series `q` and `p`.
* `has_fpower_series_at.comp` states that if two functions `g` and `f` admit power series expansions
  `q` and `p`, then `g ∘ f` admits a power series expansion given by `q.comp p`.
* `analytic_at.comp` states that the composition of analytic functions is analytic.
* `formal_multilinear_series.comp_assoc` states that composition is associative on formal
  multilinear series.

## Implementation details

The main technical difficulty is to write down things. In particular, we need to define precisely
`q.comp_along_composition p c` and to show that it is indeed a continuous multilinear
function. This requires a whole interface built on the class `composition`. Once this is set,
the main difficulty is to reorder the sums, writing the composition of the partial sums as a sum
over some subset of `Σ n, composition n`. We need to check that the reordering is a bijection,
running over difficulties due to the dependent nature of the types under consideration, that are
controlled thanks to the interface for `composition`.

The associativity of composition on formal multilinear series is a nontrivial result: it does not
follow from the associativity of composition of analytic functions, as there is no uniqueness for
the formal multilinear series representing a function (and also, it holds even when the radius of
convergence of the series is `0`). Instead, we give a direct proof, which amounts to reordering
double sums in a careful way. The change of variables is a canonical (combinatorial) bijection
`composition.sigma_equiv_sigma_pi` between `(Σ (a : composition n), composition a.length)` and
`(Σ (c : composition n), Π (i : fin c.length), composition (c.blocks_fun i))`, and is described
in more details below in the paragraph on associativity.
-/

noncomputable theory

variables {𝕜 : Type*} [nondiscrete_normed_field 𝕜]
{E : Type*} [normed_group E] [normed_space 𝕜 E]
{F : Type*} [normed_group F] [normed_space 𝕜 F]
{G : Type*} [normed_group G] [normed_space 𝕜 G]
{H : Type*} [normed_group H] [normed_space 𝕜 H]

open filter list
open_locale topological_space big_operators classical nnreal ennreal

/-! ### Composing formal multilinear series -/

namespace formal_multilinear_series

/-!
In this paragraph, we define the composition of formal multilinear series, by summing over all
possible compositions of `n`.
-/

/-- Given a formal multilinear series `p`, a composition `c` of `n` and the index `i` of a
block of `c`, we may define a function on `fin n → E` by picking the variables in the `i`-th block
of `n`, and applying the corresponding coefficient of `p` to these variables. This function is
called `p.apply_composition c v i` for `v : fin n → E` and `i : fin c.length`. -/
def apply_composition
  (p : formal_multilinear_series 𝕜 E F) {n : ℕ} (c : composition n) :
  (fin n → E) → (fin (c.length) → F) :=
λ v i, p (c.blocks_fun i) (v ∘ (c.embedding i))

lemma apply_composition_ones (p : formal_multilinear_series 𝕜 E F) (n : ℕ) :
  p.apply_composition (composition.ones n) =
    λ v i, p 1 (λ _, v (fin.cast_le (composition.length_le _) i)) :=
begin
  funext v i,
  apply p.congr (composition.ones_blocks_fun _ _),
  intros j hjn hj1,
  obtain rfl : j = 0, { linarith },
  refine congr_arg v _,
  rw [fin.ext_iff, fin.coe_cast_le, composition.ones_embedding, fin.coe_mk],
end

lemma apply_composition_single (p : formal_multilinear_series 𝕜 E F) {n : ℕ} (hn : 0 < n)
  (v : fin n → E) : p.apply_composition (composition.single n hn) v = λ j, p n v :=
begin
  ext j,
  refine p.congr (by simp) (λ i hi1 hi2, _),
  dsimp,
  congr' 1,
  convert composition.single_embedding hn ⟨i, hi2⟩,
  cases j,
  have : j_val = 0 := le_bot_iff.1 (nat.lt_succ_iff.1 j_property),
  unfold_coes,
  congr; try { assumption <|> simp },
end

@[simp] lemma remove_zero_apply_composition
  (p : formal_multilinear_series 𝕜 E F) {n : ℕ} (c : composition n) :
  p.remove_zero.apply_composition c = p.apply_composition c :=
begin
  ext v i,
  simp [apply_composition, zero_lt_one.trans_le (c.one_le_blocks_fun i), remove_zero_of_pos],
end

/-- Technical lemma stating how `p.apply_composition` commutes with updating variables. This
will be the key point to show that functions constructed from `apply_composition` retain
multilinearity. -/
lemma apply_composition_update
  (p : formal_multilinear_series 𝕜 E F) {n : ℕ} (c : composition n)
  (j : fin n) (v : fin n → E) (z : E) :
  p.apply_composition c (function.update v j z) =
    function.update (p.apply_composition c v) (c.index j)
      (p (c.blocks_fun (c.index j))
        (function.update (v ∘ (c.embedding (c.index j))) (c.inv_embedding j) z)) :=
begin
  ext k,
  by_cases h : k = c.index j,
  { rw h,
    let r : fin (c.blocks_fun (c.index j)) → fin n := c.embedding (c.index j),
    simp only [function.update_same],
    change p (c.blocks_fun (c.index j)) ((function.update v j z) ∘ r) = _,
    let j' := c.inv_embedding j,
    suffices B : (function.update v j z) ∘ r = function.update (v ∘ r) j' z,
      by rw B,
    suffices C : (function.update v (r j') z) ∘ r = function.update (v ∘ r) j' z,
      by { convert C, exact (c.embedding_comp_inv j).symm },
    exact function.update_comp_eq_of_injective _ (c.embedding _).injective _ _ },
  { simp only [h, function.update_eq_self, function.update_noteq, ne.def, not_false_iff],
    let r : fin (c.blocks_fun k) → fin n := c.embedding k,
    change p (c.blocks_fun k) ((function.update v j z) ∘ r) = p (c.blocks_fun k) (v ∘ r),
    suffices B : (function.update v j z) ∘ r = v ∘ r, by rw B,
    apply function.update_comp_eq_of_not_mem_range,
    rwa c.mem_range_embedding_iff' }
end

@[simp] lemma comp_continuous_linear_map_apply_composition {n : ℕ}
  (p : formal_multilinear_series 𝕜 F G) (f : E →L[𝕜] F) (c : composition n) (v : fin n → E) :
  (p.comp_continuous_linear_map f).apply_composition c v = p.apply_composition c (f ∘ v) :=
by simp [apply_composition]

end formal_multilinear_series

namespace continuous_multilinear_map
open formal_multilinear_series

/-- Given a formal multilinear series `p`, a composition `c` of `n` and a continuous multilinear
map `f` in `c.length` variables, one may form a multilinear map in `n` variables by applying
the right coefficient of `p` to each block of the composition, and then applying `f` to the
resulting vector. It is called `f.comp_along_composition_aux p c`.
This function admits a version as a continuous multilinear map, called
`f.comp_along_composition p c` below. -/
def comp_along_composition_aux {n : ℕ}
  (p : formal_multilinear_series 𝕜 E F) (c : composition n)
  (f : continuous_multilinear_map 𝕜 (λ (i : fin c.length), F) G) :
  multilinear_map 𝕜 (λ i : fin n, E) G :=
{ to_fun    := λ v, f (p.apply_composition c v),
  map_add'  := λ v i x y, by simp only [apply_composition_update,
    continuous_multilinear_map.map_add],
  map_smul' := λ v i c x, by simp only [apply_composition_update,
    continuous_multilinear_map.map_smul] }

/-- The norm of `f.comp_along_composition_aux p c` is controlled by the product of
the norms of the relevant bits of `f` and `p`. -/
lemma comp_along_composition_aux_bound {n : ℕ}
  (p : formal_multilinear_series 𝕜 E F) (c : composition n)
  (f : continuous_multilinear_map 𝕜 (λ (i : fin c.length), F) G) (v : fin n → E) :
  ∥f.comp_along_composition_aux p c v∥ ≤
    ∥f∥ * (∏ i, ∥p (c.blocks_fun i)∥) * (∏ i : fin n, ∥v i∥) :=
calc ∥f.comp_along_composition_aux p c v∥ = ∥f (p.apply_composition c v)∥ : rfl
... ≤ ∥f∥ * ∏ i, ∥p.apply_composition c v i∥ : continuous_multilinear_map.le_op_norm _ _
... ≤ ∥f∥ * ∏ i, ∥p (c.blocks_fun i)∥ *
        ∏ j : fin (c.blocks_fun i), ∥(v ∘ (c.embedding i)) j∥ :
  begin
    apply mul_le_mul_of_nonneg_left _ (norm_nonneg _),
    refine finset.prod_le_prod (λ i hi, norm_nonneg _) (λ i hi, _),
    apply continuous_multilinear_map.le_op_norm,
  end
... = ∥f∥ * (∏ i, ∥p (c.blocks_fun i)∥) *
        ∏ i (j : fin (c.blocks_fun i)), ∥(v ∘ (c.embedding i)) j∥ :
  by rw [finset.prod_mul_distrib, mul_assoc]
... = ∥f∥ * (∏ i, ∥p (c.blocks_fun i)∥) * (∏ i : fin n, ∥v i∥) :
  by { rw [← c.blocks_fin_equiv.prod_comp, ← finset.univ_sigma_univ, finset.prod_sigma],
       congr }

/-- Given a formal multilinear series `p`, a composition `c` of `n` and a continuous multilinear
map `f` in `c.length` variables, one may form a continuous multilinear map in `n` variables by
applying the right coefficient of `p` to each block of the composition, and then applying `f` to
the resulting vector. It is called `f.comp_along_composition p c`. It is constructed from the
analogous multilinear function `f.comp_along_composition_aux p c`, together with a norm
control to get the continuity. -/
def comp_along_composition {n : ℕ}
  (p : formal_multilinear_series 𝕜 E F) (c : composition n)
  (f : continuous_multilinear_map 𝕜 (λ (i : fin c.length), F) G) :
  continuous_multilinear_map 𝕜 (λ i : fin n, E) G :=
(f.comp_along_composition_aux p c).mk_continuous _
  (f.comp_along_composition_aux_bound p c)

@[simp] lemma comp_along_composition_apply {n : ℕ}
  (p : formal_multilinear_series 𝕜 E F) (c : composition n)
  (f : continuous_multilinear_map 𝕜 (λ (i : fin c.length), F) G) (v : fin n → E) :
  (f.comp_along_composition p c) v = f (p.apply_composition c v) := rfl

end continuous_multilinear_map

namespace formal_multilinear_series

/-- Given two formal multilinear series `q` and `p` and a composition `c` of `n`, one may
form a continuous multilinear map in `n` variables by applying the right coefficient of `p` to each
block of the composition, and then applying `q c.length` to the resulting vector. It is
called `q.comp_along_composition p c`. It is constructed from the analogous multilinear
function `q.comp_along_composition_aux p c`, together with a norm control to get
the continuity. -/
def comp_along_composition {n : ℕ}
  (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F)
  (c : composition n) : continuous_multilinear_map 𝕜 (λ i : fin n, E) G :=
(q c.length).comp_along_composition p c

@[simp] lemma comp_along_composition_apply {n : ℕ}
  (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F)
  (c : composition n) (v : fin n → E) :
  (q.comp_along_composition p c) v = q c.length (p.apply_composition c v) := rfl

/-- The norm of `q.comp_along_composition p c` is controlled by the product of
the norms of the relevant bits of `q` and `p`. -/
lemma comp_along_composition_norm {n : ℕ}
  (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F)
  (c : composition n) :
  ∥q.comp_along_composition p c∥ ≤ ∥q c.length∥ * ∏ i, ∥p (c.blocks_fun i)∥ :=
multilinear_map.mk_continuous_norm_le _
  (mul_nonneg (norm_nonneg _) (finset.prod_nonneg (λ i hi, norm_nonneg _))) _

lemma comp_along_composition_nnnorm {n : ℕ}
  (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F)
  (c : composition n) :
  nnnorm (q.comp_along_composition p c) ≤ nnnorm (q c.length) * ∏ i, nnnorm (p (c.blocks_fun i)) :=
by { rw ← nnreal.coe_le_coe, push_cast, exact q.comp_along_composition_norm p c }

/-- Formal composition of two formal multilinear series. The `n`-th coefficient in the composition
is defined to be the sum of `q.comp_along_composition p c` over all compositions of
`n`. In other words, this term (as a multilinear function applied to `v_0, ..., v_{n-1}`) is
`∑'_{k} ∑'_{i₁ + ... + iₖ = n} pₖ (q_{i_1} (...), ..., q_{i_k} (...))`, where one puts all variables
`v_0, ..., v_{n-1}` in increasing order in the dots.

In general, the composition `q ∘ p` only makes sense when the constant coefficient of `p` vanishes.
We give a general formula but which ignores the value of `p 0` instead.
-/
protected def comp (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F) :
  formal_multilinear_series 𝕜 E G :=
λ n, ∑ c : composition n, q.comp_along_composition p c

/-- The `0`-th coefficient of `q.comp p` is `q 0`. Since these maps are multilinear maps in zero
variables, but on different spaces, we can not state this directly, so we state it when applied to
arbitrary vectors (which have to be the zero vector). -/
lemma comp_coeff_zero (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F)
  (v : fin 0 → E) (v' : fin 0 → F) :
  (q.comp p) 0 v = q 0 v' :=
begin
  let c : composition 0 := composition.ones 0,
  dsimp [formal_multilinear_series.comp],
  have : {c} = (finset.univ : finset (composition 0)),
  { apply finset.eq_of_subset_of_card_le; simp [finset.card_univ, composition_card 0] },
  rw [← this, finset.sum_singleton, comp_along_composition_apply],
  symmetry, congr'
end

@[simp] lemma comp_coeff_zero'
  (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F) (v : fin 0 → E) :
  (q.comp p) 0 v = q 0 (λ i, 0) :=
q.comp_coeff_zero p v _

/-- The `0`-th coefficient of `q.comp p` is `q 0`. When `p` goes from `E` to `E`, this can be
expressed as a direct equality -/
lemma comp_coeff_zero'' (q : formal_multilinear_series 𝕜 E F)
  (p : formal_multilinear_series 𝕜 E E) :
  (q.comp p) 0 = q 0 :=
by { ext v, exact q.comp_coeff_zero p _ _ }

/-- The first coefficient of a composition of formal multilinear series is the composition of the
first coefficients seen as continuous linear maps. -/
lemma comp_coeff_one (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F)
  (v : fin 1 → E) : (q.comp p) 1 v = q 1 (λ i, p 1 v) :=
begin
  have : {composition.ones 1} = (finset.univ : finset (composition 1)) :=
    finset.eq_univ_of_card _ (by simp [composition_card]),
  simp only [formal_multilinear_series.comp, comp_along_composition_apply, ← this,
    finset.sum_singleton],
  refine q.congr (by simp) (λ i hi1 hi2, _),
  simp only [apply_composition_ones],
  exact p.congr rfl (λ j hj1 hj2, by congr)
end

lemma remove_zero_comp_of_pos (q : formal_multilinear_series 𝕜 F G)
  (p : formal_multilinear_series 𝕜 E F) {n : ℕ} (hn : 0 < n) :
  q.remove_zero.comp p n = q.comp p n :=
begin
  ext v,
  simp only [formal_multilinear_series.comp, comp_along_composition,
    continuous_multilinear_map.comp_along_composition_apply, continuous_multilinear_map.sum_apply],
  apply finset.sum_congr rfl (λ c hc, _),
  rw remove_zero_of_pos _ (c.length_pos_of_pos hn)
end

@[simp] lemma comp_remove_zero (q : formal_multilinear_series 𝕜 F G)
  (p : formal_multilinear_series 𝕜 E F) :
  q.comp p.remove_zero = q.comp p :=
by { ext n, simp [formal_multilinear_series.comp] }

/-!
### The identity formal power series

We will now define the identity power series, and show that it is a neutral element for left and
right composition.
-/

section
variables (𝕜 E)

/-- The identity formal multilinear series, with all coefficients equal to `0` except for `n = 1`
where it is (the continuous multilinear version of) the identity. -/
def id : formal_multilinear_series 𝕜 E E
| 0 := 0
| 1 := (continuous_multilinear_curry_fin1 𝕜 E E).symm (continuous_linear_map.id 𝕜 E)
| _ := 0

/-- The first coefficient of `id 𝕜 E` is the identity. -/
@[simp] lemma id_apply_one (v : fin 1 → E) : (formal_multilinear_series.id 𝕜 E) 1 v = v 0 := rfl

/-- The `n`th coefficient of `id 𝕜 E` is the identity when `n = 1`. We state this in a dependent
way, as it will often appear in this form. -/
lemma id_apply_one' {n : ℕ} (h : n = 1) (v : fin n → E) :
  (id 𝕜 E) n v = v ⟨0, h.symm ▸ zero_lt_one⟩ :=
begin
  subst n,
  apply id_apply_one
end

/-- For `n ≠ 1`, the `n`-th coefficient of `id 𝕜 E` is zero, by definition. -/
@[simp] lemma id_apply_ne_one {n : ℕ} (h : n ≠ 1) : (formal_multilinear_series.id 𝕜 E) n = 0 :=
by { cases n, { refl }, cases n, { contradiction }, refl }

end

@[simp] theorem comp_id (p : formal_multilinear_series 𝕜 E F) : p.comp (id 𝕜 E) = p :=
begin
  ext1 n,
  dsimp [formal_multilinear_series.comp],
  rw finset.sum_eq_single (composition.ones n),
  show comp_along_composition p (id 𝕜 E) (composition.ones n) = p n,
  { ext v,
    rw comp_along_composition_apply,
    apply p.congr (composition.ones_length n),
    intros,
    rw apply_composition_ones,
    refine congr_arg v _,
    rw [fin.ext_iff, fin.coe_cast_le, fin.coe_mk, fin.coe_mk], },
  show ∀ (b : composition n),
    b ∈ finset.univ → b ≠ composition.ones n → comp_along_composition p (id 𝕜 E) b = 0,
  { assume b _ hb,
    obtain ⟨k, hk, lt_k⟩ : ∃ (k : ℕ) (H : k ∈ composition.blocks b), 1 < k :=
      composition.ne_ones_iff.1 hb,
    obtain ⟨i, i_lt, hi⟩ : ∃ (i : ℕ) (h : i < b.blocks.length), b.blocks.nth_le i h = k :=
      nth_le_of_mem hk,
    let j : fin b.length := ⟨i, b.blocks_length ▸ i_lt⟩,
    have A : 1 < b.blocks_fun j := by convert lt_k,
    ext v,
    rw [comp_along_composition_apply, continuous_multilinear_map.zero_apply],
    apply continuous_multilinear_map.map_coord_zero _ j,
    dsimp [apply_composition],
    rw id_apply_ne_one _ _ (ne_of_gt A),
    refl },
  { simp }
end

@[simp] theorem id_comp (p : formal_multilinear_series 𝕜 E F) (h : p 0 = 0) : (id 𝕜 F).comp p = p :=
begin
  ext1 n,
  by_cases hn : n = 0,
  { rw [hn, h],
    ext v,
    rw [comp_coeff_zero', id_apply_ne_one _ _ zero_ne_one],
    refl },
  { dsimp [formal_multilinear_series.comp],
    have n_pos : 0 < n := bot_lt_iff_ne_bot.mpr hn,
    rw finset.sum_eq_single (composition.single n n_pos),
    show comp_along_composition (id 𝕜 F) p (composition.single n n_pos) = p n,
    { ext v,
      rw [comp_along_composition_apply, id_apply_one' _ _ (composition.single_length n_pos)],
      dsimp [apply_composition],
      refine p.congr rfl (λ i him hin, congr_arg v $ _),
      ext, simp },
    show ∀ (b : composition n),
      b ∈ finset.univ → b ≠ composition.single n n_pos → comp_along_composition (id 𝕜 F) p b = 0,
    { assume b _ hb,
      have A : b.length ≠ 1, by simpa [composition.eq_single_iff_length] using hb,
      ext v,
      rw [comp_along_composition_apply, id_apply_ne_one _ _ A],
      refl },
    { simp } }
end

/-! ### Summability properties of the composition of formal power series-/

/-- If two formal multilinear series have positive radius of convergence, then the terms appearing
in the definition of their composition are also summable (when multiplied by a suitable positive
geometric term). -/
theorem comp_summable_nnreal
  (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F)
  (hq : 0 < q.radius) (hp : 0 < p.radius) :
  ∃ r > (0 : ℝ≥0),
    summable (λ i : Σ n, composition n, nnnorm (q.comp_along_composition p i.2) * r ^ i.1) :=
begin
  /- This follows from the fact that the growth rate of `∥qₙ∥` and `∥pₙ∥` is at most geometric,
  giving a geometric bound on each `∥q.comp_along_composition p op∥`, together with the
  fact that there are `2^(n-1)` compositions of `n`, giving at most a geometric loss. -/
  rcases ennreal.lt_iff_exists_nnreal_btwn.1 (lt_min ennreal.zero_lt_one hq) with ⟨rq, rq_pos, hrq⟩,
  rcases ennreal.lt_iff_exists_nnreal_btwn.1 (lt_min ennreal.zero_lt_one hp) with ⟨rp, rp_pos, hrp⟩,
  simp only [lt_min_iff, ennreal.coe_lt_one_iff, ennreal.coe_pos] at hrp hrq rp_pos rq_pos,
  obtain ⟨Cq, hCq0, hCq⟩ : ∃ Cq > 0, ∀ n, nnnorm (q n) * rq^n ≤ Cq :=
    q.nnnorm_mul_pow_le_of_lt_radius hrq.2,
  obtain ⟨Cp, hCp1, hCp⟩ : ∃ Cp ≥ 1, ∀ n, nnnorm (p n) * rp^n ≤ Cp,
  { rcases p.nnnorm_mul_pow_le_of_lt_radius hrp.2 with ⟨Cp, -, hCp⟩,
    exact ⟨max Cp 1, le_max_right _ _, λ n, (hCp n).trans (le_max_left _ _)⟩ },
  let r0 : ℝ≥0 := (4 * Cp)⁻¹,
  have r0_pos : 0 < r0 := nnreal.inv_pos.2 (mul_pos zero_lt_four (zero_lt_one.trans_le hCp1)),
  set r : ℝ≥0 := rp * rq * r0,
  have r_pos : 0 < r := mul_pos (mul_pos rp_pos rq_pos) r0_pos,
  have I : ∀ (i : Σ (n : ℕ), composition n),
    nnnorm (q.comp_along_composition p i.2) * r ^ i.1 ≤ Cq / 4 ^ i.1,
  { rintros ⟨n, c⟩,
    have A,
    calc nnnorm (q c.length) * rq ^ n ≤ nnnorm (q c.length)* rq ^ c.length :
      mul_le_mul' le_rfl (pow_le_pow_of_le_one rq.2 hrq.1.le c.length_le)
    ... ≤ Cq : hCq _,
    have B,
    calc ((∏ i, nnnorm (p (c.blocks_fun i))) * rp ^ n)
        = ∏ i, nnnorm (p (c.blocks_fun i)) * rp ^ c.blocks_fun i :
      by simp only [finset.prod_mul_distrib, finset.prod_pow_eq_pow_sum, c.sum_blocks_fun]
    ... ≤ ∏ i : fin c.length, Cp : finset.prod_le_prod' (λ i _, hCp _)
    ... = Cp ^ c.length : by simp
    ... ≤ Cp ^ n : pow_le_pow hCp1 c.length_le,
    calc nnnorm (q.comp_along_composition p c) * r ^ n
        ≤ (nnnorm (q c.length) * ∏ i, nnnorm (p (c.blocks_fun i))) * r ^ n :
          mul_le_mul' (q.comp_along_composition_nnnorm p c) le_rfl
    ... = (nnnorm (q c.length) * rq ^ n) * ((∏ i, nnnorm (p (c.blocks_fun i))) * rp ^ n) * r0 ^ n :
          by { simp only [r, mul_pow], ac_refl }
    ... ≤ Cq * Cp ^ n * r0 ^ n : mul_le_mul' (mul_le_mul' A B) le_rfl
    ... = Cq / 4 ^ n :
      begin
        simp only [r0],
        field_simp [mul_pow, (zero_lt_one.trans_le hCp1).ne'],
        ac_refl
      end },
  refine ⟨r, r_pos, nnreal.summable_of_le I _⟩,
  simp_rw div_eq_mul_inv,
  refine summable.mul_left _ _,
  have h4 : ∀ n : ℕ, 0 < (4 ^ n : ℝ≥0)⁻¹ := λ n, nnreal.inv_pos.2 (pow_pos zero_lt_four _),
  have : ∀ n : ℕ, has_sum (λ c : composition n, (4 ^ n : ℝ≥0)⁻¹) (2 ^ (n - 1) / 4 ^ n),
  { intro n,
    convert has_sum_fintype (λ c : composition n, (4 ^ n : ℝ≥0)⁻¹),
    simp [finset.card_univ, composition_card, div_eq_mul_inv] },
  refine nnreal.summable_sigma.2 ⟨λ n, (this n).summable, (nnreal.summable_nat_add_iff 1).1 _⟩,
  convert (nnreal.summable_geometric (nnreal.div_lt_one_of_lt one_lt_two)).mul_left (1 / 4),
  ext1 n,
  rw [(this _).tsum_eq, nat.add_sub_cancel],
  field_simp [← mul_assoc, pow_succ', mul_pow, show (4 : ℝ≥0) = 2 * 2, from (two_mul 2).symm,
    mul_right_comm]
end

/-- Bounding below the radius of the composition of two formal multilinear series assuming
summability over all compositions. -/
theorem le_comp_radius_of_summable
  (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F) (r : ℝ≥0)
  (hr : summable (λ i : (Σ n, composition n), nnnorm (q.comp_along_composition p i.2) * r ^ i.1)) :
  (r : ℝ≥0∞) ≤ (q.comp p).radius :=
begin
  refine le_radius_of_bound_nnreal _
    (∑' i : (Σ n, composition n), nnnorm (comp_along_composition q p i.snd) * r ^ i.fst) (λ n, _),
  calc nnnorm (formal_multilinear_series.comp q p n) * r ^ n ≤
  ∑' (c : composition n), nnnorm (comp_along_composition q p c) * r ^ n :
    begin
      rw [tsum_fintype, ← finset.sum_mul],
      exact mul_le_mul' (nnnorm_sum_le _ _) le_rfl
    end
  ... ≤ ∑' (i : Σ (n : ℕ), composition n), nnnorm (comp_along_composition q p i.snd) * r ^ i.fst :
    nnreal.tsum_comp_le_tsum_of_inj hr sigma_mk_injective
end

/-!
### Composing analytic functions

Now, we will prove that the composition of the partial sums of `q` and `p` up to order `N` is
given by a sum over some large subset of `Σ n, composition n` of `q.comp_along_composition p`, to
deduce that the series for `q.comp p` indeed converges to `g ∘ f` when `q` is a power series for
`g` and `p` is a power series for `f`.

This proof is a big reindexing argument of a sum. Since it is a bit involved, we define first
the source of the change of variables (`comp_partial_source`), its target
(`comp_partial_target`) and the change of variables itself (`comp_change_of_variables`) before
giving the main statement in `comp_partial_sum`. -/

/-- Source set in the change of variables to compute the composition of partial sums of formal
power series.
See also `comp_partial_sum`. -/
def comp_partial_sum_source (m M N : ℕ) : finset (Σ n, (fin n) → ℕ) :=
finset.sigma (finset.Ico m M) (λ (n : ℕ), fintype.pi_finset (λ (i : fin n), finset.Ico 1 N) : _)

@[simp] lemma mem_comp_partial_sum_source_iff (m M N : ℕ) (i : Σ n, (fin n) → ℕ) :
  i ∈ comp_partial_sum_source m M N ↔
    (m ≤ i.1 ∧ i.1 < M) ∧ ∀ (a : fin i.1), 1 ≤ i.2 a ∧ i.2 a < N :=
by simp only [comp_partial_sum_source, finset.Ico.mem, fintype.mem_pi_finset, finset.mem_sigma,
  iff_self]

/-- Change of variables appearing to compute the composition of partial sums of formal
power series -/
def comp_change_of_variables (m M N : ℕ) (i : Σ n, (fin n) → ℕ)
  (hi : i ∈ comp_partial_sum_source m M N) : (Σ n, composition n) :=
begin
  rcases i with ⟨n, f⟩,
  rw mem_comp_partial_sum_source_iff at hi,
  refine ⟨∑ j, f j, of_fn (λ a, f a), λ i hi', _, by simp [sum_of_fn]⟩,
  obtain ⟨j, rfl⟩ : ∃ (j : fin n), f j = i, by rwa [mem_of_fn, set.mem_range] at hi',
  exact (hi.2 j).1
end

@[simp] lemma comp_change_of_variables_length
  (m M N : ℕ) {i : Σ n, (fin n) → ℕ} (hi : i ∈ comp_partial_sum_source m M N) :
  composition.length (comp_change_of_variables m M N i hi).2 = i.1 :=
begin
  rcases i with ⟨k, blocks_fun⟩,
  dsimp [comp_change_of_variables],
  simp only [composition.length, map_of_fn, length_of_fn]
end

lemma comp_change_of_variables_blocks_fun
  (m M N : ℕ) {i : Σ n, (fin n) → ℕ} (hi : i ∈ comp_partial_sum_source m M N) (j : fin i.1) :
  (comp_change_of_variables m M N i hi).2.blocks_fun
    ⟨j, (comp_change_of_variables_length m M N hi).symm ▸ j.2⟩ = i.2 j :=
begin
  rcases i with ⟨n, f⟩,
  dsimp [composition.blocks_fun, composition.blocks, comp_change_of_variables],
  simp only [map_of_fn, nth_le_of_fn', function.comp_app],
  apply congr_arg,
  exact fin.eta _ _
end

/-- Target set in the change of variables to compute the composition of partial sums of formal
power series, here given a a set. -/
def comp_partial_sum_target_set (m M N : ℕ) : set (Σ n, composition n) :=
{i | (m ≤ i.2.length) ∧ (i.2.length < M) ∧ (∀ (j : fin i.2.length), i.2.blocks_fun j < N)}

lemma comp_partial_sum_target_subset_image_comp_partial_sum_source
  (m M N : ℕ) (i : Σ n, composition n) (hi : i ∈ comp_partial_sum_target_set m M N) :
  ∃ j (hj : j ∈ comp_partial_sum_source m M N), i = comp_change_of_variables m M N j hj :=
begin
  rcases i with ⟨n, c⟩,
  refine ⟨⟨c.length, c.blocks_fun⟩, _, _⟩,
  { simp only [comp_partial_sum_target_set, set.mem_set_of_eq] at hi,
    simp only [mem_comp_partial_sum_source_iff, hi.left, hi.right, true_and, and_true],
    exact λ a, c.one_le_blocks' _ },
  { dsimp [comp_change_of_variables],
    rw composition.sigma_eq_iff_blocks_eq,
    simp only [composition.blocks_fun, composition.blocks, subtype.coe_eta, nth_le_map'],
    conv_lhs { rw ← of_fn_nth_le c.blocks } }
end

/-- Target set in the change of variables to compute the composition of partial sums of formal
power series, here given a a finset.
See also `comp_partial_sum`. -/
def comp_partial_sum_target (m M N : ℕ) : finset (Σ n, composition n) :=
set.finite.to_finset $ ((finset.finite_to_set _).dependent_image _).subset $
  comp_partial_sum_target_subset_image_comp_partial_sum_source m M N

@[simp] lemma mem_comp_partial_sum_target_iff {m M N : ℕ} {a : Σ n, composition n} :
  a ∈ comp_partial_sum_target m M N ↔
    m ≤ a.2.length ∧ a.2.length < M ∧ (∀ (j : fin a.2.length), a.2.blocks_fun j < N) :=
by simp [comp_partial_sum_target, comp_partial_sum_target_set]

/-- `comp_change_of_variables m M N` is a bijection between `comp_partial_sum_source m M N`
and `comp_partial_sum_target m M N`, yielding equal sums for functions that correspond to each
other under the bijection. As `comp_change_of_variables m M N` is a dependent function, stating
that it is a bijection is not directly possible, but the consequence on sums can be stated
more easily. -/
lemma comp_change_of_variables_sum {α : Type*} [add_comm_monoid α] (m M N : ℕ)
  (f : (Σ (n : ℕ), fin n → ℕ) → α) (g : (Σ n, composition n) → α)
  (h : ∀ e (he : e ∈ comp_partial_sum_source m M N),
    f e = g (comp_change_of_variables m M N e he)) :
  ∑ e in comp_partial_sum_source m M N, f e = ∑ e in comp_partial_sum_target m M N, g e :=
begin
  apply finset.sum_bij (comp_change_of_variables m M N),
  -- We should show that the correspondance we have set up is indeed a bijection
  -- between the index sets of the two sums.
  -- 1 - show that the image belongs to `comp_partial_sum_target m N N`
  { rintros ⟨k, blocks_fun⟩ H,
    rw mem_comp_partial_sum_source_iff at H,
    simp only [mem_comp_partial_sum_target_iff, composition.length, composition.blocks, H.left,
               map_of_fn, length_of_fn, true_and, comp_change_of_variables],
    assume j,
    simp only [composition.blocks_fun, (H.right _).right, nth_le_of_fn'] },
  -- 2 - show that the composition gives the `comp_along_composition` application
  { rintros ⟨k, blocks_fun⟩ H,
    rw h },
  -- 3 - show that the map is injective
  { rintros ⟨k, blocks_fun⟩ ⟨k', blocks_fun'⟩ H H' heq,
    obtain rfl : k = k',
    { have := (comp_change_of_variables_length m M N H).symm,
      rwa [heq, comp_change_of_variables_length] at this, },
    congr,
    funext i,
    calc blocks_fun i = (comp_change_of_variables m M N _ H).2.blocks_fun _  :
     (comp_change_of_variables_blocks_fun m M N H i).symm
      ... = (comp_change_of_variables m M N _ H').2.blocks_fun _ :
        begin
          apply composition.blocks_fun_congr; try { rw heq },
          refl
        end
      ... = blocks_fun' i : comp_change_of_variables_blocks_fun m M N H' i },
  -- 4 - show that the map is surjective
  { assume i hi,
    apply comp_partial_sum_target_subset_image_comp_partial_sum_source m M N i,
    simpa [comp_partial_sum_target] using hi }
end

/-- The auxiliary set corresponding to the composition of partial sums asymptotically contains
all possible compositions. -/

/-- Composing the partial sums of two multilinear series coincides with the sum over all
compositions in `comp_partial_sum_target 0 N N`. This is precisely the motivation for the
definition of `comp_partial_sum_target`. -/
lemma comp_partial_sum
  (q : formal_multilinear_series 𝕜 F G) (p : formal_multilinear_series 𝕜 E F) (N : ℕ) (z : E) :
  q.partial_sum N (∑ i in finset.Ico 1 N, p i (λ j, z)) =
    ∑ i in comp_partial_sum_target 0 N N, q.comp_along_composition p i.2 (λ j, z) :=
begin
  -- we expand the composition, using the multilinearity of `q` to expand along each coordinate.
  suffices H : ∑ n in finset.range N, ∑ r in fintype.pi_finset (λ (i : fin n), finset.Ico 1 N),
    q n (λ (i : fin n), p (r i) (λ j, z)) =
    ∑ i in comp_partial_sum_target 0 N N, q.comp_along_composition p i.2 (λ j, z),
    by simpa only [formal_multilinear_series.partial_sum,
                   continuous_multilinear_map.map_sum_finset] using H,
  -- rewrite the first sum as a big sum over a sigma type, in the finset
  -- `comp_partial_sum_target 0 N N`
  rw [finset.range_eq_Ico, finset.sum_sigma'],
  -- use `comp_change_of_variables_sum`, saying that this change of variables respects sums
  apply comp_change_of_variables_sum 0 N N,
  rintros ⟨k, blocks_fun⟩ H,
  apply congr _ (comp_change_of_variables_length 0 N N H).symm,
  intros,
  rw ← comp_change_of_variables_blocks_fun 0 N N H,
  refl
end

end formal_multilinear_series

open formal_multilinear_series

/-- If two functions `g` and `f` have power series `q` and `p` respectively at `f x` and `x`, then
`g ∘ f` admits the power series `q.comp p` at `x`. -/
theorem has_fpower_series_at.comp {g : F → G} {f : E → F}
  {q : formal_multilinear_series 𝕜 F G} {p : formal_multilinear_series 𝕜 E F} {x : E}
  (hg : has_fpower_series_at g q (f x)) (hf : has_fpower_series_at f p x) :
  has_fpower_series_at (g ∘ f) (q.comp p) x :=
begin
  /- Consider `rf` and `rg` such that `f` and `g` have power series expansion on the disks
  of radius `rf` and `rg`. -/
  rcases hg with ⟨rg, Hg⟩,
  rcases hf with ⟨rf, Hf⟩,
  /- The terms defining `q.comp p` are geometrically summable in a disk of some radius `r`. -/
  rcases q.comp_summable_nnreal p Hg.radius_pos Hf.radius_pos with ⟨r, r_pos : 0 < r, hr⟩,
  /- We will consider `y` which is smaller than `r` and `rf`, and also small enough that
  `f (x + y)` is close enough to `f x` to be in the disk where `g` is well behaved. Let
  `min (r, rf, δ)` be this new radius.-/
  have : continuous_at f x := Hf.analytic_at.continuous_at,
  obtain ⟨δ, δpos, hδ⟩ : ∃ (δ : ℝ≥0∞) (H : 0 < δ),
    ∀ {z : E}, z ∈ emetric.ball x δ → f z ∈ emetric.ball (f x) rg,
  { have : emetric.ball (f x) rg ∈ 𝓝 (f x) := emetric.ball_mem_nhds _ Hg.r_pos,
    rcases emetric.mem_nhds_iff.1 (Hf.analytic_at.continuous_at this) with ⟨δ, δpos, Hδ⟩,
    exact ⟨δ, δpos, λ z hz, Hδ hz⟩ },
  let rf' := min rf δ,
  have min_pos : 0 < min rf' r,
    by simp only [r_pos, Hf.r_pos, δpos, lt_min_iff, ennreal.coe_pos, and_self],
  /- We will show that `g ∘ f` admits the power series `q.comp p` in the disk of
  radius `min (r, rf', δ)`. -/
  refine ⟨min rf' r, _⟩,
  refine ⟨le_trans (min_le_right rf' r)
    (formal_multilinear_series.le_comp_radius_of_summable q p r hr), min_pos, λ y hy, _⟩,
  /- Let `y` satisfy `∥y∥ < min (r, rf', δ)`. We want to show that `g (f (x + y))` is the sum of
  `q.comp p` applied to `y`. -/
  -- First, check that `y` is small enough so that estimates for `f` and `g` apply.
  have y_mem : y ∈ emetric.ball (0 : E) rf :=
    (emetric.ball_subset_ball (le_trans (min_le_left _ _) (min_le_left _ _))) hy,
  have fy_mem : f (x + y) ∈ emetric.ball (f x) rg,
  { apply hδ,
    have : y ∈ emetric.ball (0 : E) δ :=
      (emetric.ball_subset_ball (le_trans (min_le_left _ _) (min_le_right _ _))) hy,
    simpa [edist_eq_coe_nnnorm_sub, edist_eq_coe_nnnorm] },
  /- Now the proof starts. To show that the sum of `q.comp p` at `y` is `g (f (x + y))`, we will
  write `q.comp p` applied to `y` as a big sum over all compositions. Since the sum is
  summable, to get its convergence it suffices to get the convergence along some increasing sequence
  of sets. We will use the sequence of sets `comp_partial_sum_target 0 n n`, along which the sum is
  exactly the composition of the partial sums of `q` and `p`, by design. To show that it converges
  to `g (f (x + y))`, pointwise convergence would not be enough, but we have uniform convergence
  to save the day. -/
  -- First step: the partial sum of `p` converges to `f (x + y)`.
  have A : tendsto (λ n, ∑ a in finset.Ico 1 n, p a (λ b, y)) at_top (𝓝 (f (x + y) - f x)),
  { have L : ∀ᶠ n in at_top, ∑ a in finset.range n, p a (λ b, y) - f x =
      ∑ a in finset.Ico 1 n, p a (λ b, y),
    { rw eventually_at_top,
      refine ⟨1, λ n hn, _⟩,
      symmetry,
      rw [eq_sub_iff_add_eq', finset.range_eq_Ico, ← Hf.coeff_zero (λi, y),
          finset.sum_eq_sum_Ico_succ_bot hn] },
    have : tendsto (λ n, ∑ a in finset.range n, p a (λ b, y) - f x) at_top (𝓝 (f (x + y) - f x)) :=
      (Hf.has_sum y_mem).tendsto_sum_nat.sub tendsto_const_nhds,
    exact tendsto.congr' L this },
  -- Second step: the composition of the partial sums of `q` and `p` converges to `g (f (x + y))`.
  have B : tendsto (λ n, q.partial_sum n (∑ a in finset.Ico 1 n, p a (λ b, y)))
    at_top (𝓝 (g (f (x + y)))),
  { -- we use the fact that the partial sums of `q` converge locally uniformly to `g`, and that
    -- composition passes to the limit under locally uniform convergence.
    have B₁ : continuous_at (λ (z : F), g (f x + z)) (f (x + y) - f x),
    { refine continuous_at.comp _ (continuous_const.add continuous_id).continuous_at,
      simp only [add_sub_cancel'_right, id.def],
      exact Hg.continuous_on.continuous_at (mem_nhds_sets (emetric.is_open_ball) fy_mem) },
    have B₂ : f (x + y) - f x ∈ emetric.ball (0 : F) rg,
      by simpa [edist_eq_coe_nnnorm, edist_eq_coe_nnnorm_sub] using fy_mem,
    rw [← nhds_within_eq_of_open B₂ emetric.is_open_ball] at A,
    convert Hg.tendsto_locally_uniformly_on.tendsto_comp B₁.continuous_within_at B₂ A,
    simp only [add_sub_cancel'_right] },
  -- Third step: the sum over all compositions in `comp_partial_sum_target 0 n n` converges to
  -- `g (f (x + y))`. As this sum is exactly the composition of the partial sum, this is a direct
  -- consequence of the second step
  have C : tendsto (λ n,
    ∑ i in comp_partial_sum_target 0 n n, q.comp_along_composition p i.2 (λ j, y))
    at_top (𝓝 (g (f (x + y)))),
  by simpa [comp_partial_sum] using B,
  -- Fourth step: the sum over all compositions is `g (f (x + y))`. This follows from the
  -- convergence along a subsequence proved in the third step, and the fact that the sum is Cauchy
  -- thanks to the summability properties.
  have D : has_sum (λ i : (Σ n, composition n),
    q.comp_along_composition p i.2 (λ j, y)) (g (f (x + y))),
  { have cau : cauchy_seq (λ (s : finset (Σ n, composition n)),
      ∑ i in s, q.comp_along_composition p i.2 (λ j, y)),
    { apply cauchy_seq_finset_of_norm_bounded _ (nnreal.summable_coe.2 hr) _,
      simp only [coe_nnnorm, nnreal.coe_mul, nnreal.coe_pow],
      rintros ⟨n, c⟩,
      calc ∥(comp_along_composition q p c) (λ (j : fin n), y)∥
      ≤ ∥comp_along_composition q p c∥ * ∏ j : fin n, ∥y∥ :
        by apply continuous_multilinear_map.le_op_norm
      ... ≤ ∥comp_along_composition q p c∥ * (r : ℝ) ^ n :
        begin
          apply mul_le_mul_of_nonneg_left _ (norm_nonneg _),
          rw [finset.prod_const, finset.card_fin],
          apply pow_le_pow_of_le_left (norm_nonneg _),
          rw [emetric.mem_ball, edist_eq_coe_nnnorm] at hy,
          have := (le_trans (le_of_lt hy) (min_le_right _ _)),
          rwa [ennreal.coe_le_coe, ← nnreal.coe_le_coe, coe_nnnorm] at this
        end },
    exact tendsto_nhds_of_cauchy_seq_of_subseq cau
          comp_partial_sum_target_tendsto_at_top C },
  -- Fifth step: the sum over `n` of `q.comp p n` can be expressed as a particular resummation of
  -- the sum over all compositions, by grouping together the compositions of the same
  -- integer `n`. The convergence of the whole sum therefore implies the converence of the sum
  -- of `q.comp p n`
  have E : has_sum (λ n, (q.comp p) n (λ j, y)) (g (f (x + y))),
  { apply D.sigma,
    assume n,
    dsimp [formal_multilinear_series.comp],
    convert has_sum_fintype _,
    simp only [continuous_multilinear_map.sum_apply],
    refl },
  exact E
end

/-- If two functions `g` and `f` are analytic respectively at `f x` and `x`, then `g ∘ f` is
analytic at `x`. -/
theorem analytic_at.comp {g : F → G} {f : E → F} {x : E}
  (hg : analytic_at 𝕜 g (f x)) (hf : analytic_at 𝕜 f x) : analytic_at 𝕜 (g ∘ f) x :=
let ⟨q, hq⟩ := hg, ⟨p, hp⟩ := hf in (hq.comp hp).analytic_at

/-!
### Associativity of the composition of formal multilinear series

In this paragraph, we prove the associativity of the composition of formal power series.
By definition,
```
(r.comp q).comp p n v
= ∑_{i₁ + ... + iₖ = n} (r.comp q)ₖ (p_{i₁} (v₀, ..., v_{i₁ -1}), p_{i₂} (...), ..., p_{iₖ}(...))
= ∑_{a : composition n} (r.comp q) a.length (apply_composition p a v)
```
decomposing `r.comp q` in the same way, we get
```
(r.comp q).comp p n v
= ∑_{a : composition n} ∑_{b : composition a.length}
  r b.length (apply_composition q b (apply_composition p a v))
```
On the other hand,
```
r.comp (q.comp p) n v = ∑_{c : composition n} r c.length (apply_composition (q.comp p) c v)
```
Here, `apply_composition (q.comp p) c v` is a vector of length `c.length`, whose `i`-th term is
given by `(q.comp p) (c.blocks_fun i) (v_l, v_{l+1}, ..., v_{m-1})` where `{l, ..., m-1}` is the
`i`-th block in the composition `c`, of length `c.blocks_fun i` by definition. To compute this term,
we expand it as `∑_{dᵢ : composition (c.blocks_fun i)} q dᵢ.length (apply_composition p dᵢ v')`,
where `v' = (v_l, v_{l+1}, ..., v_{m-1})`. Therefore, we get
```
r.comp (q.comp p) n v =
∑_{c : composition n} ∑_{d₀ : composition (c.blocks_fun 0),
  ..., d_{c.length - 1} : composition (c.blocks_fun (c.length - 1))}
  r c.length (λ i, q dᵢ.length (apply_composition p dᵢ v'ᵢ))
```
To show that these terms coincide, we need to explain how to reindex the sums to put them in
bijection (and then the terms we are summing will correspond to each other). Suppose we have a
composition `a` of `n`, and a composition `b` of `a.length`. Then `b` indicates how to group
together some blocks of `a`, giving altogether `b.length` blocks of blocks. These blocks of blocks
can be called `d₀, ..., d_{a.length - 1}`, and one obtains a composition `c` of `n` by saying that
each `dᵢ` is one single block. Conversely, if one starts from `c` and the `dᵢ`s, one can concatenate
the `dᵢ`s to obtain a composition `a` of `n`, and register the lengths of the `dᵢ`s in a composition
`b` of `a.length`.

An example might be enlightening. Suppose `a = [2, 2, 3, 4, 2]`. It is a composition of
length 5 of 13. The content of the blocks may be represented as `0011222333344`.
Now take `b = [2, 3]` as a composition of `a.length = 5`. It says that the first 2 blocks of `a`
should be merged, and the last 3 blocks of `a` should be merged, giving a new composition of `13`
made of two blocks of length `4` and `9`, i.e., `c = [4, 9]`. But one can also remember that
the new first block was initially made of two blocks of size `2`, so `d₀ = [2, 2]`, and the new
second block was initially made of three blocks of size `3`, `4` and `2`, so `d₁ = [3, 4, 2]`.

This equivalence is called `composition.sigma_equiv_sigma_pi n` below.

We start with preliminary results on compositions, of a very specialized nature, then define the
equivalence `composition.sigma_equiv_sigma_pi n`, and we deduce finally the associativity of
composition of formal multilinear series in `formal_multilinear_series.comp_assoc`.
-/

namespace composition

variable {n : ℕ}

/-- Rewriting equality in the dependent type `Σ (a : composition n), composition a.length)` in
non-dependent terms with lists, requiring that the blocks coincide. -/
lemma sigma_composition_eq_iff (i j : Σ (a : composition n), composition a.length) :
  i = j ↔ i.1.blocks = j.1.blocks ∧ i.2.blocks = j.2.blocks :=
begin
  refine ⟨by rintro rfl; exact ⟨rfl, rfl⟩, _⟩,
  rcases i with ⟨a, b⟩,
  rcases j with ⟨a', b'⟩,
  rintros ⟨h, h'⟩,
  have H : a = a', by { ext1, exact h },
  induction H, congr, ext1, exact h'
end

/-- Rewriting equality in the dependent type
`Σ (c : composition n), Π (i : fin c.length), composition (c.blocks_fun i)` in
non-dependent terms with lists, requiring that the lists of blocks coincide. -/
lemma sigma_pi_composition_eq_iff
  (u v : Σ (c : composition n), Π (i : fin c.length), composition (c.blocks_fun i)) :
  u = v ↔ of_fn (λ i, (u.2 i).blocks) = of_fn (λ i, (v.2 i).blocks) :=
begin
  refine ⟨λ H, by rw H, λ H, _⟩,
  rcases u with ⟨a, b⟩,
  rcases v with ⟨a', b'⟩,
  dsimp at H,
  have h : a = a',
  { ext1,
    have : map list.sum (of_fn (λ (i : fin (composition.length a)), (b i).blocks)) =
      map list.sum (of_fn (λ (i : fin (composition.length a')), (b' i).blocks)), by rw H,
    simp only [map_of_fn] at this,
    change of_fn (λ (i : fin (composition.length a)), (b i).blocks.sum) =
      of_fn (λ (i : fin (composition.length a')), (b' i).blocks.sum) at this,
    simpa [composition.blocks_sum, composition.of_fn_blocks_fun] using this },
  induction h,
  simp only [true_and, eq_self_iff_true, heq_iff_eq],
  ext i : 2,
  have : nth_le (of_fn (λ (i : fin (composition.length a)), (b i).blocks)) i (by simp [i.is_lt]) =
         nth_le (of_fn (λ (i : fin (composition.length a)), (b' i).blocks)) i (by simp [i.is_lt]) :=
    nth_le_of_eq H _,
  rwa [nth_le_of_fn, nth_le_of_fn] at this
end

/-- When `a` is a composition of `n` and `b` is a composition of `a.length`, `a.gather b` is the
composition of `n` obtained by gathering all the blocks of `a` corresponding to a block of `b`.
For instance, if `a = [6, 5, 3, 5, 2]` and `b = [2, 3]`, one should gather together
the first two blocks of `a` and its last three blocks, giving `a.gather b = [11, 10]`. -/
def gather (a : composition n) (b : composition a.length) : composition n :=
{ blocks := (a.blocks.split_wrt_composition b).map sum,
  blocks_pos :=
  begin
    rw forall_mem_map_iff,
    intros j hj,
    suffices H : ∀ i ∈ j, 1 ≤ i, from
      calc 0 < j.length : length_pos_of_mem_split_wrt_composition hj
        ... ≤ j.sum    : length_le_sum_of_one_le _ H,
    intros i hi,
    apply a.one_le_blocks,
    rw ← a.blocks.join_split_wrt_composition b,
    exact mem_join_of_mem hj hi,
  end,
  blocks_sum := by { rw [← sum_join, join_split_wrt_composition, a.blocks_sum] } }

lemma length_gather (a : composition n) (b : composition a.length) :
  length (a.gather b) = b.length :=
show (map list.sum (a.blocks.split_wrt_composition b)).length = b.blocks.length,
by rw [length_map, length_split_wrt_composition]

/-- An auxiliary function used in the definition of `sigma_equiv_sigma_pi` below, associating to
two compositions `a` of `n` and `b` of `a.length`, and an index `i` bounded by the length of
`a.gather b`, the subcomposition of `a` made of those blocks belonging to the `i`-th block of
`a.gather b`. -/
def sigma_composition_aux (a : composition n) (b : composition a.length)
  (i : fin (a.gather b).length) :
  composition ((a.gather b).blocks_fun i) :=
{ blocks := nth_le (a.blocks.split_wrt_composition b) i
    (by { rw [length_split_wrt_composition, ← length_gather], exact i.2 }),
  blocks_pos := assume i hi, a.blocks_pos
    (by { rw ← a.blocks.join_split_wrt_composition b,
          exact mem_join_of_mem (nth_le_mem _ _ _) hi }),
  blocks_sum := by simp only [composition.blocks_fun, nth_le_map', composition.gather] }

lemma length_sigma_composition_aux (a : composition n) (b : composition a.length)
  (i : fin b.length) :
  composition.length (composition.sigma_composition_aux a b ⟨i, (length_gather a b).symm ▸ i.2⟩) =
    composition.blocks_fun b i :=
show list.length (nth_le (split_wrt_composition a.blocks b) i _) = blocks_fun b i,
by { rw [nth_le_map_rev list.length, nth_le_of_eq (map_length_split_wrt_composition _ _)], refl }

lemma blocks_fun_sigma_composition_aux (a : composition n) (b : composition a.length)
  (i : fin b.length) (j : fin (blocks_fun b i)) :
  blocks_fun (sigma_composition_aux a b ⟨i, (length_gather a b).symm ▸ i.2⟩)
      ⟨j, (length_sigma_composition_aux a b i).symm ▸ j.2⟩ = blocks_fun a (embedding b i j) :=
show nth_le (nth_le _ _ _) _ _ = nth_le a.blocks _ _,
by { rw [nth_le_of_eq (nth_le_split_wrt_composition _ _ _), nth_le_drop', nth_le_take'], refl }

/-- Auxiliary lemma to prove that the composition of formal multilinear series is associative.

Consider a composition `a` of `n` and a composition `b` of `a.length`. Grouping together some
blocks of `a` according to `b` as in `a.gather b`, one can compute the total size of the blocks
of `a` up to an index `size_up_to b i + j` (where the `j` corresponds to a set of blocks of `a`
that do not fill a whole block of `a.gather b`). The first part corresponds to a sum of blocks
in `a.gather b`, and the second one to a sum of blocks in the next block of
`sigma_composition_aux a b`. This is the content of this lemma. -/
lemma size_up_to_size_up_to_add (a : composition n) (b : composition a.length)
  {i j : ℕ} (hi : i < b.length) (hj : j < blocks_fun b ⟨i, hi⟩) :
  size_up_to a (size_up_to b i + j) = size_up_to (a.gather b) i +
    (size_up_to (sigma_composition_aux a b ⟨i, (length_gather a b).symm ▸ hi⟩) j) :=
begin
  induction j with j IHj,
  { show sum (take ((b.blocks.take i).sum) a.blocks) =
      sum (take i (map sum (split_wrt_composition a.blocks b))),
    induction i with i IH,
    { refl },
    { have A : i < b.length := nat.lt_of_succ_lt hi,
      have B : i < list.length (map list.sum (split_wrt_composition a.blocks b)), by simp [A],
      have C : 0 < blocks_fun b ⟨i, A⟩ := composition.blocks_pos' _ _ _,
      rw [sum_take_succ _ _ B, ← IH A C],
      have : take (sum (take i b.blocks)) a.blocks =
        take (sum (take i b.blocks)) (take (sum (take (i+1) b.blocks)) a.blocks),
      { rw [take_take, min_eq_left],
        apply monotone_sum_take _ (nat.le_succ _) },
      rw [this, nth_le_map', nth_le_split_wrt_composition,
        ← take_append_drop (sum (take i b.blocks))
          ((take (sum (take (nat.succ i) b.blocks)) a.blocks)), sum_append],
      congr,
      rw [take_append_drop] } },
  { have A : j < blocks_fun b ⟨i, hi⟩ := lt_trans (lt_add_one j) hj,
    have B : j < length (sigma_composition_aux a b ⟨i, (length_gather a b).symm ▸ hi⟩),
      by { convert A, rw [← length_sigma_composition_aux], refl },
    have C : size_up_to b i + j < size_up_to b (i + 1),
    { simp only [size_up_to_succ b hi, add_lt_add_iff_left],
      exact A },
    have D : size_up_to b i + j < length a := lt_of_lt_of_le C (b.size_up_to_le _),
    have : size_up_to b i + nat.succ j = (size_up_to b i + j).succ := rfl,
    rw [this, size_up_to_succ _ D, IHj A, size_up_to_succ _ B],
    simp only [sigma_composition_aux, add_assoc, add_left_inj, fin.coe_mk],
    rw [nth_le_of_eq (nth_le_split_wrt_composition _ _ _), nth_le_drop', nth_le_take _ _ C] }
end

/--
Natural equivalence between `(Σ (a : composition n), composition a.length)` and
`(Σ (c : composition n), Π (i : fin c.length), composition (c.blocks_fun i))`, that shows up as a
change of variables in the proof that composition of formal multilinear series is associative.

Consider a composition `a` of `n` and a composition `b` of `a.length`. Then `b` indicates how to
group together some blocks of `a`, giving altogether `b.length` blocks of blocks. These blocks of
blocks can be called `d₀, ..., d_{a.length - 1}`, and one obtains a composition `c` of `n` by
saying that each `dᵢ` is one single block. The map `⟨a, b⟩ → ⟨c, (d₀, ..., d_{a.length - 1})⟩` is
the direct map in the equiv.

Conversely, if one starts from `c` and the `dᵢ`s, one can join the `dᵢ`s to obtain a composition
`a` of `n`, and register the lengths of the `dᵢ`s in a composition `b` of `a.length`. This is the
inverse map of the equiv.
-/
def sigma_equiv_sigma_pi (n : ℕ) :
  (Σ (a : composition n), composition a.length) ≃
  (Σ (c : composition n), Π (i : fin c.length), composition (c.blocks_fun i)) :=
{ to_fun := λ i, ⟨i.1.gather i.2, i.1.sigma_composition_aux i.2⟩,
  inv_fun := λ i, ⟨
    { blocks := (of_fn (λ j, (i.2 j).blocks)).join,
      blocks_pos :=
      begin
        simp only [and_imp, mem_join, exists_imp_distrib, forall_mem_of_fn_iff],
        exact λ i j hj, composition.blocks_pos _ hj
      end,
      blocks_sum := by simp [sum_of_fn, composition.blocks_sum, composition.sum_blocks_fun] },
    { blocks := of_fn (λ j, (i.2 j).length),
      blocks_pos := forall_mem_of_fn_iff.2
        (λ j, composition.length_pos_of_pos _ (composition.blocks_pos' _ _ _)),
      blocks_sum := by { dsimp only [composition.length], simp [sum_of_fn] } }⟩,
  left_inv :=
  begin
    -- the fact that we have a left inverse is essentially `join_split_wrt_composition`,
    -- but we need to massage it to take care of the dependent setting.
    rintros ⟨a, b⟩,
    rw sigma_composition_eq_iff,
    dsimp,
    split,
    { have A := length_map list.sum (split_wrt_composition a.blocks b),
      conv_rhs { rw [← join_split_wrt_composition a.blocks b,
        ← of_fn_nth_le (split_wrt_composition a.blocks b)] },
      congr,
      { exact A },
      { exact (fin.heq_fun_iff A).2 (λ i, rfl) } },
    { have B : composition.length (composition.gather a b) = list.length b.blocks :=
        composition.length_gather _ _,
      conv_rhs { rw [← of_fn_nth_le b.blocks] },
      congr' 1,
      apply (fin.heq_fun_iff B).2 (λ i, _),
      rw [sigma_composition_aux, composition.length, nth_le_map_rev list.length,
          nth_le_of_eq (map_length_split_wrt_composition _ _)], refl }
  end,
  right_inv :=
  begin
    -- the fact that we have a right inverse is essentially `split_wrt_composition_join`,
    -- but we need to massage it to take care of the dependent setting.
    rintros ⟨c, d⟩,
    have : map list.sum (of_fn (λ (i : fin (composition.length c)), (d i).blocks)) = c.blocks,
      by simp [map_of_fn, (∘), composition.blocks_sum, composition.of_fn_blocks_fun],
    rw sigma_pi_composition_eq_iff,
    dsimp,
    congr,
    { ext1,
      dsimp [composition.gather],
      rwa split_wrt_composition_join,
      simp only [map_of_fn] },
    { rw fin.heq_fun_iff,
      { assume i,
        dsimp [composition.sigma_composition_aux],
        rw [nth_le_of_eq (split_wrt_composition_join _ _ _)],
        { simp only [nth_le_of_fn'] },
        { simp only [map_of_fn] } },
      { congr,
        ext1,
        dsimp [composition.gather],
        rwa split_wrt_composition_join,
        simp only [map_of_fn] } }
  end }

end composition

namespace formal_multilinear_series
open composition

theorem comp_assoc (r : formal_multilinear_series 𝕜 G H) (q : formal_multilinear_series 𝕜 F G)
  (p : formal_multilinear_series 𝕜 E F) :
  (r.comp q).comp p = r.comp (q.comp p) :=
begin
  ext n v,
  /- First, rewrite the two compositions appearing in the theorem as two sums over complicated
  sigma types, as in the description of the proof above. -/
  let f : (Σ (a : composition n), composition a.length) → H :=
    λ c, r c.2.length (apply_composition q c.2 (apply_composition p c.1 v)),
  let g : (Σ (c : composition n), Π (i : fin c.length), composition (c.blocks_fun i)) → H :=
    λ c, r c.1.length (λ (i : fin c.1.length),
      q (c.2 i).length (apply_composition p (c.2 i) (v ∘ c.1.embedding i))),
  suffices : ∑ c, f c = ∑ c, g c,
    by simpa only [formal_multilinear_series.comp, continuous_multilinear_map.sum_apply,
      comp_along_composition_apply, continuous_multilinear_map.map_sum, finset.sum_sigma',
      apply_composition],
  /- Now, we use `composition.sigma_equiv_sigma_pi n` to change
  variables in the second sum, and check that we get exactly the same sums. -/
  rw ← (sigma_equiv_sigma_pi n).sum_comp,
  /- To check that we have the same terms, we should check that we apply the same component of
  `r`, and the same component of `q`, and the same component of `p`, to the same coordinate of
  `v`. This is true by definition, but at each step one needs to convince Lean that the types
  one considers are the same, using a suitable congruence lemma to avoid dependent type issues.
  This dance has to be done three times, one for `r`, one for `q` and one for `p`.-/
  apply finset.sum_congr rfl,
  rintros ⟨a, b⟩ _,
  dsimp [f, g, sigma_equiv_sigma_pi],
  -- check that the `r` components are the same. Based on `composition.length_gather`
  apply r.congr (composition.length_gather a b).symm,
  intros i hi1 hi2,
  -- check that the `q` components are the same. Based on `length_sigma_composition_aux`
  apply q.congr (length_sigma_composition_aux a b _).symm,
  intros j hj1 hj2,
  -- check that the `p` components are the same. Based on `blocks_fun_sigma_composition_aux`
  apply p.congr (blocks_fun_sigma_composition_aux a b _ _).symm,
  intros k hk1 hk2,
  -- finally, check that the coordinates of `v` one is using are the same. Based on
  -- `size_up_to_size_up_to_add`.
  refine congr_arg v (fin.eq_of_veq _),
  dsimp [composition.embedding],
  rw [size_up_to_size_up_to_add _ _ hi1 hj1, add_assoc],
end

end formal_multilinear_series


import algebra.lie.classical
import linear_algebra.unitary_group
import data.real.basic
import data.nat.basic
import linear_algebra.matrix.determinant
import linear_algebra.matrix.nonsingular_inverse

/-
Unrelated ideas that have come up:
* Maybe refactoring ᵀ into a typeclass
-/

open_locale matrix big_operators

variables {l : Type*}

namespace matrix

variables {m n o p q : Type*} {m' n' p' : o → Type*}
variables {R : Type*} {S : Type*} {α : Type*} {β : Type*} [has_neg R] [ring S]

lemma from_blocks_neg_2
  (A : matrix n l R) (B : matrix n m R) (C : matrix o l R) (D : matrix o m R) :
  - (from_blocks A B C D) = from_blocks (-A) (-B) (-C) (-D) :=
begin
  ext i j, cases i; cases j; simp [from_blocks]
end

lemma from_blocks_neg
  (A : matrix n l S) (B : matrix n m S) (C : matrix o l S) (D : matrix o m S) :
  (-1 : S) • (from_blocks A B C D) = - from_blocks (A) (B) (C) (D) :=
begin
  ext i j, cases i; cases j; simp [from_blocks]
end

end matrix

section

open lie_algebra.symplectic matrix

-- TODO: Open more sections to eliminate `l` as an explicit argument in most places

variables (l) [decidable_eq l] [fintype l]

def symplectic_group : submonoid (matrix (l ⊕ l) (l ⊕ l)  ℝ) :=
{ carrier := { A | A ⬝ (J l ℝ) ⬝ Aᵀ = J l ℝ},
  mul_mem' :=
  begin
    intros a b ha hb,
    --change (a * b) * (J l ℝ) * (a * b)ᵀ = J l ℝ,
    --change (a) * (J l ℝ) * (a)ᵀ = J l ℝ at ha,
    --change (b) * (J l ℝ) * (b)ᵀ = J l ℝ at hb,
    simp only [mul_eq_mul, set.mem_set_of_eq, transpose_mul] at *,
    rw ←matrix.mul_assoc,
    rw a.mul_assoc,
    rw a.mul_assoc,
    rw hb,
    exact ha,
  end,
  one_mem' := by simp }

variables {l} -- MD: I am making the l implicit whenever we know l from the context already

namespace symplectic_group

lemma mem_symplectic_group_iff {A : matrix (l ⊕ l) (l ⊕ l)  ℝ} :
  A ∈ symplectic_group l ↔ A ⬝ (J l ℝ) ⬝ Aᵀ = J l ℝ :=
by simp [symplectic_group]

instance coe_matrix : has_coe (symplectic_group l) (matrix (l ⊕ l) (l ⊕ l)  ℝ) := ⟨subtype.val⟩

instance coe_fun : has_coe_to_fun (symplectic_group l) (λ _, (l ⊕ l) → (l ⊕ l) → ℝ) :=
{ coe := λ A, A.val }

section coe_lemmas

variables (A B : symplectic_group l)

@[simp] lemma mul_val : ↑(A * B) = A ⬝ B := rfl

@[simp] lemma mul_apply : ⇑(A * B) = (A ⬝ B) := rfl

@[simp] lemma one_val : ↑(1 : symplectic_group l) = (1 : matrix (l ⊕ l) (l ⊕ l)  ℝ) := rfl

@[simp] lemma one_apply : ⇑(1 : symplectic_group l) = (1 : matrix (l ⊕ l) (l ⊕ l)  ℝ) := rfl

lemma mul_mem {A B : matrix (l ⊕ l) (l ⊕ l) ℝ} (hA : A ∈ symplectic_group l) (hB : B ∈ symplectic_group l) :
A ⬝ B ∈ symplectic_group l :=
submonoid.mul_mem _ hA hB

end coe_lemmas

variables (l)

lemma J_mem : (J l ℝ) ∈ symplectic_group l :=
begin
  rw mem_symplectic_group_iff,
  unfold J,
  rw [from_blocks_multiply, from_blocks_transpose, from_blocks_multiply],
  simp,
end

def sym_J : symplectic_group l := ⟨J l ℝ, J_mem l⟩

variables {l}

@[simp] lemma coe_J : ↑(sym_J l) = J l ℝ := rfl

@[simp] lemma J_apply : ⇑(sym_J l) = J l ℝ := rfl

lemma neg_one_transpose : (-1 : matrix l l ℝ)ᵀ = -1 := by rw [transpose_neg, transpose_one]

@[simp] lemma J_transpose : (J l ℝ)ᵀ = - (J l ℝ) :=
begin
  unfold J,
  rw [from_blocks_transpose],
  rw ←from_blocks_neg,
  rw from_blocks_smul,
  rw matrix.transpose_zero,
  rw matrix.transpose_one,
  rw neg_one_transpose,
  simp [from_blocks],
end

-- An old attempt to calculate the determinant by using permutations
-- open sum

-- def bswap (α : Type*) : (α ⊕ α) → (α ⊕ α)
--   | (sum.inl l) := inr l
--   | (sum.inr r) := inl r

-- lemma bswap_bswap {α : Type*} : (bswap α) ∘ (bswap α) = id :=
-- begin
-- ext,
-- cases x;
-- refl
-- end

-- def block_perm : equiv.perm (l ⊕ l) := { to_fun := bswap l,
--   inv_fun := bswap l,
--   left_inv :=
--   begin
--     intro x,
--     change ((bswap l) ∘ (bswap l)) x = x,
--     rw bswap_bswap,
--     refl,
--   end,
--   right_inv :=
--   begin
--     intro x,
--     change ((bswap l) ∘ (bswap l)) x = x,
--     rw bswap_bswap,
--     refl
--   end }

lemma J_squared : (J l ℝ) ⬝ (J l ℝ) = -1 :=
begin
  unfold J,
  rw from_blocks_multiply,
  simp,
  rw ← neg_zero,
  rw ← matrix.from_blocks_neg_2,
  rw ← from_blocks_one,
end

lemma J_inv : (J l ℝ)⁻¹ = -(J l ℝ) :=
begin
  refine matrix.inv_eq_right_inv _,
  rw [matrix.mul_neg, J_squared],
  exact neg_neg 1,
end

lemma neg_one : (-1 : matrix l l ℝ)  = (-1 : ℝ) • 1  := by simp only [neg_smul, one_smul]

lemma minus_powers (n : ℕ) : (-1 : ℝ)^(n + n) = 1 :=
begin
  rw neg_one_pow_eq_one_iff_even,
  exact even_add_self n,
  norm_num,
  /-induction n with n hn,
  simp only [pow_zero],
  calc (-1: ℝ) ^ (n.succ + n.succ) = (-1 : ℝ)^((n + 1) + (n + 1)) : by refl
  ...                              = (-1 : ℝ)^(n + n)*(-1)^2 : by ring_exp
  ...                              = 1 * (-1 : ℝ)^2 : by rw hn
  ...                              = (-1 : ℝ)^2 : by rw one_mul
  ...                              = 1 : by {simp only [neg_one_sq]} -/
end

variables (l)

lemma J_det : det (J l ℝ) = 1 ∨ det (J l ℝ) = - 1:=
begin
  have H : (det (J l ℝ)) * (det (J l ℝ)) = 1 := by {
    rw ← det_mul,
    rw J_squared,
    rw neg_one,
    rw det_smul,
    simp,
    rw minus_powers,
  },
  have H2 : (det (J l ℝ))^2 = 1 := by {
    unfold pow,
    unfold monoid.npow,
    unfold ring.npow,
    unfold comm_ring.npow,
    unfold npow_rec,
    rw mul_one,
    exact H,
  } ,
  rw ←sq_eq_one_iff,
  exact H2,
end

variables {l}

lemma pm_one_unit {S : Type*} [ring S] {x : S} (h : x = 1 ∨ x = -1) : is_unit x :=
begin
  cases h,
  {simp [h],},
  { rw h,
    use -1,
    simp,}
end

lemma J_det_unit : is_unit (det (J l ℝ)) := pm_one_unit (J_det l)

lemma neg_mem {A : matrix (l ⊕ l) (l ⊕ l) ℝ} (h : A ∈ symplectic_group l) :
  -A ∈ symplectic_group l :=
begin
  rw mem_symplectic_group_iff at h ⊢,
  simp [h],
end

instance : has_neg (symplectic_group l) :=
{ neg := λ A, ⟨-A, neg_mem A.2⟩}

@[simp] lemma coe_neg (A : symplectic_group l): (↑(-A) : matrix _ _ _) = -A := rfl

@[simp] lemma neg_apply (A : symplectic_group l): ⇑(-A) = -A := rfl

lemma symplectic_det {A : matrix (l ⊕ l) (l ⊕ l) ℝ} (hA : A ∈ symplectic_group l) :
  is_unit $ det A :=
begin
  rw mem_symplectic_group_iff at hA,
  apply_fun det at hA,
  simp at hA,
  have H := J_det l,
  cases H,
  { rw H at hA,
    simp at hA,
    rw ←sq at hA,
    rw sq_eq_one_iff at hA,
    exact pm_one_unit hA,
  },
  { rw H at hA,
  simp at hA,
  rw ←sq at hA,
  rw sq_eq_one_iff at hA,
  exact pm_one_unit hA,},
end

lemma transpose_mem {A : matrix (l ⊕ l) (l ⊕ l) ℝ} (hA : A ∈ symplectic_group l) : Aᵀ ∈ symplectic_group l :=
begin
  rw mem_symplectic_group_iff at ⊢ hA,
  rw transpose_transpose,
  have huA := symplectic_det hA,
  have huAT : is_unit (Aᵀ).det :=
  begin
    rw matrix.det_transpose,
    exact huA,
  end,

  calc Aᵀ ⬝ J l ℝ ⬝ A = - Aᵀ ⬝ (J l ℝ)⁻¹ ⬝ A  : by {rw J_inv, simp}
  -- ...                 = - Aᵀ ⬝ (J l ℝ)⁻¹ ⬝ A  : by
  ...                 = - Aᵀ ⬝ (A ⬝ J l ℝ ⬝ Aᵀ)⁻¹ ⬝ A : by rw hA
  ...                 = - (Aᵀ ⬝ (Aᵀ⁻¹ ⬝ (J l ℝ)⁻¹)) ⬝ A⁻¹ ⬝ A : by simp [matrix.mul_inv_rev, matrix.mul_assoc]
  ...                 = - (J l ℝ)⁻¹ : by rw [matrix.mul_nonsing_inv_cancel_left _ _ huAT, nonsing_inv_mul_cancel_right _ _ huA]
  ...                 = (J l ℝ) : by simp [J_inv]
end

lemma transpose_mem_iff {A : matrix (l ⊕ l) (l ⊕ l) ℝ} : Aᵀ ∈ symplectic_group l ↔ A ∈ symplectic_group l :=
⟨λ hA, by simpa using transpose_mem hA , transpose_mem⟩

lemma mem_symplectic_group_iff' {A : matrix (l ⊕ l) (l ⊕ l)  ℝ} :
  A ∈ symplectic_group l ↔ Aᵀ ⬝ (J l ℝ) ⬝ A = J l ℝ :=
by rw [←transpose_mem_iff, mem_symplectic_group_iff,transpose_transpose]

-- Things have kind of started following apart starting here

noncomputable def symplectic_inv {A : matrix (l ⊕ l) (l ⊕ l) ℝ} (hA : A ∈ symplectic_group l) :
  symplectic_group l :=
{ val := A⁻¹,
  property :=
  begin
    have huA := symplectic_det hA,
    have huAT : is_unit (Aᵀ).det :=
    begin
      rw matrix.det_transpose,
      exact huA,
    end,
    rw mem_symplectic_group_iff at hA ⊢,
    apply_fun (λ x, A⁻¹ ⬝ (x) ⬝ (Aᵀ)⁻¹) at hA,
    rw matrix.transpose_nonsing_inv,
    calc A⁻¹ ⬝ J l ℝ ⬝ Aᵀ⁻¹ = A⁻¹ ⬝ (A ⬝ J l ℝ ⬝ Aᵀ) ⬝ Aᵀ⁻¹ : by rw hA
    ...                     = A⁻¹ ⬝ A ⬝ J l ℝ ⬝ Aᵀ ⬝ Aᵀ⁻¹ : by simp only [matrix.mul_assoc]
    ...                     = (A⁻¹ ⬝ A) ⬝ (J l ℝ) ⬝ (Aᵀ ⬝ Aᵀ⁻¹) : by simp only [matrix.mul_assoc]
    ...                     = 1 ⬝ (J l ℝ) ⬝ (Aᵀ ⬝ Aᵀ⁻¹) : by rw A.nonsing_inv_mul huA
    ...                     = 1 ⬝ (J l ℝ) ⬝ 1 : by rw Aᵀ.mul_nonsing_inv huAT
    ...                     = J l ℝ : by simp
  end }

def computable_symp_inv {A : matrix (l ⊕ l) (l ⊕ l) ℝ} (hA : A ∈ symplectic_group l) :
  A ⬝ (- (J l ℝ) ⬝ Aᵀ ⬝ (J l ℝ)) = 1 :=
  begin
  rw mem_symplectic_group_iff at hA,
  calc A ⬝ (-J l ℝ ⬝ Aᵀ ⬝ J l ℝ) = - (A ⬝ (J l ℝ) ⬝ Aᵀ) ⬝ (J l ℝ) : by simp [matrix.mul_assoc]
  ...                            = - (J l ℝ) ⬝ (J l ℝ) : by rw hA
  ...                            = (-1 : ℝ) • ( (J l ℝ) ⬝ (J l ℝ) ) : by simp
  ...                            = - (-1) : by {rw J_squared, simp}
  ...                            = 1 : by simp,
  end

-- def computable_symp_inv {A : matrix (l ⊕ l) (l ⊕ l) ℝ} (hA : A ∈ symplectic_group l) :
--   symplectic_group l :=
-- { val := - (J l ℝ) ⬝ Aᵀ ⬝ (J l ℝ),
--   property :=
--   begin
--     rw mem_symplectic_group_iff at hA ⊢,
--     simp
--     sorry
--   end }

lemma inv_mem_aux {A : matrix (l ⊕ l) (l ⊕ l) ℝ} (hA : A ∈ symplectic_group l) :
  - (J l ℝ) ⬝ Aᵀ ⬝ (J l ℝ) ∈ symplectic_group l :=
mul_mem (mul_mem (neg_mem $ J_mem _) $ transpose_mem hA) $ J_mem _
  -- simp only [matrix.neg_mul, transpose_neg, transpose_mul, J_transpose, transpose_transpose, matrix.mul_neg, neg_neg],
  -- rw matrix.mul_assoc _ _ (J l ℝ),
  -- rw J_squared,
  -- rw matrix.transpose_mul,
  -- rw matrix.transpose_mul,
  -- rw matrix.transpose_transpose,
  -- rw matrix.transpose_neg,

lemma inv_left_mul_aux {A : matrix (l ⊕ l) (l ⊕ l) ℝ} (hA : A ∈ symplectic_group l) :
  -(J l ℝ ⬝ Aᵀ ⬝ J l ℝ ⬝ A) = 1 :=
calc -(J l ℝ ⬝ Aᵀ ⬝ J l ℝ ⬝ A) = - J l ℝ ⬝ (Aᵀ ⬝ J l ℝ ⬝ A) : by simp [matrix.mul_assoc]
...                            = - J l ℝ ⬝ J l ℝ : by {rw mem_symplectic_group_iff' at hA, rw hA}
...                            = (-1 : ℝ) • (J l ℝ ⬝ J l ℝ) : by simp only [matrix.neg_mul, neg_smul, one_smul]
...                            = (-1 : ℝ) • -1 : by rw J_squared
...                            = 1 : by simp only [neg_smul_neg, one_smul]

instance : has_inv (symplectic_group l) := {
  inv := λ A, ⟨- (J l ℝ) ⬝ Aᵀ ⬝ (J l ℝ),
    mul_mem (mul_mem (neg_mem $ J_mem _) $ transpose_mem A.2) $ J_mem _⟩,
}

@[simp] lemma coe_inv (A : symplectic_group l): (↑(A⁻¹) : matrix _ _ _) = - (J l ℝ) ⬝ Aᵀ ⬝ (J l ℝ) := rfl

@[simp] lemma inv_apply (A : symplectic_group l): ⇑(A⁻¹) = - (J l ℝ) ⬝ Aᵀ ⬝ (J l ℝ) := rfl

instance : group (symplectic_group l) := {
  mul_left_inv :=
  begin
    intro A,
    apply subtype.ext,
    simp,
    exact inv_left_mul_aux A.2,
  end,
  .. symplectic_group.has_inv,
  .. submonoid.to_monoid _
}

-- I think at this point I'm starting to realize I shouldn't be using `A ∈ symplectic l`...
noncomputable instance old : group (symplectic_group l) := {
  inv := λ A, symplectic_inv A.2,
  mul_left_inv :=
  begin
  intro A,
  apply subtype.ext,
  simp,
  -- rw matrix.nonsing_inv_mul,
  -- unfold has_inv.inv,
  -- unfold div_inv_monoid.inv,
  -- unfold symplectic_inv,

  -- Not sure how to deal with this `⟨(↑A)⁻¹, _⟩ * A = 1`
  -- unfold matrix.mul,
  show_term {change (A)⁻¹ ⬝ A = 1},
  rw matrix.nonsing_inv_mul A (symplectic_det A.prop),
  end,
  .. submonoid.to_monoid _ }

end symplectic_group

end

-- TODO: Add this back in
-- def symplectic_transpose : symplectic l → symplectic l := fun A,
-- { val := Aᵀ,
--   property :=
--   begin
--    sorry
--   end }


/-
Copyright (c) 2021 Arthur Paulino. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Arthur Paulino, Kyle Miller

! This file was ported from Lean 3 source module combinatorics.simple_graph.partition
! leanprover-community/mathlib commit 2303b3e299f1c75b07bceaaac130ce23044d1386
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
import Mathlib.Combinatorics.SimpleGraph.Coloring

/-!
# Graph partitions

This module provides an interface for dealing with partitions on simple graphs. A partition of
a graph `G`, with vertices `V`, is a set `P` of disjoint nonempty subsets of `V` such that:

* The union of the subsets in `P` is `V`.

* Each element of `P` is an independent set. (Each subset contains no pair of adjacent vertices.)

Graph partitions are graph colorings that do not name their colors.  They are adjoint in the
following sense. Given a graph coloring, there is an associated partition from the set of color
classes, and given a partition, there is an associated graph coloring from using the partition's
subsets as colors.  Going from graph colorings to partitions and back makes a coloring "canonical":
all colors are given a canonical name and unused colors are removed.  Going from partitions to
graph colorings and back is the identity.

## Main definitions

* `SimpleGraph.Partition` is a structure to represent a partition of a simple graph

* `SimpleGraph.Partition.PartsCardLe` is whether a given partition is an `n`-partition.
  (a partition with at most `n` parts).

* `SimpleGraph.Partitionable n` is whether a given graph is `n`-partite

* `SimpleGraph.Partition.toColoring` creates colorings from partitions

* `SimpleGraph.Coloring.toPartition` creates partitions from colorings

## Main statements

* `SimpleGraph.partitionable_iff_colorable` is that `n`-partitionability and
  `n`-colorability are equivalent.

-/

universe u v

namespace SimpleGraph

variable {V : Type u} (G : SimpleGraph V)

/-- A `Partition` of a simple graph `G` is a structure constituted by
* `parts`: a set of subsets of the vertices `V` of `G`
* `isPartition`: a proof that `parts` is a proper partition of `V`
* `independent`: a proof that each element of `parts` doesn't have a pair of adjacent vertices
-/
structure Partition where
  /-- `parts`: a set of subsets of the vertices `V` of `G`. -/
  parts : Set (Set V)
  /-- `isPartition`: a proof that `parts` is a proper partition of `V`. -/
  isPartition : Setoid.IsPartition parts
  /-- `independent`: a proof that each element of `parts` doesn't have a pair of adjacent vertices.
-/
  independent : ∀ s ∈ parts, IsAntichain G.Adj s
#align simple_graph.partition SimpleGraph.Partition

/-- Whether a partition `P` has at most `n` parts. A graph with a partition
satisfying this predicate called `n`-partite. (See `SimpleGraph.Partitionable`.) -/
def Partition.PartsCardLe {G : SimpleGraph V} (P : G.Partition) (n : ℕ) : Prop :=
  ∃ h : P.parts.Finite, h.toFinset.card ≤ n
#align simple_graph.partition.parts_card_le SimpleGraph.Partition.PartsCardLe

/-- Whether a graph is `n`-partite, which is whether its vertex set
can be partitioned in at most `n` independent sets. -/
def Partitionable (n : ℕ) : Prop := ∃ P : G.Partition, P.PartsCardLe n
#align simple_graph.partitionable SimpleGraph.Partitionable

namespace Partition

variable {G} (P : G.Partition)

/-- The part in the partition that `v` belongs to -/
def partOfVertex (v : V) : Set V := Classical.choose (P.isPartition.2 v)
#align simple_graph.partition.part_of_vertex SimpleGraph.Partition.partOfVertex

theorem partOfVertex_mem (v : V) : P.partOfVertex v ∈ P.parts := by
  obtain ⟨h, -⟩ := (P.isPartition.2 v).choose_spec.1
  exact h
#align simple_graph.partition.part_of_vertex_mem SimpleGraph.Partition.partOfVertex_mem

theorem mem_partOfVertex (v : V) : v ∈ P.partOfVertex v := by
  obtain ⟨⟨h1, h2⟩, _h3⟩ := (P.isPartition.2 v).choose_spec
  exact h2.1
#align simple_graph.partition.mem_part_of_vertex SimpleGraph.Partition.mem_partOfVertex

theorem partOfVertex_ne_of_adj {v w : V} (h : G.Adj v w) : P.partOfVertex v ≠ P.partOfVertex w := by
  intro hn
  have hw := P.mem_partOfVertex w
  rw [← hn] at hw
  exact P.independent _ (P.partOfVertex_mem v) (P.mem_partOfVertex v) hw (G.ne_of_adj h) h
#align simple_graph.partition.part_of_vertex_ne_of_adj SimpleGraph.Partition.partOfVertex_ne_of_adj

/-- Create a coloring using the parts themselves as the colors.
Each vertex is colored by the part it's contained in. -/
def toColoring : G.Coloring P.parts :=
  Coloring.mk (fun v ↦ ⟨P.partOfVertex v, P.partOfVertex_mem v⟩) fun hvw ↦
    by
    rw [Ne.def, Subtype.mk_eq_mk]
    exact P.partOfVertex_ne_of_adj hvw
#align simple_graph.partition.to_coloring SimpleGraph.Partition.toColoring

/-- Like `SimpleGraph.Partition.toColoring` but uses `Set V` as the coloring type. -/
def toColoring' : G.Coloring (Set V) :=
  Coloring.mk P.partOfVertex fun hvw ↦ P.partOfVertex_ne_of_adj hvw
#align simple_graph.partition.to_coloring' SimpleGraph.Partition.toColoring'

theorem to_colorable [Fintype P.parts] : G.Colorable (Fintype.card P.parts) :=
  P.toColoring.to_colorable
#align simple_graph.partition.to_colorable SimpleGraph.Partition.to_colorable

end Partition

variable {G}

/-- Creates a partition from a coloring. -/
@[simps]
def Coloring.toPartition {α : Type v} (C : G.Coloring α) : G.Partition
    where
  parts := C.colorClasses
  isPartition := C.colorClasses_isPartition
  independent := by
    rintro s ⟨c, rfl⟩
    apply C.color_classes_independent
#align simple_graph.coloring.to_partition SimpleGraph.Coloring.toPartition

/-- The partition where every vertex is in its own part. -/
@[simps]
instance : Inhabited (Partition G) := ⟨G.selfColoring.toPartition⟩

theorem partitionable_iff_colorable {n : ℕ} : G.Partitionable n ↔ G.Colorable n := by
  constructor
  · rintro ⟨P, hf, hc⟩
    have : Fintype P.parts := hf.fintype
    rw [Set.Finite.card_toFinset hf] at hc
    apply P.to_colorable.mono hc
  · rintro ⟨C⟩
    refine' ⟨C.toPartition, C.colorClasses_finite, le_trans _ (Fintype.card_fin n).le⟩
    generalize_proofs h
    haveI : Fintype C.colorClasses := C.colorClasses_finite.fintype
    rw [h.card_toFinset]
    exact C.card_colorClasses_le
#align simple_graph.partitionable_iff_colorable SimpleGraph.partitionable_iff_colorable

end SimpleGraph


/-
Copyright (c) 2022 Michael Stoll. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Michael Stoll
-/
import number_theory.legendre_symbol.add_character
import number_theory.legendre_symbol.zmod_char
import algebra.char_p.char_and_card

/-!
# Gauss sums

We define the Gauss sum associated to a multiplicative and an additive
character of a finite field and prove some results about them.

## Main definition

Let `R` be a finite commutative ring and let `R'` be another commutative ring.
If `χ` is a multiplicative character `R → R'` (type `mul_char R R'`) and `ψ`
is an additive character `R → R'` (type `add_char R R'`, which abbreviates
`(multiplicative R) →* R'`), then the *Gauss sum* of `χ` and `ψ` is `∑ a, χ a * ψ a`.

## Main results

Some important results are as follows.

* `gauss_sum_mul_gauss_sum_eq_card`: The product of the Gauss
  sums of `χ` and `ψ` and that of `χ⁻¹` and `ψ⁻¹` is the cardinality
  of the source ring `R` (if `χ` is nontrivial, `ψ` is primitive and `R` is a field).
* `gauss_sum_sq`: The square of the Gauss sum is `χ(-1)` times
  the cardinality of `R` if in addition `χ` is a quadratic character.
* `quad_gauss_sum_frob`: For a quadratic character `χ`, raising
  the Gauss sum to the `p`th power (where `p` is the characteristic of
  the target ring `R'`) multiplies it by `χ p`.
* `char.card_pow_card`: When `F` and `F'` are finite fields and `χ : F → F'`
  is a nontrivial quadratic character, then `(χ (-1) * #F)^(#F'/2) = χ (#F')`.
* `finite_field.two_pow_card`: For every finite field `F` of odd characteristic,
  we have `2^(#F/2) = χ₈(#F)` in `F`.

This machinery can be used to derive (a generalization of) the Law of
Quadratic Reciprocity.

## Tags

additive character, multiplicative character, Gauss sum
-/

universes u v

open_locale big_operators

open add_char mul_char

section gauss_sum_def

-- `R` is the domain of the characters
variables {R : Type u} [comm_ring R] [fintype R]
-- `R'` is the target of the characters
variables {R' : Type v} [comm_ring R']

/-!
### Definition and first properties
-/

/-- Definition of the Gauss sum associated to a multiplicative and an additive character. -/
def gauss_sum (χ : mul_char R R') (ψ : add_char R R') : R' := ∑ a, χ a * ψ a

/-- Replacing `ψ` by `mul_shift ψ a` and multiplying the Gauss sum by `χ a` does not change it. -/
lemma gauss_sum_mul_shift (χ : mul_char R R') (ψ : add_char R R') (a : Rˣ) :
  χ a * gauss_sum χ (mul_shift ψ a) = gauss_sum χ ψ :=
begin
  simp only [gauss_sum, mul_shift_apply, finset.mul_sum],
  simp_rw [← mul_assoc, ← map_mul],
  exact fintype.sum_bijective _ a.mul_left_bijective _ _ (λ x, rfl),
end

end gauss_sum_def

/-!
### The product of two Gauss sums
-/

section gauss_sum_prod

-- In the following, we need `R` to be a finite field and `R'` to be a domain.
variables {R : Type u} [field R] [fintype R] {R' : Type v} [comm_ring R'] [is_domain R']

-- A helper lemma for `gauss_sum_mul_gauss_sum_eq_card` below
-- Is this useful enough in other contexts to be public?
private
lemma gauss_sum_mul_aux {χ : mul_char R R'} (hχ : is_nontrivial χ) (ψ : add_char R R') (b : R) :
  ∑ a, χ (a * b⁻¹) * ψ (a - b) = ∑ c, χ c * ψ (b * (c - 1)) :=
begin
  cases eq_or_ne b 0 with hb hb,
  { -- case `b = 0`
    simp only [hb, inv_zero, mul_zero, mul_char.map_zero, zero_mul, finset.sum_const_zero,
               map_zero_one, mul_one],
    exact hχ.sum_eq_zero.symm, },
  { -- case `b ≠ 0`
    refine (fintype.sum_bijective _ (mul_left_bijective₀ b hb) _ _ $ λ x, _).symm,
    rw [mul_assoc, mul_comm x, ← mul_assoc, mul_inv_cancel hb, one_mul, mul_sub, mul_one] },
end

/-- We have `gauss_sum χ ψ * gauss_sum χ⁻¹ ψ⁻¹ = fintype.card R`
when `χ` is nontrivial and `ψ` is primitive (and `R` is a field). -/
lemma gauss_sum_mul_gauss_sum_eq_card  {χ : mul_char R R'} (hχ : is_nontrivial χ)
  {ψ : add_char R R'} (hψ : is_primitive ψ) :
  gauss_sum χ ψ * gauss_sum χ⁻¹ ψ⁻¹ = fintype.card R :=
begin
  simp only [gauss_sum, add_char.inv_apply, finset.sum_mul, finset.mul_sum, mul_char.inv_apply'],
  conv in (_ * _ * (_ * _))
    { rw [mul_mul_mul_comm, ← map_mul, ← map_add_mul, ← sub_eq_add_neg], },
  simp_rw gauss_sum_mul_aux hχ ψ,
  rw [finset.sum_comm],
  classical, -- to get `[decidable_eq R]` for `sum_mul_shift`
  simp_rw [← finset.mul_sum, sum_mul_shift _ hψ, sub_eq_zero, mul_ite, mul_zero],
  rw [finset.sum_ite_eq' finset.univ (1 : R)],
  simp only [finset.mem_univ, map_one, one_mul, if_true],
end

/-- When `χ` is a nontrivial quadratic character, then the square of `gauss_sum χ ψ`
is `χ(-1)` times the cardinality of `R`. -/
lemma gauss_sum_sq  {χ : mul_char R R'} (hχ₁ : is_nontrivial χ) (hχ₂ : is_quadratic χ)
  {ψ : add_char R R'} (hψ : is_primitive ψ) :
  (gauss_sum χ ψ) ^ 2 = χ (-1) * fintype.card R :=
begin
  rw [pow_two, ← gauss_sum_mul_gauss_sum_eq_card hχ₁ hψ, hχ₂.inv, mul_rotate'],
  congr,
  rw [mul_comm, ← gauss_sum_mul_shift _ _ (-1 : Rˣ), inv_mul_shift],
  refl,
end

end gauss_sum_prod

/-!
### Gauss sums and Frobenius
-/

section gauss_sum_frob

variables {R : Type u} [comm_ring R] [fintype R] {R' : Type v} [comm_ring R']
-- We assume that the target ring `R'` has prime characteristic `p`.
variables (p : ℕ) [fp : fact p.prime] [hch : char_p R' p]
include fp hch

/-- When `R'` has prime characteristic `p`, then the `p`th power of the Gauss sum
of `χ` and `ψ` is the Gauss sum of `χ^p` and `ψ^p`. -/
lemma gauss_sum_frob (χ : mul_char R R') (ψ : add_char R R') :
  gauss_sum χ ψ ^ p = gauss_sum (χ ^ p) (ψ ^ p) :=
begin
  rw [← frobenius_def, gauss_sum, gauss_sum, map_sum],
  simp_rw [pow_apply' χ fp.1.pos, map_mul, frobenius_def],
  refl,
end

/-- For a quadratic character `χ` and when the characteristic `p` of the target ring
is a unit in the source ring, the `p`th power of the Gauss sum of`χ` and `ψ` is
`χ p` times the original Gauss sum. -/
lemma mul_char.is_quadratic.gauss_sum_frob (hp : is_unit (p : R)) {χ : mul_char R R'}
  (hχ : is_quadratic χ) (ψ : add_char R R') :
  gauss_sum χ ψ ^ p = χ p * gauss_sum χ ψ :=
by rw [gauss_sum_frob, pow_mul_shift, hχ.pow_char p, ← gauss_sum_mul_shift χ ψ hp.unit,
       ← mul_assoc, hp.unit_spec, ← pow_two, ← pow_apply' _ (by norm_num : 0 < 2),
       hχ.sq_eq_one, ← hp.unit_spec, one_apply_coe, one_mul]

/-- For a quadratic character `χ` and when the characteristic `p` of the target ring
is a unit in the source ring and `n` is a natural number, the `p^n`th power of the Gauss
sum of`χ` and `ψ` is `χ (p^n)` times the original Gauss sum. -/
lemma mul_char.is_quadratic.gauss_sum_frob_iter (n : ℕ) (hp : is_unit (p : R))
  {χ : mul_char R R'} (hχ : is_quadratic χ) (ψ : add_char R R') :
  gauss_sum χ ψ ^ (p ^ n) = χ (p ^ n) * gauss_sum χ ψ :=
begin
  induction n with n ih,
  { rw [pow_zero, pow_one, pow_zero, mul_char.map_one, one_mul], },
  { rw [pow_succ, mul_comm p, pow_mul, ih, mul_pow, hχ.gauss_sum_frob _ hp,
        ← mul_assoc, pow_succ, mul_comm (p : R), map_mul,
        ← pow_apply' χ fp.1.pos (p ^ n), hχ.pow_char p], },
end

end gauss_sum_frob

/-!
### Values of quadratic characters
-/

section gauss_sum_values

variables {R : Type u} [comm_ring R] [fintype R] {R' : Type v} [comm_ring R'] [is_domain R']

/-- If the square of the Gauss sum of a quadratic character is `χ(-1) * #R`,
then we get, for all `n : ℕ`, the relation `(χ(-1) * #R) ^ (p^n/2) = χ(p^n)`,
where `p` is the (odd) characteristic of the target ring `R'`.
This version can be used when `R` is not a field, e.g., `ℤ/8ℤ`. -/
lemma char.card_pow_char_pow {χ : mul_char R R'} (hχ : is_quadratic χ) (ψ : add_char R R') (p n : ℕ)
  [fp : fact p.prime] [hch : char_p R' p] (hp : is_unit (p : R)) (hp' : p ≠ 2)
  (hg : (gauss_sum χ ψ) ^ 2 = χ (-1) * fintype.card R) :
  (χ (-1) * fintype.card R) ^ (p ^ n / 2) = χ (p ^ n) :=
begin
  have : gauss_sum χ ψ ≠ 0,
  { intro hf, rw [hf, zero_pow (by norm_num : 0 < 2), eq_comm, mul_eq_zero] at hg,
    exact not_is_unit_prime_of_dvd_card p
      ((char_p.cast_eq_zero_iff R' p _).mp $ hg.resolve_left (is_unit_one.neg.map χ).ne_zero) hp },
  rw ← hg, apply mul_right_cancel₀ this,
  rw [← hχ.gauss_sum_frob_iter p n hp ψ, ← pow_mul, mul_comm, ← pow_succ,
      nat.two_mul_div_two_add_one_of_odd ((fp.1.eq_two_or_odd').resolve_left hp').pow],
end

/-- When `F` and `F'` are finite fields and `χ : F → F'` is a nontrivial quadratic character,
then `(χ(-1) * #F)^(#F'/2) = χ(#F')`. -/
lemma char.card_pow_card {F : Type*} [field F] [fintype F] {F' : Type*} [field F'] [fintype F']
  {χ : mul_char F F'} (hχ₁ : is_nontrivial χ) (hχ₂ : is_quadratic χ)
  (hch₁ : ring_char F' ≠ ring_char F) (hch₂ : ring_char F' ≠ 2) :
  (χ (-1) * fintype.card F) ^ (fintype.card F' / 2) = χ (fintype.card F') :=
begin
  obtain ⟨n, hp, hc⟩ := finite_field.card F (ring_char F),
  obtain ⟨n', hp', hc'⟩ := finite_field.card F' (ring_char F'),
  let ψ := primitive_char_finite_field F F' hch₁,
  let FF' := cyclotomic_field ψ.n F',
  have hchar := algebra.ring_char_eq F' FF',
  apply (algebra_map F' FF').injective,
  rw [map_pow, map_mul, map_nat_cast, hc', hchar, nat.cast_pow],
  simp only [← mul_char.ring_hom_comp_apply],
  haveI := fact.mk hp',
  haveI := fact.mk (hchar.subst hp'),
  rw [ne, ← nat.prime_dvd_prime_iff_eq hp' hp, ← is_unit_iff_not_dvd_char, hchar] at hch₁,
  exact char.card_pow_char_pow (hχ₂.comp _) ψ.char (ring_char FF') n' hch₁ (hchar ▸ hch₂)
    (gauss_sum_sq (hχ₁.comp $ ring_hom.injective _) (hχ₂.comp _) ψ.prim),
end

end gauss_sum_values

section gauss_sum_two

/-!
### The quadratic character of 2

This section proves the following result.

For every finite field `F` of odd characteristic, we have `2^(#F/2) = χ₈(#F)` in `F`.
This can be used to show that the quadratic character of `F` takes the value
`χ₈(#F)` at `2`.

The proof uses the Gauss sum of `χ₈` and a primitive additive character on `ℤ/8ℤ`;
in this way, the result is reduced to `card_pow_char_pow`.
-/

open zmod

/-- For every finite field `F` of odd characteristic, we have `2^(#F/2) = χ₈(#F)` in `F`. -/
lemma finite_field.two_pow_card {F : Type*} [fintype F] [field F] (hF : ring_char F ≠ 2) :
  (2 : F) ^ (fintype.card F / 2) = χ₈ (fintype.card F) :=
begin
  have hp2 : ∀ (n : ℕ), (2 ^ n : F) ≠ 0 := λ n, pow_ne_zero n (ring.two_ne_zero hF),
  obtain ⟨n, hp, hc⟩ := finite_field.card F (ring_char F),

  -- we work in `FF`, the eighth cyclotomic field extension of `F`
  let FF := (polynomial.cyclotomic 8 F).splitting_field,
  haveI : finite_dimensional F FF :=
    polynomial.is_splitting_field.finite_dimensional FF (polynomial.cyclotomic 8 F),
  haveI : fintype FF := finite_dimensional.fintype_of_fintype F FF,
  have hchar := algebra.ring_char_eq F FF,
  have FFp := hchar.subst hp,
  haveI := fact.mk FFp,
  have hFF := ne_of_eq_of_ne hchar.symm hF, -- `ring_char FF ≠ 2`
  have hu : is_unit (ring_char FF : zmod 8),
  { rw [is_unit_iff_not_dvd_char, ring_char_zmod_n],
    rw [ne, ← nat.prime_dvd_prime_iff_eq FFp nat.prime_two] at hFF,
    change ¬ _ ∣ 2 ^ 3,
    exact mt FFp.dvd_of_dvd_pow hFF },

  -- there is a primitive additive character `ℤ/8ℤ → FF`, sending `a + 8ℤ ↦ τ^a`
  -- with a primitive eighth root of unity `τ`
  let ψ₈ := primitive_zmod_char 8 F (by convert hp2 3; norm_num),
  let τ : FF := ψ₈.char 1,
  have τ_spec : τ ^ 4 = -1,
  { refine (sq_eq_one_iff.1 _).resolve_left _;
    { simp only [τ, ← map_nsmul_pow],
      erw add_char.is_primitive.zmod_char_eq_one_iff 8 ψ₈.prim,
      dec_trivial } },

  -- we consider `χ₈` as a multiplicative character `ℤ/8ℤ → FF`
  let χ := χ₈.ring_hom_comp (int.cast_ring_hom FF),
  have hχ : χ (-1) = 1 := norm_num.int_cast_one,
  have hq : is_quadratic χ := is_quadratic_χ₈.comp _,

  -- we now show that the Gauss sum of `χ` and `ψ₈` has the relevant property
  have hg : gauss_sum χ ψ₈.char ^ 2 = χ (-1) * fintype.card (zmod 8),
  { rw [hχ, one_mul, card, gauss_sum],
    convert ← congr_arg (^ 2) (fin.sum_univ_eight $ λ x, (χ₈ x : FF) * τ ^ x.val),
    { ext, congr, apply pow_one },
    convert_to (0 + 1 * τ ^ 1 + 0 + (-1) * τ ^ 3 + 0 + (-1) * τ ^ 5 + 0 + 1 * τ ^ 7) ^ 2 = _,
    { simp only [χ₈_apply, matrix.cons_val_zero, matrix.cons_val_one, matrix.head_cons,
        matrix.cons_vec_bit0_eq_alt0, matrix.cons_vec_bit1_eq_alt1, matrix.cons_vec_append,
        matrix.cons_vec_alt0, matrix.cons_vec_alt1, int.cast_zero, int.cast_one, int.cast_neg,
        zero_mul], refl },
    convert_to 8 + (τ ^ 4 + 1) * (τ ^ 10 - 2 * τ ^ 8 - 2 * τ ^ 6 + 6 * τ ^ 4 + τ ^ 2 - 8) = _,
    { ring }, { rw τ_spec, norm_num } },

  -- this allows us to apply `card_pow_char_pow` to our situation
  have h := char.card_pow_char_pow hq ψ₈.char (ring_char FF) n hu hFF hg,
  rw [card, ← hchar, hχ, one_mul, ← hc, ← nat.cast_pow (ring_char F), ← hc] at h,

  -- finally, we change `2` to `8` on the left hand side
  convert_to (8 : F) ^ (fintype.card F / 2) = _,
  { rw [(by norm_num : (8 : F) = 2 ^ 2 * 2), mul_pow,
      (finite_field.is_square_iff hF $ hp2 2).mp ⟨2, pow_two 2⟩, one_mul] },
  apply (algebra_map F FF).injective,
  simp only [map_pow, map_bit0, map_one, map_int_cast],
  convert h, norm_num,
end

end gauss_sum_two


import data.real.basic
import tactic

variables x y : ℝ

-- Showing examples of how ring works

example : (x + y)^2 = x^2 + 2 * x * y + y^2 :=
begin
  ring,
end

example : (x + y)^3 = x^3 + 3 * x^2 * y + 3 * x * y^2  + y^3 :=
begin
  ring,
end

-- If didn't use the variables decleared first
example {a : ℝ} {b : ℝ} : (a + b)^2 = a^2 + 2 * a * b + b^2 :=
begin
  ring,
end