/-
Copyright (c) 2021 Andrew Yang. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Andrew Yang
-/
import category_theory.limits.filtered_colimit_commutes_finite_limit
import category_theory.limits.preserves.functor_category
import category_theory.limits.preserves.shapes.equalizers
import category_theory.limits.bicones
import category_theory.limits.comma
import category_theory.limits.preserves.finite
import category_theory.limits.shapes.finite_limits

/-!
# Representably flat functors

We define representably flat functors as functors such that the category of structured arrows
over `X` is cofiltered for each `X`. This concept is also known as flat functors as in [Elephant]
Remark C2.3.7, and this name is suggested by Mike Shulman in
https://golem.ph.utexas.edu/category/2011/06/flat_functors_and_morphisms_of.html to avoid
confusion with other notions of flatness.

This definition is equivalent to left exact functors (functors that preserves finite limits) when
`C` has all finite limits.

## Main results

* `flat_of_preserves_finite_limits`: If `F : C ⥤ D` preserves finite limits and `C` has all finite
  limits, then `F` is flat.
* `preserves_finite_limits_of_flat`: If `F : C ⥤ D` is flat, then it preserves all finite limits.
* `preserves_finite_limits_iff_flat`: If `C` has all finite limits,
  then `F` is flat iff `F` is left_exact.
* `Lan_preserves_finite_limits_of_flat`: If `F : C ⥤ D` is a flat functor between small categories,
  then the functor `Lan F.op` between presheaves of sets preserves all finite limits.
* `flat_iff_Lan_flat`: If `C`, `D` are small and `C` has all finite limits, then `F` is flat iff
  `Lan F.op : (Cᵒᵖ ⥤ Type*) ⥤ (Dᵒᵖ ⥤ Type*)` is flat.
* `preserves_finite_limits_iff_Lan_preserves_finite_limits`: If `C`, `D` are small and `C` has all
  finite limits, then `F` preserves finite limits iff `Lan F.op : (Cᵒᵖ ⥤ Type*) ⥤ (Dᵒᵖ ⥤ Type*)`
  does.

-/

universes v₁ v₂ v₃ u₁ u₂ u₃

open category_theory
open category_theory.limits
open opposite

namespace category_theory

namespace structured_arrow_cone
open structured_arrow
variables {C : Type u₁} [category.{v₁} C] {D : Type u₂} [category.{v₁} D]
variables {J : Type v₁} [small_category J]
variables {K : J ⥤ C} (F : C ⥤ D) (c : cone K)

/--
Given a cone `c : cone K` and a map `f : X ⟶ c.X`, we can construct a cone of structured
arrows over `X` with `f` as the cone point. This is the underlying diagram.
-/
@[simps]
def to_diagram : J ⥤ structured_arrow c.X K :=
{ obj := λ j, structured_arrow.mk (c.π.app j),
  map := λ j k g, structured_arrow.hom_mk g (by simpa) }

/-- Given a diagram of `structured_arrow X F`s, we may obtain a cone with cone point `X`. -/
@[simps]
def diagram_to_cone {X : D} (G : J ⥤ structured_arrow X F) : cone (G ⋙ proj X F ⋙ F) :=
{ X := X, π := { app := λ j, (G.obj j).hom } }

