/-
Copyright (c) 2017 Daniel Selsam. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Daniel Selsam

Certified graph transformation that "reparameterizes" a specific occurrence of a stochastic choice.
-/
import .util .tensor .tfacts .compute_grad .graph .tactics .ops .predicates .lemmas .env

namespace certigrad
open list

section algebra
open T

lemma mvn_transform {shape : S} (μ σ x : T shape) (H_σ : σ > 0) :
  mvn_pdf μ σ x = (prod σ⁻¹) * mvn_pdf 0 1 ((x - μ) / σ) :=
calc  mvn_pdf μ σ x
    = prod ((sqrt ((2 * pi shape) * square σ))⁻¹ * exp ((- 2⁻¹) * (square $ (x - μ) / σ))) : rfl
... = prod ((sqrt (2 * pi shape) * σ)⁻¹ * exp ((- 2⁻¹) * (square $ (x - μ) / σ))) : by rw [sqrt_mul, sqrt_square]
... = prod (((sqrt (2 * pi shape))⁻¹ * σ⁻¹) * exp ((- 2⁻¹) * (square $ (x - μ) / σ))) : by rw [T.mul_inv_pos (sqrt_pos two_pi_pos) H_σ]
... = (prod σ⁻¹) * prod ((sqrt (2 * pi shape))⁻¹ * exp ((- 2⁻¹) * (square $ (x - μ) / σ))) : by simp [prod_mul]
... = (prod σ⁻¹) * prod ((sqrt ((2 * pi shape) * square 1))⁻¹ * exp ((- 2⁻¹) * (square ((((x - μ) / σ) - 0) / 1)))) : by simp [T.div_one, square]
... = (prod σ⁻¹) * mvn_pdf 0 1 ((x - μ) / σ) : rfl

end algebra

open sprog

lemma mvn_reparam_same {shape oshape : S} {μ σ : T shape} (f : dvec T [shape] → T oshape) : σ > 0 →
E (prim (rand.op.mvn shape) ⟦μ, σ⟧) f
=
E (bind (prim (rand.op.mvn_std shape) ⟦⟧) (λ (x : dvec T [shape]), ret ⟦(x^.head * σ) + μ⟧)) f :=
assume (H_σ_pos : σ > 0),
begin
simp only [E.E_bind, E.E_ret],
dunfold E rand.op.mvn rand.op.pdf T.dintegral dvec.head rand.pdf.mvn rand.pdf.mvn_std,
simp only [λ x, mvn_transform μ σ x H_σ_pos],

assert H : ∀ (x : T shape), ((σ * x + μ + -μ) / σ) = x,
  { intro x, simp only [add_assoc, add_neg_self, add_zero], rw mul_comm, rw -T.mul_div_mul_alt, rw T.div_self H_σ_pos, rw mul_one},
definev g : T shape → T oshape := λ (x : T shape), T.mvn_pdf 0 1 ((x - μ) / σ) ⬝ f ⟦x⟧,
assert H_rhs : ∀ (x : T shape), T.mvn_pdf 0 1 x ⬝ f ⟦x * σ + μ⟧ = g (σ * x + μ),
{ intro x, dsimp, rw H, simp },

rw funext H_rhs,
rw T.integral_scale_shift_var g,
dsimp,
simp [T.smul_group]
end

def reparameterize_pre (eshape : S) : list node → env → Prop
| [] inputs := true
| (⟨⟨ref, shape⟩, [⟨μ, .(shape)⟩, ⟨σ, .(shape)⟩], operator.rand (rand.op.mvn .(shape))⟩::nodes) inputs :=
  eshape = shape ∧ σ ≠ μ ∧ 0 < env.get (σ, shape) inputs
| (⟨ref, parents, operator.det op⟩::nodes) inputs := reparameterize_pre nodes (env.insert ref (op^.f (env.get_ks parents inputs)) inputs)
| (⟨ref, parents, operator.rand op⟩::nodes) inputs := ∀ x, reparameterize_pre nodes (env.insert ref x inputs)

def reparameterize (fname : ID) : list node → list node
| [] := []

| (⟨⟨ident, shape⟩, [⟨μ, .(shape)⟩, ⟨σ, .(shape)⟩], operator.rand (rand.op.mvn .(shape))⟩::nodes) :=

 (⟨(fname, shape), [],                                       operator.rand (rand.op.mvn_std shape)⟩
::⟨(ident, shape),   [(fname, shape), (σ, shape), (μ, shape)], operator.det (ops.mul_add shape)⟩
::nodes)

| (n::nodes) := n :: reparameterize nodes

theorem reparameterize_correct (costs : list ID) :
∀ (nodes : list node) (inputs : env) (fref : reference),
  reparameterize_pre fref.2 nodes inputs →
  uniq_ids nodes inputs →
  all_parents_in_env inputs nodes →
  (¬ env.has_key fref inputs) →  fref ∉ map node.ref nodes →
  (fref.1 ∉ costs) →
E (graph.to_dist (λ env₀, ⟦sum_costs env₀ costs⟧) inputs (reparameterize fref.1 nodes)) dvec.head
=
E (graph.to_dist (λ env₀, ⟦sum_costs env₀ costs⟧) inputs nodes) dvec.head

| [] _ _ _ _ _ _ _ _ := rfl

| (⟨⟨ident, shape⟩, [⟨μ, .(shape)⟩, ⟨σ, .(shape)⟩], operator.rand (rand.op.mvn .(shape))⟩::nodes) inputs fref H_pre H_uids H_ps_in_env H_fresh₁ H_fresh₂ H_not_cost :=
begin
dunfold reparameterize,
assertv H_eshape : fref.2 = shape := H_pre^.left,
assert H_fref : fref = (fref.1, shape),
{ clear reparameterize_correct, cases fref with fref₁ fref₂, dsimp at H_eshape, rw H_eshape },

assertv H_σ_μ : σ ≠ μ := H_pre^.right^.left,

dunfold graph.to_dist operator.to_dist,
dsimp,
simp [E.E_bind],
erw (mvn_reparam_same _ H_pre^.right^.right),
simp [E.E_bind, E.E_ret],
dunfold dvec.head,
dsimp,
apply congr_arg, apply funext, intro x,

assertv H_μ_in : env.has_key (μ, shape) inputs := H_ps_in_env^.left (μ, shape) (mem_cons_self _ _),
assertv H_σ_in : env.has_key (σ, shape) inputs := H_ps_in_env^.left (σ, shape) (mem_cons_of_mem _ (mem_cons_self _ _)),
assertv H_ident_nin : ¬ env.has_key (ident, shape) inputs := H_uids^.left,

assertv H_μ_neq_ident : (μ, shape) ≠ (ident, shape) := env_in_nin_ne H_μ_in H_ident_nin,
assertv H_σ_neq_ident : (σ, shape) ≠ (ident, shape) := env_in_nin_ne H_σ_in H_ident_nin,

assertv H_μ_neq_fref : (μ, shape) ≠ (fref.1, shape) := eq.rec_on H_fref (env_in_nin_ne H_μ_in H_fresh₁),
assertv H_σ_neq_fref : (σ, shape) ≠ (fref.1, shape) := eq.rec_on H_fref (env_in_nin_ne H_σ_in H_fresh₁),
assertv H_ident_neq_fref : (ident, shape) ≠ (fref.1, shape) := eq.rec_on H_fref (mem_not_mem_neq mem_of_cons_same H_fresh₂),

dunfold env.get_ks,
tactic.dget_dinsert,

rw (env.insert_insert_flip _ _ _ H_ident_neq_fref),
dsimp,

definev fval : T shape := dvec.head x,
definev fval_inputs : env := env.insert (ident, shape)
                                        (det.op.f (ops.mul_add shape)
                                                  ⟦dvec.head x, (env.get (σ, shape) inputs : T shape), (env.get (μ, shape) inputs : T shape)⟧)
                                         inputs,

assertv H_ps_in_env_next : all_parents_in_env fval_inputs nodes := H_ps_in_env^. right _,
assertv H_fresh₁_next : ¬ env.has_key (fref.1, shape) fval_inputs :=
  eq.rec_on H_fref (env_not_has_key_insert (eq.rec_on (eq.symm H_fref) $ ne.symm H_ident_neq_fref) H_fresh₁),

assertv H_fresh₂_next : (fref.1, shape) ∉ map node.ref nodes := eq.rec_on H_fref (not_mem_of_not_mem_cons H_fresh₂),

erw (@to_dist_congr_insert costs nodes fval_inputs (fref.1, shape) fval H_ps_in_env_next H_fresh₁_next H_fresh₂_next H_not_cost),
dsimp,
dunfold det.op.f,
rw [add_comm, mul_comm],
reflexivity
end

| (⟨(ref, shape), [], operator.det op⟩::nodes) inputs fref H_pre H_uids H_ps_in_env H_fresh₁ H_fresh₂ H_not_cost :=
begin
dunfold reparameterize graph.to_dist operator.to_dist,
simp [E.E_bind, E.E_ret],
definev x : T shape := op^.f (env.get_ks [] inputs),
assertv H_pre_next : reparameterize_pre fref.2 nodes (env.insert (ref, shape) x inputs) := by apply H_pre,
assertv H_ps_in_env_next : all_parents_in_env (env.insert (ref, shape) x inputs) nodes := H_ps_in_env^.right _,
assertv H_fresh₁_next : ¬ env.has_key fref (env.insert (ref, shape) x inputs) := env_not_has_key_insert (ne_of_not_mem_cons H_fresh₂) H_fresh₁,
assertv H_fresh₂_next : fref ∉ map node.ref nodes := not_mem_of_not_mem_cons H_fresh₂,
apply (reparameterize_correct _ _ fref H_pre_next (H_uids^.right _) H_ps_in_env_next H_fresh₁_next H_fresh₂_next H_not_cost)
end

| (⟨(ref, shape), [], operator.rand op⟩::nodes) inputs fref H_pre H_uids H_ps_in_env H_fresh₁ H_fresh₂ H_not_cost :=
begin
dunfold reparameterize graph.to_dist,
simp [E.E_bind],
apply congr_arg, apply funext, intro x,
assertv H_pre_next : reparameterize_pre fref.2 nodes (env.insert (ref, shape) (dvec.head x) inputs) := by apply H_pre,
assertv H_ps_in_env_next : all_parents_in_env (env.insert (ref, shape) (dvec.head x) inputs) nodes := H_ps_in_env^.right x^.head,
assertv H_fresh₁_next : ¬ env.has_key fref (env.insert (ref, shape) x^.head inputs) := env_not_has_key_insert (ne_of_not_mem_cons H_fresh₂) H_fresh₁,
assertv H_fresh₂_next : fref ∉ map node.ref nodes := not_mem_of_not_mem_cons H_fresh₂,
apply (reparameterize_correct _ _ fref H_pre_next (H_uids^.right _) H_ps_in_env_next H_fresh₁_next H_fresh₂_next H_not_cost)
end

| (⟨(ref, shape), [(parent₁, shape₁)], operator.det op⟩::nodes) inputs fref H_pre H_uids H_ps_in_env H_fresh₁ H_fresh₂ H_not_cost :=
begin
dunfold reparameterize graph.to_dist operator.to_dist,
simp [E.E_bind, E.E_ret],
definev x : T shape := det.op.f op (env.get_ks [(parent₁, shape₁)] inputs),
assertv H_pre_next : reparameterize_pre fref.2 nodes (env.insert (ref, shape) x inputs) := by apply H_pre,
assertv H_ps_in_env_next : all_parents_in_env (env.insert (ref, shape) x inputs) nodes := H_ps_in_env^.right x,
assertv H_fresh₁_next : ¬ env.has_key fref (env.insert (ref, shape) x inputs) := env_not_has_key_insert (ne_of_not_mem_cons H_fresh₂) H_fresh₁,
assertv H_fresh₂_next : fref ∉ map node.ref nodes := not_mem_of_not_mem_cons H_fresh₂,
apply (reparameterize_correct _ _ fref H_pre_next (H_uids^.right _) H_ps_in_env_next H_fresh₁_next H_fresh₂_next H_not_cost)
end

| (⟨(ref, shape), [(parent₁, shape₁), (parent₂, shape₂)], operator.det op⟩::nodes) inputs fref H_pre H_uids H_ps_in_env H_fresh₁ H_fresh₂ H_not_cost :=
begin
dunfold reparameterize graph.to_dist operator.to_dist,
simp [E.E_bind, E.E_ret],
definev x : T shape := det.op.f op (env.get_ks [(parent₁, shape₁), (parent₂, shape₂)] inputs),
assertv H_pre_next : reparameterize_pre fref.2 nodes (env.insert (ref, shape) x inputs) := by apply H_pre,
assertv H_ps_in_env_next : all_parents_in_env (env.insert (ref, shape) x inputs) nodes := H_ps_in_env^.right x,
assertv H_fresh₁_next : ¬ env.has_key fref (env.insert (ref, shape) x inputs) := env_not_has_key_insert (ne_of_not_mem_cons H_fresh₂) H_fresh₁,
assertv H_fresh₂_next : fref ∉ map node.ref nodes := not_mem_of_not_mem_cons H_fresh₂,
apply (reparameterize_correct _ _ fref H_pre_next (H_uids^.right _) H_ps_in_env_next H_fresh₁_next H_fresh₂_next H_not_cost)
end

| (⟨(ref, shape), (parent₁, shape₁) :: (parent₂, shape₂) :: (parent₃, shape₃) :: parents, operator.det op⟩::nodes) inputs fref H_pre H_uids H_ps_in_env H_fresh₁ H_fresh₂ H_not_cost :=
begin
dunfold reparameterize graph.to_dist operator.to_dist,
simp [E.E_bind, E.E_ret],
definev x : T shape := det.op.f op (env.get_ks ((parent₁, shape₁) :: (parent₂, shape₂) :: (parent₃, shape₃) :: parents) inputs),
assertv H_pre_next : reparameterize_pre fref.2 nodes (env.insert (ref, shape) x inputs) := by apply H_pre,
assertv H_ps_in_env_next : all_parents_in_env (env.insert (ref, shape) x inputs) nodes := H_ps_in_env^.right x,
assertv H_fresh₁_next : ¬ env.has_key fref (env.insert (ref, shape) x inputs) := env_not_has_key_insert (ne_of_not_mem_cons H_fresh₂) H_fresh₁,
assertv H_fresh₂_next : fref ∉ map node.ref nodes := not_mem_of_not_mem_cons H_fresh₂,
apply (reparameterize_correct _ _ fref H_pre_next (H_uids^.right _) H_ps_in_env_next H_fresh₁_next H_fresh₂_next H_not_cost)
end

def reparam : graph → graph
| g := ⟨reparameterize (ID.str label.ε) g^.nodes, g^.costs, g^.targets, g^.inputs⟩

end certigrad


-- vim: ts=2 sw=0 sts=-1 et ai tw=70

-- Note we're not currently transitioning over to hidden namespace - we probably should,
-- though it will break things

import ..myset.basic
import ..myset.finite
import ..mylist.map
import ..logic

namespace hidden

structure topological_space (α : Type) :=
(is_open : myset (myset α))
(univ_open : is_open myset.univ)
(empty_open : is_open ∅)
(open_union_open (σ : myset (myset α)): σ ⊆ is_open → is_open ⋃₀ σ)
(open_intersection_open (U V : myset α) : is_open U → is_open V → is_open (U ∩ V))

