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

Extended non-negative reals
-/
import analysis.nnreal data.real.ennreal
noncomputable theory
open classical set lattice filter
local attribute [instance] prop_decidable
variables {α : Type*} {β : Type*} {γ : Type*}

local notation `∞` := ennreal.infinity

namespace ennreal
variables {a b c d : ennreal} {r p q : nnreal}

section topological_space
open topological_space

instance : topological_space ennreal :=
topological_space.generate_from {s | ∃a, s = {b | a < b} ∨ s = {b | b < a}}

instance : orderable_topology ennreal := ⟨rfl⟩

instance : t2_space ennreal := by apply_instance

instance : second_countable_topology ennreal :=
⟨⟨⋃q ≥ (0:ℚ), {{a : ennreal | a < nnreal.of_real q}, {a : ennreal | ↑(nnreal.of_real q) < a}},
  countable_bUnion (countable_encodable _) $ assume a ha, countable_insert (countable_singleton _),
  le_antisymm
    (generate_from_le $ λ s h, begin
      rcases h with ⟨a, hs | hs⟩;
      [ rw show s = ⋃q∈{q:ℚ | 0 ≤ q ∧ a < nnreal.of_real q}, {b | ↑(nnreal.of_real q) < b},
           from set.ext (assume b, by simp [hs, @ennreal.lt_iff_exists_rat_btwn a b, and_assoc]),
        rw show s = ⋃q∈{q:ℚ | 0 ≤ q ∧ ↑(nnreal.of_real q) < a}, {b | b < ↑(nnreal.of_real q)},
           from set.ext (assume b, by simp [hs, @ennreal.lt_iff_exists_rat_btwn b a, and_comm, and_assoc])];
      { apply is_open_Union, intro q,
        apply is_open_Union, intro hq,
        exact generate_open.basic _ (mem_bUnion hq.1 $ by simp) }
    end)
    (generate_from_le $ by simp [or_imp_distrib, is_open_lt', is_open_gt'] {contextual := tt})⟩⟩

