import challenge

/-!

The main goal of this file is to discuss the definition of
condensed abelian groups which is used in this project.

As prerequisites, we also discuss continuous maps and the
category of profinite sets.

-/

open category_theory category_theory.limits opposite
open_locale liquid_tensor_experiment

section continuous_maps

/-!
This section describes the type `C(X,Y)` of continuous maps from `X` to `Y`.
-/

/-!
Let `X` and `Y` be topological spaces.
-/
variables {X Y : Type*} [topological_space X] [topological_space Y]

/-!
Any `f : C(X,Y)` can be thought of as a function from `X` to `Y`.
-/
example (f : C(X,Y)) : X → Y :=
f

/-!
A term `f` of `C(X,Y)` is continuous, when considered as a function `X → Y`.
-/
example (f : C(X,Y)) : continuous f :=
f.continuous

/-!
Conversely, any continuous function `X → Y` yields an element of `C(X,Y)`.
-/
example (f : X → Y) (hf : continuous f) : C(X,Y) :=
⟨f,hf⟩

end continuous_maps

section profinite_sets

/-!
In this section, we discuss the category of profinite sets.
-/

/-!
`Profinite.{0}` denotes the type of profinite sets whose underlying type
lives in `Type = Type 0`.
-/
example : Type 1 := Profinite.{0}
example (X : Profinite.{0}) : Type := X

/-!
Any profinite set is a compact, Hausdorff, totally disconnected topological space.
-/
example (X : Profinite.{0}) : topological_space X := infer_instance
example (X : Profinite.{0}) : compact_space X := infer_instance
example (X : Profinite.{0}) : t2_space X := infer_instance
example (X : Profinite.{0}) : totally_disconnected_space X := infer_instance

/-!
Conversely, any topological space which is compact, Hausdorff and
totally disconnected is a profinite set.
The function `Profinite.of` is used to make an object of `Profinite.{0}`
from such a topological space.
-/
example (X : Type)
  [topological_space X]
  [compact_space X]
  [t2_space X]
  [totally_disconnected_space X] :
  Profinite.{0} :=
Profinite.of X

/-!
`Profinite.{0}` has a structure of a category, where morphisms are, by definition,
continuous maps.
-/
example (X Y : Profinite.{0}) : (X ⟶ Y : Type) = C(X,Y) :=
rfl

end profinite_sets

section proetale_topology

/-!
This section describes the Grothendieck topology on `Profinite.{0}` which is used in
our formalization of condensed mathematics.
We use the name `proetale_topology` for this Grothendieck topology.
-/
example : grothendieck_topology Profinite.{0} := proetale_topology

/-!
This example shows that the sheaf condition with respect to this Grothendieck topology,
for a presheaf of abelian groups on `Profinite.{0}`, is equivalent to the usual definition.
-/
example
-- Let `F` be a presheaf of abelian groups on `Profinite.{0}`.
  (F : Profinite.{0}ᵒᵖ ⥤ Ab.{1}) :
-- `F` is a sheaf for `proetale_topology`
  presheaf.is_sheaf proetale_topology F
-- if and only if the following condition holds:
  ↔
-- For any finite indexing type `α`,
  ∀ (α : Fintype.{0})
-- profinite set `B`,
    (B : Profinite.{0})
-- family of profinite sets `X` indexed by `α`
    (X : α → Profinite.{0})
-- which map to `B` with a family of maps `π`,
    (π : Π i, X i ⟶ B)
-- such that `π` is jointly surjective,
    (hπ : ∀ b : B, ∃ i (x : X i), π i x = b)
-- and all families of elements `x i : F (op (X i))`,
    (x : Π i, F (op (X i)))
-- which are compatible on pullbacks `X i ×_{B} X j`
    (hx : ∀ i j : α,
      F.map (pullback.fst : pullback (π i) (π j) ⟶ X i).op (x i) =
      F.map (pullback.snd : pullback (π i) (π j) ⟶ X j).op (x j)),
-- there is a unique `s : F (op B)`
    ∃! s : F (op B),
-- which restricts to `x i` over `X i` for all `i`.
      ∀ i, F.map (π i).op s = x i :=
begin
  rw presheaf.is_sheaf_iff_is_sheaf_forget proetale_topology F (forget _),
  rw [is_sheaf_iff_is_sheaf_of_type, (F ⋙ forget Ab).is_proetale_sheaf_of_types_tfae.out 0 2],
  split,
  { intros H α, exact H _ },
  { introsI H α _, exact H ⟨α⟩ }
end

/-!
The category of condensed abelian groups is denoted by `Condensed.{0} Ab.{1}`.
As a type, `Condensed.{0} Ab.{1}` is defined to be the type of sheaves from `mathlib`.
-/
example : Condensed.{0} Ab.{1} = Sheaf proetale_topology.{0} Ab.{1} :=
rfl

/-!
Any object of `Condensed.{0} Ab.{1}` yields a presheaf of abelian groups.
-/
example (F : Condensed.{0} Ab.{1}) : Profinite.{0}ᵒᵖ ⥤ Ab.{1} :=
F.1

/-!
The presheaf in the example above is indeed a sheaf.
-/
example (F : Condensed.{0} Ab.{1}) : presheaf.is_sheaf proetale_topology F.1 :=
F.2

/-!
Conversely, any presheaf of abelian groups on `Profinite.{0}`
(with the correct universe parameters) which is a sheaf for `proetale_topology`
yields an object of `Condensed.{0} Ab.{1}`.
-/
example (F : Profinite.{0}ᵒᵖ ⥤ Ab.{1}) (hF : presheaf.is_sheaf proetale_topology F) :
  Condensed.{0} Ab.{1} :=
⟨F,hF⟩

end proetale_topology


import tactic
import tactic.induction
import logic.function.iterate

noncomputable theory
open_locale classical

inductive Expr : Type
| K : Expr
| S : Expr
| M : Expr
| app : Expr → Expr → Expr
open Expr

infixl ` ~ `:100 := app

instance : inhabited Expr := ⟨K⟩

def step : Expr → Expr
| (K ~ a ~ b) := a
| (S ~ a ~ b ~ c) := a ~ c ~ (b ~ c)
| (M ~ K) := K
| (M ~ S) := S
| (M ~ a) := M ~ step a
| (a ~ b) := step a ~ b
| a := a

def reduces (e₁ e₂ : Expr) : Prop :=
∃ (n : ℕ), (step^[n]) e₁ = e₂

local infix ` > `:50 := reduces

-----

@[refl]
lemma reduces.refl {a} : a > a :=
by use [0, rfl]

@[trans]
lemma reduces.trans {a b c} (h₁ : a > b) (h₂ : b > c) : a > c :=
begin
  cases h₁ with n₁ h₁, cases h₂ with n₂ h₂, use n₂ + n₁,
  rw function.iterate_add_apply, substs h₁ h₂,
end

lemma not_reduces {a b} : ¬a > b ↔ ∀ (n : ℕ), (step^[n]) a ≠ b :=
by simp [reduces]

-----

def Expr.comb : Expr → Prop
| (_ ~ _) := false
| _ := true

lemma term_induct {P : Expr → Prop} {a : Expr}
  (h₁ : a.comb) (h₂ : P K) (h₃ : P S) (h₄ : P M) : P a :=
by cases a; simp [*, Expr.comb] at *

lemma reduces_iff_eq_of_step_eq_self {a b}
  (h : step a = a) : a > b ↔ a = b :=
begin
  split; intro h₁,
  { contrapose h₁, rw not_reduces, intro n, induction n with n ih,
    { tauto },
    { simp [h, ih] } },
  { subst h₁ },
end

lemma step_self_of_term {a : Expr} (h : a.comb) : step a = a :=
by apply term_induct h; refl

lemma reduces_iff_of_term {a b : Expr}
  (h : a.comb) : a > b ↔ a = b :=
reduces_iff_eq_of_step_eq_self (step_self_of_term h)

-----

lemma apps_reduce_iff_args_reduce {t f : Expr} {args : list (Expr × Expr)}
  (h₁ : t.comb)
  (h₂ : ∀ (p : Expr × Expr), p ∈ args → p.1 > p.2) :
  function.uncurry (λ a b, a > t ↔ b > t)
    (list.foldl (λ (a p : Expr × Expr), (a.1 ~ p.1, a.2 ~ p.2)) ⟨f, f⟩ args) :=
begin
  let m : Expr × Expr := _,
  change function.uncurry _ m,
  change _ ↔ _,
  simp_rw reduces,
  sorry
end

lemma app_reduces_iff_arg_reduces {t f a b : Expr}
  (h₁ : t.comb)
  (h₂ : a > b) :
  f ~ a > t ↔ f ~ b > t :=
begin
  apply @apps_reduce_iff_args_reduce t f [(a, b)] h₁,
  rintro p h₄, rw list.mem_singleton at h₄, subst p, exact h₂,
end

-----

def I := S ~ K ~ K

lemma I_id {a} : I ~ a > a :=
⟨2, rfl⟩

-----

def Expr.mk_diff (a : Expr) : Expr :=
K ~ a

