import tactic.find

#find (ff = _)

constants
(f : nat → nat)
(f₀ : f 0 = 0)
(f₁ : f 1 = 1)
(f₂ : f 2 = 1)
(f_next : ∀ n : nat, f (n + 3) = f (n + 2) + f (n + 1) ∨ f (n + 3) = f (n + 1) + f n)

namespace evaluation

meta def hasElement : list nat → nat → bool
| [] _ := ff
| (x :: xs) e := if e = x then tt else hasElement xs e

#eval hasElement [1,2,3] 1 -- tt
#eval hasElement [1,2,3] 4 -- ff

#reduce hasElement [1,2,3] 1 -- tt

-- def union (lx ly : list nat) : list nat :=
-- match lx, ly with
-- | [], listy := listy
-- | listx, [] := listx
-- | (x :: xs), listy := if hasElem listy x then union xs listy else union xs (x :: listy)
-- end
-- @[reducible]
meta def union : list nat → list nat → list nat
| [] ys := ys
| (x :: xs) ys := if hasElement ys x then union xs ys else union xs (x :: ys)

#print union._main

#eval union [1,2,3,4,5] [1] -- [5, 4, 3, 2, 1]
#eval union [1,2,3] [4,5,6] -- [3, 2, 1, 4, 5, 6]

#eval union [1,2,3] [] -- [3, 2, 1]
#reduce union [1,2,3] [] -- [union._main] [2, 3] [1], why reduce doesn't return the same result as eval?

end evaluation

namespace lemmas

set_option trace.simplify.rewrite true

constant hasElement : list nat → nat → bool

namespace hasElement

axiom base (n : nat) : hasElement [] n = ff
axiom step (x : nat) (xs : list nat) : hasElement (x :: xs) x = tt
axiom step_not_eq (x y : nat) (xs : list nat) (h : x ≠ y) :
  hasElement (x :: xs) y = hasElement xs y

lemma empty_not_contains_any (x : nat) : ¬ hasElement [] x = tt :=
begin
  rw hasElement.base _,
  -- change ¬ false = true, -- fail
  -- rw not.elim, -- fail
  -- rw not_false_iff, -- fail
  apply not.intro,
  exact bool.ff_ne_tt,
end

lemma single {x : nat} : hasElement [x] x :=
  hasElement.step x []

lemma two (x y : nat) : hasElement [x, y] y :=
begin
  -- if x = y => tt
  -- if x ≠ y => hasElement [y] y => tt [by single]
  by_cases (x = y),
  {
    rw ←h,
    exact hasElement.step x [x],
  },
  {
    change x ≠ y at h,
    -- rw hasElement.step_not_eq _ _ _ h, -- OK
    rw hasElement.step_not_eq,
    { exact @single y },
    { exact h },
  },
end

example : ¬ hasElement [0] 1 :=
begin
  rw hasElement.step_not_eq,
  apply empty_not_contains_any,
  exact zero_ne_one,
end

#print two

end hasElement

variables {x : ℕ} {xs ys : list ℕ}

constant union : list nat → list nat → list nat

axiom union_base (xs : list nat) :
  union [] xs = xs

axiom union_step₁ (h : hasElement ys x) :
  union (x :: xs) ys = union xs ys

axiom union_step₂ (h : hasElement ys x → false) :
  union (x :: xs) ys = union xs (x :: ys)

lemma union_single : union [x] [x] = [x] :=
begin
  have h₁ :                    union []  [x] = [x]          := union_base [x],
  have h₂ : hasElement [x] x → union [x] [x] = union [] [x] := union_step₁,
  have h₃ :                    union [x] [x] = union [] [x] := h₂ hasElement.single,
  apply eq.trans h₃,
  exact h₁,
end

lemma union_single' : union [x] [x] = [x] :=
begin
  rw union_step₁,
  rw union_base,
  apply hasElement.single,
end

#print notation >>=

open tactic
meta def trace_all : tactic unit :=
   do n ← num_goals
    , trace "num_goals ="
    , trace n
    , trace_state
    , trace "---------------------------------------------"
    , trace_result

lemma union_single'' : union [x] [x] = [x] := by {
  apply eq.trans (union_step₁ hasElement.single),
  trace_all,
  sorry
}

meta def test_tactic : tactic unit :=
do
  define `x (expr.const `nat [])
, trace "-- after assert --"
, trace_state
, n ← get_local `n
, trace (infer_type n)

example (n : nat) : n = n := by { test_tactic, sorry }

example : (1 = 2) → (2 = 3) → (3 = 1) := begin
  assume h1 h2,
  apply eq.symm,
  apply eq.subst h2,
  apply eq.subst h1,
  apply eq.refl,
  trace_result,
  trace_call_stack,
end

example (a : nat) (b : bool) (c : int) : true := by {
  do
    a ← get_local `a
  , b ← get_local `b
  , c ← get_local `c
  , infer_type a >>= trace
  , infer_type b >>= trace
  , infer_type c >>= trace
  , interactive.sorry
}

set_option trace.app_builder true
-- set_option pp.all true