-- Don't know what this does
-- attribute [class] topological_space

namespace topological_space

open classical

local attribute [instance] classical.prop_decidable

variables {α β : Type}
-- include X Y

structure sets_separate_points
(X: topological_space α) (x y: α) (U V: myset α): Prop :=
(open_left: X.is_open U)
(open_right: X.is_open V)
(contains_left: x ∈ U)
(contains_right: y ∈ V)
(disjoint: U ∩ V = ∅)

def is_hausdorff (X : topological_space α) : Prop :=
∀ x y : α, x ≠ y →
  ∃ (U V : myset α),
    sets_separate_points X x y U V

def is_closed (X : topological_space α) (A : myset α) : Prop :=
is_open X (myset.compl A)

-- An attempt to create subspaces
def subspace_topology (X : topological_space α) (Y : myset α) : topological_space (subtype Y) :=
{
  is_open := { V | ∃ U : myset α, is_open X U ∧ V = myset.subtype_restriction Y U },
  univ_open :=
  begin
    existsi myset.univ,
    split, {
      from X.univ_open,
    }, {
      apply funext,
      intro x,
      refl,
    },
  end,
  empty_open :=
  begin
    -- it does weird things if i try ∅
    existsi (λ x, false),
    split, {
      from X.empty_open,
    }, {
      apply funext,
      intro x,
      refl,
    },
  end,
  open_union_open :=
  begin
    intros σ hσ,
    let τ : myset (myset α) :=
    {W | X.is_open W ∧ ∃ V : myset (subtype Y), V ∈ σ ∧ V = Y.subtype_restriction W},
    existsi ⋃₀τ,
    split, {
      apply X.open_union_open,
      intros U hU,
      from hU.left,
    }, {
      unfold myset.subtype_restriction,
      apply myset.setext,
      intro x,
      split; assume h, {
        cases h with W hW,
        cases hW with hWσ hxW,
        cases hσ _ hWσ with U hU,
        existsi U,
        cases hU with hUopen hWU,
        suffices : U ∈ τ,
          existsi this,
          rwa hWU at hxW,
        split,
          assumption,
        existsi W,
        split; assumption,
      }, {
        cases h with V hV,
        cases hV with hVτ hxV,
        existsi Y.subtype_restriction V,
        cases hVτ with hVopen this,
        cases this with U this,
        cases this with hUσ hUV,
        rw hUV at hUσ,
        existsi hUσ,
        assumption,
      },
    },
  end,
  open_intersection_open :=
  begin
    intros U V,
    assume hU hV,
    cases hU with U' hU',
    cases hV with V' hV',
    existsi U' ∩ V',
    split, {
      apply X.open_intersection_open, {
        from hU'.left,
      }, {
        from hV'.left,
      },
    }, {
      apply myset.setext,
      intro x,
      split, {
        assume hUVx,
        split, {
          have := hUVx.left,
          rw hU'.right at this,
          from this,
        }, {
          have := hUVx.right,
          rw hV'.right at this,
          from this,
        },
      }, {
        assume hUVx,
        split, {
          rw hU'.right,
          from hUVx.left,
        }, {
          rw hV'.right,
          from hUVx.right,
        },
      },
    },
  end
}

-- We don't assume a topology and define a base, we *build* a topology from a base
def is_base (B : myset (myset α)) :=
myset.univ ∈ B ∧ ∀ b₁ b₂: myset α, b₁ ∈ B → b₂ ∈ B → b₁ ∩ b₂ ∈ B

-- Given a suitable myset, build a topological space with it as a base.
def space_from_base (B : myset (myset α)) (hB : is_base B) : topological_space α :=
{
  is_open := λ U, (∀ x : α, x ∈ U → (∃ b, b ∈ B ∧ x ∈ b ∧ b ⊆ U)),
  univ_open :=
  begin
    unfold is_base at hB,
    intro x,
    assume h,
    existsi myset.univ,
    split,
      from hB.left,
    split,
      from h,
    intro y,
    assume hy,
    assumption,
  end,
  empty_open :=
  begin
    intro x,
    assume h,
    exfalso,
    from h,
  end,
  open_union_open :=
  begin
    intro σ,
    assume hσ,
    intro x,
    assume hx,
    cases hx with Uj hUj, -- Zsak called it U_j in his proof
    cases hUj with hUj hxUj, -- ???
    have h := hσ Uj hUj,
    cases h x hxUj with b hb,
    existsi b,
    split,
      from hb.left,
    split,
      from hb.right.left,
    intro y,
    assume hy,
    existsi Uj,
    existsi hUj,
    apply hb.right.right,
    from hy,
  end,
  open_intersection_open :=
  begin
    unfold is_base at hB,
    intros U V,
    assume hU hV,
    intro x,
    assume hUV,
    cases hU x (hUV.left) with b₁ hb₁,
    cases hV x (hUV.right) with b₂ hb₂,
    existsi b₁ ∩ b₂,
    split,
      apply hB.right,
        from hb₁.left,
      from hb₂.left,
    split,
      split,
        from hb₁.right.left,
      from hb₂.right.left,
    intro y,
    assume hy,
    split,
      apply hb₁.right.right,
      from hy.left,
    apply hb₂.right.right,
    from hy.right,
  end
}

structure is_box_of
(X: topological_space α) (Y: topological_space β)
(U: myset α) (V: myset β) (B: myset (α × β)): Prop :=
(eq: B = (U × V))
(open_left: X.is_open U)
(open_right: Y.is_open V)

def product_base (X : topological_space α) (Y : topological_space β) : myset (myset (α × β)) :=
{ B | ∃ (U : myset α) (V : myset β), is_box_of X Y U V B}

theorem is_base_product_base (X : topological_space α) (Y : topological_space β) :
is_base (product_base X Y) :=
begin
  split,
    existsi myset.univ,
    existsi myset.univ,
    split,
      apply myset.setext,
      intro x,
      split; assume h,
      split; trivial,
      trivial,
      from X.univ_open,
    from Y.univ_open,
  intros b₁ b₂,
  assume hb₁ hb₂,
  cases hb₁ with U₁ this,
  cases this with V₁ this,
  cases this with hb₁ this,
  cases hb₂ with U₂ this,
  cases this with V₂ this,
  cases this with hb₂ this,
  existsi U₁ ∩ U₂,
  existsi V₁ ∩ V₂,
  split,
    -- This should be general theorem (U₁ × V₁) ∩ (U₂ × V₂) = (U₁ ∩ U₂ × V₁ ∩ V₂)
    subst b₁, subst b₂,
    apply myset.setext,
    intro x,
    split; assume h,
      cases h with h₁ h₂,
      cases h₁ with hxU₁ hxV₁,
      cases h₂ with hxU₂ hxV₂,
      split; split; assumption, -- :) semicolons are nice
    cases h with hU hV,
    cases hU with hxU₁ hxU₂,
    cases hV with hxV₁ hxV₂,
    split; split; assumption,
    apply X.open_intersection_open; assumption,
  apply Y.open_intersection_open; assumption,
end

def product_topology (X : topological_space α) (Y : topological_space β) :
topological_space (α × β) := space_from_base (product_base X Y) (is_base_product_base X Y)

def discrete_topology (α : Type) : topological_space α :=
{
  is_open := λ U, true,
  univ_open := trivial,
  empty_open := trivial,
  open_union_open := λ σ h, trivial,
  open_intersection_open := λ U V hU hV, trivial,
}

def indiscrete_topology (α : Type) : topological_space α :=
{
  is_open := λ U, U = ∅ ∨ U = myset.univ,
  univ_open := begin
    right,
    refl,
  end,
  empty_open := begin
    left,
    refl,
  end,
  open_union_open := λ σ h,
  begin
    by_cases hempty: ⋃₀ σ = ∅, {
      left, assumption,
    }, {
      right,
      rw ←myset.empty_iff_eq_empty at hempty,
      rw ←myset.exists_iff_nempty at hempty,
      cases hempty with x hx,
      apply myset.setext,
      intro y,
      dsimp only [myset.univ],
      split; assume hy, {
        trivial,
      }, {
        dsimp only [myset.sUnion],
        dsimp only [myset.sUnion] at hx,
        cases hx with S hS,
        cases hS with H hH,
        cases h S H with he hu, {
          rw he at hH,
          exfalso, from hH,
        }, {
          existsi S,
          existsi H,
          rw hu,
          trivial,
        },
      },
    },
  end,
  open_intersection_open := λ U V hU hV,
  begin
    cases hU; cases hV, {
      left,
      rw [hU, hV],
      -- Now we are proving ∅ ∩ ∅ = ∅
      apply myset.setext,
      intro x,
      split; assume h,
      exact h.left,
      exact ⟨h,h⟩,
    }, {
      left,
      apply myset.setext,
      intro x,
      subst hU, subst hV,
      split; assume h,
        from h.left,
      split,
        from h,
      unfold myset.univ,
    }, {
      left,
      apply myset.setext,
      intro x,
      subst hU, subst hV,
      split; assume h,
        from h.right,
      split,
        unfold myset.univ,
      from h,
    }, {
      right,
      apply myset.setext,
      intro x,
      subst hU, subst hV,
      split; assume h,
        unfold myset.univ,
      split; unfold myset.univ,
    },
  end
}

def interior (X : topological_space α) (A : myset α) : myset α :=
⋃₀ { U | is_open X U ∧ U ⊆ A }

theorem interior_open (X : topological_space α) (A : myset α) : is_open X (interior X A) :=
begin
  apply X.open_union_open,
  intros U hU,
  from hU.left,
end

def closure (X : topological_space α) (A : myset α) : myset α :=
⋂₀ { F | is_closed X F ∧ A ⊆ F }

theorem closure_closed (X : topological_space α) (A : myset α) : is_closed X (closure X A) :=
begin
  unfold closure,
  unfold is_closed,
  rw myset.compl_sIntersection,
  apply open_union_open,
  intros S hS,
  cases hS with T hT,
  rw ←hT.right,
  from hT.left.left,
end

theorem compl_interior_compl_eq_closure (X : topological_space α) (A : myset α) :
(X.interior A.compl).compl = X.closure A :=
begin
  unfold interior,
  rw myset.compl_sUnion,
  congr,
  apply myset.setext,
  intro S,
  split; assume hS,
    cases hS with T hT,
    rw hT.right.symm,
    split,
      unfold is_closed,
      rw myset.compl_compl,
      from hT.left.left,
    apply myset.subset_compl_symm,
    from hT.left.right,
  unfold myset.image,
  existsi S.compl,
  split,
    split,
      from hS.left,
    rw myset.compl_subset_compl,
    from hS.right,
  rw myset.compl_compl,
end

theorem closure_minimal_closed_superset
(X: topological_space α) (A B: myset α):
X.is_closed B → A ⊆ B → X.closure A ⊆ B :=
begin
  assume hBcl hAB,
  intro x,
  assume hxAcl,
  from hxAcl B ⟨hBcl, hAB⟩,
end

theorem interior_maximal_open_subset
(X: topological_space α) (U V: myset α):
X.is_open V → V ⊆ U → V ⊆ X.interior U :=
begin
  assume hVo hVU,
  intro x,
  assume hVint,
  existsi V,
  split, {
    split; assumption,
  }, {
    assumption,
  },
end

theorem self_subset_closure
(X: topological_space α) (A: myset α):
A ⊆ X.closure A :=
begin
  intro a,
  assume hAa,
  intro F,
  assume hF,
  apply hF.right,
  assumption,
end

theorem interior_subset_self
(X: topological_space α) (U: myset α):
X.interior U ⊆ U :=
begin
  intro u,
  assume hUu,
  cases hUu with V hV,
  cases hV with hV huV,
  apply hV.right,
  assumption,
end

theorem closure_of_closed
(X: topological_space α) (A: myset α) (hAc: X.is_closed A):
X.closure A = A :=
begin
  apply myset.setext,
  intro x,
  split, {
    apply closure_minimal_closed_superset, {
      assumption,
    }, {
      refl,
    },
  }, {
    apply self_subset_closure _ A,
  },
end

theorem interior_of_open
(X: topological_space α) (U: myset α) (hUo: X.is_open U):
X.interior U = U :=
begin
  apply myset.setext,
  intro x,
  split, {
    apply interior_subset_self,
  }, {
    apply interior_maximal_open_subset _ _ U, {
      assumption,
    }, {
      refl,
    },
  },
end

theorem closure_idem_iff_closed
(X: topological_space α) (A: myset α):
X.closure A = A ↔ X.is_closed A :=
begin
  split, {
    assume h,
    rw ←h,
    apply closure_closed,
  }, {
    assume h,
    apply closure_of_closed,
    assumption,
  },
end

theorem interior_idem_iff_open
(X: topological_space α) (U: myset α):
X.interior U = U ↔ X.is_open U :=
begin
  split, {
    assume h,
    rw ←h,
    apply interior_open,
  }, {
    assume h,
    apply interior_of_open,
    assumption,
  },
end

theorem closure_subset
(X: topological_space α) (A B: myset α)
(hAB: A ⊆ B):
X.closure A ⊆ X.closure B :=
begin
  apply closure_minimal_closed_superset X A (X.closure B), {
    apply closure_closed,
  }, {
    apply myset.subset_trans A B (X.closure B), assumption,
    apply self_subset_closure,
  },
end

theorem interior_subset
(X: topological_space α) (U V: myset α)
(hUV: U ⊆ V):
X.interior U ⊆ X.interior V :=
begin
  apply interior_maximal_open_subset, {
    apply interior_open,
  }, {
    transitivity U, {
      apply interior_subset_self,
    }, {
      assumption,
    },
  },
end

theorem union_two_open
(X: topological_space α) (U V: myset α)
(hUo: X.is_open U) (hVo: X.is_open V):
X.is_open (U ∪ V) :=
begin
  rw myset.union_two_sUnion,
  apply X.open_union_open,
  intro S,
  assume hUVS,
  cases hUVS with h h; rw h; assumption,
end

theorem intersect_two_closed
(X: topological_space α) (U V: myset α)
(hUo: X.is_closed U) (hVo: X.is_closed V):
X.is_closed (U ∩ V) :=
begin
  unfold is_closed,
  -- should move into myset really
  have: (U ∩ V).compl = U.compl ∪ V.compl, {
    apply myset.setext,
    intro x,
    apply decidable.not_and_iff_or_not,
  },
  rw this,
  apply union_two_open; assumption,
end

-- if A ⊆ B is closed and B ⊆ C is closed,
-- then A is closed in C. yuckkkkkkkkkk
theorem closed_in_closed
(X: topological_space α) (U: myset α)
(hUcl: X.is_closed U)
(V: myset (subtype U))
(hVcl: (subspace_topology X U).is_closed V):
X.is_closed (myset.subtype_unrestriction U V) :=
begin
  cases hVcl with V' hV',
  cases hV' with hV'o hV',
  unfold myset.subtype_restriction at hV',
  have: (U.subtype_unrestriction V) = U ∩ V'.compl, {
    apply myset.setext,
    intro x,
    split, {
      assume hx,
      cases hx with hx hxV,
      split, {
        assumption,
      }, {
        assume hV'x,
        suffices hcontr: (⟨x, hx⟩: subtype U) ∈ V.compl, {
          contradiction,
        }, {
          rw hV',
          from hV'x,
        },
      },
    }, {
      assume hx,
      existsi hx.left,
      have: V.compl.compl = {w : subtype U | ↑w ∈ V'.compl}, {
        rw hV',
        refl,
      },
      rw myset.compl_compl at this,
      rw this,
      from hx.right,
    },
  },
  rw this,
  apply X.intersect_two_closed, {
    assumption,
  }, {
    unfold is_closed,
    rw myset.compl_compl,
    assumption,
  },