lemma mk_diff_ne {a : Expr} : a.mk_diff ≠ a :=
by { rw Expr.mk_diff, induction' a; simp, rintro rfl, assumption }

-----

-- def func_coe_raw

-- inductive Func : (Expr → Expr) → Expr → Prop
-- | id : Func id I
-- | comb {a : Expr} : a.comb → Func (λ _, a) (K ~ a)
-- | app {f₁ f₂ : Expr → Expr} {e₁ e₂ : Expr} :
--   Func f₁ e₁ → Func f₂ e₂ → Func (λ a, f₁ a ~ f₂ a) (S ~ e₁ ~ e₂)

-- def func_coe_raw (f : Expr → Expr) : Expr :=
-- if h : ∃ (e : Expr), Func f e then h.some else default

-- lemma exi_func_coe : ∃ (f : (Expr → Expr) → Expr),
--   (f id = I) ∧
--   (∀ (a : Expr), a.comb → f (λ _, a) = K ~ a) ∧
--   (∀ (g₁ g₂ : Expr → Expr), f (λ a, g₁ a ~ g₂ a) = S ~ f g₁ ~ f g₂) :=
-- begin
--   refine ⟨func_coe_raw, _, _, _⟩; simp_rw func_coe_raw,
--   sorry {
--     rw dif_pos, swap,
--     {
--       use I,
--       exact Func.id,
--     },
--     generalize_proofs h,
--     have h₁ := h.some_spec,
--     induction' h₁,
--     {
--       refl,
--     },
--     {
--       cases mk_diff_ne (congr_fun induction_eq h.some.mk_diff).symm,
--     },
--     {
--       cases congr_fun induction_eq K,
--     },
--   },
--   sorry {
--     rintro a h,
--     rw dif_pos, swap,
--     {
--       use K ~ a,
--       exact Func.comb h,
--     },
--     generalize_proofs h₁,
--     have h₂ := h₁.some_spec,
--     induction' h₂,
--     {
--       cases congr_fun induction_eq (K ~ K),
--       cases h,
--     },
--     {
--       cases congr_fun induction_eq K,
--       refl,
--     },
--     {
--       cases congr_fun induction_eq K,
--       cases h,
--     },
--   },
--   {
--     rintro g₁ g₂,
--     rw dif_pos,
--   },
-- end

-----


import tableau order.lexicographic

open matrix fintype finset function pequiv partition
variables {m n : ℕ}

local notation `rvec`:2000 n := matrix (fin 1) (fin n) ℚ
local notation `cvec`:2000 m := matrix (fin m) (fin 1) ℚ
local infix ` ⬝ `:70 := matrix.mul
local postfix `ᵀ` : 1500 := transpose

namespace tableau

def pivot_col (T : tableau m n) (obj : fin m) : option (fin n) :=
option.cases_on
  (fin.find (λ c : fin n, T.to_matrix obj c ≠ 0 ∧ T.to_partition.colg c ∉ T.restricted
    ∧ c ∉ T.dead))
  (((list.fin_range n).filter (λ c : fin n, 0 < T.to_matrix obj c ∧ c ∉ T.dead)).argmin
    T.to_partition.colg)
  some

def to_lex (T : tableau m n) (c : fin n) (r' : fin m) : lex ℚ (fin (m + n)) :=
(abs (T.const r' 0 / T.to_matrix r' c), T.to_partition.rowg r')

lemma to_lex_le_iff (T : tableau m n) (c : fin n) (i i' : fin m) :
  to_lex T c i ≤ to_lex T c i' ↔
  abs (T.const i 0 / T.to_matrix i c) < abs (T.const i' 0 / T.to_matrix i' c) ∨
    (abs (T.const i 0 / T.to_matrix i c) = abs (T.const i' 0 / T.to_matrix i' c) ∧
    T.to_partition.rowg i ≤ T.to_partition.rowg i') :=
prod.lex_def _ _

def pivot_row (T : tableau m n) (obj: fin m) (c : fin n) : option (fin m) :=
let l := (list.fin_range m).filter (λ i : fin m, obj ≠ i ∧ T.to_partition.rowg i ∈ T.restricted
  ∧ T.to_matrix obj c / T.to_matrix i c < 0) in
list.argmin (to_lex T c) l

lemma pivot_col_spec {T : tableau m n} {obj : fin m} {c : fin n} :
  c ∈ pivot_col T obj → ((T.to_matrix obj c ≠ 0 ∧ T.to_partition.colg c ∉ T.restricted)
  ∨ (0 < T.to_matrix obj c ∧ T.to_partition.colg c ∈ T.restricted)) ∧ c ∉ T.dead :=
begin
  dsimp only [pivot_col],
  cases h : fin.find (λ c : fin n, T.to_matrix obj c ≠ 0 ∧ T.to_partition.colg c ∉ T.restricted
    ∧ c ∉ T.dead),
  { have := not_and.1 (fin.find_eq_none_iff.1 h c) ∘ ne_of_gt,
    simp only [list.argmin_eq_some_iff, option.mem_def, list.mem_filter, list.mem_fin_range, gt]
      at * {contextual := tt},
    tauto },
  { finish [fin.find_eq_some_iff] }
end

lemma nonpos_of_lt_pivot_col {T : tableau m n} {obj : fin m} {c j : fin n}
  (hc : c ∈ pivot_col T obj) (hcres : T.to_partition.colg c ∈ T.restricted)
  (hdead : j ∉ T.dead) (hjc : T.to_partition.colg j < T.to_partition.colg c) :
  T.to_matrix obj j ≤ 0 :=
begin
  rw [pivot_col] at hc,
  cases h : fin.find (λ c, T.to_matrix obj c ≠ 0 ∧ colg (T.to_partition) c ∉ T.restricted
    ∧ c ∉ T.dead),
  { rw h at hc,
    refine le_of_not_lt (λ hj0, _),
    exact not_le_of_gt hjc ((list.mem_argmin_iff.1 hc).2.1 j
      (list.mem_filter.2 (by simp [hj0, hdead]))) },
  { rw h at hc,
    simp [*, fin.find_eq_some_iff] at * }
end

lemma pivot_col_eq_none_aux {T : tableau m n} {obj : fin m} (hT : T.feasible) {c : fin n} :
  pivot_col T obj = none → c ∉ T.dead →
  ((T.to_matrix obj c = 0 ∧ T.to_partition.colg c ∉ T.restricted)
    ∨ (T.to_matrix obj c ≤ 0 ∧ T.to_partition.colg c ∈ T.restricted)) :=
begin
  simp only [pivot_col],
  cases h : fin.find (λ c : fin n, T.to_matrix obj c ≠ 0 ∧ T.to_partition.colg c ∉ T.restricted
    ∧ c ∉ T.dead),
  { simp only [list.filter_eq_nil, list.argmin_eq_none, not_and', list.mem_fin_range,
      true_implies_iff, not_lt, fin.find_eq_none_iff, and_imp, not_not] at *,
    assume hnonneg hdead,
    by_cases hres : T.to_partition.colg c ∈ T.restricted; simp * at * },
  { simp }
end

lemma pivot_col_eq_none {T : tableau m n} {obj : fin m} (hT : T.feasible)
  (h : pivot_col T obj = none) : T.is_optimal (T.of_col 0) (T.to_partition.rowg obj) :=
is_optimal_of_col_zero hT
(λ j hj, begin
  have := pivot_col_eq_none_aux hT h hj,
  finish [lt_irrefl]
end)

lemma pivot_row_spec {T : tableau m n} {obj i : fin m} {c : fin n} :
  i ∈ pivot_row T obj c →
  obj ≠ i ∧ T.to_partition.rowg i ∈ T.restricted ∧
  T.to_matrix obj c / T.to_matrix i c < 0 ∧
  (∀ i' : fin m, obj ≠ i' → T.to_partition.rowg i' ∈ T.restricted →
    T.to_matrix obj c / T.to_matrix i' c < 0 →
  abs (T.const i 0 / T.to_matrix i c) ≤ abs (T.const i' 0 / T.to_matrix i' c)) :=
begin
  simp only [list.mem_filter, pivot_row, option.mem_def,
    list.argmin_eq_some_iff, list.mem_fin_range, true_and, and_imp],
  simp only [to_lex_le_iff],
  intros hor hres hr0 h _,
  simp only [*, true_and, ne.def, not_false_iff],
  intros i' hoi' hres' hi0',
  cases h i' hoi' hres' hi0',
  { exact le_of_lt (by assumption) },
  { exact le_of_eq (by tauto) }
end

lemma nonneg_of_lt_pivot_row {T : tableau m n} {obj : fin m} {i i' : fin m} {c : fin n}
  (hc0 : 0 < T.to_matrix obj c) (hres : T.to_partition.rowg i' ∈ T.restricted)
  (hc : c ∈ pivot_col T obj) (hrow : i ∈ pivot_row T obj c)
  (hconst : T.const i' 0 = 0)
  (hjc : T.to_partition.rowg i' < T.to_partition.rowg i) :
  0 ≤ T.to_matrix i' c :=
if hobj : obj = i' then le_of_lt $ hobj ▸ hc0
else
le_of_not_gt $ λ hic, not_le_of_lt hjc
begin
  have := (list.argmin_eq_some_iff.1 hrow).2.1 i'
    (list.mem_filter.2 ⟨list.mem_fin_range _, hobj, hres, div_neg_of_pos_of_neg hc0 hic⟩),
  simp [hconst, not_lt_of_ge (abs_nonneg _), *, to_lex_le_iff] at *
end

lemma ne_zero_of_mem_pivot_row {T : tableau m n} {obj i : fin m} {c : fin n}
  (hrow : i ∈ pivot_row T obj c) : T.to_matrix i c ≠ 0 :=
assume hrc, by simpa [lt_irrefl, hrc] using pivot_row_spec hrow

lemma ne_zero_of_mem_pivot_col {T : tableau m n} {obj : fin m} {c : fin n}
  (hc : c ∈ pivot_col T obj) : T.to_matrix obj c ≠ 0 :=
λ h, by simpa [h, lt_irrefl] using pivot_col_spec hc

lemma pivot_row_eq_none_aux {T : tableau m n} {obj : fin m} {c : fin n}
  (hrow : pivot_row T obj c = none) (hs : c ∈ pivot_col T obj) :
  ∀ i, obj ≠ i → T.to_partition.rowg i ∈ T.restricted → 0 ≤ T.to_matrix obj c / T.to_matrix i c :=
by simpa [pivot_row, list.filter_eq_nil] using hrow

lemma pivot_row_eq_none {T : tableau m n} {obj : fin m} {c : fin n} (hT : T.feasible)
  (hrow : pivot_row T obj c = none) (hs : c ∈ pivot_col T obj) :
  T.is_unbounded_above (T.to_partition.rowg obj) :=
have hrow : ∀ i, obj ≠ i → T.to_partition.rowg i ∈ T.restricted →
    0 ≤ T.to_matrix obj c / T.to_matrix i c,
  from pivot_row_eq_none_aux hrow hs,
have hc : ((T.to_matrix obj c ≠ 0 ∧ T.to_partition.colg c ∉ T.restricted)
    ∨ (0 < T.to_matrix obj c ∧ T.to_partition.colg c ∈ T.restricted)) ∧ c ∉ T.dead,
  from pivot_col_spec hs,
have hToc : T.to_matrix obj c ≠ 0, from λ h, by simpa [h, lt_irrefl] using hc,
(lt_or_gt_of_ne hToc).elim
  (λ hToc : T.to_matrix obj c < 0, is_unbounded_above_rowg_of_nonpos hT c
    (hc.1.elim and.right (λ h, (not_lt_of_gt hToc h.1).elim)) hc.2
    (λ i hi, classical.by_cases
      (λ hoi : obj = i, le_of_lt (hoi ▸ hToc))
      (λ hoi : obj ≠ i, inv_nonpos.1 $ nonpos_of_mul_nonneg_right (hrow _ hoi hi) hToc))
    hToc)
  (λ hToc : 0 < T.to_matrix obj c, is_unbounded_above_rowg_of_nonneg hT c
    (λ i hi, classical.by_cases
      (λ hoi : obj = i, le_of_lt (hoi ▸ hToc))
      (λ hoi : obj ≠ i, inv_nonneg.1 $ nonneg_of_mul_nonneg_left (hrow _ hoi hi) hToc))
    hc.2 hToc)

def feasible_of_mem_pivot_row_and_col {T : tableau m n} {obj : fin m} (hT : T.feasible) {c}
  (hc : c ∈ pivot_col T obj) {i} (hr : i ∈ pivot_row T obj c) :
  feasible (T.pivot i c) :=
begin
  have := pivot_col_spec hc,
  have := pivot_row_spec hr,
  have := @feasible_simplex_pivot _ _ _ obj hT i c,
  tauto
end

section blands_rule

local attribute [instance, priority 0] classical.dec
variable (obj : fin m)

def fickle (T T' : tableau m n) (v : fin (m + n)) : Prop :=
T.to_partition.rowp v ≠ T'.to_partition.rowp v ∨
T.to_partition.colp v ≠ T'.to_partition.colp v

lemma fickle_symm {T T' : tableau m n} {v : fin (m + n)} :
  fickle T T' v ↔ fickle T' T v :=
by simp [fickle, eq_comm]

lemma fickle_colg_iff_ne {T T' : tableau m n} {j : fin n} :
  fickle T T' (T.to_partition.colg j) ↔
  T.to_partition.colg j ≠ T'.to_partition.colg j :=
⟨λ h h', h.elim (by rw [rowp_colg_eq_none, h', rowp_colg_eq_none]; simp)
  (by rw [colp_colg, h', colp_colg]; simp),
λ h, or.inr $ λ h', h begin
  have : T'.to_partition.colp (T.to_partition.colg j) = some j,
  { simpa [eq_comm] using h' },
  rwa [← pequiv.eq_some_iff, colp_symm_eq_some_colg, option.some_inj, eq_comm] at this
end⟩

lemma fickle_rowg_iff_ne {T T' : tableau m n} {i : fin m} :
  fickle T T' (T.to_partition.rowg i) ↔
  T.to_partition.rowg i ≠ T'.to_partition.rowg i :=
⟨λ h h', h.elim (by rw [rowp_rowg, h', rowp_rowg]; simp)
  (by rw [colp_rowg_eq_none, h', colp_rowg_eq_none]; simp),
λ h, or.inl $ λ h', h begin
  have : T'.to_partition.rowp (T.to_partition.rowg i) = some i,
  { simpa [eq_comm] using h' },
  rwa [← pequiv.eq_some_iff, rowp_symm_eq_some_rowg, option.some_inj, eq_comm] at this
end⟩

lemma not_unique_row_and_unique_col {T T' : tableau m n} {i c c'}
  (hcobj0 : 0 < T.to_matrix obj c)
  (hc'obj0 : 0 < T'.to_matrix obj c')
  (hrc0 : T.to_matrix i c < 0)
  (hflat : T.flat = T'.flat)
  (hs : T.to_partition.rowg i = T'.to_partition.colg c')
  (hrobj : T.to_partition.rowg obj = T'.to_partition.rowg obj)
  (hfickle : ∀ i, (fickle T T' (T.to_partition.rowg i)) → T.const i 0 = 0)
  (hobj : T.const obj 0 = T'.const obj 0)
  (nonpos_of_colg_eq : ∀ j, j ≠ c' →
    T'.to_partition.colg j = T.to_partition.colg c → T'.to_matrix obj j ≤ 0)
  (unique_col : ∀ j,
    (fickle T' T (T'.to_partition.colg j)) → j ≠ c' → T'.to_matrix obj j ≤ 0)
  (unique_row : ∀ i' ≠ i, T.const i' 0 = 0 → fickle T T' (T.to_partition.rowg i') →
    0 ≤ T.to_matrix i' c) :
  false :=
let objr := T.to_partition.rowg obj in
let x := λ y : ℚ, T.of_col (y • (single c 0).to_matrix) in
have hxflatT' : ∀ {y}, x y ∈ flat T', from hflat ▸ λ _, of_col_mem_flat _ _,
have hxrow : ∀ y i, x y (T.to_partition.rowg i) 0 = T.const i 0 + y * T.to_matrix i c,
  by simp [x, of_col_single_rowg],
have hxcol : ∀ {y j}, j ≠ c → x y (T.to_partition.colg j) 0 = 0,
  from λ y j hjc, by simp [x, of_col_colg, pequiv.to_matrix, single_apply_of_ne hjc.symm],
have hxcolc : ∀ {y}, x y (T.to_partition.colg c) 0 = y, by simp [x, of_col_colg, pequiv.to_matrix],
let c_star : fin (m + n) → ℚ := λ v, option.cases_on (T'.to_partition.colp v) 0
  (T'.to_matrix obj) in
have hxobj : ∀ y, x y objr 0 = T.const obj 0 + y * T.to_matrix obj c, from λ y, hxrow _ _,
have hgetr : ∀ {y v}, c_star v * x y v 0 ≠ 0 → (T'.to_partition.colp v).is_some,
  from λ y v, by cases h : T'.to_partition.colp v; dsimp [c_star]; rw h; simp,
have c_star_eq_get : ∀ {v} (hv : (T'.to_partition.colp v).is_some),
    c_star v = T'.to_matrix obj (option.get hv),
  from λ v hv, by dsimp only [c_star]; conv_lhs{rw [← option.some_get hv]}; refl,
have hsummmn : ∀ {y}, sum univ (λ j, T'.to_matrix obj j * x y (T'.to_partition.colg j) 0) =
    sum univ (λ v, c_star v * x y v 0),
  from λ y, sum_bij_ne_zero (λ j _ _, T'.to_partition.colg j) (λ _ _ _, mem_univ _)
    (λ _ _ _ _ _ _ h, T'.to_partition.injective_colg h)
    (λ v _ h0, ⟨option.get (hgetr h0), mem_univ _,
      by rw [← c_star_eq_get (hgetr h0)]; simpa using h0, by simp⟩)
    (λ _ _ h0, by dsimp [c_star]; rw [colp_colg]),
have hgetc : ∀ {y v}, c_star v * x y v 0 ≠ 0 → v ≠ T.to_partition.colg c →
    (T.to_partition.rowp v).is_some,
  from λ y v, (eq_rowg_or_colg T.to_partition v).elim
    (λ ⟨i, hi⟩, by rw [hi, rowp_rowg]; simp)
    (λ ⟨j, hj⟩ h0 hvc,
      by rw [hj, hxcol (mt (congr_arg T.to_partition.colg) (hvc ∘ hj.trans)), mul_zero] at h0;
        exact (h0 rfl).elim),
have hsummmnn : ∀ {y}, (univ.erase (T.to_partition.colg c)).sum (λ v, c_star v * x y v 0) =
    univ.sum (λ i, c_star (T.to_partition.rowg i) * x y (T.to_partition.rowg i) 0),
  from λ y, eq.symm $ sum_bij_ne_zero (λ i _ _, T.to_partition.rowg i) (by simp)
    (λ _ _ _ _ _ _ h, T.to_partition.injective_rowg h)
    (λ v hvc h0, ⟨option.get (hgetc h0 (mem_erase.1 hvc).1), mem_univ _, by simpa using h0⟩)
    (by intros; refl),
have hsumm : ∀ {y}, univ.sum (λ i, c_star (T.to_partition.rowg i) * x y (T.to_partition.rowg i) 0) =
    univ.sum (λ i, c_star (T.to_partition.rowg i) * T.const i 0) +
    y * univ.sum (λ i, c_star (T.to_partition.rowg i) * T.to_matrix i c),
  from λ y, by simp only [hxrow, mul_add, add_mul, sum_add_distrib, mul_assoc,
    mul_left_comm _ y, mul_sum.symm],
have hxobj' : ∀ y, x y objr 0 = univ.sum (λ v, c_star v * x y v 0) + T'.const obj 0,
  from λ y, by dsimp [objr]; rw [hrobj, mem_flat_iff.1 hxflatT', hsummmn],
have hy : ∀ {y}, y * T.to_matrix obj c = c_star (T.to_partition.colg c) * y +
    univ.sum (λ i, c_star (T.to_partition.rowg i) * T.const i 0) +
      y * univ.sum (λ i, c_star (T.to_partition.rowg i) * T.to_matrix i c),
  from λ y, by rw [← add_left_inj (T.const obj 0), ← hxobj, hxobj',
    ← insert_erase (mem_univ (T.to_partition.colg c)), sum_insert (not_mem_erase _ _),
    hsummmnn, hobj, hsumm, hxcolc]; simp,
have hy' : ∀ (y), y * (T.to_matrix obj c - c_star (T.to_partition.colg c) -
    univ.sum (λ i, c_star (T.to_partition.rowg i) * T.to_matrix i c)) =
    univ.sum (λ i, c_star (T.to_partition.rowg i) * T.const i 0),
  from λ y, by rw [mul_sub, mul_sub, hy]; simp [mul_comm, mul_assoc, mul_left_comm],
have h0 : T.to_matrix obj c - c_star (T.to_partition.colg c) -
    univ.sum (λ i, c_star (T.to_partition.rowg i) * T.to_matrix i c) = 0,
  by rw [← (domain.mul_left_inj (@one_ne_zero ℚ _)), hy', ← hy' 0, zero_mul, mul_zero],
have hcolnec' : T'.to_partition.colp (T.to_partition.colg c) ≠ some c',
  from λ h,
    by simpa [hs.symm] using congr_arg T'.to_partition.colg (option.eq_some_iff_get_eq.1 h).snd,
have eq_of_roweqc' : ∀ {i'}, T'.to_partition.colp (T.to_partition.rowg i') = some c' → i' = i,
  from λ i' h, by simpa [hs.symm, T.to_partition.injective_rowg.eq_iff] using
    congr_arg T'.to_partition.colg (option.eq_some_iff_get_eq.1 h).snd,
have sumpos : 0 < univ.sum (λ i, c_star (T.to_partition.rowg i) * T.to_matrix i c),
  by rw [← sub_eq_zero.1 h0]; exact add_pos_of_pos_of_nonneg hcobj0
    (begin
      simp only [c_star, neg_nonneg],
      cases h : T'.to_partition.colp (T.to_partition.colg c) with j,
      { refl },
      { exact nonpos_of_colg_eq j (mt (congr_arg some) (h ▸ hcolnec'))
          (by rw [← (option.eq_some_iff_get_eq.1 h).snd]; simp) }
    end),
have hexi : ∃ i', 0 < c_star (T.to_partition.rowg i') * T.to_matrix i' c,
  from imp_of_not_imp_not _ _ (by simpa using @sum_nonpos _ _ (@univ (fin m) _)
    (λ i', c_star (T.to_partition.rowg i') * T.to_matrix i' c) _ _) sumpos,
let ⟨i', hi'⟩ := hexi in
have hi'0 : T.const i' 0 = 0, from hfickle i'
  (fickle_rowg_iff_ne.2 $
    λ h, by dsimp [c_star] at hi'; rw [h, colp_rowg_eq_none] at hi'; simpa [lt_irrefl] using hi'),
have hi'_some : (T'.to_partition.colp (T.to_partition.rowg i')).is_some,
  from option.ne_none_iff_is_some.1 (λ h, by dsimp only [c_star] at hi'; rw h at hi';
    simpa [lt_irrefl] using hi'),
have hi' : 0 < T'.to_matrix obj (option.get hi'_some) * T.to_matrix i' c,
  by dsimp only [c_star] at hi'; rwa [← option.some_get hi'_some] at hi',
have hii : i' ≠ i, from λ hir, begin
    have : option.get hi'_some = c', from T'.to_partition.injective_colg
      (by rw [colg_get_colp_symm, ← hs, hir]),
    rw [this, hir] at hi',
    exact not_lt_of_gt hi' (mul_neg_of_pos_of_neg hc'obj0 hrc0)
  end,
have hnec' : option.get hi'_some ≠ c',
  from λ eq_c', hii $ @eq_of_roweqc' i' (eq_c' ▸ by simp),
have hic0 : T.to_matrix i' c < 0,
  from neg_of_mul_pos_right hi' (unique_col _
    (by rw [fickle_colg_iff_ne]; simp) hnec'),
not_le_of_gt hic0 (unique_row _ hii hi'0
  (by rw [fickle_rowg_iff_ne, ← colg_get_colp_symm _ _ hi'_some]; exact colg_ne_rowg _ _ _))

inductive rel : tableau m n → tableau m n → Prop
| pivot : ∀ {T}, feasible T → ∀ {i c}, c ∈ pivot_col T obj →
  i ∈ pivot_row T obj c → rel (T.pivot i c) T
| trans_pivot : ∀ {T₁ T₂ i c}, rel T₁ T₂ → c ∈ pivot_col T₁ obj →
  i ∈ pivot_row T₁ obj c → rel (T₁.pivot i c) T₂

lemma feasible_of_rel_right {T T' : tableau m n} (h : rel obj T' T) : T.feasible :=
rel.rec_on h (by tauto) (by tauto)

lemma feasible_of_rel_left {T T' : tableau m n} (h : rel obj T' T) : T'.feasible :=
rel.rec_on h (λ _ hT _ _ hc hr, feasible_of_mem_pivot_row_and_col hT hc hr)
  (λ _ _ _ _ _ hc hr hT, feasible_of_mem_pivot_row_and_col hT hc hr)

/-- Slightly stronger recursor than the default recursor -/
@[elab_as_eliminator]
lemma rel.rec_on' {obj : fin m} {C : tableau m n → tableau m n → Prop} {T T' : tableau m n}
  (hrel : rel obj T T')
  (hpivot : ∀ {T : tableau m n} {i : fin m} {c : fin n},
     feasible T → c ∈ pivot_col T obj → i ∈ pivot_row T obj c → C (pivot T i c) T)
  (hpivot_trans : ∀ {T₁ T₂ : tableau m n} {i : fin m} {c : fin n},
    rel obj (T₁.pivot i c) T₁ → rel obj T₁ T₂ →
     c ∈ pivot_col T₁ obj →
     i ∈ pivot_row T₁ obj c → C (T₁.pivot i c) T₁ → C T₁ T₂ → C (pivot T₁ i c) T₂) :
  C T T' :=
rel.rec_on hrel (λ T hT i c  hc hr, hpivot hT hc hr) (λ T₁ T₂ i c hrelT₁₂ hc hr ih, hpivot_trans
  (rel.pivot (feasible_of_rel_left obj hrelT₁₂) hc hr) hrelT₁₂ hc hr
  (hpivot (feasible_of_rel_left obj hrelT₁₂) hc hr) ih)

lemma rel.trans {obj : fin m} {T₁ T₂ T₃ : tableau m n} (h₁₂ : rel obj T₁ T₂) :
  rel obj T₂ T₃ → rel obj T₁ T₃ :=
rel.rec_on h₁₂
  (λ T i c hT hc hr hrelT, rel.trans_pivot hrelT hc hr)
  (λ T₁ T₂ i c hrelT₁₂ hc hr ih hrelT₂₃, rel.trans_pivot (ih hrelT₂₃) hc hr)

instance : is_trans (tableau m n) (rel obj) := ⟨@rel.trans _ _ obj⟩

lemma flat_eq_of_rel {T T' : tableau m n} (h : rel obj T' T) : flat T' = flat T :=
rel.rec_on' h (λ _ _ _ _ _ hr, flat_pivot (ne_zero_of_mem_pivot_row hr))
  (λ _ _ _ _ _ _ _ _, eq.trans)

lemma rowg_obj_eq_of_rel {T T' : tableau m n} (h : rel obj T T') : T.to_partition.rowg obj =
  T'.to_partition.rowg obj :=
rel.rec_on' h (λ T i c hfT hc hr, by simp [rowg_swap_of_ne _ (pivot_row_spec hr).1])
  (λ _ _ _ _ _ _ _ _, eq.trans)

lemma restricted_eq_of_rel {T T' : tableau m n} (h : rel obj T T') : T.restricted = T'.restricted :=
rel.rec_on' h (λ _ _ _ _ _ _, rfl) (λ _ _ _ _ _ _ _ _, eq.trans)

lemma dead_eq_of_rel {T T' : tableau m n} (h : rel obj T T') : T.dead = T'.dead :=
rel.rec_on' h (λ _ _ _ _ _ _, rfl) (λ _ _ _ _ _ _ _ _, eq.trans)

lemma dead_eq_of_rel_or_eq {T T' : tableau m n} (h : T = T' ∨ rel obj T T') : T.dead = T'.dead :=
h.elim (congr_arg _) $ dead_eq_of_rel _

lemma exists_mem_pivot_row_col_of_rel {T T' : tableau m n} (h : rel obj T' T) :
  ∃ i c, c ∈ pivot_col T obj ∧ i ∈ pivot_row T obj c :=
rel.rec_on' h (λ _ i c _ hc hr, ⟨i, c, hc, hr⟩) (λ _ _ _ _ _ _ _ _ _, id)

lemma exists_mem_pivot_row_of_rel {T T' : tableau m n} (h : rel obj T' T) {c : fin n}
  (hc : c ∈ pivot_col T obj) : ∃ i, i ∈ pivot_row T obj c :=
let ⟨i, c', hc', hr⟩ := exists_mem_pivot_row_col_of_rel obj h in ⟨i, by simp * at *⟩

lemma exists_mem_pivot_col_of_fickle {T₁ T₂ : tableau m n} (h : rel obj T₂ T₁) {c : fin n} :
  fickle T₁ T₂ (T₁.to_partition.colg c) →
  ∃ T₃, (T₃ = T₁ ∨ rel obj T₃ T₁) ∧ (rel obj T₂ T₃) ∧
  T₃.to_partition.colg c = T₁.to_partition.colg c ∧
  c ∈ pivot_col T₃ obj :=
rel.rec_on' h begin
    assume T i c' hT hc' hr,
    rw fickle_colg_iff_ne,
    by_cases hcc : c = c',
    { subst hcc,
      exact λ _, ⟨T, or.inl rfl, rel.pivot hT hc' hr, rfl, hc'⟩ },
    { simp [colg_swap_of_ne _ hcc] }
  end
  (λ T₁ T₂ i c hrelp₁ hrel₁₂ hc hr ihp₁ ih₁₂,
    (imp_iff_not_or.1 ih₁₂).elim
      (λ ih₁₂, (imp_iff_not_or.1 ihp₁).elim
        (λ ihp₁ hf, (fickle_colg_iff_ne.1 hf (by simp [*, fickle_colg_iff_ne] at *)).elim)
        (λ ⟨T₃, hT₃⟩ hf, ⟨T₃,
          hT₃.1.elim (λ h, h.symm ▸ or.inr hrel₁₂) (λ h, or.inr $ h.trans hrel₁₂),
          hT₃.2.1, hT₃.2.2.1.trans (by simpa [eq_comm, fickle_colg_iff_ne] using ih₁₂), hT₃.2.2.2⟩))
      (λ ⟨T₃, hT₃⟩ hf, ⟨T₃, hT₃.1, hrelp₁.trans hT₃.2.1, hT₃.2.2⟩))

lemma exists_mem_pivot_row_of_fickle {T₁ T₂ : tableau m n} (h : rel obj T₂ T₁) (i : fin m) :
  fickle T₁ T₂ (T₁.to_partition.rowg i) →
  ∃ (T₃ : tableau m n) c, (T₃ = T₁ ∨ rel obj T₃ T₁) ∧ (rel obj T₂ T₃) ∧
    T₃.to_partition.rowg i = T₁.to_partition.rowg i ∧
    c ∈ pivot_col T₃ obj ∧ i ∈ pivot_row T₃ obj c :=
rel.rec_on' h
  begin
    assume T i' c hT hc hi',
    rw fickle_rowg_iff_ne,
    by_cases hii : i = i',
    { subst hii,
      exact λ _, ⟨T, c, or.inl rfl, rel.pivot hT hc hi', rfl, hc, hi'⟩ },
    { simp [rowg_swap_of_ne _ hii] }
  end
  (λ T₁ T₂ i c hrelp₁ hrel₁₂ hc hrow ihp₁ ih₁₂,
    (imp_iff_not_or.1 ih₁₂).elim
      (λ ih₁₂, (imp_iff_not_or.1 ihp₁).elim
        (λ ihp₁ hf, (fickle_rowg_iff_ne.1 hf (by simp [*, fickle_rowg_iff_ne] at *)).elim)
        (λ ⟨T₃, c', hT₃⟩ hf, ⟨T₃, c', hT₃.1.elim (λ h, h.symm ▸ or.inr hrel₁₂)
          (λ h, or.inr $ h.trans hrel₁₂),
            hT₃.2.1,
            hT₃.2.2.1.trans (by simpa [eq_comm, fickle_rowg_iff_ne] using ih₁₂),
            by clear_aux_decl; tauto⟩))
      (λ ⟨T₃, c', hT₃⟩ _, ⟨T₃, c', hT₃.1,
        (rel.pivot (feasible_of_rel_left _ hrel₁₂) hc hrow).trans hT₃.2.1, hT₃.2.2⟩))

lemma eq_or_rel_pivot_of_rel {T₁ T₂ : tableau m n} (h : rel obj T₁ T₂) : ∀ {i j}
  (hcol : j ∈ pivot_col T₂ obj) (hrow : i ∈ pivot_row T₂ obj j),
  T₁ = T₂.pivot i j ∨ rel obj T₁ (T₂.pivot i j) :=
rel.rec_on' h (λ T i c hT hc hr r' c' hc' hr', by simp * at *)
  (λ T₁ T₂ i c hrelp₁ hrel₁₂ hc hr ihp₁ ih₁₂ r' c' hc' hr',
    (ih₁₂ hc' hr').elim
      (λ ih₁₂, or.inr $ ih₁₂ ▸ rel.pivot (feasible_of_rel_left _ hrel₁₂) hc hr)
      (λ ih₁₂, or.inr $ (rel.pivot (feasible_of_rel_left _ hrel₁₂) hc hr).trans ih₁₂))

lemma exists_mem_pivot_col_of_mem_pivot_row {T : tableau m n} (hrelTT : rel obj T T)
  {i c} (hc : c ∈ pivot_col T obj) (hrow : i ∈ pivot_row T obj c) :
  ∃ (T' : tableau m n), c ∈ pivot_col T' obj ∧ T'.to_partition.colg c =
  T.to_partition.rowg i ∧ rel obj T' T ∧ rel obj T T' :=
have hrelTTp : rel obj T (T.pivot i c),
  from (eq_or_rel_pivot_of_rel _ hrelTT hc hrow).elim (λ h, h ▸ hrelTT ) id,
let ⟨T', hT'⟩ := exists_mem_pivot_col_of_fickle obj hrelTTp $ fickle_colg_iff_ne.2 $
  (show (T.pivot i c).to_partition.colg c ≠ T.to_partition.colg c, by simp) in
⟨T', hT'.2.2.2, by simp [hT'.2.2.1], hT'.1.elim
  (λ h, h.symm ▸ rel.pivot (feasible_of_rel_left _ hrelTT) hc hrow)
  (λ h, h.trans $ rel.pivot (feasible_of_rel_left _ hrelTT) hc hrow), hT'.2.1⟩

lemma exists_mem_pivot_col_of_fickle_row {T T' : tableau m n} (hrelTT' : rel obj T T') {i : fin m}
  (hrelT'T : rel obj T' T) (hrow : fickle T T' (T.to_partition.rowg i)) :
  ∃ (T₃ : tableau m n) c, c ∈ pivot_col T₃ obj ∧ T₃.to_partition.colg c =
  T.to_partition.rowg i ∧ rel obj T₃ T ∧ rel obj T T₃ :=
let ⟨T₃, c, hT₃, hrelT₃T, hrow₃, hc, hr⟩ :=
  exists_mem_pivot_row_of_fickle obj hrelT'T _ hrow in
let ⟨T₄, hT₄⟩ := exists_mem_pivot_col_of_mem_pivot_row obj
  (show rel obj T₃ T₃, from hT₃.elim (λ h, h.symm ▸ hrelTT'.trans hrelT'T)
    (λ h, h.trans $ hrelTT'.trans hrelT₃T)) hc hr in
⟨T₄, c, hT₄.1, hT₄.2.1.trans hrow₃, hT₄.2.2.1.trans $ hT₃.elim (λ h, h.symm ▸ hrelTT'.trans hrelT'T)
  (λ h, h.trans $ hrelTT'.trans hrelT'T), hrelTT'.trans (hrelT₃T.trans hT₄.2.2.2)⟩

lemma const_obj_le_of_rel {T₁ T₂ : tableau m n} (h : rel obj T₁ T₂) :
  T₂.const obj 0 ≤ T₁.const obj 0 :=
rel.rec_on' h (λ T i c hT hc hr,
    have hr' : _ := pivot_row_spec hr,
    simplex_const_obj_le hT (by tauto) (by tauto))
  (λ _ _ _ _ _ _ _ _ h₁ h₂, le_trans h₂ h₁)

lemma const_obj_eq_of_rel_of_rel {T₁ T₂ : tableau m n} (h₁₂ : rel obj T₁ T₂)
  (h₂₁ : rel obj T₂ T₁) : T₁.const obj 0 = T₂.const obj 0 :=
le_antisymm (const_obj_le_of_rel _ h₂₁) (const_obj_le_of_rel _ h₁₂)

lemma const_eq_const_of_const_obj_eq {T₁ T₂ : tableau m n} (h₁₂ : rel obj T₁ T₂) :
  ∀ (hobj : T₁.const obj 0 = T₂.const obj 0) (i : fin m), T₁.const i 0 = T₂.const i 0 :=
rel.rec_on' h₁₂
  (λ T i c hfT hc hrow hobj i',
    have hr0 : T.const i 0 = 0, from const_eq_zero_of_const_obj_eq hfT
      (ne_zero_of_mem_pivot_col hc) (ne_zero_of_mem_pivot_row hrow)
      (pivot_row_spec hrow).1 hobj,
    if hii : i' = i
      then by simp [hii, hr0]
      else by simp [const_pivot_of_ne _ hii, hr0])
  (λ T₁ T₂ i c hrelp₁ hrel₁₂ hc hr ihp₁ ih₁₂ hobj i',
    have hobjp : (pivot T₁ i c).const obj 0 = T₁.const obj 0,
      from le_antisymm (hobj.symm ▸ const_obj_le_of_rel _ hrel₁₂)
        (const_obj_le_of_rel _ hrelp₁),
    by rw [ihp₁ hobjp, ih₁₂ (hobjp.symm.trans hobj)])

lemma const_eq_zero_of_fickle_of_rel_self {T T' : tableau m n} (hrelTT' : rel obj T T')
  (hrelT'T : rel obj T' T) (i : fin m) (hrow : fickle T T' (T.to_partition.rowg i)) :
  T.const i 0 = 0 :=
let ⟨T₃, c, hT₃₁, hT'₃, hrow₃, hc, hi⟩ := exists_mem_pivot_row_of_fickle obj hrelT'T _ hrow in
have T₃.const i 0 = 0, from const_eq_zero_of_const_obj_eq
  (feasible_of_rel_right _ hT'₃) (ne_zero_of_mem_pivot_col hc)
  (ne_zero_of_mem_pivot_row hi) (pivot_row_spec hi).1
  (const_obj_eq_of_rel_of_rel _ (rel.pivot (feasible_of_rel_right _ hT'₃) hc hi)
    ((eq_or_rel_pivot_of_rel _ hT'₃ hc hi).elim
      (λ h, h ▸ hT₃₁.elim (λ h, h.symm ▸ hrelTT') (λ h, h.trans hrelTT'))
      (λ hrelT'p, hT₃₁.elim (λ h, h.symm ▸ hrelTT'.trans (h ▸ hrelT'p))
        (λ h, h.trans $ hrelTT'.trans hrelT'p)))),
have hobj : T₃.const obj 0 = T.const obj 0,
  from hT₃₁.elim (λ h, h ▸ rfl) (λ h, const_obj_eq_of_rel_of_rel _ h (hrelTT'.trans hT'₃)),
hT₃₁.elim (λ h, h ▸ this) (λ h, const_eq_const_of_const_obj_eq obj h hobj i ▸ this)

lemma colg_mem_restricted_of_rel_self {T : tableau m n} (hrelTT : rel obj T T)
  {c} (hc : c ∈ pivot_col T obj) : T.to_partition.colg c ∈ T.restricted :=
let ⟨i, hrow⟩ := exists_mem_pivot_row_of_rel obj hrelTT hc in
let ⟨T', c', hT', hrelTT', hrowcol, _, hi'⟩ := exists_mem_pivot_row_of_fickle obj
    ((eq_or_rel_pivot_of_rel _ hrelTT hc hrow).elim
      (λ h, show rel obj T (T.pivot i c), from h ▸ hrelTT) id) _
  (fickle_rowg_iff_ne.2 $
    show (T.pivot i c).to_partition.rowg i ≠ T.to_partition.rowg i, by simp) in
(restricted_eq_of_rel _ hrelTT').symm ▸ by convert (pivot_row_spec hi').2.1; simp [hrowcol]

lemma eq_zero_of_not_mem_restricted_of_rel_self {T : tableau m n} (hrelTT : rel obj T T)
  {j} (hjres : T.to_partition.colg j ∉ T.restricted) (hdead : j ∉ T.dead) : T.to_matrix obj j = 0 :=
let ⟨r, c, hc, hr⟩ := exists_mem_pivot_row_col_of_rel obj hrelTT in
have hcres : T.to_partition.colg c ∈ T.restricted,
  from colg_mem_restricted_of_rel_self obj hrelTT hc,
by_contradiction $ λ h0,
begin
  simp [pivot_col] at hc,
  cases h : fin.find (λ c, T.to_matrix obj c ≠ 0 ∧ colg (T.to_partition) c ∉ T.restricted
    ∧ c ∉ T.dead),
  { simp [*, fin.find_eq_none_iff] at * },
  { rw h at hc, clear_aux_decl,
    have := (fin.find_eq_some_iff.1 h).1,
    simp * at * }
end

lemma rel.irrefl {obj : fin m} : ∀ (T : tableau m n), ¬ rel obj T T :=
λ T1 hrelT1,
let ⟨iT1 , cT1, hrT1, hcT1⟩ := exists_mem_pivot_row_col_of_rel obj hrelT1 in
let ⟨t, ht⟩ := finset.max_of_mem
  (show T1.to_partition.colg cT1 ∈ univ.filter (λ v, ∃ (T' : tableau m n) (c : fin n),
      rel obj T' T' ∧ c ∈ pivot_col T' obj ∧ T'.to_partition.colg c = v),
    by simp only [true_and, mem_filter, mem_univ, exists_and_distrib_left];
      exact ⟨T1, hrelT1, cT1, hrT1, rfl⟩) in
let ⟨_, T', c', hrelTT'', hcT', hct⟩ := finset.mem_filter.1 (finset.mem_of_max ht) in
have htmax : ∀ (s : fin (m + n)) (T : tableau m n),
    rel obj T T → ∀ (j : fin n), pivot_col T obj = some j →
      T.to_partition.colg j = s → s ≤ t,
  by simpa using λ s (h : s ∈ _), finset.le_max_of_mem h ht,
let ⟨i, hiT'⟩ := exists_mem_pivot_row_of_rel obj hrelTT'' hcT' in
have hrelTT''p : rel obj T' (T'.pivot i c'),
  from (eq_or_rel_pivot_of_rel obj hrelTT'' hcT' hiT').elim (λ h, h ▸ hrelTT'') id,
let ⟨T, c, hTT', hrelT'T, hT'Tr, hc, hr⟩ := exists_mem_pivot_row_of_fickle obj
  hrelTT''p i (by rw fickle_symm; simp [fickle_colg_iff_ne]) in
have hfT' : feasible T', from feasible_of_rel_left _ hrelTT'',
have hfT : feasible T, from feasible_of_rel_right _ hrelT'T,
have hrelT'pT' : rel obj (T'.pivot i c') T', from rel.pivot hfT' hcT' hiT',
have hrelTT' : rel obj T T', from hTT'.elim (λ h, h.symm ▸ hrelT'pT') (λ h, h.trans hrelT'pT'),
have hrelTT : rel obj T T, from hrelTT'.trans hrelT'T,
have hc't : T.to_partition.colg c ≤ t, from htmax _ T hrelTT _ hc rfl,
have hoT'T : T'.const obj 0 = T.const obj 0, from const_obj_eq_of_rel_of_rel _ hrelT'T hrelTT',
have hfickle : ∀ i, fickle T T' (T.to_partition.rowg i) → T.const i 0 = 0,
  from const_eq_zero_of_fickle_of_rel_self obj hrelTT' hrelT'T,
have hobj : T.const obj 0 = T'.const obj 0, from const_obj_eq_of_rel_of_rel _ hrelTT' hrelT'T,
have hflat : T.flat = T'.flat, from flat_eq_of_rel obj hrelTT',
have hrobj : T.to_partition.rowg obj = T'.to_partition.rowg obj, from rowg_obj_eq_of_rel _ hrelTT',
have hs : T.to_partition.rowg i = T'.to_partition.colg c', by simpa using hT'Tr,
have hc'res : T'.to_partition.colg c' ∈ T'.restricted,
  from hs ▸ restricted_eq_of_rel _ hrelTT' ▸ (pivot_row_spec hr).2.1,
have hc'obj0 : 0 < T'.to_matrix obj c' ∧ c' ∉ T'.dead,
  by simpa [hc'res] using pivot_col_spec hcT',
have hcres : T.to_partition.colg c ∈ T.restricted,
  from colg_mem_restricted_of_rel_self obj hrelTT hc,
have hcobj0 : 0 < to_matrix T obj c ∧ c ∉ T.dead,
  by simpa [hcres] using pivot_col_spec hc,
have hrc0 : T.to_matrix i c < 0,
  from inv_neg'.1 $ neg_of_mul_neg_left (pivot_row_spec hr).2.2.1 (le_of_lt hcobj0.1),
have nonpos_of_colg_ne : ∀ j, (fickle T' T (T'.to_partition.colg j)) → j ≠ c' →
    T'.to_matrix obj j ≤ 0,
  from λ j hj hjc',
    let ⟨T₃, hT₃⟩ := exists_mem_pivot_col_of_fickle obj hrelTT' hj in
    nonpos_of_lt_pivot_col hcT' hc'res
      (dead_eq_of_rel_or_eq obj hT₃.1 ▸ (pivot_col_spec hT₃.2.2.2).2)
      (lt_of_le_of_ne
        (hct.symm ▸ hT₃.2.2.1 ▸ htmax _ T₃ (hT₃.1.elim (λ h, h.symm ▸ hrelTT'')
          (λ h, h.trans (hrelT'T.trans hT₃.2.1))) _ hT₃.2.2.2 rfl)
        (by rwa [ne.def, T'.to_partition.injective_colg.eq_iff])),
have nonpos_of_colg_eq : ∀ j, j ≠ c' →
    T'.to_partition.colg j = T.to_partition.colg c → T'.to_matrix obj j ≤ 0,
  from λ j hjc' hj,
    if hjc : j = c
    then by clear_aux_decl; subst hjc;
      exact nonpos_of_lt_pivot_col hcT' hc'res
        (by rw [dead_eq_of_rel obj hrelT'T]; tauto)
        (lt_of_le_of_ne (hj.symm ▸ hct.symm ▸ hc't) (by simpa))
    else nonpos_of_colg_ne _ (fickle_colg_iff_ne.2 $ by simpa [hj, eq_comm] using hjc) hjc',
have unique_row : ∀ i' ≠ i, T.const i' 0 = 0 → fickle T T' (T.to_partition.rowg i') →
    0 ≤ T.to_matrix i' c,
  from λ i' hii hi0 hrow,
    let ⟨T₃, c₃, hc₃, hrow₃, hrelT₃T, hrelTT₃⟩ :=
      exists_mem_pivot_col_of_fickle_row _ hrelTT' hrelT'T hrow in
    have hrelT₃T₃ : rel obj T₃ T₃, from hrelT₃T.trans hrelTT₃,
    nonneg_of_lt_pivot_row (by exact hcobj0.1)
      (by rw [← hrow₃, ← restricted_eq_of_rel _ hrelT₃T];
        exact colg_mem_restricted_of_rel_self _ hrelT₃T₃ hc₃) hc hr hi0
      (lt_of_le_of_ne (by rw [hs, hct, ← hrow₃]; exact htmax _ _ hrelT₃T₃ _ hc₃ rfl)
        (by simpa [fickle_rowg_iff_ne] using hrow)),
not_unique_row_and_unique_col obj hcobj0.1 hc'obj0.1 hrc0 hflat hs hrobj hfickle hobj
  nonpos_of_colg_eq nonpos_of_colg_ne unique_row

noncomputable instance fintype_rel (T : tableau m n) : fintype {T' | rel obj T' T} :=
fintype.of_injective (λ T', T'.val.to_partition)
  (λ T₁ T₂ h, subtype.eq $ tableau.ext
    (by rw [flat_eq_of_rel _ T₁.2, flat_eq_of_rel _ T₂.2]) h
    (by rw [dead_eq_of_rel _ T₁.2, dead_eq_of_rel _ T₂.2])
    (by rw [restricted_eq_of_rel _ T₁.2, restricted_eq_of_rel _ T₂.2]))

lemma rel_wf (m n : ℕ) (obj : fin m): well_founded (@rel m n obj) :=
subrelation.wf
  (show subrelation (@rel m n obj) (measure (λ T, fintype.card {T' | rel obj T' T})),
    from assume T₁ T₂ h,
    set.card_lt_card (set.ssubset_iff_subset_not_subset.2 ⟨λ T' hT', hT'.trans h,
      not_forall_of_exists_not ⟨T₁, λ h', rel.irrefl _ (h' h)⟩⟩))
  (measure_wf (λ T, fintype.card {T' | rel obj T' T}))

end blands_rule

inductive termination (n : ℕ) : Type
| while {}              : termination
| unbounded (c : fin n) : termination
| optimal {}            : termination

namespace termination

lemma injective_unbounded : function.injective (@unbounded n) :=
λ _ _ h, by injection h

@[simp] lemma unbounded_inj {c c' : fin n} : unbounded c = unbounded c' ↔ c = c' :=
injective_unbounded.eq_iff

end termination

open termination

instance {n : ℕ} : has_repr $ termination n :=
⟨λ t, termination.cases_on t
  "while"
  (λ c, "unbounded " ++ repr c)
  "optimal"⟩

open termination

/-- The simplex algorithm -/
def simplex (w : tableau m n → bool) (obj : fin m) : Π (T : tableau m n) (hT : feasible T),
  tableau m n × termination n
| T := λ hT, cond (w T)
  (match pivot_col T obj, @feasible_of_mem_pivot_row_and_col _ _ _ obj hT,
      @rel.pivot m n obj _ hT with
    | none,   hc, hrel := (T, optimal)
    | some j, hc, hrel :=
      match pivot_row T obj j, @hc _ rfl, (λ i, @hrel i j rfl) with
      | none,   hrow, hrel := (T, unbounded j)
      | some i, hrow, hrel := have wf : rel obj (pivot T i j) T, from hrel _ rfl,
        simplex (T.pivot i j) (hrow rfl)
      end
    end)
  (T, while)
using_well_founded {rel_tac := λ _ _, `[exact ⟨_, rel_wf m n obj⟩],
  dec_tac := tactic.assumption}

lemma simplex_pivot {w : tableau m n → bool} {T : tableau m n} (hT : feasible T)
  (hw : w T = tt) {obj : fin m} {i : fin m} {c : fin n}
  (hc : c ∈ pivot_col T obj) (hr : i ∈ pivot_row T obj c) :
  (T.pivot i c).simplex w obj  (feasible_of_mem_pivot_row_and_col hT hc hr) =
  T.simplex w obj hT  :=
by conv_rhs { rw simplex };
  simp [hw, show _ = _, from hr, show _ = _, from hc, simplex._match_1, simplex._match_2]

lemma simplex_spec_aux (w : tableau m n → bool) (obj : fin m) :
  Π (T : tableau m n) (hT : feasible T),
  ((T.simplex w obj hT).2 = while ∧ w (T.simplex w obj hT).1 = ff) ∨
  ((T.simplex w obj hT).2 = optimal ∧ w (T.simplex w obj hT).1 = tt ∧
    pivot_col (T.simplex w obj hT).1 obj = none) ∨
  ∃ c, ((T.simplex w obj hT).2 = unbounded c ∧ w (T.simplex w obj hT).1 = tt ∧
    c ∈ pivot_col (T.simplex w obj hT).1 obj ∧
    pivot_row (T.simplex w obj hT).1 obj c = none)
| T := λ hT,
  begin
    cases hw : w T,
    { rw simplex, simp [hw] },
    { cases hc : pivot_col T obj with c,
      { rw simplex, simp [hc, hw, simplex._match_1] },
      { cases hrow : pivot_row T obj c with i,
        { rw simplex,
          simp [hrow, hc, hw, simplex._match_1, simplex._match_2] },
        { rw [← simplex_pivot hT hw hc hrow],
          exact have wf : rel obj (T.pivot i c) T, from rel.pivot hT hc hrow,
            simplex_spec_aux _ _ } } }
  end
using_well_founded {rel_tac := λ _ _, `[exact ⟨_, rel_wf m n obj⟩],
  dec_tac := tactic.assumption}

lemma simplex_while_eq_ff {w : tableau m n → bool} {T : tableau m n} {hT : feasible T}
  {obj : fin m} (hw : w T = ff) : T.simplex w obj hT = (T, while) :=
by rw [simplex, hw]; refl

lemma simplex_pivot_col_eq_none {w : tableau m n → bool} {T : tableau m n} {hT : feasible T}
  (hw : w T = tt) {obj : fin m} (hc : pivot_col T obj = none) :
  T.simplex w obj hT = (T, optimal) :=
by rw simplex; simp [hc, hw, simplex._match_1]

lemma simplex_pivot_row_eq_none {w : tableau m n → bool} {T : tableau m n} {hT : feasible T}
  {obj : fin m} (hw : w T = tt) {c} (hc : c ∈ pivot_col T obj)
  (hr : pivot_row T obj c = none) : T.simplex w obj hT = (T, unbounded c) :=
by rw simplex; simp [hw, show _ = _, from hc, hr, simplex._match_1, simplex._match_2]

lemma simplex_induction (P : tableau m n → Prop) (w : tableau m n → bool) (obj : fin m):
  Π {T : tableau m n} (hT : feasible T)  (h0 : P T)
  (hpivot : ∀ {T' i c}, w T' = tt → c ∈ pivot_col T' obj → i ∈ pivot_row T' obj c
    → feasible T' → P T' → P (T'.pivot i c)),
  P (T.simplex w obj hT).1
| T := λ hT h0 hpivot,
  begin
    cases hw : w T,
    { rwa [simplex_while_eq_ff hw] },
    { cases hc : pivot_col T obj with c,
      { rwa [simplex_pivot_col_eq_none hw hc] },
      { cases hrow : pivot_row T obj c with i,
        { rwa simplex_pivot_row_eq_none hw hc hrow },
        { rw [← simplex_pivot _ hw hc hrow],
          exact have wf : rel obj (pivot T i c) T, from rel.pivot hT hc hrow,
            simplex_induction (feasible_of_mem_pivot_row_and_col hT hc hrow)
              (hpivot hw hc hrow hT h0) @hpivot } } }
  end
using_well_founded {rel_tac := λ _ _, `[exact ⟨_, rel_wf m n obj⟩],
  dec_tac := `[tauto]}

@[simp] lemma feasible_simplex {w : tableau m n → bool} {T : tableau m n}
  {hT : feasible T} {obj : fin m} : feasible (T.simplex w obj hT).1 :=
simplex_induction feasible _ _ hT hT
  (λ _ _ _ _ hc hr _ hT', feasible_of_mem_pivot_row_and_col hT' hc hr)

@[simp] lemma simplex_simplex {w : tableau m n → bool} {T : tableau m n} {hT : feasible T}
  {obj : fin m} : (T.simplex w obj hT).1.simplex w obj feasible_simplex = T.simplex w obj hT :=
simplex_induction (λ T', ∀ (hT' : feasible T'), T'.simplex w obj hT' = T.simplex w obj hT) w _ _
  (λ _, rfl) (λ T' i c hw hc hr hT' ih hpivot, by rw [simplex_pivot hT' hw hc hr, ih]) _

/-- `simplex` does not move the row variable it is trying to maximise. -/
@[simp] lemma rowg_simplex (T : tableau m n) (hT : feasible T) (w : tableau m n → bool)
  (obj : fin m) : (T.simplex w obj hT).1.to_partition.rowg obj = T.to_partition.rowg obj :=
simplex_induction (λ T', T'.to_partition.rowg obj = T.to_partition.rowg obj) _ _ _ rfl
  (λ T' i c hw hc hr, by simp [rowg_swap_of_ne _ (pivot_row_spec hr).1])

@[simp] lemma flat_simplex (T : tableau m n) (hT : feasible T) (w : tableau m n → bool)
  (obj : fin m) : (T.simplex w obj hT).1.flat = T.flat :=
simplex_induction (λ T', T'.flat = T.flat) w obj _ rfl
  (λ T' i c hw hc hr hT' ih,
    have T'.to_matrix i c ≠ 0,
      from λ h, by simpa [h, lt_irrefl] using pivot_row_spec hr,
    by rw [flat_pivot this, ih])

@[simp] lemma restricted_simplex (T : tableau m n) (hT : feasible T) (w : tableau m n → bool)
  (obj : fin m) : (T.simplex w obj hT).1.restricted = T.restricted :=
simplex_induction (λ T', T'.restricted = T.restricted) _ _ _ rfl (by simp { contextual := tt })

@[simp] lemma dead_simplex (T : tableau m n) (hT : feasible T) (w : tableau m n → bool)
  (obj : fin m) : (T.simplex w obj hT).1.dead = T.dead :=
simplex_induction (λ T', T'.dead = T.dead) _ _ _ rfl (by simp { contextual := tt })

@[simp] lemma res_set_simplex (T : tableau m n) (hT : feasible T) (w : tableau m n → bool)
  (obj : fin m) : (T.simplex w obj hT).1.res_set = T.res_set :=
simplex_induction (λ T', T'.res_set = T.res_set) w obj _ rfl
  (λ T' i c hw hc hr, by simp [res_set_pivot (ne_zero_of_mem_pivot_row hr)] {contextual := tt})

@[simp] lemma dead_set_simplex (T : tableau m n) (hT : feasible T) (w : tableau m n → bool)
  (obj : fin m) : (T.simplex w obj hT).1.dead_set = T.dead_set :=
simplex_induction (λ T', T'.dead_set = T.dead_set) w obj _ rfl
  (λ T' i c hw hc hr,
    by simp [dead_set_pivot (ne_zero_of_mem_pivot_row hr) (pivot_col_spec hc).2] {contextual := tt})

@[simp] lemma sol_set_simplex (T : tableau m n) (hT : feasible T) (w : tableau m n → bool)
  (obj : fin m) : (T.simplex w obj hT).1.sol_set = T.sol_set :=
by simp [sol_set_eq_res_set_inter_dead_set]

@[simp] lemma of_col_simplex_zero_mem_sol_set {w : tableau m n → bool} {T : tableau m n}
  {hT : feasible T} {obj : fin m} : (T.simplex w obj hT).1.of_col 0 ∈ sol_set T :=
by rw [← sol_set_simplex, of_col_zero_mem_sol_set_iff]; exact feasible_simplex

@[simp] lemma of_col_simplex_rowg {w : tableau m n → bool} {T : tableau m n}
  {hT : feasible T} {obj : fin m} (x : cvec n) :
  (T.simplex w obj hT).1.of_col x (T.to_partition.rowg obj) =
  ((T.simplex w obj hT).1.to_matrix ⬝ x + (T.simplex w obj hT).1.const) obj :=
by rw [← of_col_rowg (T.simplex w obj hT).1 x obj, rowg_simplex]

@[simp] lemma is_unbounded_above_simplex {T : tableau m n} {hT : feasible T} {w : tableau m n → bool}
  {obj : fin m} {v : fin (m + n)} : is_unbounded_above (T.simplex w obj hT).1 v ↔
  is_unbounded_above T v := by simp [is_unbounded_above]

@[simp] lemma is_optimal_simplex {T : tableau m n} {hT : feasible T} {w : tableau m n → bool}
  {obj : fin m} {x : cvec (m + n)} {v : fin (m + n)} : is_optimal (T.simplex w obj hT).1 x v ↔
  is_optimal T x v := by simp [is_optimal]

lemma termination_eq_while_iff {T : tableau m n} {hT : feasible T} {w : tableau m n → bool}
  {obj : fin m} : (T.simplex w obj hT).2 = while ↔ w (T.simplex w obj hT).1 = ff :=
by have := simplex_spec_aux w obj T hT; finish

lemma termination_eq_optimal_iff_pivot_col_eq_none {T : tableau m n}
  {hT : feasible T} {w : tableau m n → bool} {obj : fin m} : (T.simplex w obj hT).2 = optimal ↔
  w (T.simplex w obj hT).1 = tt ∧ pivot_col (T.simplex w obj hT).1 obj = none :=
by rcases simplex_spec_aux w obj T hT with _ | ⟨_, _, _⟩ | ⟨⟨_, _⟩, _, _, _, _⟩; simp * at *

lemma termination_eq_unbounded_iff_pivot_row_eq_none {T : tableau m n} {hT : feasible T}
  {w : tableau m n → bool} {obj : fin m} {c : fin n} :
  (T.simplex w obj hT).2 = unbounded c ↔
  w (T.simplex w obj hT).1 = tt ∧ c ∈ pivot_col (T.simplex w obj hT).1 obj ∧
    pivot_row (T.simplex w obj hT).1 obj c = none :=
by split; intros; rcases simplex_spec_aux w obj T hT with
  _ | ⟨_, _, _⟩ | ⟨⟨⟨_, _⟩, _⟩, _, _, _, _⟩; simp * at *

lemma unbounded_of_termination_eq_unbounded {T : tableau m n} {hT : feasible T}
  {w : tableau m n → bool} {obj : fin m} {c : fin n} : (T.simplex w obj hT).2 = unbounded c →
  w (T.simplex w obj hT).1 = tt ∧
  is_unbounded_above T (T.to_partition.rowg obj) :=
begin
  rw termination_eq_unbounded_iff_pivot_row_eq_none,
  rintros ⟨_, hc⟩,
  simpa * using pivot_row_eq_none feasible_simplex hc.2 hc.1
end

lemma termination_eq_optimal_iff {T : tableau m n} {hT : feasible T}
  {w : tableau m n → bool} {obj : fin m} : (T.simplex w obj hT).2 = optimal ↔
  w (T.simplex w obj hT).1 = tt ∧
  is_optimal T ((T.simplex w obj hT).1.of_col 0) (T.to_partition.rowg obj) :=
begin
  rw [termination_eq_optimal_iff_pivot_col_eq_none],
  split,
  { rintros ⟨_, hc⟩,
    simpa * using pivot_col_eq_none feasible_simplex hc },
  { cases ht : (T.simplex w obj hT).2,
    { simp [*, termination_eq_while_iff] at * },
    { cases unbounded_of_termination_eq_unbounded ht,
      simp [*, not_optimal_of_unbounded_above right] },
    { simp [*, termination_eq_optimal_iff_pivot_col_eq_none] at * } }
end

lemma termination_eq_unbounded_iff {T : tableau m n} {hT : feasible T}
  {w : tableau m n → bool} {obj : fin m} {c : fin n}: (T.simplex w obj hT).2 = unbounded c ↔
  w (T.simplex w obj hT).1 = tt ∧ is_unbounded_above T (T.to_partition.rowg obj)
  ∧ c ∈ pivot_col (T.simplex w obj hT).1 obj :=
⟨λ hc, and.assoc.1 $ ⟨unbounded_of_termination_eq_unbounded hc,
  (termination_eq_unbounded_iff_pivot_row_eq_none.1 hc).2.1⟩,
begin
  have := @not_optimal_of_unbounded_above m n (T.simplex w obj hT).1 (T.to_partition.rowg obj)
    ((T.simplex w obj hT).1.of_col 0),
  cases ht : (T.simplex w obj hT).2;
  simp [termination_eq_optimal_iff, termination_eq_while_iff,
    termination_eq_unbounded_iff_pivot_row_eq_none, *] at *
end⟩

end tableau


/-
Copyright (c) 2022 Mario Carneiro. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro
-/
import Std.Logic
import Std.Data.RBMap.Basic

/-!
# Lemmas for Red-black trees

The main theorem in this file is `WF_def`, which shows that the `RBNode.WF.mk` constructor
subsumes the others, by showing that `insert` and `erase` satisfy the red-black invariants.
-/

namespace Std

namespace RBNode
open RBColor

attribute [simp] All

theorem All.trivial (H : ∀ {x : α}, p x) : ∀ {t : RBNode α}, t.All p
  | nil => _root_.trivial
  | node .. => ⟨H, All.trivial H, All.trivial H⟩

theorem All_and {t : RBNode α} : t.All (fun a => p a ∧ q a) ↔ t.All p ∧ t.All q := by
  induction t <;> simp [*, and_assoc, and_left_comm]

theorem cmpLT.trans (h₁ : cmpLT cmp x y) (h₂ : cmpLT cmp y z) : cmpLT cmp x z :=
  ⟨TransCmp.lt_trans h₁.1 h₂.1⟩

theorem cmpLT.trans_l {cmp x y} (H : cmpLT cmp x y) {t : RBNode α}
    (h : t.All (cmpLT cmp y ·)) : t.All (cmpLT cmp x ·) := h.imp fun h => H.trans h

theorem cmpLT.trans_r {cmp x y} (H : cmpLT cmp x y) {a : RBNode α}
    (h : a.All (cmpLT cmp · x)) : a.All (cmpLT cmp · y) := h.imp fun h => h.trans H

theorem cmpEq.lt_congr_left (H : cmpEq cmp x y) : cmpLT cmp x z ↔ cmpLT cmp y z :=
  ⟨fun ⟨h⟩ => ⟨TransCmp.cmp_congr_left H.1 ▸ h⟩, fun ⟨h⟩ => ⟨TransCmp.cmp_congr_left H.1 ▸ h⟩⟩

theorem cmpEq.lt_congr_right (H : cmpEq cmp y z) : cmpLT cmp x y ↔ cmpLT cmp x z :=
  ⟨fun ⟨h⟩ => ⟨TransCmp.cmp_congr_right H.1 ▸ h⟩, fun ⟨h⟩ => ⟨TransCmp.cmp_congr_right H.1 ▸ h⟩⟩

/-- The `balance1` function preserves the ordering invariants. -/
protected theorem Ordered.balance1 {l : RBNode α} {v : α} {r : RBNode α}
    (lv : l.All (cmpLT cmp · v)) (vr : r.All (cmpLT cmp v ·))
    (hl : l.Ordered cmp) (hr : r.Ordered cmp) : (balance1 l v r).Ordered cmp := by
  unfold balance1; split
  · next a x b y c =>
    have ⟨yv, _, cv⟩ := lv; have ⟨xy, yc, hx, hc⟩ := hl
    exact ⟨xy, ⟨yv, yc, yv.trans_l vr⟩, hx, cv, vr, hc, hr⟩
  · next a x b y c _ =>
    have ⟨_, _, yv, _, cv⟩ := lv; have ⟨ax, ⟨xy, xb, _⟩, ha, by_, yc, hb, hc⟩ := hl
    exact ⟨⟨xy, xy.trans_r ax, by_⟩, ⟨yv, yc, yv.trans_l vr⟩, ⟨ax, xb, ha, hb⟩, cv, vr, hc, hr⟩
  · exact ⟨lv, vr, hl, hr⟩

@[simp] theorem balance1_All {l : RBNode α} {v : α} {r : RBNode α} :
    (balance1 l v r).All p ↔ p v ∧ l.All p ∧ r.All p := by
  unfold balance1; split <;> simp [and_assoc, and_left_comm]

/-- The `balance2` function preserves the ordering invariants. -/
protected theorem Ordered.balance2 {l : RBNode α} {v : α} {r : RBNode α}
    (lv : l.All (cmpLT cmp · v)) (vr : r.All (cmpLT cmp v ·))
    (hl : l.Ordered cmp) (hr : r.Ordered cmp) : (balance2 l v r).Ordered cmp := by
  unfold balance2; split
  · next b y c z d =>
    have ⟨_, ⟨vy, vb, _⟩, _⟩ := vr; have ⟨⟨yz, _, cz⟩, zd, ⟨by_, yc, hy, hz⟩, hd⟩ := hr
    exact ⟨⟨vy, vy.trans_r lv, by_⟩, ⟨yz, yc, yz.trans_l zd⟩, ⟨lv, vb, hl, hy⟩, cz, zd, hz, hd⟩
  · next a x b y c _ =>
    have ⟨vx, va, _⟩ := vr; have ⟨ax, xy, ha, hy⟩ := hr
    exact ⟨⟨vx, vx.trans_r lv, ax⟩, xy, ⟨lv, va, hl, ha⟩, hy⟩
  · exact ⟨lv, vr, hl, hr⟩

@[simp] theorem balance2_All {l : RBNode α} {v : α} {r : RBNode α} :
    (balance2 l v r).All p ↔ p v ∧ l.All p ∧ r.All p := by
  unfold balance2; split <;> simp [and_assoc, and_left_comm]

protected theorem Ordered.setBlack {t : RBNode α} : (setBlack t).Ordered cmp ↔ t.Ordered cmp := by
  unfold setBlack; split <;> simp [Ordered]

protected

theorem setBlack_idem {t : RBNode α} : t.setBlack.setBlack = t.setBlack := by cases t <;> rfl

theorem insert_setBlack {t : RBNode α} :
    (t.insert cmp v).setBlack = (t.ins cmp v).setBlack := by
  unfold insert; split <;> simp [setBlack_idem]

protected theorem All.ins {x : α} {t : RBNode α}
  (h₁ : p x) (h₂ : t.All p) : (ins cmp x t).All p := by
  induction t <;> unfold ins <;> simp [*]
  split <;> cases ‹_=_› <;> split <;> simp at h₂ <;> simp [*]

/-- The `ins` function preserves the ordering invariants. -/
protected theorem Ordered.ins : ∀ {t : RBNode α}, t.Ordered cmp → (ins cmp x t).Ordered cmp
  | nil, _ => ⟨⟨⟩, ⟨⟩, ⟨⟩, ⟨⟩⟩
  | node red a y b, ⟨ay, yb, ha, hb⟩ => by
    unfold ins; split
    · next h => exact ⟨ay.ins ⟨h⟩, yb, ha.ins, hb⟩
    · next h => exact ⟨ay, yb.ins ⟨OrientedCmp.cmp_eq_gt.1 h⟩, ha, hb.ins⟩
    · next h => exact (⟨
        ay.imp fun ⟨h'⟩ => ⟨(TransCmp.cmp_congr_right h).trans h'⟩,
        yb.imp fun ⟨h'⟩ => ⟨(TransCmp.cmp_congr_left h).trans h'⟩, ha, hb⟩)
  | node black a y b, ⟨ay, yb, ha, hb⟩ => by
    unfold ins; split
    · next h => exact ha.ins.balance1 (ay.ins ⟨h⟩) yb hb
    · next h => exact ha.balance2 ay (yb.ins ⟨OrientedCmp.cmp_eq_gt.1 h⟩) hb.ins
    · next h => exact (⟨
        ay.imp fun ⟨h'⟩ => ⟨(TransCmp.cmp_congr_right h).trans h'⟩,
        yb.imp fun ⟨h'⟩ => ⟨(TransCmp.cmp_congr_left h).trans h'⟩, ha, hb⟩)

/-- The `insert` function preserves the ordering invariants. -/
protected theorem Ordered.insert (h : t.Ordered cmp) : (insert cmp t v).Ordered cmp := by
  unfold RBNode.insert; split <;> simp [Ordered.setBlack, h.ins (x := v)]

/--
The red-red invariant is a weakening of the red-black balance invariant which allows
the root to be red with red children, but does not allow any other violations.
It occurs as a temporary condition in the `insert` and `erase` functions.

The `p` parameter allows the `.redred` case to be dependent on an additional condition.
If it is false, then this is equivalent to the usual red-black invariant.
-/
inductive RedRed (p : Prop) : RBNode α → Nat → Prop where
  /-- A balanced tree has the red-red invariant. -/
  | balanced : Balanced t c n → RedRed p t n
  /-- A red node with balanced red children has the red-red invariant (if `p` is true). -/
  | redred : p → Balanced a c₁ n → Balanced b c₂ n → RedRed p (node red a x b) n

/-- When `p` is false, the red-red case is impossible so the tree is balanced. -/
protected theorem RedRed.of_false (h : ¬p) : RedRed p t n → ∃ c, Balanced t c n
  | .balanced h => ⟨_, h⟩
  | .redred hp .. => nomatch h hp

/-- A `red` node with the red-red invariant has balanced children. -/
protected theorem RedRed.of_red : RedRed p (node red a x b) n →
    ∃ c₁ c₂, Balanced a c₁ n ∧ Balanced b c₂ n
  | .balanced (.red ha hb) | .redred _ ha hb => ⟨_, _, ha, hb⟩

/-- The red-red invariant is monotonic in `p`. -/
protected theorem RedRed.imp (h : p → q) : RedRed p t n → RedRed q t n
  | .balanced h => .balanced h
  | .redred hp ha hb => .redred (h hp) ha hb

/-- If `t` has the red-red invariant, then setting the root to black yields a balanced tree. -/
protected theorem RedRed.setBlack : t.RedRed p n → ∃ n', (setBlack t).Balanced black n'
  | .balanced h => h.setBlack
  | .redred _ hl hr => ⟨_, hl.black hr⟩

/-- The `balance1` function repairs the balance invariant when the first argument is red-red. -/
protected theorem RedRed.balance1 {l : RBNode α} {v : α} {r : RBNode α}
    (hl : l.RedRed p n) (hr : r.Balanced c n) : ∃ c, (balance1 l v r).Balanced c (n + 1) := by
  unfold balance1; split
  · have .redred _ (.red ha hb) hc := hl; exact ⟨_, .red (.black ha hb) (.black hc hr)⟩
  · have .redred _ ha (.red hb hc) := hl; exact ⟨_, .red (.black ha hb) (.black hc hr)⟩
  · next H1 H2 => match hl with
    | .balanced hl => exact ⟨_, .black hl hr⟩
    | .redred _ (c₁ := black) (c₂ := black) ha hb => exact ⟨_, .black (.red ha hb) hr⟩
    | .redred _ (c₁ := red) (.red ..) _ => cases H1 _ _ _ _ _ rfl
    | .redred _ (c₂ := red) _ (.red ..) => cases H2 _ _ _ _ _ rfl

/-- The `balance2` function repairs the balance invariant when the second argument is red-red. -/
protected theorem RedRed.balance2 {l : RBNode α} {v : α} {r : RBNode α}
    (hl : l.Balanced c n) (hr : r.RedRed p n) : ∃ c, (balance2 l v r).Balanced c (n + 1) := by
  unfold balance2; split
  · have .redred _ (.red ha hb) hc := hr; exact ⟨_, .red (.black hl ha) (.black hb hc)⟩
  · have .redred _ ha (.red hb hc) := hr; exact ⟨_, .red (.black hl ha) (.black hb hc)⟩
  · next H1 H2 => match hr with
    | .balanced hr => exact ⟨_, .black hl hr⟩
    | .redred _ (c₁ := black) (c₂ := black) ha hb => exact ⟨_, .black hl (.red ha hb)⟩
    | .redred _ (c₁ := red) (.red ..) _ => cases H1 _ _ _ _ _ rfl
    | .redred _ (c₂ := red) _ (.red ..) => cases H2 _ _ _ _ _ rfl

/-- The `balance1` function does nothing if the first argument is already balanced. -/
theorem balance1_eq {l : RBNode α} {v : α} {r : RBNode α}
    (hl : l.Balanced c n) : balance1 l v r = node black l v r := by
  unfold balance1; split <;> first | rfl | match hl with.

/-- The `balance2` function does nothing if the second argument is already balanced. -/
theorem balance2_eq {l : RBNode α} {v : α} {r : RBNode α}
    (hr : r.Balanced c n) : balance2 l v r = node black l v r := by
  unfold balance2; split <;> first | rfl | match hr with.

/-! ## insert -/

/--
The balance invariant of the `ins` function.
The result of inserting into the tree either yields a balanced tree,
or a tree which is almost balanced except that it has a red-red violation at the root.
-/
protected theorem Balanced.ins (cmp v) {t : RBNode α}
    (h : t.Balanced c n) : (ins cmp v t).RedRed (t.isRed = red) n := by
  induction h with
  | nil => exact .balanced (.red .nil .nil)
  | @red a n b x hl hr ihl ihr =>
    unfold ins; split
    · match ins cmp v a, ihl with
      | _, .balanced .nil => exact .balanced (.red .nil hr)
      | _, .balanced (.red ha hb) => exact .redred rfl (.red ha hb) hr
      | _, .balanced (.black ha hb) => exact .balanced (.red (.black ha hb) hr)
      | _, .redred h .. => cases hl <;> cases h
    · match ins cmp v b, ihr with
      | _, .balanced .nil => exact .balanced (.red hl .nil)
      | _, .balanced (.red ha hb) => exact .redred rfl hl (.red ha hb)
      | _, .balanced (.black ha hb) => exact .balanced (.red hl (.black ha hb))
      | _, .redred h .. => cases hr <;> cases h
    · exact .balanced (.red hl hr)
  | @black a ca n b cb x hl hr ihl ihr =>
    unfold ins; split
    · exact have ⟨c, h⟩ := ihl.balance1 hr; .balanced h
    · exact have ⟨c, h⟩ := ihr.balance2 hl; .balanced h
    · exact .balanced (.black hl hr)

/--
The `insert` function is balanced if the input is balanced.
(We lose track of both the color and the black-height of the result,
so this is only suitable for use on the root of the tree.)
-/
theorem Balanced.insert {t : RBNode α} (h : t.Balanced c n) :
    ∃ c' n', (insert cmp t v).Balanced c' n' := by
  unfold insert; match ins cmp v t, h.ins cmp v with
  | _, .balanced h => split <;> [exact ⟨_, h.setBlack⟩, exact ⟨_, _, h⟩]
  | _, .redred _ ha hb => have .node red .. := t; exact ⟨_, _, .black ha hb⟩

protected theorem All.setRed {t : RBNode α} (h : t.All p) : (setRed t).All p := by
  unfold setRed; split <;> simp_all

/-- The `setRed` function preserves the ordering invariants. -/
protected theorem Ordered.setRed {t : RBNode α} : (setRed t).Ordered cmp ↔ t.Ordered cmp := by
  unfold setRed; split <;> simp [Ordered]

protected theorem All.balLeft
    (hl : l.All p) (hv : p v) (hr : r.All p) : (balLeft l v r).All p := by
  unfold balLeft; split <;> simp_all; split <;> simp_all [All.setRed]

/-- The `balLeft` function preserves the ordering invariants. -/
protected theorem Ordered.balLeft {l : RBNode α} {v : α} {r : RBNode α}
    (lv : l.All (cmpLT cmp · v)) (vr : r.All (cmpLT cmp v ·))
    (hl : l.Ordered cmp) (hr : r.Ordered cmp) : (balLeft l v r).Ordered cmp := by
  unfold balLeft; split
  · exact ⟨lv, vr, hl, hr⟩
  split
  · exact hl.balance2 lv vr hr
  · have ⟨vy, va, _⟩ := vr.2.1; have ⟨⟨yz, _, bz⟩, zc, ⟨ay, yb, ha, hb⟩, hc⟩ := hr
    exact ⟨⟨vy, vy.trans_r lv, ay⟩, balance2_All.2 ⟨yz, yb, (yz.trans_l zc).setRed⟩,
      ⟨lv, va, hl, ha⟩, hb.balance2 bz zc.setRed (Ordered.setRed.2 hc)⟩
  · exact ⟨lv, vr, hl, hr⟩

/-- The balancing properties of the `balLeft` function. -/
protected theorem Balanced.balLeft (hl : l.RedRed True n) (hr : r.Balanced cr (n + 1)) :
    (balLeft l v r).RedRed (cr = red) (n + 1) := by
  unfold balLeft; split
  · next a x b => exact
    let ⟨ca, cb, ha, hb⟩ := hl.of_red
    match cr with
    | red => .redred rfl (.black ha hb) hr
    | black => .balanced (.red (.black ha hb) hr)
  · next H => exact match hl with
    | .redred .. => nomatch H _ _ _ rfl
    | .balanced hl => match hr with
      | .black ha hb =>
        let ⟨c, h⟩ := RedRed.balance2 hl (.redred trivial ha hb); .balanced h
      | .red (.black ha hb) (.black hc hd) =>
        let ⟨c, h⟩ := RedRed.balance2 hb (.redred trivial hc hd); .redred rfl (.black hl ha) h

protected theorem All.balRight
    (hl : l.All p) (hv : p v) (hr : r.All p) : (balRight l v r).All p := by
  unfold balRight; split <;> simp_all; split <;> simp_all [All.setRed]

/-- The `balRight` function preserves the ordering invariants. -/
protected theorem Ordered.balRight {l : RBNode α} {v : α} {r : RBNode α}
    (lv : l.All (cmpLT cmp · v)) (vr : r.All (cmpLT cmp v ·))
    (hl : l.Ordered cmp) (hr : r.Ordered cmp) : (balRight l v r).Ordered cmp := by
  unfold balRight; split
  · exact ⟨lv, vr, hl, hr⟩
  split
  · exact hl.balance1 lv vr hr
  · have ⟨yv, _, cv⟩ := lv.2.2; have ⟨ax, ⟨xy, xb, _⟩, ha, by_, yc, hb, hc⟩ := hl
    exact ⟨balance1_All.2 ⟨xy, (xy.trans_r ax).setRed, by_⟩, ⟨yv, yc, yv.trans_l vr⟩,
      (Ordered.setRed.2 ha).balance1 ax.setRed xb hb, cv, vr, hc, hr⟩
  · exact ⟨lv, vr, hl, hr⟩

/-- The balancing properties of the `balRight` function. -/
protected theorem Balanced.balRight (hl : l.Balanced cl (n + 1)) (hr : r.RedRed True n) :
    (balRight l v r).RedRed (cl = red) (n + 1) := by
  unfold balRight; split
  · next b y c => exact
    let ⟨cb, cc, hb, hc⟩ := hr.of_red
    match cl with
    | red => .redred rfl hl (.black hb hc)
    | black => .balanced (.red hl (.black hb hc))
  · next H => exact match hr with
    | .redred .. => nomatch H _ _ _ rfl
    | .balanced hr => match hl with
      | .black hb hc =>
        let ⟨c, h⟩ := RedRed.balance1 (.redred trivial hb hc) hr; .balanced h
      | .red (.black ha hb) (.black hc hd) =>
        let ⟨c, h⟩ := RedRed.balance1 (.redred trivial ha hb) hc; .redred rfl h (.black hd hr)

protected theorem All.append (hl : l.All p) (hr : r.All p) : (append l r).All p := by
  unfold append; split <;> simp [*]
  · have ⟨hx, ha, hb⟩ := hl; have ⟨hy, hc, hd⟩ := hr
    have := hb.append hc; split <;> simp_all
  · have ⟨hx, ha, hb⟩ := hl; have ⟨hy, hc, hd⟩ := hr
    have := hb.append hc; split <;> simp_all [All.balLeft]
  · simp_all [hl.append hr.2.1]
  · simp_all [hl.2.2.append hr]
termination_by _ => l.size + r.size

/-- The `append` function preserves the ordering invariants. -/
protected theorem Ordered.append {l : RBNode α} {v : α} {r : RBNode α}
    (lv : l.All (cmpLT cmp · v)) (vr : r.All (cmpLT cmp v ·))
    (hl : l.Ordered cmp) (hr : r.Ordered cmp) : (append l r).Ordered cmp := by
  unfold append; split
  · exact hr
  · exact hl
  · have ⟨xv, _, bv⟩ := lv; have ⟨ax, xb, ha, hb⟩ := hl
    have ⟨vy, vc, _⟩ := vr; have ⟨cy, yd, hc, hd⟩ := hr
    have : _ ∧ _ ∧ _ := ⟨hb.append bv vc hc, xb.append (xv.trans_l vc), (vy.trans_r bv).append cy⟩
    split
    · next H =>
      have ⟨⟨b'z, c'z, hb', hc'⟩, ⟨xz, xb', _⟩, zy, _, c'y⟩ := H ▸ this
      have az := xz.trans_r ax; have zd := zy.trans_l yd
      exact ⟨⟨xz, az, b'z⟩, ⟨zy, c'z, zd⟩, ⟨ax, xb', ha, hb'⟩, c'y, yd, hc', hd⟩
    · have ⟨hbc, xbc, bcy⟩ := this; have xy := xv.trans vy
      exact ⟨ax, ⟨xy, xbc, xy.trans_l yd⟩, ha, bcy, yd, hbc, hd⟩
  · have ⟨xv, _, bv⟩ := lv; have ⟨ax, xb, ha, hb⟩ := hl
    have ⟨vy, vc, _⟩ := vr; have ⟨cy, yd, hc, hd⟩ := hr
    have : _ ∧ _ ∧ _ := ⟨hb.append bv vc hc, xb.append (xv.trans_l vc), (vy.trans_r bv).append cy⟩
    split
    · next H =>
      have ⟨⟨b'z, c'z, hb', hc'⟩, ⟨xz, xb', _⟩, zy, _, c'y⟩ := H ▸ this
      have az := xz.trans_r ax; have zd := zy.trans_l yd
      exact ⟨⟨xz, az, b'z⟩, ⟨zy, c'z, zd⟩, ⟨ax, xb', ha, hb'⟩, c'y, yd, hc', hd⟩
    · have ⟨hbc, xbc, bcy⟩ := this; have xy := xv.trans vy
      exact ha.balLeft ax ⟨xy, xbc, xy.trans_l yd⟩ ⟨bcy, yd, hbc, hd⟩
  · have ⟨vx, vb, _⟩ := vr; have ⟨bx, yc, hb, hc⟩ := hr
    exact ⟨(vx.trans_r lv).append bx, yc, hl.append lv vb hb, hc⟩
  · have ⟨xv, _, bv⟩ := lv; have ⟨ax, xb, ha, hb⟩ := hl
    exact ⟨ax, xb.append (xv.trans_l vr), ha, hb.append bv vr hr⟩
termination_by _ => l.size + r.size

/-- The balance properties of the `append` function. -/
protected theorem Balanced.append {l r : RBNode α}
    (hl : l.Balanced c₁ n) (hr : r.Balanced c₂ n) :
    (l.append r).RedRed (c₁ = black → c₂ ≠ black) n := by
  unfold append; split
  · exact .balanced hr
  · exact .balanced hl
  · next b c _ _ =>
    have .red ha hb := hl; have .red hc hd := hr
    have ⟨_, IH⟩ := (hb.append hc).of_false (· rfl rfl); split
    · next e =>
      have .red hb' hc' := e ▸ IH
      exact .redred (fun.) (.red ha hb') (.red hc' hd)
    · next bcc _ H =>
      match bcc, append b c, IH, H with
      | black, _, IH, _ => exact .redred (fun.) ha (.red IH hd)
      | red, _, .red .., H => cases H _ _ _ rfl
  · next b c _ _ =>
    have .black ha hb := hl; have .black hc hd := hr
    have IH := hb.append hc; split
    · next e => match e ▸ IH with
      | .balanced (.red hb' hc') | .redred _ hb' hc' =>
        exact .balanced (.red (.black ha hb') (.black hc' hd))
    · next H =>
      match append b c, IH, H with
      | bc, .balanced hbc, _ =>
        unfold balLeft; split
        · have .red ha' hb' := ha
          exact .balanced (.red (.black ha' hb') (.black hbc hd))
        · exact have ⟨c, h⟩ := RedRed.balance2 ha (.redred trivial hbc hd); .balanced h
      | _, .redred .., H => cases H _ _ _ rfl
  · have .red hc hd := hr; have IH := hl.append hc
    have .black ha hb := hl; have ⟨c, IH⟩ := IH.of_false (· rfl rfl)
    exact .redred (fun.) IH hd
  · have .red ha hb := hl; have IH := hb.append hr
    have .black hc hd := hr; have ⟨c, IH⟩ := IH.of_false (· rfl rfl)
    exact .redred (fun.) ha IH
termination_by _ => l.size + r.size

/-! ## erase -/

/--
The invariant of the `del` function.
* If the input tree is black, then the result of deletion is a red-red tree with
  black-height lowered by 1.
* If the input tree is red or nil, then the result of deletion is a balanced tree with
  some color and the same black-height.
-/
def DelProp (p : RBColor) (t : RBNode α) (n : Nat) : Prop :=
  match p with
  | black => ∃ n', n = n' + 1 ∧ RedRed True t n'
  | red => ∃ c, Balanced t c n

/-- The `DelProp` property is a strengthened version of the red-red invariant. -/
theorem DelProp.redred (h : DelProp c t n) : ∃ n', RedRed (c = black) t n' := by
  unfold DelProp at h
  exact match c, h with
  | red, ⟨_, h⟩ => ⟨_, .balanced h⟩
  | black, ⟨_, _, h⟩ => ⟨_, h.imp fun _ => rfl⟩

protected theorem All.del : ∀ {t : RBNode α}, t.All p → (del cut t).All p
  | .nil, h => h
  | .node .., ⟨hy, ha, hb⟩ => by
    unfold del; split
    · split
      · exact ha.del.balLeft hy hb
      · exact ⟨hy, ha.del, hb⟩
    · split
      · exact ha.balRight hy hb.del
      · exact ⟨hy, ha, hb.del⟩
    · exact ha.append hb

/-- The `del` function preserves the ordering invariants. -/
protected theorem Ordered.del : ∀ {t : RBNode α}, t.Ordered cmp → (del cut t).Ordered cmp
  | .nil, _ => ⟨⟩
  | .node _ a y b, ⟨ay, yb, ha, hb⟩ => by
    unfold del; split
    · split
      · exact ha.del.balLeft ay.del yb hb
      · exact ⟨ay.del, yb, ha.del, hb⟩
    · split
      · exact ha.balRight ay yb.del hb.del
      · exact ⟨ay, yb.del, ha, hb.del⟩
    · exact ha.append ay yb hb

/-- The `del` function has the `DelProp` property. -/
protected theorem Balanced.del {t : RBNode α} (h : t.Balanced c n) :
    (t.del cut).DelProp t.isBlack n := by
  induction h with
  | nil => exact ⟨_, .nil⟩
  | @black a _ n b _ _ ha hb iha ihb =>
    refine ⟨_, rfl, ?_⟩
    unfold del; split
    · exact match a, n, iha with
      | .nil, _, ⟨c, ha⟩ | .node red .., _, ⟨c, ha⟩ => .redred ⟨⟩ ha hb
      | .node black .., _, ⟨n, rfl, ha⟩ => (hb.balLeft ha).imp fun _ => ⟨⟩
    · exact match b, n, ihb with
      | .nil, _, ⟨c, hb⟩ | .node .red .., _, ⟨c, hb⟩ => .redred ⟨⟩ ha hb
      | .node black .., _, ⟨n, rfl, hb⟩ => (ha.balRight hb).imp fun _ => ⟨⟩
    · exact (ha.append hb).imp fun _ => ⟨⟩
  | @red a n b _ ha hb iha ihb =>
    unfold del; split
    · exact match a, n, iha with
      | .nil, _, _ => ⟨_, .red ha hb⟩
      | .node black .., _, ⟨n, rfl, ha⟩ => (hb.balLeft ha).of_false (fun.)
    · exact match b, n, ihb with
      | .nil, _, _ => ⟨_, .red ha hb⟩
      | .node black .., _, ⟨n, rfl, hb⟩ => (ha.balRight hb).of_false (fun.)
    · exact (ha.append hb).of_false (· rfl rfl)

/-- The `erase` function preserves the ordering invariants. -/
protected theorem Ordered.erase {t : RBNode α} (h : t.Ordered cmp) : (erase cut t).Ordered cmp :=
  Ordered.setBlack.2 h.del

/-- The `erase` function preserves the balance invariants. -/
protected theorem Balanced.erase {t : RBNode α}
    (h : t.Balanced c n) : ∃ n, (t.erase cut).Balanced black n :=
  have ⟨_, h⟩ := h.del.redred; h.setBlack

/-- The well-formedness invariant implies the ordering and balance properties. -/
theorem WF.out {t : RBNode α} (h : t.WF cmp) : t.Ordered cmp ∧ ∃ c n, t.Balanced c n := by
  induction h with
  | mk o h => exact ⟨o, _, _, h⟩
  | insert _ ih => have ⟨o, _, _, h⟩ := ih; exact ⟨o.insert, h.insert⟩
  | erase _ ih => have ⟨o, _, _, h⟩ := ih; exact ⟨o.erase, _, h.erase⟩

/--
The well-formedness invariant for a red-black tree is exactly the `mk` constructor,
because the other constructors of `WF` are redundant.
-/
@[simp] theorem WF_iff {t : RBNode α} : t.WF cmp ↔ t.Ordered cmp ∧ ∃ c n, t.Balanced c n :=
  ⟨fun h => h.out, fun ⟨o, _, _, h⟩ => .mk o h⟩

/-- The `map` function preserves the balance invariants. -/
protected theorem Balanced.map {t : RBNode α} : t.Balanced c n → (t.map f).Balanced c n
  | .nil => .nil
  | .red hl hr => .red hl.map hr.map
  | .black hl hr => .black hl.map hr.map

/-- The property of a map function `f` which ensures the `map` operation is valid. -/
class IsMonotone (cmpα cmpβ) (f : α → β) : Prop where
  /-- If `x < y` then `f x < f y`. -/
  lt_mono : cmpLT cmpα x y → cmpLT cmpβ (f x) (f y)

/-- Sufficient condition for `map` to preserve an `All` quantifier. -/
protected theorem All.map {f : α → β} (H : ∀ {x}, p x → q (f x)) :
    ∀ {t : RBNode α}, t.All p → (t.map f).All q
  | nil, _ => ⟨⟩
  | node .., ⟨hx, ha, hb⟩ => ⟨H hx, ha.map H, hb.map H⟩

/-- The `map` function preserves the order invariants if `f` is monotone. -/
protected theorem Ordered.map (f : α → β) [IsMonotone cmpα cmpβ f] :
    ∀ {t : RBNode α}, t.Ordered cmpα → (t.map f).Ordered cmpβ
  | nil, _ => ⟨⟩
  | node _ a x b, ⟨ax, xb, ha, hb⟩ => by
    refine ⟨ax.map ?_, xb.map ?_, ha.map f, hb.map f⟩ <;> exact IsMonotone.lt_mono

end RBNode

namespace RBSet
export RBNode (IsMonotone)

/--
`O(n)`. Map a function on every value in the tree.
This requires `IsMonotone` on the function in order to preserve the order invariant.
-/
@[inline] def map (f : α → β) [IsMonotone cmpα cmpβ f] (t : RBSet α cmpα) : RBSet β cmpβ :=
  ⟨t.1.map f, have ⟨h₁, _, _, h₂⟩ := t.2.out; .mk (h₁.map _) h₂.map⟩

end RBSet

namespace RBMap
export RBNode (IsMonotone)

namespace Imp

/--
Applies `f` to the second component.
We extract this as a function so that `IsMonotone (mapSnd f)` can be an instance.
-/
@[inline] def mapSnd (f : α → β → γ) := fun (a, b) => (a, f a b)

instance (cmp : α → α → Ordering) (f : α → β → γ) :
    IsMonotone (byKey Prod.fst cmp) (byKey Prod.fst cmp) (mapSnd f) where
  lt_mono | ⟨h⟩ => ⟨@fun _ => @h {
    symm := fun (a₁, b₁) (a₂, b₂) =>
      OrientedCmp.symm (cmp := byKey Prod.fst cmp) (a₁, f a₁ b₁) (a₂, f a₂ b₂)
    le_trans := @fun (a₁, b₁) (a₂, b₂) (a₃, b₃) =>
      TransCmp.le_trans (cmp := byKey Prod.fst cmp)
        (x := (a₁, f a₁ b₁)) (y := (a₂, f a₂ b₂)) (z := (a₃, f a₃ b₃))
  }⟩

end Imp

/-- `O(n)`. Map a function on the values in the map. -/
def mapVal (f : α → β → γ) (t : RBMap α β cmp) : RBMap α γ cmp := t.map (Imp.mapSnd f)

end RBMap