example (a b c : ℕ) : true := by {
  do
    a ← get_local `a
  , x ← mk_app `nat.succ [a] -- x : expr
  , r ← to_expr ```(%%x + b * c + (λ n, n) 100) -- to_expr : pexpr → tactic expr
  , trace r -- a.succ + b * c + (λ (n : ℕ), n) 100
  , infer_type r >>= trace -- ℕ
  , interactive.trivial -- or exact_dec_trivial
}

example (a b c : nat) (H1 : a = b) (H2 : b = c) : a = c :=
by do
  --  refine ```(eq.trans H1 _),
   interactive.refine ```(eq.trans H1 _),
   trace_state,
   assumption

end lemmas

meta def solve (n : nat) (f : nat → nat) :=
match n, f with
| 0, f := 0
| _, _ := 1
end

-- solve n = [f 0, f 1, f 2, ... f (n-1)]
-- solve 0 = [[]]
-- solve 1 = [[0]]
-- solve 2 = [[0], [1]]
-- solve 3 = [[0], [1], [1]]
-- solve 4 = [[0], [1], [1], [1, 2]]
-- solve 5 = [[0], [1], [1], [1, 2], X],
--    where X = (1 + [1, 2]) ∪ (1 + 1) = [2, 3] ∪ [2] = [2, 3]

#print add_lt_add_of_le_of_lt -- a ≤ b → c < d → a + c < b + d

#print add_pos_of_nonneg_of_pos -- 0 ≤ a → 0 < b → 0 < a + b
-- (ha : 0 ≤ a) (hb : 0 < b),
  -- zero_add 0 ▸ add_lt_add_of_le_of_lt ha hb

#print notation ▸ -- eq.subst #1 #0
#print notation + -- has_add.add #1 #0
#print eq.subst


example : True := by
  simp (config := (fun (c : Lean.Meta.Simp.Config) => { c with arith := true }) {})


import for_mathlib.open_embeddings

import sheaves.sheaf_of_topological_rings
import sheaves.stalk_of_rings

universe variable u

open topological_space

variables {X : Type u} {Y : Type u} {Z : Type u}
variables [topological_space X] [topological_space Y] [topological_space Z]

namespace presheaf_of_rings
variables {F : presheaf_of_rings X} {G : presheaf_of_rings Y} {H : presheaf_of_rings Z}

structure f_map
  (F : presheaf_of_rings X) (G : presheaf_of_rings Y) :=
(f : X → Y)
(hf : continuous f)
(f_flat : ∀ V : opens Y, G V → F (hf.comap V))
(f_flat_is_ring_hom : ∀ V : opens Y, is_ring_hom (f_flat V))
(presheaf_f_flat : ∀ V W : opens Y, ∀ (hWV : W ⊆ V),
  ∀ s : G V, F.res _ _ (hf.comap_mono hWV) (f_flat V s) = f_flat W (G.res V W hWV s))

attribute [instance] f_map.f_flat_is_ring_hom

def f_map_id (F : presheaf_of_rings X) : presheaf_of_rings.f_map F F :=
{ f := λ x, x,
  hf := continuous_id,
  f_flat := λ U, F.res _ _ (λ _ hx, hx),
  f_flat_is_ring_hom := λ U, presheaf_of_rings.res_is_ring_hom _ _ _ _,
  presheaf_f_flat :=  λ U V hVU s, begin
      rw ←F.to_presheaf.Hcomp',
      rw ←F.to_presheaf.Hcomp',
    end }

def f_map.comp (a : presheaf_of_rings.f_map F G) (b : presheaf_of_rings.f_map G H) :
presheaf_of_rings.f_map F H :=
{ f := λ x, b.f (a.f x),
  hf := b.hf.comp a.hf,
  f_flat := λ V s, (a.f_flat (b.hf.comap V)) ((b.f_flat V) s),
  f_flat_is_ring_hom := λ V, show (is_ring_hom ((a.f_flat (b.hf.comap V)) ∘ (b.f_flat V))), from is_ring_hom.comp _ _,
  presheaf_f_flat := λ V W hWV s,
  begin
    rw ←b.presheaf_f_flat V W hWV s,
    rw ←a.presheaf_f_flat (b.hf.comap V) (b.hf.comap W) (b.hf.comap_mono hWV),
    refl,
  end }

@[simp] lemma f_map.id_comp (a : presheaf_of_rings.f_map F G) :
  (presheaf_of_rings.f_map_id F).comp a = a :=
begin
  cases a, delta presheaf_of_rings.f_map_id presheaf_of_rings.f_map.comp,
  congr, funext V s, dsimp,
  show _ = id _, apply congr_fun, exact F.to_presheaf.Hid _,
end

@[simp] lemma f_map.comp_id (a : presheaf_of_rings.f_map F G) :
  a.comp (presheaf_of_rings.f_map_id G) = a :=
begin
  cases a with f hf f_flat f_flat_is_ring_hom presheaf_f_flat,
  delta presheaf_of_rings.f_map_id presheaf_of_rings.f_map.comp,
  congr, funext V s, dsimp,
  rw ← presheaf_f_flat,
  show _ = id _, apply congr_fun, exact F.to_presheaf.Hid _,
end

end presheaf_of_rings

namespace presheaf_of_topological_rings
variables {F : presheaf_of_topological_rings X} {G : presheaf_of_topological_rings Y} {H : presheaf_of_topological_rings Z}

structure f_map (F : presheaf_of_topological_rings X) (G : presheaf_of_topological_rings Y) :=
(f : X → Y)
(hf : continuous f)
(f_flat : ∀ V : opens Y, G V → F (hf.comap V))
[f_flat_is_ring_hom : ∀ V : opens Y, is_ring_hom (f_flat V)]
(cont_f_flat : ∀ V : opens Y, continuous (f_flat V))
(presheaf_f_flat : ∀ V W : opens Y, ∀ (hWV : W ⊆ V),
  ∀ s : G V, F.res _ _ (hf.comap_mono hWV) (f_flat V s) = f_flat W (G.res V W hWV s))

instance f_map_flat.is_ring_hom (f : presheaf_of_topological_rings.f_map F G) (V : opens Y) :
  is_ring_hom (f.f_flat V) := f.f_flat_is_ring_hom V

attribute [instance] presheaf_of_topological_rings.f_map.f_flat_is_ring_hom

def f_map.to_presheaf_of_rings_f_map
  {X : Type u} [topological_space X] {Y : Type u} [topological_space Y]
  {F : presheaf_of_topological_rings X} {G : presheaf_of_topological_rings Y}
  (f : presheaf_of_topological_rings.f_map F G) :
  presheaf_of_rings.f_map F.to_presheaf_of_rings G.to_presheaf_of_rings :=
{ ..f}

@[ext]
lemma presheaf_of_topological_rings.f_map.ext
  {X : Type u} [topological_space X] {Y : Type u} [topological_space Y]
  {F : presheaf_of_topological_rings X} {G : presheaf_of_topological_rings Y}
  (a b : F.f_map G) (h : a.to_presheaf_of_rings_f_map = b.to_presheaf_of_rings_f_map) :
  a = b :=
begin
  cases a, cases b,
  dsimp [f_map.to_presheaf_of_rings_f_map] at h,
  injections,
  simp [*]
end

@[simp] lemma f_map.to_presheaf_of_rings_f_map_f (f : presheaf_of_topological_rings.f_map F G) :
  f.to_presheaf_of_rings_f_map.f = f.f := rfl

def f_map_id (F : presheaf_of_topological_rings X) :
  presheaf_of_topological_rings.f_map F F :=
{ cont_f_flat := λ U, begin
      show continuous (((F.to_presheaf_of_rings).to_presheaf).res U (continuous.comap continuous_id U) _),
      convert continuous_id,
      { simp [continuous.comap_id U] },
      { simp [continuous.comap_id U] },
      convert heq_of_eq (F.Hid U),
        rw continuous.comap_id U,
      exact continuous.comap_id U,
    end,
  ..presheaf_of_rings.f_map_id F.to_presheaf_of_rings }

@[simp] lemma f_map.to_presheaf_of_rings_f_map_id (F : presheaf_of_topological_rings X) :
  (presheaf_of_topological_rings.f_map_id F).to_presheaf_of_rings_f_map =
  presheaf_of_rings.f_map_id F.to_presheaf_of_rings := rfl

@[simp] lemma f_map_id_apply (F : presheaf_of_topological_rings X) (x : X) :
  (presheaf_of_topological_rings.f_map_id F).f x = x := rfl

def f_map.comp (a : presheaf_of_topological_rings.f_map F G) (b : presheaf_of_topological_rings.f_map G H) :
  presheaf_of_topological_rings.f_map F H :=
{ cont_f_flat := λ V, (a.cont_f_flat _).comp (b.cont_f_flat _),
  .. a.to_presheaf_of_rings_f_map.comp b.to_presheaf_of_rings_f_map }

@[simp] lemma f_map.comp_f (a : presheaf_of_topological_rings.f_map F G) (b : presheaf_of_topological_rings.f_map G H) :
  (a.comp b).f = b.f ∘ a.f := rfl

@[simp] lemma f_map.comp_to_presheaf_of_rings_f_map
  (a : presheaf_of_topological_rings.f_map F G) (b : presheaf_of_topological_rings.f_map G H) :
  (a.comp b).to_presheaf_of_rings_f_map =
  a.to_presheaf_of_rings_f_map.comp b.to_presheaf_of_rings_f_map := rfl

@[simp] lemma f_map.id_comp (a : presheaf_of_topological_rings.f_map F G) :
  (presheaf_of_topological_rings.f_map_id F).comp a = a :=
by ext; simp

@[simp] lemma f_map.comp_id (a : presheaf_of_topological_rings.f_map F G) :
  a.comp (presheaf_of_topological_rings.f_map_id G) = a :=
by ext; simp

end presheaf_of_topological_rings

open_locale classical

/-- The map on stalks induced from an f-map -/
noncomputable def stalk_map {F : presheaf_of_rings X} {G : presheaf_of_rings Y}
  (f : F.f_map G) (x : X) :
  stalk_of_rings G (f.f x) → stalk_of_rings F x :=
to_stalk.rec G (f.f x) (stalk_of_rings F x)
  (λ V hfx s, ⟦⟨f.hf.comap V, hfx, f.f_flat V s⟩⟧)
  (λ V W H r hfx, quotient.sound begin
    use [f.hf.comap V, hfx, set.subset.refl _, f.hf.comap_mono H],
    erw F.to_presheaf.Hid,
    symmetry,
    apply f.presheaf_f_flat
  end )

namespace stalk_map
variables {F : presheaf_of_rings X} {G : presheaf_of_rings Y} {H : presheaf_of_rings Z}
variables (f : F.f_map G) (g : G.f_map H)

instance (F : presheaf_of_rings X) (x : X) :
  comm_ring (quotient (stalk.setoid (F.to_presheaf) x)) :=
stalk_of_rings_is_comm_ring F x

instance f_flat_is_ring_hom (x : X) (V : opens Y) (hfx : f.f x ∈ V) :
  is_ring_hom (λ (s : G.F V), (⟦⟨f.hf.comap V, hfx, f.f_flat V s⟩⟧ : stalk_of_rings F x)) :=
begin
  show is_ring_hom ((to_stalk F x (f.hf.comap V) hfx) ∘ (f.f_flat V)),
  refine is_ring_hom.comp _ _,
end

instance (x : X) : is_ring_hom (stalk_map f x) := to_stalk.rec_is_ring_hom _ _ _ _ _

@[simp] lemma stalk_map_id (F : presheaf_of_rings X) (x : X) (s : stalk_of_rings F x) :
  stalk_map (presheaf_of_rings.f_map_id F) x s = s :=
begin
  induction s,
    apply quotient.sound,
    use s.U,
    use s.HxU,
    use (le_refl s.U),
    use (le_refl s.U),
    symmetry,
    convert (F.to_presheaf.Hcomp' _ _ _ _ _ s.s),
  refl,
end

@[simp] lemma stalk_map_id' (F : presheaf_of_rings X) (x : X) :
  stalk_map (presheaf_of_rings.f_map_id F) x = id := by ext; apply stalk_map_id

lemma stalk_map_comp (x : X) (s : stalk_of_rings H (g.f (f.f x))) :
  stalk_map (f.comp g) x s = stalk_map f x (stalk_map g (f.f x) s) :=
begin
  induction s,
    apply quotient.sound,
    use f.hf.comap (g.hf.comap s.U),
    use s.HxU,
    existsi _, swap, intros t ht, exact ht,
    existsi _, swap, intros t ht, exact ht,
    refl,
  refl,
end

@[simp] lemma stalk_map_comp' (x : X) :
  stalk_map (f.comp g) x = (stalk_map f x) ∘ (stalk_map g (f.f x)) :=
by ext; apply stalk_map_comp

end stalk_map

namespace presheaf_of_rings

def restrict (U : opens X) (G : presheaf_of_rings X) :
  presheaf_of_rings U :=
{ F := λ V, G.F (topological_space.opens.map U V),
  res := λ V W HWV, G.res _ _ (topological_space.opens.map_mono HWV),
  Hid := λ V, G.Hid (topological_space.opens.map U V),
  Hcomp := λ V₁ V₂ V₃ H12 H23, G.Hcomp (topological_space.opens.map U V₁)
    (topological_space.opens.map U V₂) (topological_space.opens.map U V₃)
    (topological_space.opens.map_mono H12) (topological_space.opens.map_mono H23),
  Fring := λ V, G.Fring (topological_space.opens.map U V),
  res_is_ring_hom := λ V W HWV, G.res_is_ring_hom (topological_space.opens.map U V)
    (topological_space.opens.map U W) (topological_space.opens.map_mono HWV) }

variables {U : opens X} (G : presheaf_of_rings X) (u : U)

noncomputable def restrict_stalk_map :
  stalk_of_rings (G.restrict U) u → stalk_of_rings G u :=
to_stalk.rec (G.restrict U) u (stalk_of_rings G u)
  (λ V hu, to_stalk G u (topological_space.opens.map U V) ( opens.map_mem_of_mem hu))
  (λ W V HWV s huW, quotient.sound (begin
    use [(topological_space.opens.map U W), opens.map_mem_of_mem huW],
    use [(set.subset.refl (topological_space.opens.map U W)), topological_space.opens.map_mono HWV],
    rw G.Hid (topological_space.opens.map U W),
    refl,
  end))

instance : is_ring_hom (G.restrict_stalk_map u) :=
by delta restrict_stalk_map; apply_instance

end presheaf_of_rings

def presheaf_of_topological_rings.restrict  (U : opens X) (G : presheaf_of_topological_rings X) :
  presheaf_of_topological_rings U :=
{ Ftop := λ V, G.Ftop (topological_space.opens.map U V),
  Ftop_ring := λ V, G.Ftop_ring (topological_space.opens.map U V),
  res_continuous := λ V W HWV, G.res_continuous (topological_space.opens.map U V)
    (topological_space.opens.map U W) (topological_space.opens.map_mono HWV),
..presheaf_of_rings.restrict U G.to_presheaf_of_rings }


/-
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 Lean.ScopedEnvExtension
import Lean.Util.Recognizers
import Lean.Meta.LevelDefEq
import Lean.Meta.DiscrTree
import Lean.Meta.AppBuilder
import Lean.Meta.Tactic.AuxLemma
namespace Lean.Meta

/--
  The fields `levelParams` and `proof` are used to encode the proof of the simp lemma.
  If the `proof` is a global declaration `c`, we store `Expr.const c []` at `proof` without the universe levels, and `levelParams` is set to `#[]`
  When using the lemma, we create fresh universe metavariables.
  Motivation: most simp lemmas are global declarations, and this approach is faster and saves memory.

  The field `levelParams` is not empty only when we elaborate an expression provided by the user, and it contains universe metavariables.
  Then, we use `abstractMVars` to abstract the universe metavariables and create new fresh universe parameters that are stored at the field `levelParams`.
-/
structure SimpLemma where
  keys        : Array DiscrTree.Key := #[]
  levelParams : Array Name := #[] -- non empty for local universe polymorhic proofs.
  proof       : Expr
  priority    : Nat  := eval_prio default
  post        : Bool := true
  perm        : Bool := false -- true is lhs and rhs are identical modulo permutation of variables
  name?       : Option Name := none -- for debugging and tracing purposes
  deriving Inhabited

def SimpLemma.getName (s : SimpLemma) : Name :=
  match s.name? with
  | some n => n
  | none   => "<unknown>"

instance : ToFormat SimpLemma where
  format s :=
    let perm := if s.perm then ":perm" else ""
    let name := format s.getName
    let prio := f!":{s.priority}"
    name ++ prio ++ perm

instance : ToMessageData SimpLemma where
  toMessageData s := format s

instance : BEq SimpLemma where
  beq e₁ e₂ := e₁.proof == e₂.proof

structure SimpLemmas where
  pre         : DiscrTree SimpLemma := DiscrTree.empty
  post        : DiscrTree SimpLemma := DiscrTree.empty
  lemmaNames  : Std.PHashSet Name := {}
  toUnfold    : Std.PHashSet Name := {}
  erased      : Std.PHashSet Name := {}
  deriving Inhabited

def addSimpLemmaEntry (d : SimpLemmas) (e : SimpLemma) : SimpLemmas :=
  if e.post then
    { d with post := d.post.insertCore e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
  else
    { d with pre := d.pre.insertCore e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
where
  updateLemmaNames (s : Std.PHashSet Name) : Std.PHashSet Name :=
    match e.name? with
    | none => s
    | some name => s.insert name

def SimpLemmas.addDeclToUnfold (d : SimpLemmas) (declName : Name) : SimpLemmas :=
  { d with toUnfold := d.toUnfold.insert declName }

def SimpLemmas.isDeclToUnfold (d : SimpLemmas) (declName : Name) : Bool :=
  d.toUnfold.contains declName

def SimpLemmas.isLemma (d : SimpLemmas) (declName : Name) : Bool :=
  d.lemmaNames.contains declName

def SimpLemmas.eraseCore [Monad m] [MonadError m] (d : SimpLemmas) (declName : Name) : m SimpLemmas := do
  return { d with erased := d.erased.insert declName, lemmaNames := d.lemmaNames.erase declName, toUnfold := d.toUnfold.erase declName }

def SimpLemmas.erase [Monad m] [MonadError m] (d : SimpLemmas) (declName : Name) : m SimpLemmas := do
  unless d.isLemma declName || d.isDeclToUnfold declName do
    throwError "'{declName}' does not have [simp] attribute"
  d.eraseCore declName

private partial def isPerm : Expr → Expr → MetaM Bool
  | Expr.app f₁ a₁ _, Expr.app f₂ a₂ _ => isPerm f₁ f₂ <&&> isPerm a₁ a₂
  | Expr.mdata _ s _, t => isPerm s t
  | s, Expr.mdata _ t _ => isPerm s t
  | s@(Expr.mvar ..), t@(Expr.mvar ..) => isDefEq s t
  | Expr.forallE n₁ d₁ b₁ _, Expr.forallE n₂ d₂ b₂ _ => isPerm d₁ d₂ <&&> withLocalDeclD n₁ d₁ fun x => isPerm (b₁.instantiate1 x) (b₂.instantiate1 x)
  | Expr.lam n₁ d₁ b₁ _, Expr.lam n₂ d₂ b₂ _ => isPerm d₁ d₂ <&&> withLocalDeclD n₁ d₁ fun x => isPerm (b₁.instantiate1 x) (b₂.instantiate1 x)
  | Expr.letE n₁ t₁ v₁ b₁ _, Expr.letE n₂ t₂ v₂ b₂ _ =>
    isPerm t₁ t₂ <&&> isPerm v₁ v₂ <&&> withLetDecl n₁ t₁ v₁ fun x => isPerm (b₁.instantiate1 x) (b₂.instantiate1 x)
  | Expr.proj _ i₁ b₁ _, Expr.proj _ i₂ b₂ _ => i₁ == i₂ <&&> isPerm b₁ b₂
  | s, t => s == t

private partial def shouldPreprocess (type : Expr) : MetaM Bool :=
  forallTelescopeReducing type fun xs result => return !result.isEq

private partial def preprocess (e type : Expr) (inv : Bool) : MetaM (List (Expr × Expr)) := do
  let type ← whnf type
  if type.isForall then
    forallTelescopeReducing type fun xs type => do
      let e := mkAppN e xs
      let ps ← preprocess e type inv
      ps.mapM fun (e, type) =>
        return (← mkLambdaFVars xs e, ← mkForallFVars xs type)
  else if let some (_, lhs, rhs) := type.eq? then
    if inv then
      let type ← mkEq rhs lhs
      let e    ← mkEqSymm e
      return [(e, type)]
    else
      return [(e, type)]
  else if let some (lhs, rhs) := type.iff? then
    if inv then
      let type ← mkEq rhs lhs
      let e    ← mkEqSymm (← mkPropExt e)
      return [(e, type)]
    else
      let type ← mkEq lhs rhs
      let e    ← mkPropExt e
      return [(e, type)]
  else if let some (_, lhs, rhs) := type.ne? then
    if inv then
      throwError "invalid '←' modifier in rewrite rule to 'False'"
    let type ← mkEq (← mkEq lhs rhs) (mkConst ``False)
    let e    ← mkEqFalse e
    return [(e, type)]
  else if let some p := type.not? then
    if inv then
      throwError "invalid '←' modifier in rewrite rule to 'False'"
    let type ← mkEq p (mkConst ``False)
    let e    ← mkEqFalse e
    return [(e, type)]
  else if let some (type₁, type₂) := type.and? then
    let e₁ := mkProj ``And 0 e
    let e₂ := mkProj ``And 1 e
    return (← preprocess e₁ type₁ inv) ++ (← preprocess e₂ type₂ inv)
  else
    if inv then
      throwError "invalid '←' modifier in rewrite rule to 'True'"
    let type ← mkEq type (mkConst ``True)
    let e    ← mkEqTrue e
    return [(e, type)]

private def checkTypeIsProp (type : Expr) : MetaM Unit :=
  unless (← isProp type) do
    throwError "invalid 'simp', proposition expected{indentExpr type}"

private def mkSimpLemmaCore (e : Expr) (levelParams : Array Name) (proof : Expr) (post : Bool) (prio : Nat) (name? : Option Name) : MetaM SimpLemma := do
  let type ← instantiateMVars (← inferType e)
  withNewMCtxDepth do
    let (xs, _, type) ← withReducible <| forallMetaTelescopeReducing type
    let type ← whnfR type
    let (keys, perm) ←
      match type.eq? with
      | some (_, lhs, rhs) => pure (← DiscrTree.mkPath lhs, ← isPerm lhs rhs)
      | none => throwError "unexpected kind of 'simp' theorem{indentExpr type}"
    return { keys := keys, perm := perm, post := post, levelParams := levelParams, proof := proof, name? := name?, priority := prio }

private def mkSimpLemmasFromConst (declName : Name) (post : Bool) (inv : Bool) (prio : Nat) : MetaM (Array SimpLemma) := do
  let cinfo ← getConstInfo declName
  let val := mkConst declName (cinfo.levelParams.map mkLevelParam)
  withReducible do
    let type ← inferType val
    checkTypeIsProp type
    if inv || (← shouldPreprocess type) then
      let mut r := #[]
      for (val, type) in (← preprocess val type inv) do
        let auxName ← mkAuxLemma cinfo.levelParams type val
        r := r.push <| (← mkSimpLemmaCore (mkConst auxName (cinfo.levelParams.map mkLevelParam)) #[] (mkConst auxName) post prio declName)
      return r
    else
      #[← mkSimpLemmaCore (mkConst declName (cinfo.levelParams.map mkLevelParam)) #[] (mkConst declName) post prio declName]

inductive SimpEntry where
  | lemma    : SimpLemma → SimpEntry
  | toUnfold : Name → SimpEntry
  deriving Inhabited

abbrev SimpExtension := SimpleScopedEnvExtension SimpEntry SimpLemmas

def SimpExtension.getLemmas (ext : SimpExtension) : CoreM SimpLemmas :=
  return ext.getState (← getEnv)

def addSimpLemma (ext : SimpExtension) (declName : Name) (post : Bool) (inv : Bool) (attrKind : AttributeKind) (prio : Nat) : MetaM Unit := do
  let simpLemmas ← mkSimpLemmasFromConst declName post inv prio
  for simpLemma in simpLemmas do
    ext.add (SimpEntry.lemma simpLemma) attrKind

def mkSimpAttr (attrName : Name) (attrDescr : String) (ext : SimpExtension) : IO Unit :=
  registerBuiltinAttribute {
    name  := attrName
    descr := attrDescr
    add   := fun declName stx attrKind =>
      let go : MetaM Unit := do
        let info ← getConstInfo declName
        if (← isProp info.type) then
          let post :=
            if stx[1].isNone then true else stx[1][0].getKind == ``Lean.Parser.Tactic.simpPost
          let prio ← getAttrParamOptPrio stx[2]
          addSimpLemma ext declName post (inv := false) attrKind prio
        else if info.hasValue then
          ext.add (SimpEntry.toUnfold declName) attrKind
        else
          throwError "invalid 'simp', it is not a proposition nor a definition (to unfold)"
      discard <| go.run {} {}
    erase := fun declName => do
      let s ← ext.getState (← getEnv)
      let s ← s.erase declName
      modifyEnv fun env => ext.modifyState env fun _ => s
  }

def mkSimpExt (extName : Name) : IO SimpExtension :=
  registerSimpleScopedEnvExtension {
    name     := extName
    initial  := {}
    addEntry := fun d e =>
      match e with
      | SimpEntry.lemma e => addSimpLemmaEntry d e
      | SimpEntry.toUnfold n => d.addDeclToUnfold n
  }

def registerSimpAttr (attrName : Name) (attrDescr : String) (extName : Name := attrName.appendAfter "Ext") : IO SimpExtension := do
  let ext ← mkSimpExt extName
  mkSimpAttr attrName attrDescr ext
  return ext

builtin_initialize simpExtension : SimpExtension ← registerSimpAttr `simp "simplification theorem"

def getSimpLemmas : CoreM SimpLemmas :=
  simpExtension.getLemmas

/- Auxiliary method for adding a global declaration to a `SimpLemmas` datastructure. -/
def SimpLemmas.addConst (s : SimpLemmas) (declName : Name) (post : Bool := true) (inv : Bool := false) (prio : Nat := eval_prio default) : MetaM SimpLemmas := do
  let simpLemmas ← mkSimpLemmasFromConst declName post inv prio
  return simpLemmas.foldl addSimpLemmaEntry s

def SimpLemma.getValue (simpLemma : SimpLemma) : MetaM Expr := do
  if simpLemma.proof.isConst && simpLemma.levelParams.isEmpty then
    let info ← getConstInfo simpLemma.proof.constName!
    if info.levelParams.isEmpty then
      return simpLemma.proof
    else
      return simpLemma.proof.updateConst! (← info.levelParams.mapM (fun _ => mkFreshLevelMVar))
  else
    let us ← simpLemma.levelParams.mapM fun _ => mkFreshLevelMVar
    simpLemma.proof.instantiateLevelParamsArray simpLemma.levelParams us

private def preprocessProof (val : Expr) (inv : Bool) : MetaM (Array Expr) := do
  let type ← inferType val
  checkTypeIsProp type
  let ps ← preprocess val type inv
  return ps.toArray.map fun (val, _) => val

/- Auxiliary method for creating simp lemmas from a proof term `val`. -/
def mkSimpLemmas (levelParams : Array Name) (proof : Expr) (post : Bool := true) (inv : Bool := false) (prio : Nat := eval_prio default) (name? : Option Name := none): MetaM (Array SimpLemma) :=
  withReducible do
    (← preprocessProof proof inv).mapM fun val => mkSimpLemmaCore val levelParams val post prio name?

/- Auxiliary method for adding a local simp lemma to a `SimpLemmas` datastructure. -/
def SimpLemmas.add (s : SimpLemmas) (levelParams : Array Name) (proof : Expr) (inv : Bool := false) (post : Bool := true) (prio : Nat := eval_prio default) (name? : Option Name := none): MetaM SimpLemmas := do
  if proof.isConst then
    s.addConst proof.constName! post inv prio
  else
    let simpLemmas ← mkSimpLemmas levelParams proof post inv prio (← getName? proof)
    return simpLemmas.foldl addSimpLemmaEntry s
where
  getName? (e : Expr) : MetaM (Option Name) := do
    match name? with
    | some _ => return name?
    | none   =>
      let f := e.getAppFn
      if f.isConst then
        return f.constName!
      else if f.isFVar then
        let localDecl ← getFVarLocalDecl f
        return localDecl.userName
      else
        return none

end Lean.Meta


import .run hp.core

namespace hp.writeup

open hp

/-- Class predicate (CP). A class predicate is a predicate that can be used to declare a term.
So for example, you can say "Let `f : A → B` be a function" or "Let H be a normal subgroup of G" or "Let x, y and z be primes".
But you _can't_ write "Let `f` be injective on `A`", you have to write "Let `f` be a function which is injective on `A`"
  the parameter `A` means it doesn't work as an adjective: "Let `f` be an injective function on `A`".
There are currently 4 different kinds of CP:
- Adjective a predicate that should be treated as an adjective modifier.
  Examples are: "__open__ set", "__continuous__ function", "__normal__ subgroup".
- Fold Adjective is an adjective that applies to a collection of subjects.
  That is, it is a predicate of type `list subject → Prop`.
  Example: "__parallel__ lines".
- SymbolicPostfix are math-mode symbols that can be used to introduce a variable. For example `ε > 0`, `x y ∈ A`.
- ClassNoun are nouns that the given subject _is_. Eg "f is a __function__", "H is a __subgroup of G__".
-/
@[derive_prisms, derive decidable_eq, derive has_to_tactic_format]
meta inductive cp
| Adjective (r : run) -- adjectives must be strictly one word.
| FoldAdjective (r : run)
| SymbolicPostfix (e : run)
| ClassNoun (s p : run) -- "group"/"groups"
-- | ClassRelational (s p : run) (rhs : expr) -- "is/are subset[s] of" "element of" "function of". `s` and `p` come with `rhs` baked in but you need `rhs` for merging.
-- [todo] definite article tricks; "the least X such that Y"
-- [todo] ClassRelation packing; eg
  -- "(σ : ℕ → ℕ) (_ : ∀ n, prime (σ n))" ↝ "let σ be a sequence of prime numbers"
  -- "(X : set ℕ) (_ : ∀ x ∈ X, prime x)" ↝ "let X be a set of prime numbers"

/-- Class predicate collection. A list of class predicates to be rendered together.-/
@[derive_setters, derive decidable_eq, derive has_to_tactic_format]
meta structure cpc :=
(subjects : list expr)
(predicates : list cp)

/-- Class predicate collection __list__. -/
@[reducible]
meta def cpcs := list cpc

/-- An apply tree is a record of how an `apply` operation was performed.
Note that in HP, apply operations also unpack ∃ and ∧ propositions and automatically try discharging with `assumption`.

- Match the result was matched directly with the goal, with the given metavariables assigned
- ExistsElim value a: the applyer had the signature `∃ (x : X), P`
- AndElim value a : applyer had the signature `X ∧ Y` or `Y ∧ X`
- ApplyAssigned arg a; applyer had signature `Π (x : X), P` and `x` was assigned a value.
- ApplyGoal g a; applyer had signature `Π (x : X), P` but `x` was not assigned and hence appears as a new metavariable.
-/
@[derive_prisms, derive decidable_eq, derive has_to_tactic_format]
meta inductive ApplyTree : Type
| Match (result : expr) (goal : stub) (setters : list (stub × expr))
| ExistsElim (value : source) : ApplyTree -> ApplyTree
| AndElim (value : source) : ApplyTree -> ApplyTree
| ApplyAssigned (argument : expr) : ApplyTree -> ApplyTree
| ApplyGoal (name : name) (bi : binder_info) (stub : expr) : ApplyTree → ApplyTree

namespace ApplyTree

meta def getChild : ApplyTree → option ApplyTree
| (Match result goal setters) := none
| (ExistsElim value a) := some a
| (AndElim value a) := some a
| (ApplyAssigned argument a) := some a
| (ApplyGoal name bi stub a) := some a

meta def getAssigned : ApplyTree → list expr
| a := (id <| (some list.cons <*> (prod.fst <$> as_ApplyAssigned a))) $ ([] <| (getAssigned <$> getChild a))

meta def getGoals : ApplyTree → list expr
| a := (id <| (some list.cons <*> (as_ApplyGoal a >>= λ ⟨_,_,s,_⟩, some s))) $ ([] <| (getGoals <$> getChild a))

meta def getSources : ApplyTree → list source
| a := (id <| (some list.cons <*> ((as_ExistsElim a <|> as_AndElim a) >>= some ∘ prod.fst))) $ ([] <| (getSources <$> getChild a))

end ApplyTree

/-- A Statement corresponds to an intra-sentence proposition.
The proposition may be object-level or meta/context-level.

A Reason is some additional information that is used to justify statements.
-/
@[derive_prisms]
meta mutual inductive Reason, Statement

with Reason : Type
| BySetting (setters : list (stub × expr)) : Reason
-- "since A is open, "
| Since (smt : Statement) : Reason
| ofRun : run → Reason
-- "by applying norm_num, ", "by applying `ring`"...
| Tactic : string → Reason
-- don't include an explanation
| Omit

with Statement : Type
-- escape hatch for just writing out as string
| ofRun : run → Statement
| TermStatement (type : expr):  Statement
| CPC (cpc : cpc) : Statement
| Forall (binders : list cpc) (result : Statement) : Statement
| Exists (binders : list cpc) (result : Statement) : Statement
| Implies (premiss concl : Statement) : Statement
/-- Same as `Implies` but with ordering flipped. -/
| Whenever (premiss concl : Statement) : Statement
| And : list Statement → Statement
| By (s : Statement) (r : Reason) : Statement
/-- "`s` for some `...bs` where `where[bs]`". This is different to Exists because the 'binders' are already present in the context.
It's a meta-level statement about the tactic context rather than a standalone statement.
The 'where' are just other facts that might be relevant. -/
| ForSome (s : Statement) (binders : list cpc) (where : Statement) : Statement
-- it suffices to find `decls` such that `s`.
| Suffices (decls : list cpc) (s : Statement) : Statement
| Have (s : Statement)
| Provided (s1 s2 : Statement) : Statement
| Either (ss : list Statement) : Statement
| WeAreDone : Statement

section pp

open tactic Reason Statement format
meta mutual def Reason.pp, Statement.pp

with Reason.pp : Reason → tactic format
| (BySetting setters) := nest_join "BySetting" $ [pp setters]
| (Since smt) := nest_join "Since" $ [Statement.pp smt]
| (ofRun r) := nest_join "ofRun" $ [pp r]
| (Tactic t) := format.nest_join "Tactic" $ [pp t]
| (Omit) := format.nest_join "Omit" []

with Statement.pp : Statement → tactic format
| (ofRun r) := nest_join "ofRun" $ [pp r]
| (TermStatement t) := nest_join "TermStatement" [pp t]
| (CPC c) := nest_join "CPC" [pp c]
| (Either ss) := nest_join "Either" $ list.map Statement.pp ss
| (And ss) := nest_join "And" $ list.map Statement.pp ss
| (Provided a b) := nest_join "Provided" $ list.map Statement.pp [a, b]
| (Exists cpc b) := nest_join "Exists" $ [pp cpc, Statement.pp b]
| (Forall cpc b) := nest_join "Forall" $ [pp cpc, Statement.pp b]
| (Implies a b) := nest_join "Implies" $ list.map Statement.pp [a, b]
| (Whenever a b) := nest_join "Whenever" $ list.map Statement.pp [a, b]
| (By s r) := nest_join "By" $ [Statement.pp s, Reason.pp r]
| (ForSome s bs where) := nest_join "ForSome" $ [Statement.pp s, pp bs, Statement.pp where]
| (Suffices bs where) := nest_join "Suffices" $ [pp bs, Statement.pp where]
| (Have s) := nest_join "Have" $ [Statement.pp s]
| (WeAreDone) := nest_join "WeAreDone" $ []

meta instance Statement.has_pp : has_to_tactic_format Statement := ⟨Statement.pp⟩
meta instance Reason.has_pp : has_to_tactic_format Reason := ⟨Reason.pp⟩

end pp

meta def Reason.is_Omit : Reason → bool
| (Reason.Omit) := tt
| _ := ff

@[derive_prisms, derive has_to_tactic_format]
meta inductive Sentence -- [todo] merge with Statement?
/-- Introduce some variables and assumptions. -/
| Let (decls : list cpc) (where : Statement) : Sentence
/-- Restate the target. -/
| WeNeedToShow : Statement → Sentence
/-- Just print the statement. -/
| BareAssert : Statement → Sentence
| ReasonedAssert: Reason → Statement → Sentence
| WeAreDone
| Suffices (s : Statement) (r : Reason)
| Therefore : Sentence → Sentence
| WeMustChoose (decls : list cpc) : Statement → Sentence
/-- "In the case that ${s}: " -/
| InCase (s : Statement) : Sentence
| LineBreak : Sentence

-- | Since (s : Statement) (r : Reason) (result : Statement) : Sentence

/-- A high-level proof 'move'. The act type for the writeup procedure.
Given a stream of acts, one can produce a writeup.
[todo] rename to `act`. -/
@[derive_prisms, derive decidable_eq, derive has_to_tactic_format]
meta inductive act
| ProofDone
| Intro (h : list hyp)
| Existsi (v : stub) (prop : stub)
| Andi (l : stub) (r : stub)
| Apply (target : stub) (src : source) (results : ApplyTree)
| ExpandTarget
| TargetTactic (target : stub) (tactic_label : string) (results : ApplyTree)
/-- You can add a scope to an act to indicate that the microplanner should place it in its own paragraph.-/
| Scope : binder → act → act
| Cases : list binder → act

open tactic

section

variables {m : Type → Type} [monad m] (f : telescope → expr → m expr)

meta def ApplyTree.mmap_exprs : telescope → ApplyTree → m ApplyTree
| Γ (ApplyTree.Match result goal setters) := pure ApplyTree.Match <*> assignable.mmap_children f Γ result <*> assignable.mmap_children f Γ goal <*> assignable.mmap_children f Γ setters
| Γ (ApplyTree.ExistsElim value a) := pure ApplyTree.ExistsElim <*> assignable.mmap_children f Γ value <*> ApplyTree.mmap_exprs Γ a
| Γ (ApplyTree.AndElim value a) := pure ApplyTree.AndElim <*> assignable.mmap_children f Γ value <*> ApplyTree.mmap_exprs Γ a
| Γ (ApplyTree.ApplyAssigned argument a) := pure ApplyTree.ApplyAssigned <*> assignable.mmap_children f Γ argument <*> ApplyTree.mmap_exprs Γ a
| Γ (ApplyTree.ApplyGoal name bi stub a) := pure ApplyTree.ApplyGoal <*> pure name <*> pure bi <*> assignable.mmap_children f Γ  stub <*> ApplyTree.mmap_exprs Γ a

meta def act.mmap_exprs : telescope → act → m act
| Γ act.ProofDone := pure act.ProofDone
| Γ (act.Intro a) := pure act.Intro <*> assignable.mmap_children f Γ a
| Γ (act.Existsi v prop) := pure act.Existsi <*> assignable.mmap_children f Γ v <*> assignable.mmap_children f Γ prop
| Γ (act.Andi a b) := pure act.Andi <*> assignable.mmap_children f Γ a <*> assignable.mmap_children f Γ b
| Γ (act.Apply target src results) := pure act.Apply <*> assignable.mmap_children f Γ target <*> assignable.mmap_children f Γ src <*> ApplyTree.mmap_exprs f Γ results
| Γ act.ExpandTarget := pure act.ExpandTarget
| Γ (act.TargetTactic target tactic_label results) := pure act.TargetTactic <*> assignable.mmap_children f Γ target <*> pure tactic_label <*> ApplyTree.mmap_exprs f Γ results
| Γ (act.Scope n i) := pure ( act.Scope n ) <*> act.mmap_exprs Γ i
| Γ (act.Cases s) := pure (act.Cases) <*> assignable.mmap_children f Γ s

end

meta instance act.assignable : assignable act := ⟨@act.mmap_exprs⟩
meta instance ApplyTree.assignable : assignable ApplyTree := ⟨@ApplyTree.mmap_exprs⟩

meta def cp.pp : cp → tactic format
| (cp.Adjective r)       := (format.compose "Adjective ")       <$> tactic.pp r
| (cp.FoldAdjective r)   := (format.compose "FoldAdjective ")   <$> tactic.pp r
| (cp.SymbolicPostfix e) := (format.compose "SymbolicPostfix ") <$> tactic.pp e
| (cp.ClassNoun s p)     := pure (λ s p, "ClassNoun " ++ s ++ " " ++ p) <*> tactic.pp s <*> tactic.pp p
-- | (cp.ClassRelational s p rhs) :=  pure (λ s p r, "ClassRelational " ++ s ++ " " ++ p ++ " " ++ r) <*> tactic.pp s <*> tactic.pp p <*> tactic.pp rhs

meta def cp.to_string : cp → string
| (cp.Adjective r)       := ("Adjective ")       ++ to_string r
| (cp.FoldAdjective r)   := ("FoldAdjective ")   ++ to_string r
| (cp.SymbolicPostfix e) := ("SymbolicPostfix ") ++ to_string e
| (cp.ClassNoun s p)     := "ClassNoun " ++ to_string s ++ " " ++ to_string p
-- | (cp.ClassRelational s p rhs) :=  "ClassRelational " ++ to_string s ++ " " ++ to_string p ++ " " ++ to_string rhs

meta instance cp.has_pp : has_to_tactic_format cp := ⟨cp.pp⟩
meta instance cp.has_to_string : has_to_string cp := ⟨cp.to_string⟩

meta def cpc.to_string : cpc → string
| ⟨ss, ps⟩ := "⟨" ++ to_string ss ++ ", " ++ to_string ps ++ "⟩"

meta def cpc.pp : cpc → tactic format
| ⟨ss, ps⟩ := do
  ss ← tactic.pp ss,
  ps ← tactic.pp ps,
  pure $ "⟨" ++ ss ++ ", " ++ ps ++ "⟩"

meta instance cpc.has_pp : has_to_tactic_format cpc := ⟨cpc.pp⟩

meta instance cpc.has_to_string : has_to_string cpc := ⟨cpc.to_string⟩

meta def contains_subject : expr → cpc → bool
| e ⟨xs, _⟩ := e ∈ xs

section
variables {m : Type → Type} [monad m]
meta def cp.mmap_exprs  (f : telescope → expr → m expr) : telescope → cp → m cp
| Γ (cp.Adjective r) := (pure cp.Adjective) <*> ((Γ ⍄ f) r)
| Γ (cp.FoldAdjective r) := (pure cp.FoldAdjective) <*> ((Γ ⍄ f) r)
| Γ (cp.SymbolicPostfix r) := (pure cp.SymbolicPostfix) <*> ((Γ ⍄ f) r)
| Γ (cp.ClassNoun s p) := pure (cp.ClassNoun s p)

meta instance cp.has_assignable : assignable cp := ⟨@cp.mmap_exprs⟩

meta def cpc.mmap_exprs (f : telescope → expr → m expr) : telescope → cpc → m cpc
| Γ ⟨subjects, cps⟩ := (pure cpc.mk) <*> (list.mmap (f Γ) subjects) <*> (list.mmap (Γ ⍄ f) cps)

meta instance cpc.has_assignable : assignable cpc := ⟨@cpc.mmap_exprs⟩

-- meta def cpc.cpcs_assignable : assignable cpcs := by apply_instance

end
end hp.writeup


import .expr_zipper
open tactic
universe u
/-- ltb is the same as lt but it goes to a bool rather than a prop to make deriving decidability easier. -/
class has_ltb (α : Type u) := (ltb : α → α → bool)
instance lt_of_ltb (α : Type u) [has_ltb α] : has_lt α := ⟨λ x y, has_ltb.ltb x y⟩
instance dec_lt_of_ltb (α : Type u) [has_ltb α] : decidable_rel ((<) : α → α → Prop) := by apply_instance
instance le_of_ltb (α : Type u) [has_ltb α] : has_le α := ⟨λ x y, x < y ∨ x = y⟩
instance dec_le_of_ltb (α : Type u) [has_ltb α] [decidable_eq α] : decidable_rel ((≤) : α → α → Prop)  := by apply_instance

namespace derivations
/-- An 'inductive argument' is either just a normal argument or a recursive argument.
So for example `@list.cons {α}` is a constructor with two arguments,
one being a normal with type α and the second being recursive.
-/
meta inductive ind_arg
| normal : expr → ind_arg
| recursive (arg : expr) (rr : expr) : ind_arg

meta def map_ind_arg (f : pexpr → pexpr): ind_arg → pexpr
| (ind_arg.normal e) := f $ to_pexpr e
| (ind_arg.recursive _ rr) := to_pexpr rr

meta def ind_arg.pp : ind_arg → tactic format
| (ind_arg.normal a) := pp a >>= pure ∘ (++ "normal ")
| (ind_arg.recursive a b) := pure (λ a b, "rec " ++ a ++ " " ++ b) <*> pp a <*> pp b

meta instance : has_to_tactic_format ind_arg := ⟨ind_arg.pp⟩

/-- Induction but the induction hypotheses are bundled with their respective arguments.
Contrast this with the result of `induction` where all of the arguments are given, followed by the induction hyps
and you don't know immediately whether you are looking at an ind hyp, a normal arg or a recursive arg.
 -/
meta def induction_but_the_recursors_are_kept_with_the_arguments (x : expr)
    : tactic $ list (name × list ind_arg × list (name × expr) ) := do
    xT ← infer_type x,
    cs ← induction x,
    gs ← get_goals,
    es ← (list.zip cs gs).mmap (λ c, do
        ⟨⟨n,ls, ss⟩, g⟩ ← pure c,
        set_goals [g],
        e ← resolve_name n >>= pure ∘ pexpr.mk_explicit >>= to_expr,
        T ← infer_type e,
        (ctx,b) ← pure $ telescope.of_pis T,
        args ← pure $ ls.take $ ctx.length,
        indos ← pure $ ls.drop $ ctx.length,
        ⟨[], acc⟩ ← args.mfoldl (λ p a, do
            ⟨indos, acc⟩ ← pure (p : (list expr) × (list ind_arg)),
            T ← infer_type a,
            if expr.get_app_fn xT = expr.get_app_fn T then do
                h::indos ← pure indos,
                pure (indos, ind_arg.recursive a h :: acc) else
            pure (indos, ind_arg.normal a :: acc)
        ) (indos, []),
        acc ← pure $ acc.reverse,
        pure (n, acc, ss)
    ),
    set_goals gs,
    pure es

/-- Automatically derive `has_to_string` for an inductive datatype. -/
@[derive_handler] meta def to_string_handler :=
instance_derive_handler ``has_to_string $
    do
        e ← get_env,
        split,
        x ← intro `x,
        cs ← induction_but_the_recursors_are_kept_with_the_arguments x,
        cs.mmap (λ c, do
            ⟨n,ls,ss⟩ ← pure c,
            constructor_name ← pure $ reflect $ n.components.ilast,
            if ls.empty then refine $ ```(to_string %%constructor_name) else do
            str ← pure $ list.foldl (λ x y, ```(%%x ++ %%y)) ```("") $ list.intersperse ```( " " ) $ list.map (λ x, ```( "(" ++ %%x ++ ")" )) $ ls.map (map_ind_arg (λ x, ```( to_string %%x))),
            refine $ ```(to_string %%constructor_name ++ " " ++ %%str)
        ),
        pure ()

/-- Helper method for has_to_tactic_format_handler -/
private meta def pp_cases : ind_arg → tactic unit
|(ind_arg.normal x) :=
    refine ```(tactic.pp %%x)
    <|> refine ```(pure $ to_fmt %%x)
    <|> refine ```(pure $ format.of_string $ to_string %%x)
    <|> (do
        ppx ← tactic.pp x,
        T ← tactic.infer_type x,
        ppT ← tactic.pp T,
        msg : format ← pure $ to_fmt "Couldn't find a way of showing " ++ ppx ++ " : " ++ ppT,
        tactic.fail $ msg)
|(ind_arg.recursive x r) := refine ```(%%r)

@[derive_handler] meta def has_to_tactic_format_handler :=
instance_derive_handler ``has_to_tactic_format $ do
    split,
    x ← intro `x,
    cs ← induction_but_the_recursors_are_kept_with_the_arguments x,
    cs.mmap (λ x, do
        ⟨n,ls,ss⟩ ← pure x,
        constructor_name ← pure $ reflect $ n.components.ilast,
        cnp ← pure $ ```(tactic.pp %%constructor_name), -- : tactic format
        if ls.empty then refine cnp else do
        pps ← ls.mmap (λ x, do
            g ← mk_meta_var `(tactic format),
            gs ← get_goals,
            set_goals [g],
            pp_cases x,
            set_goals gs,
            instantiate_mvars g),
        ppl ← pure $ pps.foldr (λ x acc, ```(%%x :: %%acc)) ```([]),
        ppl ← pure $ ```(list.mmap id %%ppl), -- : tactic (list format)
        refine ```(pure (λ cnp ppl, cnp ++ (to_fmt "{")++ (format.group $ format.nest 1 $ format.join $ list.intersperse ("," ++ format.line) ppl) ++ "}") <*> %%cnp <*> %%ppl)
    ),
    pure ()

private meta def lt_cases : (ind_arg × ind_arg) → pexpr → pexpr
| (ind_arg.normal a_h, ind_arg.normal b_h) acc :=
    ```((%%a_h < %%b_h) ∨ ((%%a_h = %%b_h) ∧ %%acc))
| (ind_arg.recursive a_h ai, ind_arg.recursive b_h bi) acc :=
    ```((%%ai %%b_h) ∨ ((%%a_h = %%b_h) ∧ %%acc))
| _ acc := acc

/-- Derive a total ordering for a given inductive datatype. This just puts an ordering on the constructor names and then
compares two terms with the same ctor lexically on their arguments. Useful for enabling the datatype to be
stored in rtrees and other structures which require some total ordering on data.
You should only use this if you don't care about the semantics of `<`.
-/
@[derive_handler] meta def lt_derive_handler :=
instance_derive_handler ``has_lt $ do
    split,
    a ← intro `a,
    a_cs ← induction_but_the_recursors_are_kept_with_the_arguments a,
    a_goals ← get_goals,
    (list.zip a_cs a_goals).mmap (λ c, do
        ⟨⟨a_cn,a_hs,a_subs⟩,g⟩ ← pure c,
        b ← intro `b,
        b_cs ← induction_but_the_recursors_are_kept_with_the_arguments b,
        b_goals ← get_goals,
        bz ← pure $ list.zip b_cs b_goals,
        bz.mmap(λ c, do
            ⟨⟨b_cn,b_hs,b_subs⟩,g⟩ ← pure c,
            if a_cn < b_cn then exact `(true) else
            if a_cn > b_cn then exact `(false) else do
            r ← pure $ list.foldr lt_cases ```(false) $ list.zip a_hs b_hs,
            r ← to_expr r,
            exact r
        )
    ),
    set_goals a_goals,
    pure ()

private meta def ltb_cases : (ind_arg × ind_arg) → pexpr → pexpr
| (ind_arg.normal a_h, ind_arg.normal b_h) acc :=
    ```(bor (%%a_h < %%b_h) (band (to_bool $ %%a_h = %%b_h) %%acc))
| (ind_arg.recursive a_h ai, ind_arg.recursive b_h bi) acc :=
    ```(bor (%%ai %%b_h) (band (to_bool $ %%a_h = %%b_h) %%acc))
| _ acc := acc

@[derive_handler] meta def mk_ltb_instance := instance_derive_handler ``has_ltb $ do
    split,
    a ← intro `a,
    (expr.const typename _) ← infer_type a >>= pure ∘ expr.get_app_fn,
    e ← get_env,
    all_cases ← pure $ environment.constructors_of e typename,
    a_cs ← induction_but_the_recursors_are_kept_with_the_arguments a,
    a_goals ← get_goals,
    (list.zip a_cs a_goals).mmap (λ c, do
        ⟨⟨a_cn,a_hs,a_subs⟩,g⟩ ← pure c,
        a_idx ← pure $ list.find_index (= a_cn) all_cases,
        b ← intro `b,
        b_cs ← induction_but_the_recursors_are_kept_with_the_arguments b,
        b_goals ← get_goals,
        bz ← pure $ list.zip b_cs b_goals,
        bz.mmap(λ c, do
            ⟨⟨b_cn,b_hs,b_subs⟩,g⟩ ← pure c,
            b_idx ← pure $ list.find_index (= b_cn) all_cases,
            if a_idx < b_idx then exact `(tt) else
            if a_idx > b_idx then exact `(ff) else do
            r ← pure $ list.foldr ltb_cases ```(ff) $ list.zip a_hs b_hs,
            r ← to_expr r,
            exact r
        )
    ),
    set_goals a_goals,
    pure ()

/-!
# An example:

```
@[derive decidable_eq, derive has_ltb, derive has_to_tactic_format]
inductive qwerty
| nil : qwerty
| asdf (n : nat) : qwerty
| br (x: qwerty) (n : nat) (y : qwerty) : qwerty
| s : qwerty → qwerty

open qwerty
#eval to_bool $  nil < (asdf 3)
#eval to_bool $ s (asdf 3) < s (asdf 4)
#eval to_bool (br (asdf 3) 2 nil < br (asdf 3) 2 nil)
#eval to_bool $ ((asdf 3) < (asdf 4))
run_cmd (trace $ br (asdf 3) 2 nil)
```
-/
end derivations


/-
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

Normalizing casts inside expressions.
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.tactic.converter.interactive
import Mathlib.tactic.hint
import Mathlib.PostPort

universes l u_1 u_2

namespace Mathlib

/-!
# 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`
-/

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.
-/
end tactic

namespace norm_cast

/--
Output a trace message if `trace.norm_cast` is enabled.
-/
/--
`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
-/
inductive label
where
| elim : label
| move : label
| squash : label

namespace label

/-- Convert `label` into `string`. -/
protected def to_string : label → string :=
  sorry

protected instance has_to_string : has_to_string label :=
  has_to_string.mk label.to_string

protected instance has_repr : has_repr label :=
  has_repr.mk label.to_string

/-- Convert `string` into `label`. -/
def of_string : string → Option label :=
  sorry

end label

/-- Count how many coercions are at the top of the expression. -/
/-- Count how many coercions are inside the expression, including the top ones. -/
/-- Count how many coercions are inside the expression, excluding the top ones. -/
/--
Classifies a declaration of type `ty` as a `norm_cast` rule.
-/
/-- The cache for `norm_cast` attribute stores three `simp_lemma` objects. -/
/-- Empty `norm_cast_cache`. -/
/-- `add_elim cache e` adds `e` as an `elim` lemma to `cache`. -/
/-- `add_move cache e` adds `e` as a `move` lemma to `cache`. -/
/-- `add_squash cache e` adds `e` as an `squash` lemma to `cache`. -/
/--
The type of the `norm_cast` attribute.
The optional label is used to overwrite the classifier.
-/
/--
Efficient getter for the `@[norm_cast]` attribute parameter that does not call `eval_expr`.

See Note [user attribute parameters].
-/
/--
`add_lemma cache decl` infers the proper `norm_cast` attribute for `decl` and adds it to `cache`.
-/
-- special lemmas to handle the ≥, > and ≠ operators

/--
`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.
-/
-- names has the declarations in reverse order

--some special lemmas to handle binary relations

/--
The `norm_cast` attribute.
-/
/-- Classify a declaration as a `norm_cast` rule. -/
/--
Gets the `norm_cast` classification label for a declaration. Applies the
override specified on the attribute, if necessary.
-/
end norm_cast

namespace tactic.interactive

/--
`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
```
-/
end tactic.interactive

namespace norm_cast

/-- Prove `a = b` using the given simp set. -/
/-- Prove `a = b` by simplifying using move and squash lemmas. -/
/--
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
-/
/--
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
-/
/--
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.
-/
/-!
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.
-/
/--
If possible, rewrite `((n : ℕ) : α)` to `(n : α)` where `n` is a numeral.
Returns a pair of the new expression and proof that they are equal.
-/
/-- A local variant on `simplify_top_down`. -/
/--
The core simplification routine of `norm_cast`.
-/
/--
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'`.
-/
end norm_cast

namespace tactic

/-- `aux_mod_cast e` runs `norm_cast` on `e` and returns the result. If `include_goal` is true, it
also normalizes the goal. -/
/-- `exact_mod_cast e` runs `norm_cast` on the goal and `e`, and tries to use `e` to close the goal. -/
/-- `apply_mod_cast e` runs `norm_cast` on the goal and `e`, and tries to apply `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. -/
end tactic

namespace tactic.interactive

/--
Normalize casts at the given locations by moving them "upwards".
As opposed to simp, norm_cast can be used without necessarily closing the goal.
-/
/--
Rewrite with the given rules and normalize casts between steps.
-/
/--
Normalize the goal and the given expression, then close the goal with exact.
-/
/--
Normalize the goal and the given expression, then apply the expression to the goal.
-/
/--
Normalize the goal and every expression in the local context, then close the goal with assumption.
-/
end tactic.interactive

namespace conv.interactive

/-- the converter version of `norm_cast' -/
end conv.interactive

-- TODO: move this elsewhere?

theorem ite_cast {α : Sort u_1} {β : Sort u_2} [has_lift_t α β] {c : Prop} [Decidable c] {a : α} {b : α} : ↑(ite c a b) = ite c ↑a ↑b := sorry

/--
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>.
-/
/--
The `norm_cast` attribute should be given to lemmas that describe the


/-
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 Lean.Parser.Term
import Lean.Meta.Closure
import Lean.Meta.Check
import Lean.Elab.Command
import Lean.Elab.DefView
import Lean.Elab.PreDefinition
import Lean.Elab.DeclarationRange

namespace Lean.Elab
open Lean.Parser.Term

/- DefView after elaborating the header. -/
structure DefViewElabHeader where
  ref           : Syntax
  modifiers     : Modifiers
  kind          : DefKind
  shortDeclName : Name
  declName      : Name
  levelNames    : List Name
  numParams     : Nat
  type          : Expr -- including the parameters
  valueStx      : Syntax
  deriving Inhabited

namespace Term

open Meta

private def checkModifiers (m₁ m₂ : Modifiers) : TermElabM Unit := do
  unless m₁.isUnsafe == m₂.isUnsafe do
    throwError "cannot mix unsafe and safe definitions"
  unless m₁.isNoncomputable == m₂.isNoncomputable do
    throwError "cannot mix computable and non-computable definitions"
  unless m₁.isPartial == m₂.isPartial do
    throwError "cannot mix partial and non-partial definitions"

private def checkKinds (k₁ k₂ : DefKind) : TermElabM Unit := do
  unless k₁.isExample == k₂.isExample do
    throwError "cannot mix examples and definitions" -- Reason: we should discard examples
  unless k₁.isTheorem == k₂.isTheorem do
    throwError "cannot mix theorems and definitions" -- Reason: we will eventually elaborate theorems in `Task`s.

private def check (prevHeaders : Array DefViewElabHeader) (newHeader : DefViewElabHeader) : TermElabM Unit := do
  if newHeader.kind.isTheorem && newHeader.modifiers.isUnsafe then
    throwError "'unsafe' theorems are not allowed"
  if newHeader.kind.isTheorem && newHeader.modifiers.isPartial then
    throwError "'partial' theorems are not allowed, 'partial' is a code generation directive"
  if newHeader.kind.isTheorem && newHeader.modifiers.isNoncomputable then
    throwError "'theorem' subsumes 'noncomputable', code is not generated for theorems"
  if newHeader.modifiers.isNoncomputable && newHeader.modifiers.isUnsafe then
    throwError "'noncomputable unsafe' is not allowed"
  if newHeader.modifiers.isNoncomputable && newHeader.modifiers.isPartial then
    throwError "'noncomputable partial' is not allowed"
  if newHeader.modifiers.isPartial && newHeader.modifiers.isUnsafe then
    throwError "'unsafe' subsumes 'partial'"
  if h : 0 < prevHeaders.size then
    let firstHeader := prevHeaders.get ⟨0, h⟩
    try
      unless newHeader.levelNames == firstHeader.levelNames do
        throwError "universe parameters mismatch"
      checkModifiers newHeader.modifiers firstHeader.modifiers
      checkKinds newHeader.kind firstHeader.kind
    catch
       | Exception.error ref msg => throw (Exception.error ref m!"invalid mutually recursive definitions, {msg}")
       | ex => throw ex
  else
    pure ()

private def registerFailedToInferDefTypeInfo (type : Expr) (ref : Syntax) : TermElabM Unit :=
  registerCustomErrorIfMVar type ref "failed to infer definition type"

private def elabHeaders (views : Array DefView) : TermElabM (Array DefViewElabHeader) := do
  let mut headers := #[]
  for view in views do
    let newHeader ← withRef view.ref do
      let ⟨shortDeclName, declName, levelNames⟩ ← expandDeclId (← getCurrNamespace) (← getLevelNames) view.declId view.modifiers
      addDeclarationRanges declName view.ref
      applyAttributesAt declName view.modifiers.attrs AttributeApplicationTime.beforeElaboration
      withDeclName declName <| withAutoBoundImplicit <| withLevelNames levelNames <|
        elabBinders view.binders.getArgs fun xs => do
          let refForElabFunType := view.value
          let type ← match view.type? with
            | some typeStx =>
              let type ← elabType typeStx
              registerFailedToInferDefTypeInfo type typeStx
              pure type
            | none =>
              let hole := mkHole refForElabFunType
              let type ← elabType hole
              registerFailedToInferDefTypeInfo type refForElabFunType
              pure type
          Term.synthesizeSyntheticMVarsNoPostponing
          let type ← mkForallFVars xs type
          let type ← mkForallFVars (← read).autoBoundImplicits.toArray type
          let type ← instantiateMVars type
          let xs ← addAutoBoundImplicits xs
          let levelNames ← getLevelNames
          if view.type?.isSome then
            let pendingMVarIds ← getMVars type
            discard <| logUnassignedUsingErrorInfos pendingMVarIds <|
              m!"\nwhen the resulting type of a declaration is explicitly provided, all holes (e.g., `_`) in the header are resolved before the declaration body is processed"
          let newHeader := {
            ref           := view.ref,
            modifiers     := view.modifiers,
            kind          := view.kind,
            shortDeclName := shortDeclName,
            declName      := declName,
            levelNames    := levelNames,
            numParams     := xs.size,
            type          := type,
            valueStx      := view.value : DefViewElabHeader }
          check headers newHeader
          pure newHeader
    headers := headers.push newHeader
  pure headers

private partial def withFunLocalDecls {α} (headers : Array DefViewElabHeader) (k : Array Expr → TermElabM α) : TermElabM α :=
  let rec loop (i : Nat) (fvars : Array Expr) := do
    if h : i < headers.size then
      let header := headers.get ⟨i, h⟩
      withLocalDecl header.shortDeclName BinderInfo.auxDecl header.type fun fvar => loop (i+1) (fvars.push fvar)
    else
      k fvars
  loop 0 #[]

private def expandWhereDeclsAsStructInst : Macro
  | `(whereDecls|where $[$decls:letRecDecl$[;]?]*) => do
    let letIdDecls ← decls.mapM fun stx => match stx with
      | `(letRecDecl|$attrs:attributes $decl:letDecl) => Macro.throwErrorAt stx "attributes are 'where' elements are currently not supported here"
      | `(letRecDecl|$decl:letPatDecl)  => Macro.throwErrorAt stx "patterns are not allowed here"
      | `(letRecDecl|$decl:letEqnsDecl) => expandLetEqnsDecl decl
      | `(letRecDecl|$decl:letIdDecl)   => pure decl
      | _                               => Macro.throwUnsupported
    let structInstFields ← letIdDecls.mapM fun
      | stx@`(letIdDecl|$id:ident $[$binders]* $[: $ty?]? := $val) => withRef stx do
        let mut val := val
        if let some ty := ty? then
          val ← `(($val : $ty))
        val ← if binders.size > 0 then `(fun $[$binders]* => $val:term) else val
        `(structInstField|$id:ident := $val)
      | _ => Macro.throwUnsupported
    `({ $[$structInstFields,]* })
  | _ => Macro.throwUnsupported

/-
Recall that
```
def declValSimple    := leading_parser " :=\n" >> termParser >> optional Term.whereDecls
def declValEqns      := leading_parser Term.matchAltsWhereDecls
def declVal          := declValSimple <|> declValEqns <|> Term.whereDecls
```
-/
private def declValToTerm (declVal : Syntax) : MacroM Syntax := withRef declVal do
  if declVal.isOfKind `Lean.Parser.Command.declValSimple then
    expandWhereDeclsOpt declVal[2] declVal[1]
  else if declVal.isOfKind `Lean.Parser.Command.declValEqns then
    expandMatchAltsWhereDecls declVal[0]
  else if declVal.isOfKind `Lean.Parser.Term.whereDecls then
    expandWhereDeclsAsStructInst declVal
  else if declVal.isMissing then
    Macro.throwErrorAt declVal "declaration body is missing"
  else
    Macro.throwErrorAt declVal "unexpected declaration body"

private def elabFunValues (headers : Array DefViewElabHeader) : TermElabM (Array Expr) :=
  headers.mapM fun header => withDeclName header.declName $ withLevelNames header.levelNames do
    let valStx ← liftMacroM $ declValToTerm header.valueStx
    forallBoundedTelescope header.type header.numParams fun xs type => do
      let val ← elabTermEnsuringType valStx type
      mkLambdaFVars xs val

private def collectUsed (headers : Array DefViewElabHeader) (values : Array Expr) (toLift : List LetRecToLift)
    : StateRefT CollectFVars.State MetaM Unit := do
  headers.forM fun header => collectUsedFVars header.type
  values.forM collectUsedFVars
  toLift.forM fun letRecToLift => do
    collectUsedFVars letRecToLift.type
    collectUsedFVars letRecToLift.val

private def removeUnusedVars (vars : Array Expr) (headers : Array DefViewElabHeader) (values : Array Expr) (toLift : List LetRecToLift)
    : TermElabM (LocalContext × LocalInstances × Array Expr) := do
  let (_, used) ← (collectUsed headers values toLift).run {}
  removeUnused vars used

private def withUsed {α} (vars : Array Expr) (headers : Array DefViewElabHeader) (values : Array Expr) (toLift : List LetRecToLift)
    (k : Array Expr → TermElabM α) : TermElabM α := do
  let (lctx, localInsts, vars) ← removeUnusedVars vars headers values toLift
  withLCtx lctx localInsts $ k vars

private def isExample (views : Array DefView) : Bool :=
  views.any (·.kind.isExample)

private def isTheorem (views : Array DefView) : Bool :=
  views.any (·.kind.isTheorem)

private def instantiateMVarsAtHeader (header : DefViewElabHeader) : TermElabM DefViewElabHeader := do
  let type ← instantiateMVars header.type
  pure { header with type := type }

private def instantiateMVarsAtLetRecToLift (toLift : LetRecToLift) : TermElabM LetRecToLift := do
  let type ← instantiateMVars toLift.type
  let val ← instantiateMVars toLift.val
  pure { toLift with type := type, val := val }

private def typeHasRecFun (type : Expr) (funFVars : Array Expr) (letRecsToLift : List LetRecToLift) : Option FVarId :=
  let occ? := type.find? fun e => match e with
    | Expr.fvar fvarId _ => funFVars.contains e || letRecsToLift.any fun toLift => toLift.fvarId == fvarId
    | _ => false
  match occ? with
  | some (Expr.fvar fvarId _) => some fvarId
  | _ => none

private def getFunName (fvarId : FVarId) (letRecsToLift : List LetRecToLift) : TermElabM Name := do
  match (← findLocalDecl? fvarId) with
  | some decl => pure decl.userName
  | none =>
    /- Recall that the FVarId of nested let-recs are not in the current local context. -/
    match letRecsToLift.findSome? fun toLift => if toLift.fvarId == fvarId then some toLift.shortDeclName else none with
    | none   => throwError "unknown function"
    | some n => pure n

/-
Ensures that the of let-rec definition types do not contain functions being defined.
In principle, this test can be improved. We could perform it after we separate the set of functions is strongly connected components.
However, this extra complication doesn't seem worth it.
-/
private def checkLetRecsToLiftTypes (funVars : Array Expr) (letRecsToLift : List LetRecToLift) : TermElabM Unit :=
  letRecsToLift.forM fun toLift =>
    match typeHasRecFun toLift.type funVars letRecsToLift with
    | none        => pure ()
    | some fvarId => do
      let fnName ← getFunName fvarId letRecsToLift
      throwErrorAt toLift.ref "invalid type in 'let rec', it uses '{fnName}' which is being defined simultaneously"

namespace MutualClosure

/- A mapping from FVarId to Set of FVarIds. -/
abbrev UsedFVarsMap := NameMap NameSet

/-
Create the `UsedFVarsMap` mapping that takes the variable id for the mutually recursive functions being defined to the set of
free variables in its definition.

For `mainFVars`, this is just the set of section variables `sectionVars` used.
For nested let-rec functions, we collect their free variables.

Recall that a `let rec` expressions are encoded as follows in the elaborator.
```lean
let rec
  f : A := t,
  g : B := s;
body
```
is encoded as
```lean
let f : A := ?m₁;
let g : B := ?m₂;
body
```
where `?m₁` and `?m₂` are synthetic opaque metavariables. That are assigned by this module.
We may have nested `let rec`s.
```lean
let rec f : A :=
    let rec g : B := t;
    s;
body
```
is encoded as
```lean
let f : A := ?m₁;
body
```
and the body of `f` is stored the field `val` of a `LetRecToLift`. For the example above,
we would have a `LetRecToLift` containing:
```
{
  mvarId := m₁,
  val    := `(let g : B := ?m₂; body)
  ...
}
```
Note that `g` is not a free variable at `(let g : B := ?m₂; body)`. We recover the fact that
`f` depends on `g` because it contains `m₂`
-/
private def mkInitialUsedFVarsMap (mctx : MetavarContext) (sectionVars : Array Expr) (mainFVarIds : Array FVarId) (letRecsToLift : List LetRecToLift)
    : UsedFVarsMap := do
  let mut sectionVarSet := {}
  for var in sectionVars do
    sectionVarSet := sectionVarSet.insert var.fvarId!
  let mut usedFVarMap := {}
  for mainFVarId in mainFVarIds do
    usedFVarMap := usedFVarMap.insert mainFVarId sectionVarSet
  for toLift in letRecsToLift do
    let state := Lean.collectFVars {} toLift.val
    let state := Lean.collectFVars state toLift.type
    let mut set := state.fvarSet
    /- toLift.val may contain metavariables that are placeholders for nested let-recs. We should collect the fvarId
       for the associated let-rec because we need this information to compute the fixpoint later. -/
    let mvarIds := (toLift.val.collectMVars {}).result
    for mvarId in mvarIds do
      match letRecsToLift.findSome? fun (toLift : LetRecToLift) => if toLift.mvarId == mctx.getDelayedRoot mvarId then some toLift.fvarId else none with
      | some fvarId => set := set.insert fvarId
      | none        => pure ()
    usedFVarMap := usedFVarMap.insert toLift.fvarId set
  pure usedFVarMap

/-
The let-recs may invoke each other. Example:
```
let rec
  f (x : Nat) := g x + y
  g : Nat → Nat
    | 0   => 1
    | x+1 => f x + z
```
`y` is free variable in `f`, and `z` is a free variable in `g`.
To close `f` and `g`, `y` and `z` must be in the closure of both.
That is, we need to generate the top-level definitions.
```
def f (y z x : Nat) := g y z x + y
def g (y z : Nat) : Nat → Nat
  | 0 => 1
  | x+1 => f y z x + z
```
-/
namespace FixPoint

structure State where
  usedFVarsMap : UsedFVarsMap := {}
  modified     : Bool         := false

abbrev M := ReaderT (List FVarId) $ StateM State

private def isModified : M Bool := do pure (← get).modified
private def resetModified : M Unit := modify fun s => { s with modified := false }
private def markModified : M Unit := modify fun s => { s with modified := true }
private def getUsedFVarsMap : M UsedFVarsMap := do pure (← get).usedFVarsMap
private def modifyUsedFVars (f : UsedFVarsMap → UsedFVarsMap) : M Unit := modify fun s => { s with usedFVarsMap := f s.usedFVarsMap }

-- merge s₂ into s₁
private def merge (s₁ s₂ : NameSet) : M NameSet :=
  s₂.foldM (init := s₁) fun s₁ k => do
    if s₁.contains k then
      pure s₁
    else
      markModified
      pure $ s₁.insert k

private def updateUsedVarsOf (fvarId : FVarId) : M Unit := do
  let usedFVarsMap ← getUsedFVarsMap
  match usedFVarsMap.find? fvarId with
  | none         => pure ()
  | some fvarIds =>
    let fvarIdsNew ← fvarIds.foldM (init := fvarIds) fun fvarIdsNew fvarId' =>
      if fvarId == fvarId' then
        pure fvarIdsNew
      else
        match usedFVarsMap.find? fvarId' with
        | none => pure fvarIdsNew
          /- We are being sloppy here `otherFVarIds` may contain free variables that are
             not in the context of the let-rec associated with fvarId.
             We filter these out-of-context free variables later. -/
        | some otherFVarIds => merge fvarIdsNew otherFVarIds
    modifyUsedFVars fun usedFVars => usedFVars.insert fvarId fvarIdsNew

private partial def fixpoint : Unit → M Unit
  | _ => do
    resetModified
    let letRecFVarIds ← read
    letRecFVarIds.forM updateUsedVarsOf
    if (← isModified) then
      fixpoint ()

def run (letRecFVarIds : List FVarId) (usedFVarsMap : UsedFVarsMap) : UsedFVarsMap :=
  let (_, s) := ((fixpoint ()).run letRecFVarIds).run { usedFVarsMap := usedFVarsMap }
  s.usedFVarsMap

end FixPoint

abbrev FreeVarMap := NameMap (Array FVarId)

private def mkFreeVarMap
    (mctx : MetavarContext) (sectionVars : Array Expr) (mainFVarIds : Array FVarId)
    (recFVarIds : Array FVarId) (letRecsToLift : List LetRecToLift) : FreeVarMap := do
  let usedFVarsMap  := mkInitialUsedFVarsMap mctx sectionVars mainFVarIds letRecsToLift
  let letRecFVarIds := letRecsToLift.map fun toLift => toLift.fvarId
  let usedFVarsMap  := FixPoint.run letRecFVarIds usedFVarsMap
  let mut freeVarMap := {}
  for toLift in letRecsToLift do
    let lctx       := toLift.lctx
    let fvarIdsSet := (usedFVarsMap.find? toLift.fvarId).get!
    let fvarIds    := fvarIdsSet.fold (init := #[]) fun fvarIds fvarId =>
      if lctx.contains fvarId && !recFVarIds.contains fvarId then
        fvarIds.push fvarId
      else
        fvarIds
    freeVarMap := freeVarMap.insert toLift.fvarId fvarIds
  pure freeVarMap

structure ClosureState where
  newLocalDecls : Array LocalDecl := #[]
  localDecls    : Array LocalDecl := #[]
  newLetDecls   : Array LocalDecl := #[]
  exprArgs      : Array Expr      := #[]

private def pickMaxFVar? (lctx : LocalContext) (fvarIds : Array FVarId) : Option FVarId :=
  fvarIds.getMax? fun fvarId₁ fvarId₂ => (lctx.get! fvarId₁).index < (lctx.get! fvarId₂).index

private def preprocess (e : Expr) : TermElabM Expr := do
  let e ← instantiateMVars e
  -- which let-decls are dependent. We say a let-decl is dependent if its lambda abstraction is type incorrect.
  Meta.check e
  pure e

/- Push free variables in `s` to `toProcess` if they are not already there. -/
private def pushNewVars (toProcess : Array FVarId) (s : CollectFVars.State) : Array FVarId :=
  s.fvarSet.fold (init := toProcess) fun toProcess fvarId =>
    if toProcess.contains fvarId then toProcess else toProcess.push fvarId

private def pushLocalDecl (toProcess : Array FVarId) (fvarId : FVarId) (userName : Name) (type : Expr) (bi := BinderInfo.default)
    : StateRefT ClosureState TermElabM (Array FVarId) := do
  let type ← preprocess type
  modify fun s => { s with
    newLocalDecls := s.newLocalDecls.push $ LocalDecl.cdecl arbitrary fvarId userName type bi,
    exprArgs      := s.exprArgs.push (mkFVar fvarId)
  }
  pure $ pushNewVars toProcess (collectFVars {} type)

private partial def mkClosureForAux (toProcess : Array FVarId) : StateRefT ClosureState TermElabM Unit := do
  let lctx ← getLCtx
  match pickMaxFVar? lctx toProcess with
  | none        => pure ()
  | some fvarId =>
    trace[Elab.definition.mkClosure] "toProcess: {toProcess.map mkFVar}, maxVar: {mkFVar fvarId}"
    let toProcess := toProcess.erase fvarId
    let localDecl ← getLocalDecl fvarId
    match localDecl with
    | LocalDecl.cdecl _ _ userName type bi =>
      let toProcess ← pushLocalDecl toProcess fvarId userName type bi
      mkClosureForAux toProcess
    | LocalDecl.ldecl _ _ userName type val _ =>
      let zetaFVarIds ← getZetaFVarIds
      if !zetaFVarIds.contains fvarId then
        /- Non-dependent let-decl. See comment at src/Lean/Meta/Closure.lean -/
        let toProcess ← pushLocalDecl toProcess fvarId userName type
        mkClosureForAux toProcess
      else
        /- Dependent let-decl. -/
        let type ← preprocess type
        let val  ← preprocess val
        modify fun s => { s with
          newLetDecls   := s.newLetDecls.push $ LocalDecl.ldecl arbitrary fvarId userName type val false,
          /- We don't want to interleave let and lambda declarations in our closure. So, we expand any occurrences of fvarId
             at `newLocalDecls` and `localDecls` -/
          newLocalDecls := s.newLocalDecls.map (replaceFVarIdAtLocalDecl fvarId val),
          localDecls := s.localDecls.map (replaceFVarIdAtLocalDecl fvarId val)
        }
        mkClosureForAux (pushNewVars toProcess (collectFVars (collectFVars {} type) val))

private partial def mkClosureFor (freeVars : Array FVarId) (localDecls : Array LocalDecl) : TermElabM ClosureState := do
  let (_, s) ← (mkClosureForAux freeVars).run { localDecls := localDecls }
  pure { s with
    newLocalDecls := s.newLocalDecls.reverse,
    newLetDecls   := s.newLetDecls.reverse,
    exprArgs      := s.exprArgs.reverse
  }

structure LetRecClosure where
  ref        : Syntax
  localDecls : Array LocalDecl
  closed     : Expr -- expression used to replace occurrences of the let-rec FVarId
  toLift     : LetRecToLift

private def mkLetRecClosureFor (toLift : LetRecToLift) (freeVars : Array FVarId) : TermElabM LetRecClosure := do
  let lctx := toLift.lctx
  withLCtx lctx toLift.localInstances do
  lambdaTelescope toLift.val fun xs val => do
    let type ← instantiateForall toLift.type xs
    let lctx ← getLCtx
    let s ← mkClosureFor freeVars $ xs.map fun x => lctx.get! x.fvarId!
    let type := Closure.mkForall s.localDecls $ Closure.mkForall s.newLetDecls type
    let val  := Closure.mkLambda s.localDecls $ Closure.mkLambda s.newLetDecls val
    let c    := mkAppN (Lean.mkConst toLift.declName) s.exprArgs
    assignExprMVar toLift.mvarId c
    return {
      ref        := toLift.ref
      localDecls := s.newLocalDecls
      closed     := c
      toLift     := { toLift with val := val, type := type }
    }

private def mkLetRecClosures (letRecsToLift : List LetRecToLift) (freeVarMap : FreeVarMap) : TermElabM (List LetRecClosure) :=
  letRecsToLift.mapM fun toLift => mkLetRecClosureFor toLift (freeVarMap.find? toLift.fvarId).get!

/- Mapping from FVarId of mutually recursive functions being defined to "closure" expression. -/
abbrev Replacement := NameMap Expr

def insertReplacementForMainFns (r : Replacement) (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainFVars : Array Expr) : Replacement :=
  mainFVars.size.fold (init := r) fun i r =>
    r.insert mainFVars[i].fvarId! (mkAppN (Lean.mkConst mainHeaders[i].declName) sectionVars)

def insertReplacementForLetRecs (r : Replacement) (letRecClosures : List LetRecClosure) : Replacement :=
  letRecClosures.foldl (init := r) fun r c =>
    r.insert c.toLift.fvarId c.closed

def Replacement.apply (r : Replacement) (e : Expr) : Expr :=
  e.replace fun e => match e with
    | Expr.fvar fvarId _ => match r.find? fvarId with
      | some c => some c
      | _      => none
    | _ => none

def pushMain (preDefs : Array PreDefinition) (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainVals : Array Expr)
    : TermElabM (Array PreDefinition) :=
  mainHeaders.size.foldM (init := preDefs) fun i preDefs => do
    let header := mainHeaders[i]
    let val  ← mkLambdaFVars sectionVars mainVals[i]
    let type ← mkForallFVars sectionVars header.type
    return preDefs.push {
      ref         := getDeclarationSelectionRef header.ref
      kind        := header.kind
      declName    := header.declName
      levelParams := [], -- we set it later
      modifiers   := header.modifiers
      type        := type
      value       := val
    }

def pushLetRecs (preDefs : Array PreDefinition) (letRecClosures : List LetRecClosure) (kind : DefKind) (modifiers : Modifiers) : Array PreDefinition :=
  letRecClosures.foldl (init := preDefs) fun preDefs c =>
    let type := Closure.mkForall c.localDecls c.toLift.type
    let val  := Closure.mkLambda c.localDecls c.toLift.val
    preDefs.push {
      ref         := c.ref
      kind        := kind
      declName    := c.toLift.declName
      levelParams := [] -- we set it later
      modifiers   := { modifiers with attrs := c.toLift.attrs }
      type        := type
      value       := val
    }

def getKindForLetRecs (mainHeaders : Array DefViewElabHeader) : DefKind :=
  if mainHeaders.any fun h => h.kind.isTheorem then DefKind.«theorem»
  else DefKind.«def»

def getModifiersForLetRecs (mainHeaders : Array DefViewElabHeader) : Modifiers := {
  isNoncomputable := mainHeaders.any fun h => h.modifiers.isNoncomputable,
  isPartial       := mainHeaders.any fun h => h.modifiers.isPartial,
  isUnsafe        := mainHeaders.any fun h => h.modifiers.isUnsafe
}

/-
- `sectionVars`:   The section variables used in the `mutual` block.
- `mainHeaders`:   The elaborated header of the top-level definitions being defined by the mutual block.
- `mainFVars`:     The auxiliary variables used to represent the top-level definitions being defined by the mutual block.
- `mainVals`:      The elaborated value for the top-level definitions
- `letRecsToLift`: The let-rec's definitions that need to be lifted
-/
def main (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainFVars : Array Expr) (mainVals : Array Expr) (letRecsToLift : List LetRecToLift)
    : TermElabM (Array PreDefinition) := do
  -- Store in recFVarIds the fvarId of every function being defined by the mutual block.
  let mainFVarIds := mainFVars.map Expr.fvarId!
  let recFVarIds  := (letRecsToLift.toArray.map fun toLift => toLift.fvarId) ++ mainFVarIds
  -- Compute the set of free variables (excluding `recFVarIds`) for each let-rec.
  let mctx ← getMCtx
  let freeVarMap := mkFreeVarMap mctx sectionVars mainFVarIds recFVarIds letRecsToLift
  resetZetaFVarIds
  withTrackingZeta do
    -- By checking `toLift.type` and `toLift.val` we populate `zetaFVarIds`. See comments at `src/Lean/Meta/Closure.lean`.
    letRecsToLift.forM fun toLift => withLCtx toLift.lctx toLift.localInstances do Meta.check toLift.type; Meta.check toLift.val
    let letRecClosures ← mkLetRecClosures letRecsToLift freeVarMap
    -- mkLetRecClosures assign metavariables that were placeholders for the lifted declarations.
    let mainVals    ← mainVals.mapM (instantiateMVars ·)
    let mainHeaders ← mainHeaders.mapM instantiateMVarsAtHeader
    let letRecClosures ← letRecClosures.mapM fun closure => do pure { closure with toLift := (← instantiateMVarsAtLetRecToLift closure.toLift) }
    -- Replace fvarIds for functions being defined with closed terms
    let r              := insertReplacementForMainFns {} sectionVars mainHeaders mainFVars
    let r              := insertReplacementForLetRecs r letRecClosures
    let mainVals       := mainVals.map r.apply
    let mainHeaders    := mainHeaders.map fun h => { h with type := r.apply h.type }
    let letRecClosures := letRecClosures.map fun c => { c with toLift := { c.toLift with type := r.apply c.toLift.type, val := r.apply c.toLift.val } }
    let letRecKind     := getKindForLetRecs mainHeaders
    let letRecMods     := getModifiersForLetRecs mainHeaders
    pushMain (pushLetRecs #[] letRecClosures letRecKind letRecMods) sectionVars mainHeaders mainVals

end MutualClosure

private def getAllUserLevelNames (headers : Array DefViewElabHeader) : List Name :=
  if h : 0 < headers.size then
    -- Recall that all top-level functions must have the same levels. See `check` method above
    (headers.get ⟨0, h⟩).levelNames
  else
    []

/-- Eagerly convert universe metavariables occurring in theorem headers to universe parameters. -/
private def levelMVarToParamHeaders (views : Array DefView) (headers : Array DefViewElabHeader) : TermElabM (Array DefViewElabHeader) := do
  let rec process : StateRefT Nat TermElabM (Array DefViewElabHeader) := do
    let mut newHeaders := #[]
    for view in views, header in headers do
      if view.kind.isTheorem then
        newHeaders := newHeaders.push { header with type := (← levelMVarToParam' header.type) }
      else
        newHeaders := newHeaders.push header
    return newHeaders
  let newHeaders ← process.run' 1
  newHeaders.mapM fun header => return { header with type := (← instantiateMVars header.type) }

def elabMutualDef (vars : Array Expr) (views : Array DefView) : TermElabM Unit :=
  if isExample views then
    withoutModifyingEnv go
  else
    go
where go := do
  let scopeLevelNames ← getLevelNames
  let headers ← elabHeaders views
  let headers ← levelMVarToParamHeaders views headers
  let allUserLevelNames := getAllUserLevelNames headers
  withFunLocalDecls headers fun funFVars => do
    let values ← elabFunValues headers
    Term.synthesizeSyntheticMVarsNoPostponing
    let values ← values.mapM (instantiateMVars ·)
    let headers ← headers.mapM instantiateMVarsAtHeader
    let letRecsToLift ← getLetRecsToLift
    let letRecsToLift ← letRecsToLift.mapM instantiateMVarsAtLetRecToLift
    checkLetRecsToLiftTypes funFVars letRecsToLift
    withUsed vars headers values letRecsToLift fun vars => do
      let preDefs ← MutualClosure.main vars headers funFVars values letRecsToLift
      let preDefs ← levelMVarToParamPreDecls preDefs
      let preDefs ← instantiateMVarsAtPreDecls preDefs
      let preDefs ← fixLevelParams preDefs scopeLevelNames allUserLevelNames
      addPreDefinitions preDefs

end Term
namespace Command

def elabMutualDef (ds : Array Syntax) : CommandElabM Unit := do
  let views ← ds.mapM fun d => do
    let modifiers ← elabModifiers d[0]
    mkDefView modifiers d[1]
  runTermElabM none fun vars => Term.elabMutualDef vars views

end Command
end Lean.Elab


import Structure.Generic.Axioms.Universes

import mathlib4_experiments.Data.Equiv.Basic

set_option autoBoundImplicitLocal false
--set_option pp.universes true

universes u v w

-- We additionally want "sort-like" types to have some concept of "functors" that map instances. Here, we
-- need to reconcile two conflicting requirements:
--
-- * We want a functor `F : α ⟶ β` with `α β : V` to be an instance of `V`, so that we can chain functors
--   as in `α ⟶ β ⟶ γ`.
--
-- * We axiomatically assert the existence of certain functors such as identity and composition, which we
--   assume to map instances in a specific way. We want this mapping to be a definitional equality so that
--   e.g. applying the identity functor is a trivial operation.
--   This implies that "functoriality" should be extra structure on functions, so that we can say that a
--   given function "is functorial".
--
-- In order to meet both requirements, we define both kinds of functors, and write `α ⟶ β` for the first
-- kind and `α ⟶' β` for a bundled version of the second kind, i.e. a function that is functorial. We
-- assert the equivalence of these two kinds of functors as an axiom.
--
-- In the case of `Bundled C`, we can make sure that `α ⟶' β` is actually an instance of the type class
-- `C`, so we can define `α ⟶ β` to be the same as `α ⟶' β`. However, in the case of `Sort u`, `α ⟶' β`
-- is not an instance of `Sort u` if `u = 0` (i.e. `Sort u` is actually `Prop`).
--
-- Moreover, `α ⟶' β` is defined so that `α` and `β` can live in different universes.

class HasExternalFunctors (U : Universe.{u}) (V : Universe.{v}) : Type (max u v) where
(IsFun {α : U} {β : V} : (α → β) → Sort (max u v))

structure BundledFunctor {U : Universe.{u}} {V : Universe.{v}} [h : HasExternalFunctors U V]
                         (α : U) (β : V) : Sort (max 1 u v) where
(f     : α → β)
(isFun : h.IsFun f)

namespace BundledFunctor

  infixr:20 " ⟶' " => BundledFunctor

  variable {U V : Universe} [h : HasExternalFunctors U V]

  instance coeFun (α : U) (β : V) : CoeFun (α ⟶' β) (λ _ => α → β) := ⟨BundledFunctor.f⟩

  def mkFun {α : U} {β : V} {f : α → β} (hf : h.IsFun f) : α ⟶' β := ⟨f, hf⟩

end BundledFunctor

class HasInternalFunctors (U : Universe.{u}) extends HasExternalFunctors U U : Type u where
(Fun                : U → U → U)
(funEquiv (α β : U) : ⌈Fun α β⌉ ≃ (α ⟶' β))

namespace HasInternalFunctors

  infixr:20 " ⟶ " => HasInternalFunctors.Fun

  variable {U : Universe} [h : HasInternalFunctors U]

  def toBundled   {α β : U} (F : α ⟶  β) : α ⟶' β := (h.funEquiv α β).toFun  F
  def fromBundled {α β : U} (F : α ⟶' β) : α ⟶  β := (h.funEquiv α β).invFun F

  @[simp] theorem fromToBundled {α β : U} (F : α ⟶  β) : fromBundled (toBundled F) = F := (h.funEquiv α β).leftInv  F
  @[simp] theorem toFromBundled {α β : U} (F : α ⟶' β) : toBundled (fromBundled F) = F := (h.funEquiv α β).rightInv F

  def funCoe {α β : U} (F : α ⟶ β) : α → β := (toBundled F).f
  instance (α β : U) : CoeFun ⌈α ⟶ β⌉ (λ _ => α → β) := ⟨funCoe⟩

  -- Workaround for cases where `coeFun` doesn't work.
  notation:max F:max "⟮" x:0 "⟯" => HasInternalFunctors.funCoe F x

  def isFun {α β : U} (F : α ⟶ β) : h.IsFun (funCoe F) := (toBundled F).isFun

  theorem toBundled.eff {α β : U} (F : α ⟶ β) (a : α) : (toBundled F) a = F a := rfl

  @[simp] theorem fromBundled.coe {α β : U} (F : α ⟶' β) : funCoe (fromBundled F) = F.f :=
  congrArg BundledFunctor.f (toFromBundled F)
  @[simp] theorem fromBundled.eff {α β : U} (F : α ⟶' β) (a : α) : (fromBundled F) a = F a :=
  congrFun (fromBundled.coe F) a

  def mkFun {α β : U} {f : α → β} (hf : h.IsFun f) : α ⟶ β := fromBundled (BundledFunctor.mkFun hf)

  @[simp] theorem mkFun.eff {α β : U} {f : α → β} (hf : h.IsFun f) (a : α) :
    (mkFun hf) a = f a :=
  fromBundled.eff (BundledFunctor.mkFun hf) a

end HasInternalFunctors

@[simp] theorem elimRec {α : Sort u} {a a' : α} {ha : a = a'}
                        {T : α → Sort v} {x : T a}
                        {β : Sort w} {f : {a : α} → T a → β} :
  @f a' (ha ▸ x) = f x :=
by subst ha; rfl

-- The following axioms are equivalent to asserting the existence of five functors with specified behavior:
-- id    : `α ⟶ α,                           a ↦ a`
-- const : `β ⟶ (α ⟶ β),                     c ↦ (a ↦ c)`
-- app   : `α ⟶ (α ⟶ β) ⟶ β,                 a ↦ (F ↦ F a)`
-- dup   : `(α ⟶ α ⟶ β) ⟶ (α ⟶ β),           F ↦ (a ↦ F a a)`
-- comp  : `(α ⟶ β) ⟶ (β ⟶ γ) ⟶ (α ⟶ γ),     F ↦ (G ↦ (a ↦ G (F a)))`
--
-- In `DerivedFunctors.lean`, we construct several other functors such as
-- swap  : `(α ⟶ β ⟶ γ) ⟶ (β ⟶ α ⟶ γ),       F ↦ (b ↦ (a ↦ F a b))`
-- subst : `(α ⟶ β ⟶ γ) ⟶ (α ⟶ β) ⟶ (α ⟶ γ), F ↦ (G ↦ (a ↦ F a (G a)))`
-- Using these, we can give a general algorithm for proving that a function is functorial.

class HasIdFun (U : Universe) [h : HasExternalFunctors U U] where
(idIsFun (α : U) : h.IsFun (λ a : α => a))

namespace HasIdFun

  variable {U : Universe} [HasExternalFunctors U U] [h : HasIdFun U]

  def idFun' (α : U) : α ⟶' α := BundledFunctor.mkFun (h.idIsFun α)

end HasIdFun

class HasConstFun (U V : Universe) [h : HasExternalFunctors U V] where
(constIsFun (α : U) {β : V} (c : β) : h.IsFun (λ a : α => c))

namespace HasConstFun

  variable {U V : Universe} [HasExternalFunctors U V] [h : HasConstFun U V]

  def constFun' (α : U) {β : V} (c : β) : α ⟶' β := BundledFunctor.mkFun (h.constIsFun α c)

end HasConstFun

class HasCompFun (U V W : Universe) [HasExternalFunctors U V] [HasExternalFunctors V W] [h : HasExternalFunctors U W] where
(compIsFun {α : U} {β : V} {γ : W} (F : α ⟶' β) (G : β ⟶' γ) : h.IsFun (λ a : α => G (F a)))

namespace HasCompFun

  variable {U V W : Universe} [HasExternalFunctors U V] [HasExternalFunctors V W] [HasExternalFunctors U W] [h : HasCompFun U V W]

  def compFun' {α : U} {β : V} {γ : W} (F : α ⟶' β) (G : β ⟶' γ) : α ⟶' γ := BundledFunctor.mkFun (h.compIsFun F G)

  def revCompFun' {α : U} {β : V} {γ : W} (G : β ⟶' γ) (F : α ⟶' β) : α ⟶' γ := compFun' F G
  infixr:90 " ⊙' " => HasCompFun.revCompFun'

end HasCompFun

class HasLinearFunOp (U : Universe) [h : HasInternalFunctors U] extends HasIdFun U, HasCompFun U U U where
(appIsFun        {α : U} (a : α) (β : U)        : h.IsFun (λ F : α ⟶ β     => F a))
(appFunIsFun     (α β : U)                      : h.IsFun (λ a : α         => HasInternalFunctors.mkFun (appIsFun a β)))
(compFunIsFun    {α β : U} (F : α ⟶' β) (γ : U) : h.IsFun (λ G : β ⟶ γ     => HasInternalFunctors.mkFun (compIsFun F (HasInternalFunctors.toBundled G))))
(compFunFunIsFun (α β γ : U)                    : h.IsFun (λ F : α ⟶ β     => HasInternalFunctors.mkFun (compFunIsFun (HasInternalFunctors.toBundled F) γ)))

namespace HasLinearFunOp

  variable {U : Universe} [HasInternalFunctors U] [h : HasLinearFunOp U]

  def idFun' (α : U) : α ⟶' α := HasIdFun.idFun' α
  def idFun  (α : U) : α ⟶  α := HasInternalFunctors.fromBundled (idFun' α)

  @[simp] theorem idFun.eff (α : U) (a : α) : (idFun α) a = a :=
  by apply HasInternalFunctors.fromBundled.eff

  def appFun' {α : U} (a : α) (β : U) : (α ⟶ β) ⟶' β := BundledFunctor.mkFun (h.appIsFun a β)
  def appFun  {α : U} (a : α) (β : U) : (α ⟶ β) ⟶  β := HasInternalFunctors.fromBundled (appFun' a β)

  @[simp] theorem appFun.eff {α : U} (a : α) (β : U) (F : α ⟶ β) : (appFun a β) F = F a :=
  by apply HasInternalFunctors.fromBundled.eff

  def appFunFun' (α β : U) : α ⟶' (α ⟶ β) ⟶ β := BundledFunctor.mkFun (h.appFunIsFun α β)
  def appFunFun  (α β : U) : α ⟶  (α ⟶ β) ⟶ β := HasInternalFunctors.fromBundled (appFunFun' α β)

  @[simp] theorem appFunFun.eff (α β : U) (a : α) : (appFunFun α β) a = appFun a β :=
  by apply HasInternalFunctors.fromBundled.eff
  @[simp] theorem appFunFun.effEff (α β : U) (a : α) (F : α ⟶ β) : ((appFunFun α β) a) F = F a :=
  by simp

  def compFun' {α β γ : U} (F : α ⟶' β) (G : β ⟶' γ) : α ⟶' γ := HasCompFun.compFun' F G
  def compFun  {α β γ : U} (F : α ⟶  β) (G : β ⟶  γ) : α ⟶  γ :=
  HasInternalFunctors.fromBundled (compFun' (HasInternalFunctors.toBundled F) (HasInternalFunctors.toBundled G))

  @[simp] theorem compFun.eff {α β γ : U} (F : α ⟶ β) (G : β ⟶ γ) (a : α) : (compFun F G) a = G (F a) :=
  by apply HasInternalFunctors.fromBundled.eff

  def compFunFun' {α β : U} (F : α ⟶' β) (γ : U) : (β ⟶ γ) ⟶' (α ⟶ γ) := BundledFunctor.mkFun (h.compFunIsFun F γ)
  def compFunFun  {α β : U} (F : α ⟶  β) (γ : U) : (β ⟶ γ) ⟶  (α ⟶ γ) :=
  HasInternalFunctors.fromBundled (compFunFun' (HasInternalFunctors.toBundled F) γ)

  @[simp] theorem compFunFun.eff {α β : U} (F : α ⟶ β) (γ : U) (G : β ⟶ γ) : (compFunFun F γ) G = compFun F G :=
  by apply HasInternalFunctors.fromBundled.eff
  @[simp] theorem compFunFun.effEff {α β : U} (F : α ⟶ β) (γ : U) (G : β ⟶ γ) (a : α) : ((compFunFun F γ) G) a = G (F a) :=
  by simp

  def compFunFunFun' (α β γ : U) : (α ⟶ β) ⟶' (β ⟶ γ) ⟶ (α ⟶ γ) := BundledFunctor.mkFun (h.compFunFunIsFun α β γ)
  def compFunFunFun  (α β γ : U) : (α ⟶ β) ⟶  (β ⟶ γ) ⟶ (α ⟶ γ) := HasInternalFunctors.fromBundled (compFunFunFun' α β γ)

  @[simp] theorem compFunFunFun.eff (α β γ : U) (F : α ⟶ β) : (compFunFunFun α β γ) F = compFunFun F γ :=
  by apply HasInternalFunctors.fromBundled.eff
  @[simp] theorem compFunFunFun.effEff (α β γ : U) (F : α ⟶ β) (G : β ⟶ γ) : ((compFunFunFun α β γ) F) G = compFun F G :=
  by simp
  @[simp] theorem compFunFunFun.effEffEff (α β γ : U) (F : α ⟶ β) (G : β ⟶ γ) (a : α) : (((compFunFunFun α β γ) F) G) a = G (F a) :=
  by simp

end HasLinearFunOp

class HasSubLinearFunOp (U : Universe) [h : HasInternalFunctors U] extends HasConstFun U U where
(constFunIsFun (α β : U) : h.IsFun (λ c : β => HasInternalFunctors.mkFun (constIsFun α c)))

namespace HasSubLinearFunOp

  variable {U : Universe} [HasInternalFunctors U] [h : HasSubLinearFunOp U]

  def constFun' (α : U) {β : U} (c : β) : α ⟶' β := HasConstFun.constFun' α c
  def constFun  (α : U) {β : U} (c : β) : α ⟶  β := HasInternalFunctors.fromBundled (constFun' α c)

  @[simp] theorem constFun.eff (α : U) {β : U} (c : β) (a : α) : (constFun α c) a = c :=
  by apply HasInternalFunctors.fromBundled.eff

  def constFunFun' (α β : U) : β ⟶' (α ⟶ β) := BundledFunctor.mkFun (h.constFunIsFun α β)
  def constFunFun  (α β : U) : β ⟶  (α ⟶ β) := HasInternalFunctors.fromBundled (constFunFun' α β)

  @[simp] theorem constFunFun.eff (α β : U) (c : β) : (constFunFun α β) c = constFun α c :=
  by apply HasInternalFunctors.fromBundled.eff
  @[simp] theorem constFunFun.effEff (α β : U) (c : β) (a : α) : ((constFunFun α β) c) a = c :=
  by simp

end HasSubLinearFunOp

class HasAffineFunOp (U : Universe) [h : HasInternalFunctors U] extends HasLinearFunOp U, HasSubLinearFunOp U

class HasNonLinearFunOp (U : Universe) [h : HasInternalFunctors U] where
(dupIsFun    {α β : U} (F : α ⟶' α ⟶ β) : h.IsFun (λ a : α         => F a a))
(dupFunIsFun (α β : U)                  : h.IsFun (λ F : α ⟶ α ⟶ β => HasInternalFunctors.mkFun (dupIsFun (HasInternalFunctors.toBundled F))))

namespace HasNonLinearFunOp

  variable {U : Universe} [HasInternalFunctors U] [h : HasNonLinearFunOp U]

  def dupFun' {α β : U} (F : α ⟶' α ⟶ β) : α ⟶' β := BundledFunctor.mkFun (h.dupIsFun F)
  def dupFun  {α β : U} (F : α ⟶  α ⟶ β) : α ⟶  β :=
  HasInternalFunctors.fromBundled (dupFun' (HasInternalFunctors.toBundled F))

  @[simp] theorem dupFun.eff {α β : U} (F : α ⟶ α ⟶ β) (a : α) : (dupFun F) a = F a a :=
  by apply HasInternalFunctors.fromBundled.eff

  def dupFunFun' (α β : U) : (α ⟶ α ⟶ β) ⟶' (α ⟶ β) := BundledFunctor.mkFun (h.dupFunIsFun α β)
  def dupFunFun  (α β : U) : (α ⟶ α ⟶ β) ⟶  (α ⟶ β) := HasInternalFunctors.fromBundled (dupFunFun' α β)

  @[simp] theorem dupFunFun.eff (α β : U) (F : α ⟶ α ⟶ β) : (dupFunFun α β) F = dupFun F :=
  by apply HasInternalFunctors.fromBundled.eff
  @[simp] theorem dupFunFun.effEff (α β : U) (F : α ⟶ α ⟶ β) (a : α) : ((dupFunFun α β) F) a = F a a :=
  by simp

end HasNonLinearFunOp

class HasFullFunOp (U : Universe) [h : HasInternalFunctors U] extends HasAffineFunOp U, HasNonLinearFunOp U

class HasFunOp (U : Universe.{u}) extends HasInternalFunctors U, HasFullFunOp U : Type u