end

theorem closed_union_closed
(X: topological_space α) (U V: myset α)
(hUo: X.is_closed U) (hVo: X.is_closed V):
X.is_closed (U ∪ V) :=
begin
  unfold is_closed,
  have: (U ∪ V).compl = U.compl ∩ V.compl, {
    apply myset.setext,
    intro x,
    apply decidable.not_or_iff_and_not,
  },
  rw this,
  apply X.open_intersection_open; assumption,
end

theorem closed_intersection_closed
(X: topological_space α) (σ: myset (myset α))
(hSc: σ ⊆ X.is_closed):
X.is_closed ⋂₀ σ :=
begin
  rw ←myset.compl_compl (⋂₀σ),
  rw myset.compl_sIntersection,
  unfold is_closed,
  rw myset.compl_compl,
  apply open_union_open,
  intros S hS,
  unfold myset.image at hS,
  cases hS with T hT,
  rw hT.right.symm,
  from hSc _ hT.left,
end

-- finite intersections of open set are open,
-- stated with lists
theorem finite_open_intersection_open
(X: topological_space α) (σ : mylist (myset α))
(hso: mylist.for_all X.is_open σ):
X.is_open (mylist.reduce_d myset.intersection myset.univ σ) :=
begin
  induction σ, {
    from X.univ_open,
  }, {
    apply X.open_intersection_open, {
      from hso.left,
    }, {
      apply σ_ih,
      from hso.right,
    },
  },
end

theorem topological_space_eq:
∀ (X: topological_space α) (Y: topological_space α),
X.is_open = Y.is_open → X = Y
| ⟨Xo, _, _, _, _⟩ ⟨Yo, _, _, _, _⟩ rfl := rfl

theorem closure_intersects_all_open
(X: topological_space α) (A: myset α):
X.closure A =
{x : α | ∀ U: myset α, X.is_open U → x ∈ U → U ∩ A ≠ ∅} :=
begin
  apply myset.setext,
  intro x,
  split; assume hx, {
    intro U,
    assume hUo,
    assume hxU,
    assume hUAdisj,
    suffices h: x ∈ X.closure A ∩ U.compl, {
      from h.right hxU,
    }, {
      apply hx,
      split, {
        apply intersect_two_closed, {
          from closure_closed _ _,
        }, {
          unfold is_closed,
          rw myset.compl_compl,
          assumption,
        },
      }, {
        intro a,
        assume haA,
        split, {
          intro A',
          assume hA',
          apply hA'.right,
          assumption,
        }, {
          assume hUa,
          have: a ∈ U ∩ A, {
            split; assumption,
          },
          rw hUAdisj at this,
          from this,
        },
      },
    },
  }, {
    intro A',
    assume hA',
    by_contradiction h,
    suffices hcontr: A'.compl ∩ A ≠ ∅, {
      apply hcontr,
      apply myset.setext,
      intro x,
      split; assume hAx, {
        cases hAx with hA'x hAx,
        have := hA'.right x hAx,
        contradiction,
      }, {
        exfalso, from hAx,
      },
    },
    apply hx, {
      from hA'.left,
    }, {
      from h,
    },
  },
end

def is_dense (X : topological_space α) (A: myset α): Prop :=
closure X A = myset.univ

structure is_open_enclosure
(X: topological_space α) (U: myset α) (x: α)
(V: myset α): Prop :=
(is_open: X.is_open V)
(contains: x ∈ V)
(is_subset: V ⊆ U)

def is_neighbourhood (X: topological_space α) (U: myset α) (x: α): Prop :=
∃ V: myset α, is_open_enclosure X U x V

-- turns out neighbourhoods are useful!!!
-- I like this theorem because it means I don't have to unironically write
-- down unions all the time
theorem open_iff_neighbourhood_forall
(X: topological_space α) (U: myset α):
X.is_open U ↔ ∀ x: α, x ∈ U → is_neighbourhood X U x :=
begin
  split, {
    assume hUo,
    intro x,
    assume hxU,
    existsi U,
    split, {
      assumption,
    }, {
      assumption,
    }, {
      refl,
    },
  }, {
    assume hUnbhd,
    have hunion := X.open_union_open
      {V | X.is_open V ∧ V ⊆ U}
      begin
        intro V,
        assume hV,
        from hV.left,
      end,
    suffices heq: (⋃₀{V : myset α | X.is_open V ∧ V ⊆ U}) = U, {
      rw ←heq,
      assumption,
    }, {
      apply myset.setext,
      intro x,
      split; assume h, {
        cases h with V hV,
        cases hV with hV hxV,
        apply hV.right,
        assumption,
      }, {
        cases hUnbhd x h with V hV,
        existsi V,
        split, split, {
          from hV.is_open,
        }, {
          from hV.is_subset,
        }, {
          from hV.contains,
        },
      },
    },
  },
end

theorem basis_set_open
(B: myset (myset α)) (hB: is_base B) (U: myset α) (hUB: U ∈ B):
(space_from_base B hB).is_open U :=
begin
  intro x,
  assume hxU,
  existsi U,
  split, {
    assumption,
  }, split, {
    assumption,
  }, {
    refl,
  },
end

theorem hausdorff_iff_diagonal_closed
(X: topological_space α):
is_hausdorff X ↔ (product_topology X X).is_closed (λ xy, xy.fst = xy.snd) :=
begin
  split, {
    assume haus,
    unfold is_closed,
    rw open_iff_neighbourhood_forall,
    intro xy,
    assume hxydiag,
    have := haus xy.fst xy.snd hxydiag,
    cases this with U hUV,
    cases hUV with V hUV,
    existsi U × V,
    split, {
      apply basis_set_open,
      existsi U,
      existsi V,
      split, {
        refl,
      }, {
        from hUV.open_left,
      }, {
        from hUV.open_right,
      },
    }, {
      split, {
        from hUV.contains_left,
      }, {
        from hUV.contains_right,
      },
    }, {
      intro uv,
      assume huvUV,
      assume huvint,
      have := hUV.disjoint,
      rw ←myset.empty_iff_eq_empty at this,
      apply this uv.fst,
      split, {
        from huvUV.left,
      }, {
        change uv.fst = uv.snd at huvint,
        rw huvint,
        from huvUV.right,
      },
    },
  }, {
    assume hdiag,
    intros x y,
    assume hxy,
    cases hdiag ⟨x, y⟩ hxy with W hW,
    cases hW with hWpr hW,
    cases hWpr with U hUV,
    cases hUV with V hUV,
    existsi U,
    existsi V,
    rw hUV.eq at hW,
    split, {
      from hUV.open_left,
    }, {
      from hUV.open_right,
    }, {
      from hW.left.left,
    }, {
      from hW.left.right,
    }, {
      rw ←myset.empty_iff_eq_empty,
      intro uv,
      assume huv,
      apply hW.right ⟨uv, uv⟩, {
        from huv,
      }, {
        from rfl,
      },
    },
  },
end

theorem hausdorff_point_closed
(X: topological_space α) (h_ausdorff: is_hausdorff X)
(x: α):
is_closed X ({y: α | y = x}) :=
begin
  apply (open_iff_neighbourhood_forall _ _).mpr,
  intro y,
  assume hynex,
  cases h_ausdorff y x hynex with U hUV,
  cases hUV with V hUV,
  existsi U,
  split, {
    from hUV.open_left,
  }, {
    from hUV.contains_left,
  }, {
    intro z,
    assume hUz,
    assume hzx,
    have := hUV.disjoint,
    rw ←myset.empty_iff_eq_empty at this,
    apply this x,
    split, {
      have: z = x := hzx,
      rw ←this,
      assumption,
    }, {
      from hUV.contains_right,
    },
  },
end

theorem discrete_subspace
(S: myset α):
discrete_topology (subtype S) =
subspace_topology (discrete_topology α) S :=
begin
  apply topological_space_eq,
  apply myset.setext,
  intro X,
  split; assume h, {
    existsi myset.subtype_unrestriction S X,
    split, {
      trivial,
    }, {
      apply myset.setext,
      intro x,
      split, {
        assume hx,
        existsi x.property,
        have: (⟨↑x, x.property⟩: subtype S) = x, {
          apply subtype.eq,
          refl,
        },
        rw this,
        from hx,
      }, {
        assume hx,
        cases hx with hX hx,
        have: (⟨↑x, hX⟩: subtype S) = x, {
          apply subtype.eq,
          refl,
        },
        rw ←this,
        from hx,
      },
    },
  }, {
    trivial,
  },
end

theorem indiscrete_subspace
(S: myset α):
indiscrete_topology (subtype S) =
subspace_topology (indiscrete_topology α) S :=
begin
  apply topological_space_eq,
  apply myset.setext,
  intro X,
  split; assume hXo, {
    cases hXo with hXe hXu, {
      rw hXe,
      apply topological_space.empty_open,
    }, {
      rw hXu,
      apply topological_space.univ_open,
    },
  }, {
    cases hXo with U hXo,
    cases hXo with hUo hXU,
    cases hUo with hUe hUu, {
      rw hUe at hXU,
      left,
      rw hXU,
      refl,
    }, {
      rw hUu at hXU,
      right,
      rw hXU,
      refl,
    },
  },
end

theorem empty_discrete (X: topological_space α):
discrete_topology (subtype (∅: myset α)) =
subspace_topology X ∅ :=
begin
  apply topological_space_eq,
  apply myset.setext,
  intro U,
  split; {
    assume hUo,
    have: U = ∅, {
      apply myset.setext,
      intro x,
      split; assume h, {
        exfalso, from x.property,
      }, {
        exfalso, from h,
      },
    },
    rw this,
    apply topological_space.empty_open,
  },
end

-- lol same proof
theorem empty_indiscrete (X: topological_space α):
indiscrete_topology (subtype (∅: myset α)) =
subspace_topology X ∅ :=
begin
  apply topological_space_eq,
  apply myset.setext,
  intro U,
  split; {
    assume hUo,
    have: U = ∅, {
      apply myset.setext,
      intro x,
      split; assume h, {
        exfalso, from x.property,
      }, {
        exfalso, from h,
      },
    },
    rw this,
    apply topological_space.empty_open,
  },
end

theorem singleton_indiscrete (X: topological_space α) (x: α):
indiscrete_topology (subtype (myset.singleton x)) =
subspace_topology X (myset.singleton x) :=
begin
  apply topological_space_eq,
  apply myset.setext,
  intro S,
  by_cases hSe: S = ∅, {
    rw hSe,
    split; assume h,
    from topological_space.empty_open _,
    from topological_space.empty_open _,
  }, {
    change S ≠ ∅ at hSe,
    rw ←myset.exists_iff_neq_empty at hSe,
    cases hSe with y hy,
    have: S = myset.univ, {
      apply myset.setext,
      intro z,
      split; assume h, {
        trivial,
      }, {
        have: y = z, {
          let x': subtype {y : α | y = x} := ⟨x, rfl⟩,
          transitivity x', {
            have := y.property,
            apply subtype.eq,
            from this,
          }, {
            symmetry,
            have := z.property,
            apply subtype.eq,
            from this,
          },
        },
        rw this at hy,
        assumption,
      },
    },
    rw this,
    split; assume h,
    from topological_space.univ_open _,
    from topological_space.univ_open _,
  },
end

theorem singleton_discrete (X: topological_space α) (x: α):
discrete_topology (subtype (myset.singleton x)) =
subspace_topology X (myset.singleton x) :=
begin
  apply topological_space_eq,
  apply myset.setext,
  intro S,
  by_cases hSe: S = ∅, {
    rw hSe,
    split; assume h,
    from topological_space.empty_open _,
    from topological_space.empty_open _,
  }, {
    change S ≠ ∅ at hSe,
    rw ←myset.exists_iff_neq_empty at hSe,
    cases hSe with y hy,
    have: S = myset.univ, {
      apply myset.setext,
      intro z,
      split; assume h, {
        trivial,
      }, {
        have: y = z, {
          let x': subtype {y : α | y = x} := ⟨x, rfl⟩,
          transitivity x', {
            have := y.property,
            apply subtype.eq,
            from this,
          }, {
            symmetry,
            have := z.property,
            apply subtype.eq,
            from this,
          },
        },
        rw this at hy,
        assumption,
      },
    },
    rw this,
    split; assume h,
    from topological_space.univ_open _,
    from topological_space.univ_open _,
  },
end

theorem hausdorff_subspace
(X: topological_space α)
(h_ausdorff: is_hausdorff X)
(Y: myset α):
is_hausdorff (X.subspace_topology Y) :=
begin
  intros x y,
  assume hxney,
  have := h_ausdorff x.val y.val
    begin
      assume hxy,
      apply hxney,
      apply subtype.eq,
      assumption,
    end,
  cases this with U hUV,
  cases hUV with V hUV,
  existsi Y.subtype_restriction U,
  existsi Y.subtype_restriction V,
  split, {
    existsi U,
    split, {
      from hUV.open_left,
    }, {
      refl,
    },
  }, {
    existsi V,
    split, {
      from hUV.open_right,
    }, {
      refl,
    },
  }, {
    from hUV.contains_left,
  }, {
    from hUV.contains_right,
  }, {
    rw ←myset.empty_iff_eq_empty,
    intro z,
    assume hz,
    apply myset.empty_iff_eq_empty.mpr hUV.disjoint z.val,
    split, {
      from hz.left,
    }, {
      from hz.right,
    },
  },
end

end topological_space
end hidden


/-
Copyright (c) 2019 Neil Strickland. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Neil Strickland

-/

import data.vector data.list.sort data.fintype.basic algebra.big_operators
import data.fin_extra order.sort_rank data.heq_extra data.enumeration
import combinatorics.partition.basic combinatorics.stirling

/-
 Elsewhere, we have defined a type (partition α) for any finite
 type α, representing the set of partitions of α.  In particular,
 for any natural number n we have a finite type (fin n) of numbers
 less than n, and thus a type (partition (fin n)) of partitions.
 From this, we can extract a subtype of partitions of rank r
 (where the rank is defined to be the number of blocks).
 However, this is defined in a way that is not well-designed for
 efficient computation. The aim of this file is to define an equivalent
 type (fin_map n r) that has better computational properties.

 Specifically, an element of (fin_map n r) is essentially
 a structure consisting of maps b : (fin n) → (fin r) and
 t : (fin r) → (fin n) and proofs of three properties: that b ∘ t is
 the identity, that t ∘ b is everywhere greater than or equal to
 the identity, and that t is strictly increasing.  Given any such
 pair, we can partition (fin n) using the fibres of the map b.  It
 is not hard to see that each partition of rank r arises from a
 unique pair (b,t) in this way.

 The word "essentially" in the above paragraph covers the following
 technicality.  For any type α, maps (fin k) → α biject in an obvious
 way with lists in α of length k.  Such lists can be considered as
 elements of the subtype (vector α k).  Lists and vectors are standard
 data structures that can be represented efficiently in the Lean
 virtual machine, especially when the entries are just natural numbers.
 For this reason, in the definition of (fin_map n r), we
 encode b and t as vectors rather than functions.
-/

namespace combinatorics.partition
open fin fintype combinatorics combinatorics.partition

