/-
Copyright (c) 2017 Mario Carneiro. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro, Jeremy Avigad, Simon Hudon
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.data.rel
import Mathlib.PostPort

universes u l u_1 u_2 u_3

namespace Mathlib

/-- `roption α` is the type of "partial values" of type `α`. It
  is similar to `option α` except the domain condition can be an
  arbitrary proposition, not necessarily decidable. -/
structure roption (α : Type u) where
  dom : Prop
  get : dom → α

namespace roption

/-- Convert an `roption α` with a decidable domain to an option -/
def to_option {α : Type u_1} (o : roption α) [Decidable (dom o)] : Option α :=
  dite (dom o) (fun (h : dom o) => some (get o h)) fun (h : ¬dom o) => none

/-- `roption` extensionality -/
theorem ext' {α : Type u_1} {o : roption α} {p : roption α} (H1 : dom o ↔ dom p)
    (H2 : ∀ (h₁ : dom o) (h₂ : dom p), get o h₁ = get p h₂) : o = p :=
  sorry

/-- `roption` eta expansion -/
@[simp] theorem eta {α : Type u_1} (o : roption α) : (mk (dom o) fun (h : dom o) => get o h) = o :=
  sorry

/-- `a ∈ o` means that `o` is defined and equal to `a` -/
protected def mem {α : Type u_1} (a : α) (o : roption α) := ∃ (h : dom o), get o h = a

protected instance has_mem {α : Type u_1} : has_mem α (roption α) := has_mem.mk roption.mem

theorem mem_eq {α : Type u_1} (a : α) (o : roption α) : a ∈ o = ∃ (h : dom o), get o h = a := rfl

theorem dom_iff_mem {α : Type u_1} {o : roption α} : dom o ↔ ∃ (y : α), y ∈ o := sorry

theorem get_mem {α : Type u_1} {o : roption α} (h : dom o) : get o h ∈ o := Exists.intro h rfl

/-- `roption` extensionality -/
theorem ext {α : Type u_1} {o : roption α} {p : roption α} (H : ∀ (a : α), a ∈ o ↔ a ∈ p) : o = p :=
  sorry

/-- The `none` value in `roption` has a `false` domain and an empty function. -/
def none {α : Type u_1} : roption α := mk False False._oldrec

protected instance inhabited {α : Type u_1} : Inhabited (roption α) := { default := none }

@[simp] theorem not_mem_none {α : Type u_1} (a : α) : ¬a ∈ none :=
  fun (h : a ∈ none) => Exists.fst h

/-- The `some a` value in `roption` has a `true` domain and the
  function returns `a`. -/
def some {α : Type u_1} (a : α) : roption α := mk True fun (_x : True) => a

theorem mem_unique {α : Type u_1} : relator.left_unique has_mem.mem := sorry