lemma embedding_coe : embedding (coe : nnreal → ennreal) :=
and.intro (assume a b, coe_eq_coe.1) $
begin
  refine le_antisymm _ _,
  { rw [orderable_topology.topology_eq_generate_intervals nnreal],
    refine generate_from_le (assume s ha, _),
    rcases ha with ⟨a, rfl | rfl⟩,
    exact ⟨{b : ennreal | ↑a < b}, @is_open_lt' ennreal ennreal.topological_space _ _ _, by simp⟩,
    exact ⟨{b : ennreal | b < ↑a}, @is_open_gt' ennreal ennreal.topological_space _ _ _, by simp⟩, },
  { rw [orderable_topology.topology_eq_generate_intervals ennreal,
      induced_le_iff_le_coinduced],
    refine generate_from_le (assume s ha, _),
    rcases ha with ⟨a, rfl | rfl⟩,
    show is_open {b : nnreal | a < ↑b},
    { cases a; simp [none_eq_top, some_eq_coe, is_open_lt'] },
    show is_open {b : nnreal | ↑b < a},
    { cases a; simp [none_eq_top, some_eq_coe, is_open_gt', is_open_const] } }
end

lemma is_open_ne_top : is_open {a : ennreal | a ≠ ⊤} :=
is_open_neg (is_closed_eq continuous_id continuous_const)

lemma coe_image_univ_mem_nhds : (coe : nnreal → ennreal) '' univ ∈ (nhds (r : ennreal)).sets :=
have {a : ennreal | a ≠ ⊤} = (coe : nnreal → ennreal) '' univ,
  from set.ext $ assume a, by cases a; simp [none_eq_top, some_eq_coe],
this ▸ mem_nhds_sets is_open_ne_top coe_ne_top

lemma tendsto_coe {f : filter α} {m : α → nnreal} {a : nnreal} :
  tendsto (λa, (m a : ennreal)) f (nhds ↑a) ↔ tendsto m f (nhds a) :=
embedding_coe.tendsto_nhds_iff.symm

lemma continuous_coe {α} [topological_space α] {f : α → nnreal} :
continuous (λa, (f a : ennreal)) ↔ continuous f :=
embedding_coe.continuous_iff.symm

lemma nhds_coe {r : nnreal} : nhds (r : ennreal) = (nhds r).map coe :=
by rw [embedding_coe.2, map_nhds_induced_eq coe_image_univ_mem_nhds]

lemma nhds_coe_coe {r p : nnreal} : nhds ((r : ennreal), (p : ennreal)) =
  (nhds (r, p)).map (λp:nnreal×nnreal, (p.1, p.2)) :=
begin
  rw [(embedding_prod_mk embedding_coe embedding_coe).map_nhds_eq],
  rw [← univ_prod_univ, ← prod_image_image_eq],
  exact prod_mem_nhds_sets coe_image_univ_mem_nhds coe_image_univ_mem_nhds
end

lemma tendsto_to_nnreal {a : ennreal} : a ≠ ⊤ →
  tendsto (ennreal.to_nnreal) (nhds a) (nhds a.to_nnreal) :=
begin
  cases a; simp [some_eq_coe, none_eq_top, nhds_coe, tendsto_map'_iff, (∘)],
  exact tendsto_id
end

lemma coe_nat : ∀n:ℕ, (n : ennreal) = (n : nnreal)
| 0       := rfl
| (n + 1) := by simp [coe_nat n]

lemma tendsto_nhds_top {m : α → ennreal} {f : filter α}
  (h : ∀n:ℕ, {a | ↑n < m a} ∈ f.sets) : tendsto m f (nhds ⊤) :=
tendsto_nhds_generate_from $ assume s hs,
match s, hs with
| _, ⟨none,   or.inl rfl⟩, hr := (lt_irrefl ⊤ hr).elim
| _, ⟨some r, or.inl rfl⟩, hr :=
  let ⟨n, hrn⟩ := exists_nat_gt r in
  mem_sets_of_superset (h n) $ assume a hnma, show ↑r < m a, from
    lt_trans (show (r : ennreal) < n, from (coe_nat n).symm ▸ coe_lt_coe.2 hrn) hnma
| _, ⟨a,      or.inr rfl⟩, hr := (not_top_lt $ show ⊤ < a, from hr).elim
end

/-
lemma nhds_top : nhds (⊤ : ennreal) = (⨅i:ℕ, principal {e : ennreal | ↑i < e}) :=
_
-/

/- MOVE TO nnreal
lemma nhds_of_real_eq_map_of_real_nhds_nonneg {r : ℝ} :
  nhds (of_real r) = (nhds r ⊓ principal {x | 0 ≤ x}).map of_real :=
by rw [nhds_of_real_eq_map_of_real_nhds hr];
from by_cases
  (assume : r = 0,
    le_antisymm
      (assume s (hs : {a | of_real a ∈ s} ∈ (nhds r ⊓ principal {x | 0 ≤ x}).sets),
        let ⟨t₁, ht₁, t₂, ht₂, ht⟩ := mem_inf_sets.mp hs in
        show {a | of_real a ∈ s} ∈ (nhds r).sets,
          from (nhds r).sets_of_superset ht₁ $ assume a ha,
          match le_total 0 a with
          | or.inl h := have a ∈ t₂, from ht₂ h, ht ⟨ha, this⟩
          | or.inr h :=
            have r ∈ t₁ ∩ t₂, from ⟨mem_of_nhds ht₁, ht₂ (le_of_eq ‹r = 0›.symm)⟩,
            have of_real 0 ∈ s, from ‹r = 0› ▸ ht this,
            by simp [of_real_of_nonpos h]; assumption
          end)
      (map_mono inf_le_left))
  (assume : r ≠ 0,
    have 0 < r, from lt_of_le_of_ne hr this.symm,
    have nhds r ⊓ principal {x : ℝ | 0 ≤ x} = nhds r,
      from inf_of_le_left $ le_principal_iff.mpr $ le_mem_nhds this,
    by simp [*])
-/

instance : topological_add_monoid ennreal :=
⟨ continuous_iff_tendsto.2 $
  have hl : ∀a:ennreal, tendsto (λ (p : ennreal × ennreal), p.fst + p.snd) (nhds (⊤, a)) (nhds ⊤), from
    assume a, tendsto_nhds_top $ assume n,
    have set.prod {a | ↑n < a } univ ∈ (nhds ((⊤:ennreal), a)).sets, from
      prod_mem_nhds_sets (lt_mem_nhds $ (coe_nat n).symm ▸ coe_lt_top) univ_mem_sets,
    begin filter_upwards [this] assume ⟨a₁, a₂⟩ ⟨h₁, h₂⟩, lt_of_lt_of_le h₁ (le_add_right $ le_refl _) end,
  begin
    rintro ⟨a₁, a₂⟩,
    cases a₁, { simp [none_eq_top, hl a₂], },
    cases a₂, { simp [none_eq_top, some_eq_coe, nhds_swap (a₁ : ennreal) ⊤, tendsto_map'_iff, (∘), hl ↑a₁] },
    simp [some_eq_coe, nhds_coe_coe, tendsto_map'_iff, (∘)],
    simp only [coe_add.symm, tendsto_coe, tendsto_add']
  end ⟩

protected lemma tendsto_mul' (ha : a ≠ 0) (hb : b ≠ 0) :
  tendsto (λp:ennreal×ennreal, p.1 * p.2) (nhds (a, b)) (nhds (a * b)) :=
have ht : ∀b:ennreal, b ≠ 0 → tendsto (λp:ennreal×ennreal, p.1 * p.2) (nhds ((⊤:ennreal), b)) (nhds ⊤),
begin
  refine assume b hb, tendsto_nhds_top $ assume n, _,
  rcases dense (zero_lt_iff_ne_zero.2 hb) with ⟨ε', hε', hεb'⟩,
  rcases ennreal.lt_iff_exists_coe.1 hεb' with ⟨ε, rfl, h⟩,
  rcases exists_nat_gt (↑n / ε) with ⟨m, hm⟩,
  have hε : ε > 0, from coe_lt_coe.1 hε',
  refine mem_sets_of_superset (prod_mem_nhds_sets (lt_mem_nhds $ @coe_lt_top m) (lt_mem_nhds $ h)) _,
  rintros ⟨a₁, a₂⟩ ⟨h₁, h₂⟩,
  dsimp at h₁ h₂ ⊢,
  calc (n:ennreal) = ↑(((n:nnreal) / ε) * ε) :
    begin
      simp [nnreal.div_def],
      rw [mul_assoc, ← coe_mul, nnreal.inv_mul_cancel, coe_one, coe_nat, mul_one],
      exact zero_lt_iff_ne_zero.1 hε
    end
    ... < (↑m * ε : nnreal) : coe_lt_coe.2 $ mul_lt_mul hm (le_refl _) hε (nat.cast_nonneg _)
    ... ≤ a₁ * a₂ : by rw [coe_mul]; exact canonically_ordered_semiring.mul_le_mul
      (le_of_lt h₁)
      (le_of_lt h₂)
end,
begin
  cases a, { simp [none_eq_top, ht b hb, top_mul, hb] },
  cases b, {
    have ha' : a ≠ 0, from mt coe_eq_coe.2 ha,
    simp [*, nhds_swap (a : ennreal) ⊤, none_eq_top, some_eq_coe, top_mul, tendsto_map'_iff, (∘), mul_comm] },
  simp [some_eq_coe, nhds_coe_coe, tendsto_map'_iff, (∘)],
  simp only [coe_mul.symm, tendsto_coe, tendsto_mul']
end

protected lemma tendsto_mul {f : filter α} {ma : α → ennreal} {mb : α → ennreal} {a b : ennreal}
  (hma : tendsto ma f (nhds a)) (ha : a ≠ 0) (hmb : tendsto mb f (nhds b)) (hb : b ≠ 0) :
  tendsto (λa, ma a * mb a) f (nhds (a * b)) :=
show tendsto ((λp:ennreal×ennreal, p.1 * p.2) ∘ (λa, (ma a, mb a))) f (nhds (a * b)), from
tendsto.comp (tendsto_prod_mk_nhds hma hmb) (ennreal.tendsto_mul' ha hb)

protected lemma tendsto_mul_right {f : filter α} {m : α → ennreal} {a b : ennreal}
  (hm : tendsto m f (nhds b)) (hb : b ≠ 0) : tendsto (λb, a * m b) f (nhds (a * b)) :=
by_cases
  (assume : a = 0, by simp [this, tendsto_const_nhds])
  (assume ha : a ≠ 0, ennreal.tendsto_mul tendsto_const_nhds ha hm hb)

lemma Sup_add {s : set ennreal} (hs : s ≠ ∅) : Sup s + a = ⨆b∈s, b + a :=
have Sup ((λb, b + a) '' s) = Sup s + a,
  from is_lub_iff_Sup_eq.mp $ is_lub_of_is_lub_of_tendsto
    (assume x _ y _ h, add_le_add' h (le_refl _))
    is_lub_Sup
    hs
    (tendsto_add (tendsto_id' inf_le_left) tendsto_const_nhds),
by simp [Sup_image, -add_comm] at this; exact this.symm

lemma supr_add {ι : Sort*} {s : ι → ennreal} [h : nonempty ι] : supr s + a = ⨆b, s b + a :=
let ⟨x⟩ := h in
calc supr s + a = Sup (range s) + a : by simp [Sup_range]
  ... = (⨆b∈range s, b + a) : Sup_add $ ne_empty_iff_exists_mem.mpr ⟨s x, x, rfl⟩
  ... = _ : by simp [supr_range, -mem_range]

lemma add_supr {ι : Sort*} {s : ι → ennreal} [h : nonempty ι] : a + supr s = ⨆b, a + s b :=
by rw [add_comm, supr_add]; simp

lemma supr_add_supr {ι : Sort*} [nonempty ι]
  {f g : ι → ennreal} (h : ∀i j, ∃k, f i + g j ≤ f k + g k) :
  supr f + supr g = (⨆ a, f a + g a) :=
begin
  refine le_antisymm _ (supr_le $ λ a, add_le_add' (le_supr _ _) (le_supr _ _)),
  simpa [add_supr, supr_add] using
    λ i j, show f i + g j ≤ ⨆ a, f a + g a, from
    let ⟨k, hk⟩ := h i j in le_supr_of_le k hk,
end

protected lemma tendsto_coe_sub : ∀{b:ennreal}, tendsto (λb:ennreal, ↑r - b) (nhds b) (nhds (↑r - b)) :=
begin
  refine (forall_ennreal.2 $ and.intro (assume a, _) _),
  { simp [@nhds_coe a, tendsto_map'_iff, (∘), tendsto_coe, coe_sub.symm],
    exact nnreal.tendsto_sub tendsto_const_nhds tendsto_id },
  simp,
  exact (tendsto_cong tendsto_const_nhds $ mem_sets_of_superset (lt_mem_nhds $ @coe_lt_top r) $
    by simp [le_of_lt] {contextual := tt})
end

lemma sub_supr {ι : Sort*} [hι : nonempty ι] {b : ι → ennreal} (hr : a < ⊤) :
  a - (⨆i, b i) = (⨅i, a - b i) :=
let ⟨i⟩ := hι in
let ⟨r, eq, _⟩ := lt_iff_exists_coe.mp hr in
have Inf ((λb, ↑r - b) '' range b) = ↑r - (⨆i, b i),
  from is_glb_iff_Inf_eq.mp $ is_glb_of_is_lub_of_tendsto
    (assume x _ y _, sub_le_sub (le_refl _))
    is_lub_supr
    (ne_empty_of_mem ⟨i, rfl⟩)
    (tendsto.comp (tendsto_id' inf_le_left) ennreal.tendsto_coe_sub),
by rw [eq, ←this]; simp [Inf_image, infi_range, -mem_range]; exact le_refl _

end topological_space

section tsum

variables {f g : α → ennreal}

protected lemma is_sum_coe {f : α → nnreal} {r : nnreal} :
  is_sum (λa, (f a : ennreal)) ↑r ↔ is_sum f r :=
have (λs:finset α, s.sum (coe ∘ f)) = (coe : nnreal → ennreal) ∘ (λs:finset α, s.sum f),
  from funext $ assume s, ennreal.coe_finset_sum.symm,
by unfold is_sum; rw [this, tendsto_coe]

protected lemma tsum_coe_eq {f : α → nnreal} (h : is_sum f r) : (∑a, (f a : ennreal)) = r :=
tsum_eq_is_sum $ ennreal.is_sum_coe.2 $ h

protected lemma tsum_coe {f : α → nnreal} : has_sum f → (∑a, (f a : ennreal)) = ↑(tsum f)
| ⟨r, hr⟩ := by rw [tsum_eq_is_sum hr, ennreal.tsum_coe_eq hr]

protected lemma is_sum : is_sum f (⨆s:finset α, s.sum f) :=
tendsto_orderable.2
  ⟨assume a' ha',
    let ⟨s, hs⟩ := lt_supr_iff.mp ha' in
    mem_at_top_sets.mpr ⟨s, assume t ht, lt_of_lt_of_le hs $ finset.sum_le_sum_of_subset ht⟩,
  assume a' ha',
    univ_mem_sets' $ assume s,
    have s.sum f ≤ ⨆(s : finset α), s.sum f,
      from le_supr (λ(s : finset α), s.sum f) s,
    lt_of_le_of_lt this ha'⟩

@[simp] protected lemma has_sum : has_sum f := ⟨_, ennreal.is_sum⟩

protected lemma tsum_eq_supr_sum : (∑a, f a) = (⨆s:finset α, s.sum f) :=
tsum_eq_is_sum ennreal.is_sum

protected lemma tsum_sigma {β : α → Type*} (f : Πa, β a → ennreal) :
  (∑p:Σa, β a, f p.1 p.2) = (∑a b, f a b) :=
tsum_sigma (assume b, ennreal.has_sum) ennreal.has_sum

protected lemma tsum_prod {f : α → β → ennreal} : (∑p:α×β, f p.1 p.2) = (∑a, ∑b, f a b) :=
let j : α × β → (Σa:α, β) := λp, sigma.mk p.1 p.2 in
let i : (Σa:α, β) → α × β := λp, (p.1, p.2) in
let f' : (Σa:α, β) → ennreal := λp, f p.1 p.2 in
calc (∑p:α×β, f' (j p)) = (∑p:Σa:α, β, f p.1 p.2) :
    tsum_eq_tsum_of_iso j i (assume ⟨a, b⟩, rfl) (assume ⟨a, b⟩, rfl)
   ... = (∑a, ∑b, f a b) : ennreal.tsum_sigma f

protected lemma tsum_comm {f : α → β → ennreal} : (∑a, ∑b, f a b) = (∑b, ∑a, f a b) :=
let f' : α×β → ennreal := λp, f p.1 p.2 in
calc (∑a, ∑b, f a b) = (∑p:α×β, f' p) : ennreal.tsum_prod.symm
  ... = (∑p:β×α, f' (prod.swap p)) :
    (tsum_eq_tsum_of_iso prod.swap (@prod.swap α β) (assume ⟨a, b⟩, rfl) (assume ⟨a, b⟩, rfl)).symm
  ... = (∑b, ∑a, f' (prod.swap (b, a))) : @ennreal.tsum_prod β α (λb a, f' (prod.swap (b, a)))

protected lemma tsum_add : (∑a, f a + g a) = (∑a, f a) + (∑a, g a) :=
tsum_add ennreal.has_sum ennreal.has_sum

protected lemma tsum_le_tsum (h : ∀a, f a ≤ g a) : (∑a, f a) ≤ (∑a, g a) :=
tsum_le_tsum h ennreal.has_sum ennreal.has_sum

protected lemma tsum_eq_supr_nat {f : ℕ → ennreal} :
  (∑i:ℕ, f i) = (⨆i:ℕ, (finset.range i).sum f) :=
calc _ = (⨆s:finset ℕ, s.sum f) : ennreal.tsum_eq_supr_sum
  ... = (⨆i:ℕ, (finset.range i).sum f) : le_antisymm
    (supr_le_supr2 $ assume s,
      let ⟨n, hn⟩ := finset.exists_nat_subset_range s in
      ⟨n, finset.sum_le_sum_of_subset hn⟩)
    (supr_le_supr2 $ assume i, ⟨finset.range i, le_refl _⟩)

protected lemma le_tsum (a : α) : f a ≤ (∑a, f a) :=
calc f a = ({a} : finset α).sum f : by simp
  ... ≤ (⨆s:finset α, s.sum f) : le_supr (λs:finset α, s.sum f) _
  ... = (∑a, f a) : by rw [ennreal.tsum_eq_supr_sum]

protected lemma mul_tsum : (∑i, a * f i) = a * (∑i, f i) :=
if h : ∀i, f i = 0 then by simp [h] else
let ⟨i, (hi : f i ≠ 0)⟩ := classical.not_forall.mp h in
have sum_ne_0 : (∑i, f i) ≠ 0, from ne_of_gt $
  calc 0 < f i : lt_of_le_of_ne (zero_le _) hi.symm
    ... ≤ (∑i, f i) : ennreal.le_tsum _,
have tendsto (λs:finset α, s.sum ((*) a ∘ f)) at_top (nhds (a * (∑i, f i))),
  by rw [← show (*) a ∘ (λs:finset α, s.sum f) = λs, s.sum ((*) a ∘ f),
         from funext $ λ s, finset.mul_sum];
  exact ennreal.tendsto_mul_right (is_sum_tsum ennreal.has_sum) sum_ne_0,
tsum_eq_is_sum this

protected lemma tsum_mul : (∑i, f i * a) = (∑i, f i) * a :=
by simp [mul_comm, ennreal.mul_tsum]

@[simp] lemma tsum_supr_eq {α : Type*} (a : α) {f : α → ennreal} :
  (∑b:α, ⨆ (h : a = b), f b) = f a :=
le_antisymm
  (by rw [ennreal.tsum_eq_supr_sum]; exact supr_le (assume s,
    calc s.sum (λb, ⨆ (h : a = b), f b) ≤ (finset.singleton a).sum (λb, ⨆ (h : a = b), f b) :
        finset.sum_le_sum_of_ne_zero $ assume b _ hb,
          suffices a = b, by simpa using this.symm,
          classical.by_contradiction $ assume h,
            by simpa [h] using hb
      ... = f a : by simp))
  (calc f a ≤ (⨆ (h : a = a), f a) : le_supr (λh:a=a, f a) rfl
    ... ≤ (∑b:α, ⨆ (h : a = b), f b) : ennreal.le_tsum _)

end tsum

end ennreal

namespace nnreal

lemma exists_le_is_sum_of_le {f g : β → nnreal} {r : nnreal} (hgf : ∀b, g b ≤ f b) (hfr : is_sum f r) :
  ∃p≤r, is_sum g p :=
have (∑b, (g b : ennreal)) ≤ r,
begin
  refine is_sum_le (assume b, _) (is_sum_tsum ennreal.has_sum) (ennreal.is_sum_coe.2 hfr),
  exact ennreal.coe_le_coe.2 (hgf _)
end,
let ⟨p, eq, hpr⟩ := ennreal.le_coe_iff.1 this in
⟨p, hpr, ennreal.is_sum_coe.1 $ eq ▸ is_sum_tsum ennreal.has_sum⟩

lemma has_sum_of_le {f g : β → nnreal} (hgf : ∀b, g b ≤ f b) : has_sum f → has_sum g
| ⟨r, hfr⟩ := let ⟨p, _, hp⟩ := exists_le_is_sum_of_le hgf hfr in has_sum_spec hp

end nnreal

lemma has_sum_of_nonneg_of_le {f g : β → ℝ} (hg : ∀b, 0 ≤ g b) (hgf : ∀b, g b ≤ f b) (hf : has_sum f) :
  has_sum g :=
let f' (b : β) : nnreal := ⟨f b, le_trans (hg b) (hgf b)⟩ in
let g' (b : β) : nnreal := ⟨g b, hg b⟩ in
have has_sum f', from nnreal.has_sum_coe.1 hf,
have has_sum g', from nnreal.has_sum_of_le (assume b, (nnreal.coe_le (g' b) (f' b)).2 $ hgf b) this,
show has_sum (λb, g' b : β → ℝ), from nnreal.has_sum_coe.2 this


/-
Copyright (c) 2020 Frédéric Dupuis. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Frédéric Dupuis
-/
import linear_algebra.affine_space.affine_map
import topology.algebra.group.basic
import topology.algebra.mul_action

/-!
# Topological properties of affine spaces and maps

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

For now, this contains only a few facts regarding the continuity of affine maps in the special
case when the point space and vector space are the same.

TODO: Deal with the case where the point spaces are different from the vector spaces. Note that
we do have some results in this direction under the assumption that the topologies are induced by
(semi)norms.
-/

namespace affine_map

variables {R E F : Type*}
variables [add_comm_group E] [topological_space E]
variables [add_comm_group F] [topological_space F] [topological_add_group F]

section ring

variables [ring R] [module R E] [module R F]

/-- An affine map is continuous iff its underlying linear map is continuous. See also
`affine_map.continuous_linear_iff`. -/
lemma continuous_iff {f : E →ᵃ[R] F} :
  continuous f ↔ continuous f.linear :=
begin
  split,
  { intro hc,
    rw decomp' f,
    have := hc.sub continuous_const,
    exact this, },
  { intro hc,
    rw decomp f,
    have := hc.add continuous_const,
    exact this }
end

/-- The line map is continuous. -/
@[continuity]
lemma line_map_continuous [topological_space R] [has_continuous_smul R F] {p v : F} :
  continuous ⇑(line_map p v : R →ᵃ[R] F) :=
continuous_iff.mpr $ (continuous_id.smul continuous_const).add $
  @continuous_const _ _ _ _ (0 : F)

end ring

section comm_ring

variables [comm_ring R] [module R F] [has_continuous_const_smul R F]

@[continuity]
lemma homothety_continuous (x : F) (t : R) : continuous $ homothety x t :=
begin
  suffices : ⇑(homothety x t) = λ y, t • (y - x) + x, { rw this, continuity, },
  ext y,
  simp [homothety_apply],
end

end comm_ring

section field

variables [field R] [module R F] [has_continuous_const_smul R F]

lemma homothety_is_open_map (x : F) (t : R) (ht : t ≠ 0) : is_open_map $ homothety x t :=
begin
  apply is_open_map.of_inverse (homothety_continuous x t⁻¹);
  intros e;
  simp [← affine_map.comp_apply, ← homothety_mul, ht],
end

end field

end affine_map


/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura, Jannis Limperg, Mario Carneiro
-/
import Std.Classes.Order
import Std.Data.Nat.Lemmas
import Std.Tactic.NoMatch

namespace Std.BinomialHeapImp

/--
A `HeapNode` is one of the internal nodes of the binomial heap.
It is always a perfect binary tree, with the depth of the tree stored in the `Heap`.
However the interpretation of the two pointers is different: we view the `child`
as going to the first child of this node, and `sibling` goes to the next sibling
of this tree. So it actually encodes a forest where each node has children
`node.child`, `node.child.sibling`, `node.child.sibling.sibling`, etc.

Each edge in this forest denotes a `le a b` relation that has been checked, so
the root is smaller than everything else under it.
-/
inductive HeapNode (α : Type u) where
  /-- An empty forest, which has depth `0`. -/
  | nil : HeapNode α
  /-- A forest of rank `r + 1` consists of a root `a`,
  a forest `child` of rank `r` elements greater than `a`,
  and another forest `sibling` of rank `r`. -/
  | node (a : α) (child sibling : HeapNode α) : HeapNode α
  deriving Repr

/--
The "real size" of the node, counting up how many values of type `α` are stored.
This is `O(n)` and is intended mainly for specification purposes.
For a well formed `HeapNode` the size is always `2^n - 1` where `n` is the depth.
-/
@[simp] def HeapNode.realSize : HeapNode α → Nat
  | .nil => 0
  | .node _ c s => c.realSize + 1 + s.realSize

/-- A node containing a single element `a`. -/
def HeapNode.singleton (a : α) : HeapNode α := .node a .nil .nil

/--
`O(log n)`. The rank, or the number of trees in the forest.
It is also the depth of the forest.
-/
@[inline] def HeapNode.rankTR (s : HeapNode α) : Nat := go s 0 where
  /-- Computes `s.rank + r` -/
  go : HeapNode α → Nat → Nat
  | .nil, r => r
  | .node _ _ s, r => go s (r + 1)

/--
A `Heap` is the top level structure in a binomial heap.
It consists of a forest of `HeapNode`s with strictly increasing ranks.
-/
inductive Heap (α : Type u) where
  /-- An empty heap. -/
  | nil : Heap α
  /-- A cons node contains a tree of root `val`, children `node` and rank `rank`,
  and then `next` which is the rest of the forest. -/
  | cons (rank : Nat) (val : α) (node : HeapNode α) (next : Heap α) : Heap α
  deriving Repr

/--
`O(n)`. The "real size" of the heap, counting up how many values of type `α` are stored.
This is intended mainly for specification purposes.
Prefer `Heap.size`, which is the same for well formed heaps.
-/
@[simp] def Heap.realSize : Heap α → Nat
  | .nil => 0
  | .cons _ _ c s => c.realSize + 1 + s.realSize

/-- `O(log n)`. The number of elements in the heap. -/
def Heap.size : Heap α → Nat
  | .nil => 0
  | .cons r _ _ s => 1 <<< r + s.realSize

/-- `O(1)`. Is the heap empty? -/
def Heap.isEmpty : Heap α → Bool
  | .nil => true
  | _    => false

/-- `O(1)`. The heap containing a single value `a`. -/
def Heap.singleton (a : α) : Heap α := .cons 0 a .nil .nil

/-- `O(1)`. Auxiliary for `Heap.merge`: Is the minimum rank in `Heap` strictly larger than `n`? -/
@[simp] def Heap.rankGT : Heap α → Nat → Prop
  | .nil, _ => True
  | .cons r .., n => n < r

instance : Decidable (Heap.rankGT s n) :=
  match s with
  | .nil => inferInstanceAs (Decidable True)
  | .cons .. => inferInstanceAs (Decidable (_ < _))

/-- `O(log n)`. The number of trees in the forest. -/
@[simp] def Heap.length : Heap α → Nat
  | .nil => 0
  | .cons _ _ _ r => r.length + 1

/--
`O(1)`. Auxiliary for `Heap.merge`: combines two heap nodes of the same rank
into one with the next larger rank.
-/
@[specialize] def combine (le : α → α → Bool) (a₁ a₂ : α) (n₁ n₂ : HeapNode α) : α × HeapNode α :=
  if le a₁ a₂ then (a₁, .node a₂ n₂ n₁) else (a₂, .node a₁ n₁ n₂)

/--
Merge two forests of binomial trees. The forests are assumed to be ordered
by rank and `merge` maintains this invariant.
-/
@[specialize] def Heap.merge (le : α → α → Bool) : Heap α → Heap α → Heap α
  | .nil, h  => h
  | h,  .nil => h
  | s₁@(.cons r₁ a₁ n₁ t₁), s₂@(.cons r₂ a₂ n₂ t₂) =>
    if r₁ < r₂ then .cons r₁ a₁ n₁ (merge le t₁ s₂)
    else if r₂ < r₁ then .cons r₂ a₂ n₂ (merge le s₁ t₂)
    else
      let (a, n) := combine le a₁ a₂ n₁ n₂
      let r := r₁ + 1
      if t₁.rankGT r then if t₂.rankGT r
        then .cons r a n (merge le t₁ t₂)
        else merge le (.cons r a n t₁) t₂
      else if t₂.rankGT r
        then merge le t₁ (.cons r a n t₂)
        else .cons r a n (merge le t₁ t₂)
termination_by _ s₁ s₂ => s₁.length + s₂.length
decreasing_by simp_wf; simp_arith [*]

/--
`O(log n)`. Convert a `HeapNode` to a `Heap` by reversing the order of the nodes
along the `sibling` spine.
-/
def HeapNode.toHeap (s : HeapNode α) : Heap α := go s s.rankTR .nil where
  /-- Computes `s.toHeap ++ res` tail-recursively, assuming `n = s.rank`. -/
  go : HeapNode α → Nat → Heap α → Heap α
  | .nil, _, res => res
  | .node a c s, n, res => go s (n - 1) (.cons (n - 1) a c res)

/-- `O(log n)`. Get the smallest element in the heap, including the passed in value `a`. -/
@[specialize] def Heap.headD (le : α → α → Bool) (a : α) : Heap α → α
  | .nil => a
  | .cons _ b _ hs => headD le (if le a b then a else b) hs

/-- `O(log n)`. Get the smallest element in the heap, if it has an element. -/
@[inline] def Heap.head? (le : α → α → Bool) : Heap α → Option α
  | .nil => none
  | .cons _ h _ hs => some <| headD le h hs

/--
The return type of `FindMin`, which encodes various quantities needed to
reconstruct the tree in `deleteMin`.
-/
structure FindMin (α) where
  /-- The list of elements prior to the minimum element, encoded as a "difference list". -/
  before : Heap α → Heap α := id
  /-- The minimum element. -/
  val : α
  /-- The children of the minimum element. -/
  node : HeapNode α
  /-- The forest after the minimum element. -/
  next : Heap α

/--
`O(log n)`. Find the minimum element, and return a data structure `FindMin` with information
needed to reconstruct the rest of the binomial heap.
-/
@[specialize] def Heap.findMin (le : α → α → Bool) (k : Heap α → Heap α) :
    Heap α → FindMin α → FindMin α
  | .nil, res => res
  | .cons r a c s, res =>
    -- It is important that we check `le res.val a` here, not the other way
    -- around. This ensures that head? and findMin find the same element even
    -- when we have `le res.val a` and `le a res.val` (i.e. le is not antisymmetric).
    findMin le (k ∘ .cons r a c) s <| if le res.val a then res else ⟨k, a, c, s⟩

/-- `O(log n)`. Find and remove the the minimum element from the binomial heap. -/
def Heap.deleteMin (le : α → α → Bool) : Heap α → Option (α × Heap α)
  | .nil => none
  | .cons r a c s =>
    let { before, val, node, next } := findMin le (.cons r a c) s ⟨id, a, c, s⟩
    some (val, node.toHeap.merge le (before next))

/-- `O(log n)`. Get the tail of the binomial heap after removing the minimum element. -/
@[inline] def Heap.tail? (le : α → α → Bool) (h : Heap α) : Option (Heap α) :=
  deleteMin le h |>.map (·.snd)

/-- `O(log n)`. Remove the minimum element of the heap. -/
@[inline] def Heap.tail (le : α → α → Bool) (h : Heap α) : Heap α :=
  tail? le h |>.getD .nil

theorem Heap.realSize_merge (le) (s₁ s₂ : Heap α) :
    (s₁.merge le s₂).realSize = s₁.realSize + s₂.realSize := by
  unfold merge; split
  · simp
  · simp
  · rename_i r₁ a₁ n₁ t₁ r₂ a₂ n₂ t₂
    have IH₁ r a n := realSize_merge le t₁ (cons r a n t₂)
    have IH₂ r a n := realSize_merge le (cons r a n t₁) t₂
    have IH₃ := realSize_merge le t₁ t₂
    split; { simp [IH₁, Nat.add_assoc] }
    split; { simp [IH₂, Nat.add_assoc, Nat.add_left_comm] }
    split; simp only; rename_i a n eq
    have : n.realSize = n₁.realSize + 1 + n₂.realSize := by
      rw [combine] at eq; split at eq <;> cases eq <;>
        simp [Nat.add_assoc, Nat.add_left_comm, Nat.add_comm]
    split <;> split <;> simp [IH₁, IH₂, IH₃, this, Nat.add_assoc, Nat.add_left_comm]
termination_by _ => s₁.length + s₂.length
decreasing_by simp_wf; simp_arith [*]

private def FindMin.HasSize (res : FindMin α) (n : Nat) : Prop :=
  ∃ m,
    (∀ s, (res.before s).realSize = m + s.realSize) ∧
    n = m + res.node.realSize + res.next.realSize + 1

theorem Heap.realSize_findMin {s : Heap α}
    (m) (hk : ∀ s, (k s).realSize = m + s.realSize)
    (eq : n = m + s.realSize) (hres : res.HasSize n) :
    (s.findMin le k res).HasSize n :=
  match s with
  | .nil => hres
  | .cons r a c s => by
    simp [findMin]
    refine realSize_findMin (m + c.realSize + 1)
      (by simp [hk, Nat.add_assoc]) (by simp [eq, Nat.add_assoc]) ?_
    split; { exact hres }
    exact ⟨m, hk, by simp [eq, Nat.add_assoc, Nat.add_comm, Nat.add_left_comm]⟩

theorem HeapNode.realSize_toHeap (s : HeapNode α) : s.toHeap.realSize = s.realSize := go s where
  go {n res} : ∀ s : HeapNode α, (toHeap.go s n res).realSize = s.realSize + res.realSize
  | .nil => (Nat.zero_add _).symm
  | .node a c s => by simp [toHeap.go, go, Nat.add_assoc, Nat.add_left_comm]

theorem Heap.realSize_deleteMin {s : Heap α} (eq : s.deleteMin le = some (a, s')) :
    s.realSize = s'.realSize + 1 := by
  cases s with cases eq | cons r a c s => ?_
  have : (s.findMin le (cons r a c) ⟨id, a, c, s⟩).HasSize (c.realSize + s.realSize + 1) :=
    Heap.realSize_findMin (c.realSize + 1) (by simp) (Nat.add_right_comm ..) ⟨0, by simp⟩
  revert this
  match s.findMin le (cons r a c) ⟨id, a, c, s⟩ with
  | { before, val, node, next } =>
    intro ⟨m, ih₁, ih₂⟩; dsimp only at ih₁ ih₂
    rw [realSize, Nat.add_right_comm, ih₂]
    simp only [realSize_merge, HeapNode.realSize_toHeap, ih₁, Nat.add_assoc, Nat.add_left_comm]

theorem Heap.realSize_tail? {s : Heap α} : s.tail? le = some s' →
    s.realSize = s'.realSize + 1 := by
  simp only [Heap.tail?]; intro eq
  match eq₂ : s.deleteMin le, eq with
  | some (a, tl), rfl => exact realSize_deleteMin eq₂

theorem Heap.realSize_tail (le) (s : Heap α) : (s.tail le).realSize = s.realSize - 1 := by
  simp only [Heap.tail]
  match eq : s.tail? le with
  | none => cases s with cases eq | nil => rfl
  | some tl => simp [Heap.realSize_tail? eq]; rfl

/--
`O(n log n)`. Monadic fold over the elements of a heap in increasing order,
by repeatedly pulling the minimum element out of the heap.
-/
@[specialize] def Heap.foldM [Monad m] (le : α → α → Bool) (s : Heap α)
    (init : β) (f : β → α → m β) : m β :=
  match eq : s.deleteMin le with
  | none => pure init
  | some (hd, tl) =>
    have : tl.realSize < s.realSize := by simp_arith [Heap.realSize_deleteMin eq]
    do foldM le tl (← f init hd) f
termination_by _ => s.realSize

/--
`O(n log n)`. Fold over the elements of a heap in increasing order,
by repeatedly pulling the minimum element out of the heap.
-/
@[inline] def Heap.fold (le : α → α → Bool) (s : Heap α) (init : β) (f : β → α → β) : β :=
  Id.run <| s.foldM le init f

/-- `O(n log n)`. Convert the heap to an array in increasing order. -/
@[inline] def Heap.toArray (le : α → α → Bool) (s : Heap α) : Array α := fold le s #[] Array.push

/-- `O(n log n)`. Convert the heap to a list in increasing order. -/
@[inline] def Heap.toList (le : α → α → Bool) (s : Heap α) : List α := (s.toArray le).toList

section
variable [Monad m] (nil : β) (join : α → β → β → m β)

/-- `O(n)`. Fold a monadic function over the tree structure to accumulate a value. -/
@[specialize] def HeapNode.foldTreeM : HeapNode α → m β
  | .nil => pure nil
  | .node a c s => do join a (← c.foldTreeM) (← s.foldTreeM)

/-- `O(n)`. Fold a monadic function over the tree structure to accumulate a value. -/
@[specialize] def Heap.foldTreeM : Heap α → m β
  | .nil => pure nil
  | .cons _ a c s => do join a (← c.foldTreeM nil join) (← s.foldTreeM)

end

/-- `O(n)`. Fold a function over the tree structure to accumulate a value. -/
@[inline] def Heap.foldTree (nil : β) (join : α → β → β → β) (s : Heap α) : β :=
  Id.run <| s.foldTreeM nil join

/-- `O(n)`. Convert the heap to a list in arbitrary order. -/
def Heap.toListUnordered (s : Heap α) : List α :=
  s.foldTree id (fun a c s l => a :: c (s l)) []

/-- `O(n)`. Convert the heap to an array in arbitrary order. -/
def Heap.toArrayUnordered (s : Heap α) : Array α :=
  s.foldTree id (fun a c s r => s (c (r.push a))) #[]

/--
The well formedness predicate for a heap node.
It asserts that:
* If `a` is added at the top to make the forest into a tree, the resulting tree
  is a `le`-min-heap (if `le` is well-behaved)
* When interpreting `child` and `sibling` as left and right children of a binary tree,
  it is a perfect binary tree with depth `r`
-/
def HeapNode.WellFormed (le : α → α → Bool) (a : α) : HeapNode α → Nat → Prop
  | .nil, r => r = 0
  | .node b c s, r => ∃ r', r = r' + 1 ∧ (∀ [TotalBLE le], le a b) ∧
    c.WellFormed le b r' ∧ s.WellFormed le a r'

/--
The well formedness predicate for a binomial heap.
It asserts that:
* It consists of a list of well formed trees with the specified ranks
* The ranks are in strictly increasing order, and all are at least `n`
-/
def Heap.WellFormed (le : α → α → Bool) (n : Nat) : Heap α → Prop
  | .nil => True
  | .cons r a c s => n ≤ r ∧ c.WellFormed le a r ∧ s.WellFormed le (r+1)

theorem Heap.WellFormed.nil : Heap.nil.WellFormed le n := trivial

theorem Heap.WellFormed.singleton : (Heap.singleton a).WellFormed le 0 := ⟨by decide, rfl, ⟨⟩⟩

theorem Heap.WellFormed.of_rankGT
  (hlt : s.rankGT n) (h : (s : Heap _).WellFormed le n') : s.WellFormed le (n+1) :=
  match s with
  | .nil => trivial
  | .cons .. => let ⟨_, h₂, h₃⟩ := h; ⟨hlt, h₂, h₃⟩

theorem Heap.WellFormed.of_le
    (hle : n ≤ n') (h : (s : Heap _).WellFormed le n') : s.WellFormed le n :=
  match s with
  | .nil => trivial
  | .cons .. => let ⟨h₁, h₂, h₃⟩ := h; ⟨Nat.le_trans hle h₁, h₂, h₃⟩

theorem Heap.rankGT.le_trans (h : Heap.rankGT s n) (h' : n' ≤ n) : s.rankGT n' :=
  match s with
  | .nil => trivial
  | .cons .. => Nat.lt_of_le_of_lt h' h

theorem Heap.WellFormed.rankGT (h : Heap.WellFormed lt (n+1) s) : s.rankGT n :=
  match s with
  | .nil => trivial
  | .cons .. => Nat.lt_of_succ_le h.1

theorem Heap.WellFormed.merge'
    (h₁ : s₁.WellFormed le n) (h₂ : s₂.WellFormed le n) :
    (merge le s₁ s₂).WellFormed le n ∧
    ((s₁.rankGT n ↔ s₂.rankGT n) → (merge le s₁ s₂).rankGT n) := by
  unfold merge; split
  · exact ⟨h₂, fun h => h.1 h₁⟩
  · exact ⟨h₁, fun h => h.2 h₂⟩
  · rename_i r₁ a₁ n₁ t₁ r₂ a₂ n₂ t₂
    let ⟨hr₁, hn₁, ht₁⟩ := h₁
    let ⟨hr₂, hn₂, ht₂⟩ := h₂
    split <;> rename_i lt₁
    · refine ⟨⟨hr₁, hn₁, And.left (merge' ht₁ ⟨lt₁, hn₂, ht₂⟩)⟩, fun h => ?_⟩
      exact h.2 <| Nat.lt_of_le_of_lt hr₁ lt₁
    split <;> rename_i lt₂
    · refine ⟨⟨hr₂, hn₂, And.left (merge' ⟨lt₂, hn₁, ht₁⟩ ht₂)⟩, fun h => ?_⟩
      exact h.1 <| Nat.lt_of_le_of_lt hr₂ lt₂
    cases Nat.le_antisymm (Nat.ge_of_not_lt lt₂) (Nat.ge_of_not_lt lt₁)
    split; rename_i a n eq
    have : n.WellFormed le a (r₁+1) := by
      unfold combine at eq; split at eq <;> cases eq <;> rename_i h
      · exact ⟨r₁, rfl, h, hn₂, hn₁⟩
      · exact ⟨r₁, rfl, TotalBLE.total.resolve_left h, hn₁, hn₂⟩
    simp only; split <;> split <;> rename_i hl₁ hl₂
    · exact ⟨⟨Nat.le_succ_of_le hr₁, this,
        (merge' (ht₁.of_rankGT hl₁) (ht₂.of_rankGT hl₂)).1⟩,
        fun _ => Nat.lt_succ_of_le hr₁⟩
    · let ⟨ih₁, ih₂⟩ := merge' (s₁ := .cons ..)
        ⟨Nat.le_succ_of_le hr₁, this, ht₁.of_rankGT hl₁⟩
        (ht₂.of_le (Nat.le_succ_of_le hr₁))
      exact ⟨ih₁, fun _ => ih₂ ⟨fun _ => ht₂.rankGT.le_trans hr₁, fun h => Nat.lt_succ_of_le hr₁⟩⟩
    · let ⟨ih₁, ih₂⟩ := merge' (s₂ := .cons ..) (ht₁.of_le (Nat.le_succ_of_le hr₁))
        ⟨Nat.le_succ_of_le hr₁, this, ht₂.of_rankGT hl₂⟩
      exact ⟨ih₁, fun _ => ih₂ ⟨fun h => Nat.lt_succ_of_le hr₁, fun _ => ht₁.rankGT.le_trans hr₁⟩⟩
    · let ⟨ih₁, ih₂⟩ := merge' ht₁ ht₂
      exact ⟨⟨Nat.le_succ_of_le hr₁, this, ih₁.of_rankGT (ih₂ (iff_of_false hl₁ hl₂))⟩,
        fun _ => Nat.lt_succ_of_le hr₁⟩
termination_by _ => s₁.length + s₂.length
decreasing_by simp_wf; simp_arith [*]

theorem Heap.WellFormed.merge
    (h₁ : s₁.WellFormed le n) (h₂ : s₂.WellFormed le n) :
    (merge le s₁ s₂).WellFormed le n := (merge' h₁ h₂).1

theorem Heap.findMin_val : ((s : Heap α).findMin le k res).val = s.headD le res.val :=
  match s with
  | .nil => rfl
  | .cons r a c s => by rw [findMin, headD]; split <;> apply findMin_val

theorem Heap.deleteMin_fst : ((s : Heap α).deleteMin le).map (·.1) = s.head? le :=
  match s with
  | .nil => rfl
  | .cons r a c s => by simp only [deleteMin, findMin_val, Option.map, head?]

/--
`O(log n)`. The rank, or the number of trees in the forest.
This is the same as `rankTR` but it is not tail recursive.
-/
@[simp] def HeapNode.rank : HeapNode α → Nat
  | .nil => 0
  | .node _ _ s => s.rank + 1

@[simp] theorem HeapNode.rankTR_eq (s : HeapNode α) : s.rankTR = s.rank := go s 0 where
  go : ∀ s n, rankTR.go s n = rank s + n
  | .nil, _ => (Nat.zero_add ..).symm
  | .node .., _ => by simp_arith only [rankTR.go, go, rank]

@[simp] theorem HeapNode.WellFormed.rank_eq :
    ∀ {n} {s : HeapNode α}, s.WellFormed le a n → s.rank = n
  | _, .nil, h => h.symm
  | _, .node .., ⟨_, rfl, _, _, h⟩ => congrArg Nat.succ (rank_eq h)

@[simp] theorem HeapNode.WellFormed.realSize_eq :
    ∀ {n} {s : HeapNode α}, s.WellFormed le a n → s.realSize + 1 = 2 ^ n
  | _, .nil, rfl => rfl
  | _, .node .., ⟨_, rfl, _, c, s⟩ => by
    rw [realSize, realSize_eq c, Nat.pow_succ, Nat.mul_succ]
    simp [Nat.add_assoc, realSize_eq s]

@[simp] theorem Heap.WellFormed.size_eq :
    ∀ {s : Heap α}, s.WellFormed le n → s.size = s.realSize
  | .nil, _ => rfl
  | .cons .., ⟨_, h₁, h₂⟩ => by
    simp [size, Nat.shiftLeft, size_eq h₂, Nat.pow_succ, Nat.mul_succ]
    simp [Nat.add_assoc, Nat.one_shiftLeft, h₁.realSize_eq, h₂.size_eq]

theorem HeapNode.WellFormed.toHeap {s : HeapNode α}
    (h : s.WellFormed le a n) : s.toHeap.WellFormed le 0 := by
  rw [toHeap, rankTR_eq]
  exact go h trivial
where
  go {res} : ∀ {n s}, s.WellFormed le a n → res.WellFormed le s.rank →
    (HeapNode.toHeap.go s s.rank res).WellFormed le 0
  | _, .nil, _, hr => hr
  | _, .node a c s, ⟨n, rfl, _, h, h'⟩, hr =>
    go (s := s) h' ⟨Nat.le_refl _, by rw [← h'.rank_eq] at h; exact h, hr⟩

/--
The well formedness predicate for a `FindMin` value.
This is not actually a predicate, as it contains an additional data value
`rank` corresponding to the rank of the returned node, which is omitted from `findMin`.
-/
structure FindMin.WellFormed (le : α → α → Bool) (res : FindMin α) where
  /-- The rank of the minimum element -/
  rank : Nat
  /-- `before` is a difference list which can be appended to a binomial heap
  with ranks at least `rank` to produce another well formed heap. -/
  before : ∀ {s}, s.WellFormed le rank → (res.before s).WellFormed le 0
  /-- `node` is a well formed forest of rank `rank` with `val` at the root. -/
  node : res.node.WellFormed le res.val rank
  /-- `next` is a binomial heap with ranks above `rank + 1`. -/
  next : res.next.WellFormed le (rank + 1)

theorem Heap.WellFormed.findMin {s : Heap α}
    (h : s.WellFormed le n) (hr : res.WellFormed le)
    (hk : ∀ {s}, s.WellFormed le n → (k s).WellFormed le 0) :
    Nonempty (((s : Heap α).findMin le k res).WellFormed le) :=
  match s with
  | .nil => ⟨hr⟩
  | .cons r a c s => by
    let ⟨h₁, h₂, h₃⟩ := h
    simp [Heap.findMin]; split
    · exact findMin h₃ hr (fun h => hk ⟨h₁, h₂, h⟩)
    · exact findMin h₃ ⟨_, fun h => hk (h.of_le h₁), h₂, h₃⟩ (fun h => hk ⟨h₁, h₂, h⟩)

theorem Heap.WellFormed.deleteMin {s : Heap α}
    (h : s.WellFormed le n) (eq : s.deleteMin le = some (a, s')) : s'.WellFormed le 0 := by
  cases s with cases eq | cons r a c s => ?_
  have : Nonempty ((s.findMin le (cons r a c) ⟨id, a, c, s⟩).WellFormed le) :=
    let ⟨_, h₂, h₃⟩ := h
    h₃.findMin ⟨_, fun h => h.of_le (Nat.zero_le _), h₂, h₃⟩
      fun h => ⟨Nat.zero_le _, h₂, h⟩
  revert this
  match s.findMin le (cons r a c) ⟨id, a, c, s⟩ with
  | { before, val, node, next } =>
    intro ⟨⟨_, hk, ih₁, ih₂⟩⟩
    exact ih₁.toHeap.merge <| hk (ih₂.of_le (Nat.le_succ _))

theorem Heap.WellFormed.tail?
    (hwf : (s : Heap α).WellFormed le n) : s.tail? le = some tl →
  tl.WellFormed le 0 := by
  simp only [Heap.tail?]; intro eq
  match eq₂ : s.deleteMin le, eq with
  | some (a, tl), rfl => exact hwf.deleteMin eq₂

theorem Heap.WellFormed.tail
    (hwf : (s : Heap α).WellFormed le n) : (s.tail le).WellFormed le 0 := by
  simp only [Heap.tail]
  match eq : s.tail? le with
  | none => exact Heap.WellFormed.nil
  | some tl => exact hwf.tail? eq

end BinomialHeapImp

open BinomialHeapImp

/--
A [binomial heap](https://en.wikipedia.org/wiki/Binomial_heap) is a data structure which supports
the following primary operations:

* `insert : α → BinomialHeap α → BinomialHeap α`: add an element to the heap
* `deleteMin : BinomialHeap α → Option (α × BinomialHeap α)`:
  remove the minimum element from the heap
* `merge : BinomialHeap α → BinomialHeap α → BinomialHeap α`: combine two heaps

The first two operations are known as a "priority queue", so this could be called
a "mergeable priority queue". The standard choice for a priority queue is a binary heap,
which supports `insert` and `deleteMin` in `O(log n)`, but `merge` is `O(n)`.
With a `BinomialHeap`, all three operations are `O(log n)`.
-/
def BinomialHeap (α : Type u) (le : α → α → Bool) :=
  { h : Heap α // h.WellFormed le 0 }

/-- `O(1)`. Make a new empty binomial heap. -/
@[inline] def mkBinomialHeap (α : Type u) (le : α → α → Bool) : BinomialHeap α le :=
  ⟨.nil, Heap.WellFormed.nil⟩

namespace BinomialHeap
variable {α : Type u} {le : α → α → Bool}

/-- `O(1)`. Make a new empty binomial heap. -/
@[inline] def empty : BinomialHeap α le := mkBinomialHeap α le

instance : Inhabited (BinomialHeap α le) := ⟨.empty⟩

/-- `O(1)`. Is the heap empty? -/
@[inline] def isEmpty (b : BinomialHeap α le) : Bool := b.1.isEmpty

/-- `O(log n)`. The number of elements in the heap. -/
@[inline] def size (b : BinomialHeap α le) : Nat := b.1.size

/-- `O(1)`. Make a new heap containing `a`. -/
@[inline] def singleton (a : α) : BinomialHeap α le :=
  ⟨Heap.singleton a, Heap.WellFormed.singleton⟩

/-- `O(log n)`. Merge the contents of two heaps. -/
@[inline] def merge : BinomialHeap α le → BinomialHeap α le → BinomialHeap α le
  | ⟨b₁, h₁⟩, ⟨b₂, h₂⟩ => ⟨b₁.merge le b₂, h₁.merge h₂⟩

/-- `O(log n)`. Add element `a` to the given heap `h`. -/
@[inline] def insert (a : α) (h : BinomialHeap α le) : BinomialHeap α le :=
  merge (singleton a) h

/-- `O(n log n)`. Construct a heap from a list by inserting all the elements. -/
def ofList (le : α → α → Bool) (as : List α) : BinomialHeap α le :=
  as.foldl (flip insert) empty

/-- `O(n log n)`. Construct a heap from a list by inserting all the elements. -/
def ofArray (le : α → α → Bool) (as : Array α) : BinomialHeap α le :=
  as.foldl (flip insert) empty

/-- `O(log n)`. Remove and return the minimum element from the heap. -/
@[inline] def deleteMin (b : BinomialHeap α le) : Option (α × BinomialHeap α le) :=
  match eq : b.1.deleteMin le with
  | none => none
  | some (a, tl) => some (a, ⟨tl, b.2.deleteMin eq⟩)

/-- `O(log n)`. Returns the smallest element in the heap, or `none` if the heap is empty. -/
@[inline] def head? (b : BinomialHeap α le) : Option α := b.1.head? le

/-- `O(log n)`. Returns the smallest element in the heap, or panics if the heap is empty. -/
@[inline] def head! [Inhabited α] (b : BinomialHeap α le) : α := b.head?.get!

/-- `O(log n)`. Returns the smallest element in the heap, or `default` if the heap is empty. -/
@[inline] def headI [Inhabited α] (b : BinomialHeap α le) : α := b.head?.getD default

/-- `O(log n)`. Removes the smallest element from the heap, or `none` if the heap is empty. -/
@[inline] def tail? (b : BinomialHeap α le) : Option (BinomialHeap α le) :=
  match eq : b.1.tail? le with
  | none => none
  | some tl => some ⟨tl, b.2.tail? eq⟩

/-- `O(log n)`. Removes the smallest element from the heap, if possible. -/
@[inline] def tail (b : BinomialHeap α le) : BinomialHeap α le := ⟨b.1.tail le, b.2.tail⟩

/-- `O(n log n)`. Convert the heap to a list in increasing order. -/
@[inline] def toList (b : BinomialHeap α le) : List α := b.1.toList le

/-- `O(n log n)`. Convert the heap to an array in increasing order. -/
@[inline] def toArray (b : BinomialHeap α le) : Array α := b.1.toArray le

/-- `O(n)`. Convert the heap to a list in arbitrary order. -/
@[inline] def toListUnordered (b : BinomialHeap α le) : List α := b.1.toListUnordered

/-- `O(n)`. Convert the heap to an array in arbitrary order. -/
@[inline] def toArrayUnordered (b : BinomialHeap α le) : Array α := b.1.toArrayUnordered


/-
Copyright (c) 2023 Eric Wieser. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Eric Wieser
-/
import algebra.dual_number
import algebra.quaternion

/-!
# Dual quaternions

Similar to the way that rotations in 3D space can be represented by quaternions of unit length,
rigid motions in 3D space can be represented by dual quaternions of unit length.

## Main results

* `quaternion.dual_number_equiv`: quaternions over dual numbers or dual
  numbers over quaternions are equivalent constructions.

## References

* <https://en.wikipedia.org/wiki/Dual_quaternion>
-/

variables {R : Type*} [comm_ring R]

namespace quaternion

/-- The dual quaternions can be equivalently represented as a quaternion with dual coefficients,
or as a dual number with quaternion coefficients.

See also `matrix.dual_number_equiv` for a similar result. -/
def dual_number_equiv :
  quaternion (dual_number R) ≃ₐ[R] dual_number (quaternion R) :=
{ to_fun := λ q,
    (⟨q.re.fst, q.im_i.fst, q.im_j.fst, q.im_k.fst⟩,
     ⟨q.re.snd, q.im_i.snd, q.im_j.snd, q.im_k.snd⟩),
  inv_fun := λ d,
    ⟨(d.fst.re, d.snd.re), (d.fst.im_i, d.snd.im_i),
     (d.fst.im_j, d.snd.im_j), (d.fst.im_k, d.snd.im_k)⟩,
  left_inv := λ ⟨⟨r, rε⟩, ⟨i, iε⟩, ⟨j, jε⟩, ⟨k, kε⟩⟩, rfl,
  right_inv := λ ⟨⟨r, i, j, k⟩, ⟨rε, iε, jε, kε⟩⟩, rfl,
  map_mul' := begin
    rintros ⟨⟨xr, xrε⟩, ⟨xi, xiε⟩, ⟨xj, xjε⟩, ⟨xk, xkε⟩⟩,
    rintros ⟨⟨yr, yrε⟩, ⟨yi, yiε⟩, ⟨yj, yjε⟩, ⟨yk, ykε⟩⟩,
    ext : 1,
    { refl },
    { dsimp,
      congr' 1; ring },
  end,
  map_add' := begin
    rintros ⟨⟨xr, xrε⟩, ⟨xi, xiε⟩, ⟨xj, xjε⟩, ⟨xk, xkε⟩⟩,
    rintros ⟨⟨yr, yrε⟩, ⟨yi, yiε⟩, ⟨yj, yjε⟩, ⟨yk, ykε⟩⟩,
    refl
  end,
  commutes' := λ r, rfl }

/-! Lemmas characterizing `quaternion.dual_number_equiv`. -/

-- `simps` can't work on `dual_number` because it's not a structure
@[simp] lemma re_fst_dual_number_equiv (q : quaternion (dual_number R)) :
  (dual_number_equiv q).fst.re = q.re.fst := rfl
@[simp] lemma im_i_fst_dual_number_equiv (q : quaternion (dual_number R)) :
  (dual_number_equiv q).fst.im_i = q.im_i.fst := rfl
@[simp] lemma im_j_fst_dual_number_equiv (q : quaternion (dual_number R)) :
  (dual_number_equiv q).fst.im_j = q.im_j.fst := rfl
@[simp] lemma im_k_fst_dual_number_equiv (q : quaternion (dual_number R)) :
  (dual_number_equiv q).fst.im_k = q.im_k.fst := rfl
@[simp] lemma re_snd_dual_number_equiv (q : quaternion (dual_number R)) :
  (dual_number_equiv q).snd.re = q.re.snd := rfl
@[simp] lemma im_i_snd_dual_number_equiv (q : quaternion (dual_number R)) :
  (dual_number_equiv q).snd.im_i = q.im_i.snd := rfl
@[simp] lemma im_j_snd_dual_number_equiv (q : quaternion (dual_number R)) :
  (dual_number_equiv q).snd.im_j = q.im_j.snd := rfl
@[simp] lemma im_k_snd_dual_number_equiv (q : quaternion (dual_number R)) :
  (dual_number_equiv q).snd.im_k = q.im_k.snd := rfl
@[simp] lemma fst_re_dual_number_equiv_symm (d : dual_number (quaternion R)) :
  (dual_number_equiv.symm d).re.fst = d.fst.re := rfl
@[simp] lemma fst_im_i_dual_number_equiv_symm (d : dual_number (quaternion R)) :
  (dual_number_equiv.symm d).im_i.fst = d.fst.im_i := rfl
@[simp] lemma fst_im_j_dual_number_equiv_symm (d : dual_number (quaternion R)) :
  (dual_number_equiv.symm d).im_j.fst = d.fst.im_j := rfl
@[simp] lemma fst_im_k_dual_number_equiv_symm (d : dual_number (quaternion R)) :
  (dual_number_equiv.symm d).im_k.fst = d.fst.im_k := rfl
@[simp] lemma snd_re_dual_number_equiv_symm (d : dual_number (quaternion R)) :
  (dual_number_equiv.symm d).re.snd = d.snd.re := rfl
@[simp] lemma snd_im_i_dual_number_equiv_symm (d : dual_number (quaternion R)) :
  (dual_number_equiv.symm d).im_i.snd = d.snd.im_i := rfl
@[simp] lemma snd_im_j_dual_number_equiv_symm (d : dual_number (quaternion R)) :
  (dual_number_equiv.symm d).im_j.snd = d.snd.im_j := rfl
@[simp] lemma snd_im_k_dual_number_equiv_symm (d : dual_number (quaternion R)) :
  (dual_number_equiv.symm d).im_k.snd = d.snd.im_k := rfl

end quaternion


/-
Copyright (c) 2019 Johan Commelin. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Johan Commelin, Simon Hudon, Scott Morrison
-/
import Mathlib.Data.Equiv.Basic

/-!

# Functors can be applied to `equiv`s.

```
def functor.map_equiv (f : Type u → Type v) [functor f] [is_lawful_functor f] :
  α ≃ β → f α ≃ f β
```

-/

open Equiv

namespace Functor

variable (f : Type u → Type v) [Functor f] [LawfulFunctor f]

-- this is in control.basic in Lean 3
theorem map_map (m : α → β) (g : β → γ) (x : f α) :
  g <$> (m <$> x) = (g ∘ m) <$> x :=
(comp_map _ _ _).symm

/-- Apply a functor to an `equiv`. -/
def map_equiv (h : α ≃ β) : f α ≃ f β where
  toFun    := map h
  invFun   := map h.symm
  left_inv x := by simp [map_map]
  right_inv x := by simp [map_map]

@[simp]
lemma map_equiv_apply (h : α ≃ β) (x : f α) :
  (map_equiv f h : f α ≃ f β) x = map h x := rfl

@[simp]
lemma map_equiv_symm_apply (h : α ≃ β) (y : f β) :
  (map_equiv f h : f α ≃ f β).symm y = map h.symm y := rfl

end Functor


import UniqueList

universe u

namespace Nat
  theorem repeat_assoc {α : Type u} :
    (f : α → α) →
    (n : Nat) →
    (a : α) →
    Nat.repeat f n (f a) = f (Nat.repeat f n a)
  | f, 0, a => Eq.refl _
  | f, n+1, a => by
    let h : Nat.repeat f (n + 1) (f a) = Nat.repeat f n (f (f a)) := Eq.refl _
    let h' : Nat.repeat f (n + 1) a = Nat.repeat f n (f a) := Eq.refl _
    rw [h, h']
    rw [repeat_assoc f n (f a)]

  theorem repeat_add {α : Type u} :
    (f : α → α) →
    (m n : Nat) →
    (a : α) →
    Nat.repeat f m (Nat.repeat f n a) = Nat.repeat f (m + n) a
  | f, m, 0, a => by rfl
  | f, m, n+1, a => repeat_add f m n (f a)

end Nat

def FinArray (size : Nat) :=
  {xs : Array (Fin size) // xs.size = size}

namespace FinArray
  def get {size : Nat} (array : FinArray size) (id : Fin size) : Fin size :=
    array.val.get (
      let in_range : id.val < Array.size array.val := by
        rw [array.property]
        exact id.isLt
      ⟨id.val, in_range⟩
    )

  def set {size : Nat} (array : FinArray size) (loc : Fin size) (new : Fin size) : FinArray size :=
    ⟨
      array.val.set (
        let in_range : loc.val < Array.size array.val := by
          rw [array.property]
          exact loc.isLt
        ⟨loc.val, in_range⟩
      ) new,
      (array.val.size_set _ _).trans array.property
    ⟩

  theorem get_set {size : Nat} :
    ∀ array : FinArray size,
    ∀ loc val : Fin size,
    (array.set loc val).get loc = val
  | ⟨⟨list⟩, len_proof⟩, ⟨loc, loc_lt_size⟩, val =>
    let list_len_proof : list.length = size := len_proof
    let list_loc : Fin (list.length) := ⟨loc, by
      rw [list_len_proof]
      exact loc_lt_size
    ⟩
    list.get_set list_loc val

  theorem get_set_ne {size : Nat} :
    ∀ array : FinArray size,
    ∀ setloc getloc val : Fin size,
    setloc ≠ getloc →
    array.get getloc = (array.set setloc val).get getloc
  | ⟨⟨list⟩, len_proof⟩, ⟨setloc, hsetloc⟩, ⟨getloc, hgetloc⟩, val, setloc_ne_getloc =>
    let list_len_proof : list.length = size := len_proof
    let list_setloc : Fin (list.length) := ⟨setloc, by
      rw [list_len_proof]
      exact hsetloc
    ⟩
    let list_getloc : Fin (list.length) := ⟨getloc, by
      rw [list_len_proof]
      exact hgetloc
    ⟩
    let list_setloc_ne_list_getloc : list_setloc ≠ list_getloc := (
      λ equal =>
      let h : setloc = getloc := Fin.val_eq_of_eq equal
      setloc_ne_getloc (Fin.eq_of_val_eq h)
    )
    list.get_set_ne list_setloc list_getloc val list_setloc_ne_list_getloc
end FinArray

def UnionFindLinks (size : Nat) :=
  {
    parent_array : FinArray size //
    ∀ id, ∃ steps,
      let is_root := λ x => parent_array.get x = x
      is_root (Nat.repeat parent_array.get steps id) ∧
      ∀ n, n < steps → ¬(is_root (Nat.repeat parent_array.get n id))
  }

namespace UnionFindLinks

  def parent {size : Nat} (uf : UnionFindLinks size) (id : Fin size) : Fin size :=
    uf.val.get id

  def is_root {size : Nat} (uf : UnionFindLinks size) (id : Fin size) : Prop :=
    uf.parent id = id

  def rooted_in_steps {size : Nat} (uf : UnionFindLinks size) (id : Fin size) (steps : Nat) : Prop :=
    uf.is_root (Nat.repeat uf.parent steps id)

  def rooted_in_exact_steps {size : Nat} (uf : UnionFindLinks size) (id : Fin size) (steps : Nat) : Prop :=
    uf.is_root (Nat.repeat uf.parent steps id) ∧
    ∀ n : Nat, n < steps → ¬(uf.is_root (Nat.repeat uf.parent n id))

  def roots {size : Nat} (uf : UnionFindLinks size) : Type :=
    {id : Fin size // uf.is_root id}

  def is_descendency {size : Nat} (uf : UnionFindLinks size) : Fin size → List (Fin size) → Prop
  | _x, [] => True
  | final, (x::xs) => (
    uf.parent x = final ∧ uf.is_descendency x xs
  )

  theorem root_ancestor_root {size : Nat} (uf : UnionFindLinks size) :
    ∀ id steps, uf.is_root id → uf.is_root (Nat.repeat uf.parent steps id) := by
    intro id steps
    induction steps
    case zero =>
      intro h; exact h
    case succ steps h_ind =>
      intro id_is_root
      let hroot := h_ind id_is_root
      let h_eq : uf.parent (Nat.repeat (uf.parent) steps id) = (Nat.repeat (uf.parent) steps id) := hroot
      let associate_parent := Nat.repeat_assoc uf.parent steps id
      rw [associate_parent.symm] at h_eq
      let h_eq : Nat.repeat (parent uf) (steps+1) id = Nat.repeat (parent uf) steps id := h_eq
      rw [h_eq]
      exact hroot

  theorem exact_rooted {size : Nat} (uf : UnionFindLinks size) :
    ∀ id, ∃ steps, uf.rooted_in_exact_steps id steps :=
      uf.property

  def from_exact_rooted {size : Nat} :
    (parent_array : FinArray size) →
    (∀ id, ∃ steps,
      let is_root := λ x => parent_array.get x = x
      is_root (Nat.repeat parent_array.get steps id) ∧
      ∀ n, n < steps → ¬(is_root (Nat.repeat parent_array.get n id))
    ) →
    UnionFindLinks size :=
    λ parent_array is_exact_rooted => ⟨parent_array, is_exact_rooted⟩

  theorem rooted {size : Nat} (uf : UnionFindLinks size) :
    ∀ id, ∃ steps, uf.rooted_in_steps id steps :=
      λ id =>
      (uf.exact_rooted id).elim (
        λ steps ⟨rooted_in_steps, _⟩ =>
        ⟨steps, rooted_in_steps⟩
      )

  def from_rooted {size : Nat} :
    (parent_array : FinArray size) →
    (∀ id, ∃ steps,
      let r := Nat.repeat parent_array.get steps id
      parent_array.get r = r) →
    UnionFindLinks size
  | parent_array, rooted => ⟨parent_array,
      λ id =>
      (rooted id).elim (by
        intro steps
        induction steps
        case zero =>
          intro h_is_root
          exact ⟨0, h_is_root,
            λ n n_lt_zero =>
            False.elim (Nat.not_lt_zero n n_lt_zero)
          ⟩
        case succ steps h_ind =>
          intro root_in_succ_steps
          exact (
            let l := Nat.repeat parent_array.get steps id
            if h_steps_root : parent_array.get l = l then
              h_ind h_steps_root
            else
              let is_root := λ x => parent_array.get x = x
              let no_roots_below_aux : ∀ (k : Nat), k < steps + 1 → ¬is_root (Nat.repeat parent_array.get (steps - k) id) := by
                intro k
                induction k
                case zero =>
                  intro _
                  exact h_steps_root
                case succ k h_ind =>
                  intro succ_k_lt_succ_steps
                  let k_lt_steps := Nat.lt_of_succ_lt_succ succ_k_lt_succ_steps
                  let steps_sub_k_ne_zero := Nat.sub_ne_zero_of_lt k_lt_steps
                  let h : steps - (k + 1) = steps - k - 1 := Nat.sub_succ steps k
                  let h' : steps - k - 1 + 1 = steps - k := Nat.succ_pred steps_sub_k_ne_zero
                  intro root_in_succ_n_steps
                  rw [h] at root_in_succ_n_steps
                  rw [h'.symm] at h_ind
                  let k_lt_succ_steps := Nat.lt_of_succ_lt succ_k_lt_succ_steps
                  let ind_h : ¬is_root ( (Nat.repeat parent_array.get (steps - k - 1) (parent_array.get id)) ) := h_ind k_lt_succ_steps
                  let h'' :
                    Nat.repeat parent_array.get (steps - k - 1) (parent_array.get id) =
                    parent_array.get (Nat.repeat parent_array.get (steps - k - 1) id) :=
                    Nat.repeat_assoc parent_array.get (steps - k - 1) id
                  rw [h''] at ind_h
                  rw [root_in_succ_n_steps] at ind_h
                  exact ind_h root_in_succ_n_steps

              let no_roots_below : ∀ (n : Nat), n < steps + 1 → ¬is_root (Nat.repeat parent_array.get n id) := by
                intro n n_lt_succ_steps
                let k := steps - n
                let k_le_steps : k ≤ steps := Nat.sub_le steps n
                let n_le_steps : n ≤ steps := Nat.le_of_lt_succ n_lt_succ_steps
                let k_lt_succ_steps : k < steps + 1 := Nat.lt_succ_of_le k_le_steps
                let aux_output := no_roots_below_aux k k_lt_succ_steps
                let steps_eq_n_add_k : steps = n + k := (Nat.eq_add_of_sub_eq n_le_steps (Eq.refl _)).trans (Nat.add_comm k n)
                let steps_sub_k_eq_n : steps - k = n := Nat.sub_eq_of_eq_add steps_eq_n_add_k
                rw [steps_sub_k_eq_n] at aux_output
                exact aux_output

              ⟨steps+1, root_in_succ_steps, no_roots_below⟩
          )
      )
    ⟩

  theorem no_cycles {size : Nat} (uf : UnionFindLinks size) :
    ∀ id steps,
      Nat.repeat uf.parent (steps+1) id = id →
      uf.is_root id :=
    λ id steps repeat_in_succ_steps =>
    if hroot : uf.parent id = id then
      hroot
    else
      let cycles_in_succ_steps : ∀ n : Nat, Nat.repeat uf.parent ((steps+1) * n) id = id := by
        intro n
        induction n
        case zero =>
          rfl
        case succ n h_ind =>
          let h : (steps+1) * (n+1) = (steps+1) * n + (steps+1) := by rfl
          rw [h]
          rw [(Nat.repeat_add uf.parent ((steps+1) * n) (steps+1) id).symm]
          rw [repeat_in_succ_steps]
          rw [h_ind]

      let multiple_not_root : ∀ n : Nat, ¬uf.is_root (Nat.repeat uf.parent ((steps+1) * n) id) := by
        intro n
        rw [cycles_in_succ_steps n]
        exact hroot

      let never_reaching_root : ∀ n : Nat, ¬uf.is_root (Nat.repeat uf.parent n id) :=
        λ n n_is_root =>
        let rem := n % (steps+1)
        let to_next := (steps+1) - rem
        let quot := n / (steps+1)
        let n_sum : (steps+1) * quot + rem = n := Nat.div_add_mod n (steps+1)
        let total_sum : n + to_next = (steps+1) * (quot+1) := by
          rw [n_sum.symm]
          rw [Nat.add_assoc, Nat.add_comm rem to_next]
          let h : to_next = steps + 1 - rem := by rfl
          rw [h]
          let rem_not_large : rem ≤ steps+1 :=
            Nat.le_of_lt (Nat.mod_lt n (Nat.zero_lt_succ steps))
          rw [Nat.sub_add_cancel rem_not_large]
          rfl
        let multiple_is_root := uf.root_ancestor_root (Nat.repeat (parent uf) n id) to_next n_is_root
        (by
          rw [Nat.repeat_add uf.parent to_next n, Nat.add_comm, total_sum] at multiple_is_root
          exact multiple_not_root (quot+1) multiple_is_root
        )

      (uf.rooted id).elim (
        λ n rooted_in_n =>
        False.elim (
          never_reaching_root n rooted_in_n
        )
      )

  /-
  def from_no_cycles {size : Nat} :
    (parent_array : FinArray size) →
    (∀ id steps,
      Nat.repeat parent_array.get (steps+1) id = id →
      parent_array.get id = id) →
    UnionFindLinks size
  | parent_array, no_cycles => ⟨parent_array, sorry⟩
  -/

  theorem path_containing_ancestor_end_ends_repeat {size : Nat} (uf : UnionFindLinks size) :
    ∀ id final path, uf.is_descendency final path → path.prop_contains id →
    (∃ steps, Nat.repeat uf.parent steps final = id) →
    ∃ steps, Nat.repeat uf.parent (steps+1) id = id
    | _, _, [] => by intro _ _ _; contradiction
    | id, final, (x::xs) => by
      intro is_desc id_in_xs final_desc_id
      cases id_in_xs
      case inl x_eq_id =>
        exact final_desc_id.elim ( by
          intro n final_to_id_in_n_steps
          rw [is_desc.left.symm, x_eq_id] at final_to_id_in_n_steps
          exact ⟨n, final_to_id_in_n_steps⟩
        )
      case inr id_in_xs =>
        let h := uf.path_containing_ancestor_end_ends_repeat id x xs is_desc.right id_in_xs
        exact final_desc_id.elim ( by
          intro n final_to_id_in_n_steps
          rw [is_desc.left.symm] at final_to_id_in_n_steps
          exact h ⟨n+1, final_to_id_in_n_steps⟩
        )

  theorem path_containing_end_ends_root {size : Nat} (uf : UnionFindLinks size) :
    ∀ id path, uf.is_descendency id path →
    path.prop_contains id → uf.is_root id :=
      λ id path is_desc id_in_path =>
      let id_is_repeat := uf.path_containing_ancestor_end_ends_repeat id id path is_desc id_in_path ⟨0, Eq.refl _⟩
      id_is_repeat.elim (uf.no_cycles id)

  theorem non_repeating_path {size : Nat} (uf : UnionFindLinks size) :
    ∀ id path, uf.is_descendency id path → ¬uf.is_root id → (id::path).unique := by
      intro id path is_desc not_root
      constructor
      case left =>
        intro id_in_path
        apply not_root
        let id_is_repeat := uf.path_containing_ancestor_end_ends_repeat id id path is_desc id_in_path ⟨0, Eq.refl _⟩
        exact id_is_repeat.elim (uf.no_cycles id)
      case right =>
        cases path
        case nil =>
          trivial
        case cons x xs =>
          let x_not_root : ¬uf.is_root x := by
            intro x_is_root
            rw [is_desc.left.symm] at not_root
            rw [x_is_root] at not_root
            exact not_root x_is_root
          exact uf.non_repeating_path x xs is_desc.right x_not_root

  def root_distance_aux {size : Nat} (uf : UnionFindLinks size) :
    (id : Fin size) → (n : Nat) →
    (∃ desc : List (Fin size),
      uf.is_descendency id desc ∧
      desc.length + n = size
    ) →
    {
      path_length : Nat //
      uf.rooted_in_exact_steps id path_length
    } :=
    λ id n desc_ex =>
    if hroot : uf.parent id = id then
      ⟨0, hroot,
        λ _ h _ =>
        Nat.not_lt_zero _ h
      ⟩
    else
      let parent := uf.parent id
      let next_n := n-1

      let n_ne_zero : n ≠ 0 := desc_ex.elim (
        λ desc ⟨h_desc, h_length⟩ =>

        let next_desc := id :: desc

        λ n_eq_zero =>
        let desc_length_size : desc.length = size := by rw [n_eq_zero] at h_length; exact h_length
        let next_desc_larger : next_desc.length = size + 1 := congrArg Nat.succ desc_length_size
        let next_desc_too_big : next_desc.length > size := by
          rw [next_desc_larger]
          exact Nat.lt_succ_self size
        let next_desc_unique : next_desc.unique := uf.non_repeating_path id desc h_desc hroot
        pigeonhole_principle next_desc_unique next_desc_too_big
      )

      let next_desc_ex :
        ∃ desc : List (Fin size),
          uf.is_descendency parent desc ∧
          desc.length + next_n = size := desc_ex.elim (
        λ desc ⟨h_desc, h_length⟩ =>

        let next_desc := id :: desc
        let next_is_desc : uf.is_descendency parent next_desc :=
          ⟨by rfl, h_desc⟩
        let next_right_length : next_desc.length + next_n = size := by
          rw [(Nat.succ_pred n_ne_zero).symm] at h_length
          let h_length : desc.length + (next_n + 1) = size := h_length
          rw [Nat.add_comm next_n 1, (Nat.add_assoc _ _ _).symm] at h_length
          exact h_length

        ⟨next_desc, next_is_desc, next_right_length⟩
      )

      --not needed apparently? I guess the n_ne_zero already does it?
      --let term : next_n < n := Nat.pred_lt n_ne_zero

      let ⟨parent_root_dist, root_in_dist_steps, not_less_steps⟩ := uf.root_distance_aux parent next_n next_desc_ex
      let our_root_dist := parent_root_dist + 1
      let our_root_dist_proof : uf.is_root (Nat.repeat uf.parent our_root_dist id) := root_in_dist_steps
      let our_first_dist_proof : ∀ (n : Nat),
        n < our_root_dist →
        ¬uf.is_root (Nat.repeat uf.parent n id) := by
        intro n hn
        cases n
        case zero =>
          exact hroot
        case succ n =>
          let hn := Nat.lt_of_succ_lt_succ hn
          exact not_less_steps n hn
      ⟨our_root_dist, our_root_dist_proof, our_first_dist_proof⟩

  def root_distance {size : Nat} (uf : UnionFindLinks size) (id : Fin size) : Nat :=
    (uf.root_distance_aux id size ⟨[], trivial, (by rw [Nat.add_comm]; rfl)⟩).val

  theorem root_distance_property {size : Nat} (uf : UnionFindLinks size) (id : Fin size) : uf.rooted_in_exact_steps id (uf.root_distance id) :=
    (uf.root_distance_aux id size ⟨[], trivial, (by rw [Nat.add_comm]; rfl)⟩).property

  theorem parent_nearer_root {size : Nat} (uf : UnionFindLinks size) :
    (id : Fin size) →
    (¬uf.is_root id) →
    uf.root_distance (uf.parent id) < uf.root_distance id :=
    λ id not_root =>
    if h : uf.root_distance (uf.parent id) < uf.root_distance id then
      h
    else False.elim (
      let root_distance_not_zero : uf.root_distance id ≠ 0 := by
        intro is_zero
        apply not_root
        let h := (uf.root_distance_property id).left
        rw [is_zero] at h
        exact h
      let h : uf.root_distance id ≤ uf.root_distance (uf.parent id)     := Nat.ge_of_not_lt h
      let h : uf.root_distance id < uf.root_distance (uf.parent id) + 1 := Nat.lt_succ_of_le h
      let h : uf.root_distance id - 1 < uf.root_distance (uf.parent id) := Nat.lt_of_succ_lt_succ (by
        let h' : Nat.succ (uf.root_distance id - 1) = uf.root_distance id := Nat.succ_pred root_distance_not_zero
        rw [h']
        exact h
      )

      (uf.root_distance_property (uf.parent id)).right (uf.root_distance id - 1) h (by
        let h := (uf.root_distance_property id).left
        rw [(Nat.succ_pred root_distance_not_zero).symm] at h
        exact h
      )
    )

  def set_to_root_raw {size : Nat} (uf : UnionFindLinks size) :
    Fin size → uf.roots → FinArray size :=
    λ set_loc new_root =>
    uf.val.set set_loc new_root.val

  def set_to_root_gives_rooted {size : Nat} (uf : UnionFindLinks size)
    (set_loc : Fin size) (new_root : uf.roots) (id : Fin size) :
    let new_FinArray := uf.set_to_root_raw set_loc new_root
    ∃ steps,
      let r := Nat.repeat new_FinArray.get steps id
      new_FinArray.get r = r :=

    let new_FinArray := uf.set_to_root_raw set_loc new_root
    if hchanged : id = set_loc then
      let get_eq_new_root : new_FinArray.get set_loc = new_root.val := uf.val.get_set set_loc new_root.val
      let get_eq_new_root' : new_FinArray.get id = new_root.val := by
        rw [hchanged.symm] at get_eq_new_root
        exact get_eq_new_root
      let get_is_a_root : new_FinArray.get (new_FinArray.get id) = new_FinArray.get id := by
        rw [get_eq_new_root']
        exact if h : set_loc = new_root.val then by
          rw [h] at get_eq_new_root
          exact get_eq_new_root
        else by
          let h' : uf.val.get new_root.val = new_FinArray.get new_root.val := uf.val.get_set_ne set_loc new_root.val new_root.val h
          rw [h'.symm]
          exact new_root.property
      ⟨1, get_is_a_root⟩
    else
      let same_get : uf.parent id = new_FinArray.get id := uf.val.get_set_ne set_loc id new_root.val (λ x => hchanged x.symm)
      if hroot : uf.parent id = id then
        ⟨0, by rw [same_get] at hroot; exact hroot⟩
      else
        --required for termination
        let _hterm : uf.root_distance (uf.parent id) < uf.root_distance id :=
          uf.parent_nearer_root id hroot

        (uf.set_to_root_gives_rooted set_loc new_root (uf.parent id)).elim (
          λ steps hsteps =>
          ⟨steps+1, by rw [same_get] at hsteps; exact hsteps⟩
        )
  termination_by set_to_root_gives_rooted uf set_loc new_root id => uf.root_distance id

  def set_to_root {size : Nat} (uf : UnionFindLinks size) :
    Fin size → uf.roots → UnionFindLinks size :=
    λ set_loc new_root =>
    let new_FinArray := uf.set_to_root_raw set_loc new_root
    let is_rooted := uf.set_to_root_gives_rooted set_loc new_root
    UnionFindLinks.from_rooted new_FinArray is_rooted

  def find_aux {size : Nat} (uf : UnionFindLinks size) (id : Fin size) :
    Σ (new_uf : UnionFindLinks size),
      {_root : new_uf.roots //
      (∀ a, uf.is_root a ↔ new_uf.is_root a)} :=
    if hroot : uf.parent id = id then (
      let uf_equiv_uf : ∀ (a : Fin size), uf.is_root a ↔ uf.is_root a :=
        λ a => ⟨λ b => b, λ b => b⟩
      ⟨uf, ⟨id, hroot⟩, uf_equiv_uf⟩
    ) else
      --termination chacking
      let _hterm : uf.root_distance (uf.parent id) < uf.root_distance id :=
        uf.parent_nearer_root id hroot

      let ⟨new_uf, root, old_equiv_new⟩ := uf.find_aux (uf.parent id)

      let out_uf : UnionFindLinks size :=
        new_uf.set_to_root id root

      let new_equiv_out : ∀ a, new_uf.is_root a ↔ out_uf.is_root a := by
        intro a
        constructor
        case mp =>
          intro a_is_new_root
          let id_ne_a : id ≠ a := by
            intro id_eq_a
            rw [id_eq_a] at hroot
            exact hroot ((old_equiv_new a).mpr a_is_new_root)
          let parent_unchanged : new_uf.parent a = out_uf.parent a := new_uf.val.get_set_ne id a root.val id_ne_a
          exact parent_unchanged.symm.trans a_is_new_root
        case mpr =>
          intro a_is_out_root
          by_cases id = a
          case inl id_eq_a =>
            let id_is_out_root : out_uf.parent id = id := by
              rw [id_eq_a]
              exact a_is_out_root
            let parent_is_root : out_uf.parent id = root.val := new_uf.val.get_set id root.val
            let id_eq_root : id = root.val := id_is_out_root.symm.trans parent_is_root
            rw [id_eq_root] at hroot
            let root_is_old_root : uf.is_root root.val := (old_equiv_new root.val).mpr root.property
            let _ := hroot root_is_old_root
            contradiction
          case inr id_ne_a =>
            let parent_unchanged : new_uf.parent a = out_uf.parent a := new_uf.val.get_set_ne id a root.val id_ne_a
            exact parent_unchanged.trans a_is_out_root

      let root : out_uf.roots := ⟨root.val, (new_equiv_out root.val).mp root.property⟩

      let old_equiv_out : ∀ a, uf.is_root a ↔ out_uf.is_root a :=
        λ a => (old_equiv_new a).trans (new_equiv_out a)

      ⟨out_uf, root, old_equiv_out⟩
  termination_by find_aux uf id => uf.root_distance id

  def find {size:Nat} : Fin size → StateM (UnionFindLinks size) (Fin size)
  | id, uf => (
    let ⟨new_uf, root, _⟩ := uf.find_aux id
    ⟨root.val, new_uf⟩
  )
end UnionFindLinks

--while ranks should store the number of children for each root,
--this isn't really an important invariant or anything, and is only
--there to help the program run faster, so no stored proof
--required for lean to accept.
structure UnionFind (size : Nat) : Type where
  (links : UnionFindLinks size)
  (ranks : {xs : Array Nat // xs.size = size})

namespace UnionFind
  def union {size : Nat} : Fin size → Fin size → StateM (UnionFind size) (Fin size) :=
    λ id1 id2 uf =>
    let uf_links := uf.links
    let ⟨uf_links, root1, _    ⟩ := uf_links.find_aux id1
    let ⟨uf_links, root2, equiv⟩ := uf_links.find_aux id2

    let root1 : uf_links.roots := ⟨root1.val, (equiv root1.val).mp root1.property⟩

    let rank1 := uf.ranks.val.get ⟨
      root1.val,
      by rw [uf.ranks.property]; exact root1.val.isLt
    ⟩

    let rank2 := uf.ranks.val.get ⟨
      root2.val,
      by rw [uf.ranks.property]; exact root2.val.isLt
    ⟩

    let ⟨bigroot, smallroot⟩ :=
      if rank1 < rank2 then
        (root2, root1)
      else
        (root1, root2)

    let uf_links := uf_links.set_to_root smallroot.val bigroot

    let bigroot_array_loc : Fin uf.ranks.val.size := ⟨
      bigroot.val.val,
      by rw [uf.ranks.property]; exact bigroot.val.isLt
    ⟩

    let new_ranks := uf.ranks.val.set bigroot_array_loc (rank1 + rank2)
    let same_size := uf.ranks.val.size_set bigroot_array_loc (rank1 + rank2)
    let expected_size : new_ranks.size = size := same_size.trans uf.ranks.property

    ⟨bigroot.val, uf_links, new_ranks, expected_size⟩

  def find {size : Nat} : Fin size → StateM (UnionFind size) (Fin size) :=
    λ id uf =>
    let ⟨root, uf_links⟩ := do {
      UnionFindLinks.find id
    }.run uf.links
    ⟨root, uf_links, uf.ranks⟩
end UnionFind