section block_max

variables {n : ℕ} {p : partition (fin n)}

/-
 As a basic ingredient of all the constructions in this file, we need
 to order the blocks of a given partition by declaring that b < c
 iff (largest element of b) < (largest element of c), and we need to
 sort the set of blocks according to this order.
-/

/-
 We start with a function that extracts the largest element of each
 block.  For this to make sense we need to know that each block is
 nonempty; this is proved by the lemma block_not_empty from
 partition.lean.
-/
def block_max (b : p.block_type) : fin n :=
 (finset_largest_element b.val (block_not_empty b.property)).val

/-
 This defines a lemma showing that (block_max b) has the expected
 defining properties (it lies in b and is maximal subject to that).
-/
def block_max_spec (b : p.block_type) :=
 (finset_largest_element b.val (block_not_empty b.property)).property

/-
 The block containing (block_max b) is just b.
-/
lemma block_max_block (b : p.block_type) :
 b = p.block_alt (block_max b) :=
  block_eq_of_mem (block_max_spec b).left

variable (p)

/-
 We now define our order on the set of blocks, packaging it as an
 instance of the linear_order typeclass.
-/

instance block_order : linear_order p.block_type := {
 le := λ b c ,(block_max b) ≤ (block_max c),
 lt := λ b c ,(block_max b) < (block_max c),
 lt_iff_le_not_le := λ b c, @lt_iff_le_not_le (fin n) _ (block_max b) (block_max c),
 le_refl := λ b,le_refl (block_max b),
 le_trans := λ (a b c : p.block_type)
   (a_le_b : block_max a ≤ block_max b)
   (b_le_c : block_max b ≤ block_max c),
    le_trans a_le_b b_le_c,
 le_antisymm := λ (a b : p.block_type)
   (a_le_b : block_max a ≤ block_max b)
   (b_le_a : block_max b ≤ block_max a),
   begin
    have s_eq_s : block_max a = block_max b :=
     le_antisymm a_le_b b_le_a,
    exact
     ((block_max_block a).trans
      (congr_arg p.block_alt s_eq_s)).trans
     (block_max_block b).symm,
   end,
 le_total := λ (a b : p.block_type),
   le_total (block_max a) (block_max b),
 decidable_le := λ b c, fin.decidable_le (block_max b) (block_max c)
}

/-
 If the number of blocks is r, then the following function returns the
 unique order-equivalence from the set of blocks to (fin r).
-/
def block_rank_equiv {r : ℕ} (p_rank : p.rank = r) :=
 fintype.rank_equiv ((block_type_card p).trans p_rank)

end block_max

/-
 We now define the type (fin_map n r).  The elements are
 structures with data fields block and top, together with three
 propositional fields (which will be ignored in compiled code for
 the Lean virtual machine).  The block field is a vector of length
 n with entries in (fin r).  This gives rise to a map
 (fin n) → (fin r), which is called block.nth in Lean notation.
 Similarly, the top field is a vector of length r with entries in
 (fin n), corresponding to a map top.nth : (fin r) → (fin n).
-/
structure fin_map (n r : ℕ) :=
(block : vector (fin r) n)
(top : vector (fin n) r)
(block_top : ∀ j, block.nth (top.nth j) = j)
(top_block : ∀ i, top.nth (block.nth i) ≥ i)
(top_mono : ∀ j1 j2, j1 < j2 → top.nth j1 < top.nth j2)

/- -------------------------------------------------------------- -/
/- -------------------------------------------------------------- -/
/- -------------------------------------------------------------- -/
/- -------------------------------------------------------------- -/

namespace fin_map
open fin_map

def tot (n : ℕ) := Σ (r : ℕ), fin_map n r

def to_tot {n r : ℕ} (p : fin_map n r) : tot n := ⟨r,p⟩

variables {n r : ℕ}

/-
 The next definition tells Lean how to convert a term of type
 `(fin_map n r)` to a string, for display to the user.
 Lean already knows how to display terms of type `(fin r)`, and
 also knows a rule for how to display lists if it already knows
 how to display each entry.  If `p` is of type
 `(fin_map n r)` then `p.block` is a vector and so
 `p.block.val` is a list.  We tell Lean to use the string
 representation of `p.block.val` as a representation for `p`
 itself.  (Lemma `eq_of_block_eq` below will verify that
 `p.block.val` determines all of `p`.)
-/
instance (n r : ℕ) : has_repr (fin_map n r) :=
 ⟨λ p, repr p.block.val⟩

/-
 One of our axioms is that the map top : (fin r) → (fin n) should
 preserve strict inequalities.  We now prove the obvious consequence
 that it also preserves weak inequalities.
-/
lemma top_mono_weak (p : fin_map n r) (j1 j2 : fin r) :
 j1 ≤ j2 → p.top.nth j1 ≤ p.top.nth j2 :=
begin
 by_cases e : j1 = j2,
 {rw[e],intro,exact le_refl _,},
 {intro j1_le_j2,
  exact le_of_lt (p.top_mono j1 j2 (lt_of_le_of_ne j1_le_j2 e)),
 }
end

lemma top_mono_iff (p : fin_map n r) (j1 j2 : fin r) :
 j1 < j2 ↔ p.top.nth j1 < p.top.nth j2 :=
begin
 split,
 {exact top_mono p j1 j2},
 {intro top_lt,
  by_cases h : j1 < j2,
  {exact h},
  {exfalso,
   let top_ge := top_mono_weak p j2 j1 (le_of_not_lt h),
   exact not_lt_of_ge top_ge top_lt,
  }
 }
end

@[ext]
lemma eq_of_block_eq (p1 p2 : fin_map n r)
 (e_block : ∀ i, p1.block.nth i = p2.block.nth i) : p1 = p2 :=
begin
 rcases p1,
 rcases p2,
 simp at e_block,
 have e_block1 : p1_block = p2_block := vector.ext e_block,
 have e_top : ∀ j, p1_top.nth j = p2_top.nth j :=
 begin
  intro j,
  let h0 := (congr_arg p1_top.nth ((e_block (p2_top.nth j)).trans
             (p2_block_top j))).subst (p1_top_block (p2_top.nth j)),
  let h1 := (congr_arg p2_top.nth ((e_block (p1_top.nth j)).symm.trans
             (p1_block_top j))).subst (p2_top_block (p1_top.nth j)),
  exact le_antisymm h1 h0,
 end,
 have e_top1 : p1_top = p2_top := vector.ext e_top,
 simp[e_block1,e_top1],
end

/-
 This function converts an element of (fin_map n r) to
 the corresponding partition of the set (fin n).
-/
def to_partition (p : fin_map n r) :
 partition (fin n) := partition.fiber_partition p.block.nth

lemma to_partition_top (p : fin_map n r) (i : fin n) :
 block_max ((to_partition p).block_alt i) =
  p.top.nth (p.block.nth i) :=
begin
 let q := to_partition p,
 let b := q.block_alt i,
 have i_in_b : i ∈ b.val := partition.mem_block _ i,
 let j := p.block.nth i,
 let i1 := p.top.nth j,
 have i_blocks0 : p.block.nth i1 = j := p.block_top j,
 have i_blocks1 : q.block_alt i1 = b :=
  (partition.fiber_partition_blocks_alt p.block.nth i1 i).mpr i_blocks0,
 have i_blocks2 : q.block i1 = b.val := partition.block_veq_of_eq i_blocks1,
 have i1_in_b : i1 ∈ b.val := i_blocks2.subst (partition.mem_block q i1),
 let m := block_max b,
 let m_in_b : m ∈ b.val := (block_max_spec b).left,
 let i1_le_m : i1 ≤ m := (block_max_spec b).right i1 i1_in_b,
 have m_block : b = q.block_alt m := partition.block_eq_of_mem m_in_b,
 let h0 := (congr_arg p.top.nth
           ((partition.fiber_partition_blocks_alt p.block.nth i1 m).mp
            (i_blocks1.trans m_block))),
 let h1 := p.top_block m,
 let h2 := h0.symm.trans (congr_arg p.top.nth i_blocks0),
 let m_le_i1 : m ≤ i1 := h2.subst h1,
 exact le_antisymm m_le_i1 i1_le_m,
end

/-
 Given an element of (fin_map n r), the above definition
 gives us a partition of (fin n), and thus a set of blocks.  The
 next definition gives a bijection from that block set to (fin r).
 This bijection is packaged using the framework set up in the
 standard library file data.equiv.lean, so we have a structure
 involving maps to_fun and inv_fun together with proofs (named
 left_inv and right_inv) that their composites are equal
 (pointwise) to identity functions.
-/
def to_partition_blocks (p : fin_map n r) :
 (to_partition p).block_type ≃ fin r :=
begin
 let q := (to_partition p),
 let to_fun : q.block_type → fin r :=
  λ b, p.block.nth (block_max b),
 let inv_fun : fin r → q.block_type :=
  λ j, q.block_alt (p.top.nth j),
 have left_inv : ∀ b, inv_fun (to_fun b) = b :=
 begin
  intro b,
  let i1 := block_max b,
  let j1 := p.block.nth i1,
  let i2 := p.top.nth j1,
  let e : p.block.nth i2 = p.block.nth i1 := p.block_top j1,
  let h0 : q.block_alt i1 = q.block_alt i2 :=
   subtype.eq ((partition.fiber_partition_blocks p.block.nth i1 i2).mpr e.symm),
  let h1 : b = q.block_alt i1 := block_max_block b,
  exact (h1.trans h0).symm
 end,
 have right_inv : ∀ j, to_fun (inv_fun j) = j :=
 begin
  intro j,
  let i1 := p.top.nth j,
  let b := q.block_alt i1,
  have i1_in_b : i1 ∈ b.val := q.mem_block i1,
  let i2 := block_max b,
  let h0 := block_max_spec b,
  have i2_in_b : i2 ∈ b.val := h0.left,
  have i1_le_i2 : i1 ≤ i2 := h0.right i1 i1_in_b,
  have i2_block : q.block i1 = q.block i2 :=
   partition.block_alt_eq.mpr (block_max_block b),
  let h1 : p.block.nth i1 = p.block.nth i2 :=
    (partition.fiber_partition_blocks p.block.nth i1 i2).mp i2_block,
  let h2 : p.block.nth i1 = j := p.block_top j,
  let h3 : j = p.block.nth i2 := h2.symm.trans h1,
  let h4 : p.top.nth (p.block.nth i2) = i1 := (congr_arg p.top.nth h3).symm,
  let h5 : p.top.nth (p.block.nth i2) ≥ i2 := (p.top_block i2),
  let i1_ge_i2 : i1 ≥ i2 := h4.subst h5,
  let i1_eq_i2 := le_antisymm i1_le_i2 i1_ge_i2,
  have h6 : to_fun (inv_fun j) = p.block.nth i2 := rfl,
  exact (h6.trans (congr_arg p.block.nth i1_eq_i2.symm)).trans h2,
 end,
 exact ⟨to_fun,inv_fun,left_inv,right_inv⟩
end

/-
 From the above bijection it is clear that the relevant partition
 of (fin n) has rank r.
-/
lemma to_partition_rank (p : fin_map n r) :
 (to_partition p).rank = r :=
begin
 let q := to_partition p,
 let e0 := (partition.block_type_card q).symm,
 let e1 := fintype.card_congr (to_partition_blocks p),
 let e2 := fintype.card_fin r,
 have e3 := subsingleton.elim (fin.fintype r) (@enumeration_fintype (fin r) (fin.enumeration r)),
 rw[e3] at e2,
 simp[e0,e1,e2],
end

lemma to_partition_block_type_card (p : fin_map n r) :
 fintype.card (to_partition p).block_type = r :=
  (partition.block_type_card (to_partition p)).trans (to_partition_rank p)

lemma to_partition_blocks_ordered (p : fin_map n r) (j1 j2 : fin r) :
 j1 < j2 ↔ (to_partition_blocks p).inv_fun j1 < (to_partition_blocks p).inv_fun j2 :=
begin
 let q := to_partition p,
 let i1 := p.top.nth j1,
 let i2 := p.top.nth j2,
 let b1 := q.block_alt i1,
 let b2 := q.block_alt i2,
 let max1 : block_max b1 = i1 :=
  (to_partition_top p i1).trans (congr_arg p.top.nth (p.block_top j1)),
 let max2 : block_max b2 = i2 :=
  (to_partition_top p i2).trans (congr_arg p.top.nth (p.block_top j2)),
 let h : j1 < j2 ↔ i1 < i2 := top_mono_iff p j1 j2,
 rw[max1.symm,max2.symm] at h,
 exact h,
end

lemma to_partition_blocks_ordered_alt (p : fin_map n r)
 (b1 b2 : (to_partition p).block_type) :
  b1 < b2 ↔ (to_partition_blocks p).to_fun b1 < (to_partition_blocks p).to_fun b2 :=
begin
 let j1 := (to_partition_blocks p).to_fun b1,
 let j2 := (to_partition_blocks p).to_fun b2,
 let h := (to_partition_blocks_ordered p j1 j2).symm,
 rw[(to_partition_blocks p).left_inv] at h,
 rw[(to_partition_blocks p).left_inv] at h,
 exact h,
end

lemma block_rank_equiv_eq (p : fin_map n r) :
 (block_rank_equiv (to_partition p) (to_partition_rank p)) =
  (to_partition_blocks p) :=
begin
 let q := to_partition p,
 let f := block_rank_equiv q (to_partition_rank p),
 let g := to_partition_blocks p,
 let h : fin r → fin r := λ j, f.to_fun (g.inv_fun j),
 have h_mono : ∀ j1 j2, j1 < j2 → (h j1) < (h j2) :=
 begin
  intros j1 j2 j1_lt_j2,
  let r0 : g.inv_fun j1 < g.inv_fun j2 :=
   (to_partition_blocks_ordered p j1 j2).mp j1_lt_j2,
  exact (@fintype.rank_lt q.block_type _ _ _ r
         (to_partition_block_type_card p) (g.inv_fun j1) (g.inv_fun j2)).mp r0,
 end,
 have h_id : ∀ j, h j = j := fin.rigid h_mono,
 ext b, congr' 1,
 exact (congr_arg f.to_fun (g.left_inv b)).symm.trans (h_id (g.to_fun b)),
end

/-
 This function accepts a partition of (fin n) of rank r and
 produces the corresponding element of (fin_map n r).
-/
def of_partition (q : partition (fin n)) (e : q.rank = r) :
 fin_map n r :=
begin
 let r0 := fintype.card q.block_type,
 have e0 : r0 = r := q.block_type_card.trans e,
 let f := block_rank_equiv q e,
 let block_map : fin n → fin r := λ i, f.to_fun (q.block_alt i),
 let top_map : fin r → fin n := λ j, block_max (f.inv_fun j),
 let block_vec := vector.of_fn block_map,
 let top_vec := vector.of_fn top_map,
 have block_top : ∀ j, block_vec.nth (top_vec.nth j) = j :=
 begin
  intro j,dsimp[block_vec,top_vec],
  simp only [vector.nth_of_fn top_map,vector.nth_of_fn block_map],
  dsimp[top_map,block_map],
  let h0 := (block_max_block (f.symm j)).symm,
  rw[h0],
  exact f.right_inv j,
 end,
 have top_block : ∀ i, top_vec.nth (block_vec.nth i) ≥ i :=
 begin
  intro i,dsimp[top_vec,block_vec],
  simp only [vector.nth_of_fn top_map,vector.nth_of_fn block_map],
  dsimp[top_map,block_map],
  let h0 : f.symm (f _) = _ := f.left_inv (q.block_alt i),
  rw[h0],
  let h1 := (block_max_spec (q.block_alt i)).right i
             (partition.mem_block q i),
  exact h1,
 end,
 have top_mono : ∀ j1 j2, j1 < j2 → top_vec.nth j1 < top_vec.nth j2 :=
 begin
  intros j1 j2 j1_lt_j2,
  simp only[top_vec,vector.nth_of_fn,top_map],
  exact (@fintype.seq_lt q.block_type
   (combinatorics.partition.block_order q) _ _ r e0 j1 j2).mp j1_lt_j2,
 end,
 exact {
   block := block_vec,
   top := top_vec,
   block_top := block_top,
   top_block := top_block,
   top_mono := top_mono
 },