theorem get_eq_of_mem {α : Type u_1} {o : roption α} {a : α} (h : a ∈ o) (h' : dom o) :
    get o h' = a :=
  mem_unique (Exists.intro h' rfl) h

@[simp] theorem get_some {α : Type u_1} {a : α} (ha : dom (some a)) : get (some a) ha = a := rfl

theorem mem_some {α : Type u_1} (a : α) : a ∈ some a := Exists.intro trivial rfl

@[simp] theorem mem_some_iff {α : Type u_1} {a : α} {b : α} : b ∈ some a ↔ b = a := sorry

theorem eq_some_iff {α : Type u_1} {a : α} {o : roption α} : o = some a ↔ a ∈ o := sorry

theorem eq_none_iff {α : Type u_1} {o : roption α} : o = none ↔ ∀ (a : α), ¬a ∈ o := sorry

theorem eq_none_iff' {α : Type u_1} {o : roption α} : o = none ↔ ¬dom o :=
  { mp := fun (e : o = none) => Eq.symm e ▸ id,
    mpr := fun (h : ¬dom o) => iff.mpr eq_none_iff fun (a : α) (h' : a ∈ o) => h (Exists.fst h') }

theorem some_ne_none {α : Type u_1} (x : α) : some x ≠ none :=
  id
    fun (h : some x = none) =>
      id (eq.mpr (id (Eq._oldrec (Eq.refl (dom none)) (Eq.symm h))) trivial)

theorem ne_none_iff {α : Type u_1} {o : roption α} : o ≠ none ↔ ∃ (x : α), o = some x := sorry

theorem eq_none_or_eq_some {α : Type u_1} (o : roption α) : o = none ∨ ∃ (x : α), o = some x :=
  sorry

@[simp] theorem some_inj {α : Type u_1} {a : α} {b : α} : some a = some b ↔ a = b :=
  function.injective.eq_iff
    fun (a b : α) (h : some a = some b) => congr_fun (eq_of_heq (and.right (mk.inj h))) trivial

@[simp] theorem some_get {α : Type u_1} {a : roption α} (ha : dom a) : some (get a ha) = a :=
  Eq.symm (iff.mpr eq_some_iff (Exists.intro ha rfl))

theorem get_eq_iff_eq_some {α : Type u_1} {a : roption α} {ha : dom a} {b : α} :
    get a ha = b ↔ a = some b :=
  sorry

protected instance none_decidable {α : Type u_1} : Decidable (dom none) := decidable.false

protected instance some_decidable {α : Type u_1} (a : α) : Decidable (dom (some a)) :=
  decidable.true

def get_or_else {α : Type u_1} (a : roption α) [Decidable (dom a)] (d : α) : α :=
  dite (dom a) (fun (ha : dom a) => get a ha) fun (ha : ¬dom a) => d

@[simp] theorem get_or_else_none {α : Type u_1} (d : α) : get_or_else none d = d := dif_neg id

@[simp] theorem get_or_else_some {α : Type u_1} (a : α) (d : α) : get_or_else (some a) d = a :=
  dif_pos trivial

@[simp] theorem mem_to_option {α : Type u_1} {o : roption α} [Decidable (dom o)] {a : α} :
    a ∈ to_option o ↔ a ∈ o :=
  sorry

/-- Convert an `option α` into an `roption α` -/
def of_option {α : Type u_1} : Option α → roption α := sorry

@[simp] theorem mem_of_option {α : Type u_1} {a : α} {o : Option α} : a ∈ of_option o ↔ a ∈ o :=
  sorry

@[simp] theorem of_option_dom {α : Type u_1} (o : Option α) :
    dom (of_option o) ↔ ↥(option.is_some o) :=
  sorry

theorem of_option_eq_get {α : Type u_1} (o : Option α) :
    of_option o = mk (↥(option.is_some o)) option.get :=
  sorry

protected instance has_coe {α : Type u_1} : has_coe (Option α) (roption α) := has_coe.mk of_option

@[simp] theorem mem_coe {α : Type u_1} {a : α} {o : Option α} : a ∈ ↑o ↔ a ∈ o := mem_of_option

@[simp] theorem coe_none {α : Type u_1} : ↑none = none := rfl

@[simp] theorem coe_some {α : Type u_1} (a : α) : ↑(some a) = some a := rfl

protected theorem induction_on {α : Type u_1} {P : roption α → Prop} (a : roption α)
    (hnone : P none) (hsome : ∀ (a : α), P (some a)) : P a :=
  or.elim (classical.em (dom a)) (fun (h : dom a) => some_get h ▸ hsome (get a h))
    fun (h : ¬dom a) => Eq.symm (iff.mpr eq_none_iff' h) ▸ hnone

protected instance of_option_decidable {α : Type u_1} (o : Option α) :
    Decidable (dom (of_option o)) :=
  sorry

@[simp] theorem to_of_option {α : Type u_1} (o : Option α) : to_option (of_option o) = o :=
  option.cases_on o (Eq.refl (to_option (of_option none)))
    fun (o : α) => Eq.refl (to_option (of_option (some o)))

@[simp] theorem of_to_option {α : Type u_1} (o : roption α) [Decidable (dom o)] :
    of_option (to_option o) = o :=
  ext fun (a : α) => iff.trans mem_of_option mem_to_option

def equiv_option {α : Type u_1} : roption α ≃ Option α :=
  equiv.mk (fun (o : roption α) => to_option o) of_option sorry sorry

protected instance order_bot {α : Type u_1} : order_bot (roption α) :=
  order_bot.mk none (fun (x y : roption α) => ∀ (i : α), i ∈ x → i ∈ y)
    (partial_order.lt._default fun (x y : roption α) => ∀ (i : α), i ∈ x → i ∈ y) sorry sorry sorry
    sorry

protected instance preorder {α : Type u_1} : preorder (roption α) :=
  partial_order.to_preorder (roption α)

theorem le_total_of_le_of_le {α : Type u_1} {x : roption α} {y : roption α} (z : roption α)
    (hx : x ≤ z) (hy : y ≤ z) : x ≤ y ∨ y ≤ x :=
  sorry

/-- `assert p f` is a bind-like operation which appends an additional condition
  `p` to the domain and uses `f` to produce the value. -/
def assert {α : Type u_1} (p : Prop) (f : p → roption α) : roption α :=
  mk (∃ (h : p), dom (f h)) fun (ha : ∃ (h : p), dom (f h)) => get (f sorry) sorry

/-- The bind operation has value `g (f.get)`, and is defined when all the
  parts are defined. -/
protected def bind {α : Type u_1} {β : Type u_2} (f : roption α) (g : α → roption β) : roption β :=
  assert (dom f) fun (b : dom f) => g (get f b)

/-- The map operation for `roption` just maps the value and maintains the same domain. -/
def map {α : Type u_1} {β : Type u_2} (f : α → β) (o : roption α) : roption β :=
  mk (dom o) (f ∘ get o)

theorem mem_map {α : Type u_1} {β : Type u_2} (f : α → β) {o : roption α} {a : α} :
    a ∈ o → f a ∈ map f o :=
  sorry

@[simp] theorem mem_map_iff {α : Type u_1} {β : Type u_2} (f : α → β) {o : roption α} {b : β} :
    b ∈ map f o ↔ ∃ (a : α), ∃ (H : a ∈ o), f a = b :=
  sorry

@[simp] theorem map_none {α : Type u_1} {β : Type u_2} (f : α → β) : map f none = none := sorry

@[simp] theorem map_some {α : Type u_1} {β : Type u_2} (f : α → β) (a : α) :
    map f (some a) = some (f a) :=
  iff.mpr eq_some_iff (mem_map f (mem_some a))

theorem mem_assert {α : Type u_1} {p : Prop} {f : p → roption α} {a : α} (h : p) :
    a ∈ f h → a ∈ assert p f :=
  sorry

@[simp] theorem mem_assert_iff {α : Type u_1} {p : Prop} {f : p → roption α} {a : α} :
    a ∈ assert p f ↔ ∃ (h : p), a ∈ f h :=
  sorry

theorem assert_pos {α : Type u_1} {p : Prop} {f : p → roption α} (h : p) : assert p f = f h := sorry

theorem assert_neg {α : Type u_1} {p : Prop} {f : p → roption α} (h : ¬p) : assert p f = none :=
  sorry

theorem mem_bind {α : Type u_1} {β : Type u_2} {f : roption α} {g : α → roption β} {a : α} {b : β} :
    a ∈ f → b ∈ g a → b ∈ roption.bind f g :=
  sorry

@[simp] theorem mem_bind_iff {α : Type u_1} {β : Type u_2} {f : roption α} {g : α → roption β}
    {b : β} : b ∈ roption.bind f g ↔ ∃ (a : α), ∃ (H : a ∈ f), b ∈ g a :=
  sorry

@[simp] theorem bind_none {α : Type u_1} {β : Type u_2} (f : α → roption β) :
    roption.bind none f = none :=
  sorry

@[simp] theorem bind_some {α : Type u_1} {β : Type u_2} (a : α) (f : α → roption β) :
    roption.bind (some a) f = f a :=
  sorry

theorem bind_some_eq_map {α : Type u_1} {β : Type u_2} (f : α → β) (x : roption α) :
    roption.bind x (some ∘ f) = map f x :=
  sorry

theorem bind_assoc {α : Type u_1} {β : Type u_2} {γ : Type u_3} (f : roption α) (g : α → roption β)
    (k : β → roption γ) :
    roption.bind (roption.bind f g) k = roption.bind f fun (x : α) => roption.bind (g x) k :=
  sorry

@[simp] theorem bind_map {α : Type u_1} {β : Type u_2} {γ : Type u_3} (f : α → β) (x : roption α)
    (g : β → roption γ) : roption.bind (map f x) g = roption.bind x fun (y : α) => g (f y) :=
  sorry

@[simp] theorem map_bind {α : Type u_1} {β : Type u_2} {γ : Type u_3} (f : α → roption β)
    (x : roption α) (g : β → γ) :
    map g (roption.bind x f) = roption.bind x fun (y : α) => map g (f y) :=
  sorry

theorem map_map {α : Type u_1} {β : Type u_2} {γ : Type u_3} (g : β → γ) (f : α → β)
    (o : roption α) : map g (map f o) = map (g ∘ f) o :=
  sorry

protected instance monad : Monad roption :=
  { toApplicative :=
      { toFunctor := { map := map, mapConst := fun (α β : Type u_1) => map ∘ function.const β },
        toPure := { pure := some },
        toSeq :=
          { seq :=
              fun (α β : Type u_1) (f : roption (α → β)) (x : roption α) =>
                roption.bind f fun (_x : α → β) => map _x x },
        toSeqLeft :=
          { seqLeft :=
              fun (α β : Type u_1) (a : roption α) (b : roption β) =>
                (fun (α β : Type u_1) (f : roption (α → β)) (x : roption α) =>
                    roption.bind f fun (_x : α → β) => map _x x)
                  β α (map (function.const β) a) b },
        toSeqRight :=
          { seqRight :=
              fun (α β : Type u_1) (a : roption α) (b : roption β) =>
                (fun (α β : Type u_1) (f : roption (α → β)) (x : roption α) =>
                    roption.bind f fun (_x : α → β) => map _x x)
                  β β (map (function.const α id) a) b } },
    toBind := { bind := roption.bind } }

protected instance is_lawful_monad : is_lawful_monad roption :=
  is_lawful_monad.mk bind_some bind_assoc

theorem map_id' {α : Type u_1} {f : α → α} (H : ∀ (x : α), f x = x) (o : roption α) : map f o = o :=
  eq.mpr (id (Eq._oldrec (Eq.refl (map f o = o)) ((fun (this : f = id) => this) (funext H))))
    (id_map o)

@[simp] theorem bind_some_right {α : Type u_1} (x : roption α) : roption.bind x some = x := sorry

@[simp] theorem pure_eq_some {α : Type u_1} (a : α) : pure a = some a := rfl

@[simp] theorem ret_eq_some {α : Type u_1} (a : α) : return a = some a := rfl

@[simp] theorem map_eq_map {α : Type u_1} {β : Type u_1} (f : α → β) (o : roption α) :
    f <$> o = map f o :=
  rfl

@[simp] theorem bind_eq_bind {α : Type u_1} {β : Type u_1} (f : roption α) (g : α → roption β) :
    f >>= g = roption.bind f g :=
  rfl

theorem bind_le {β : Type u_2} {α : Type u_2} (x : roption α) (f : α → roption β) (y : roption β) :
    x >>= f ≤ y ↔ ∀ (a : α), a ∈ x → f a ≤ y :=
  sorry

protected instance monad_fail : monad_fail roption :=
  monad_fail.mk fun (_x : Type u_1) (_x_1 : string) => none

/- `restrict p o h` replaces the domain of `o` with `p`, and is well defined when
  `p` implies `o` is defined. -/

def restrict {α : Type u_1} (p : Prop) (o : roption α) : (p → dom o) → roption α := sorry

@[simp] theorem mem_restrict {α : Type u_1} (p : Prop) (o : roption α) (h : p → dom o) (a : α) :
    a ∈ restrict p o h ↔ p ∧ a ∈ o :=
  sorry

/-- `unwrap o` gets the value at `o`, ignoring the condition.
  (This function is unsound.) -/
theorem assert_defined {α : Type u_1} {p : Prop} {f : p → roption α} (h : p) :
    dom (f h) → dom (assert p f) :=
  exists.intro

theorem bind_defined {α : Type u_1} {β : Type u_2} {f : roption α} {g : α → roption β} (h : dom f) :
    dom (g (get f h)) → dom (roption.bind f g) :=
  assert_defined

@[simp] theorem bind_dom {α : Type u_1} {β : Type u_2} {f : roption α} {g : α → roption β} :
    dom (roption.bind f g) ↔ ∃ (h : dom f), dom (g (get f h)) :=
  iff.rfl

end roption

/-- `pfun α β`, or `α →. β`, is the type of partial functions from
  `α` to `β`. It is defined as `α → roption β`. -/
def pfun (α : Type u_1) (β : Type u_2) := α → roption β

infixr:25 " →. " => Mathlib.pfun

namespace pfun

protected instance inhabited {α : Type u_1} {β : Type u_2} : Inhabited (α →. β) :=
  { default := fun (a : α) => roption.none }

/-- The domain of a partial function -/
def dom {α : Type u_1} {β : Type u_2} (f : α →. β) : set α :=
  set_of fun (a : α) => roption.dom (f a)

theorem mem_dom {α : Type u_1} {β : Type u_2} (f : α →. β) (x : α) :
    x ∈ dom f ↔ ∃ (y : β), y ∈ f x :=
  sorry

theorem dom_eq {α : Type u_1} {β : Type u_2} (f : α →. β) :
    dom f = set_of fun (x : α) => ∃ (y : β), y ∈ f x :=
  set.ext (mem_dom f)

/-- Evaluate a partial function -/
def fn {α : Type u_1} {β : Type u_2} (f : α →. β) (x : α) (h : dom f x) : β := roption.get (f x) h

/-- Evaluate a partial function to return an `option` -/
def eval_opt {α : Type u_1} {β : Type u_2} (f : α →. β) [D : decidable_pred (dom f)] (x : α) :
    Option β :=
  roption.to_option (f x)

/-- Partial function extensionality -/
theorem ext' {α : Type u_1} {β : Type u_2} {f : α →. β} {g : α →. β}
    (H1 : ∀ (a : α), a ∈ dom f ↔ a ∈ dom g)
    (H2 : ∀ (a : α) (p : dom f a) (q : dom g a), fn f a p = fn g a q) : f = g :=
  funext fun (a : α) => roption.ext' (H1 a) (H2 a)

theorem ext {α : Type u_1} {β : Type u_2} {f : α →. β} {g : α →. β}
    (H : ∀ (a : α) (b : β), b ∈ f a ↔ b ∈ g a) : f = g :=
  funext fun (a : α) => roption.ext (H a)

/-- Turn a partial function into a function out of a subtype -/
def as_subtype {α : Type u_1} {β : Type u_2} (f : α →. β) (s : ↥(dom f)) : β := fn f ↑s sorry

/-- The set of partial functions `α →. β` is equivalent to
the set of pairs `(p : α → Prop, f : subtype p → β)`. -/
def equiv_subtype {α : Type u_1} {β : Type u_2} :
    (α →. β) ≃ sigma fun (p : α → Prop) => Subtype p → β :=
  equiv.mk (fun (f : α →. β) => sigma.mk (fun (a : α) => roption.dom (f a)) (as_subtype f))
    (fun (f : sigma fun (p : α → Prop) => Subtype p → β) (x : α) =>
      roption.mk (sigma.fst f x) fun (h : sigma.fst f x) => sigma.snd f { val := x, property := h })
    sorry sorry

theorem as_subtype_eq_of_mem {α : Type u_1} {β : Type u_2} {f : α →. β} {x : α} {y : β}
    (fxy : y ∈ f x) (domx : x ∈ dom f) : as_subtype f { val := x, property := domx } = y :=
  roption.mem_unique (roption.get_mem (as_subtype._proof_1 f { val := x, property := domx })) fxy

/-- Turn a total function into a partial function -/
protected def lift {α : Type u_1} {β : Type u_2} (f : α → β) : α →. β :=
  fun (a : α) => roption.some (f a)

protected instance has_coe {α : Type u_1} {β : Type u_2} : has_coe (α → β) (α →. β) :=
  has_coe.mk pfun.lift

@[simp] theorem lift_eq_coe {α : Type u_1} {β : Type u_2} (f : α → β) : pfun.lift f = ↑f := rfl

@[simp] theorem coe_val {α : Type u_1} {β : Type u_2} (f : α → β) (a : α) :
    coe f a = roption.some (f a) :=
  rfl

/-- The graph of a partial function is the set of pairs
  `(x, f x)` where `x` is in the domain of `f`. -/
def graph {α : Type u_1} {β : Type u_2} (f : α →. β) : set (α × β) :=
  set_of fun (p : α × β) => prod.snd p ∈ f (prod.fst p)

def graph' {α : Type u_1} {β : Type u_2} (f : α →. β) : rel α β := fun (x : α) (y : β) => y ∈ f x

/-- The range of a partial function is the set of values
  `f x` where `x` is in the domain of `f`. -/
def ran {α : Type u_1} {β : Type u_2} (f : α →. β) : set β :=
  set_of fun (b : β) => ∃ (a : α), b ∈ f a

/-- Restrict a partial function to a smaller domain. -/
def restrict {α : Type u_1} {β : Type u_2} (f : α →. β) {p : set α} (H : p ⊆ dom f) : α →. β :=
  fun (x : α) => roption.restrict (x ∈ p) (f x) H

@[simp] theorem mem_restrict {α : Type u_1} {β : Type u_2} {f : α →. β} {s : set α} (h : s ⊆ dom f)
    (a : α) (b : β) : b ∈ restrict f h a ↔ a ∈ s ∧ b ∈ f a :=
  sorry

def res {α : Type u_1} {β : Type u_2} (f : α → β) (s : set α) : α →. β :=
  restrict (pfun.lift f) (set.subset_univ s)

theorem mem_res {α : Type u_1} {β : Type u_2} (f : α → β) (s : set α) (a : α) (b : β) :
    b ∈ res f s a ↔ a ∈ s ∧ f a = b :=
  sorry

theorem res_univ {α : Type u_1} {β : Type u_2} (f : α → β) : res f set.univ = ↑f := rfl

theorem dom_iff_graph {α : Type u_1} {β : Type u_2} (f : α →. β) (x : α) :
    x ∈ dom f ↔ ∃ (y : β), (x, y) ∈ graph f :=
  roption.dom_iff_mem

theorem lift_graph {α : Type u_1} {β : Type u_2} {f : α → β} {a : α} {b : β} :
    (a, b) ∈ graph ↑f ↔ f a = b :=
  sorry

/-- The monad `pure` function, the total constant `x` function -/
protected def pure {α : Type u_1} {β : Type u_2} (x : β) : α →. β := fun (_x : α) => roption.some x

/-- The monad `bind` function, pointwise `roption.bind` -/
def bind {α : Type u_1} {β : Type u_2} {γ : Type u_3} (f : α →. β) (g : β → α →. γ) : α →. γ :=
  fun (a : α) => roption.bind (f a) fun (b : β) => g b a

/-- The monad `map` function, pointwise `roption.map` -/
def map {α : Type u_1} {β : Type u_2} {γ : Type u_3} (f : β → γ) (g : α →. β) : α →. γ :=
  fun (a : α) => roption.map f (g a)

protected instance monad {α : Type u_1} : Monad (pfun α) :=
  { toApplicative :=
      { toFunctor := { map := map, mapConst := fun (α_1 β : Type u_2) => map ∘ function.const β },
        toPure := { pure := pfun.pure },
        toSeq :=
          { seq :=
              fun (α_1 β : Type u_2) (f : α →. α_1 → β) (x : α →. α_1) =>
                bind f fun (_x : α_1 → β) => map _x x },
        toSeqLeft :=
          { seqLeft :=
              fun (α_1 β : Type u_2) (a : α →. α_1) (b : α →. β) =>
                (fun (α_2 β : Type u_2) (f : α →. α_2 → β) (x : α →. α_2) =>
                    bind f fun (_x : α_2 → β) => map _x x)
                  β α_1 (map (function.const β) a) b },
        toSeqRight :=
          { seqRight :=
              fun (α_1 β : Type u_2) (a : α →. α_1) (b : α →. β) =>
                (fun (α_2 β : Type u_2) (f : α →. α_2 → β) (x : α →. α_2) =>
                    bind f fun (_x : α_2 → β) => map _x x)
                  β β (map (function.const α_1 id) a) b } },
    toBind := { bind := bind } }

protected instance is_lawful_monad {α : Type u_1} : is_lawful_monad (pfun α) :=
  is_lawful_monad.mk
    (fun (β γ : Type u_2) (x : β) (f : β → α →. γ) =>
      funext fun (a : α) => roption.bind_some a (f x))
    fun (β γ δ : Type u_2) (f : α →. β) (g : β → α →. γ) (k : γ → α →. δ) =>
      funext fun (a : α) => roption.bind_assoc (f a) (fun (b : β) => g b a) fun (b : γ) => k b a

theorem pure_defined {α : Type u_1} {β : Type u_2} (p : set α) (x : β) : p ⊆ dom (pfun.pure x) :=
  set.subset_univ p

theorem bind_defined {α : Type u_1} {β : Type u_2} {γ : Type u_2} (p : set α) {f : α →. β}
    {g : β → α →. γ} (H1 : p ⊆ dom f) (H2 : ∀ (x : β), p ⊆ dom (g x)) : p ⊆ dom (f >>= g) :=
  fun (a : α) (ha : a ∈ p) => Exists.intro (H1 ha) (H2 (roption.get (f a) (H1 ha)) ha)

def fix {α : Type u_1} {β : Type u_2} (f : α →. β ⊕ α) : α →. β :=
  fun (a : α) =>
    roption.assert (acc (fun (x y : α) => sum.inr x ∈ f y) a)
      fun (h : acc (fun (x y : α) => sum.inr x ∈ f y) a) =>
        well_founded.fix_F
          (fun (a : α) (IH : (y : α) → sum.inr y ∈ f a → roption β) =>
            roption.assert (roption.dom (f a))
              fun (hf : roption.dom (f a)) =>
                (fun (_x : β ⊕ α) (e : roption.get (f a) hf = _x) =>
                    sum.cases_on _x
                      (fun (b : β) (e : roption.get (f a) hf = sum.inl b) => roption.some b)
                      (fun (a' : α) (e : roption.get (f a) hf = sum.inr a') => IH a' sorry) e)
                  (roption.get (f a) hf) sorry)
          a h

theorem dom_of_mem_fix {α : Type u_1} {β : Type u_2} {f : α →. β ⊕ α} {a : α} {b : β}
    (h : b ∈ fix f a) : roption.dom (f a) :=
  sorry

theorem mem_fix_iff {α : Type u_1} {β : Type u_2} {f : α →. β ⊕ α} {a : α} {b : β} :
    b ∈ fix f a ↔ sum.inl b ∈ f a ∨ ∃ (a' : α), sum.inr a' ∈ f a ∧ b ∈ fix f a' :=
  sorry

def fix_induction {α : Type u_1} {β : Type u_2} {f : α →. β ⊕ α} {b : β} {C : α → Sort u_3} {a : α}
    (h : b ∈ fix f a)
    (H : (a : α) → b ∈ fix f a → ((a' : α) → b ∈ fix f a' → sum.inr a' ∈ f a → C a') → C a) : C a :=
  sorry

end pfun

namespace pfun

def image {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set α) : set β := rel.image (graph' f) s

theorem image_def {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set α) :
    image f s = set_of fun (y : β) => ∃ (x : α), ∃ (H : x ∈ s), y ∈ f x :=
  rfl

theorem mem_image {α : Type u_1} {β : Type u_2} (f : α →. β) (y : β) (s : set α) :
    y ∈ image f s ↔ ∃ (x : α), ∃ (H : x ∈ s), y ∈ f x :=
  iff.rfl

theorem image_mono {α : Type u_1} {β : Type u_2} (f : α →. β) {s : set α} {t : set α} (h : s ⊆ t) :
    image f s ⊆ image f t :=
  rel.image_mono (graph' f) h

theorem image_inter {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set α) (t : set α) :
    image f (s ∩ t) ⊆ image f s ∩ image f t :=
  rel.image_inter (graph' f) s t

theorem image_union {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set α) (t : set α) :
    image f (s ∪ t) = image f s ∪ image f t :=
  rel.image_union (graph' f) s t

def preimage {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) : set α :=
  rel.preimage (fun (x : α) (y : β) => y ∈ f x) s

theorem preimage_def {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) :
    preimage f s = set_of fun (x : α) => ∃ (y : β), ∃ (H : y ∈ s), y ∈ f x :=
  rfl

theorem mem_preimage {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) (x : α) :
    x ∈ preimage f s ↔ ∃ (y : β), ∃ (H : y ∈ s), y ∈ f x :=
  iff.rfl

theorem preimage_subset_dom {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) :
    preimage f s ⊆ dom f :=
  sorry

theorem preimage_mono {α : Type u_1} {β : Type u_2} (f : α →. β) {s : set β} {t : set β}
    (h : s ⊆ t) : preimage f s ⊆ preimage f t :=
  rel.preimage_mono (fun (x : α) (y : β) => y ∈ f x) h

theorem preimage_inter {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) (t : set β) :
    preimage f (s ∩ t) ⊆ preimage f s ∩ preimage f t :=
  rel.preimage_inter (fun (x : α) (y : β) => y ∈ f x) s t

theorem preimage_union {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) (t : set β) :
    preimage f (s ∪ t) = preimage f s ∪ preimage f t :=
  rel.preimage_union (fun (x : α) (y : β) => y ∈ f x) s t

theorem preimage_univ {α : Type u_1} {β : Type u_2} (f : α →. β) : preimage f set.univ = dom f :=
  sorry

def core {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) : set α := rel.core (graph' f) s

theorem core_def {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) :
    core f s = set_of fun (x : α) => ∀ (y : β), y ∈ f x → y ∈ s :=
  rfl

theorem mem_core {α : Type u_1} {β : Type u_2} (f : α →. β) (x : α) (s : set β) :
    x ∈ core f s ↔ ∀ (y : β), y ∈ f x → y ∈ s :=
  iff.rfl

theorem compl_dom_subset_core {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) :
    dom fᶜ ⊆ core f s :=
  fun (x : α) (hx : x ∈ (dom fᶜ)) (y : β) (fxy : graph' f x y) =>
    absurd (iff.mpr (mem_dom f x) (Exists.intro y fxy)) hx

theorem core_mono {α : Type u_1} {β : Type u_2} (f : α →. β) {s : set β} {t : set β} (h : s ⊆ t) :
    core f s ⊆ core f t :=
  rel.core_mono (graph' f) h

theorem core_inter {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) (t : set β) :
    core f (s ∩ t) = core f s ∩ core f t :=
  rel.core_inter (graph' f) s t

theorem mem_core_res {α : Type u_1} {β : Type u_2} (f : α → β) (s : set α) (t : set β) (x : α) :
    x ∈ core (res f s) t ↔ x ∈ s → f x ∈ t :=
  sorry

theorem core_res {α : Type u_1} {β : Type u_2} (f : α → β) (s : set α) (t : set β) :
    core (res f s) t = sᶜ ∪ f ⁻¹' t :=
  sorry

theorem core_restrict {α : Type u_1} {β : Type u_2} (f : α → β) (s : set β) :
    core (↑f) s = f ⁻¹' s :=
  sorry

theorem preimage_subset_core {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) :
    preimage f s ⊆ core f s :=
  sorry

theorem preimage_eq {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) :
    preimage f s = core f s ∩ dom f :=
  sorry

theorem core_eq {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) :
    core f s = preimage f s ∪ (dom fᶜ) :=
  sorry

theorem preimage_as_subtype {α : Type u_1} {β : Type u_2} (f : α →. β) (s : set β) :
    as_subtype f ⁻¹' s = subtype.val ⁻¹' preimage f s :=
  sorry

end Mathlib


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

/-!
# Extra facts about `pprod`
-/

variables {α : Sort*} {β : Sort*}

@[simp] lemma pprod.mk.eta {p : pprod α β} : pprod.mk p.1 p.2 = p :=
pprod.cases_on p (λ a b, rfl)


/-
Copyright (c) 2020 Gabriel Ebner. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Gabriel Ebner
-/
import Mathlib.Tactic.Lint.Basic
import Mathlib.Tactic.OpenPrivate
import Mathlib.Util.LibraryNote
open Lean Meta

namespace Mathlib.Tactic.Lint

/-!
# Linter for simplification lemmas

This files defines several linters that prevent common mistakes when declaring simp lemmas:

 * `simpNF` checks that the left-hand side of a simp lemma is not simplified by a different lemma.
 * `simpVarHead` checks that the head symbol of the left-hand side is not a variable.
 * `simpComm` checks that commutativity lemmas are not marked as simplification lemmas.
-/

structure SimpTheoremInfo where
  hyps : Array Expr
  isConditional : Bool
  lhs : Expr
  rhs : Expr

def isConditionalHyps (eq : Expr) : List Expr → MetaM Bool
  | [] => pure false
  | h :: hs => do
    let ldecl ← getFVarLocalDecl h
    if !ldecl.binderInfo.isInstImplicit
        && !(← hs.anyM fun h' => do pure $ (← inferType h').containsFVar h.fvarId!)
        && !eq.containsFVar h.fvarId! then
      return true
    isConditionalHyps eq hs

open private preprocess from Lean.Meta.Tactic.Simp.SimpTheorems in
def withSimpTheoremInfos (ty : Expr) (k : SimpTheoremInfo → MetaM α) : MetaM (Array α) := withReducible do
  (← preprocess (← mkSorry ty true) ty (inv := false) (isGlobal := true))
      |>.toArray.mapM fun (_, ty') => do
    forallTelescopeReducing ty' fun hyps eq => do
      let some (_, lhs, rhs) := eq.eq? | throwError "not an equality {eq}"
      k {
        hyps, lhs, rhs
        isConditional := ← isConditionalHyps eq hyps.toList
      }

/-- Checks whether two expressions are equal for the simplifier. That is,
they are reducibly-definitional equal, and they have the same head symbol. -/
def isSimpEq (a b : Expr) (whnfFirst := true) : MetaM Bool := withReducible do
  let a ← if whnfFirst then whnf a else pure a
  let b ← if whnfFirst then whnf b else pure b
  if a.getAppFn.constName? != b.getAppFn.constName? then return false
  isDefEq a b

def checkAllSimpTheoremInfos (ty : Expr) (k : SimpTheoremInfo → MetaM (Option MessageData)) : MetaM (Option MessageData) := do
  let errors := (← withSimpTheoremInfos ty fun i => do (← k i).mapM addMessageContextFull).filterMap id
  if errors.isEmpty then
    return none
  else
    return MessageData.joinSep errors.toList Format.line

def isSimpTheorem (declName : Name) : MetaM Bool := do
  pure $ (← getSimpTheorems).lemmaNames.contains declName

open Lean.Meta.DiscrTree
partial def trieElements : Trie α → StateT (Array α) Id Unit
  | Trie.node vs children => do
    modify (· ++ vs)
    for (_, child) in children do
      trieElements child

def elements (d : DiscrTree α) : StateT (Array α) Id Unit := do
  for (_, child) in d.root.toList do
    trieElements child

open Std

-- In Lean 4, the simplifier adds auxiliary lemmas if the declaration is not an equation.
-- For example, for `decl : a ↔ b` it generates `decl._auxLemma.1 : a = b`.
-- This function computes the map ``{`decl._auxLemma.1 ↦ `decl}``
def constToSimpDeclMap (ctx : Simp.Context) : HashMap Name Name := Id.run do
  let mut map : HashMap Name Name := {}
  for sls' in ctx.simpTheorems do
    for sls in [sls'.pre, sls'.post] do
      for sl in ((elements sls).run #[]).2 do
        if let some declName := sl.name? then
          if let some auxDeclName := sl.proof.getAppFn.constName? then
            map := map.insert auxDeclName declName
  return map

def isEqnLemma? (n : Name) : Option Name :=
  if n.isStr && n.getString!.startsWith "_eq_" then
    n.getPrefix
  else
    none

def heuristicallyExtractSimpTheoremsCore (ctx : Simp.Context) (constToSimpDecl : HashMap Name Name) (prf : Expr) : Array Name := Id.run do
  let mut cnsts : HashSet Name := {}
  for c in prf.getUsedConstants do
    if ctx.simpTheorems.isDeclToUnfold c then
      cnsts := cnsts.insert c
    else if ctx.congrTheorems.lemmas.contains c then
      cnsts := cnsts.insert c
    else if let some c' := constToSimpDecl.find? c then
      cnsts := cnsts.insert c'
    else if let some c' := isEqnLemma? c then
      cnsts := cnsts.insert c'
  return cnsts.toArray

@[inline] def heuristicallyExtractSimpTheorems (ctx : Simp.Context) (prf : Expr) : Array Name :=
  heuristicallyExtractSimpTheoremsCore ctx (constToSimpDeclMap ctx) prf

def decorateError (msg : MessageData) (k : MetaM α) : MetaM α := do
  try k catch e => throw e

def formatLemmas (lems : Array Name) : CoreM MessageData := do
  toMessageData <$> lems.mapM mkConstWithLevelParams

/-- A linter for simp lemmas whose lhs is not in simp-normal form, and which hence never fire. -/
@[mathlibLinter] def simpNF : Linter where
  noErrorsFound := "All left-hand sides of simp lemmas are in simp-normal form."
  errorsFound := "SOME SIMP LEMMAS ARE NOT IN SIMP-NORMAL FORM.
see note [simp-normal form] for tips how to debug this.
https://leanprover-community.github.io/mathlib_docs/notes.html#simp-normal%20form"

  test := fun declName => do
    unless ← isSimpTheorem declName do return none
    -- TODO: equation lemmas
    let ctx ← Simp.Context.mkDefault
    checkAllSimpTheoremInfos (← getConstInfo declName).type fun {lhs, rhs, isConditional, ..} => do
    let ⟨lhs', prf1⟩ ← decorateError "simplify fails on left-hand side:" <| simp lhs ctx
    let prf1_lems := heuristicallyExtractSimpTheorems ctx (prf1.getD (mkBVar 0))
    if prf1_lems.contains declName then return none
    let ⟨rhs', prf2⟩ ← decorateError "simplify fails on right-hand side:" <| simp rhs ctx
    let lhs'_eq_rhs' ← isSimpEq lhs' rhs' (whnfFirst := false)
    let lhs_in_nf ← isSimpEq lhs' lhs
    if lhs'_eq_rhs' then do
      if prf1.isNone then return none -- TODO: cannot detect used rfl-lemmas
      let used_lemmas := heuristicallyExtractSimpTheorems ctx <|
        mkApp (prf1.getD (mkBVar 0)) (prf2.getD (mkBVar 0))
      return m!"simp can prove this:
  by simp only {← formatLemmas used_lemmas}
One of the lemmas above could be a duplicate.
If that's not the case try reordering lemmas or adding @[priority].
"
    else if ¬ lhs_in_nf then do
      return m!"Left-hand side simplifies from
  {lhs}
to
  {lhs'}
using
  {← formatLemmas prf1_lems}
Try to change the left-hand side to the simplified term!
"
    else if !isConditional && lhs == lhs' then
      return m!"Left-hand side does not simplify.
You need to debug this yourself using `set_option trace.Meta.Tactic.simp.rewrite true`"
    else
      return none

library_note "simp-normal form" /--
This note gives you some tips to debug any errors that the simp-normal form linter raises.

The reason that a lemma was considered faulty is because its left-hand side is not in simp-normal
form.
These lemmas are hence never used by the simplifier.

This linter gives you a list of other simp lemmas: look at them!

Here are some tips depending on the error raised by the linter:

  1. 'the left-hand side reduces to XYZ':
     you should probably use XYZ as the left-hand side.

  2. 'simp can prove this':
     This typically means that lemma is a duplicate, or is shadowed by another lemma:

     2a. Always put more general lemmas after specific ones:
      ```
      @[simp] lemma zero_add_zero : 0 + 0 = 0 := rfl
      @[simp] lemma add_zero : x + 0 = x := rfl
      ```

      And not the other way around!  The simplifier always picks the last matching lemma.

     2b. You can also use `@[priority]` instead of moving simp-lemmas around in the file.

      Tip: the default priority is 1000.
      Use `@[priority 1100]` instead of moving a lemma down,
      and `@[priority 900]` instead of moving a lemma up.

     2c. Conditional simp lemmas are tried last. If they are shadowed
         just remove the `simp` attribute.

     2d. If two lemmas are duplicates, the linter will complain about the first one.
         Try to fix the second one instead!
         (You can find it among the other simp lemmas the linter prints out!)

  3. 'try_for tactic failed, timeout':
     This typically means that there is a loop of simp lemmas.
     Try to apply squeeze_simp to the right-hand side (removing this lemma from the simp set) to see
     what lemmas might be causing the loop.

     Another trick is to `set_option trace.simplify.rewrite true` and
     then apply `try_for 10000 { simp }` to the right-hand side.  You will
     see a periodic sequence of lemma applications in the trace message.
-/

/--
A linter for simp lemmas whose lhs has a variable as head symbol,
and which hence never fire.
-/
@[mathlibLinter] def simpVarHead : Linter where
  noErrorsFound :=
    "No left-hand sides of a simp lemma has a variable as head symbol."
  errorsFound := "LEFT-HAND SIDE HAS VARIABLE AS HEAD SYMBOL.
Some simp lemmas have a variable as head symbol of the left-hand side:"
  test := fun declName => do
    unless ← isSimpTheorem declName do return none
    checkAllSimpTheoremInfos (← getConstInfo declName).type fun {lhs, ..} => do
    let headSym := lhs.getAppFn
    unless headSym.isFVar do return none
    return m!"Left-hand side has variable as head symbol: {headSym}"

/-- A linter for commutativity lemmas that are marked simp. -/
@[mathlibLinter] def simpComm : Linter where
  noErrorsFound := "No commutativity lemma is marked simp."
  errorsFound := "COMMUTATIVITY LEMMA IS SIMP.
Some commutativity lemmas are simp lemmas:"
  test := fun declName => withReducible do
    unless ← isSimpTheorem declName do return none
    let ty := (← getConstInfo declName).type
    forallTelescopeReducing ty fun xs ty => do
    let some (_, lhs, rhs) := ty.eq? | return none
    unless lhs.getAppFn.constName? == rhs.getAppFn.constName? do return none
    let (mvars, bis, ty') ← forallMetaTelescopeReducing ty
    let some (_, lhs', rhs') := ty'.eq? | return none
    unless ← isDefEq rhs lhs' do return none
    unless ← withNewMCtxDepth (isDefEq rhs lhs') do return none
    -- ensure that the second application makes progress:
    if ← isDefEq lhs' rhs' then return none
    pure m!"should not be marked simp"


/-
Copyright 2021 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
 -/
import measure_theory.measurable_space

import measure_theory.measure_space
import measure_theory.outer_measure
import measure_theory.lebesgue_measure
import measure_theory.integration

import measure_theory.borel_space
import data.set.countable
import formal_ml.nnreal
import formal_ml.sum
import formal_ml.lattice
import formal_ml.measurable_space
import formal_ml.classical
import data.equiv.list
import formal_ml.prod_measure
import formal_ml.finite_pi_measure
import formal_ml.probability_space
import formal_ml.random_variable_identical
import formal_ml.prod_probability_space

/-
  This file proves that two IID distributions are identical.

  This result is subtle (i.e., why isn't it just by definition).
  IID variables can be constructed as IID, at which point they are
  obviously IID. However, they can be observed to be IID.

  Notice that, given X_1 and X_2, and Y_1 and Y_2, all coin flips.

  if X_1 is identical to Y_1, and X_2 is identical to Y_2, it is
  not enough to show that every event has the same probability. For
  example, maybe `Pr[ X_1 =ᵣ X_2] ≠ Pr[Y_1 =ᵣ Y_2]`. This kind of
  reasoning gives us a deeper understanding of the proof of VCD, which
  are based upon constructing an IID distribution through shuffling
  and subsampling.

-/

def pi_base {δ:Type*} {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
  (S:set (Π (a:δ), π a)):Prop := ∃ (f:Π (a:δ), set (π a)), (∀ i, measurable_set (f i)) ∧
  set.pi set.univ f = S

lemma pi_base.closed_inter
  {δ:Type*} {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
  (S T:set (Π (a:δ), π a)):pi_base S → pi_base T → pi_base (S ∩ T) :=
begin
  intros hS hT,
  cases hS with fS hS,
  cases hT with fT hT,
  apply exists.intro (λ i, (fS i) ∩ (fT i)),
  split,
  { intros i, apply measurable_set.inter, apply hS.left, apply hT.left },
  rw ← hS.right,
  rw ← hT.right,
  ext ω, split; intros h1; simp at h1; simp [h1],
end

lemma pi_base.empty
  {δ:Type*} [nonempty δ] {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]:
  pi_base (∅:set (Π (a : δ), π a)) := begin
  apply exists.intro (λ (d:δ), (∅:set (π d))),
  simp [set.pi],
end

lemma pi_base.univ
  {δ:Type*} [nonempty δ] {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]:
  pi_base (set.univ:set (Π (a : δ), π a)) := begin
  apply exists.intro (λ (d:δ), (set.univ:set (π d))),
  simp [set.pi],
end

lemma pi_base.closed_compl
  {δ:Type*} [fintype δ] [nonempty δ] {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
    (s : set (Π (a : δ), π a)):
    pi_base s →
    (set.disjoint_union_closure pi_base) sᶜ  := begin
  classical,
  intros h1,
  simp [pi_base] at h1,
  cases h1 with f h1,
  let f':(δ → bool) → set (Π (a : δ), π a) :=
    (λ x, if (∀ (d:δ), x d) then ∅  else
   (set.pi (@set.univ δ) (λ (d:δ), if (x d) then (f d) else (f d)ᶜ))),
  begin
    have h2:sᶜ = set.Union f',
    { rw ← h1.right, simp [f'], ext; split; intros h2_1;
      simp at h2_1; simp [h2_1],
      { let i:δ → bool := (λ (d:δ), x d ∈ (f d)),
        apply exists.intro i,
        rw if_neg,
        simp [set.pi, i],
        intros d,
        cases classical.em (x d ∈ f d) with h2_2 h2_2,
        rw if_pos h2_2, apply h2_2,
        rw if_neg h2_2, apply h2_2,
        intros contra,
        cases h2_1 with d h2_1,
        apply h2_1,
        have contra_2 := contra d,
        simp [i] at contra_2,
        apply contra_2 },
      { cases h2_1 with i h2_1,
        cases classical.em (∀ (d : δ), (i d)) with h2_3 h2_3,
        { exfalso, rw if_pos h2_3 at h2_1, apply h2_1 },
        rw if_neg h2_3 at h2_1,
        rw classical.not_forall_iff_exists_not at h2_3,
        cases h2_3 with d h2_3,
        apply exists.intro d,
        simp [set.pi] at h2_1,
        have h2_4 := h2_1 d,
        rw if_neg h2_3 at h2_4,
        apply h2_4 } },
    rw h2,
    apply set.disjoint_union_closure_intro,
    intros b,
    { simp [f',pi_base],
      cases classical.em (∀ (d : δ), (b d)) with h3 h3,
      rw if_pos,
      apply pi_base.empty,
      apply h3,
      rw if_neg,
      { apply exists.intro (λ (d : δ), ite ↥(b d) (f d) (f d)ᶜ),
        rw and.comm,
        split,
        refl,
        intros d,
        simp,
        cases classical.em (b d) with h4 h4,
        rw if_pos,
        apply h1.left,
        apply h4,
        rw if_neg,
        apply measurable_set.compl,
        apply h1.left, apply h4 },
      apply h3 },
    intros b b' h_ne,
    simp [function.on_fun], rw disjoint_iff,
    simp,
    rw ← set.subset_empty_iff,
    rw set.subset_def,
    intros ω h_contra,
    exfalso,
    simp [f'] at h_contra,
    cases h_contra with h_contra_1 h_contra_2,
    cases classical.em (∀ (d : δ), (b d)) with h5 h5,
    { rw if_pos at h_contra_1,
      apply h_contra_1,
      apply h5 },
    rw if_neg h5 at h_contra_1,
    cases classical.em (∀ (d : δ), (b' d)) with h6 h6,
    { rw if_pos h6 at h_contra_2,
      apply h_contra_2 },
    rw if_neg h6 at h_contra_2,
    apply h_ne,
    simp [set.pi] at h_contra_1,
    simp [set.pi] at h_contra_2,
    ext i,
    have h_contra_1_i := h_contra_1 i,
    have h_contra_2_i := h_contra_2 i,
    destruct (b i); destruct (b' i);
    intros h_b'_i h_b_i;
    simp [h_b_i, h_b'_i];
    simp [h_b_i] at h_contra_1_i;
    simp [h_b'_i] at h_contra_2_i,
    apply h_contra_1_i h_contra_2_i,
    apply h_contra_2_i h_contra_1_i,
  end
end

lemma set.pi_as_preimage {δ:Type*} {π:δ → Type*} {f:Π (a:δ), set (π a)}:
  (set.pi set.univ f) = (⋂ (i:δ), (λ (d:Π (a:δ), π a), d i) ⁻¹' (f i)) := begin
  ext,split;intros A1; simp at A1; simp [A1],
end

lemma pi_base.measurable_set {δ:Type*} [encodable δ] {π:δ → Type*}
  [M:Π (a:δ), measurable_space (π a)]
  (S:set (Π (a:δ), π a)):pi_base S → measurable_set S :=
begin
  intros hS,
  cases hS with fS hS,
  cases hS with h_meas h_def,
  rw ← h_def,
  rw set.pi_as_preimage,
  apply measurable_set.Inter,
  intros b,
  have h_proj_meas:measurable (λ (d : Π (a : δ), π a), d b),
  { apply measurable_pi_apply },
  apply h_proj_meas,
  apply h_meas,
end

lemma pi_base.covers_generate {δ:Type*} [encodable δ] {π:δ → Type*}
  [M:Π (a:δ), measurable_space (π a)]
  (S:set (Π (a:δ), π a)):
  S ∈
     (⋃ (i : δ),
       (measurable_space.comap (λ (b : Π (a : δ), (λ (a : δ), π a) a), b i)
          ((λ (a : δ), M a) i)).measurable_set') →
  pi_base S := begin
  intros h1,
  simp [pi_base],
  simp at h1,
  cases h1 with d h1,
  simp [measurable_space.comap] at h1,
  cases h1 with s h1,
  cases h1 with h1 h3,
  have h2:∀ (d':δ), ∃ (T:set (π d')), measurable_set T ∧
     (∀ (s':(Π (a:δ), π a)), s' ∈ S → s' d' ∈ T) ∧ (d' = d → T == s),
  { intros d',
    cases classical.em (d = d') with h2_1 h2_1,
    { subst d', apply exists.intro s,
      split,
      apply h1,
      split,
      intros s' h4, rw ← h3 at h4,
      simp at h4, apply h4, intros h5, refl },
    { apply exists.intro set.univ,
      split,
      apply measurable_set.univ,
      split,
      intros s' h6,
      simp,
      intros h2_1_not, exfalso, apply h2_1, rw h2_1_not, },
   },
   have h7 := classical.axiom_of_choice h2,
   cases h7 with f h7,
   apply exists.intro f,
   split,
   { intros d',
     apply (h7 d').left },
   { ext ω,
     simp [set.pi],split; intros h8,
     { rw ← h3,
       have h9:d = d := rfl,
       have h10 := (h7 d).right.right h9,
       simp at h10, rw ← h10,
       simp, apply h8 },
     { intros d',
       have h11 := (h7 d').right.left ω h8, apply h11 } },
end

lemma pi_base_eq_measurable_space_pi {δ:Type*} [encodable δ] {π:δ → Type*}
  [M:Π (a:δ), measurable_space (π a)]:
 @measurable_space.pi δ π M  = measurable_space.generate_from pi_base :=
begin
  apply le_antisymm;
  apply measurable_space.generate_from_le;
  intros a h_a,
  { simp [measurable_space.generate_from],
    apply measurable_space.generate_measurable.basic,
    apply pi_base.covers_generate,
    simp at h_a, simp, cases h_a with i h_a, apply exists.intro i, apply h_a },
  { apply pi_base.measurable_set, apply h_a },
end

lemma pi_base_eq {δ:Type*} [fintype δ] {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
  (S:measurable_setB (@measurable_space.pi δ π M)):pi_base S.val →
  ∃ (f:Π (b:δ), measurable_setB (M b)), (set.pi_measurable set.univ f) = S := begin
  intros h1,
  unfold pi_base at h1,
  cases h1 with g h1,
  cases h1 with h1 h2,
  let f := (λ (b:δ), @measurable_setB.mk _ _ (g b) (h1 b)),
  begin
     apply exists.intro f,
     apply subtype.eq,
     rw ← h2,
     simp [set.pi, set.pi_measurable, f],
     ext ω,split;intros h1; simp at h1; simp [h1]; intros i; apply h1,
  end
end

lemma pi_base_is_semialgebra {δ:Type*} [fintype δ] [nonempty δ] {π:δ → Type*}
  [M:Π (a:δ), measurable_space (π a)]:set.is_semialgebra (@pi_base δ π M) := {
  univ := pi_base.univ,
  empty := pi_base.empty,
  compl := pi_base.closed_compl,
  inter := pi_base.closed_inter,
}

lemma pi_union_all {Ω δ:Type*} [fintype δ]
  {P:probability_space Ω}
  {π:δ → Type*} [M:Π (a:δ), measurable_space (π a)]
  (f:Π (b:δ), measurable_setB (M b))
  {X:Π (b:δ), P →ᵣ (M b)}:
  (pi.random_variable_combine X ∈ᵣ (set.pi_measurable set.univ f)) =
  (∀ᵣ (b:δ), (X b) ∈ᵣ (f b)) := begin
  apply event.eq,
  simp [pi.random_variable_combine, set.pi_measurable, pi.measurable_fun],
  ext ω, split; intros h1; simp at h1; simp [h1],
end

lemma random_variable_independent_all {Ω β:Type*} {P:probability_space Ω} [fintype β]
  {γ:β → Type*} {M:Π (b:β), measurable_space (γ b)}
   (S:Π (b:β), measurable_setB (M b))
   {X:Π (b:β), P →ᵣ (M b)}:
   random_variable_independent X →
   Pr[∀ᵣ (b : β), X b ∈ᵣ S b] =  finset.univ.prod (λ (b:β), Pr[X b ∈ᵣ S b]) := begin
  intros h1,
  simp [random_variable_independent] at h1,
  have h2 := h1 S,
  simp [independent_events] at h2,
  have h3 := h2 finset.univ,
  rw h3,
  refl,
end

lemma IID_identical_joint_base' {Ω₁ Ω₂ β γ:Type*} [fintype β]
  {P₁:probability_space Ω₁}
  {P₂:probability_space Ω₂}
  {M:measurable_space γ}
  {X₁:β → P₁ →ᵣ M}  {X₂:β  → P₂ →ᵣ M}
  (S:measurable_setB (@measurable_space.pi β (λ _, γ) (λ _, M))):
  pi_base S.val →
  (∀ (b:β), random_variable_identical (X₁ b) (X₂ b)) →
  (random_variables_IID X₁) →
  (random_variables_IID X₂) →
  (Pr[pi.random_variable_combine X₁ ∈ᵣ S] =
   Pr[pi.random_variable_combine X₂ ∈ᵣ S]) :=
begin
  intros h1 h2 h3 h4,
  have h5:=pi_base_eq S h1,
  cases h5 with f h5,
  rw ← h5,
  rw pi_union_all,
  rw pi_union_all,
  rw random_variable_independent_all,
  rw random_variable_independent_all,
  have h6:(λ (b : β), Pr[X₁ b ∈ᵣ f b]) = (λ (b : β), Pr[X₂ b ∈ᵣ f b]),
  { ext1 b, apply h2, },
  rw h6,
  apply h4.left,
  apply h3.left,
end

lemma IID_identical_joint' {Ω₁ Ω₂ β γ:Type*} [fintype β] [nonempty β]
  {P₁:probability_space Ω₁}
  {P₂:probability_space Ω₂}
  {M:measurable_space γ}
  {X₁:β → P₁ →ᵣ M}  {X₂:β  → P₂ →ᵣ M}:
  (∀ (b:β), random_variable_identical (X₁ b) (X₂ b)) →
  (random_variables_IID X₁) →
  (random_variables_IID X₂) →
  random_variable_identical (pi.random_variable_combine X₁)
                            (pi.random_variable_combine X₂)
  :=
begin
  intros h1 h2 h3,
  haveI:encodable β := fintype.encodable β,
  apply random_variable_identical_on_semialgebra' pi_base,
  --apply pi_base
  apply pi_base.closed_inter,
  apply pi_base.closed_compl,
  apply pi_base.empty,
  apply pi_base.univ,
  {
    apply pi_base_eq_measurable_space_pi },
  intros T h_T,
  apply IID_identical_joint_base',
  apply h_T,
  apply h1,
  apply h2,
  apply h3,
end

lemma IID_identical_joint {Ω₁ Ω₂ β γ:Type*} [fintype β] [nonempty β]
  {P₁:probability_space Ω₁}
  {P₂:probability_space Ω₂}
  {M:measurable_space γ}
  {X₁:β → P₁ →ᵣ M}  {X₂:β  → P₂ →ᵣ M}:
  (∃ (b₁ b₂:β), random_variable_identical (X₁ b₁) (X₂ b₂)) →
  (random_variables_IID X₁) →
  (random_variables_IID X₂) →
  random_variable_identical (pi.random_variable_combine X₁)
                            (pi.random_variable_combine X₂) :=
begin
  intros h1 h2 h3,
  apply IID_identical_joint',
  intros b,
  cases h1 with b₁ h1,
  cases h1 with b₂ h1,
  have h4:random_variable_identical (X₁ b) (X₁ b₁),
  { apply h2.right },
  apply random_variable_identical.trans h4,
  apply random_variable_identical.trans h1,
  apply h3.right,
  apply h2,
  apply h3,
end


example (P Q R : Type) : (P → (Q → R)) → ((P → Q) → (P → R)) := by
  intro p_qr
  intro p_q
  intro p
  apply p_qr
  case a =>
    exact p
  case a =>
    apply p_q
    exact p

example (P Q R : Type) : (P → (Q → R)) → ((P → Q) → (P → R)) := by
  intro p_qr
  intro p_q
  intro p
  apply p_qr
  . exact p
  apply p_q
  exact p


/-
Copyright (c) 2019 Paul-Nicolas Madelaine. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Paul-Nicolas Madelaine, Robert Y. Lewis
-/
import tactic.converter.interactive
import tactic.hint

/-!
# A tactic for normalizing casts inside expressions

This tactic normalizes casts inside expressions.
It can be thought of as a call to the simplifier with a specific set of lemmas to
move casts upwards in the expression.
It has special handling of numerals and a simple heuristic to help moving
casts "past" binary operators.
Contrary to simp, it should be safe to use as a non-terminating tactic.

The algorithm implemented here is described in the paper
<https://lean-forward.github.io/norm_cast/norm_cast.pdf>.

## Important definitions
* `tactic.interactive.norm_cast`
* `tactic.interactive.push_cast`
* `tactic.interactive.exact_mod_cast`
* `tactic.interactive.apply_mod_cast`
* `tactic.interactive.rw_mod_cast`
* `tactic.interactive.assumption_mod_cast`
-/

setup_tactic_parser

namespace tactic

/--
Runs `mk_instance` with a time limit.

This is a work around to the fact that in some cases
mk_instance times out instead of failing,
for example: `has_lift_t ℤ ℕ`

`mk_instance_fast` is used when we assume the type class search
should end instantly.
-/
meta def mk_instance_fast (e : expr) (timeout := 1000) : tactic expr :=
try_for timeout (mk_instance e)

end tactic

namespace norm_cast

open tactic expr

declare_trace norm_cast

/--
Output a trace message if `trace.norm_cast` is enabled.
-/
meta def trace_norm_cast {α} [has_to_tactic_format α] (msg : string) (a : α) : tactic unit :=
when_tracing `norm_cast $ do
a ← pp a,
trace ("[norm_cast] " ++ msg ++ a : format)

mk_simp_attribute push_cast "The `push_cast` simp attribute uses `norm_cast` lemmas
to move casts toward the leaf nodes of the expression."

/--
`label` is a type used to classify `norm_cast` lemmas.
* elim lemma:   LHS has 0 head coes and ≥ 1 internal coe
* move lemma:   LHS has 1 head coe and 0 internal coes,    RHS has 0 head coes and ≥ 1 internal coes
* squash lemma: LHS has ≥ 1 head coes and 0 internal coes, RHS has fewer head coes
-/
@[derive [decidable_eq, has_reflect, inhabited]]
inductive label
| elim   : label
| move   : label
| squash : label

namespace label

/-- Convert `label` into `string`. -/
protected def to_string : label → string
| elim   := "elim"
| move   := "move"
| squash := "squash"

instance : has_to_string label := ⟨label.to_string⟩
instance : has_repr label := ⟨label.to_string⟩
meta instance : has_to_format label := ⟨λ l, l.to_string⟩

/-- Convert `string` into `label`. -/
def of_string : string -> option label
| "elim" := some elim
| "move" := some move
| "squash" := some squash
| _ := none

end label

open label

/-- Count how many coercions are at the top of the expression. -/
meta def count_head_coes : expr → ℕ
| `(coe %%e) := count_head_coes e + 1
| `(coe_sort %%e) := count_head_coes e + 1
| `(coe_fn %%e) := count_head_coes e + 1
| _ := 0

/-- Count how many coercions are inside the expression, including the top ones. -/
meta def count_coes : expr → tactic ℕ
| `(coe %%e) := (+1) <$> count_coes e
| `(coe_sort %%e) := (+1) <$> count_coes e
| `(coe_fn %%e) := (+1) <$> count_coes e
| (app `(coe_fn %%e) x) := (+) <$> count_coes x <*> (+1) <$> count_coes e
| (expr.lam n bi t e) := do
  l ← mk_local' n bi t,
  count_coes $ e.instantiate_var l
| e := do
  as ← e.get_simp_args,
  list.sum <$> as.mmap count_coes

/-- Count how many coercions are inside the expression, excluding the top ones. -/
private meta def count_internal_coes (e : expr) : tactic ℕ := do
ncoes ← count_coes e,
pure $ ncoes - count_head_coes e

/--
Classifies a declaration of type `ty` as a `norm_cast` rule.
-/
meta def classify_type (ty : expr) : tactic label := do
(_, ty) ← open_pis ty,
(lhs, rhs) ← match ty with
  | `(%%lhs = %%rhs) := pure (lhs, rhs)
  | `(%%lhs ↔ %%rhs) := pure (lhs, rhs)
  | _ := fail "norm_cast: lemma must be = or ↔"
  end,
lhs_coes ← count_coes lhs,
when (lhs_coes = 0) $ fail "norm_cast: badly shaped lemma, lhs must contain at least one coe",
let lhs_head_coes := count_head_coes lhs,
lhs_internal_coes ← count_internal_coes lhs,
let rhs_head_coes := count_head_coes rhs,
rhs_internal_coes ← count_internal_coes rhs,
if lhs_head_coes = 0 then
  return elim
else if lhs_head_coes = 1 then do
  when (rhs_head_coes ≠ 0) $ fail "norm_cast: badly shaped lemma, rhs can't start with coe",
  if rhs_internal_coes = 0 then
    return squash
  else
    return move
else if rhs_head_coes < lhs_head_coes then do
  return squash
else do
  fail "norm_cast: badly shaped shaped squash lemma, rhs must have fewer head coes than lhs"

/-- The cache for `norm_cast` attribute stores three `simp_lemma` objects. -/
meta structure norm_cast_cache :=
(up : simp_lemmas)
(down : simp_lemmas)
(squash : simp_lemmas)

/-- Empty `norm_cast_cache`. -/
meta def empty_cache : norm_cast_cache :=
{ up     := simp_lemmas.mk,
  down   := simp_lemmas.mk,
  squash := simp_lemmas.mk, }

meta instance : inhabited norm_cast_cache := ⟨empty_cache⟩

/-- `add_elim cache e` adds `e` as an `elim` lemma to `cache`. -/
meta def add_elim (cache : norm_cast_cache) (e : expr) : tactic norm_cast_cache :=
do
  new_up ← cache.up.add e,
  return
  { up     := new_up,
    down   := cache.down,
    squash := cache.squash, }

/-- `add_move cache e` adds `e` as a `move` lemma to `cache`. -/
meta def add_move (cache : norm_cast_cache) (e : expr) : tactic norm_cast_cache :=
do
  new_up ← cache.up.add e tt,
  new_down ← cache.down.add e,
  return
  { up     := new_up,
    down   := new_down,
    squash := cache.squash, }

/-- `add_squash cache e` adds `e` as an `squash` lemma to `cache`. -/
meta def add_squash (cache : norm_cast_cache) (e : expr) : tactic norm_cast_cache :=
do
  new_squash ← cache.squash.add e,
  new_down ← cache.down.add e,
  return
  { up     := cache.up,
    down   := new_down,
    squash := new_squash, }

/--
The type of the `norm_cast` attribute.
The optional label is used to overwrite the classifier.
-/
meta def norm_cast_attr_ty : Type := user_attribute norm_cast_cache (option label)

/--
Efficient getter for the `@[norm_cast]` attribute parameter that does not call `eval_expr`.

See Note [user attribute parameters].
-/
meta def get_label_param (attr : norm_cast_attr_ty) (decl : name) : tactic (option label) := do
p ← attr.get_param_untyped decl,
match p with
| `(none) := pure none
| `(some label.elim) := pure label.elim
| `(some label.move) := pure label.move
| `(some label.squash) := pure label.squash
| _ := fail p
end

/--
`add_lemma cache decl` infers the proper `norm_cast` attribute for `decl` and adds it to `cache`.
-/
meta def add_lemma (attr : norm_cast_attr_ty) (cache : norm_cast_cache) (decl : name) :
  tactic norm_cast_cache :=
do
  e ← mk_const decl,
  param ← get_label_param attr decl,
  l ← param <|> (infer_type e >>= classify_type),
  match l with
  | elim   := add_elim cache e
  | move   := add_move cache e
  | squash := add_squash cache e
  end

-- special lemmas to handle the ≥, > and ≠ operators
private lemma ge_from_le {α} [has_le α] : ∀ (x y : α), x ≥ y ↔ y ≤ x := λ _ _, iff.rfl
private lemma gt_from_lt {α} [has_lt α] : ∀ (x y : α), x > y ↔ y < x := λ _ _, iff.rfl
private lemma ne_from_not_eq {α} : ∀ (x y : α), x ≠ y ↔ ¬(x = y) := λ _ _, iff.rfl

/--
`mk_cache names` creates a `norm_cast_cache`. It infers the proper `norm_cast` attributes
for names in `names`, and collects the lemmas attributed with specific `norm_cast` attributes.
-/
meta def mk_cache (attr : thunk norm_cast_attr_ty) (names : list name) :
  tactic norm_cast_cache := do
-- names has the declarations in reverse order
cache ← names.mfoldr (λ name cache, add_lemma (attr ()) cache name) empty_cache,

--some special lemmas to handle binary relations
let up := cache.up,
up ← up.add_simp ``ge_from_le,
up ← up.add_simp ``gt_from_lt,
up ← up.add_simp ``ne_from_not_eq,

let down := cache.down,
down ← down.add_simp ``coe_coe,

pure { up := up, down := down, squash := cache.squash }

/--
The `norm_cast` attribute.
-/
@[user_attribute] meta def norm_cast_attr : user_attribute norm_cast_cache (option label) :=
{ name      := `norm_cast,
  descr     := "attribute for norm_cast",
  parser    :=
    (do some l ← (label.of_string ∘ to_string) <$> ident, return l)
      <|> return none,
  after_set := some (λ decl prio persistent, do
    param ← get_label_param norm_cast_attr decl,
    match param with
    | some l :=
      when (l ≠ elim) $ simp_attr.push_cast.set decl () tt prio
    | none := do
      e ← mk_const decl,
      ty ← infer_type e,
      l ← classify_type ty,
      norm_cast_attr.set decl l persistent prio
    end),
  before_unset := some $ λ _ _, tactic.skip,
  cache_cfg := { mk_cache := mk_cache norm_cast_attr, dependencies := [] } }

/-- Classify a declaration as a `norm_cast` rule. -/
meta def make_guess (decl : name) : tactic label :=
do
  e ← mk_const decl,
  ty ← infer_type e,
  classify_type ty

/--
Gets the `norm_cast` classification label for a declaration. Applies the
override specified on the attribute, if necessary.
-/
meta def get_label (decl : name) : tactic label :=
do
  param ← get_label_param norm_cast_attr decl,
  param <|> make_guess decl

end norm_cast

namespace tactic.interactive
open norm_cast

/--
`push_cast` rewrites the expression to move casts toward the leaf nodes.
For example, `↑(a + b)` will be written to `↑a + ↑b`.
Equivalent to `simp only with push_cast`.
Can also be used at hypotheses.

`push_cast` can also be used at hypotheses and with extra simp rules.

```lean
example (a b : ℕ) (h1 : ((a + b : ℕ) : ℤ) = 10) (h2 : ((a + b + 0 : ℕ) : ℤ) = 10) :
  ((a + b : ℕ) : ℤ) = 10 :=
begin
  push_cast,
  push_cast at h1,
  push_cast [int.add_zero] at h2,
end
```
-/
meta def push_cast (hs : parse tactic.simp_arg_list) (l : parse location) : tactic unit :=
tactic.interactive.simp none none tt hs [`push_cast] l {discharger := tactic.assumption}

end tactic.interactive

namespace norm_cast
open tactic expr

/-- Prove `a = b` using the given simp set. -/
meta def prove_eq_using (s : simp_lemmas) (a b : expr) : tactic expr := do
(a', a_a', _) ← simplify s [] a {fail_if_unchanged := ff},
(b', b_b', _) ← simplify s [] b {fail_if_unchanged := ff},
on_exception (trace_norm_cast "failed: " (to_expr ``(%%a' = %%b') >>= pp)) $
  is_def_eq a' b' reducible,
b'_b ← mk_eq_symm b_b',
mk_eq_trans a_a' b'_b

/-- Prove `a = b` by simplifying using move and squash lemmas. -/
meta def prove_eq_using_down (a b : expr) : tactic expr := do
cache ← norm_cast_attr.get_cache,
trace_norm_cast "proving: " (to_expr ``(%%a = %%b) >>= pp),
prove_eq_using cache.down a b

/--
This is the main heuristic used alongside the elim and move lemmas.
The goal is to help casts move past operators by adding intermediate casts.
An expression of the shape: op (↑(x : α) : γ) (↑(y : β) : γ)
is rewritten to:            op (↑(↑(x : α) : β) : γ) (↑(y : β) : γ)
when (↑(↑(x : α) : β) : γ) = (↑(x : α) : γ) can be proven with a squash lemma
-/
meta def splitting_procedure : expr → tactic (expr × expr)
| (app (app op x) y) :=
(do
  `(@coe %%α %%δ %%coe1 %%xx) ← return x,
  `(@coe %%β %%γ %%coe2 %%yy) ← return y,
  success_if_fail $ is_def_eq α β,
  is_def_eq δ γ,

  (do
    coe3 ← mk_app `has_lift_t [α, β] >>= mk_instance_fast,
    new_x ← to_expr ``(@coe %%β %%δ %%coe2 (@coe %%α %%β %%coe3 %%xx)),
    let new_e := app (app op new_x) y,
    eq_x ← prove_eq_using_down x new_x,
    pr ← mk_congr_arg op eq_x,
    pr ← mk_congr_fun pr y,
    return (new_e, pr)
  ) <|> (do
    coe3 ← mk_app `has_lift_t [β, α] >>= mk_instance_fast,
    new_y ← to_expr ``(@coe %%α %%δ %%coe1 (@coe %%β %%α %%coe3 %%yy)),
    let new_e := app (app op x) new_y,
    eq_y ← prove_eq_using_down y new_y,
    pr ← mk_congr_arg (app op x) eq_y,
    return (new_e, pr)
  )
) <|> (do
  `(@coe %%α %%β %%coe1 %%xx) ← return x,
  `(@has_one.one %%β %%h1) ← return y,
  h2 ← to_expr ``(has_one %%α) >>= mk_instance_fast,
  new_y ← to_expr ``(@coe %%α %%β %%coe1 (@has_one.one %%α %%h2)),
  eq_y ← prove_eq_using_down y new_y,
  let new_e := app (app op x) new_y,
  pr ← mk_congr_arg (app op x) eq_y,
  return (new_e, pr)
 ) <|> (do
  `(@coe %%α %%β %%coe1 %%xx) ← return x,
  `(@has_zero.zero %%β %%h1) ← return y,
  h2 ← to_expr ``(has_zero %%α) >>= mk_instance_fast,
  new_y ← to_expr ``(@coe %%α %%β %%coe1 (@has_zero.zero %%α %%h2)),
  eq_y ← prove_eq_using_down y new_y,
  let new_e := app (app op x) new_y,
  pr ← mk_congr_arg (app op x) eq_y,
  return (new_e, pr)
) <|> (do
  `(@has_one.one %%β %%h1) ← return x,
  `(@coe %%α %%β %%coe1 %%xx) ← return y,
  h1 ← to_expr ``(has_one %%α) >>= mk_instance_fast,
  new_x ← to_expr ``(@coe %%α %%β %%coe1 (@has_one.one %%α %%h1)),
  eq_x ← prove_eq_using_down x new_x,
  let new_e := app (app op new_x) y,
  pr ← mk_congr_arg (lam `x binder_info.default β (app (app op (var 0)) y)) eq_x,
  return (new_e, pr)
) <|> (do
  `(@has_zero.zero %%β %%h1) ← return x,
  `(@coe %%α %%β %%coe1 %%xx) ← return y,
  h1 ← to_expr ``(has_zero %%α) >>= mk_instance_fast,
  new_x ← to_expr ``(@coe %%α %%β %%coe1 (@has_zero.zero %%α %%h1)),
  eq_x ← prove_eq_using_down x new_x,
  let new_e := app (app op new_x) y,
  pr ← mk_congr_arg (lam `x binder_info.default β (app (app op (var 0)) y)) eq_x,
  return (new_e, pr)
)
| _ := failed

/--
Discharging function used during simplification in the "squash" step.

TODO: norm_cast takes a list of expressions to use as lemmas for the discharger
TODO: a tactic to print the results the discharger fails to proove
-/
private meta def prove : tactic unit :=
assumption

/--
Core rewriting function used in the "squash" step, which moves casts upwards
and eliminates them.

It tries to rewrite an expression using the elim and move lemmas.
On failure, it calls the splitting procedure heuristic.
-/
meta def upward_and_elim (s : simp_lemmas) (e : expr) : tactic (expr × expr) :=
(do
  r ← mcond (is_prop e) (return `iff) (return `eq),
  (new_e, pr) ← s.rewrite e prove r,
  pr ← match r with
  | `iff := mk_app `propext [pr]
  | _    := return pr
  end,
  return (new_e, pr)
) <|> splitting_procedure e

/-!
The following auxiliary functions are used to handle numerals.
-/

/--
If possible, rewrite `(n : α)` to `((n : ℕ) : α)` where `n` is a numeral and `α ≠ ℕ`.
Returns a pair of the new expression and proof that they are equal.
-/
meta def numeral_to_coe (e : expr) : tactic (expr × expr) :=
do
  α ← infer_type e,
  success_if_fail $ is_def_eq α `(ℕ),
  n ← e.to_nat,
  h1 ← mk_app `has_lift_t [`(ℕ), α] >>= mk_instance_fast,
  let new_e : expr := reflect n,
  new_e ← to_expr ``(@coe ℕ %%α %%h1 %%new_e),
  pr ← prove_eq_using_down e new_e,
  return (new_e, pr)

/--
If possible, rewrite `((n : ℕ) : α)` to `(n : α)` where `n` is a numeral.
Returns a pair of the new expression and proof that they are equal.
-/
meta def coe_to_numeral (e : expr) : tactic (expr × expr) :=
do
  `(@coe ℕ %%α %%h1 %%e') ← return e,
  n ← e'.to_nat,
  -- replace e' by normalized numeral
  is_def_eq (reflect n) e' reducible,
  let e := e.app_fn (reflect n),
  new_e ← expr.of_nat α n,
  pr ← prove_eq_using_down e new_e,
  return (new_e, pr)

/-- A local variant on `simplify_top_down`. -/
private meta def simplify_top_down' {α} (a : α) (pre : α → expr → tactic (α × expr × expr))
  (e : expr) (cfg : simp_config := {}) : tactic (α × expr × expr) :=
ext_simplify_core a cfg simp_lemmas.mk (λ _, failed)
  (λ a _ _ _ e, do
    (new_a, new_e, pr) ← pre a e,
    guard (¬ new_e =ₐ e),
    return (new_a, new_e, some pr, ff))
  (λ _ _ _ _ _, failed)
  `eq e

/--
The core simplification routine of `norm_cast`.
-/
meta def derive (e : expr) : tactic (expr × expr) :=
do
  cache ← norm_cast_attr.get_cache,
  e ← instantiate_mvars e,
  let cfg : simp_config :=
  { zeta := ff,
    beta := ff,
    eta  := ff,
    proj := ff,
    iota := ff,
    iota_eqn := ff,
    fail_if_unchanged := ff },
  let e0 := e,

  -- step 1: pre-processing of numerals
  ((), e1, pr1) ← simplify_top_down' () (λ _ e, prod.mk () <$> numeral_to_coe e) e0 cfg,
  trace_norm_cast "after numeral_to_coe: " e1,

  -- step 2: casts are moved upwards and eliminated
  ((), e2, pr2) ← simplify_bottom_up () (λ _ e, prod.mk () <$> upward_and_elim cache.up e) e1 cfg,
  trace_norm_cast "after upward_and_elim: " e2,

  -- step 3: casts are squashed
  (e3, pr3, _) ← simplify cache.squash [] e2 cfg,
  trace_norm_cast "after squashing: " e3,

  -- step 4: post-processing of numerals
  ((), e4, pr4) ← simplify_top_down' () (λ _ e, prod.mk () <$> coe_to_numeral e) e3 cfg,
  trace_norm_cast "after coe_to_numeral: " e4,

  let new_e := e4,
  guard (¬ new_e =ₐ e),
  pr ← mk_eq_trans pr1 pr2,
  pr ← mk_eq_trans pr pr3,
  pr ← mk_eq_trans pr pr4,
  return (new_e, pr)

/--
A small variant of `push_cast` suited for non-interactive use.

`derive_push_cast extra_lems e` returns an expression `e'` and a proof that `e = e'`.
-/
meta def derive_push_cast (extra_lems : list simp_arg_type) (e : expr) : tactic (expr × expr) :=
do (s, _) ← mk_simp_set tt [`push_cast] extra_lems,
   (e, prf, _) ← simplify (s.erase [`nat.cast_succ]) [] e
                  {fail_if_unchanged := ff} `eq tactic.assumption,
   return (e, prf)

end norm_cast

namespace tactic
open expr norm_cast

/-- `aux_mod_cast e` runs `norm_cast` on `e` and returns the result. If `include_goal` is true, it
also normalizes the goal. -/
meta def aux_mod_cast (e : expr) (include_goal : bool := tt) : tactic expr :=
match e with
| local_const _ lc _ _ := do
  e ← get_local lc,
  replace_at derive [e] include_goal,
  get_local lc
| e := do
  t ← infer_type e,
  e ← assertv `this t e,
  replace_at derive [e] include_goal,
  get_local `this
end

/-- `exact_mod_cast e` runs `norm_cast` on the goal and `e`, and tries to use `e` to close the
goal. -/
meta def exact_mod_cast (e : expr) : tactic unit :=
decorate_error "exact_mod_cast failed:" $ do
  new_e ← aux_mod_cast e,
  exact new_e

/-- `apply_mod_cast e` runs `norm_cast` on the goal and `e`, and tries to apply `e`. -/
meta def apply_mod_cast (e : expr) : tactic (list (name × expr)) :=
decorate_error "apply_mod_cast failed:" $ do
  new_e ← aux_mod_cast e,
  apply new_e

/-- `assumption_mod_cast` runs `norm_cast` on the goal. For each local hypothesis `h`, it also
normalizes `h` and tries to use that to close the goal. -/
meta def assumption_mod_cast : tactic unit :=
decorate_error "assumption_mod_cast failed:" $ do
  let cfg : simp_config :=
  { fail_if_unchanged := ff,
    canonize_instances := ff,
    canonize_proofs := ff,
    proj := ff },
  replace_at derive [] tt,
  ctx ← local_context,
  ctx.mfirst (λ h, aux_mod_cast h ff >>= tactic.exact)

end tactic

namespace tactic.interactive
open tactic norm_cast

/--
Normalize casts at the given locations by moving them "upwards".
As opposed to simp, norm_cast can be used without necessarily closing the goal.
-/
meta def norm_cast (loc : parse location) : tactic unit :=
do
  ns ← loc.get_locals,
  tt ← replace_at derive ns loc.include_goal | fail "norm_cast failed to simplify",
  when loc.include_goal $ try tactic.reflexivity,
  when loc.include_goal $ try tactic.triv,
  when (¬ ns.empty) $ try tactic.contradiction

/--
Rewrite with the given rules and normalize casts between steps.
-/
meta def rw_mod_cast (rs : parse rw_rules) (loc : parse location) : tactic unit :=
decorate_error "rw_mod_cast failed:" $ do
  let cfg_norm : simp_config := {},
  let cfg_rw : rewrite_cfg := {},
  ns ← loc.get_locals,
  monad.mapm' (λ r : rw_rule, do
    save_info r.pos,
    replace_at derive ns loc.include_goal,
    rw ⟨[r], none⟩ loc {}
  ) rs.rules,
  replace_at derive ns loc.include_goal,
  skip

/--
Normalize the goal and the given expression, then close the goal with exact.
-/
meta def exact_mod_cast (e : parse texpr) : tactic unit :=
do
  e ← i_to_expr e <|> do
  { ty ← target,
    e ← i_to_expr_strict ``(%%e : %%ty),
    pty ← pp ty, ptgt ← pp e,
    fail ("exact_mod_cast failed, expression type not directly " ++
    "inferrable. Try:\n\nexact_mod_cast ...\nshow " ++
    to_fmt pty ++ ",\nfrom " ++ ptgt : format) },
  tactic.exact_mod_cast e

/--
Normalize the goal and the given expression, then apply the expression to the goal.
-/
meta def apply_mod_cast (e : parse texpr) : tactic unit :=
do
  e ← i_to_expr_for_apply e,
  concat_tags $ tactic.apply_mod_cast e

/--
Normalize the goal and every expression in the local context, then close the goal with assumption.
-/
meta def assumption_mod_cast : tactic unit :=
tactic.assumption_mod_cast

end tactic.interactive

namespace conv.interactive
open conv
open norm_cast (derive)

/-- the converter version of `norm_cast' -/
meta def norm_cast : conv unit := replace_lhs derive

end conv.interactive

-- TODO: move this elsewhere?
@[norm_cast] lemma ite_cast {α β} [has_lift_t α β]
  {c : Prop} [decidable c] {a b : α} :
  ↑(ite c a b) = ite c (↑a : β) (↑b : β) :=
by by_cases h : c; simp [h]

@[norm_cast] lemma dite_cast {α β} [has_lift_t α β]
  {c : Prop} [decidable c] {a : c → α} {b : ¬ c → α} :
  ↑(dite c a b) = dite c (λ h, (↑(a h) : β)) (λ h, (↑(b h) : β)) :=
by by_cases h : c; simp [h]

add_hint_tactic "norm_cast at *"

/--
The `norm_cast` family of tactics is used to normalize casts inside expressions.
It is basically a simp tactic with a specific set of lemmas to move casts
upwards in the expression.
Therefore it can be used more safely as a non-terminating tactic.
It also has special handling of numerals.

For instance, given an assumption
```lean
a b : ℤ
h : ↑a + ↑b < (10 : ℚ)
```

writing `norm_cast at h` will turn `h` into
```lean
h : a + b < 10
```

You can also use `exact_mod_cast`, `apply_mod_cast`, `rw_mod_cast`
or `assumption_mod_cast`.
Writing `exact_mod_cast h` and `apply_mod_cast h` will normalize the goal and
`h` before using `exact h` or `apply h`.
Writing `assumption_mod_cast` will normalize the goal and for every
expression `h` in the context it will try to normalize `h` and use
`exact h`.
`rw_mod_cast` acts like the `rw` tactic but it applies `norm_cast` between steps.

`push_cast` rewrites the expression to move casts toward the leaf nodes.
This uses `norm_cast` lemmas in the forward direction.
For example, `↑(a + b)` will be written to `↑a + ↑b`.
It is equivalent to `simp only with push_cast`.
It can also be used at hypotheses with `push_cast at h`
and with extra simp lemmas with `push_cast [int.add_zero]`.

```lean
example (a b : ℕ) (h1 : ((a + b : ℕ) : ℤ) = 10) (h2 : ((a + b + 0 : ℕ) : ℤ) = 10) :
  ((a + b : ℕ) : ℤ) = 10 :=
begin
  push_cast,
  push_cast at h1,
  push_cast [int.add_zero] at h2,
end
```

The implementation and behavior of the `norm_cast` family is described in detail at
<https://lean-forward.github.io/norm_cast/norm_cast.pdf>.
-/
add_tactic_doc
{ name := "norm_cast",
  category   := doc_category.tactic,
  decl_names := [``tactic.interactive.norm_cast, ``tactic.interactive.rw_mod_cast,
                 ``tactic.interactive.apply_mod_cast, ``tactic.interactive.assumption_mod_cast,
                 ``tactic.interactive.exact_mod_cast, ``tactic.interactive.push_cast],
  tags       := ["coercions", "simplification"] }

/--
The `norm_cast` attribute should be given to lemmas that describe the
behaviour of a coercion in regard to an operator, a relation, or a particular
function.

It only concerns equality or iff lemmas involving `↑`, `⇑` and `↥`, describing the behavior of
the coercion functions.
It does not apply to the explicit functions that define the coercions.

Examples:
```lean
@[norm_cast] theorem coe_nat_inj' {m n : ℕ} : (↑m : ℤ) = ↑n ↔ m = n

@[norm_cast] theorem coe_int_denom (n : ℤ) : (n : ℚ).denom = 1

@[norm_cast] theorem cast_id : ∀ n : ℚ, ↑n = n

@[norm_cast] theorem coe_nat_add (m n : ℕ) : (↑(m + n) : ℤ) = ↑m + ↑n

@[norm_cast] theorem cast_sub [add_group α] [has_one α] {m n} (h : m ≤ n) :
  ((n - m : ℕ) : α) = n - m

@[norm_cast] theorem coe_nat_bit0 (n : ℕ) : (↑(bit0 n) : ℤ) = bit0 ↑n

@[norm_cast] theorem cast_coe_nat (n : ℕ) : ((n : ℤ) : α) = n

@[norm_cast] theorem cast_one : ((1 : ℚ) : α) = 1
```

Lemmas tagged with `@[norm_cast]` are classified into three categories: `move`, `elim`, and
`squash`. They are classified roughly as follows:

* elim lemma:   LHS has 0 head coes and ≥ 1 internal coe
* move lemma:   LHS has 1 head coe and 0 internal coes,    RHS has 0 head coes and ≥ 1 internal coes
* squash lemma: LHS has ≥ 1 head coes and 0 internal coes, RHS has fewer head coes

`norm_cast` uses `move` and `elim` lemmas to factor coercions toward the root of an expression
and to cancel them from both sides of an equation or relation. It uses `squash` lemmas to clean
up the result.

Occasionally you may want to override the automatic classification.
You can do this by giving an optional `elim`, `move`, or `squash` parameter to the attribute.

```lean
@[simp, norm_cast elim] lemma nat_cast_re (n : ℕ) : (n : ℂ).re = n :=
by rw [← of_real_nat_cast, of_real_re]
```

Don't do this unless you understand what you are doing.

A full description of the tactic, and the use of each lemma category, can be found at
<https://lean-forward.github.io/norm_cast/norm_cast.pdf>.
-/
add_tactic_doc
{ name := "norm_cast attributes",
  category   := doc_category.attr,
  decl_names := [``norm_cast.norm_cast_attr],
  tags       := ["coercions", "simplification"] }


/-
Copyright (c) 2022 Andrew Yang. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Andrew Yang
-/
import algebraic_geometry.morphisms.ring_hom_properties
import topology.local_at_target

/-!

# Open immersions

A morphism is an open immersions if the underlying map of spaces is an open embedding
`f : X ⟶ U ⊆ Y`, and the sheaf map `Y(V) ⟶ f _* X(V)` is an iso for each `V ⊆ U`.

Most of the theories are developed in `algebraic_geometry/open_immersion`, and we provide the
remaining theorems analogous to other lemmas in `algebraic_geometry/morphisms/*`.

-/

noncomputable theory

open category_theory category_theory.limits opposite topological_space

universe u

namespace algebraic_geometry

variables {X Y Z : Scheme.{u}} (f : X ⟶ Y) (g : Y ⟶ Z)

lemma is_open_immersion_iff_stalk {f : X ⟶ Y} :
  is_open_immersion f ↔
    open_embedding f.1.base ∧ ∀ x, is_iso (PresheafedSpace.stalk_map f.1 x) :=
begin
  split,
  { intro h, exactI ⟨h.1, infer_instance⟩ },
  { rintro ⟨h₁, h₂⟩, exactI is_open_immersion.of_stalk_iso f h₁ }
end

lemma is_open_immersion_stable_under_composition :
  morphism_property.stable_under_composition @is_open_immersion :=
begin
  introsI X Y Z f g h₁ h₂, apply_instance
end

lemma is_open_immersion_respects_iso :
  morphism_property.respects_iso @is_open_immersion :=
begin
  apply is_open_immersion_stable_under_composition.respects_iso,
  intros _ _ _, apply_instance
end

lemma is_open_immersion_is_local_at_target : property_is_local_at_target @is_open_immersion :=
begin
  constructor,
  { exact is_open_immersion_respects_iso },
  { introsI, apply_instance },
  { intros X Y f 𝒰 H,
    rw is_open_immersion_iff_stalk,
    split,
    { apply (open_embedding_iff_open_embedding_of_supr_eq_top
        𝒰.supr_opens_range f.1.base.2).mpr,
      intro i,
      have := ((is_open_immersion_respects_iso.arrow_iso_iff
        (morphism_restrict_opens_range f (𝒰.map i))).mpr (H i)).1,
      rwa [arrow.mk_hom, morphism_restrict_val_base] at this },
    { intro x,
      have := arrow.iso_w (morphism_restrict_stalk_map f ((𝒰.map $ 𝒰.f $ f.1 x).opens_range)
        ⟨x, 𝒰.covers _⟩),
      dsimp only [arrow.mk_hom] at this,
      rw this,
      haveI : is_open_immersion (f ∣_ (𝒰.map $ 𝒰.f $ f.1 x).opens_range) :=
        (is_open_immersion_respects_iso.arrow_iso_iff
          (morphism_restrict_opens_range f (𝒰.map _))).mpr (H _),
      apply_instance } }
end

lemma is_open_immersion.open_cover_tfae {X Y : Scheme.{u}} (f : X ⟶ Y) :
  tfae [is_open_immersion f,
    ∃ (𝒰 : Scheme.open_cover.{u} Y), ∀ (i : 𝒰.J),
      is_open_immersion (pullback.snd : (𝒰.pullback_cover f).obj i ⟶ 𝒰.obj i),
    ∀ (𝒰 : Scheme.open_cover.{u} Y) (i : 𝒰.J),
      is_open_immersion (pullback.snd : (𝒰.pullback_cover f).obj i ⟶ 𝒰.obj i),
    ∀ (U : opens Y.carrier), is_open_immersion (f ∣_ U),
    ∀ {U : Scheme} (g : U ⟶ Y) [is_open_immersion g],
      is_open_immersion (pullback.snd : pullback f g ⟶ _),
    ∃ {ι : Type u} (U : ι → opens Y.carrier) (hU : supr U = ⊤),
      ∀ i, is_open_immersion (f ∣_ (U i))] :=
is_open_immersion_is_local_at_target.open_cover_tfae f

lemma is_open_immersion.open_cover_iff {X Y : Scheme.{u}}
  (𝒰 : Scheme.open_cover.{u} Y) (f : X ⟶ Y) :
  is_open_immersion f ↔ ∀ i, is_open_immersion (pullback.snd : pullback f (𝒰.map i) ⟶ _) :=
is_open_immersion_is_local_at_target.open_cover_iff f 𝒰

lemma is_open_immersion_stable_under_base_change :
  morphism_property.stable_under_base_change @is_open_immersion :=
morphism_property.stable_under_base_change.mk is_open_immersion_respects_iso $
  by { introsI X Y Z f g H, apply_instance }

end algebraic_geometry


/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Std.ShareCommon
import Lean.MetavarContext
import Lean.Environment
import Lean.Util.FoldConsts
import Lean.Meta.Basic
import Lean.Meta.Check

/-

This module provides functions for "closing" open terms and
creating auxiliary definitions. Here, we say a term is "open" if
it contains free/meta-variables.

The "closure" is performed by lambda abstracting the
free/meta-variables. Recall that in dependent type theory
lambda abstracting a let-variable may produce type incorrect terms.
For example, given the context
```lean
(n : Nat := 20)
(x : Vector α n)
(y : Vector α 20)
```
the term `x = y` is correct. However, its closure using lambda abstractions
is not.
```lean
fun (n : Nat) (x : Vector α n) (y : Vector α 20) => x = y
```
A previous version of this module would address this issue by
always use let-expressions to abstract let-vars. In the example above,
it would produce
```lean
let n : Nat := 20; fun (x : Vector α n) (y : Vector α 20) => x = y
```
This approach produces correct result, but produces unsatisfactory
results when we want to create auxiliary definitions.
For example, consider the context
```lean
(x : Nat)
(y : Nat := fact x)
```
and the term `h (g y)`, now suppose we want to create an auxiliary definition for `y`.
 The previous version of this module would compute the auxiliary definition
```lean
def aux := fun (x : Nat) => let y : Nat := fact x; h (g y)
```
and would return the term `aux x` as a substitute for `h (g y)`.
This is correct, but we will re-evaluate `fact x` whenever we use `aux`.
In this module, we produce
```lean
def aux := fun (y : Nat) => h (g y)
```
Note that in this particular case, it is safe to lambda abstract the let-varible `y`.
This module uses the following approach to decide whether it is safe or not to lambda
abstract a let-variable.
1) We enable zeta-expansion tracking in `MetaM`. That is, whenever we perform type checking
   if a let-variable needs to zeta expanded, we store it in the set `zetaFVarIds`.
   We say a let-variable is zeta expanded when we replace it with its value.
2) We use the `MetaM` type checker `check` to type check the expression we want to close,
   and the type of the binders.
3) If a let-variable is not in `zetaFVarIds`, we lambda abstract it.

Remark: We still use let-expressions for let-variables in `zetaFVarIds`, but we move the
`let` inside the lambdas. The idea is to make sure the auxiliary definition does not have
an interleaving of `lambda` and `let` expressions. Thus, if the let-variable occurs in
the type of one of the lambdas, we simply zeta-expand it there.
As a final example consider the context
```lean
(x_1 : Nat)
(x_2 : Nat)
(x_3 : Nat)
(x   : Nat := fact (10 + x_1 + x_2 + x_3))
(ty  : Type := Nat → Nat)
(f   : ty := fun x => x)
(n   : Nat := 20)
(z   : f 10)
```
and we use this module to compute an auxiliary definition for the term
```lean
(let y  : { v : Nat // v = n } := ⟨20, rfl⟩; y.1 + n + f x, z + 10)
```
we obtain
```lean
def aux (x : Nat) (f : Nat → Nat) (z : Nat) : Nat×Nat :=
let n : Nat := 20;
(let y : {v // v=n} := {val := 20, property := ex._proof_1}; y.val+n+f x, z+10)
```

BTW, this module also provides the `zeta : Bool` flag. When set to true, it
expands all let-variables occurring in the target expression.
-/

namespace Lean.Meta
namespace Closure

structure ToProcessElement where
  fvarId : FVarId
  newFVarId : FVarId
  deriving Inhabited

structure Context where
  zeta : Bool

structure State where
  visitedLevel          : LevelMap Level := {}
  visitedExpr           : ExprStructMap Expr := {}
  levelParams           : Array Name := #[]
  nextLevelIdx          : Nat := 1
  levelArgs             : Array Level := #[]
  newLocalDecls         : Array LocalDecl := #[]
  newLocalDeclsForMVars : Array LocalDecl := #[]
  newLetDecls           : Array LocalDecl := #[]
  nextExprIdx           : Nat := 1
  exprMVarArgs          : Array Expr := #[]
  exprFVarArgs          : Array Expr := #[]
  toProcess             : Array ToProcessElement := #[]

abbrev ClosureM := ReaderT Context $ StateRefT State MetaM

@[inline] def visitLevel (f : Level → ClosureM Level) (u : Level) : ClosureM Level := do
  if !u.hasMVar && !u.hasParam then
    pure u
  else
    let s ← get
    match s.visitedLevel.find? u with
    | some v => pure v
    | none   => do
      let v ← f u
      modify fun s => { s with visitedLevel := s.visitedLevel.insert u v }
      pure v

@[inline] def visitExpr (f : Expr → ClosureM Expr) (e : Expr) : ClosureM Expr := do
  if !e.hasLevelParam && !e.hasFVar && !e.hasMVar then
    pure e
  else
    let s ← get
    match s.visitedExpr.find? e with
    | some r => pure r
    | none   =>
      let r ← f e
      modify fun s => { s with visitedExpr := s.visitedExpr.insert e r }
      pure r

def mkNewLevelParam (u : Level) : ClosureM Level := do
  let s ← get
  let p := (`u).appendIndexAfter s.nextLevelIdx
  modify fun s => { s with levelParams := s.levelParams.push p, nextLevelIdx := s.nextLevelIdx + 1, levelArgs := s.levelArgs.push u }
  pure $ mkLevelParam p

partial def collectLevelAux : Level → ClosureM Level
  | u@(Level.succ v _)      => return u.updateSucc! (← visitLevel collectLevelAux v)
  | u@(Level.max v w _)     => return u.updateMax! (← visitLevel collectLevelAux v) (← visitLevel collectLevelAux w)
  | u@(Level.imax v w _)    => return u.updateIMax! (← visitLevel collectLevelAux v) (← visitLevel collectLevelAux w)
  | u@(Level.mvar mvarId _) => mkNewLevelParam u
  | u@(Level.param _ _)     => mkNewLevelParam u
  | u@(Level.zero _)        => pure u

def collectLevel (u : Level) : ClosureM Level := do
  -- u ← instantiateLevelMVars u
  visitLevel collectLevelAux u

def preprocess (e : Expr) : ClosureM Expr := do
  let e ← instantiateMVars e
  let ctx ← read
  -- If we are not zeta-expanding let-decls, then we use `check` to find
  -- which let-decls are dependent. We say a let-decl is dependent if its lambda abstraction is type incorrect.
  if !ctx.zeta then
    check e
  pure e

/--
  Remark: This method does not guarantee unique user names.
  The correctness of the procedure does not rely on unique user names.
  Recall that the pretty printer takes care of unintended collisions. -/
def mkNextUserName : ClosureM Name := do
  let s ← get
  let n := (`_x).appendIndexAfter s.nextExprIdx
  modify fun s => { s with nextExprIdx := s.nextExprIdx + 1 }
  pure n

def pushToProcess (elem : ToProcessElement) : ClosureM Unit :=
  modify fun s => { s with toProcess := s.toProcess.push elem }

partial def collectExprAux (e : Expr) : ClosureM Expr := do
  let collect (e : Expr) := visitExpr collectExprAux e
  match e with
  | Expr.proj _ _ s _    => return e.updateProj! (← collect s)
  | Expr.forallE _ d b _ => return e.updateForallE! (← collect d) (← collect b)
  | Expr.lam _ d b _     => return e.updateLambdaE! (← collect d) (← collect b)
  | Expr.letE _ t v b _  => return e.updateLet! (← collect t) (← collect v) (← collect b)
  | Expr.app f a _       => return e.updateApp! (← collect f) (← collect a)
  | Expr.mdata _ b _     => return e.updateMData! (← collect b)
  | Expr.sort u _        => return e.updateSort! (← collectLevel u)
  | Expr.const c us _    => return e.updateConst! (← us.mapM collectLevel)
  | Expr.mvar mvarId _   =>
    let mvarDecl ← getMVarDecl mvarId
    let type ← preprocess mvarDecl.type
    let type ← collect type
    let newFVarId ← mkFreshFVarId
    let userName ← mkNextUserName
    modify fun s => { s with
      newLocalDeclsForMVars := s.newLocalDeclsForMVars.push $ LocalDecl.cdecl arbitrary newFVarId userName type BinderInfo.default,
      exprMVarArgs          := s.exprMVarArgs.push e
    }
    return mkFVar newFVarId
  | Expr.fvar fvarId _ =>
    match (← read).zeta, (← getLocalDecl fvarId).value? with
    | true, some value => collect (← preprocess value)
    | _,    _          =>
      let newFVarId ← mkFreshFVarId
      pushToProcess ⟨fvarId, newFVarId⟩
      return mkFVar newFVarId
  | e => pure e

def collectExpr (e : Expr) : ClosureM Expr := do
  let e ← preprocess e
  visitExpr collectExprAux e

partial def pickNextToProcessAux (lctx : LocalContext) (i : Nat) (toProcess : Array ToProcessElement) (elem : ToProcessElement)
    : ToProcessElement × Array ToProcessElement :=
  if h : i < toProcess.size then
    let elem' := toProcess.get ⟨i, h⟩
    if (lctx.get! elem.fvarId).index < (lctx.get! elem'.fvarId).index then
      pickNextToProcessAux lctx (i+1) (toProcess.set ⟨i, h⟩ elem) elem'
    else
      pickNextToProcessAux lctx (i+1) toProcess elem
  else
    (elem, toProcess)

def pickNextToProcess? : ClosureM (Option ToProcessElement) := do
  let lctx ← getLCtx
  let s ← get
  if s.toProcess.isEmpty then
    pure none
  else
    modifyGet fun s =>
      let elem      := s.toProcess.back
      let toProcess := s.toProcess.pop
      let (elem, toProcess) := pickNextToProcessAux lctx 0 toProcess elem
      (some elem, { s with toProcess := toProcess })

def pushFVarArg (e : Expr) : ClosureM Unit :=
  modify fun s => { s with exprFVarArgs := s.exprFVarArgs.push e }

def pushLocalDecl (newFVarId : FVarId) (userName : Name) (type : Expr) (bi := BinderInfo.default) : ClosureM Unit := do
  let type ← collectExpr type
  modify fun s => { s with newLocalDecls := s.newLocalDecls.push <| LocalDecl.cdecl arbitrary newFVarId userName type bi }

partial def process : ClosureM Unit := do
  match (← pickNextToProcess?) with
  | none => pure ()
  | some ⟨fvarId, newFVarId⟩ =>
    let localDecl ← getLocalDecl fvarId
    match localDecl with
    | LocalDecl.cdecl _ _ userName type bi =>
      pushLocalDecl newFVarId userName type bi
      pushFVarArg (mkFVar fvarId)
      process
    | LocalDecl.ldecl _ _ userName type val _ =>
      let zetaFVarIds ← getZetaFVarIds
      if !zetaFVarIds.contains fvarId then
        /- Non-dependent let-decl

            Recall that if `fvarId` is in `zetaFVarIds`, then we zeta-expanded it
            during type checking (see `check` at `collectExpr`).

            Our type checker may zeta-expand declarations that are not needed, but this
            check is conservative, and seems to work well in practice. -/
        pushLocalDecl newFVarId userName type
        pushFVarArg (mkFVar fvarId)
        process
      else
        /- Dependent let-decl -/
        let type ← collectExpr type
        let val  ← collectExpr val
        modify fun s => { s with newLetDecls := s.newLetDecls.push <| LocalDecl.ldecl arbitrary newFVarId userName type val false }
        /- We don't want to interleave let and lambda declarations in our closure. So, we expand any occurrences of newFVarId
           at `newLocalDecls` -/
        modify fun s => { s with newLocalDecls := s.newLocalDecls.map (replaceFVarIdAtLocalDecl newFVarId val) }
        process

@[inline] def mkBinding (isLambda : Bool) (decls : Array LocalDecl) (b : Expr) : Expr :=
  let xs := decls.map LocalDecl.toExpr
  let b  := b.abstract xs
  decls.size.foldRev (init := b) fun i b =>
    let decl := decls[i]
    match decl with
    | LocalDecl.cdecl _ _ n ty bi  =>
      let ty := ty.abstractRange i xs
      if isLambda then
        Lean.mkLambda n bi ty b
      else
        Lean.mkForall n bi ty b
    | LocalDecl.ldecl _ _ n ty val nonDep =>
      if b.hasLooseBVar 0 then
        let ty  := ty.abstractRange i xs
        let val := val.abstractRange i xs
        mkLet n ty val b nonDep
      else
        b.lowerLooseBVars 1 1

def mkLambda (decls : Array LocalDecl) (b : Expr) : Expr :=
  mkBinding true decls b

def mkForall (decls : Array LocalDecl) (b : Expr) : Expr :=
  mkBinding false decls b

structure MkValueTypeClosureResult where
  levelParams : Array Name
  type        : Expr
  value       : Expr
  levelArgs   : Array Level
  exprArgs    : Array Expr

def mkValueTypeClosureAux (type : Expr) (value : Expr) : ClosureM (Expr × Expr) := do
  resetZetaFVarIds
  withTrackingZeta do
    let type  ← collectExpr type
    let value ← collectExpr value
    process
    pure (type, value)

def mkValueTypeClosure (type : Expr) (value : Expr) (zeta : Bool) : MetaM MkValueTypeClosureResult := do
  let ((type, value), s) ← ((mkValueTypeClosureAux type value).run { zeta := zeta }).run {}
  let newLocalDecls := s.newLocalDecls.reverse ++ s.newLocalDeclsForMVars
  let newLetDecls   := s.newLetDecls.reverse
  let type  := mkForall newLocalDecls (mkForall newLetDecls type)
  let value := mkLambda newLocalDecls (mkLambda newLetDecls value)
  pure {
    type        := type,
    value       := value,
    levelParams := s.levelParams,
    levelArgs   := s.levelArgs,
    exprArgs    := s.exprFVarArgs.reverse ++ s.exprMVarArgs
  }

end Closure

/--
  Create an auxiliary definition with the given name, type and value.
  The parameters `type` and `value` may contain free and meta variables.
  A "closure" is computed, and a term of the form `name.{u_1 ... u_n} t_1 ... t_m` is
  returned where `u_i`s are universe parameters and metavariables `type` and `value` depend on,
  and `t_j`s are free and meta variables `type` and `value` depend on. -/
def mkAuxDefinition (name : Name) (type : Expr) (value : Expr) (zeta : Bool := false) (compile : Bool := true) : MetaM Expr := do
  trace[Meta.debug] "{name} : {type} := {value}"
  let result ← Closure.mkValueTypeClosure type value zeta
  let env ← getEnv
  let decl := Declaration.defnDecl {
    name        := name,
    levelParams := result.levelParams.toList,
    type        := result.type,
    value       := result.value,
    hints       := ReducibilityHints.regular (getMaxHeight env result.value + 1),
    safety      := if env.hasUnsafe result.type || env.hasUnsafe result.value then DefinitionSafety.unsafe else DefinitionSafety.safe
  }
  trace[Meta.debug] "{name} : {result.type} := {result.value}"
  addDecl decl
  if compile then
    compileDecl decl
  return mkAppN (mkConst name result.levelArgs.toList) result.exprArgs

/-- Similar to `mkAuxDefinition`, but infers the type of `value`. -/
def mkAuxDefinitionFor (name : Name) (value : Expr) : MetaM Expr := do
  let type ← inferType value
  let type := type.headBeta
  mkAuxDefinition name type value

end Lean.Meta