/--
Given a cone `c : cone K` and a map `f : X ⟶ F.obj c.X`, we can construct a cone of structured
arrows over `X` with `f` as the cone point.
-/
@[simps]
def to_cone {X : D} (f : X ⟶ F.obj c.X) :
  cone (to_diagram (F.map_cone c) ⋙ map f ⋙ pre _ K F) :=
{ X := mk f, π := { app := λ j, hom_mk (c.π.app j) rfl,
                    naturality' := λ j k g, by { ext, dsimp, simp } } }

end structured_arrow_cone

section representably_flat
variables {C : Type u₁} [category.{v₁} C] {D : Type u₂} [category.{v₂} D]

/--
A functor `F : C ⥤ D` is representably-flat functor if the comma category `(X/F)`
is cofiltered for each `X : C`.
-/
class representably_flat (F : C ⥤ D) : Prop :=
(cofiltered : ∀ (X : D), is_cofiltered (structured_arrow X F))

attribute [instance] representably_flat.cofiltered

end representably_flat

section has_limit
variables {C : Type u₁} [category.{v₁} C] {D : Type u₂} [category.{v₁} D]

@[priority 100]
instance cofiltered_of_has_finite_limits [has_finite_limits C] : is_cofiltered C :=
{ cocone_objs := λ A B, ⟨limits.prod A B, limits.prod.fst, limits.prod.snd, trivial⟩,
  cocone_maps :=  λ A B f g, ⟨equalizer f g, equalizer.ι f g, equalizer.condition f g⟩,
  nonempty := ⟨⊤_ C⟩ }

lemma flat_of_preserves_finite_limits [has_finite_limits C] (F : C ⥤ D)
  [preserves_finite_limits F] : representably_flat F := ⟨λ X,
begin
  haveI : has_finite_limits (structured_arrow X F) :=
    { out := λ J _ _, by { resetI, apply_instance } },
  apply_instance
end⟩

namespace preserves_finite_limits_of_flat
open structured_arrow
open structured_arrow_cone
variables {J : Type v₁} [small_category J] [fin_category J] {K : J ⥤ C}
variables (F : C ⥤ D) [representably_flat F] {c : cone K} (hc : is_limit c) (s : cone (K ⋙ F))
include hc

/--
(Implementation).
Given a limit cone `c : cone K` and a cone `s : cone (K ⋙ F)` with `F` representably flat,
`s` can factor through `F.map_cone c`.
-/
noncomputable def lift : s.X ⟶ F.obj c.X :=
let s' := is_cofiltered.cone (to_diagram s ⋙ structured_arrow.pre _ K F) in
s'.X.hom ≫ (F.map $ hc.lift $
  (cones.postcompose ({ app := λ X, 𝟙 _, naturality' := by simp }
      : (to_diagram s ⋙ pre s.X K F) ⋙ proj s.X F ⟶ K)).obj $
  (structured_arrow.proj s.X F).map_cone s')

lemma fac (x : J) : lift F hc s ≫ (F.map_cone c).π.app x = s.π.app x :=
by simpa [lift, ←functor.map_comp]

lemma uniq {K : J ⥤ C} {c : cone K} (hc : is_limit c)
  (s : cone (K ⋙ F)) (f₁ f₂ : s.X ⟶ F.obj c.X)
  (h₁ : ∀ (j : J), f₁ ≫ (F.map_cone c).π.app j = s.π.app j)
  (h₂ : ∀ (j : J), f₂ ≫ (F.map_cone c).π.app j = s.π.app j) : f₁ = f₂ :=
begin
  -- We can make two cones over the diagram of `s` via `f₁` and `f₂`.
  let α₁ : to_diagram (F.map_cone c) ⋙ map f₁ ⟶ to_diagram s :=
  { app := λ X, eq_to_hom (by simp [←h₁]), naturality' := λ _ _ _, by { ext, simp } },
  let α₂ : to_diagram (F.map_cone c) ⋙ map f₂ ⟶ to_diagram s :=
  { app := λ X, eq_to_hom (by simp [←h₂]), naturality' := λ _ _ _, by { ext, simp } },
  let c₁ : cone (to_diagram s ⋙ pre s.X K F) :=
    (cones.postcompose (whisker_right α₁ (pre s.X K F) : _)).obj (to_cone F c f₁),
  let c₂ : cone (to_diagram s ⋙ pre s.X K F) :=
    (cones.postcompose (whisker_right α₂ (pre s.X K F) : _)).obj (to_cone F c f₂),

  -- The two cones can then be combined and we may obtain a cone over the two cones since
  -- `structured_arrow s.X F` is cofiltered.
  let c₀ := is_cofiltered.cone (bicone_mk _ c₁ c₂),
  let g₁ : c₀.X ⟶ c₁.X := c₀.π.app (bicone.left),
  let g₂ : c₀.X ⟶ c₂.X := c₀.π.app (bicone.right),

  -- Then `g₁.right` and `g₂.right` are two maps from the same cone into the `c`.
  have : ∀ (j : J), g₁.right ≫ c.π.app j = g₂.right ≫ c.π.app j,
  { intro j,
    injection c₀.π.naturality (bicone_hom.left  j) with _ e₁,
    injection c₀.π.naturality (bicone_hom.right j) with _ e₂,
    simpa using e₁.symm.trans e₂ },
  have : c.extend g₁.right = c.extend g₂.right,
  { unfold cone.extend, congr' 1, ext x, apply this },

  -- And thus they are equal as `c` is the limit.
  have : g₁.right = g₂.right,
  calc g₁.right = hc.lift (c.extend g₁.right) : by { apply hc.uniq (c.extend _), tidy }
            ... = hc.lift (c.extend g₂.right) : by { congr, exact this }
            ... = g₂.right                    : by { symmetry, apply hc.uniq (c.extend _), tidy },

  -- Finally, since `fᵢ` factors through `F(gᵢ)`, the result follows.
  calc f₁ = 𝟙 _ ≫ f₁                  : by simp
      ... = c₀.X.hom ≫ F.map g₁.right : g₁.w
      ... = c₀.X.hom ≫ F.map g₂.right : by rw this
      ... = 𝟙 _ ≫ f₂                  : g₂.w.symm
      ... = f₂                         : by simp
end

end preserves_finite_limits_of_flat

/-- Representably flat functors preserve finite limits. -/
noncomputable
def preserves_finite_limits_of_flat (F : C ⥤ D) [representably_flat F] :
  preserves_finite_limits F := ⟨λ J _ _, by exactI ⟨λ K, ⟨λ c hc,
{ lift := preserves_finite_limits_of_flat.lift F hc,
  fac' := preserves_finite_limits_of_flat.fac F hc,
  uniq' := λ s m h, by
  { apply preserves_finite_limits_of_flat.uniq F hc,
    exact h,
    exact preserves_finite_limits_of_flat.fac F hc s } }⟩⟩⟩

/--
If `C` is finitely cocomplete, then `F : C ⥤ D` is representably flat iff it preserves
finite limits.
-/
noncomputable
def preserves_finite_limits_iff_flat [has_finite_limits C] (F : C ⥤ D) :
  representably_flat F ≃ preserves_finite_limits F :=
{ to_fun := λ _, by exactI preserves_finite_limits_of_flat F,
  inv_fun := λ _, by exactI flat_of_preserves_finite_limits F,
  left_inv := λ _, proof_irrel _ _,
  right_inv := λ x, by { cases x, unfold preserves_finite_limits_of_flat, congr } }

end has_limit

section small_category
variables {C D : Type u₁} [small_category C] [small_category D]

/--
(Implementation)
The evaluation of `Lan F` at `X` is the colimit over the costructured arrows over `X`.
-/
noncomputable
def Lan_evaluation_iso_colim (E : Type u₂) [category.{u₁} E] (F : C ⥤ D) (X : D)
  [∀ (X : D), has_colimits_of_shape (costructured_arrow F X) E] :
  Lan F ⋙ (evaluation D E).obj X ≅
  ((whiskering_left _ _ E).obj (costructured_arrow.proj F X)) ⋙ colim :=
nat_iso.of_components (λ G, colim.map_iso (iso.refl _))
begin
  intros G H i,
  ext,
  simp only [functor.comp_map, colimit.ι_desc_assoc, functor.map_iso_refl, evaluation_obj_map,
    whiskering_left_obj_map, category.comp_id, Lan_map_app, category.assoc],
  erw [colimit.ι_pre_assoc (Lan.diagram F H X) (costructured_arrow.map j.hom),
    category.id_comp, category.comp_id, colimit.ι_map],
  cases j,
  cases j_right,
  congr,
  rw [costructured_arrow.map_mk, category.id_comp, costructured_arrow.mk]
end

/--
If `F : C ⥤ D` is a representably flat functor between small categories, then the functor
`Lan F.op` that takes presheaves over `C` to presheaves over `D` preserves finite limits.
-/
noncomputable
instance Lan_preserves_finite_limits_of_flat (F : C ⥤ D) [representably_flat F] :
  preserves_finite_limits (Lan F.op : _ ⥤ (Dᵒᵖ ⥤ Type u₁)) :=
⟨λ J _ _, begin
  resetI,
  apply preserves_limits_of_shape_of_evaluation (Lan F.op : (Cᵒᵖ ⥤ Type u₁) ⥤ (Dᵒᵖ ⥤ Type u₁)) J,
  intro K,
  haveI : is_filtered (costructured_arrow F.op K) :=
    is_filtered.of_equivalence (structured_arrow_op_equivalence F (unop K)),
  exact preserves_limits_of_shape_of_nat_iso (Lan_evaluation_iso_colim _ _ _).symm
end⟩

instance Lan_flat_of_flat (F : C ⥤ D) [representably_flat F] :
  representably_flat (Lan F.op : _ ⥤ (Dᵒᵖ ⥤ Type u₁)) := flat_of_preserves_finite_limits _

variable [has_finite_limits C]

noncomputable
instance Lan_preserves_finite_limits_of_preserves_finite_limits (F : C ⥤ D)
  [preserves_finite_limits F] : preserves_finite_limits (Lan F.op : _ ⥤ (Dᵒᵖ ⥤ Type u₁)) :=
begin
  haveI := flat_of_preserves_finite_limits F,
  apply_instance
end

lemma flat_iff_Lan_flat (F : C ⥤ D) :
  representably_flat F ↔ representably_flat (Lan F.op : _ ⥤ (Dᵒᵖ ⥤ Type u₁)) :=
⟨λ H, by exactI infer_instance, λ H,
begin
  resetI,
  haveI := preserves_finite_limits_of_flat (Lan F.op : _ ⥤ (Dᵒᵖ ⥤ Type u₁)),
  haveI : preserves_finite_limits F :=
    ⟨λ _ _ _, by exactI preserves_limit_of_Lan_presesrves_limit _ _⟩,
  apply flat_of_preserves_finite_limits
end⟩

/--
If `C` is finitely complete, then `F : C ⥤ D` preserves finite limits iff
`Lan F.op : (Cᵒᵖ ⥤ Type*) ⥤ (Dᵒᵖ ⥤ Type*)` preserves finite limits.
-/
noncomputable
def preserves_finite_limits_iff_Lan_preserves_finite_limits (F : C ⥤ D) :
  preserves_finite_limits F ≃ preserves_finite_limits (Lan F.op : _ ⥤ (Dᵒᵖ ⥤ Type u₁)) :=
{ to_fun := λ _, by exactI infer_instance,
  inv_fun := λ _, ⟨λ _ _ _, by exactI preserves_limit_of_Lan_presesrves_limit _ _⟩,
  left_inv := λ x, by { cases x, unfold preserves_finite_limits_of_flat, congr },
  right_inv := λ x,
  begin
    cases x,
    unfold preserves_finite_limits_of_flat,
    congr,
    unfold category_theory.Lan_preserves_finite_limits_of_preserves_finite_limits
      category_theory.Lan_preserves_finite_limits_of_flat, congr
  end }

end small_category
end category_theory


def g (xs ys : List Nat) : Nat :=
  match xs, ys with
  | [a, b], _ => Nat.succ (a+b)
  | _, [b, c] => Nat.succ b
  | _, _   => 1

example (a b : Bool) (x y z : Nat) (xs : List Nat) (h1 : (if a then x else y) = 0) (h2 : xs.head! = 0) : g [x] xs = 1 := by
  simp [g]
  repeat any_goals (split at *)
  any_goals (first | decide | contradiction | injections)
  next b c _ =>
    show Nat.succ b = 1
    simp [List.head!] at h2; simp [h2]
  next b c _ =>
    show Nat.succ b = 1
    simp [List.head!] at h2; simp [h2]

example (a : Bool) (h1 : (if a then x else y) = 1) : x + y > 0 := by
  split at h1
  · subst h1; rw [Nat.succ_add]; apply Nat.zero_lt_succ
  · subst h1; apply Nat.zero_lt_succ

def f (x : Nat) : Nat :=
  match x with
  | 100 => 0
  | 200 => 0
  | _   => 1

example (h1 : f x = 0) (h2 : x > 300) : False := by
  simp [f] at h1
  split at h1
  · contradiction
  · contradiction
  · contradiction

example (h1 : f x = 0) (h2 : x > 300) : False := by
  simp [f] at h1
  split at h1 <;> contradiction


import GMLInit.Data.Index.Basic
import GMLInit.Data.Index.Bind
import GMLInit.Data.Index.Map
import GMLInit.Logic.Decidable
import GMLInit.Logic.ListConnectives

namespace Any

protected def get : {ps : List Prop} → [DecidableList ps] → Any ps → Index ps
| _, .cons (isTrue _) _, _ => .head
| _::ps, .cons (isFalse fh) inst, h =>
  have : Any ps :=
    match h with
    | .head hh => absurd hh fh
    | .tail ht => ht
  .tail (@Any.get ps inst this)

theorem get_prop : {ps : List Prop} → [DecidableList ps] → (h : Any ps) → h.get.val
| _, .cons (isTrue hh) _, _ => hh
| _, .cons (isFalse fh) _, h =>
  get_prop $ match h with
    | .head hh => absurd hh fh
    | .tail ht => ht

end Any

namespace List

def dedup {α} (s : Setoid α) [DecidableRel s.r] : List α → List α
| [] => []
| x :: xs =>
  if Any (xs.map (s.r x))
  then dedup s xs
  else x :: dedup s xs

end List

namespace Index
variable {α} (s : Setoid α) [DecidableRel s.r] {xs : List α}

def dedup : {xs : List α} → Index xs → Index (xs.dedup s)
| x :: xs, head =>
  if h : Any (xs.map (s.r x))
  then
    have : (x :: xs).dedup s = xs.dedup s := if_pos h
    this ▸ dedup (h.get.unmap _)
  else
    have : (x :: xs).dedup s = x :: xs.dedup s := if_neg h
    this ▸ head
| x :: xs, tail i =>
  if h : Any (xs.map (s.r x))
  then
    have : (x :: xs).dedup s = xs.dedup s := if_pos h
    this ▸ dedup i
  else
    have : (x :: xs).dedup s = x :: xs.dedup s := if_neg h
    this ▸ tail (dedup i)

theorem val_dedup (i : Index xs) : s.r i.val (i.dedup s).val := by
  induction xs with
  | nil => contradiction
  | cons x xs ih =>
    match i with
    | head =>
      clean unfold dedup
      split
      next h =>
        rw [val_ndrec]
        transitivity (h.get.unmap (s.r x)).val using s.r
        · rw [val_head, ←val_unmap (s.r x)]
          exact Any.get_prop ..
        · exact ih ..
      next => rw [val_ndrec]; reflexivity using s.r
    | tail i =>
      clean unfold dedup
      split
      next => rw [val_ndrec]; exact ih ..
      next => rw [val_ndrec]; exact ih ..

def undedup : {xs : List α} → Index (xs.dedup s) → Index xs
| x :: xs, i =>
  if h : Any (xs.map (s.r x))
  then
    have : (x :: xs).dedup s = xs.dedup s := if_pos h
    tail (undedup (this ▸ i))
  else
    have : (x :: xs).dedup s = x :: xs.dedup s := if_neg h
    match this ▸ i with
    | head => head
    | tail i => tail (undedup i)

theorem val_undedup (i : Index (xs.dedup s)) : (i.undedup s).val = i.val := by
  induction xs with
  | nil => contradiction
  | cons x xs ih =>
    simp only [undedup]
    split
    next ha =>
      have : (x :: xs).dedup s = xs.dedup s := if_pos ha
      rw [val_tail]
      transitivity (this ▸ i).val
      · exact ih ..
      · exact val_ndrec ..
    next ha =>
      split
      next h =>
        rw [eqNdrec_symm] at h
        rw [h, val_ndrec]
      next h =>
        rw [eqNdrec_symm] at h
        rw [h, val_ndrec, val_tail, val_tail, ih]

theorem dedup_undedup {xs : List α} (i : Index (xs.dedup s)) : (i.undedup s).dedup s = i := by
  induction xs with
  | nil => contradiction
  | cons x xs ih =>
    simp only [undedup]
    split
    next ha => rw [dedup, dif_pos ha, ih, eqNdrec_symm]
    next ha =>
      split
      next h => rw [dedup, dif_neg ha, eqNdrec_symm, ←h]
      next h => rw [dedup, dif_neg ha, eqNdrec_symm, ih, ←h]

theorem undedup_dedup {xs : List α} (i : Index xs) : s.r ((i.dedup s).undedup s).val i.val := by
  symmetry
  rw [val_undedup]
  exact val_dedup ..

theorem dedup_eq_of_rel {xs : List α} {i : Index xs} {j : Index (xs.dedup s)} (h : s.r i.val j.val) : i.dedup s = j := by
  induction xs with
  | nil => contradiction
  | cons x xs ih =>
    match i with
    | head =>
      rw [val_head] at h
      rw [dedup]
      split
      next ha =>
        rw [eqNdrec_symm]
        apply ih
        rw [val_ndrec]
        transitivity x
        · symmetry
          rw [←val_unmap (s.r x)]
          exact Any.get_prop ..
        · exact h
      next ha =>
        have : (x :: xs).dedup s = x :: xs.dedup s := if_neg ha
        rw [eqNdrec_symm]
        match hj : this ▸ j with
        | head =>
          rw [←hj]
        | tail j =>
          rw [eqNdrec_symm] at hj
          rw [hj, val_ndrec, val_tail] at h
          absurd ha
          apply Any.introIdx ((j.undedup s).map (s.r x))
          rw [val_map, val_undedup]
          exact h
    | tail i =>
      rw [val_tail] at h
      rw [dedup]
      split
      next ha =>
        rw [eqNdrec_symm]
        apply ih
        rw [val_ndrec]
        exact h
      next ha =>
        have : (x :: xs).dedup s = x :: xs.dedup s := if_neg ha
        rw [eqNdrec_symm]
        match hj : this ▸ j with
        | head =>
          rw [eqNdrec_symm] at hj
          rw [hj, val_ndrec, val_head] at h
          absurd ha
          apply Any.introIdx (i.map (s.r x))
          rw [val_map]
          symmetry
          exact h
        | tail j =>
          rw [eqNdrec_symm] at hj
          rw [hj, val_ndrec, val_tail] at h
          rw [hj, ih h]
          elim_casts

theorem dedup_eq_iff_rel {xs : List α} (i : Index xs) (j : Index (xs.dedup s)) : i.dedup s = j ↔ s.r i.val j.val := by
  constr
  · intro | rfl => exact val_dedup ..
  · exact dedup_eq_of_rel s

theorem dedup_eq_dedup_of_rel {xs : List α} {i j : Index xs} (h : s.r i.val j.val) : i.dedup s = j.dedup s := by
  apply dedup_eq_of_rel
  transitivity j.val
  · exact h
  · exact val_dedup ..

end Index


import algebra.big_operators.ring
import topology.algebra.infinite_sum
import number_theory.arithmetic_function
import data.complex.basic
import analysis.special_functions.pow
import analysis.normed.group.infinite_sum
import order.filter.basic
import order.filter.at_top_bot
import analysis.calculus.fderiv
import measure_theory.integral.integrable_on
import measure_theory.integral.interval_integral
import topology.metric_space.basic

/-!
# Dirchlet Series and the Prime Number Theorem

Take `f` to have image `ℂ`. Many defined functions have image `ℕ` or `ℤ`, but they
have canonical coercions so this should be fine.
-/

open finset filter measure_theory interval_integral metric
open_locale big_operators arithmetic_function uniformity filter

-- variables {α : Type*} {β : Type*} {ι : Type*} [uniform_space β]

-- def uniform_cauchy_seq_on
--   (F : ι → α → β) (p : filter ι) (s : set α) :=
--   ∀ u : set (β × β), u ∈ 𝓤 β → (∀ᶠ (n : ι × ι) in (p ×ᶠ p), ∀ (x : α), x ∈ s → (F n.fst x, F n.snd x) ∈ u)

-- lemma uniform_cauchy_seq_on_iff [complete_space β] [nonempty β] (F : ι → α → β) (p : filter ι) (s : set α) :
--   uniform_cauchy_seq_on F p s ↔ tendsto_uniformly_on F (λ x : α, lim p (λ n : ι, F n x)) p s :=
-- begin

-- end

lemma bah
(F : ℕ → ℝ → ℝ)
(f : ℝ → ℝ)
(G : ℕ → ℝ)
(g : ℝ)
(s : set ℝ)
(hf : tendsto_uniformly_on F f at_top s)
(hg : tendsto G at_top (nhds g)) :
tendsto_uniformly_on (λ n : ℕ, λ x : ℝ, F n x + G n) (λ x : ℝ, f x + g) at_top s :=
begin
  sorry,
end

lemma mul_cancel_inv_left₀ {a b : ℝ} (ha : a ≠ 0) : a⁻¹ * (a * b) = b :=
 begin
  conv { congr, congr, skip, rw ←inv_inv a, },
  have : a⁻¹ ≠ 0, simp [ha],
  rw mul_inv_cancel_left₀ this,
end

lemma norm_add_three_le {a b c : ℝ} : ∥a + b + c∥ ≤ ∥a∥ + ∥b∥ + ∥c∥ :=
begin
  refine le_trans (norm_add_le _ _) _,
  exact add_le_add (norm_add_le _ _) rfl.le,
end

lemma norm_sub_comm {a b : ℝ} : ∥a - b∥ = ∥b - a∥ :=
begin
  have : b - a = (-1) * (a - b), ring,
  rw [this, norm_mul],
  simp,
end

/-- You could rearrange this lemma to actually choose g for the user, but I don't
understand that syntax so would welcome help! -/
lemma uniform_convergence_of_uniform_cauchy
(f : ℕ → ℝ → ℝ)
(g : ℝ → ℝ)
(s : set ℝ)
(hfg : ∀ x : ℝ, x ∈ s → tendsto (λ n, f n x) at_top (nhds (g x)))
(hu : ∀ ε : ℝ, ε > 0 → ∃ N : ℕ, ∀ m : ℕ, m ≥ N → ∀ n : ℕ, n ≥ N → ∀ x : ℝ, x ∈ s → ∥f m x - f n x∥ < ε) :
tendsto_uniformly_on f g at_top s :=
begin
  rw tendsto_uniformly_on_iff,
  intros ε hε,
  have half_ε : (2⁻¹ * ε) > 0, simp [hε.lt],
  rw eventually_at_top,

  specialize hu (2⁻¹ * ε) half_ε,
  cases hu with N hN,

  use N,
  intros n hn x hx,

  specialize hfg x hx,
  rw metric.tendsto_nhds at hfg,
  specialize hfg (2⁻¹ * ε) half_ε,
  rw eventually_at_top at hfg,
  cases hfg with N2 hN2,

  let m := max N N2,

  specialize hN n hn m (by simp) x hx,
  specialize hN2 m (by simp),
  rw dist_eq_norm,
  rw dist_eq_norm at hN2,
  rw norm_sub_comm,
  have : f n x - g x = (f n x - f m x) + (f m x - g x), ring,
  rw this,
  have : ∥(f n x - f m x) + (f m x - g x)∥ ≤ ∥f n x - f m x∥ + ∥f m x - g x∥, exact norm_add_le _ _,
  refine lt_of_le_of_lt this _,

  have : ∥f n x - f m x∥ + ∥f m x - g x∥ < (2⁻¹ * ε) + (2⁻¹ * ε),
  exact add_lt_add hN hN2,
  exact lt_of_lt_of_eq this (by ring),
end

lemma uniform_cauchy_of_uniform_convergence
(f : ℕ → ℝ → ℝ)
(g : ℝ → ℝ)
(s : set ℝ)
(hu : tendsto_uniformly_on f g at_top s) :
∀ ε : ℝ, ε > 0 → ∃ N : ℕ, ∀ m : ℕ, m ≥ N → ∀ n : ℕ, n ≥ N → ∀ x : ℝ, x ∈ s → ∥f m x - f n x∥ < ε :=
begin
  intros ε hε,
  have half_ε : (2⁻¹ * ε) > 0, simp [hε.lt],
  rw tendsto_uniformly_on_iff at hu,
  specialize hu (2⁻¹ * ε) half_ε,
  rw eventually_at_top at hu,
  cases hu with N hN,
  use N,
  intros m hm n hn x hx,

  have hmm := hN m hm x hx,
  have hnn := hN n hn x hx,
  rw dist_eq_norm at *,

  have : f m x - f n x = f m x - g x + (g x - f n x), ring,
  rw this,
  refine lt_of_le_of_lt (norm_add_le _ _) _,
  rw norm_sub_comm at hmm,
  refine lt_of_lt_of_eq (add_lt_add hmm hnn) _,
  ring,
end

lemma dumb { a b c : ℝ } : (-1) * a / (b - c) = a / (c - b) :=
begin
  ring_nf,
  rw mul_comm,
  rw ←mul_neg,
  rw neg_inv,
  rw neg_sub,
  rw mul_comm,
end

lemma fdfdfd { a b : ℝ} : min a b = a ∨ min a b = b :=
begin
  exact min_choice a b
end

lemma min_eq_max_iff_eq {a b : ℝ} : min a b = max a b ↔ a = b := begin
  refine ⟨(λ heq, _), (λ h, by simp [h])⟩,
  {
    cases min_choice a b,
    rw h at heq,
    rw min_eq_left_iff at h,
    have := max_eq_left_iff.mp heq.symm,
    exact ge_antisymm this h,

    rw h at heq,
    rw min_eq_right_iff at h,
    have := max_eq_right_iff.mp heq.symm,
    exact le_antisymm this h,
  },
end

lemma mean_value_theorem_for_differences
(f : ℕ → ℝ → ℝ)
(f' : ℕ → ℝ → ℝ)
(g : ℝ → ℝ)
(g' : ℝ → ℝ)
(x r R : ℝ)
(hrR : r < R)
(hf : ∀ (n : ℕ), ∀ (y : ℝ), y ∈ ball x R → has_deriv_at (f n) (f' n y) y)
(hfg : ∀ (y : ℝ), y ∈ closed_ball x r → tendsto (λ n, f n y) at_top (nhds (g y)))
(hfg' : tendsto_uniformly_on f' g' at_top (closed_ball x r)) :
∀ (m n : ℕ) (x' y : ℝ), x' ∈ closed_ball x r → y ∈ closed_ball x r → x' ≠ y →
    ∃ (z : ℝ), z ∈ closed_ball x r ∧ (x' - y)⁻¹ * ((f m x' - f n x') - (f m y - f n y)) = f' m z - f' n z :=
begin
  intros m n y' z' hy' hz' h,
  let y := min y' z',
  let z := max y' z',
  have hy : y ∈ closed_ball x r,
  {
    cases le_total y' z' with h h,
    have : y = y', by simp [h],
    rw this, exact hy',
    have : y = z', by simp [h],
    rw this, exact hz',
  },
  have hz : z ∈ closed_ball x r, {
    cases le_total y' z' with h h,
    have : z = z', by simp [h],
    rw this, exact hz',
    have : z = y', by simp [h],
    rw this, exact hy',
  },
  have hyz : y < z, { exact lt_of_le_of_ne min_le_max (λ hb, h (min_eq_max_iff_eq.mp hb)), },
  have hfc : continuous_on (f m - f n) (set.Icc y z), sorry,
  have hff' : ∀ (x : ℝ), x ∈ set.Ioo y z → has_deriv_at (f m - f n) (f' m x - f' n x) x, sorry,
  have mvt := exists_has_deriv_at_eq_slope (f m - f n) (f' m - f' n) hyz hfc hff',
  rcases mvt with ⟨c, hc, hc'⟩,
  use c,
  split,
  { sorry, },
  have : (y' - z')⁻¹ * (f m y' - f n y' - (f m z' - f n z')) = (y - z)⁻¹ * (f m y - f n y - (f m z - f n z)),
  { sorry, },
  rw this,
  simp at hc',
  have : (f m z - f n z - (f m y - f n y)) / (z - y) = (y - z)⁻¹ * (f m y - f n y - (f m z - f n z)),
  {
    have : f m z - f n z - (f m y - f n y) = (-1) * (f m y - f n y - (f m z - f n z)), ring,
    rw this,
    have : (-1) * (f m y - f n y - (f m z - f n z)) / (z - y) = (f m y - f n y - (f m z - f n z)) / (y - z), exact dumb,
    rw this,
    ring,
  },
  rw ←this,
  exact hc'.symm,
end

lemma difference_quotients_converge
(f : ℕ → ℝ → ℝ)
(g : ℝ → ℝ)
(x r : ℝ)
(hfg : ∀ (y : ℝ), y ∈ closed_ball x r → tendsto (λ n, f n y) at_top (nhds (g y))) :
∀ y : ℝ, y ∈ closed_ball x r → ∀ z : ℝ, z ∈ closed_ball x r → tendsto (λ n : ℕ, ∥z - y∥⁻¹ * ((f n z) - (f n y))) at_top (nhds (∥z - y∥⁻¹ * ((g z) - (g y)))) :=
begin
  intros y hy z hz,
  apply tendsto.const_mul,
  exact tendsto.sub (hfg z hz) (hfg y hy),
end

lemma difference_quotients_converge_uniformly
(f : ℕ → ℝ → ℝ)
(f' : ℕ → ℝ → ℝ)
(g : ℝ → ℝ)
(g' : ℝ → ℝ)
(x r R : ℝ)
(hrR : r < R)
(hf : ∀ (n : ℕ), ∀ (y : ℝ), y ∈ ball x R → has_deriv_at (f n) (f' n y) y)
(hfg : ∀ (y : ℝ), y ∈ closed_ball x r → tendsto (λ n, f n y) at_top (nhds (g y)))
(hfg' : tendsto_uniformly_on f' g' at_top (closed_ball x r)) :
∀ y : ℝ, y ∈ closed_ball x r → tendsto_uniformly_on (λ n : ℕ, λ z : ℝ, ∥z - y∥⁻¹ * ((f n z) - (f n y))) (λ z : ℝ, ∥z - y∥⁻¹ * ((g z) - (g y))) at_top ((closed_ball x r)) :=
begin
  -- Proof strategy: Rewrite the Cauchy sequence of difference quotients as
  -- a difference quotient. Then apply the mean value theorem and the uniform
  -- convergence of the derivatives
  intros y hy,
  apply uniform_convergence_of_uniform_cauchy,
  refine difference_quotients_converge _ _ _ _ hfg y hy,
  intros ε hε,

  cases uniform_cauchy_of_uniform_convergence _ _ _ hfg' ε hε with N hN,
  use N,
  intros m hm n hn z hz,
  rw [←mul_sub, ←norm_inv, norm_mul, norm_norm, ←norm_mul],
  by_cases hzy : z = y,
  { simp [hzy, hε.lt], },
  have hmvt := mean_value_theorem_for_differences f f' g g' x r R hrR hf hfg hfg' m n z y hz hy hzy,
  rcases hmvt with ⟨ξ, hξ, hξ'⟩,

  have : (f m z - f n z - (f m y - f n y)) = (f m z - f m y - (f n z - f n y)), ring,
  rw [←this, hξ'],
  exact hN m hm n hn ξ hξ,
end

lemma uniform_convergence_of_uniform_convergence_derivatives
(f : ℕ → ℝ → ℝ)
(f' : ℕ → ℝ → ℝ)
(g : ℝ → ℝ)
(g' : ℝ → ℝ)
(x r R : ℝ)
(hrpos : 0 < r)
(hrR : r < R)
(hf : ∀ (n : ℕ), ∀ (y : ℝ), y ∈ ball x R → has_deriv_at (f n) (f' n y) y)
(hfg : ∀ (y : ℝ), y ∈ closed_ball x r → tendsto (λ n, f n y) at_top (nhds (g y)))
(hfg' : tendsto_uniformly_on f' g' at_top (closed_ball x r)) :
tendsto_uniformly_on f g at_top (closed_ball x r) :=
begin
  refine uniform_convergence_of_uniform_cauchy _ _ _ hfg _,
  intros ε hε,
  have hxcb : x ∈ closed_ball x r, { rw mem_closed_ball, simp, exact hrpos.le, },
  have := metric.cauchy_seq_iff.mp (hfg x hxcb).cauchy_seq,
  have two_inv_pos : 0 < (2 : ℝ)⁻¹, simp,
  have ε_over_two_pos : 0 < (2⁻¹ * ε),
  { exact mul_pos two_inv_pos hε.lt, },
  specialize this (2⁻¹ * ε) ε_over_two_pos.gt,
  cases this with N1 hN1,

  have foo := uniform_cauchy_of_uniform_convergence _ _ _ hfg',
  have : 0 < (2⁻¹ * r⁻¹ * ε), {
    exact mul_pos (mul_pos (by norm_num) (by simp [hrpos])) hε.lt,
  },
  specialize foo (2⁻¹ * r⁻¹ * ε) this.gt,
  cases foo with N2 hN2,

  let N := max N1 N2,
  use N,
  intros m hm n hn y hy,

  have : f m y - f n y = (f m y - f n y) - (f m x - f n x) + (f m x - f n x), ring,
  rw this,
  have : ∥f m y - f n y - (f m x - f n x) + (f m x - f n x)∥ ≤ ∥f m y - f n y - (f m x - f n x)∥ + ∥f m x - f n x∥, exact norm_add_le _ _,
  refine lt_of_le_of_lt this _,

  have : ∥f m y - f n y - (f m x - f n x)∥ = ∥y - x∥ * (∥y - x∥⁻¹ * ∥f m y - f n y - (f m x - f n x)∥), {
    by_cases hyxx : y = x,
    { simp [hyxx], },
    have hxyy : y - x ≠ 0, exact λ H, hyxx (sub_eq_zero.mp H),
    have hxyy' : ∥y - x∥ ≠ 0, simp [hxyy],
    exact (mul_inv_cancel_left₀ hxyy' _).symm,
  },
  rw this,
  by_cases h : y = x,
  { simp only [h, sub_self, norm_zero, mul_zero, zero_add],
    refine lt_trans (hN1 m (ge_trans hm (le_max_left N1 N2).ge) n (ge_trans hn (le_max_left N1 N2).ge)) _,
    rw mul_lt_iff_lt_one_left hε.lt,
    norm_num, },

  rcases mean_value_theorem_for_differences f f' g g' x r R hrR hf hfg hfg' m n y x hy hxcb h with ⟨z, hz, hz'⟩,
  conv {to_lhs, congr, congr, skip, rw ←norm_inv, rw ←norm_mul, rw hz', },
  specialize hN2 m (ge_trans hm (by simp)) n (ge_trans hn (by simp)) z hz,
  have hyx : ∥y - x∥ ≤ r, { rw mem_closed_ball at hy, exact hy, },
  specialize hN1 m (ge_trans hm (by simp)) n (ge_trans hn (by simp)),
  rw dist_eq_norm at hN1,

  have : ε = (2⁻¹ * ε) + (2⁻¹ * ε), ring,
  rw this,

  have : ∥y - x∥ * ∥f' m z - f' n z∥ < 2⁻¹ * ε, {
    sorry,
  },
  exact add_lt_add this hN1,
end

/-! (d/dx) lim_{n → ∞} f_n x = lim_{n → ∞} f'_n x on a closed ball when the f'_n are
continuous and converge _unifomrly_ to their limit -/
lemma swap_limit_and_derivative
(f : ℕ → ℝ → ℝ)
(f' : ℕ → ℝ → ℝ)
(g : ℝ → ℝ)
(g' : ℝ → ℝ)
(x r R : ℝ)
(hrR : r < R)
(hf : ∀ (n : ℕ), ∀ (y : ℝ), y ∈ ball x R → has_deriv_at (f n) (f' n y) y)
(hfg : ∀ (y : ℝ), y ∈ closed_ball x r → tendsto (λ n, f n y) at_top (nhds (g y)))
(hfg' : tendsto_uniformly_on f' g' at_top (closed_ball x r)) :
∀ y : ℝ, y ∈ ball x r → has_deriv_at g (g' y) y :=
begin
  -- We do the famous "ε / 3 proof" which will involve several bouts of utilizing
  -- uniform continuity. First we setup our goal in terms of ε and δ
  intros y hy,
  rw has_deriv_at_iff_tendsto,
  rw tendsto_nhds_nhds,

  -- Now some important auxiliary facts such as:
  have hrpos : 0 < r, {
    rw mem_ball at hy,
    calc 0 ≤ dist y x : dist_nonneg ... < r : hy,
  },

  have hball_mem : ∀ z, z ∈ ball x r → z ∈ ball x R,
  { intros z hz,
    rw mem_ball,
    rw mem_ball at hz,
    calc dist z x < r : hz ... < R : hrR, },

  have hbig_ball_mem : ∀ z, z ∈ closed_ball x r → z ∈ ball x R,
  exact (λ z hz, by calc dist z x ≤ r : hz ... < R : hrR),

  have hyR : y ∈ ball x R, exact hball_mem y hy,

  have hyc : y ∈ closed_ball x r,
  { rw mem_ball at hy,
    exact hy.le, },

  -- The closed ball is compact
  have hball : is_compact (closed_ball x r),
  exact compact_iff_closed_bounded.mpr ⟨is_closed_ball, bounded_closed_ball⟩,

  -- has_deriv_at implies continuity of primal
  have hfc : ∀ (n : ℕ), continuous_on (f n) (closed_ball x r),
  exact λ n z hz, (hf n z (hbig_ball_mem z hz)).continuous_at.continuous_within_at,

  -- continuity of primal implies uniform continuity
  have hfuc : ∀ (n : ℕ), uniform_continuous_on (f n) (closed_ball x r),
  exact (λ n, hball.uniform_continuous_on_of_continuous (hfc n)),

  -- difference of differentiable functions is differentiable
  have hfmndiff : ∀ (m n : ℕ) (y : ℝ), y ∈ ball x R → has_deriv_at (f m - f n) (f' m y - f' n y) y,
  exact (λ m n y hy, (hf m y hy).sub (hf n y hy)),

  -- mean value theorem applied to differences
  have hfmnmvt := mean_value_theorem_for_differences f f' g g' x r R hrR hf hfg hfg',

  -- uniform convergence of the derivatives implies uniform convergence of the primal
  have hfguc := uniform_convergence_of_uniform_convergence_derivatives f f' g g' x r R hrpos hrR hf hfg hfg', -- tendsto_uniformly_on f g at_top (closed_ball x r),

  -- convergence of the primal and uniform convergence of the derivatives implies
  -- uniform convergence of the difference quotients
  have hdiff := difference_quotients_converge_uniformly f f' g g' x r R hrR hf hfg hfg' y hyc,

  -- The first (ε / 3) comes from the convergence of the derivatives
  intros ε hε,
  have : 0 < (3 : ℝ)⁻¹, simp, linarith,
  have ε_over_three_pos : 0 < (3⁻¹ * ε),
  { exact mul_pos this hε.lt, },
  rw tendsto_uniformly_on_iff at hfg',
  specialize hfg' (3⁻¹ * ε) ε_over_three_pos.gt,
  rw eventually_at_top at hfg',
  rcases hfg' with ⟨N1, hN1⟩,

  -- The second (ε / 3) comes from the uniform convergence of the difference quotients
  rw tendsto_uniformly_on_iff at hdiff,
  specialize hdiff (3⁻¹ * ε) ε_over_three_pos.gt,
  rw eventually_at_top at hdiff,
  rcases hdiff with ⟨N2, hN2⟩,

  -- These two N determine our final N
  let N := max N1 N2,

  -- The final (ε / 3) comes from the definition of a derivative
  specialize hf N y hyR,
  rw has_deriv_at_iff_tendsto at hf,
  rw tendsto_nhds_nhds at hf,
  specialize hf (3⁻¹ * ε) ε_over_three_pos.gt,
  rcases hf with ⟨δ', hδ', hf⟩,

  -- Choose our final δ
  let δ := min (r - dist y x) δ',
  have hδ : δ > 0, {
    refine lt_min _ hδ'.lt,
    rw sub_pos,
    exact hy,
  },

  -- Start the final manipulation
  use [δ, hδ],
  intros x' hx',
  have hxc : x' ∈ closed_ball x r, {
    rw mem_closed_ball,
    apply le_of_lt,
    have foo : dist x' y < r - dist y x, calc dist x' y < δ : hx' ... ≤ r - dist y x : by simp [δ],
    have ff : dist x' y + dist y x < r, linarith [foo],
    have fff : dist x' x ≤ dist x' y + dist y x, exact dist_triangle _ _ _,
    calc dist x' x ≤ dist x' y + dist y x : fff ... < r : ff,
  },
  have hxy : dist x' y < δ', calc dist x' y < δ : hx' ... ≤ δ' : by simp [δ],
  specialize hf hxy,

  -- There's a technical issue where we need to rule out the case y = x'
  by_cases hy' : y = x',
  { simp [hy', hε.lt], },

  have hx'y : x' - y ≠ 0, exact λ H, hy' (sub_eq_zero.mp H).symm,

  -- Now our three inequalities come from `hf`, `hN1`, and `hN2`. Specialize
  -- to make this clear
  specialize hN1 N (by simp) y hyc,
  specialize hN2 N (by simp) x' hxc,
  rw dist_eq_norm at *,
  simp only [algebra.id.smul_eq_mul, sub_zero, norm_mul, norm_inv, norm_norm],

  -- Begin algebraic manipulations
  have : ∥x' - y∥⁻¹ * ∥g x' - g y - (x' - y) * g' y∥ =
    ∥(x' - y)⁻¹ * (g x' - g y) - g' y∥,
  { rw [←norm_inv, ←norm_mul],
    congr,
    rw [mul_sub, mul_cancel_inv_left₀ hx'y], },
  rw this,

  -- Add zero a couple times and regroup
  have : (x' - y)⁻¹ * (g x' - g y) - g' y =
    (x' - y)⁻¹ * (g x' - g y) - (x' - y)⁻¹ * (f N x' - f N y) +
    (x' - y)⁻¹ * (f N x' - f N y) - (f' N y) +
    (f' N y) - g' y,
  { ring, },
  rw this,

  -- Triangle inequality twice
  have hregroup : (x' - y)⁻¹ * (g x' - g y) - (x' - y)⁻¹ * (f N x' - f N y) +
    (x' - y)⁻¹ * (f N x' - f N y) - (f' N y) +
    (f' N y) - g' y =
    ((x' - y)⁻¹ * (g x' - g y) - (x' - y)⁻¹ * (f N x' - f N y)) +
    ((x' - y)⁻¹ * (f N x' - f N y) - (f' N y)) +
    ((f' N y) - g' y),
  { ring, },
  have : ∥(x' - y)⁻¹ * (g x' - g y) - (x' - y)⁻¹ * (f N x' - f N y) +
    (x' - y)⁻¹ * (f N x' - f N y) - (f' N y) +
    (f' N y) - g' y∥ ≤
    ∥(x' - y)⁻¹ * (g x' - g y) - (x' - y)⁻¹ * (f N x' - f N y)∥ +
    ∥(x' - y)⁻¹ * (f N x' - f N y) - (f' N y)∥ +
    ∥(f' N y) - g' y∥,
  { rw hregroup,
    exact norm_add_three_le, },
  refine lt_of_le_of_lt this _,

  -- Get contributing factors to the right shape
  rw norm_sub_comm at hN1,
  rw [←mul_sub, norm_mul, ←norm_inv, norm_norm, ←norm_mul, mul_sub] at hN2,
  simp only [algebra.id.smul_eq_mul, sub_zero, norm_mul, norm_inv, norm_norm] at hf,
  rw [←norm_inv, ←norm_mul, mul_sub, mul_cancel_inv_left₀ hx'y] at hf,

  -- Final inequalities
  refine lt_of_lt_of_le (add_lt_add_of_lt_of_lt (add_lt_add_of_lt_of_lt hN2 hf) hN1) _,
  apply le_of_eq,
  ring,
end

variables {α : Type*} [add_comm_monoid α]

def head_sum (f : ℕ → α) : (ℕ → α) := (λ n : ℕ, ∑ i in range n, f i)

variables [topological_space α]

def nat_summable (f : ℕ → α) : Prop := ∃ (a : α), tendsto (head_sum f) at_top (nhds a)

namespace nat
namespace arithmetic_function
variables {R : Type*} [has_zero R] [has_abs R]
(f : arithmetic_function ℂ) (g h : ℂ → ℂ)
(s t : ℂ) (r : ℝ)

/-- A Dirichlet series of a function `f` at `s` is itself a function from `ℕ` to `ℂ`
which returns the `n`th term of the sum ∑ (f n) / n ^ s -/
noncomputable def dirichlet_series := (λ n : ℕ, (f n) / ((n : ℂ) ^ s))

noncomputable def dirichlet_series' := ∑' i, (λ n : ℕ, (f n) / ((n : ℂ) ^ s)) i

noncomputable def dirichlet_head_sum := (λ n : ℕ, λ s : ℂ, ∑ i in range n, (dirichlet_series f s i))

/- Should this be `real.log` or `complex.log`? -/
noncomputable def dderiv : arithmetic_function ℂ := {
  to_fun := λ n : ℕ, (f n) * (real.log n),
  map_zero' := by simp,
}

namespace dirichlet_series
localized "notation `D` := nat.arithmetic_function.dirichlet_series" in dirichlet_series
localized "notation `D'` := nat.arithmetic_function.dirichlet_series'" in dirichlet_series
localized "notation `S` := nat.arithmetic_function.dirichlet_head_sum" in dirichlet_series

/-! ### Definitions of convergence -/
/-- The Dirichlet series is convergent at a point -/
def convergent_at : Prop := nat_summable (D f s)

/-- The Dirichlet series is convergent in a right half-plane -/
def convergent : Prop := ∀ s : ℂ, r < s.re → convergent_at f s

/-- The Dirichlet series is absolutely convergent at a point -/
def abs_convergent_at : Prop := summable (D f s)

/-- The Dirichlet series is absolutely convergent in a right half-plane -/
def abs_convergent : Prop := ∀ s : ℂ, r < s.re → abs_convergent_at f s

/-- The traditional definition of absolutely convergent. Equivalent to our notion of
absolute convergence. See `abs_convergent_at_iff_norm_convergent_at` -/
def norm_convergent_at : Prop := summable (λ n : ℕ, ∥(D f s) n ∥)

/-- The traditional definition of absolutely convergent on a right half-plane.
Equivalent to our notion of
absolute convergence. See `abs_convergent_iff_norm_convergent` -/
def norm_convergent : Prop := ∀ s : ℂ, r < s.re → norm_convergent_at f s

/-- Uniform convergence on a closed ball around a point -/
def uniform_convergent_at : Prop := tendsto_uniformly_on (S f) (D' f) at_top (closed_ball s r)

/-- The set of all complex numbers strictly to the right of `r` -/
def half_plane : set ℂ := { z : ℂ | r < z.re }

/-- The set of all complex numbers at or to the right of `r` -/
def closed_half_plane : set ℂ := { z : ℂ | r ≤ z.re }

/-! ### Relationships between the various convergence modes

Currently enumerating the main theorems which will be turned into
many more useful lemmas later
-/

/-- The notion of norm convergence and absolute convergence are equivalent -/
lemma abs_convergent_at_iff_norm_convergent_at : abs_convergent_at f s ↔ norm_convergent_at f s :=
begin
  sorry,
end

/-- The notion of norm convergence and absolute convergence are equivalent, but
sometimes one may be easier to use than the other -/
lemma abs_convergent_iff_norm_convergent : abs_convergent f r ↔ norm_convergent f r :=
⟨
  λ h s hs, (abs_convergent_at_iff_norm_convergent_at f s).mp (h s hs),
  λ h s hs, (abs_convergent_at_iff_norm_convergent_at f s).mpr (h s hs),
⟩

/-- Convergence at a point implies convergence to the right of that point -/
lemma convergent_of_convergent_at
(hfs : convergent_at f s) :
convergent f s.re :=
begin
  sorry
end

/-- Norm convergence at a point implies norm convergence to the right of that point -/
lemma norm_convergent_of_norm_convergent_at
(hfs : norm_convergent_at f s) :
norm_convergent f s.re :=
begin
  intros t ht,
  refine summable_of_norm_bounded _ hfs _,
  unfold dirichlet_series,
  intros i,
  by_cases hi : i = 0,
  { simp [hi], },
  rw [real.norm_eq_abs, complex.norm_eq_abs, complex.norm_eq_abs, complex.abs_abs,
    complex.abs_div, complex.abs_div],
  refine div_le_div (complex.abs_nonneg _) rfl.le _ _,
  { rw complex.abs_pos,
    intros h,
    rw complex.cpow_eq_zero_iff at h,
    rcases h with ⟨hi', _⟩,
    norm_cast at hi', },
  { have : 1 ≤ i, exact one_le_iff_ne_zero.mpr hi,
    have : 0 < i, linarith [this],
    have aa : 0 < (i : ℝ), { norm_cast, exact this, },
    have bb : (i : ℂ) = ((i : ℝ) : ℂ), simp,
    rw bb,
    rw complex.abs_cpow_eq_rpow_re_of_pos aa,
    rw complex.abs_cpow_eq_rpow_re_of_pos aa,
    refine real.rpow_le_rpow_of_exponent_le _ ht.le,
    norm_cast,
    exact one_le_iff_ne_zero.mpr hi, },
end

/-- Absolute convergence at a point implies absolute convergence to the right of that point -/
lemma abs_convergent_of_abs_convergent_at
(hfs : abs_convergent_at f s) :
abs_convergent f s.re :=
begin
  rw abs_convergent_at_iff_norm_convergent_at at hfs,
  rw abs_convergent_iff_norm_convergent,
  exact norm_convergent_of_norm_convergent_at _ _ hfs,
end

/-- Convergence implies absolute convergence eventually -/
lemma abs_convergent_of_convergent
(hfs : convergent f r) :
abs_convergent f (r + 1) :=
begin
  sorry
end

lemma uniform_convergent_of_convergent
(hfs : convergent f r)
(hs : r < s.re)
:
uniform_convergent_at f s (s.re - r) :=
begin
  sorry,
end

/-! ### Proving convergence -/

lemma abs_convergent_of_eventually_bounded
(hf : ∀ᶠ (n : ℕ) in at_top, complex.abs (f n) ≤ r) : abs_convergent f 1 :=
begin
  sorry
end

lemma abs_convergent_of_bounded
(hf : ∀ (n : ℕ), complex.abs (f n) ≤ r) : abs_convergent f 1 :=
begin
  sorry,
end

/-! ### Differentiability and Convergence -/

/-- Convergence implies holomorphic on the open right half-plane -/
lemma derivative_of_convergent
(hfs : convergent f r)
(hs : r < s.re) :
has_deriv_at (D' f) (D' f.dderiv s) s :=
begin
  sorry,
end

/-- Convergence implies holomorphic on the open right half-plane -/
lemma differentiable_at_of_convergent
(hfs : convergent f r) (hs : r < s.re) :
differentiable_at ℂ (D' f) s :=
begin
  sorry,
end

/-- Holomorphic extension implies convergence -/
lemma convergent_of_differentiable_on
(hg : differentiable_on ℂ g $ half_plane r)
(hfg : ∀ (z : ℂ), z ∈ half_plane r → D' f z = g z) :
convergent f r :=
begin
  sorry,
end

/-! ### Important integrals -/

-- noncomputable def tmp_lo := (λ n : ℕ, (n : ℝ)⁻¹)
-- def tmp_hi := (λ n : ℕ, (n : ℝ))
-- noncomputable def tmp (S : ℝ → ℂ) := λ n : ℕ, (∫ (x : ℝ) in (tmp_lo n)..(tmp_hi n), S x)
-- noncomputable def imp_int_zero_inf (S : ℝ → ℂ) := lim at_top (tmp S)

-- lemma as_integral
-- (hfr : abs_convergent f r)
-- (hrs : r < s.re)
-- :
-- D' f s = s * ∫ x in set.Ioi (0 : ℝ), head_sum f ⌊x⌋₊ / x ^ (s + 1)
-- :=
-- begin
--   sorry
-- end

-- lemma useful_integral
-- {S : ℝ → ℂ}
-- (hbounded : ∀ (x : ℝ), complex.abs (S x) ≤ r)
-- (hint : ∀ (a b : ℝ), interval_integrable S real.measure_space.volume a b)
-- :
-- differentiable_on ℂ (λ z : ℂ, ∫ x in set.Ioi (0 : ℝ), (S x) * complex.exp (-z * s)) $ half_plane 0
-- :=
-- begin
--   sorry,
-- end

-- lemma useful_integral_diff

-- lemma norm_convergent_at.of_norm_convergent_at_re_le
-- (hfs : norm_convergent_at f s) (hst : s.re ≤ t.re) : norm_convergent_at f t :=
-- begin
--   unfold norm_convergent_at,
--   unfold dirichlet_series,
--   refine summable_of_norm_bounded _ hfs _,
--   unfold dirichlet_series,
--   intros i,
--   by_cases hi : i = 0,
--   { simp [hi], },
--   rw real.norm_eq_abs, rw complex.norm_eq_abs, rw complex.norm_eq_abs, rw complex.abs_abs,
--   rw complex.abs_div, rw complex.abs_div,
--   apply div_le_div,
--   exact complex.abs_nonneg _,
--   exact rfl.le,
--   rw complex.abs_pos,
--   intros h,
--   rw complex.cpow_eq_zero_iff at h,
--   rcases h with ⟨hi', _⟩,
--   norm_cast at hi',
--   have : 1 ≤ i, exact one_le_iff_ne_zero.mpr hi,
--   have : 0 < i, linarith [this],
--   have aa : 0 < (i : ℝ), { norm_cast, exact this, },
--   have bb : (i : ℂ) = ((i : ℝ) : ℂ), simp,
--   rw bb,
--   rw complex.abs_cpow_eq_rpow_re_of_pos aa,
--   rw complex.abs_cpow_eq_rpow_re_of_pos aa,
--   refine real.rpow_le_rpow_of_exponent_le _ hst,
--   norm_cast,
--   exact one_le_iff_ne_zero.mpr hi,
-- end

-- lemma abs_convergent_at.of_abs_convergent_at_re_le
-- (hfs : abs_convergent_at f s) (hst : s.re ≤ t.re) : abs_convergent_at f t :=
-- begin
--   rw abs_convergent_at_iff_norm_convergent_at at hfs,
--   rw abs_convergent_at_iff_norm_convergent_at,
--   exact hfs.of_norm_convergent_at_re_le f s t hst,
-- end

/-! ### Special functions -/

theorem zeta.abs_convergent : abs_convergent ζ 1 :=
begin
  refine abs_convergent_of_bounded _ 1 _,
  intros i,
  simp [zeta],
  by_cases hi : i = 0,
  simp [hi],
  simp [hi],
end

theorem moebius.abs_convergent : abs_convergent μ 1 :=
begin
  refine abs_convergent_of_bounded _ 1 _,
  intros i,
  simp [moebius],
  by_cases hi : squarefree i,
  simp [hi],
  simp [hi],
end

end dirichlet_series
end arithmetic_function
end nat


/-
Copyright (c) 2019 Reid Barton. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Reid Barton, Johan Commelin
-/
import category_theory.adjunction.basic
import category_theory.limits.creates

/-!
# Adjunctions and limits

A left adjoint preserves colimits (`category_theory.adjunction.left_adjoint_preserves_colimits`),
and a right adjoint preserves limits (`category_theory.adjunction.right_adjoint_preserves_limits`).

Equivalences create and reflect (co)limits.
(`category_theory.adjunction.is_equivalence_creates_limits`,
`category_theory.adjunction.is_equivalence_creates_colimits`,
`category_theory.adjunction.is_equivalence_reflects_limits`,
`category_theory.adjunction.is_equivalence_reflects_colimits`,)

In `category_theory.adjunction.cocones_iso` we show that
when `F ⊣ G`,
the functor associating to each `Y` the cocones over `K ⋙ F` with cone point `Y`
is naturally isomorphic to
the functor associating to each `Y` the cocones over `K` with cone point `G.obj Y`.
-/

open opposite

namespace category_theory.adjunction
open category_theory
open category_theory.functor
open category_theory.limits

universes u₁ u₂ v

variables {C : Type u₁} [category.{v} C] {D : Type u₂} [category.{v} D]

variables {F : C ⥤ D} {G : D ⥤ C} (adj : F ⊣ G)
include adj

section preservation_colimits
variables {J : Type v} [small_category J] (K : J ⥤ C)

/--
The right adjoint of `cocones.functoriality K F : cocone K ⥤ cocone (K ⋙ F)`.

Auxiliary definition for `functoriality_is_left_adjoint`.
-/
def functoriality_right_adjoint : cocone (K ⋙ F) ⥤ cocone K :=
(cocones.functoriality _ G) ⋙
  (cocones.precompose (K.right_unitor.inv ≫ (whisker_left K adj.unit) ≫ (associator _ _ _).inv))

local attribute [reducible] functoriality_right_adjoint

/--
The unit for the adjunction for `cocones.functoriality K F : cocone K ⥤ cocone (K ⋙ F)`.

Auxiliary definition for `functoriality_is_left_adjoint`.
-/
@[simps] def functoriality_unit :
  𝟭 (cocone K) ⟶ cocones.functoriality _ F ⋙ functoriality_right_adjoint adj K :=
{ app := λ c, { hom := adj.unit.app c.X } }

/--
The counit for the adjunction for `cocones.functoriality K F : cocone K ⥤ cocone (K ⋙ F)`.

Auxiliary definition for `functoriality_is_left_adjoint`.
-/
@[simps] def functoriality_counit :
  functoriality_right_adjoint adj K ⋙ cocones.functoriality _ F ⟶ 𝟭 (cocone (K ⋙ F)) :=
{ app := λ c, { hom := adj.counit.app c.X } }

/-- The functor `cocones.functoriality K F : cocone K ⥤ cocone (K ⋙ F)` is a left adjoint. -/
def functoriality_is_left_adjoint :
  is_left_adjoint (cocones.functoriality K F) :=
{ right := functoriality_right_adjoint adj K,
  adj := mk_of_unit_counit
  { unit := functoriality_unit adj K,
    counit := functoriality_counit adj K } }

/--
A left adjoint preserves colimits.

See https://stacks.math.columbia.edu/tag/0038.
-/
def left_adjoint_preserves_colimits : preserves_colimits F :=
{ preserves_colimits_of_shape := λ J 𝒥,
  { preserves_colimit := λ F,
    by exactI
    { preserves := λ c hc, is_colimit.iso_unique_cocone_morphism.inv
        (λ s, @equiv.unique _ _ (is_colimit.iso_unique_cocone_morphism.hom hc _)
          (((adj.functoriality_is_left_adjoint _).adj).hom_equiv _ _)) } } }.

omit adj

@[priority 100] -- see Note [lower instance priority]
instance is_equivalence_preserves_colimits (E : C ⥤ D) [is_equivalence E] : preserves_colimits E :=
left_adjoint_preserves_colimits E.adjunction

@[priority 100] -- see Note [lower instance priority]
instance is_equivalence_reflects_colimits (E : D ⥤ C) [is_equivalence E] : reflects_colimits E :=
{ reflects_colimits_of_shape := λ J 𝒥, by exactI
  { reflects_colimit := λ K,
    { reflects := λ c t,
      begin
        have l := (is_colimit_of_preserves E.inv t).map_cocone_equiv E.as_equivalence.unit_iso.symm,
        refine (((is_colimit.precompose_inv_equiv K.right_unitor _).symm) l).of_iso_colimit _,
        tidy,
      end } } }

@[priority 100] -- see Note [lower instance priority]
instance is_equivalence_creates_colimits (H : D ⥤ C) [is_equivalence H] : creates_colimits H :=
{ creates_colimits_of_shape := λ J 𝒥, by exactI
  { creates_colimit := λ F,
    { lifts := λ c t,
      { lifted_cocone := H.map_cocone_inv c,
        valid_lift := H.map_cocone_map_cocone_inv c } } } }

-- verify the preserve_colimits instance works as expected:
example (E : C ⥤ D) [is_equivalence E]
  (c : cocone K) (h : is_colimit c) : is_colimit (E.map_cocone c) :=
preserves_colimit.preserves h

lemma has_colimit_comp_equivalence (E : C ⥤ D) [is_equivalence E] [has_colimit K] :
  has_colimit (K ⋙ E) :=
has_colimit.mk
{ cocone := E.map_cocone (colimit.cocone K),
  is_colimit := preserves_colimit.preserves (colimit.is_colimit K) }

lemma has_colimit_of_comp_equivalence (E : C ⥤ D) [is_equivalence E] [has_colimit (K ⋙ E)] :
  has_colimit K :=
@has_colimit_of_iso _ _ _ _ (K ⋙ E ⋙ inv E) K
(@has_colimit_comp_equivalence _ _ _ _ _ _ (K ⋙ E) (inv E) _ _)
((functor.right_unitor _).symm ≪≫ iso_whisker_left K (E.as_equivalence.unit_iso))

/-- Transport a `has_colimits_of_shape` instance across an equivalence. -/
lemma has_colimits_of_shape_of_equivalence (E : C ⥤ D) [is_equivalence E]
  [has_colimits_of_shape J D] : has_colimits_of_shape J C :=
⟨λ F, by exactI has_colimit_of_comp_equivalence F E⟩

/-- Transport a `has_colimits` instance across an equivalence. -/
lemma has_colimits_of_equivalence (E : C ⥤ D) [is_equivalence E] [has_colimits D] :
  has_colimits C :=
⟨λ J hJ, by exactI has_colimits_of_shape_of_equivalence E⟩

end preservation_colimits

section preservation_limits
variables {J : Type v} [small_category J] (K : J ⥤ D)

/--
The left adjoint of `cones.functoriality K G : cone K ⥤ cone (K ⋙ G)`.

Auxiliary definition for `functoriality_is_right_adjoint`.
-/
def functoriality_left_adjoint : cone (K ⋙ G) ⥤ cone K :=
(cones.functoriality _ F) ⋙ (cones.postcompose
    ((associator _ _ _).hom ≫ (whisker_left K adj.counit) ≫ K.right_unitor.hom))

local attribute [reducible] functoriality_left_adjoint

/--
The unit for the adjunction for`cones.functoriality K G : cone K ⥤ cone (K ⋙ G)`.

Auxiliary definition for `functoriality_is_right_adjoint`.
-/
@[simps] def functoriality_unit' :
  𝟭 (cone (K ⋙ G)) ⟶ functoriality_left_adjoint adj K ⋙ cones.functoriality _ G :=
{ app := λ c, { hom := adj.unit.app c.X, } }

/--
The counit for the adjunction for`cones.functoriality K G : cone K ⥤ cone (K ⋙ G)`.

Auxiliary definition for `functoriality_is_right_adjoint`.
-/
@[simps] def functoriality_counit' :
  cones.functoriality _ G ⋙ functoriality_left_adjoint adj K ⟶ 𝟭 (cone K) :=
{ app := λ c, { hom := adj.counit.app c.X, } }

/-- The functor `cones.functoriality K G : cone K ⥤ cone (K ⋙ G)` is a right adjoint. -/
def functoriality_is_right_adjoint :
  is_right_adjoint (cones.functoriality K G) :=
{ left := functoriality_left_adjoint adj K,
  adj := mk_of_unit_counit
  { unit := functoriality_unit' adj K,
    counit := functoriality_counit' adj K } }

/--
A right adjoint preserves limits.

See https://stacks.math.columbia.edu/tag/0038.
-/
def right_adjoint_preserves_limits : preserves_limits G :=
{ preserves_limits_of_shape := λ J 𝒥,
  { preserves_limit := λ K,
    by exactI
    { preserves := λ c hc, is_limit.iso_unique_cone_morphism.inv
        (λ s, @equiv.unique _ _ (is_limit.iso_unique_cone_morphism.hom hc _)
          (((adj.functoriality_is_right_adjoint _).adj).hom_equiv _ _).symm) } } }.

omit adj

@[priority 100] -- see Note [lower instance priority]
instance is_equivalence_preserves_limits (E : D ⥤ C) [is_equivalence E] : preserves_limits E :=
right_adjoint_preserves_limits E.inv.adjunction

@[priority 100] -- see Note [lower instance priority]
instance is_equivalence_reflects_limits (E : D ⥤ C) [is_equivalence E] : reflects_limits E :=
{ reflects_limits_of_shape := λ J 𝒥, by exactI
  { reflects_limit := λ K,
    { reflects := λ c t,
      begin
        have := (is_limit_of_preserves E.inv t).map_cone_equiv E.as_equivalence.unit_iso.symm,
        refine (((is_limit.postcompose_hom_equiv K.left_unitor _).symm) this).of_iso_limit _,
        tidy,
      end } } }

@[priority 100] -- see Note [lower instance priority]
instance is_equivalence_creates_limits (H : D ⥤ C) [is_equivalence H] : creates_limits H :=
{ creates_limits_of_shape := λ J 𝒥, by exactI
  { creates_limit := λ F,
    { lifts := λ c t,
      { lifted_cone := H.map_cone_inv c,
        valid_lift := H.map_cone_map_cone_inv c } } } }

-- verify the preserve_limits instance works as expected:
example (E : D ⥤ C) [is_equivalence E]
  (c : cone K) [h : is_limit c] : is_limit (E.map_cone c) :=
preserves_limit.preserves h

lemma has_limit_comp_equivalence (E : D ⥤ C) [is_equivalence E] [has_limit K] :
  has_limit (K ⋙ E) :=
has_limit.mk
{ cone := E.map_cone (limit.cone K),
  is_limit := preserves_limit.preserves (limit.is_limit K) }

lemma has_limit_of_comp_equivalence (E : D ⥤ C) [is_equivalence E] [has_limit (K ⋙ E)] :
  has_limit K :=
@has_limit_of_iso _ _ _ _ (K ⋙ E ⋙ inv E) K
(@has_limit_comp_equivalence _ _ _ _ _ _ (K ⋙ E) (inv E) _ _)
((iso_whisker_left K E.as_equivalence.unit_iso.symm) ≪≫ (functor.right_unitor _))

/-- Transport a `has_limits_of_shape` instance across an equivalence. -/
lemma has_limits_of_shape_of_equivalence (E : D ⥤ C) [is_equivalence E] [has_limits_of_shape J C] :
  has_limits_of_shape J D :=
⟨λ F, by exactI has_limit_of_comp_equivalence F E⟩

/-- Transport a `has_limits` instance across an equivalence. -/
lemma has_limits_of_equivalence (E : D ⥤ C) [is_equivalence E] [has_limits C] : has_limits D :=
⟨λ J hJ, by exactI has_limits_of_shape_of_equivalence E⟩

end preservation_limits

/-- auxiliary construction for `cocones_iso` -/
@[simps]
def cocones_iso_component_hom {J : Type v} [small_category J] {K : J ⥤ C}
  (Y : D) (t : ((cocones J D).obj (op (K ⋙ F))).obj Y) :
  (G ⋙ (cocones J C).obj (op K)).obj Y :=
{ app := λ j, (adj.hom_equiv (K.obj j) Y) (t.app j),
  naturality' := λ j j' f, by { erw [← adj.hom_equiv_naturality_left, t.naturality], dsimp, simp } }

/-- auxiliary construction for `cocones_iso` -/
@[simps]
def cocones_iso_component_inv {J : Type v} [small_category J] {K : J ⥤ C}
  (Y : D) (t : (G ⋙ (cocones J C).obj (op K)).obj Y) :
  ((cocones J D).obj (op (K ⋙ F))).obj Y :=
{ app := λ j, (adj.hom_equiv (K.obj j) Y).symm (t.app j),
  naturality' := λ j j' f,
  begin
    erw [← adj.hom_equiv_naturality_left_symm, ← adj.hom_equiv_naturality_right_symm, t.naturality],
    dsimp, simp
  end }

/--
When `F ⊣ G`,
the functor associating to each `Y` the cocones over `K ⋙ F` with cone point `Y`
is naturally isomorphic to
the functor associating to each `Y` the cocones over `K` with cone point `G.obj Y`.
-/
-- Note: this is natural in K, but we do not yet have the tools to formulate that.
def cocones_iso {J : Type v} [small_category J] {K : J ⥤ C} :
  (cocones J D).obj (op (K ⋙ F)) ≅ G ⋙ ((cocones J C).obj (op K)) :=
nat_iso.of_components (λ Y,
{ hom := cocones_iso_component_hom adj Y,
  inv := cocones_iso_component_inv adj Y, })
(by tidy)

/-- auxiliary construction for `cones_iso` -/
@[simps]
def cones_iso_component_hom {J : Type v} [small_category J] {K : J ⥤ D}
  (X : Cᵒᵖ) (t : (functor.op F ⋙ (cones J D).obj K).obj X) :
  ((cones J C).obj (K ⋙ G)).obj X :=
{ app := λ j, (adj.hom_equiv (unop X) (K.obj j)) (t.app j),
  naturality' := λ j j' f,
  begin
    erw [← adj.hom_equiv_naturality_right, ← t.naturality, category.id_comp, category.id_comp],
    refl
  end }

/-- auxiliary construction for `cones_iso` -/
@[simps]
def cones_iso_component_inv {J : Type v} [small_category J] {K : J ⥤ D}
  (X : Cᵒᵖ) (t : ((cones J C).obj (K ⋙ G)).obj X) :
  (functor.op F ⋙ (cones J D).obj K).obj X :=
{ app := λ j, (adj.hom_equiv (unop X) (K.obj j)).symm (t.app j),
  naturality' := λ j j' f,
  begin
    erw [← adj.hom_equiv_naturality_right_symm, ← t.naturality, category.id_comp, category.id_comp]
  end }

-- Note: this is natural in K, but we do not yet have the tools to formulate that.
/--
When `F ⊣ G`,
the functor associating to each `X` the cones over `K` with cone point `F.op.obj X`
is naturally isomorphic to
the functor associating to each `X` the cones over `K ⋙ G` with cone point `X`.
-/
def cones_iso {J : Type v} [small_category J] {K : J ⥤ D} :
  F.op ⋙ ((cones J D).obj K) ≅ (cones J C).obj (K ⋙ G) :=
nat_iso.of_components (λ X,
{ hom := cones_iso_component_hom adj X,
  inv := cones_iso_component_inv adj X, } )
(by tidy)

end category_theory.adjunction


import category_theory.base
import category_theory.replete

open category_theory
open category_theory.category
local notation f ` ∘ `:80 g:80 := g ≫ f

universes v u

namespace homotopy_theory.weak_equivalences

class has_weak_equivalences (C : Type u) [category C] :=
(is_weq : Π ⦃a b : C⦄, (a ⟶ b) → Prop)

def is_weq {C : Type u} [category C] [has_weak_equivalences C] ⦃a b : C⦄ (f : a ⟶ b) :=
has_weak_equivalences.is_weq f

-- TODO: should this be a Prop mix-in?
class category_with_weak_equivalences (C : Type u) [category.{v} C]
  extends has_weak_equivalences C :=
[weq_replete_wide : replete_wide_subcategory.{v} C is_weq]
(weq_of_comp_weq_left : ∀ ⦃a b c : C⦄ {f : a ⟶ b} {g : b ⟶ c},
  is_weq f → is_weq (g ∘ f) → is_weq g)
(weq_of_comp_weq_right : ∀ ⦃a b c : C⦄ {f : a ⟶ b} {g : b ⟶ c},
  is_weq g → is_weq (g ∘ f) → is_weq f)

instance (C : Type u) [category.{v} C] [category_with_weak_equivalences C] :
  replete_wide_subcategory.{v} C is_weq :=
category_with_weak_equivalences.weq_replete_wide

section
variables {C : Type u} [category.{v} C] [category_with_weak_equivalences C]

lemma weq_id (a : C) : is_weq (𝟙 a) := mem_id a
lemma weq_comp {a b c : C} {f : a ⟶ b} {g : b ⟶ c} :
  is_weq f → is_weq g → is_weq (g ∘ f) := mem_comp
lemma weq_iso {a b : C} (i : a ≅ b) : is_weq i.hom := mem_iso i

lemma weq_iff_weq_inv {a b : C} {f : a ⟶ b} {g : b ⟶ a} (h : f ≫ g = 𝟙 _) :
  is_weq f ↔ is_weq g :=
begin
  split; intro H;
  { have : is_weq (g ∘ f) := by convert weq_id _,
    apply category_with_weak_equivalences.weq_of_comp_weq_left H this <|>
    apply category_with_weak_equivalences.weq_of_comp_weq_right H this }
end

end

-- The two-out-of-six property.
class homotopical_category (C : Type u) [category.{v} C]
  extends category_with_weak_equivalences C :=
(two_out_of_six : ∀ ⦃a b c d : C⦄ {f : a ⟶ b} {g : b ⟶ c} {h : c ⟶ d},
  is_weq (h ∘ g) → is_weq (g ∘ f) → is_weq g)

section
variables {C : Type u} [category.{v} C] [homotopical_category C]

lemma weq_two_out_of_six_g {a b c d : C} {f : a ⟶ b} {g : b ⟶ c} {h : c ⟶ d}
  (hg : is_weq (h ∘ g)) (gf : is_weq (g ∘ f)) : is_weq g :=
homotopical_category.two_out_of_six hg gf

lemma weq_two_out_of_six_f {a b c d : C} {f : a ⟶ b} {g : b ⟶ c} {h : c ⟶ d}
  (hg : is_weq (h ∘ g)) (gf : is_weq (g ∘ f)) : is_weq f :=
have wg : is_weq g := weq_two_out_of_six_g hg gf,
category_with_weak_equivalences.weq_of_comp_weq_right wg gf

lemma weq_two_out_of_six_h {a b c d : C} {f : a ⟶ b} {g : b ⟶ c} {h : c ⟶ d}
  (hg : is_weq (h ∘ g)) (gf : is_weq (g ∘ f)) : is_weq h :=
have wg : is_weq g := weq_two_out_of_six_g hg gf,
category_with_weak_equivalences.weq_of_comp_weq_left wg hg
end

section isomorphisms
variables {C : Type u} [category.{v} C]

def is_iso ⦃a b : C⦄ (f : a ⟶ b) : Prop := ∃ i : a ≅ b, i.hom = f

lemma iso_iso ⦃a b : C⦄ (i : a ≅ b) : is_iso i.hom := ⟨i, rfl⟩
lemma iso_comp ⦃a b c : C⦄ {f : a ⟶ b} {g : b ⟶ c} :
  is_iso f → is_iso g → is_iso (g ∘ f) :=
assume ⟨i, hi⟩ ⟨j, hj⟩, ⟨i.trans j, by rw [←hi, ←hj]; refl⟩

lemma iso_of_comp_iso_left ⦃a b c : C⦄ {f : a ⟶ b} {g : b ⟶ c} :
  is_iso f → is_iso (g ∘ f) → is_iso g :=
assume ⟨i, hi⟩ ⟨j, hj⟩,
  ⟨i.symm.trans j, show j.hom ∘ i.inv = g, by rw [hj, ←hi]; simp⟩
lemma iso_of_comp_iso_right ⦃a b c : C⦄ {f : a ⟶ b} {g : b ⟶ c} :
  is_iso g → is_iso (g ∘ f) → is_iso f :=
assume ⟨i, hi⟩ ⟨j, hj⟩,
  ⟨j.trans i.symm, show i.inv ∘ j.hom = f, by rw [hj, ←hi]; simp⟩

lemma iso_two_out_of_six ⦃a b c d : C⦄ {f : a ⟶ b} {g : b ⟶ c} {h : c ⟶ d} :
  is_iso (h ∘ g) → is_iso (g ∘ f) → is_iso g :=
assume ⟨i, hi⟩ ⟨j, hj⟩,
  let g' := i.inv ∘ h in
  have g'g : g' ∘ g = 𝟙 _, by rw [←assoc, ←hi]; simp,
  let g'' := f ∘ j.inv in
  have gg'' : g ∘ g'' = 𝟙 _, by rw [assoc, ←hj]; simp,
  have g' = g'', from calc
    g' = g' ∘ (g ∘ g'')  : by rw gg''; simp
   ... = (g' ∘ g) ∘ g''  : by simp
   ... = g''             : by rw g'g; simp; refl,
  ⟨⟨g, g', g'g, by rw this; exact gg''⟩, rfl⟩

instance is_iso.replete_wide_subcategory : replete_wide_subcategory.{v} C is_iso :=
replete_wide_subcategory.mk' iso_iso iso_comp

def isomorphisms_as_weak_equivalences : category_with_weak_equivalences C :=
{ is_weq := is_iso,
  weq_of_comp_weq_left := iso_of_comp_iso_left,
  weq_of_comp_weq_right := iso_of_comp_iso_right }

def isomorphisms_as_homotopical_category : homotopical_category C :=
{ two_out_of_six := iso_two_out_of_six,
  .. isomorphisms_as_weak_equivalences }

end isomorphisms

section preimage
-- TODO: generalize to different universes?
variables {C D : Type u} [category.{v} C] [category.{v} D]
variables (F : C ↝ D)

def preimage_weq (weqD : has_weak_equivalences D) : has_weak_equivalences C :=
{ is_weq := λ a b f, is_weq (F &> f) }

instance preimage_weq.replete_wide_subcategory [weqD : category_with_weak_equivalences D] :
  replete_wide_subcategory.{v} C (preimage_weq F weqD.to_has_weak_equivalences).is_weq :=
replete_wide_subcategory.mk'
    (λ a b i, weq_iso (F.map_iso i))
    (λ a b c f g hf hg, show is_weq (F &> (g ∘ f)),
      by rw F.map_comp; exact weq_comp hf hg)

def preimage_with_weak_equivalences [weqD : category_with_weak_equivalences D] :
  category_with_weak_equivalences C :=
{ to_has_weak_equivalences := preimage_weq F weqD.to_has_weak_equivalences,
  weq_of_comp_weq_left := λ a b c f g hf hgf, begin
    change is_weq (F &> (g ∘ f)) at hgf, rw F.map_comp at hgf,
    exact category_with_weak_equivalences.weq_of_comp_weq_left hf hgf
  end,
  weq_of_comp_weq_right := λ a b c f g hg hgf, begin
    change is_weq (F &> (g ∘ f)) at hgf, rw F.map_comp at hgf,
    exact category_with_weak_equivalences.weq_of_comp_weq_right hg hgf
  end }

end preimage

end homotopy_theory.weak_equivalences


/-
Copyright (c) 2018 Scott Morrison. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Scott Morrison, Bhavik Mehta
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.category_theory.limits.limits
import Mathlib.category_theory.discrete_category
import Mathlib.PostPort

universes v u u_1 u₂

namespace Mathlib

namespace category_theory.limits

-- We don't need an analogue of `pair` (for binary products), `parallel_pair` (for equalizers),

-- or `(co)span`, since we already have `discrete.functor`.

/-- A fan over `f : β → C` consists of a collection of maps from an object `P` to every `f b`. -/
/-- A cofan over `f : β → C` consists of a collection of maps from every `f b` to an object `P`. -/
def fan {β : Type v} {C : Type u} [category C] (f : β → C) :=
  cone (discrete.functor f)

def cofan {β : Type v} {C : Type u} [category C] (f : β → C) :=
  cocone (discrete.functor f)

/-- A fan over `f : β → C` consists of a collection of maps from an object `P` to every `f b`. -/
@[simp] theorem fan.mk_X {β : Type v} {C : Type u} [category C] {f : β → C} (P : C) (p : (b : β) → P ⟶ f b) : cone.X (fan.mk P p) = P :=
  Eq.refl (cone.X (fan.mk P p))

/-- A cofan over `f : β → C` consists of a collection of maps from every `f b` to an object `P`. -/
@[simp] theorem cofan.mk_X {β : Type v} {C : Type u} [category C] {f : β → C} (P : C) (p : (b : β) → f b ⟶ P) : cocone.X (cofan.mk P p) = P :=
  Eq.refl (cocone.X (cofan.mk P p))

/-- An abbreviation for `has_limit (discrete.functor f)`. -/
def has_product {β : Type v} {C : Type u} [category C] (f : β → C) :=
  has_limit (discrete.functor f)

/-- An abbreviation for `has_colimit (discrete.functor f)`. -/
def has_coproduct {β : Type v} {C : Type u} [category C] (f : β → C) :=
  has_colimit (discrete.functor f)

/-- An abbreviation for `has_limits_of_shape (discrete f)`. -/
/-- An abbreviation for `has_colimits_of_shape (discrete f)`. -/
def has_products_of_shape (β : Type v) (C : Type u_1) [category C] :=
  has_limits_of_shape (discrete β)

def has_coproducts_of_shape (β : Type v) (C : Type u_1) [category C] :=
  has_colimits_of_shape (discrete β)

/-- `pi_obj f` computes the product of a family of elements `f`. (It is defined as an abbreviation
   for `limit (discrete.functor f)`, so for most facts about `pi_obj f`, you will just use general facts
   about limits.) -/
/-- `sigma_obj f` computes the coproduct of a family of elements `f`. (It is defined as an abbreviation
def pi_obj {β : Type v} {C : Type u} [category C] (f : β → C) [has_product f] : C :=
  limit (discrete.functor f)

   for `colimit (discrete.functor f)`, so for most facts about `sigma_obj f`, you will just use general facts
   about colimits.) -/
def sigma_obj {β : Type v} {C : Type u} [category C] (f : β → C) [has_coproduct f] : C :=
  colimit (discrete.functor f)

prefix:20 "∏ " => Mathlib.category_theory.limits.pi_obj

prefix:20 "∐ " => Mathlib.category_theory.limits.sigma_obj

/-- The `b`-th projection from the pi object over `f` has the form `∏ f ⟶ f b`. -/
def pi.π {β : Type v} {C : Type u} [category C] (f : β → C) [has_product f] (b : β) : ∏ f ⟶ f b :=
  limit.π (discrete.functor f) b

/-- The `b`-th inclusion into the sigma object over `f` has the form `f b ⟶ ∐ f`. -/
def sigma.ι {β : Type v} {C : Type u} [category C] (f : β → C) [has_coproduct f] (b : β) : f b ⟶ ∐ f :=
  colimit.ι (discrete.functor f) b

/-- The fan constructed of the projections from the product is limiting. -/
def product_is_product {β : Type v} {C : Type u} [category C] (f : β → C) [has_product f] : is_limit (fan.mk (∏ f) (pi.π f)) :=
  is_limit.of_iso_limit (limit.is_limit (discrete.functor f))
    (cones.ext (iso.refl (cone.X (limit.cone (discrete.functor fun (b : β) => f b)))) sorry)

/-- A collection of morphisms `P ⟶ f b` induces a morphism `P ⟶ ∏ f`. -/
def pi.lift {β : Type v} {C : Type u} [category C] {f : β → C} [has_product f] {P : C} (p : (b : β) → P ⟶ f b) : P ⟶ ∏ f :=
  limit.lift (discrete.functor fun (b : β) => f b) (fan.mk P p)

/-- A collection of morphisms `f b ⟶ P` induces a morphism `∐ f ⟶ P`. -/
def sigma.desc {β : Type v} {C : Type u} [category C] {f : β → C} [has_coproduct f] {P : C} (p : (b : β) → f b ⟶ P) : ∐ f ⟶ P :=
  colimit.desc (discrete.functor fun (b : β) => f b) (cofan.mk P p)

/--
Construct a morphism between categorical products (indexed by the same type)
from a family of morphisms between the factors.
-/
def pi.map {β : Type v} {C : Type u} [category C] {f : β → C} {g : β → C} [has_product f] [has_product g] (p : (b : β) → f b ⟶ g b) : ∏ f ⟶ ∏ g :=
  lim_map (discrete.nat_trans p)

/--
Construct an isomorphism between categorical products (indexed by the same type)
from a family of isomorphisms between the factors.
-/
def pi.map_iso {β : Type v} {C : Type u} [category C] {f : β → C} {g : β → C} [has_products_of_shape β C] (p : (b : β) → f b ≅ g b) : ∏ f ≅ ∏ g :=
  functor.map_iso lim (discrete.nat_iso p)

/--
Construct a morphism between categorical coproducts (indexed by the same type)
from a family of morphisms between the factors.
-/
def sigma.map {β : Type v} {C : Type u} [category C] {f : β → C} {g : β → C} [has_coproduct f] [has_coproduct g] (p : (b : β) → f b ⟶ g b) : ∐ f ⟶ ∐ g :=
  colim_map (discrete.nat_trans p)

/--
Construct an isomorphism between categorical coproducts (indexed by the same type)
from a family of isomorphisms between the factors.
-/
def sigma.map_iso {β : Type v} {C : Type u} [category C] {f : β → C} {g : β → C} [has_coproducts_of_shape β C] (p : (b : β) → f b ≅ g b) : ∐ f ≅ ∐ g :=
  functor.map_iso colim (discrete.nat_iso p)

-- TODO: show this is an iso iff G preserves the product of f.

/-- The comparison morphism for the product of `f`. -/
def pi_comparison {β : Type v} {C : Type u} [category C] {D : Type u₂} [category D] (G : C ⥤ D) (f : β → C) [has_product f] [has_product fun (b : β) => functor.obj G (f b)] : functor.obj G (∏ f) ⟶ ∏ fun (b : β) => functor.obj G (f b) :=
  pi.lift fun (b : β) => functor.map G (pi.π f b)

@[simp] theorem pi_comparison_comp_π_assoc {β : Type v} {C : Type u} [category C] {D : Type u₂} [category D] (G : C ⥤ D) (f : β → C) [has_product f] [has_product fun (b : β) => functor.obj G (f b)] (b : β) {X' : D} (f' : functor.obj G (f b) ⟶ X') : pi_comparison G f ≫ pi.π (fun (b : β) => functor.obj G (f b)) b ≫ f' = functor.map G (pi.π f b) ≫ f' := sorry

@[simp] theorem map_lift_pi_comparison_assoc {β : Type v} {C : Type u} [category C] {D : Type u₂} [category D] (G : C ⥤ D) (f : β → C) [has_product f] [has_product fun (b : β) => functor.obj G (f b)] (P : C) (g : (j : β) → P ⟶ f j) {X' : D} (f' : (∏ fun (b : β) => functor.obj G (f b)) ⟶ X') : functor.map G (pi.lift g) ≫ pi_comparison G f ≫ f' = (pi.lift fun (j : β) => functor.map G (g j)) ≫ f' := sorry

-- TODO: show this is an iso iff G preserves the coproduct of f.

/-- The comparison morphism for the coproduct of `f`. -/
def sigma_comparison {β : Type v} {C : Type u} [category C] {D : Type u₂} [category D] (G : C ⥤ D) (f : β → C) [has_coproduct f] [has_coproduct fun (b : β) => functor.obj G (f b)] : (∐ fun (b : β) => functor.obj G (f b)) ⟶ functor.obj G (∐ f) :=
  sigma.desc fun (b : β) => functor.map G (sigma.ι f b)

@[simp] theorem ι_comp_sigma_comparison_assoc {β : Type v} {C : Type u} [category C] {D : Type u₂} [category D] (G : C ⥤ D) (f : β → C) [has_coproduct f] [has_coproduct fun (b : β) => functor.obj G (f b)] (b : β) {X' : D} (f' : functor.obj G (∐ f) ⟶ X') : sigma.ι (fun (b : β) => functor.obj G (f b)) b ≫ sigma_comparison G f ≫ f' = functor.map G (sigma.ι f b) ≫ f' := sorry

@[simp] theorem sigma_comparison_map_desc {β : Type v} {C : Type u} [category C] {D : Type u₂} [category D] (G : C ⥤ D) (f : β → C) [has_coproduct f] [has_coproduct fun (b : β) => functor.obj G (f b)] (P : C) (g : (j : β) → f j ⟶ P) : sigma_comparison G f ≫ functor.map G (sigma.desc g) = sigma.desc fun (j : β) => functor.map G (g j) := sorry

/-- An abbreviation for `Π J, has_limits_of_shape (discrete J) C` -/
/-- An abbreviation for `Π J, has_colimits_of_shape (discrete J) C` -/
def has_products (C : Type u) [category C] :=
  ∀ (J : Type v), has_limits_of_shape (discrete J) C

def has_coproducts (C : Type u) [category C] :=
  ∀ (J : Type v), has_colimits_of_shape (discrete J) C


/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Meta.AppBuilder
import Lean.Meta.Instances

namespace Lean.Meta

/-- Create `SizeOf` local instances for applicable parameters, and execute `k` using them. -/
private partial def mkLocalInstances {α} (params : Array Expr) (k : Array Expr → MetaM α) : MetaM α :=
  loop 0 #[]
where
  loop (i : Nat) (insts : Array Expr) : MetaM α := do
    if i < params.size then
      let param := params[i]
      let paramType ← inferType param
      let instType? ← forallTelescopeReducing paramType fun xs _ => do
        let type := mkAppN param xs
        try
          let sizeOf ← mkAppM `SizeOf #[type]
          let instType ← mkForallFVars xs sizeOf
          return some instType
        catch _ =>
          return none
      match instType? with
      | none => loop (i+1) insts
      | some instType =>
        let instName ← mkFreshUserName `inst
        withLocalDecl instName BinderInfo.instImplicit instType fun inst =>
          loop (i+1) (insts.push inst)
    else
      k insts

/--
  Return `some x` if `fvar` has type of the form `... -> motive ... fvar` where `motive` in `motiveFVars`.
  That is, `x` "produces" one of the recursor motives.
-/
private def isInductiveHypothesis? (motiveFVars : Array Expr) (fvar : Expr) : MetaM (Option Expr) := do
  forallTelescopeReducing (← inferType fvar) fun _ type =>
    if type.isApp && motiveFVars.contains type.getAppFn then
      return some type.appArg!
    else
      return none

private def isInductiveHypothesis (motiveFVars : Array Expr) (fvar : Expr) : MetaM Bool :=
  return (← isInductiveHypothesis? motiveFVars fvar).isSome

/--
  Let `motiveFVars` be free variables for each motive in a kernel recursor, and `minorFVars` the free variables for a minor premise.
  Then, return `some idx` if `minorFVars[idx]` has a type of the form `... -> motive ... fvar` for some `motive` in `motiveFVars`.
-/
private def isRecField? (motiveFVars : Array Expr) (minorFVars : Array Expr) (fvar : Expr) : MetaM (Option Nat) := do
  let mut idx := 0
  for minorFVar in minorFVars do
    if let some fvar' ← isInductiveHypothesis? motiveFVars minorFVar then
      if fvar == fvar' then
        return some idx
    idx := idx + 1
  return none

private partial def mkSizeOfMotives {α} (motiveFVars : Array Expr) (k : Array Expr → MetaM α) : MetaM α :=
  loop 0 #[]
where
  loop (i : Nat) (motives : Array Expr) : MetaM α := do
    if i < motiveFVars.size then
      let type ← inferType motiveFVars[i]
      let motive ← forallTelescopeReducing type fun xs _ => do
        mkLambdaFVars xs <| mkConst ``Nat
      trace[Meta.sizeOf] "motive: {motive}"
      loop (i+1) (motives.push motive)
    else
      k motives

private partial def mkSizeOfMinors {α} (motiveFVars : Array Expr) (minorFVars : Array Expr) (minorFVars' : Array Expr) (k : Array Expr → MetaM α) : MetaM α :=
  assert! minorFVars.size == minorFVars'.size
  loop 0 #[]
where
  loop (i : Nat) (minors : Array Expr) : MetaM α := do
    if i < minorFVars.size then
      forallTelescopeReducing (← inferType minorFVars[i]) fun xs _ => do
      forallBoundedTelescope (← inferType minorFVars'[i]) xs.size fun xs' _ => do
        let mut minor ← mkNumeral (mkConst ``Nat) 1
        for x in xs, x' in xs' do
          unless (← isInductiveHypothesis motiveFVars x) do
          unless (← whnf (← inferType x)).isForall do -- we suppress higher-order fields
            match (← isRecField? motiveFVars xs x) with
            | some idx => minor ← mkAdd minor xs'[idx]
            | none     => minor ← mkAdd minor (← mkAppM ``SizeOf.sizeOf #[x'])
        minor ← mkLambdaFVars xs' minor
        trace[Meta.sizeOf] "minor: {minor}"
        loop (i+1) (minors.push minor)
    else
      k minors

/--
  Create a "sizeOf" function with name `declName` using the recursor `recName`.
-/
partial def mkSizeOfFn (recName : Name) (declName : Name): MetaM Unit := do
  trace[Meta.sizeOf] "recName: {recName}"
  let recInfo : RecursorVal ← getConstInfoRec recName
  forallTelescopeReducing recInfo.type fun xs type =>
    let levelParams := recInfo.levelParams.tail! -- universe parameters for declaration being defined
    let params := xs[:recInfo.numParams]
    let motiveFVars := xs[recInfo.numParams : recInfo.numParams + recInfo.numMotives]
    let minorFVars := xs[recInfo.getFirstMinorIdx : recInfo.getFirstMinorIdx + recInfo.numMinors]
    let indices := xs[recInfo.getFirstIndexIdx : recInfo.getFirstIndexIdx + recInfo.numIndices]
    let major := xs[recInfo.getMajorIdx]
    let nat := mkConst ``Nat
    mkLocalInstances params fun localInsts =>
    mkSizeOfMotives motiveFVars fun motives => do
      let us := levelOne :: levelParams.map mkLevelParam -- universe level parameters for `rec`-application
      let recFn := mkConst recName us
      let val := mkAppN recFn (params ++ motives)
      forallBoundedTelescope (← inferType val) recInfo.numMinors fun minorFVars' _ =>
      mkSizeOfMinors motiveFVars minorFVars minorFVars' fun minors => do
        let sizeOfParams := params ++ localInsts ++ indices ++ #[major]
        let sizeOfType ← mkForallFVars sizeOfParams nat
        let val := mkAppN val (minors ++ indices ++ #[major])
        trace[Meta.sizeOf] "val: {val}"
        let sizeOfValue ← mkLambdaFVars sizeOfParams val
        addDecl <| Declaration.defnDecl {
          name        := declName
          levelParams := levelParams
          type        := sizeOfType
          value       := sizeOfValue
          safety      := DefinitionSafety.safe
          hints       := ReducibilityHints.abbrev
        }

/--
  Create `sizeOf` functions for all inductive datatypes in the mutual inductive declaration containing `typeName`
  The resulting array contains the generated functions names. The `NameMap` maps recursor names into the generated function names.
  There is a function for each element of the mutual inductive declaration, and for auxiliary recursors for nested inductive types.
-/
def mkSizeOfFns (typeName : Name) : MetaM (Array Name × NameMap Name) := do
  let indInfo ← getConstInfoInduct typeName
  let recInfo ← getConstInfoRec (mkRecName typeName)
  let numExtra := recInfo.numMotives - indInfo.all.length -- numExtra > 0 for nested inductive types
  let mut result := #[]
  let baseName := indInfo.all.head! ++ `_sizeOf -- we use the first inductive type as the base name for `sizeOf` functions
  let mut i := 1
  let mut recMap : NameMap Name := {}
  for indTypeName in indInfo.all do
    let sizeOfName := baseName.appendIndexAfter i
    let recName := mkRecName indTypeName
    mkSizeOfFn recName sizeOfName
    recMap := recMap.insert recName sizeOfName
    result := result.push sizeOfName
    i := i + 1
  for j in [:numExtra] do
    let recName := (mkRecName indInfo.all.head!).appendIndexAfter (j+1)
    let sizeOfName := baseName.appendIndexAfter i
    mkSizeOfFn recName sizeOfName
    recMap := recMap.insert recName sizeOfName
    result := result.push sizeOfName
    i := i + 1
  return (result, recMap)

def mkSizeOfSpecLemmaName (ctorName : Name) : Name :=
  ctorName ++ `sizeOf_spec

def mkSizeOfSpecLemmaInstance (ctorApp : Expr) : MetaM Expr :=
  matchConstCtor ctorApp.getAppFn (fun _ => throwError "failed to apply 'sizeOf' spec, constructor expected{indentExpr ctorApp}") fun ctorInfo ctorLevels => do
    let ctorArgs     := ctorApp.getAppArgs
    let ctorFields   := ctorArgs[ctorArgs.size - ctorInfo.numFields:]
    let lemmaName  := mkSizeOfSpecLemmaName ctorInfo.name
    let lemmaInfo  ← getConstInfo lemmaName
    let lemmaArity ← forallTelescopeReducing lemmaInfo.type fun xs _ => return xs.size
    let lemmaArgMask := mkArray (lemmaArity - ctorInfo.numFields) (none (α := Expr))
    let lemmaArgMask := lemmaArgMask ++ ctorFields.toArray.map some
    mkAppOptM lemmaName lemmaArgMask

/- SizeOf spec theorem for nested inductive types -/
namespace SizeOfSpecNested

structure Context where
  indInfo    : InductiveVal
  sizeOfFns  : Array Name
  ctorName   : Name
  params     : Array Expr
  localInsts : Array Expr
  recMap     : NameMap Name -- mapping from recursor name into `_sizeOf_<idx>` function name (see `mkSizeOfFns`)

abbrev M := ReaderT Context MetaM

def throwUnexpected {α} (msg : MessageData) : M α := do
  throwError "failed to generate sizeOf theorem for {(← read).ctorName} (use `set_option genSizeOfSpec false` to disable theorem generation), {msg}"

def throwFailed {α} : M α := do
  throwError "failed to generate sizeOf theorem for {(← read).ctorName}, (use `set_option genSizeOfSpec false` to disable theorem generation)"

/-- Convert a recursor application into a `_sizeOf_<idx>` application. -/
private def recToSizeOf (e : Expr) : M Expr := do
  matchConstRec e.getAppFn (fun _ => throwFailed) fun info us => do
    match (← read).recMap.find? info.name with
    | none => throwUnexpected m!"expected recursor application {indentExpr e}"
    | some sizeOfName =>
      let args    := e.getAppArgs
      let indices := args[info.getFirstIndexIdx : info.getFirstIndexIdx + info.numIndices]
      let major   := args[info.getMajorIdx]
      return mkAppN (mkConst sizeOfName us.tail!) ((← read).params ++ (← read).localInsts ++ indices ++ #[major])

mutual
  /-- Construct minor premise proof for `mkSizeOfAuxLemmaProof`. `ys` contains fields and inductive hypotheses for the minor premise. -/
  private partial def mkMinorProof (ys : Array Expr) (lhs rhs : Expr) : M Expr := do
    trace[Meta.sizeOf.minor] "{lhs} =?= {rhs}"
    if (← isDefEq lhs rhs) then
      mkEqRefl rhs
    else
      match (← whnfI lhs).natAdd?, (← whnfI rhs).natAdd? with
      | some (a₁, b₁), some (a₂, b₂) =>
        let p₁ ← mkMinorProof ys a₁ a₂
        let p₂ ← mkMinorProofStep ys b₁ b₂
        mkCongr (← mkCongrArg (mkConst ``Nat.add) p₁) p₂
      | _, _ =>
        throwUnexpected m!"expected 'Nat.add' application, lhs is {indentExpr lhs}\nrhs is{indentExpr rhs}"

  /--
    Helper method for `mkMinorProof`. The proof step is one of the following
    - Reflexivity
    - Assumption (i.e., using an inductive hypotheses from `ys`)
    - `mkSizeOfAuxLemma` application. This case happens when we have multiple levels of nesting
  -/
  private partial def mkMinorProofStep (ys : Array Expr) (lhs rhs : Expr) : M Expr := do
    if (← isDefEq lhs rhs) then
      mkEqRefl rhs
    else
      let lhs ← recToSizeOf lhs
      trace[Meta.sizeOf.minor.step] "{lhs} =?= {rhs}"
      let target ← mkEq lhs rhs
      for y in ys do
        if (← isDefEq (← inferType y) target) then
          return y
      mkSizeOfAuxLemma lhs rhs

  /-- Construct proof of auxiliary lemma. See `mkSizeOfAuxLemma` -/
  private partial def mkSizeOfAuxLemmaProof (info : InductiveVal) (lhs rhs : Expr) : M Expr := do
    let lhsArgs := lhs.getAppArgs
    let sizeOfBaseArgs := lhsArgs[:lhsArgs.size - info.numIndices - 1]
    let indicesMajor := lhsArgs[lhsArgs.size - info.numIndices - 1:]
    let sizeOfLevels := lhs.getAppFn.constLevels!
    /- Auxiliary function for constructing an `_sizeOf_<idx>` for `ys`,
       where `ys` are the indices + major.
       Recall that if `info.name` is part of a mutually inductive declaration, then the resulting application
       is not necessarily a `lhs.getAppFn` application.
       The result is an application of one of the `(← read),sizeOfFns` functions.
       We use this auxiliary function to builtin the motive of the recursor. -/
    let rec mkSizeOf (ys : Array Expr) : M Expr := do
      for sizeOfFn in (← read).sizeOfFns do
        let candidate := mkAppN (mkAppN (mkConst sizeOfFn sizeOfLevels) sizeOfBaseArgs) ys
        if (← isTypeCorrect candidate) then
          return candidate
      throwFailed
    let major := lhs.appArg!
    let majorType ← whnf (← inferType major)
    let majorTypeArgs := majorType.getAppArgs
    match majorType.getAppFn.const? with
    | none => throwFailed
    | some (_, us) =>
      let recName := mkRecName info.name
      let recInfo ← getConstInfoRec recName
      let r := mkConst recName (levelZero :: us)
      let r := mkAppN r majorTypeArgs[:info.numParams]
      forallBoundedTelescope (← inferType r) recInfo.numMotives fun motiveFVars _ => do
        let mut r := r
        -- Add motives
        for motiveFVar in motiveFVars do
          let motive ← forallTelescopeReducing (← inferType motiveFVar) fun ys _ => do
            let lhs ← mkSizeOf ys
            let rhs ← mkAppM ``SizeOf.sizeOf #[ys.back]
            mkLambdaFVars ys (← mkEq lhs rhs)
          r := mkApp r motive
        forallBoundedTelescope (← inferType r) recInfo.numMinors fun minorFVars _ => do
          let mut r := r
          -- Add minors
          for minorFVar in minorFVars do
            let minor ← forallTelescopeReducing (← inferType minorFVar) fun ys target => do
              let target ← whnf target
              match target.eq? with
              | none => throwFailed
              | some (_, lhs, rhs) =>
                if (← isDefEq lhs rhs) then
                  mkLambdaFVars ys (← mkEqRefl rhs)
                else
                  let lhs ← unfoldDefinition lhs -- Unfold `_sizeOf_<idx>`
                  -- rhs is of the form `sizeOf (ctor ...)`
                  let ctorApp := rhs.appArg!
                  let specLemma ← mkSizeOfSpecLemmaInstance ctorApp
                  let specEq ← whnf (← inferType specLemma)
                  match specEq.eq? with
                  | none => throwFailed
                  | some (_, rhs, rhsExpanded) =>
                    let lhs_eq_rhsExpanded ← mkMinorProof ys lhs rhsExpanded
                    let rhsExpanded_eq_rhs ← mkEqSymm specLemma
                    mkLambdaFVars ys (← mkEqTrans lhs_eq_rhsExpanded rhsExpanded_eq_rhs)
            r := mkApp r minor
          -- Add indices and major
          return mkAppN r indicesMajor

  /--
    Generate proof for `C._sizeOf_<idx> t = sizeOf t` where `C._sizeOf_<idx>` is a auxiliary function
    generated for a nested inductive type in `C`.
    For example, given
    ```lean
    inductive Expr where
      | app (f : String) (args : List Expr)
    ```
    We generate the auxiliary function `Expr._sizeOf_1 : List Expr → Nat`.
    To generate the `sizeOf` spec lemma
    ```
    sizeOf (Expr.app f args) = 1 + sizeOf f + sizeOf args
    ```
    we need an auxiliary lemma for showing `Expr._sizeOf_1 args = sizeOf args`.
    Recall that `sizeOf (Expr.app f args)` is definitionally equal to `1 + sizeOf f + Expr._sizeOf_1 args`, but
    `Expr._sizeOf_1 args` is **not** definitionally equal to `sizeOf args`. We need a proof by induction.
  -/
  private partial def mkSizeOfAuxLemma (lhs rhs : Expr) : M Expr := do
    trace[Meta.sizeOf.aux] "{lhs} =?= {rhs}"
    match lhs.getAppFn.const? with
    | none => throwFailed
    | some (fName, us) =>
      let thmLevelParams ← us.mapM fun
        | Level.param n _ => return n
        | _ => throwFailed
      let thmName  := fName.appendAfter "_eq"
      if (← getEnv).contains thmName then
        -- Auxiliary lemma has already been defined
        return mkAppN (mkConst thmName us) lhs.getAppArgs
      else
        -- Define auxiliary lemma
        -- First, generalize indices
        let x := lhs.appArg!
        let xType ← whnf (← inferType x)
        matchConstInduct xType.getAppFn (fun _ => throwFailed) fun info _ => do
          let params := xType.getAppArgs[:info.numParams]
          forallTelescopeReducing (← inferType (mkAppN xType.getAppFn params)) fun indices _ => do
            let majorType := mkAppN (mkAppN xType.getAppFn params) indices
            withLocalDeclD `x majorType fun major => do
              let lhsArgs := lhs.getAppArgs
              let lhsArgsNew := lhsArgs[:lhsArgs.size - 1 - indices.size] ++ indices ++ #[major]
              let lhsNew := mkAppN lhs.getAppFn lhsArgsNew
              let rhsNew ← mkAppM ``SizeOf.sizeOf #[major]
              let eq ← mkEq lhsNew rhsNew
              let thmParams := lhsArgsNew
              let thmType ← mkForallFVars thmParams eq
              let thmValue ← mkSizeOfAuxLemmaProof info lhsNew rhsNew
              let thmValue ← mkLambdaFVars thmParams thmValue
              trace[Meta.sizeOf] "thmValue: {thmValue}"
              addDecl <| Declaration.thmDecl {
                name        := thmName
                levelParams := thmLevelParams
                type        := thmType
                value       := thmValue
              }
              return mkAppN (mkConst thmName us) lhs.getAppArgs

end

/- Prove SizeOf spec lemma of the form `sizeOf <ctor-application> = 1 + sizeOf <field_1> + ... + sizeOf <field_n> -/
partial def main (lhs rhs : Expr) : M Expr := do
  if (← isDefEq lhs rhs) then
    mkEqRefl rhs
  else
    /- Expand lhs and rhs to obtain `Nat.add` applications -/
    let lhs ← whnfI lhs            -- Expand `sizeOf (ctor ...)` into `_sizeOf_<idx>` application
    let lhs ← unfoldDefinition lhs -- Unfold `_sizeOf_<idx>` application into `HAdd.hAdd` application
    loop lhs rhs
where
  loop (lhs rhs : Expr) : M Expr := do
    trace[Meta.sizeOf.loop] "{lhs} =?= {rhs}"
    if (← isDefEq lhs rhs) then
      mkEqRefl rhs
    else
      match (← whnfI lhs).natAdd?, (← whnfI rhs).natAdd? with
      | some (a₁, b₁), some (a₂, b₂) =>
        let p₁ ← loop a₁ a₂
        let p₂ ← step b₁ b₂
        mkCongr (← mkCongrArg (mkConst ``Nat.add) p₁) p₂
      | _, _ =>
        throwUnexpected m!"expected 'Nat.add' application, lhs is {indentExpr lhs}\nrhs is{indentExpr rhs}"

  step (lhs rhs : Expr) : M Expr := do
    if (← isDefEq lhs rhs) then
      mkEqRefl rhs
    else
      let lhs ← recToSizeOf lhs
      mkSizeOfAuxLemma lhs rhs

end SizeOfSpecNested

private def mkSizeOfSpecTheorem (indInfo : InductiveVal) (sizeOfFns : Array Name) (recMap : NameMap Name) (ctorName : Name) : MetaM Unit := do
  let ctorInfo ← getConstInfoCtor ctorName
  let us := ctorInfo.levelParams.map mkLevelParam
  let simpAttr ← ofExcept <| getAttributeImpl (← getEnv) `simp
  forallTelescopeReducing ctorInfo.type fun xs _ => do
    let params := xs[:ctorInfo.numParams]
    let fields := xs[ctorInfo.numParams:]
    let ctorApp := mkAppN (mkConst ctorName us) xs
    mkLocalInstances params fun localInsts => do
      let lhs ← mkAppM ``SizeOf.sizeOf #[ctorApp]
      let mut rhs ← mkNumeral (mkConst ``Nat) 1
      for field in fields do
        unless (← whnf (← inferType field)).isForall do
          rhs ← mkAdd rhs (← mkAppM ``SizeOf.sizeOf #[field])
      let target ← mkEq lhs rhs
      let thmName   := mkSizeOfSpecLemmaName ctorName
      let thmParams := params ++ localInsts ++ fields
      let thmType ← mkForallFVars thmParams target
      let thmValue ←
        if indInfo.isNested then
          SizeOfSpecNested.main lhs rhs |>.run {
            indInfo := indInfo, sizeOfFns := sizeOfFns, ctorName := ctorName, params := params, localInsts := localInsts, recMap := recMap
          }
        else
          mkEqRefl rhs
      let thmValue ← mkLambdaFVars thmParams thmValue
      trace[Meta.sizeOf] "sizeOf spec theorem: {thmName}"
      addDecl <| Declaration.thmDecl {
        name        := thmName
        levelParams := ctorInfo.levelParams
        type        := thmType
        value       := thmValue
      }
      simpAttr.add thmName default AttributeKind.global

private def mkSizeOfSpecTheorems (indTypeNames : Array Name) (sizeOfFns : Array Name) (recMap : NameMap Name) : MetaM Unit := do
  for indTypeName in indTypeNames do
    let indInfo ← getConstInfoInduct indTypeName
    for ctorName in indInfo.ctors do
      mkSizeOfSpecTheorem indInfo sizeOfFns recMap ctorName
  return ()

register_builtin_option genSizeOf : Bool := {
  defValue := true
  descr    := "generate `SizeOf` instance for inductive types and structures"
}

register_builtin_option genSizeOfSpec : Bool := {
  defValue := true
  descr    := "generate `SizeOf` specificiation theorems for automatically generated instances"
}

def mkSizeOfInstances (typeName : Name) : MetaM Unit := do
  if (← getEnv).contains ``SizeOf && genSizeOf.get (← getOptions) && !(← isInductivePredicate typeName) then
    let indInfo ← getConstInfoInduct typeName
    unless indInfo.isUnsafe do
      let (fns, recMap) ← mkSizeOfFns typeName
      for indTypeName in indInfo.all, fn in fns do
        let indInfo ← getConstInfoInduct indTypeName
        forallTelescopeReducing indInfo.type fun xs _ =>
          let params := xs[:indInfo.numParams]
          let indices := xs[indInfo.numParams:]
          mkLocalInstances params fun localInsts => do
            let us := indInfo.levelParams.map mkLevelParam
            let indType := mkAppN (mkConst indTypeName us) xs
            let sizeOfIndType ← mkAppM ``SizeOf #[indType]
            withLocalDeclD `m indType fun m => do
              let v ← mkLambdaFVars #[m] <| mkAppN (mkConst fn us) (params ++ localInsts ++ indices ++ #[m])
              let sizeOfMk ← mkAppM ``SizeOf.mk #[v]
              let instDeclName := indTypeName ++ `_sizeOf_inst
              let instDeclType ← mkForallFVars (xs ++ localInsts) sizeOfIndType
              let instDeclValue ← mkLambdaFVars (xs ++ localInsts) sizeOfMk
              addDecl <| Declaration.defnDecl {
                name        := instDeclName
                levelParams := indInfo.levelParams
                type        := instDeclType
                value       := instDeclValue
                safety      := DefinitionSafety.safe
                hints       := ReducibilityHints.abbrev
              }
              addInstance instDeclName AttributeKind.global (eval_prio default)
      if genSizeOfSpec.get (← getOptions) then
        mkSizeOfSpecTheorems indInfo.all.toArray fns recMap

builtin_initialize
  registerTraceClass `Meta.sizeOf

end Lean.Meta