end

lemma of_partition_block (q : partition (fin n)) (e : q.rank = r) (i : fin n) :
 ((of_partition q e).block.nth i) =
  ((block_rank_equiv q e).to_fun (q.block_alt i)) :=
begin
 cases e,
 simp[of_partition],
end

lemma of_partition_to_partition (p : fin_map n r) :
 of_partition (to_partition p) (to_partition_rank p) = p :=
begin
 ext i,
 simp[of_partition],
 rw[block_rank_equiv_eq p],
 dsimp[to_partition_blocks],
 rw[to_partition_top p i,p.block_top (p.block.nth i)]
end

lemma to_partition_of_partition (q : partition (fin n)) (e : q.rank = r) :
 to_partition (of_partition q e) = q :=
begin
 let p := of_partition q e,
 let r := to_partition p,
 ext i1 i2,
 let b1 := q.block_alt i1,
 let b2 := q.block_alt i2,
 let j1 := p.block.nth i1,
 let j2 := p.block.nth i2,
 have h0 : i2 ∈ r.block i1 ↔ j2 = j1 :=
 begin
  let h1 := partition.fiber_partition_blocks p.block.nth i2 i1,
  split; intro h2,
  {
    have h3 : i2 ∈ (r.block_alt i1).val := h2,
    exact h1.mp (partition.block_val_eq_of_mem h3).symm,
  },{
    exact (h1.mpr h2).subst (partition.mem_block r i2),
  }
 end,
 rw[h0],
 have h1 : i2 ∈ q.block i1 ↔ b2 = b1 :=
 begin
  split,
  {intro i2_in_b1,
   let i2_in_b1_alt : i2 ∈ b1.val := i2_in_b1,
   exact (partition.block_eq_of_mem i2_in_b1_alt).symm,
  },{
   intro b2_eq_b1,
   exact (partition.block_veq_of_eq b2_eq_b1).subst (partition.mem_block q i2),
  }
 end,
 rw[h1],
 let f := block_rank_equiv q e,
 have e1a : j1 = f.to_fun b1 := of_partition_block q e i1,
 have e2a : j2 = f.to_fun b2 := of_partition_block q e i2,
 have e1b : b1 = f.inv_fun j1 := ((congr_arg f.inv_fun e1a).trans (f.left_inv b1)).symm,
 have e2b : b2 = f.inv_fun j2 := ((congr_arg f.inv_fun e2a).trans (f.left_inv b2)).symm,
 split,
 {intro j1_eq_j2,rw[e1b,e2b,j1_eq_j2],},
 {intro b1_eq_b2,rw[e1a,e2a,b1_eq_b2],},
end

def partition_equiv : (fin_map n r) ≃ { q : partition (fin n) // q.rank = r} :=
{
 to_fun := λ p, ⟨to_partition p,to_partition_rank p⟩,
 inv_fun := λ q, of_partition q.val q.property,
 left_inv := λ p, of_partition_to_partition p,
 right_inv := begin
  intro q,
  cases q,
  exact subtype.eq (to_partition_of_partition q_val q_property),
 end
}

def partition_equiv_tot : tot n ≃ partition (fin n) :=
{
 to_fun := λ rp, to_partition rp.2,
 inv_fun := λ q, ⟨q.rank,@of_partition n q.rank q rfl⟩,
 right_inv := λ q, @to_partition_of_partition n q.rank q rfl,
 left_inv :=
 begin
  intro rp,cases rp with r p,
  let q := @to_partition n r p,
  let r_eq := to_partition_rank p,
  let p_eq := of_partition_to_partition p,
  let C : ∀ (r1 : ℕ) (u1 : q.rank = r1), Prop :=
   λ r1 u1, @of_partition n q.rank q rfl == @of_partition n r1 q u1,
  let c : C q.rank rfl := by simp[C],
  simp[r_eq,p_eq],
  exact (@eq.dcases_on ℕ q.rank C r r_eq c).trans (heq_of_eq p_eq),
 end
}

/- A convenience function for manipulating vectors. -/
def bump {p q : ℕ} (v : vector (fin p) q) : vector (fin p.succ) q.succ :=
 vector.cons (0 : fin p.succ) (vector.map fin.succ v)

/- Three obvious lemmas about vectors and the corresponding functions. -/
lemma vector.zth {α : Type*} {n : ℕ} (a : α) (v : vector α n) (h : 0 < n.succ) :
 (vector.cons a v).nth ⟨0,h⟩ = a :=
  vector.nth_cons_zero a v

lemma vector.sth {α : Type*} {n : ℕ} (a : α) (v : vector α n)
 (i0 : ℕ) (h0 : i0 < n) (h1 : i0.succ < n.succ):
 (vector.cons a v).nth ⟨i0.succ,h1⟩ = v.nth ⟨i0,h0⟩ :=
  vector.nth_cons_succ a v ⟨i0,h0⟩

lemma vector.nth_map {α β : Type*} (f : α → β) {n : ℕ} (v : vector α n) (i : fin n) :
 (vector.map f v).nth i = f (v.nth i) :=
begin
 cases v,cases i,
 dsimp[vector.map,vector.nth],
 apply list.nth_le_map
end

lemma bump.zth {p q : ℕ} (v : vector (fin p) q) : (bump v).nth 0 = 0 :=
 by { apply vector.zth }

lemma bump.sth {p q : ℕ} (v : vector (fin p) q) (j : fin q) : (bump v).nth j.succ = (v.nth j).succ :=
 by { dsimp[bump], rw[vector.nth_cons_succ,vector.nth_map] }
/-
 Suppose that p is a partition of (fin n), of rank r.  We can then
 produce a partition of (fin (n + 1)) of rank r + 1, as follows:
 we take 0 as a block on its own, and then shift all the blocks of p
 up by one.  The resulting partition is called (add p).
-/
def add (p : fin_map n r) : fin_map n.succ r.succ :=
begin
 let block := bump p.block,
 let top := bump p.top,
 have block_top : ∀ j, block.nth (top.nth j) = j :=
 begin
  intro j,cases j with j_val j_is_lt,cases j_val with j0_val; dsimp[block,top,bump],
  {simp[vector.zth]},
  {
   have j0_is_lt : j0_val < r := nat.lt_of_succ_lt_succ j_is_lt,
   let j0 : fin r := ⟨j0_val,j0_is_lt⟩,
   simp [vector.sth _ _ j0_val j0_is_lt,vector.nth_map,p.block_top j0,fin.succ],
  }
 end,
 have top_block : ∀ i, top.nth (block.nth i) ≥ i :=
 begin
  intro i,cases i with i_val i_is_lt,cases i_val with i0_val; dsimp[block,top,bump],
  {simp[vector.zth]},
  {
   have i0_is_lt : i0_val < n := nat.lt_of_succ_lt_succ i_is_lt,
   let i0 : fin n := ⟨i0_val,i0_is_lt⟩,
   simp[vector.sth _ _ i0_val i0_is_lt,vector.nth_map],
   exact fin.succ_le_succ (p.top_block i0),
  }
 end,
 have top_mono : ∀ j1 j2, j1 < j2 → top.nth j1 < top.nth j2 :=
 begin
  intros j1 j2 j1_lt_j2,
  cases j2 with j2_val j2_is_lt,cases j2_val with j20_val,
  {exact false.elim (nat.not_lt_zero j1.val j1_lt_j2),},
  {have j20_is_lt : j20_val < r := nat.lt_of_succ_lt_succ j2_is_lt,
   cases j1 with j1_val j1_is_lt,cases j1_val with j10_val,
   {simp[top,bump,vector.zth,vector.sth _ _ j20_val j20_is_lt,vector.nth_map,fin.zero_lt_succ],},
   {have j10_is_lt : j10_val < r := nat.lt_of_succ_lt_succ j1_is_lt,
     simp[top,bump,
          vector.sth _ _ j10_val j10_is_lt,
          vector.sth _ _ j20_val j20_is_lt,
          vector.nth_map],
    apply p.top_mono,
    exact fin.lt_of_succ_lt_succ j1_lt_j2,
   }
  }
 end,
 exact ⟨block,top,block_top,top_block,top_mono⟩
end

lemma add_block_zero (p : fin_map n r) : p.add.block.nth 0 = 0 :=
 by { dsimp[add], apply bump.zth }

lemma add_top_zero (p : fin_map n r) : p.add.top.nth 0 = 0 :=
 by { dsimp[add], apply bump.zth }

lemma add_block_succ (p : fin_map n r) (i : fin n) : p.add.block.nth i.succ = (p.block.nth i).succ :=
 by { dsimp[add], apply bump.sth }

lemma add_top_succ (p : fin_map n r) (i : fin r) : p.add.top.nth i.succ = (p.top.nth i).succ :=
 by { dsimp[add], apply bump.sth }

/-
 Suppose that p is a partition of (fin n), of rank r, and that u is
 in (fin r).  We can then produce a partition of (fin (n + 1)) of rank r,
 as follows: we shift all the blocks of p up by one, and then add 0 as
 an extra element in the u'th block.  The resulting partition is called
 (join p u).
-/
def join (p : fin_map n r) (u : fin r) : fin_map n.succ r :=
begin
 let block := vector.cons u p.block,
 let top := vector.map fin.succ p.top,
 have block_top : ∀ j, block.nth (top.nth j) = j :=
 begin
  intro j,dsimp[top,block],
  simp[vector.nth_map,p.block_top j],
 end,
 have top_block : ∀ i, top.nth (block.nth i) ≥ i :=
 begin
  intro i, cases i with i_val i_is_lt, cases i_val with i0_val,
  { change top.nth (block.nth 0) ≥ 0, rw[vector.nth_cons_zero], exact fin.zero_le _,},
  {have i0_is_lt : i0_val < n := nat.lt_of_succ_lt_succ i_is_lt,
   let h0 := fin.succ_le_succ (p.top_block ⟨i0_val,i0_is_lt⟩),
   dsimp[block,top],
   simp[vector.sth _ _ i0_val i0_is_lt,vector.nth_map],
   exact h0,
   }
 end,
 have top_mono : ∀ j1 j2, j1 < j2 → top.nth j1 < top.nth j2 :=
 begin
  intros j1 j2 j1_lt_j2,
  simp[top,vector.nth_map],
  exact p.top_mono j1 j2 j1_lt_j2
 end,
 exact ⟨block,top,block_top,top_block,top_mono⟩
end

lemma join_block_zero (p : fin_map n r) (u : fin r) : (p.join u).block.nth 0 = u :=
 by { dsimp[join], apply vector.zth, }

lemma join_block_succ (p : fin_map n r) (u : fin r) (i : fin n) :
 (p.join u).block.nth i.succ = p.block.nth i :=
  by { dsimp[join], apply vector.nth_cons_succ }

lemma join_top (p : fin_map n r) (u : fin r) (j : fin r) :
 (p.join u).top.nth j = (p.top.nth j).succ :=
  by { dsimp[join], apply vector.nth_map }

def add_cond (q : fin_map n.succ r.succ) := (q.top.nth 0 = 0)

def join_cond (q : fin_map n.succ r.succ) := (q.top.nth 0 ≠ 0)

lemma add_add_cond (p : fin_map n r) :
 add_cond (add p) :=
begin
 dsimp[add_cond,add,bump],
 exact vector.zth 0 _ _,
end

lemma join_join_cond (p : fin_map n r.succ) (u : fin r.succ) :
 join_cond (join p u) :=
begin
 intro e,
 exact fin.zero_ne_succ (e.symm.trans (vector.nth_map fin.succ p.top 0)),
end

/-
 Let q be a partition of (fin (n + 1)), with rank r + 1.  Suppose also
 that the map q.top.nth : (fin (r + 1)) → (fin (n + 1)) sends 0 to 0.
 It is then not hard to see that there is a unique p in
 (fin_map n r) with (add p) = q.  The present function produces
 this p.
-/
def add_res (q : fin_map n.succ r.succ) (e : add_cond q) :
 {p : fin_map n r // q = add p} :=
begin
 let block_map_aux : ∀ (i0 : fin n), { j0 : fin r // j0.succ = q.block.nth i0.succ } :=
 begin
  dsimp[add_cond] at e,
  intro i0,
  let j := q.block.nth i0.succ,
  by_cases e1 : j = 0,
  {exfalso,
   let h0 : i0.succ ≤ q.top.nth j := q.top_block i0.succ,
   let h1 : 0 < q.top.nth j :=
    @lt_of_lt_of_le (fin n.succ) _ 0 i0.succ _ fin.zero_lt_succ h0,
    rw[e1,e] at h1,exact fin.not_lt_zero (0 : fin n.succ) h1,
  },{
   exact ⟨ fin.pred j e1, fin.succ_pred j e1⟩,
  }
 end,
 let block_map := λ i0, (block_map_aux i0).val,
 let block_map_spec : ∀ i0, (block_map i0).succ = q.block.nth i0.succ :=
  λ i0,(block_map_aux i0).property,
 let top_map_aux : ∀ (j0 : fin r), { i0 : fin n // i0.succ = q.top.nth j0.succ }:=
 begin
  dsimp[add_cond] at e,
  intro j0,
  let i := q.top.nth j0.succ,
  let e1 := (ne_of_lt (q.top_mono 0 j0.succ fin.zero_lt_succ)).symm,
  rw[e] at e1,
  exact ⟨fin.pred i e1, fin.succ_pred i e1⟩
 end,
 let top_map := λ j0, (top_map_aux j0).val,
 let top_map_spec : ∀ j0, (top_map j0).succ = q.top.nth j0.succ :=
  λ j0,(top_map_aux j0).property,
 let block_vec := vector.of_fn block_map,
 let top_vec := vector.of_fn top_map,
 have block_top : ∀ j0, block_vec.nth (top_vec.nth j0) = j0 :=
 begin
  intro j0,
  dsimp[block_vec,top_vec],
  simp[vector.nth_of_fn],
  apply fin.succ_inj.mp,
  let h0 := block_map_spec (top_map j0),
  let h1 := congr_arg q.block.nth (top_map_spec j0),
  let h2 := q.block_top j0.succ,
  let h3 := h0.trans (h1.trans h2),
  exact h3,
 end,
 have top_block : ∀ i0, top_vec.nth (block_vec.nth i0) ≥ i0 :=
 begin
  intro i0,
  dsimp[block_vec,top_vec],
  simp[vector.nth_of_fn],
  apply fin.le_of_succ_le_succ,
  let h0 := top_map_spec (block_map i0),
  let h1 := congr_arg q.top.nth (block_map_spec i0),
  let h2 : i0.succ ≤ q.top.nth (q.block.nth i0.succ) := (q.top_block i0.succ),
  let h3 := (h0.trans h1).symm.subst h2,
  exact h3,
 end,
 have top_mono : ∀ j10 j20, j10 < j20 → top_vec.nth j10 < top_vec.nth j20 :=
 begin
  intros j10 j20 j10_lt_j20,
  dsimp[block_vec,top_vec],
  simp[vector.nth_of_fn],
  apply fin.lt_of_succ_lt_succ,rw[top_map_spec,top_map_spec],
  exact q.top_mono j10.succ j20.succ (fin.succ_lt_succ j10_lt_j20),
 end,
 let p : fin_map n r := ⟨block_vec,top_vec,block_top,top_block,top_mono⟩,
 have eq : q = add p :=
 begin
  ext i,cases i with i_val i_is_lt, cases i_val with i0_val,
  { have : (0 : fin n.succ) = ⟨0,i_is_lt⟩ := rfl, rw[← this],
    have h₀ := (congr_arg q.block.nth e).symm,
    have h₁ := q.block_top 0,
    have h₂ := h₀.trans h₁,
    rw[h₂,fin_map.add],simp only[bump.zth]
  },{
   dsimp[add,p,bump],
   let i0_is_lt := nat.lt_of_succ_lt_succ i_is_lt,
   have h₀ := vector.sth (0 : fin r.succ) (vector.map fin.succ block_vec) i0_val i0_is_lt i_is_lt,
   rw[h₀,vector.nth_map],dsimp[block_vec],rw[vector.nth_of_fn,block_map_spec],
   refl
  }
 end,
 exact ⟨p,eq⟩,
end

lemma res_add (p : fin_map n r) (e : add_cond (add p)) :
 (add_res (add p) e) = ⟨p,rfl⟩ :=
begin
 apply subtype.eq,
 cases (p.add.add_res e) with p₀ hp₀,simp only [],
 have : ∀ (i : fin n), p.block.nth i = p₀.block.nth i := λ i,
  begin
   apply fin.succ_inj.mp,
   rw[← add_block_succ,← add_block_succ,hp₀],
  end,
 ext i,
 rw[this]
end

/-
 Let q be a partition of (fin (n + 1)), with rank r + 1.  Suppose also
 that the map q.top.nth : (fin (r + 1)) → (fin (n + 1)) does not send 0
 to 0.  It is then not hard to see that there is a unique p in
 (fin_map n (r + 1)) and u in (fin (r + 1)) with (join p u) = q.
 The present function produces this p and u.
-/
def join_res  (q : fin_map n.succ r.succ) (e : q.top.nth 0 ≠ 0) :
 { pu : (fin_map n r.succ) × (fin r.succ) // q = (join pu.1 pu.2) } :=
begin
 have e0 : 0 < q.top.nth 0 :=
  lt_of_le_of_ne (fin.zero_le (q.top.nth 0)) e.symm,
 let block_map : fin n → fin r.succ := λ i0, q.block.nth i0.succ,
 let top_map_aux : ∀ (j : fin r.succ), { i0 : fin n // i0.succ = q.top.nth j }:=
 begin
  intro j,
  let i := q.top.nth j,
  let e1 : 0 < q.top.nth j := lt_of_lt_of_le e0 (top_mono_weak q 0 j (fin.zero_le j)),
  let e2 := (ne_of_lt e1).symm,
  exact ⟨fin.pred i e2, fin.succ_pred i e2⟩
 end,
 let top_map := λ j, (top_map_aux j).val,
 let top_map_spec : ∀ j, (top_map j).succ = q.top.nth j :=
  λ j,(top_map_aux j).property,
 let block_vec := vector.of_fn block_map,
 let top_vec := vector.of_fn top_map,
 have block_top : ∀ j, block_vec.nth (top_vec.nth j) = j :=
 begin
  intro j,
  dsimp[block_vec,top_vec],
  simp[vector.nth_of_fn,block_map],
  exact q.block_top j,
 end,
 have top_block : ∀ i0, top_vec.nth (block_vec.nth i0) ≥ i0 :=
 begin
  intro i0,
  dsimp[block_vec,top_vec],
  simp[vector.nth_of_fn,block_map],
  apply fin.le_of_succ_le_succ,
  rw[top_map_spec],
  exact q.top_block i0.succ,
 end,
 have top_mono : ∀ j1 j2, j1 < j2 → top_vec.nth j1 < top_vec.nth j2 :=
 begin
  intros j1 j2 j1_lt_j2,
  dsimp[block_vec,top_vec],
  simp[vector.nth_of_fn],
  exact q.top_mono j1 j2 j1_lt_j2
 end,
 let p : fin_map n r.succ :=
  ⟨block_vec,top_vec,block_top,top_block,top_mono⟩,
 let u : fin r.succ := q.block.nth 0,
 have eq : q = join p u :=
 begin
  ext i,
  cases i with i_val i_is_lt, cases i_val with i0_val,
  { change _ = ((p.join u).block.nth 0 : ℕ),
    rw[join_block_zero p u], refl
  },
  { let i0 : fin n := ⟨i0_val,nat.lt_of_succ_lt_succ i_is_lt⟩,
    have : i0.succ = ⟨i0_val.succ,i_is_lt⟩ := rfl,
    rw[← this,join_block_succ],
    dsimp[p,block_vec],rw[vector.nth_of_fn],dsimp[block_map],refl
  }
 end,
 exact ⟨⟨p,u⟩,eq⟩,
end

lemma res_join (p : fin_map n r.succ) (u : fin r.succ) (e : join_cond (join p u)) :
 (join_res (join p u) e) = ⟨⟨p,u⟩,rfl⟩ :=
begin
 apply subtype.eq,
 ext i; simp[join_res,join],
end

variables (n r)

def empty : (fin_map 0 0) := {
  block := vector.nil,
  top := vector.nil,
  block_top := begin intro i, exact fin.elim0 i, end,
  top_block := begin intro j, exact fin.elim0 j, end,
  top_mono := begin intros j1 j2, exact fin.elim0 j1, end,
}

end fin_map

/- -------------------------------------------------------------- -/
/- -------------------------------------------------------------- -/
/- -------------------------------------------------------------- -/
/- -------------------------------------------------------------- -/

@[derive decidable_eq]
inductive fin_ind_partition : ℕ → ℕ → Type
| empty : (fin_ind_partition 0 0)
| add : ∀ {n r : ℕ}, (fin_ind_partition n r) → (fin_ind_partition n.succ r.succ)
| join : ∀ {n r : ℕ}, (fin_ind_partition n r) → (fin r) → (fin_ind_partition n.succ r)

namespace fin_ind_partition

def {u} elim0 {r : ℕ} {α : fin_ind_partition 0 r.succ → Sort u} :
  Π (x : fin_ind_partition 0 r.succ), α x :=
begin
 intro x, cases x,
end

def {u} elim1 {n : ℕ} {α : fin_ind_partition n.succ 0 → Sort u} :
  Π (x : fin_ind_partition n.succ 0), α x :=
begin
 intro x, cases x,
 have i : fin 0 := by { assumption }, exact i.elim0
end

def tot (n : ℕ) := Σ (r : ℕ), fin_ind_partition n r

def to_tot {n r : ℕ} (p : fin_ind_partition n r) : tot n := ⟨r,p⟩

private def repr_aux : ∀{n r : ℕ}, (fin_ind_partition n r) → string
| 0 0 empty := "O"
| _ _ (add p) := "A" ++ (repr_aux p)
| _ _ (join p u) := "J" ++ (repr u) ++ (repr_aux p)

instance (n r : ℕ) : has_repr (fin_ind_partition n r) :=
 ⟨@repr_aux n r⟩

lemma rank_bound : ∀ {n r : ℕ} (p : fin_ind_partition n r), r ≤ n
| 0 0 empty := le_refl 0
| _ _ (@add n r p) := nat.succ_le_succ (rank_bound p)
| _ _ (@join n r p u) := le_trans (rank_bound p) (nat.le_succ n)

def to_map_partition : ∀ {n r : ℕ} (p : fin_ind_partition n r), fin_map n r
| 0 0 empty := fin_map.empty
| _ _ (add p) := fin_map.add (to_map_partition p)
| _ _ (join p u) := fin_map.join (to_map_partition p) u

def of_map_partition : ∀ {n r : ℕ} (q : fin_map n r), fin_ind_partition n r :=
begin
 intro n,
 induction n with n0 ih_n; intros r q; cases r with r0,
 {exact empty,},
 {have := (q.top.nth 0), exact fin.elim0 this},
 {have := (q.block.nth 0), exact fin.elim0 this},
 {
   by_cases h : q.top.nth 0 = 0,
   {exact add (ih_n (fin_map.add_res q h).val),},
   {
     let pr := (fin_map.join_res q h).val,
     exact join (ih_n pr.1) pr.2,
   }
 }
end

lemma eq_empty (p : fin_ind_partition 0 0) : p = empty :=
 by { cases p, refl}

lemma of_eq_add {n0 r0 : ℕ} {q : fin_map n0.succ r0.succ} (h : q.top.nth 0 = 0) :
 of_map_partition q = add (of_map_partition (fin_map.add_res q h)) :=
begin
 dsimp[of_map_partition], rw[dif_pos h],
end

lemma of_eq_join {n0 r0 : ℕ} {q : fin_map n0.succ r0.succ} (h : q.top.nth 0 ≠ 0) :
 of_map_partition q = join (of_map_partition (fin_map.join_res q h).1.1) (fin_map.join_res q h).1.2 :=
begin
 dsimp[of_map_partition], rw[dif_neg h],
end

lemma of_add {n r : ℕ} (q : fin_map n r) :
 of_map_partition q.add = (of_map_partition q).add :=
begin
 have h0 : q.add.top.nth 0 = 0 := fin_map.add_add_cond q,
 rw[of_eq_add h0,q.res_add h0],refl
end

lemma of_join {n r : ℕ} (q : fin_map n r) (u : fin r) :
 of_map_partition (q.join u) = (of_map_partition q).join u :=
begin
 cases r with r0,
 {exact fin.elim0 u,},
 {
   have h0 : (q.join u).top.nth 0 ≠ 0 := fin_map.join_join_cond q u,
   rw[of_eq_join h0,q.res_join u h0]
 }
end

lemma of_to_map_partition : ∀ {n r : ℕ} (p : fin_ind_partition n r),
 of_map_partition (to_map_partition p) = p
| 0 0 empty := rfl
| _ _ (@add n r p) :=
  (of_add (to_map_partition p)).trans
   (congr_arg add (of_to_map_partition p))
| _ _ (@join n r p u) :=
  (of_join (to_map_partition p) u).trans
   (congr_fun (congr_arg join (of_to_map_partition p)) u)

lemma to_of_map_partition : ∀ {n r : ℕ} (q : fin_map n r),
 to_map_partition (of_map_partition q) = q
| 0 0 q := by {ext i, exact fin.elim0 i,}
| 0 (nat.succ r) q := begin have := q.top.nth 0, exact fin.elim0 this end
| (nat.succ n) 0 q := begin have := q.block.nth 0, exact fin.elim0 this end
| (nat.succ n) (nat.succ r) q :=
  begin
   by_cases h : q.top.nth 0 = 0,
   {
     let q0e0 := fin_map.add_res q h,
     let q0 := q0e0.val,
     let e0 : q = q0.add := q0e0.property,
     exact calc
      to_map_partition (of_map_partition q)
          = to_map_partition (of_map_partition q0.add) : by rw[← e0]
      ... = to_map_partition (of_map_partition q0).add : by rw[of_add q0]
      ... = (to_map_partition (of_map_partition q0)).add : rfl
      ... = q0.add : by rw[to_of_map_partition q0]
      ... = q : by rw[e0],
   },{
     let q0ue0 := fin_map.join_res q h,
     let q0 := q0ue0.val.1,
     let u := q0ue0.val.2,
     let e0 : q = q0.join u := q0ue0.property,
     exact calc
      to_map_partition (of_map_partition q)
          = to_map_partition (of_map_partition (q0.join u)) : by rw[← e0]
      ... = to_map_partition ((of_map_partition q0).join u) : by rw[of_join q0 u]
      ... = (to_map_partition (of_map_partition q0)).join u : rfl
      ... = q0.join u : by rw[to_of_map_partition q0]
      ... = q : by rw[e0],
   }
  end

def map_equiv (n r : ℕ) : fin_ind_partition n r ≃ fin_map n r :=
{
  to_fun := @to_map_partition n r,
  inv_fun := @of_map_partition n r,
  left_inv := @of_to_map_partition n r,
  right_inv := @to_of_map_partition n r
}

private def tot_to_fun {n : ℕ} : fin_ind_partition.tot n → fin_map.tot n
| ⟨r,p⟩ := ⟨r,(map_equiv n r).to_fun p⟩

private def tot_inv_fun {n : ℕ} : fin_map.tot n → fin_ind_partition.tot n
| ⟨r,q⟩ := ⟨r,(map_equiv n r).inv_fun q⟩

private def tot_left_inv {n : ℕ} : ∀ rp : fin_ind_partition.tot n,
  tot_inv_fun (tot_to_fun rp) = rp
| ⟨r,p⟩ := by { simp[tot_inv_fun,tot_to_fun,(map_equiv n r).left_inv p] }

private def tot_right_inv {n : ℕ} : ∀ rq : fin_map.tot n,
  tot_to_fun (tot_inv_fun rq) = rq
| ⟨r,q⟩ := by { simp[tot_inv_fun,tot_to_fun,(map_equiv n r).right_inv q] }

def map_equiv_tot (n : ℕ) : fin_ind_partition.tot n ≃ fin_map.tot n :=
{
  to_fun := @tot_to_fun n,
  inv_fun := @tot_inv_fun n,
  left_inv := @tot_left_inv n,
  right_inv := @tot_right_inv n,
}

instance enum : ∀ (n r : ℕ), enumeration (fin_ind_partition n r)
| 0 0 := {
   elems := [empty],
   nodup := list.nodup_singleton empty,
   complete := λ p, begin
    cases p,exact list.mem_cons_self empty list.nil
   end
  }
| 0 (r + 1) := {
   elems := list.nil,
   nodup := list.nodup_nil,
   complete := λ p, p.elim0
  }
| (n + 1) 0 :=  {
   elems := list.nil,
   nodup := list.nodup_nil,
   complete := λ p, p.elim1
  }
| (n + 1) (r + 1) := begin
   let e0 := enum n r,
   let e1 := enum n (r + 1),
   let L_add := e0.elems.map add,
   let R := fin.elems_list (r + 1),
   let f : fin (r + 1) → list (fin_ind_partition (n + 1) (r + 1)) :=
    λ u, e1.elems.map (λ p, join p u),
   let LL_join := R.map f,
   let L_join := list.join LL_join,
   let elems := list.append L_add L_join,
   let complete : ∀ p : fin_ind_partition (n + 1) (r + 1), p ∈ elems :=
   begin
    intro p,
    cases p with _ _ p0 _ _ p0 u,
    {
     have p0_mem : p0 ∈ e0.elems := (@enumeration.complete _ e0 p0),
     exact list.mem_append_left _ (list.mem_map_of_mem add p0_mem),
    },{
     have p0_mem : p0 ∈ e1.elems := (@enumeration.complete _ e1 p0),
     let h0 := fin.elems_list_complete u,
     let h1 : f u ∈ LL_join := list.mem_map_of_mem f h0,
     let h2 : join p0 u ∈ f u := list.mem_map_of_mem (λ q,join q u) p0_mem,
     let h3 : join p0 u ∈ L_join := list.mem_join_of_mem h1 h2,
     exact list.mem_append_right _ h3,
    }
   end,
   let nodup : list.nodup elems :=
   begin
    have add_inj : function.injective (@add n r) := by simp[function.injective],
    have L_add_nodup : L_add.nodup := list.nodup.map add_inj e0.nodup,
    have LL_nodup : ∀ u, (f u).nodup :=
     λ u, list.nodup.map (by simp[function.injective]) e1.nodup,
    have LL_disjoint_aux : ∀ u v : fin (r + 1), u ≠ v → list.disjoint (f u) (f v) :=
    begin
     intros u v u_ne_v p p_in_f_u p_in_f_v,
     rcases (list.mem_map.mp p_in_f_u) with ⟨pu,⟨pu_in_elems,pu_eq⟩⟩,
     rcases (list.mem_map.mp p_in_f_v) with ⟨pv,⟨pv_in_elems,pv_eq⟩⟩,
     injection (pu_eq.trans pv_eq.symm) with _ u_eq_v,
     exact u_ne_v u_eq_v,
    end,
    have LL_disjoint : list.pairwise list.disjoint LL_join :=
    begin
     apply (@list.pairwise_map _ (fin (r + 1)) list.disjoint f R).mpr,
     exact list.pairwise.imp LL_disjoint_aux (fin.elems_list_nodup (r + 1)),
    end,
    have L_join_nodup : L_join.nodup :=
    begin
     apply list.nodup_join.mpr,
     split,
     {intros l l_in_LL,
      rcases (list.mem_map.mp l_in_LL) with ⟨u,⟨u_in_R,f_u_eq_l⟩⟩,
      rw[← f_u_eq_l],
      exact (LL_nodup u),
     },{
      exact LL_disjoint,
     }
    end,
    have L_disjoint : list.disjoint L_add L_join :=
    begin
     intros p p_in_L_add p_in_L_join,
     rcases (list.mem_map.mp p_in_L_add) with ⟨p0,⟨p0_in_elems,p0_eq⟩⟩,
     rcases list.mem_join.mp p_in_L_join with ⟨l,⟨l_in_LL_join,p_in_l⟩⟩,
     rcases (list.mem_map.mp l_in_LL_join) with ⟨u,⟨u_in_R,f_u_eq_l⟩⟩,
     rw[← f_u_eq_l] at p_in_l,
     rcases (list.mem_map.mp p_in_l) with ⟨p1,⟨p1_in_elems,p1_eq⟩⟩,
     injection (p0_eq.trans p1_eq.symm),
    end,
    exact list.nodup_append.mpr ⟨L_add_nodup,L_join_nodup,L_disjoint⟩,
   end,
   exact {
     elems := elems, nodup := nodup, complete := complete
   },
  end

lemma card_step (n r : ℕ) :
 (fin_ind_partition.enum n.succ r.succ).elems.length =
   (fin_ind_partition.enum n r).elems.length +
   (fin_ind_partition.enum n r.succ).elems.length * r.succ :=
begin
 let e0 := fin_ind_partition.enum n r,
 let e1 := fin_ind_partition.enum n (r + 1),
 let e2 := fin_ind_partition.enum (n + 1) (r + 1),
 let n0 := e0.elems.length,
 let n1 := e1.elems.length,
 let n2 := e2.elems.length,
 let L_add := e0.elems.map add,
 let R := fin.elems_list (r + 1),
 let f : fin (r + 1) → list (fin_ind_partition (n + 1) (r + 1)) :=
    λ u, e1.elems.map (λ p, join p u),
 let LL_join := R.map f,
 let L_join := list.join LL_join,
 have h0 : n2 = L_add.length + L_join.length := list.length_append L_add L_join,
 have h1 : L_add.length = n0 := list.length_map add _,
 let h2 : LL_join.length = r.succ :=
  (list.length_map f R).trans (fin.elems_list_length r.succ),
 have h3 : ∀ u, (f u).length = n1 :=
  λ u,list.length_map (λ p, join p u) _,
 let h4 : LL_join.map list.length = list.repeat n1 r.succ :=
 begin
  apply list.eq_repeat.mpr,
  split,
  {exact ((list.length_map list.length LL_join).trans h2),},
  {intros k k_in_ks,
   rcases (list.mem_map.mp k_in_ks) with ⟨l,⟨l_in_LL_join,l_len⟩⟩,
   rcases (list.mem_map.mp l_in_LL_join) with ⟨u,⟨u_in_R,f_u_eq_l⟩⟩,
   rw[← l_len,← f_u_eq_l],
   exact h3 u,
  }
 end,
 have h5 : L_join.length = n1 * r.succ := (list.length_join LL_join).trans
   ((congr_arg list.sum h4).trans (list.sum_const_nat n1 r.succ)),
 exact calc
  n2 = L_add.length + L_join.length : h0
  ... = n0 + n1 * r.succ : by rw[h1,h5]
end

lemma card_stirling : ∀ (n r : ℕ),
 (fintype.card (fin_ind_partition n r)) = partition.number.stirling2 n r
| 0 0 := rfl
| 0 (r + 1) := rfl
| (n + 1) 0 := rfl
| (n + 1) (r + 1) :=
  begin
   rw[partition.number.stirling2,enumeration.enum_card,card_step n r],
   rw[← enumeration.enum_card,← enumeration.enum_card],
   rw[card_stirling n r,card_stirling n (r + 1)],
  end

instance tot_enum (n : ℕ) : enumeration (tot n) :=
begin
 let e0 := fin_ind_partition.enum n,
 let elems0 : ℕ → list (tot n) := λ r, (e0 r).elems.map to_tot,
 let elems1 := ((list.range n.succ).map elems0),
 let elems := elems1.join,
 let nodup0 : ∀ r, (elems0 r).nodup :=
 begin
  intro r,
  have tot_inj : function.injective (@to_tot n r) :=
  begin
   intros p1 p2 tot_eq,
   injection tot_eq with r_eq p_heq,
   cases r_eq,
   exact eq_of_heq p_heq,
  end,
  exact list.nodup.map tot_inj (e0 r).nodup,
 end,
 let nodup1 : ∀ l, l ∈ elems1 → (list.nodup l) :=
 begin
  intros l l_in_elems1,
  rcases (list.mem_map.mp l_in_elems1) with ⟨r,⟨r_mem,l_eq⟩⟩,
  rw[← l_eq],
  exact nodup0 r,
 end,
 let disjoint0 : ∀ r1 r2, r1 ≠ r2 → list.disjoint (elems0 r1) (elems0 r2) :=
 begin
  intros r1 r2 r_ne q q_in_l1 q_in_l2,
  rcases (list.mem_map.mp q_in_l1) with ⟨p1,⟨p1_in_elems,p1_eq⟩⟩,
  rcases (list.mem_map.mp q_in_l2) with ⟨p2,⟨p2_in_elems,p2_eq⟩⟩,
  let r_eq : r1 = r2 := (congr_arg sigma.fst p1_eq).trans
                        (congr_arg sigma.fst p2_eq).symm,
  exact r_ne r_eq,
 end,
 let disjoint1 := @list.pairwise.imp ℕ ne
  (λ r1 r2, list.disjoint (elems0 r1) (elems0 r2))
   _ disjoint0 (list.nodup_range n.succ),
 let disjoint2 := (list.pairwise_map elems0).mpr disjoint1,
 let nodup := list.nodup_join.mpr ⟨nodup1,disjoint2⟩,
 let complete : ∀ rp : tot n, rp ∈ elems :=
 begin
  intro rp,
  rcases rp with ⟨r,p⟩,
  have rp_eq : to_tot p = ⟨r,p⟩ := rfl,
  rw[← rp_eq],
  have mem0 : p ∈ (e0 r).elems := @enumeration.complete _ (e0 r) p,
  have mem1 : to_tot p ∈ elems0 r :=
   list.mem_map_of_mem to_tot mem0,
  have mem2 : elems0 r ∈ elems1 := list.mem_map_of_mem elems0
   (list.mem_range.mpr (nat.lt_succ_of_le (rank_bound p))),
  exact list.mem_join.mpr ⟨elems0 r,⟨mem2,mem1⟩⟩,
 end,
 exact {elems := elems, nodup := nodup, complete := complete},
end

lemma card_bell (n : ℕ) :
 fintype.card (tot n) = partition.number.bell n :=
begin
 let f := λ (r : ℕ), list.map to_tot (enumeration.elems : list (fin_ind_partition n r)),
 have h : list.length ∘ f = partition.number.stirling2 n :=
 begin
  ext r,simp,rw[← enumeration.enum_card],exact card_stirling n r,
 end,
 rw[enumeration.enum_card,partition.number.bell],
 dsimp[enumeration.elems],
 rw[list.length_join,list.map_map,h],
end

end fin_ind_partition

namespace fin_map

instance enum (n r : ℕ) : enumeration (fin_map n r) :=
 enumeration.of_equiv (fin_ind_partition n r)
  (fin_ind_partition.map_equiv n r)

lemma card_stirling (n r : ℕ) :
 (fintype.card (fin_map n r)) = partition.number.stirling2 n r :=
  (fintype.card_congr (fin_ind_partition.map_equiv n r)).symm.trans
           (fin_ind_partition.card_stirling n r)

instance enum_tot (n : ℕ) : enumeration (fin_map.tot n) :=
 enumeration.of_equiv (fin_ind_partition.tot n)
  (fin_ind_partition.map_equiv_tot n)

lemma card_bell (n : ℕ) :
 (fintype.card (fin_map.tot n)) = partition.number.bell n :=
  (fintype.card_congr (fin_ind_partition.map_equiv_tot n)).symm.trans
           (fin_ind_partition.card_bell n)

end fin_map

lemma card_stirling (n r : ℕ) :
 fintype.card { p : partition (fin n) // p.rank = r } = partition.number.stirling2 n r :=
  (fintype.card_congr (@fin_map.partition_equiv n r)).symm.trans
           (fin_map.card_stirling n r)

lemma card_bell (n : ℕ) :
 fintype.card (partition (fin n)) = partition.number.bell n :=
  (fintype.card_congr (@fin_map.partition_equiv_tot n)).symm.trans
           (fin_map.card_bell n)

lemma card_bell' (α : Type*) [fintype α] [decidable_eq α] :
 fintype.card (partition α) = partition.number.bell (fintype.card α) :=
begin
 refine trunc.induction_on (fintype.trunc_equiv_fin α) _,
 intro equiv_fin,
 exact (fintype.card_congr (isofunctor equiv_fin)).trans
            (card_bell (fintype.card α)),
end

end combinatorics.partition


import algebra.big_operators.basic
import data.finset.powerset
import data.finset.sups

/-!
# The four functions theorem and FKG inequality

-/

open_locale big_operators finset_family

variables {α β : Type*}

section lattice
variables [lattice α] [has_mul β] [has_le β]

/-- Four functions `f₁`, `f₂`, `f₃`, `f₄` are log-super modular if -/
def log_super_modular (f₁ f₂ f₃ f₄ : α → β) : Prop := ∀ a b, f₁ a * f₂ b ≤ f₃ (a ⊔ b) * f₄ (a ⊓ b)

end lattice

protected lemma linear_order.log_super_modular [linear_order α] [comm_semigroup β] [preorder β]
  (f : α → β) :
  log_super_modular f f f f :=
λ a b, by cases le_total a b; simp [h, mul_comm]

section finset
variables [decidable_eq α] [ordered_semiring β] {f f₁ f₂ f₃ f₄ g μ : finset α → β}

/-- The **four functions theorem** -/
lemma four_functions_theorem' {u : finset α}  (h₁ : ∀ s ⊆ u, 0 ≤ f₁ s) (h₂ : ∀ s ⊆ u, 0 ≤ f₂ s)
  (h₃ : ∀ s ⊆ u, 0 ≤ f₃ s) (h₄ : ∀ s ⊆ u, 0 ≤ f₄ s)
  (h : ∀ a b, f₁ a * f₂ b ≤ f₃ (a ⊔ b) * f₄ (a ⊓ b)) {s t : finset (finset α)}
  (hs : s ⊆ u.powerset) (ht : t ⊆ u.powerset) :
  (∑ a in s, f₁ a) * (∑ b in t, f₂ b) ≤ (∑ a in s ⊻ t, f₃ a) * (∑ b in s ⊼ t, f₄ b) :=
begin
  induction u using finset.induction with a u h ih generalizing f₁ f₂ f₃ f₄,
  { simp only [finset.powerset_empty, finset.subset_singleton_iff] at hs ht,
    obtain rfl | rfl := hs; obtain rfl | rfl := ht; simpa using h ∅ ∅ <|> simp },
  sorry

end

end finset

variables [distrib_lattice α] [decidable_eq α] [ordered_semiring β]
  {f f₁ f₂ f₃ f₄ g μ : α → β}

/-- The **four functions theorem** -/
lemma four_functions_theorem {s t u : finset α} (h₁ : ∀ a ∈ u, 0 ≤ f₁ a) (h₂ : ∀ a ∈ u,  0 ≤ f₂)
  (h₃ : ∀ a ∈ u, 0 ≤ f₃) (h₄ : ∀ a ∈ u, 0 ≤ f₄)
  (h : ∀ a b, f₁ a * f₂ b ≤ f₃ (a ⊔ b) * f₄ (a ⊓ b)) (hs : s ⊆ u) (ht : t ⊆ u) :
  (∑ a in s, f₁ a) * (∑ b in t, f₂ b) ≤ (∑ a in s ∪ t, f₃ a) * (∑ b in s ∩ t, f₄ b) :=
begin
  sorry,
end

lemma fkg [fintype α] (hμ : 0 ≤ μ) (hf : monotone f) (hg : monotone g) :
  (∑ a, f a * μ a) * (∑ a, g a * μ a) ≤ (∑ a, f a * g a * μ a) * (∑ a, μ a) :=
begin
  sorry
end


/-
Copyright (c) 2020 Aaron Anderson. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Aaron Anderson, Jalex Stark, Kyle Miller
-/
import combinatorics.simple_graph.adj_matrix
import linear_algebra.matrix.charpoly.coeff
import data.int.modeq
import data.zmod.basic
import tactic.interval_cases

/-!
# The Friendship Theorem

## Definitions and Statement
- A `friendship` graph is one in which any two distinct vertices have exactly one neighbor in common
- A `politician`, at least in the context of this problem, is a vertex in a graph which is adjacent
  to every other vertex.
- The friendship theorem (Erdős, Rényi, Sós 1966) states that every finite friendship graph has a
  politician.

## Proof outline
The proof revolves around the theory of adjacency matrices, although some steps could equivalently
be phrased in terms of counting walks.
- Assume `G` is a finite friendship graph.
- First we show that any two nonadjacent vertices have the same degree
- Assume for contradiction that `G` does not have a politician.
- Conclude from the last two points that `G` is `d`-regular for some `d : ℕ`.
- Show that `G` has `d ^ 2 - d + 1` vertices.
- By casework, show that if `d = 0, 1, 2`, then `G` has a politician.
- If `3 ≤ d`, let `p` be a prime factor of `d - 1`.
- If `A` is the adjacency matrix of `G` with entries in `ℤ/pℤ`, we show that `A ^ p` has trace `1`.
- This gives a contradiction, as `A` has trace `0`, and thus `A ^ p` has trace `0`.

## References
- [P. Erdős, A. Rényi, V. Sós, *On A Problem of Graph Theory*][erdosrenyisos]
- [C. Huneke, *The Friendship Theorem*][huneke2002]

-/

open_locale classical big_operators
noncomputable theory

open finset simple_graph matrix

universes u v
variables {V : Type u} {R : Type v} [semiring R]

section friendship_def
variables (G : simple_graph V)

/--
This property of a graph is the hypothesis of the friendship theorem:
every pair of nonadjacent vertices has exactly one common friend,
a vertex to which both are adjacent.
-/
def friendship [fintype V] : Prop := ∀ ⦃v w : V⦄, v ≠ w → fintype.card (G.common_neighbors v w) = 1

/--
A politician is a vertex that is adjacent to all other vertices.
-/
def exists_politician : Prop := ∃ (v : V), ∀ (w : V), v ≠ w → G.adj v w

end friendship_def

variables [fintype V] {G : simple_graph V} {d : ℕ} (hG : friendship G)
include hG

namespace friendship

variables (R)
/-- One characterization of a friendship graph is that there is exactly one walk of length 2
  between distinct vertices. These walks are counted in off-diagonal entries of the square of
  the adjacency matrix, so for a friendship graph, those entries are all 1. -/
theorem adj_matrix_sq_of_ne {v w : V} (hvw : v ≠ w) :
  ((G.adj_matrix R) ^ 2) v w = 1 :=
begin
  rw [sq, ← nat.cast_one, ← hG hvw],
  simp [common_neighbors, neighbor_finset_eq_filter, finset.filter_filter, finset.filter_inter,
    and_comm],
end

/-- This calculation amounts to counting the number of length 3 walks between nonadjacent vertices.
  We use it to show that nonadjacent vertices have equal degrees. -/
lemma adj_matrix_pow_three_of_not_adj {v w : V} (non_adj : ¬ G.adj v w) :
  ((G.adj_matrix R) ^ 3) v w = degree G v :=
begin
  rw [pow_succ, mul_eq_mul, adj_matrix_mul_apply, degree, card_eq_sum_ones, nat.cast_sum],
  apply sum_congr rfl,
  intros x hx,
  rw [adj_matrix_sq_of_ne _ hG, nat.cast_one],
  rintro ⟨rfl⟩,
  rw mem_neighbor_finset at hx,
  exact non_adj hx,
end

variable {R}

/-- As `v` and `w` not being adjacent implies
  `degree G v = ((G.adj_matrix R) ^ 3) v w` and `degree G w = ((G.adj_matrix R) ^ 3) v w`,
  the degrees are equal if `((G.adj_matrix R) ^ 3) v w = ((G.adj_matrix R) ^ 3) w v`

  This is true as the adjacency matrix is symmetric. -/
lemma degree_eq_of_not_adj {v w : V} (hvw : ¬ G.adj v w) :
  degree G v = degree G w :=
begin
  rw [← nat.cast_id (G.degree v), ← nat.cast_id (G.degree w),
      ← adj_matrix_pow_three_of_not_adj ℕ hG hvw,
      ← adj_matrix_pow_three_of_not_adj ℕ hG (λ h, hvw (G.symm h))],
  conv_lhs {rw ← transpose_adj_matrix},
  simp only [pow_succ, sq, mul_eq_mul, ← transpose_mul, transpose_apply],
  simp only [← mul_eq_mul, mul_assoc],
end

/-- Let `A` be the adjacency matrix of a graph `G`.
  If `G` is a friendship graph, then all of the off-diagonal entries of `A^2` are 1.
  If `G` is `d`-regular, then all of the diagonal entries of `A^2` are `d`.
  Putting these together determines `A^2` exactly for a `d`-regular friendship graph. -/
theorem adj_matrix_sq_of_regular (hd : G.is_regular_of_degree d) :
  ((G.adj_matrix R) ^ 2) = λ v w, if v = w then d else 1 :=
begin
  ext v w, by_cases h : v = w,
  { rw [h, sq, mul_eq_mul, adj_matrix_mul_self_apply_self, hd], simp, },
  { rw [adj_matrix_sq_of_ne R hG h, if_neg h], },
end

lemma adj_matrix_sq_mod_p_of_regular {p : ℕ} (dmod : (d : zmod p) = 1)
  (hd : G.is_regular_of_degree d) :
  (G.adj_matrix (zmod p)) ^ 2 = λ _ _, 1 :=
by simp [adj_matrix_sq_of_regular hG hd, dmod]

section nonempty

variable [nonempty V]

/-- If `G` is a friendship graph without a politician (a vertex adjacent to all others), then
  it is regular. We have shown that nonadjacent vertices of a friendship graph have the same degree,
  and if there isn't a politician, we can show this for adjacent vertices by finding a vertex
  neither is adjacent to, and then using transitivity. -/
theorem is_regular_of_not_exists_politician (hG' : ¬exists_politician G) :
  ∃ (d : ℕ), G.is_regular_of_degree d :=
begin
  have v := classical.arbitrary V,
  use G.degree v,
  intro x,
  by_cases hvx : G.adj v x, swap, { exact (degree_eq_of_not_adj hG hvx).symm, },
  dunfold exists_politician at hG',
  push_neg at hG',
  rcases hG' v with ⟨w, hvw', hvw⟩,
  rcases hG' x with ⟨y, hxy', hxy⟩,
  by_cases hxw : G.adj x w,
    swap, { rw degree_eq_of_not_adj hG hvw, exact degree_eq_of_not_adj hG hxw },
  rw degree_eq_of_not_adj hG hxy,
  by_cases hvy : G.adj v y,
    swap, { exact (degree_eq_of_not_adj hG hvy).symm },
  rw degree_eq_of_not_adj hG hvw,
  apply degree_eq_of_not_adj hG,
  intro hcontra,
  rcases finset.card_eq_one.mp (hG hvw') with ⟨⟨a, ha⟩, h⟩,
  have key : ∀ {x}, x ∈ G.common_neighbors v w → x = a,
  { intros x hx,
    have h' := mem_univ (subtype.mk x hx),
    rw [h, mem_singleton] at h',
    injection h', },
  apply hxy',
  rw [key ((mem_common_neighbors G).mpr ⟨hvx, G.symm hxw⟩),
      key ((mem_common_neighbors G).mpr ⟨hvy, G.symm hcontra⟩)],
end

/-- Let `A` be the adjacency matrix of a `d`-regular friendship graph, and let `v` be a vector
  all of whose components are `1`. Then `v` is an eigenvector of `A ^ 2`, and we can compute
  the eigenvalue to be `d * d`, or as `d + (fintype.card V - 1)`, so those quantities must be equal.

  This essentially means that the graph has `d ^ 2 - d + 1` vertices. -/
lemma card_of_regular (hd : G.is_regular_of_degree d) :
  d + (fintype.card V - 1) = d * d :=
begin
  have v := classical.arbitrary V,
  transitivity ((G.adj_matrix ℕ) ^ 2).mul_vec (λ _, 1) v,
  { rw [adj_matrix_sq_of_regular hG hd, mul_vec, dot_product, ← insert_erase (mem_univ v)],
    simp only [sum_insert, mul_one, if_true, nat.cast_id, eq_self_iff_true,
               mem_erase, not_true, ne.def, not_false_iff, add_right_inj, false_and],
    rw [finset.sum_const_nat, card_erase_of_mem (mem_univ v), mul_one], { refl },
    intros x hx, simp [(ne_of_mem_erase hx).symm], },
  { rw [sq, mul_eq_mul, ← mul_vec_mul_vec],
    simp [adj_matrix_mul_vec_const_apply_of_regular hd, neighbor_finset,
          card_neighbor_set_eq_degree, hd v], }
end

/-- The size of a `d`-regular friendship graph is `1 mod (d-1)`, and thus `1 mod p` for a
  factor `p ∣ d-1`. -/
lemma card_mod_p_of_regular {p : ℕ} (dmod : (d : zmod p) = 1) (hd : G.is_regular_of_degree d) :
  (fintype.card V : zmod p) = 1 :=
begin
  have hpos : 0 < fintype.card V := fintype.card_pos_iff.mpr infer_instance,
  rw [← nat.succ_pred_eq_of_pos hpos, nat.succ_eq_add_one, nat.pred_eq_sub_one],
  simp only [add_left_eq_self, nat.cast_add, nat.cast_one],
  have h := congr_arg (λ n, (↑n : zmod p)) (card_of_regular hG hd),
  revert h, simp [dmod],
end

end nonempty

omit hG

lemma adj_matrix_sq_mul_const_one_of_regular (hd : G.is_regular_of_degree d) :
  (G.adj_matrix R) * (λ _ _, 1) = λ _ _, d :=
by { ext x, simp [← hd x, degree] }

lemma adj_matrix_mul_const_one_mod_p_of_regular {p : ℕ} (dmod : (d : zmod p) = 1)
  (hd : G.is_regular_of_degree d) :
  (G.adj_matrix (zmod p)) * (λ _ _, 1) = λ _ _, 1 :=
by rw [adj_matrix_sq_mul_const_one_of_regular hd, dmod]

include hG

/-- Modulo a factor of `d-1`, the square and all higher powers of the adjacency matrix
  of a `d`-regular friendship graph reduce to the matrix whose entries are all 1. -/
lemma adj_matrix_pow_mod_p_of_regular {p : ℕ} (dmod : (d : zmod p) = 1)
  (hd : G.is_regular_of_degree d) {k : ℕ} (hk : 2 ≤ k) :
  (G.adj_matrix (zmod p)) ^ k = λ _ _, 1 :=
begin
  iterate 2 {cases k with k, { exfalso, linarith, }, },
  induction k with k hind,
  { exact adj_matrix_sq_mod_p_of_regular hG dmod hd, },
  rw [pow_succ, hind (nat.le_add_left 2 k)],
  exact adj_matrix_mul_const_one_mod_p_of_regular dmod hd,
end

variable [nonempty V]

/-- This is the main proof. Assuming that `3 ≤ d`, we take `p` to be a prime factor of `d-1`.
  Then the `p`th power of the adjacency matrix of a `d`-regular friendship graph must have trace 1
  mod `p`, but we can also show that the trace must be the `p`th power of the trace of the original
  adjacency matrix, which is 0, a contradiction.
-/
lemma false_of_three_le_degree (hd : G.is_regular_of_degree d) (h : 3 ≤ d) : false :=
begin
  -- get a prime factor of d - 1
  let p : ℕ := (d - 1).min_fac,
  have p_dvd_d_pred := (zmod.nat_coe_zmod_eq_zero_iff_dvd _ _).mpr (d - 1).min_fac_dvd,
  have dpos : 0 < d := by linarith,
  have d_cast : ↑(d - 1) = (d : ℤ) - 1 := by norm_cast,
  haveI : fact p.prime := ⟨nat.min_fac_prime (by linarith)⟩,
  have hp2 : 2 ≤ p := (fact.out p.prime).two_le,
  have dmod : (d : zmod p) = 1,
  { rw [← nat.succ_pred_eq_of_pos dpos, nat.succ_eq_add_one, nat.pred_eq_sub_one],
    simp only [add_left_eq_self, nat.cast_add, nat.cast_one],
    exact p_dvd_d_pred, },
  have Vmod := card_mod_p_of_regular hG dmod hd,
  -- now we reduce to a trace calculation
  have := zmod.trace_pow_card (G.adj_matrix (zmod p)),
  contrapose! this, clear this,
  -- the trace is 0 mod p when computed one way
  rw [trace_adj_matrix, zero_pow (fact.out p.prime).pos],
  -- but the trace is 1 mod p when computed the other way
  rw adj_matrix_pow_mod_p_of_regular hG dmod hd hp2,
  dunfold fintype.card at Vmod,
  simp only [matrix.trace, diag_apply, mul_one, nsmul_eq_mul, linear_map.coe_mk, sum_const],
  rw [Vmod, ← nat.cast_one, zmod.nat_coe_zmod_eq_zero_iff_dvd, nat.dvd_one,
    nat.min_fac_eq_one_iff],
  linarith,
end

/-- If `d ≤ 1`, a `d`-regular friendship graph has at most one vertex, which is
  trivially a politician. -/
lemma exists_politician_of_degree_le_one (hd : G.is_regular_of_degree d) (hd1 : d ≤ 1) :
  exists_politician G :=
begin
  have sq : d * d = d := by { interval_cases d; norm_num },
  have h := card_of_regular hG hd,
  rw sq at h,
  have : fintype.card V ≤ 1,
  { cases fintype.card V with n,
    { exact zero_le _, },
    { have : n = 0,
      { rw [nat.succ_sub_succ_eq_sub, tsub_zero] at h,
        linarith },
      subst n, } },
  use classical.arbitrary V,
  intros w h, exfalso,
  apply h,
  apply fintype.card_le_one_iff.mp this,
end

/-- If `d = 2`, a `d`-regular friendship graph has 3 vertices, so it must be complete graph,
  and all the vertices are politicians. -/
lemma neighbor_finset_eq_of_degree_eq_two (hd : G.is_regular_of_degree 2) (v : V) :
  G.neighbor_finset v = finset.univ.erase v :=
begin
  apply finset.eq_of_subset_of_card_le,
  { rw finset.subset_iff,
    intro x,
    rw [mem_neighbor_finset, finset.mem_erase],
    exact λ h, ⟨(G.ne_of_adj h).symm, finset.mem_univ _⟩ },
  convert_to 2 ≤ _,
  { convert_to _ = fintype.card V - 1,
    { have hfr:= card_of_regular hG hd,
      linarith },
    { exact finset.card_erase_of_mem (finset.mem_univ _), }, },
  { dsimp [is_regular_of_degree, degree] at hd,
    rw hd, }
end

lemma exists_politician_of_degree_eq_two (hd : G.is_regular_of_degree 2) :
  exists_politician G :=
begin
  have v := classical.arbitrary V,
  use v,
  intros w hvw,
  rw [← mem_neighbor_finset, neighbor_finset_eq_of_degree_eq_two hG hd v, finset.mem_erase],
  exact ⟨hvw.symm, finset.mem_univ _⟩,
end

lemma exists_politician_of_degree_le_two (hd : G.is_regular_of_degree d) (h : d ≤ 2) :
  exists_politician G :=
begin
  interval_cases d,
  iterate 2 { apply exists_politician_of_degree_le_one hG hd, norm_num },
  { exact exists_politician_of_degree_eq_two hG hd },
end

end friendship

/-- **Friendship theorem**: We wish to show that a friendship graph has a politician (a vertex
  adjacent to all others). We proceed by contradiction, and assume the graph has no politician.
  We have already proven that a friendship graph with no politician is `d`-regular for some `d`,
  and now we do casework on `d`.
  If the degree is at most 2, we observe by casework that it has a politician anyway.
  If the degree is at least 3, the graph cannot exist. -/
theorem friendship_theorem [nonempty V] : exists_politician G :=
begin
  by_contradiction npG,
  rcases hG.is_regular_of_not_exists_politician npG with ⟨d, dreg⟩,
  cases lt_or_le d 3 with dle2 dge3,
  { exact npG (hG.exists_politician_of_degree_le_two dreg (nat.lt_succ_iff.mp dle2)) },
  { exact hG.false_of_three_le_degree dreg dge3 },
end