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

/-!
# A tactic for normalizing casts inside expressions

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

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

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

setup_tactic_parser

namespace tactic

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

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

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

end tactic

namespace norm_cast

open tactic expr

declare_trace norm_cast

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

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

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

namespace label

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

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

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

end label

open label

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

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

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

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

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

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

meta instance : inhabited norm_cast_cache := ⟨empty_cache⟩

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

end norm_cast

namespace tactic.interactive
open norm_cast

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

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

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

end tactic.interactive

namespace norm_cast
open tactic expr

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

end norm_cast

namespace tactic
open expr norm_cast

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

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

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

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

end tactic

namespace tactic.interactive
open tactic norm_cast

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

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

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

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

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

end tactic.interactive

namespace conv.interactive
open conv
open norm_cast (derive)

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

end conv.interactive

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

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

add_hint_tactic "norm_cast at *"

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

/-!

# Open immersions

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

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

-/

noncomputable theory

open category_theory category_theory.limits opposite topological_space

universe u

namespace algebraic_geometry

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

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

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

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

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

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

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

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

end algebraic_geometry


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

/-

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

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

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

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

namespace Lean.Meta
namespace Closure

structure ToProcessElement where
  fvarId : FVarId
  newFVarId : FVarId
  deriving Inhabited

structure Context where
  zeta : Bool

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

abbrev ClosureM := ReaderT Context $ StateRefT State MetaM

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

end Closure

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

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

end Lean.Meta


import Mathlib.Tactic.Linarith
import Aesop
-- A translation from SSA + Regions to Tree, with no proofs.
-- this allows for easy unfolding of the semantics.
-- If we carry around proofs of well formedness, the dependent typing
-- of the well-formedness leads to stuck terms.
-- Thus, we eschew the need to have proofs, and simply YOLO translate from
-- the origina SSA + Regions program into Tree.

namespace SSARgnVar2Tree

inductive RgnName
| r0
| r1
| r2
| r3
| r4
deriving DecidableEq

inductive VarName
| null
| x0
| x1
| x2
| x3
| x4
| x5
| x6
| x7
| y1
| y2
| y3
| y4
| z1
| z2
| z3
| z4
deriving DecidableEq

section StxSem

-- | The environment is a mapping from variable names to values.
abbrev Env (k : Type) (α : Type) := k → α

-- | Extend the environment with a new variable.
@[simp]
def Env.extend [DecidableEq k] (name : k) (v : α) (e: Env k α) : Env k α :=
  fun name' => if name = name' then v else e name'

@[simp]
def Env.empty [Inhabited α] : Env k α :=
  fun _ => default

@[simp]
def Env.map (f : α → β) (e : Env k α) : Env k β :=
  fun name => f (e name)

notation "∅" => Env.empty
notation e "[" name "↦" v "]" => Env.extend name v e

abbrev VarEnv (α : Type) := Env VarName α
abbrev RgnEnv (α : Type) := Env RgnName α

-- The input semantics given by the user.
class UserSemantics (opcode : Type)  where
  -- | Arguments given as (<args>, <rgn1>, ... <rgnN>)
  -- | Consider not allowing users to not be access all of 'Env'.
  opcodeEval (op: opcode) (vals : Int × Int) (rgns : (Int × Int → Int)) : Int

attribute [simp] UserSemantics.opcodeEval

inductive ASTKind : Type where
| O : ASTKind
| Os : ASTKind
| R : ASTKind
deriving Inhabited, DecidableEq

-- | The operations of the language.
inductive AST (opcode : Type): ASTKind → Type where
| assign (ret : VarName) (op : opcode)
  (args : VarName × VarName)
  (rgns : AST opcode .R) : AST opcode .O
| ops1 (op : AST opcode .O) :  AST opcode .Os
| opsmany (op : AST opcode .O) (ops : AST opcode .Os) : AST opcode .Os
| rgn (args : VarName × VarName) (body : AST opcode .Os) : AST opcode .R
| rgnvar (var : RgnName) : AST opcode .R
| rgn0 : AST opcode .R

def AST.retname : AST opcode .O → VarName
| .assign (ret := ret) .. => ret

instance [Inhabited opcode] : Inhabited (AST opcode .O) where
  default := .assign .x0 default (.x0, .x0) .rgn0

instance [Inhabited opcode] : Inhabited (AST opcode .Os) where
  default := .ops1 default

instance [Inhabited opcode] : Inhabited (AST opcode .R) where
  default := .rgn (.x0, .x0) default

@[simp]
def Ops.ofList [Inhabited opcode] : List (AST opcode .O) → AST opcode .Os
| [] => panic! "need non-empty list"
| [x] => .ops1 x
| x :: xs => .opsmany x (Ops.ofList xs)

@[reducible, simp]
instance [Inhabited opcode] : Coe (List (AST opcode .O)) (AST opcode .Os) := ⟨Ops.ofList⟩

@[simp, reducible]
def ASTKind.eval : ASTKind → Type
| .O => Int
| .Os => Int
| .R => (Int × Int → Int)
-- evaluate an operation with repect to a particular user semantics.
@[simp]
def AST.eval [S: UserSemantics opcode]
  {astk: ASTKind} (e: VarEnv Int)
  (re: RgnEnv (Int × Int → Int)): AST opcode astk → astk.eval × VarEnv Int
| .assign ret op args r =>
    let (arg1, arg2) := args
    let retval := S.opcodeEval op (e arg1, e arg2) (r.eval e re).fst
    (retval, e.extend ret retval)
| .ops1 op =>
   let (out, env) := op.eval e re
   (out, env)
| .opsmany op ops =>
    let e' := (op.eval e re).snd
    ops.eval e' re
| .rgnvar v => (re v, e)
| .rgn args body =>
    (fun vals =>
      let e := Env.empty
      let e1 := e.extend args.fst vals.fst
      let e2 := e1.extend args.snd vals.snd
      let (outval, _e) := body.eval e2 re
      outval, e)
| .rgn0 => (fun _ => default, e)

end StxSem

section Tree

inductive CtreeKind
| O -- op
| R -- region (higher order)
deriving Inhabited

-- closed trees, leaves are integers.
inductive Ctree  (opcode : Type) : CtreeKind → Type where
| binop
  (op : opcode)
  (lhs : Ctree opcode .O)
  (rhs : Ctree opcode .O)
  (rgns: Ctree opcode .R) : Ctree opcode .O
| rgn (f : Int × Int → Ctree opcode .O) : Ctree opcode .R
| rgn0 : Ctree opcode .R
| leaf (val : Int) : Ctree opcode .O

instance : Inhabited (Ctree opcode .O) where
  default := .leaf 10

instance : Inhabited (Ctree opcode .R) where
  default := .rgn0

-- convert an AST into a closed tree under the given environment
-- note that translation into a closed tree needs an environment,
-- to learn the values of variables.
-- This version has an `_` in the name since it needs a `Ctree.VarEnv`, not an
-- `Env`. We will writea helper that converts `Env` into `Ctree.VarEnv`.
@[simp, reducible]
def ASTKind.toCTree (opcode : Type) : ASTKind → Type
| .O => Ctree opcode .O
| .Os => Ctree opcode .O
| .R => Ctree opcode .R

@[simp]
def AST.toCtree_ {astk: ASTKind} (e : VarEnv (Ctree opcode .O))
  (re: RgnEnv (Ctree opcode .R)):
  AST opcode astk → (astk.toCTree opcode) × VarEnv (Ctree opcode .O)
| .assign ret opcode (u, v) r =>
    let rval := (r.toCtree_ e re).fst
    let t := .binop opcode (e u) (e v) rval
    (t, e.extend ret t)
| .rgnvar var => let e := Env.empty; (re var, e)
| .rgn0 => let e := Env.empty; (.rgn0, e)
| .rgn args body =>
    let e := Env.empty -- NOTE: regions are now isolated from above
    (.rgn fun vals =>
      let e1 := e.extend args.fst (.leaf vals.fst)
      let e2 := e1.extend args.snd (.leaf vals.snd)
      (body.toCtree_ e2 re).fst, e)
| .ops1 op =>
  let (val, e) := op.toCtree_ e re
  (val, e)
| .opsmany os o =>
    let (_, e) := os.toCtree_ e re
    o.toCtree_ e re

-- wrap every element in a (.leaf) constructor
@[simp]
def Ctree.VarEnv.ofEnv (e: VarEnv Int) : VarEnv (Ctree opcode .O) :=
   fun name => .leaf (e name)

-- TODO: should these be coercions?
@[simp]
def Ctree.RgnEnv.ofEnv (re: RgnEnv (Int × Int → Int)) :
  RgnEnv (Ctree opcode .R) :=
   fun name => Ctree.rgn (fun args => .leaf (re name args))

-- this converts an AST into a Ctree, given an environment
-- and an AST.
@[simp]
def Op.toCtree (a: AST opcode .O) (e: VarEnv Int)
  (re: RgnEnv (Int × Int → Int)) : Ctree opcode .O :=
    (a.toCtree_ (Ctree.VarEnv.ofEnv e) (Ctree.RgnEnv.ofEnv re)).fst

@[simp]
def Ops.toCtree (a: AST opcode .Os)
  (e: VarEnv Int) (re: RgnEnv (Int × Int → Int)) : Ctree opcode .O :=
    (a.toCtree_ (Ctree.VarEnv.ofEnv e) (Ctree.RgnEnv.ofEnv re)).fst

@[simp]
def Region.toCtree (a: AST opcode .R)
  (e: VarEnv Int)
  (re: RgnEnv (Int × Int → Int)): Ctree opcode .R :=
    (a.toCtree_ (Ctree.VarEnv.ofEnv e) (Ctree.RgnEnv.ofEnv re)).fst

-- evaluate a Ctree. note that this needs no environment.
@[simp]
def CtreeKind.eval : CtreeKind → Type
| .O => Int
| .R => Int × Int → Int

-- Note: is this literally "just" staging the partial evaluation against the environment?
def Ctree.eval [UserSemantics opcode] : Ctree opcode treek → treek.eval
| .binop o l r rs =>
  UserSemantics.opcodeEval o (l.eval, r.eval) rs.eval
| .leaf v => v
| .rgn0 => fun _ => default
| .rgn f => fun args => (f args).eval

end Tree

namespace MultipleInstructionTree
inductive Opcode
| add
| mul
| loop : Opcode
| ite : Opcode
| run : Opcode
| runnot : Opcode
| not : Opcode
| const : Int → Opcode
deriving Inhabited

def loopSemantics (n : Nat) (f : Int → Int) (v : Int) : Int :=
  match n with
  | 0 => v -- inline id
  | n + 1 => f (loopSemantics n f v) -- inline ∘

@[simp]
instance :  UserSemantics Opcode  where
  opcodeEval
  | .not, ⟨a, _⟩, _ => if a = 0 then 1 else 0
  | .mul, ⟨a, b⟩, _ => a * b
  | .add, ⟨a, b⟩, _ => a + b
  | .const i, ⟨_a, _b⟩, _ => i
  | .run, ⟨v, w⟩, r => r ⟨v, w⟩
  | .runnot, ⟨v, w⟩, r => r ⟨if v = 0 then 1 else 0, w⟩ -- execute: 'r(not v, w)'
  | .loop, ⟨n, init⟩, r => loopSemantics n.toNat (fun i => r ⟨i, 0⟩) init
  | _, _, _ => 42

def x_add_4_times_mul_val_eq (env: VarEnv Int):
  let p : AST Opcode .Os :=
      Ops.ofList [
      .assign .x1 .add (.x0, .x0) .rgn0,
      .assign .x2 .add (.x1, .x1) .rgn0
      ]
  let q : AST Opcode .Os := Ops.ofList [
        .assign .x1 (.const 4) (.x0, .x0) .rgn0
      , .assign .x2 .mul (.x1, .x0) .rgn0
    ]
  (Ops.toCtree p env renv).eval = (Ops.toCtree q env renv).eval := by {
    simp only [Ops.ofList, AST.eval, Ops.toCtree,  AST.toCtree_];
    -- see that there are environments, which are folded away when calling
    -- Ctree.eval.
    simp[Ctree.eval];
    linarith
  }

def run_inline :
  let p : AST Opcode .R :=
    AST.rgn ⟨.x0, .x1⟩ $ Ops.ofList [
      .assign .x2 .add (.x0, .x1) .rgn0 -- x2 := x0 + x1
      ]
  let q : AST Opcode .R :=
    AST.rgn ⟨.x0, .x1⟩ $ Ops.ofList [
      .assign .x2 .run (.x0, .x1) (AST.rgn ⟨.x5, .x6⟩ (Ops.ofList [
        -- x2 := run (x0, x1) { ^(x5, x6): return x5 + x6 }
        .assign .y1 .add (.x5, .x6) .rgn0
      ]))
    ]
  (Region.toCtree p Env.empty Env.empty).eval =
  (Region.toCtree q Env.empty Env.empty).eval := by {
    simp
    dsimp only[Ctree.eval]
  }

-- trying to convert an AST to a Ctree at any environment
-- is equivalent to converting it in an empty environment.
@[simp]
theorem AST.toCtree_rgn_equiv_empty
  (r: AST opcode .R)
  [Inhabited opcode] :
  AST.toCtree_ env renv r = AST.toCtree_ Env.empty renv r := by {
  simp; cases r <;> simp;
}

def runnot_inline:
  let p : AST Opcode .R :=
    AST.rgn ⟨.x0, .x1⟩ $ Ops.ofList [
      .assign .y1 .run ⟨.x0, .x1⟩ (.rgnvar .r1) -- y1 := r(x0, x1)
    ]
  (Region.toCtree p Env.empty Env.empty).eval =
  (Region.toCtree (.rgnvar .r1 : AST Opcode .R) Env.empty Env.empty).eval := by {
    simp
    simp[Ctree.eval];
    -- this stil cannot be simplified away, since this is in fact false.
    -- We need some notion of the fact that z1, z2 is free in 'r',
    -- so that we can say that we can remove 'z2, z1' in the 'q'
    -- environment when we run 'r'.
    -- if we have this, then we are good, and we can prove the theorem.
    -- but how do we convince people that use MLIR that this is a sensible
    -- thing to reason about?

  }

end MultipleInstructionTree

end SSARgnVar2Tree


namespace tactic.interactive
  open expr tactic

  private meta def congr_aux : expr → expr → tactic unit
  | (app f₁ a₁) (app f₂ a₂) :=
  do apply ``(congr),
     swap, reflexivity <|> swap,
     congr_aux f₁ f₂
  | _ _ := try reflexivity

  /-- Given a goal of form `f a1 ... an = f' a1' ... an'`, this tactic breaks it down to subgoals
      `f = f'`, `a1 = a1'`, ... Subgoals provable by reflexivity are dispensed automatically. -/
  meta def congruence : tactic unit :=
  do ```(%%lhs = %%rhs) ← target | fail "goal is not an equality",
     congr_aux lhs rhs

  /-- Given a goal that equates two structure values, this tactic breaks it down to subgoals equating each
      pair of fields. -/
  meta def congr_struct : tactic unit :=
  do ```(%%lhs = %%rhs) ← target | fail "goal is not an equality",
     ty ← infer_type lhs,
     [ctor] ← get_constructors_for ty | fail "equated type is not a structure",
     tactic.cases lhs,
     tactic.cases rhs,
     congruence
end tactic.interactive

structure X := ( x : unit ) ( y :  unit )

lemma test1 ( a b : X ) : a = b :=
begin
  congr_struct,
    -- x y x_1 y_1 : ℕ
    -- ⊢ x = x_1

    -- x y x_1 y_1 : ℕ
    -- ⊢ y = y_1
  {
      induction x,
    induction x_1,
    reflexivity
  },
  {
    induction y,
    induction y_1,
    reflexivity
  }
  -- Great!
end

def f ( a : X ) : X := { x := a.y, y := a.x }

lemma test2 ( a : X ) : a = f (f a) :=
begin
  congr_struct,
  -- breaks because cases.lhs messes up the right hand side!
end

structure Y := ( x : nat ) ( y : nat )
def g ( x : nat ) : Y := { x := x, y := x }
def h ( a : Y ) : Y := { x := a.y, y := a.x }

lemma test3 ( x : nat ) : g 0 = h ( g 0 ) :=
begin
  congr_struct,

end


import .fol .abel

universe u

-- section weekdays
-- @[derive has_reflect]
-- inductive weekday : Type
-- | monday          : weekday
-- | another_day     : weekday → weekday

-- open weekday

-- meta def dump_weekday (f : weekday) : tactic unit :=
-- tactic.trace $ to_string (expr.to_raw_fmt (reflect f).to_expr)

-- -- run_cmd dump_weekday (another_day (another_day monday))
-- --(app (const weekday.another_day []) (app (const weekday.another_day []) (const weekday.monday [])))

-- inductive weekday' : Type
-- | monday          : weekday'
-- | another_day     : weekday' → weekday'

-- open weekday'
-- meta instance has_reflect_weekday' : has_reflect weekday'
-- | weekday'.monday                  := `(monday)
-- | (weekday'.another_day x)         := `(λ l, weekday'.another_day l).subst $
--                                         by haveI := has_reflect_weekday'; exact (reflect x)

-- meta def dump_weekday' (f : weekday') : tactic unit :=
-- tactic.trace $ to_string (expr.to_raw_fmt (reflect f).to_expr)

-- -- run_cmd dump_weekday' (another_day (another_day monday))
-- -- (app (const weekday'.another_day []) (app (const weekday'.another_day []) (const weekday'.monday [])))
-- end weekdays
-- -- meta instance has_reflect_preterm {L : Language.{u}} : Π{n : ℕ}, has_reflect (preterm L n)
-- -- | 0 (var k) := `(@preterm.var L).subst (reflect k)

-- -- @[derive has_reflect]

-- open fol abel

-- section preterm_aux
-- inductive preterm_aux (L : Language.{u}) : Type u
-- | var     : ℕ → preterm_aux
-- | func    : ∀ k : ℕ, L.functions k → preterm_aux
-- | app     : preterm_aux → preterm_aux → preterm_aux

-- def to_aux {L : Language.{u}} : ∀ {l : ℕ},  preterm L l → preterm_aux L
-- | 0 (var n)      := preterm_aux.var _ n
-- | k (func f)     := preterm_aux.func _ f
-- | k (app t₁ t₂)  := preterm_aux.app (to_aux t₁) (to_aux t₂)

-- def L_abel_plus' (t₁ t₂ : preterm L_abel 0) : preterm L_abel 0 :=
-- @term_of_function L_abel 2 (abel_functions.plus : L_abel.functions 2) t₁ t₂
-- end preterm_aux

-- local infix ` +' `:100 := L_abel_plus'

-- local notation ` zero ` := (func abel_functions.zero : preterm L_abel 0)

-- section L_abel_term_biopsy

-- def sample1 : preterm L_abel 0 := (zero +' zero)

-- def sample2 : preterm L_abel 0 := zero

-- -- #reduce sample2

-- open expr
-- meta def sample2_expr : expr :=
-- mk_app (const `preterm.func list.nil) ([(const `L_abel list.nil), `(0), const `abel_functions.zero list.nil] : list expr)

-- end L_abel_term_biopsy

-- section simpler_biopsy

-- inductive my_inductive : Type
-- | a : my_inductive
-- | b : my_inductive
-- | f : my_inductive → my_inductive

-- open my_inductive
-- def sample3 : my_inductive := f a

-- open expr

-- meta def sample3_expr : expr :=
-- app (const `my_inductive.f list.nil) (const `my_inductive.a list.nil)

-- def sample3_again : my_inductive := by tactic.exact (sample3_expr)

-- example : sample3 = sample3_again := rfl

-- end simpler_biopsy

-- namespace tactic
-- namespace interactive
-- open interactive interactive.types expr

-- def my_test_term : preterm L_abel 0 := (zero +' zero)

-- end interactive
-- end tactic

-- section test
-- -- def my_term : preterm L_abel 0 := sorry

-- -- #check tactic.interactive.rcases

-- end test

-- section sample4

-- /-- Note: this is the same as `dfin` -/
-- inductive my_indexed_family : ℕ → Type u
-- | z {} : my_indexed_family 0
-- | s : ∀ {k}, my_indexed_family k → my_indexed_family (k+1)

-- -- meta example : ∀ {n}, has_reflect (my_indexed_family n)
-- -- | 0 z := `(z)

-- open my_indexed_family

-- def sample4 : my_indexed_family 1 := s z

-- -- #check tactic.eval_expr

-- end sample4

-- section sample4

-- inductive dfin'' : ℕ → Type
-- | fz {n} : dfin'' (n+1)
-- | fs {n} : dfin'' n → dfin'' (n+1)

-- inductive dfin' : ℕ → Type u
-- | gz {n} :  dfin' (n+1)
-- | gs {n} :  dfin' n → dfin' (n+1)

-- open dfin dfin'

-- meta instance dfin.reflect : ∀ {n}, has_reflect (dfin'' n)
-- | _ dfin''.fz := `(dfin''.fz)
-- | _ (dfin''.fs n) := `(dfin''.fs).subst (dfin.reflect n)

-- -- /- errors all over---why doesn't reflect like universe parameters? -/
-- -- meta instance dfin'.reflect : ∀ {n}, has_reflect (dfin' n)
-- -- | _ fz := `(fz)
-- -- | _ (fs n) := `(fs).subst (dfin'.reflect n)

-- end sample4

-- section reflect_preterm

-- /- Language with a single constant symbol -/
-- inductive L_pt_functions : ℕ → Type
-- | pt : L_pt_functions 0

-- def L_pt : Language.{0} := ⟨L_pt_functions, λ _, empty⟩

-- def pt_preterm : preterm L_pt 0 := preterm.func L_pt_functions.pt

-- meta def pt_preterm_reflected : expr :=
-- expr.mk_app (expr.const `preterm.func [level.zero]) [ (expr.const `L_pt list.nil), `(0), (expr.const `L_pt_functions.pt list.nil)]

-- set_option trace.app_builder true

-- -- meta def pt_preterm_reflected' : expr := by tactic.mk_app "preterm.func" [(expr.const `L_pt []), `(0), (expr.const `L_pt_functions.pt [])]

-- #check tactic.mk_app

-- meta def pt_preterm_reflected'' : tactic expr :=
-- tactic.to_expr ```(preterm.func L_pt_functions.pt : preterm L_pt 0)

-- def pt_preterm' : preterm L_pt 0 := by pt_preterm_reflected'' >>= tactic.exact

-- -- def pt_preterm' : preterm L_pt 0 := by tactic.exact pt_preterm_reflected

-- example : pt_preterm = pt_preterm' := rfl
--   -- infer type failed, incorrect number of universe levels

-- -- want: example : pt_preterm = pt_preterm' := rfl

-- end reflect_preterm

-- namespace hewwo
-- section reflect_preterm2
-- def L_pt.pt' : L_pt.functions 0 := L_pt_functions.pt

-- #reduce (by apply_instance : reflected L_pt.pt')
-- -- `(L_pt.pt')

-- meta def pt_preterm_reflected : tactic expr :=
-- tactic.mk_app ``preterm.func [`(L_pt.pt')]

-- def pt_preterm' : preterm L_pt 0 := by pt_preterm_reflected >>= tactic.exact

-- #eval tactic.trace (@expr.to_raw_fmt tt `(L_pt.pt'))

-- #check reflect

-- end reflect_preterm2
-- end hewwo

-- section reflect_preterm3

-- inductive L_pt_func_functions : ℕ → Type
-- | pt  : L_pt_func_functions 0
-- | foo : L_pt_func_functions 1

-- open L_pt_func_functions

-- def L_pt_func : Language.{0} :=
-- ⟨L_pt_func_functions, λ _, ulift empty⟩

-- -- def foo_pt_term : preterm L_pt_func 0 :=
-- -- preterm.app (preterm.func L_pt_func_functions.foo) (preterm.func L_pt_func_functions.pt)

-- -- def foo_pt_term_reflected : expr :=
-- -- begin
-- --   tactic.mk_app ``preterm.func [(by tactic.mk_app `preterm.func [`(L_pt_func_functions.foo)]), (by tactic.mk_app `preterm.func [`(L_pt_func_functions.pt)])]
-- -- end

-- -- def foo' : preterm L_pt_func 1 := preterm.func L_pt_func_functions.foo

-- -- #reduce (by apply_instance : reflected L_pt_func_functions.foo)

-- set_option trace.app_builder true

-- def my_foo : L_pt_func.functions 1 := L_pt_func_functions.foo

-- def my_pt : L_pt_func.functions 0 := L_pt_func_functions.pt

-- -- meta def foo_pt_term_reflected : tactic expr := tactic.mk_app ``preterm.func [`()]

-- meta def foo_pt_term_reflected' : tactic expr :=
-- do e₁ <- tactic.mk_app ``preterm.func [`(my_foo)],
--    e₂ <- tactic.mk_app ``preterm.func [`(my_pt)],
--    tactic.mk_app ``preterm.app [e₁, e₂]
-- -- #print foo_pt_term_reflected'

-- -- meta def bar : tactic expr :=
-- --   tactic.mk_app ``preterm.func [`(foo_mask)]

-- set_option trace.app_builder true

-- def foo_pt_term_reflected : preterm L_pt_func 0 := by (foo_pt_term_reflected' >>= tactic.exact)

-- -- #reduce foo_pt_term_reflected

-- end reflect_preterm3


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

! This file was ported from Lean 3 source module tactic.explode_widget
! leanprover-community/mathlib commit d13b3a4a392ea7273dfa4727dbd1892e26cfd518
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
import Mathbin.Tactic.Explode
import Mathbin.Tactic.InteractiveExpr

/-!
# `#explode_widget` command

Render a widget that displays an `#explode` proof, providing more
interactivity such as jumping to definitions and exploding constants
occurring in the exploded proofs.
-/

open Widget Tactic Tactic.Explode

unsafe instance widget.string_to_html {α} : Coe String (html α) :=
  ⟨fun s => s⟩
#align widget.string_to_html widget.string_to_html

namespace Tactic

namespace ExplodeWidget

open WidgetOverride.InteractiveExpression

open TaggedFormat

open Widget.Html Widget.Attr

/-- Redefine some of the style attributes for better formatting. -/
unsafe def get_block_attrs {γ} : sf → tactic (sf × List (attr γ))
  | sf.block i a => do
    let s : attr γ :=
      style [("display", "inline-block"), ("white-space", "pre-wrap"), ("vertical-align", "top")]
    let (a, rest) ← get_block_attrs a
    pure (a, s :: rest)
  | sf.highlight c a => do
    let (a, rest) ← get_block_attrs a
    pure (a, cn c :: rest)
  | a => pure (a, [])
#align tactic.explode_widget.get_block_attrs tactic.explode_widget.get_block_attrs

/-- Explode button for subsequent exploding. -/
unsafe def insert_explode {γ} : expr → tactic (List (html (action γ)))
  | expr.const n _ =>
    (do
        pure <|
            [h "button"
                [cn "pointer ba br3 mr1",
                  on_click fun _ =>
                    action.effect <| widget.effect.insert_text ("#explode_widget " ++ n),
                  attr.val "title" "explode"]
                ["💥"]]) <|>
      pure []
  | e => pure []
#align tactic.explode_widget.insert_explode tactic.explode_widget.insert_explode

/-- Render a subexpression as a list of html elements.
-/
unsafe def view {γ} (tooltip_component : tc subexpr (action γ))
    (click_address : Option Expr.Address) (select_address : Option Expr.Address) :
    subexpr → sf → tactic (List (html (action γ)))
  | ⟨ce, current_address⟩, sf.tag_expr ea e m => do
    let new_address := current_address ++ ea
    let select_attrs : List (attr (action γ)) :=
      if some new_address = select_address then [className "highlight"] else []
    let click_attrs : List (attr (action γ)) ←
      if some new_address = click_address then do
          let content ← tc.to_html tooltip_component (e, new_address)
          let efmt : String ← format.to_string <$> tactic.pp e
          let gd_btn ← goto_def_button e
          let epld_btn ← insert_explode e
          pure
              [tooltip <|
                  h "div" []
                    [h "div" [cn "fr"]
                        (gd_btn ++ epld_btn ++
                          [h "button"
                              [cn "pointer ba br3 mr1",
                                on_click fun _ => action.effect <| widget.effect.copy_text efmt,
                                attr.val "title" "copy expression to clipboard"]
                              ["📋"],
                            h "button"
                              [cn "pointer ba br3", on_click fun _ => action.on_close_tooltip,
                                attr.val "title" "close"]
                              ["×"]]),
                      content]]
        else pure []
    let (m, block_attrs) ← get_block_attrs m
    let as := [className "expr-boundary", key ea] ++ select_attrs ++ click_attrs ++ block_attrs
    let inner ← view (e, new_address) m
    pure [h "span" as inner]
  | ca, sf.compose x y => pure (· ++ ·) <*> view ca x <*> view ca y
  | ca, sf.of_string s =>
    pure
      [h "span"
          [on_mouse_enter fun _ => action.on_mouse_enter ca, on_click fun _ => action.on_click ca,
            key s]
          [html.of_string s]]
  | ca, b@(sf.block _ _) => do
    let (a, attrs) ← get_block_attrs b
    let inner ← view ca a
    pure [h "span" attrs inner]
  | ca, b@(sf.highlight _ _) => do
    let (a, attrs) ← get_block_attrs b
    let inner ← view ca a
    pure [h "span" attrs inner]
#align tactic.explode_widget.view tactic.explode_widget.view

/-- Make an interactive expression. -/
unsafe def mk {γ} (tooltip : tc subexpr γ) : tc expr γ :=
  let tooltip_comp :=
    (component.with_should_update fun x y : tactic_state × expr × Expr.Address => x.2.2 ≠ y.2.2) <|
      component.map_action action.on_tooltip_action tooltip
  (component.filter_map_action fun _ (a : Sum γ widget.effect) =>
      Sum.casesOn a some fun _ => none) <|
    (component.with_effects fun _ (a : Sum γ widget.effect) =>
        match a with
        | Sum.inl g => []
        | Sum.inr s => [s]) <|
      tc.mk_simple (action γ) (Option subexpr × Option subexpr) (fun e => pure <| (none, none))
        (fun e ⟨ca, sa⟩ act =>
          pure <|
            match act with
            | action.on_mouse_enter ⟨e, ea⟩ => ((ca, some (e, ea)), none)
            | action.on_mouse_leave_all => ((ca, none), none)
            | action.on_click ⟨e, ea⟩ =>
              if some (e, ea) = ca then ((none, sa), none) else ((some (e, ea), sa), none)
            | action.on_tooltip_action g => ((none, sa), some <| Sum.inl g)
            | action.on_close_tooltip => ((none, sa), none)
            | action.effect e => ((ca, sa), some <| Sum.inr <| e))
        fun e ⟨ca, sa⟩ => do
        let m ← sf.of_eformat <$> tactic.pp_tagged e
        let m := m.elim_part_apps
        let m := m.flatten
        let m := m.tag_expr [] e
        let v ← view tooltip_comp (Prod.snd <$> ca) (Prod.snd <$> sa) ⟨e, []⟩ m
        pure <|
            [h "span"
                  [className "expr", key e, on_mouse_leave fun _ => action.on_mouse_leave_all] <|
                v]
#align tactic.explode_widget.mk tactic.explode_widget.mk

/-- Render the implicit arguments for an expression in fancy, little pills. -/
unsafe def implicit_arg_list (tooltip : tc subexpr Empty) (e : expr) : tactic <| html Empty := do
  let fn ← mk tooltip <| expr.get_app_fn e
  let args ← List.mapM (mk tooltip) <| expr.get_app_args e
  pure <|
      h "div" []
        (h "span" [className "bg-blue br3 ma1 ph2 white"] [fn] ::
          List.map (fun a => h "span" [className "bg-gray br3 ma1 ph2 white"] [a]) args)
#align tactic.explode_widget.implicit_arg_list tactic.explode_widget.implicit_arg_list

/-- Component for the type tooltip.
-/
unsafe def type_tooltip : tc subexpr Empty :=
  tc.stateless fun ⟨e, ea⟩ => do
    let y ← tactic.infer_type e
    let y_comp ← mk type_tooltip y
    let implicit_args ← implicit_arg_list type_tooltip e
    pure
        [h "div" [style [("minWidth", "12rem")]]
            [h "div" [cn "pl1"] [y_comp], h "hr" [] [], implicit_args]]
#align tactic.explode_widget.type_tooltip tactic.explode_widget.type_tooltip

/-- Component that shows a type.
-/
unsafe def show_type_component : tc expr Empty :=
  tc.stateless fun x => do
    let y ← infer_type x
    let y_comp ← mk type_tooltip <| y
    pure y_comp
#align tactic.explode_widget.show_type_component tactic.explode_widget.show_type_component

/-- Component that shows a constant.
-/
unsafe def show_constant_component : tc expr Empty :=
  tc.stateless fun x => do
    let y_comp ← mk type_tooltip x
    pure y_comp
#align tactic.explode_widget.show_constant_component tactic.explode_widget.show_constant_component

/-- Search for an entry that has the specified line number.
-/
unsafe def lookup_lines : entries → Nat → entry
  | ⟨_, []⟩, n => ⟨default, 0, 0, Status.sintro, thm.string "", []⟩
  | ⟨rb, hd :: tl⟩, n => if hd.line = n then hd else lookup_lines ⟨rb, tl⟩ n
#align tactic.explode_widget.lookup_lines tactic.explode_widget.lookup_lines

/-- Render a row that shows a goal.
-/
unsafe def goal_row (e : expr) (show_expr := true) : tactic (List (html Empty)) := do
  let t ← explode_widget.show_type_component e
  return <|
      [h "td" [cn "ba bg-dark-green tc"] "Goal",
        h "td" [cn "ba tc"] (if show_expr then [html.of_name e, " : ", t] else t)]
#align tactic.explode_widget.goal_row tactic.explode_widget.goal_row

/-- Render a row that shows the ID of a goal.
-/
unsafe def id_row {γ} (l : Nat) : tactic (List (html γ)) :=
  return <| [h "td" [cn "ba bg-dark-green tc"] "ID", h "td" [cn "ba tc"] (toString l)]
#align tactic.explode_widget.id_row tactic.explode_widget.id_row

/-- Render a row that shows the rule or theorem being applied.
-/
unsafe def rule_row : thm → tactic (List (html Empty))
  | thm.expr e => do
    let t ← explode_widget.show_constant_component e
    return <| [h "td" [cn "ba bg-dark-green tc"] "Rule", h "td" [cn "ba tc"] t]
  | t => return <| [h "td" [cn "ba bg-dark-green tc"] "Rule", h "td" [cn "ba tc"] t.toString]
#align tactic.explode_widget.rule_row tactic.explode_widget.rule_row

/-- Render a row that contains the sub-proofs, i.e., the proofs of the
arguments.
-/
unsafe def proof_row {γ} (args : List (html γ)) : List (html γ) :=
  [h "td" [cn "ba bg-dark-green tc"] "Proofs",
    h "td" [cn "ba tc"]
      [h "details" [] <| h "summary" [attr.style [("color", "orange")]] "Details" :: args]]
#align tactic.explode_widget.proof_row tactic.explode_widget.proof_row

/-- Combine the goal row, id row, rule row and proof row to make them a table.
-/
unsafe def assemble_table {γ} (gr ir rr) : List (html γ) → html γ
  | [] => h "table" [cn "collapse"] [h "tbody" [] [h "tr" [] gr, h "tr" [] ir, h "tr" [] rr]]
  | pr =>
    h "table" [cn "collapse"]
      [h "tbody" [] [h "tr" [] gr, h "tr" [] ir, h "tr" [] rr, h "tr" [] pr]]
#align tactic.explode_widget.assemble_table tactic.explode_widget.assemble_table

/-- Render a table for a given entry.
-/
unsafe def assemble (es : entries) : entry → tactic (html Empty)
  | ⟨e, l, d, status.sintro, t, ref⟩ => do
    let gr ← goal_row e
    let ir ← id_row l
    let rr ← rule_row <| thm.string "Assumption"
    return <| assemble_table gr ir rr []
  | ⟨e, l, d, status.intro, t, ref⟩ => do
    let gr ← goal_row e
    let ir ← id_row l
    let rr ← rule_row <| thm.string "Assumption"
    return <| assemble_table gr ir rr []
  | ⟨e, l, d, st, t, ref⟩ => do
    let gr ← goal_row e false
    let ir ← id_row l
    let rr ← rule_row t
    let el : List entry := List.map (lookup_lines es) ref
    let ls ← Monad.mapM assemble el
    let pr := proof_row <| ls.intersperse (h "br" [] [])
    return <| assemble_table gr ir rr pr
#align tactic.explode_widget.assemble tactic.explode_widget.assemble

/-- Render a widget from given entries.
-/
unsafe def explode_component (es : entries) : tactic (html Empty) :=
  let concl := lookup_lines es (es.l.length - 1)
  assemble es concl
#align tactic.explode_widget.explode_component tactic.explode_widget.explode_component

/-- Explode a theorem and return entries.
-/
unsafe def explode_entries (n : Name) (hide_non_prop := true) : tactic entries := do
  let expr.const n _ ← resolve_name n |
    fail "cannot resolve name"
  let d ← get_decl n
  let v ←
    match d with
      | declaration.defn _ _ _ v _ _ => return v
      | declaration.thm _ _ _ v => return v.get
      | _ => fail "not a definition"
  let t ← pp d.type
  explode_expr v hide_non_prop
#align tactic.explode_widget.explode_entries tactic.explode_widget.explode_entries

end ExplodeWidget

open ExplodeWidget

/- ./././Mathport/Syntax/Translate/Tactic/Mathlib/Core.lean:38:34: unsupported: setup_tactic_parser -/
/-- User command of the explode widget.
-/
@[user_command]
unsafe def explode_widget_cmd (_ : parse <| tk "#explode_widget") : lean.parser Unit := do
  let ⟨li, co⟩ ← cur_pos
  let n ← ident
  let es ← explode_entries n
  let comp ←
    parser.of_tactic do
        let html ← explode_component es
        let c ← pure <| component.stateless fun _ => [html]
        pure <| component.ignore_props <| component.ignore_action <| c
  save_widget ⟨li, co - "#explode_widget".length - 1⟩ comp
  trace "successfully rendered widget"
  skip
#align tactic.explode_widget_cmd tactic.explode_widget_cmd

end Tactic


/-
Copyright (c) 2020 Mario Carneiro. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro
-/
import data.fintype.basic

/-!
# Derive handler for `fintype` instances

This file introduces a derive handler to automatically generate `fintype`
instances for structures and inductives.

## Implementation notes

To construct a fintype instance, we need 3 things:

  1. A list `l` of elements
  2. A proof that `l` has no duplicates
  3. A proof that every element in the type is in `l`

Now fintype is defined as a finset which enumerates all elements, so steps (1) and (2) are
bundled together. It is possible to use finset operations that remove duplicates to avoid the need
to prove (2), but this adds unnecessary functions to the constructed term, which makes it more
expensive to compute the list, and it also adds a dependence on decidable equality for the type,
which we want to avoid.

Because we will rely on fintype instances for constructor arguments, we can't actually build a list
directly, so (1) and (2) are necessarily somewhat intertwined. The inductive types we will be
proving instances for look something like this:

```
@[derive fintype]
inductive foo
| zero : foo
| one : bool → foo
| two : ∀ x : fin 3, bar x → foo
```

The list of elements that we generate is
```
{foo.zero}
∪ (finset.univ : bool).map (λ b, finset.one b)
∪ (finset.univ : Σ' x : fin 3, bar x).map (λ ⟨x, y⟩, finset.two x y)
```
except that instead of `∪`, that is `finset.union`, we use `finset.disj_union` which doesn't
require any deduplication, but does require a proof that the two parts of the union are disjoint.
We use `finset.cons` to append singletons like `foo.zero`.

The proofs of disjointness would be somewhat expensive since there are quadratically many of them,
so instead we use a "discriminant" function. Essentially, we define
```
def foo.enum : foo → ℕ
| foo.zero := 0
| (foo.one _) := 1
| (foo.two _ _) := 2
```
and now the existence of this function implies that foo.zero is not foo.two and so on because they
map to different natural numbers. We can prove that sets of natural numbers are mutually disjoint
more easily because they have a linear order: `0 < 1 < 2` so `0 ≠ 2`.

To package this argument up, we define `finset_above foo foo.enum n` to be a finset `s` together
with a proof that all elements `a ∈ s` have `n ≤ enum a`. Now we only have to prove that
`enum foo.zero = 0`, `enum (foo.one _) = 1`, etc. (linearly many proofs, all `rfl`) in order to
prove that all variants are mutually distinct.

We mirror the `finset.cons` and `finset.disj_union` functions into `finset_above.cons` and
`finset_above.union`, and this forms the main part of the finset construction.

This only handles distinguishing variants of a finset. Now we must enumerate the elements of a
variant, for example `{foo.one ff, foo.one tt}`, while at the same time proving that all these
elements have discriminant `1` in this case. To do that, we use the `finset_in` type, which
is a finset satisfying a property `P`, here `λ a, foo.enum a = 1`.

We could use `finset.bind` many times to construct the finset but it turns out to be somewhat
complicated to get good side goals for a naturally nodup version of `finset.bind` in the same way
as we did with `finset.cons` and `finset.union`. Instead, we tuple up all arguments into one type,
leveraging the `fintype` instance on `psigma`, and then define a map from this type to the
inductive type that untuples them and applies the constructor. The injectivity property of the
constructor ensures that this function is injective, so we can use `finset.map` to apply it. This
is the content of the constructor `finset_in.mk`.

That completes the proofs of (1) and (2). To prove (3), we perform one case analysis over the
inductive type, proving theorems like
```
foo.one a ∈ {foo.zero}
  ∪ (finset.univ : bool).map (λ b, finset.one b)
  ∪ (finset.univ : Σ' x : fin 3, bar x).map (λ ⟨x, y⟩, finset.two x y)
```
by seeking to the relevant disjunct and then supplying the constructor arguments. This part of the
proof is quadratic, but quite simple. (We could do it in `O(n log n)` if we used a balanced tree
for the unions.)

The tactics perform the following parts of this proof scheme:
* `mk_sigma` constructs the type `Γ` in `finset_in.mk`
* `mk_sigma_elim` constructs the function `f` in `finset_in.mk`
* `mk_sigma_elim_inj` proves that `f` is injective
* `mk_sigma_elim_eq` proves that `∀ a, enum (f a) = k`
* `mk_finset` constructs the finset `S = {foo.zero} ∪ ...` by recursion on the variants
* `mk_finset_total` constructs the proof `|- foo.zero ∈ S; |- foo.one a ∈ S; |- foo.two a b ∈ S`
  by recursion on the subgoals coming out of the initial `cases`
* `mk_fintype_instance` puts it all together to produce a proof of `fintype foo`.
  The construction of `foo.enum` is also done in this function.

-/

namespace derive_fintype

/-- A step in the construction of `finset.univ` for a finite inductive type.
We will set `enum` to the discriminant of the inductive type, so a `finset_above`
represents a finset that enumerates all elements in a tail of the constructor list. -/
def finset_above (α) (enum : α → ℕ) (n : ℕ) :=
{s : finset α // ∀ x ∈ s, n ≤ enum x}

/-- Construct a fintype instance from a completed `finset_above`. -/
def mk_fintype {α} (enum : α → ℕ) (s : finset_above α enum 0) (H : ∀ x, x ∈ s.1) :
  fintype α := ⟨s.1, H⟩

/-- This is the case for a simple variant (no arguments) in an inductive type. -/
def finset_above.cons {α} {enum : α → ℕ} (n)
  (a : α) (h : enum a = n) (s : finset_above α enum (n+1)) : finset_above α enum n :=
begin
  refine ⟨finset.cons a s.1 _, _⟩,
  { intro h',
    have := s.2 _ h', rw h at this,
    exact nat.not_succ_le_self n this },
  { intros x h', rcases finset.mem_cons.1 h' with rfl | h',
    { exact ge_of_eq h },
    { exact nat.le_of_succ_le (s.2 _ h') } }
end

theorem finset_above.mem_cons_self {α} {enum : α → ℕ} {n a h s} :
  a ∈ (@finset_above.cons α enum n a h s).1 := multiset.mem_cons_self _ _

theorem finset_above.mem_cons_of_mem {α} {enum : α → ℕ} {n a h s b} :
  b ∈ (s : finset_above _ _ _).1 → b ∈ (@finset_above.cons α enum n a h s).1 :=
multiset.mem_cons_of_mem

/-- The base case is when we run out of variants; we just put an empty finset at the end. -/
def finset_above.nil {α} {enum : α → ℕ} (n) : finset_above α enum n := ⟨∅, by rintro _ ⟨⟩⟩

instance (α enum n) : inhabited (finset_above α enum n) := ⟨finset_above.nil _⟩

/-- This is a finset covering a nontrivial variant (with one or more constructor arguments).
The property `P` here is `λ a, enum a = n` where `n` is the discriminant for the current
variant. -/
@[nolint has_nonempty_instance]
def finset_in {α} (P : α → Prop) := {s : finset α // ∀ x ∈ s, P x}

/-- To construct the finset, we use an injective map from the type `Γ`, which will be the
sigma over all constructor arguments. We use sigma instances and existing fintype instances
to prove that `Γ` is a fintype, and construct the function `f` that maps `⟨a, b, c, ...⟩`
to `C_n a b c ...` where `C_n` is the nth constructor, and `mem` asserts
`enum (C_n a b c ...) = n`. -/
def finset_in.mk {α} {P : α → Prop} (Γ) [fintype Γ]
  (f : Γ → α) (inj : function.injective f) (mem : ∀ x, P (f x)) : finset_in P :=
⟨finset.univ.map ⟨f, inj⟩,
 λ x h, by rcases finset.mem_map.1 h with ⟨x, _, rfl⟩; exact mem x⟩

theorem finset_in.mem_mk {α} {P : α → Prop} {Γ} {s : fintype Γ} {f : Γ → α} {inj mem a}
  (b) (H : f b = a) : a ∈ (@finset_in.mk α P Γ s f inj mem).1 :=
finset.mem_map.2 ⟨_, finset.mem_univ _, H⟩

/-- For nontrivial variants, we split the constructor list into a `finset_in` component for the
current constructor and a `finset_above` for the rest. -/
def finset_above.union {α} {enum : α → ℕ} (n)
  (s : finset_in (λ a, enum a = n)) (t : finset_above α enum (n+1)) : finset_above α enum n :=
begin
  refine ⟨finset.disj_union s.1 t.1 _, _⟩,
  { rw finset.disjoint_left,
    intros a hs ht,
    have := t.2 _ ht, rw s.2 _ hs at this,
    exact nat.not_succ_le_self n this },
  { intros x h', rcases finset.mem_disj_union.1 h' with h' | h',
    { exact ge_of_eq (s.2 _ h') },
    { exact nat.le_of_succ_le (t.2 _ h') } }
end

theorem finset_above.mem_union_left {α} {enum : α → ℕ} {n s t a}
  (H : a ∈ (s : finset_in _).1) : a ∈ (@finset_above.union α enum n s t).1 :=
multiset.mem_add.2 (or.inl H)

theorem finset_above.mem_union_right {α} {enum : α → ℕ} {n s t a}
  (H : a ∈ (t : finset_above _ _ _).1) : a ∈ (@finset_above.union α enum n s t).1 :=
multiset.mem_add.2 (or.inr H)

end derive_fintype

namespace tactic

open derive_fintype tactic expr

namespace derive_fintype

/-- Construct the term `Σ' (a:A) (b:B a) (c:C a b), unit` from
`Π (a:A) (b:B a), C a b → T` (the type of a constructor). -/
meta def mk_sigma : expr → tactic expr
| (expr.pi n bi d b) := do
  p ← mk_local' n bi d,
  e ← mk_sigma (expr.instantiate_var b p),
  tactic.mk_app ``psigma [d, bind_lambda e p]
| _ := pure `(unit)

/-- Prove the goal `(Σ' (a:A) (b:B a) (c:C a b), unit) → T`
(this is the function `f` in `finset_in.mk`) using recursive `psigma.elim`,
finishing with the constructor. The two arguments are the type of the constructor,
and the constructor term itself; as we recurse we add arguments
to the constructor application and destructure the pi type of the constructor. We return the number
of `psigma.elim` applications constructed, which is the number of constructor arguments. -/
meta def mk_sigma_elim : expr → expr → tactic ℕ
| (expr.pi n bi d b) c := do
  refine ``(@psigma.elim %%d _ _ _),
  i ← intro_fresh n,
  (+ 1) <$> mk_sigma_elim (expr.instantiate_var b i) (c i)
| _ c := do intro1, exact c $> 0

/-- Prove the goal `a, b |- f a = f b → g a = g b` where `f` is the function we constructed in
`mk_sigma_elim`, and `g` is some other term that gets built up and eventually closed by
reflexivity. Here `a` and `b` have sigma types so the proof approach is to case on `a` and `b`
until the goal reduces to `C_n a1 ... am = C_n b1 ... bm → ⟨a1, ..., am⟩ = ⟨b1, ..., bm⟩`, at which
point cases on the equality reduces the problem to reflexivity.

The arguments are the number `m` returned from `mk_sigma_elim`, and the hypotheses `a,b` that we
need to case on. -/
meta def mk_sigma_elim_inj : ℕ → expr → expr → tactic unit
| (m+1) x y := do
  [(_, [x1, x2])] ← cases x,
  [(_, [y1, y2])] ← cases y,
  mk_sigma_elim_inj m x2 y2
| 0 x y := do
  cases x, cases y,
  is ← intro1 >>= injection,
  is.mmap' cases,
  reflexivity

/-- Prove the goal `a |- enum (f a) = n`, where `f` is the function constructed in `mk_sigma_elim`,
and `enum` is a function that reduces to `n` on the constructor `C_n`. Here we just have to case on
`a` `m` times, and then `reflexivity` finishes the proof. -/
meta def mk_sigma_elim_eq : ℕ → expr → tactic unit
| (n+1) x := do
  [(_, [x1, x2])] ← cases x,
  mk_sigma_elim_eq n x2
| 0 x := reflexivity

/-- Prove the goal `|- finset_above T enum k`, where `T` is the inductive type and `enum` is the
discriminant function. The arguments are `args`, the parameters to the inductive type (and all
constructors), `k`, the index of the current variant, and `cs`, the list of constructor names.
This uses `finset_above.cons` for basic variants and `finset_above.union` for variants with
arguments, using the auxiliary functions `mk_sigma`, `mk_sigma_elim`, `mk_sigma_elim_inj`,
`mk_sigma_elim_eq` to close subgoals. -/
meta def mk_finset (ls : list level) (args : list expr) : ℕ → list name → tactic unit
| k (c::cs) := do
  let e := (expr.const c ls).mk_app args,
  t ← infer_type e,
  if is_pi t then do
    to_expr ``(finset_above.union %%(reflect k)) tt ff >>=
      (λ c, apply c {new_goals := new_goals.all}),
    Γ ← mk_sigma t,
    to_expr ``(finset_in.mk %%Γ) tt ff >>= (λ c, apply c {new_goals := new_goals.all}),
    n ← mk_sigma_elim t e,
    intro1 >>= (λ x, intro1 >>= mk_sigma_elim_inj n x),
    intro1 >>= mk_sigma_elim_eq n,
    mk_finset (k+1) cs
  else do
    c ← to_expr ``(finset_above.cons %%(reflect k) %%e) tt ff,
    apply c {new_goals := new_goals.all}, reflexivity,
    mk_finset (k+1) cs
| k [] := applyc ``finset_above.nil

/-- Prove the goal `|- Σ' (a:A) (b: B a) (c:C a b), unit` given a list of terms `a, b, c`. -/
meta def mk_sigma_mem : list expr → tactic unit
| (x::xs) := fconstructor >> exact x >> mk_sigma_mem xs
| [] := fconstructor $> ()

/-- This function is called to prove `a : T |- a ∈ S.1` where `S` is the `finset_above` constructed
by `mk_finset`, after the initial cases on `a : T`, producing a list of subgoals. For each case,
we have to navigate past all the variants that don't apply (which is what the `tac` input tactic
does), and then call either `finset_above.mem_cons_self` for trivial variants or
`finset_above.mem_union_left` and `finset_in.mem_mk` for nontrivial variants. Either way the proof
is quite simple. -/
meta def mk_finset_total : tactic unit → list (name × list expr) → tactic unit
| tac [] := done
| tac ((_, xs) :: gs) := do
  tac,
  b ← succeeds (applyc ``finset_above.mem_cons_self),
  if b then
    mk_finset_total (tac >> applyc ``finset_above.mem_cons_of_mem) gs
  else do
    applyc ``finset_above.mem_union_left,
    applyc ``finset_in.mem_mk {new_goals := new_goals.all},
    mk_sigma_mem xs,
    reflexivity,
    mk_finset_total (tac >> applyc ``finset_above.mem_union_right) gs

end derive_fintype

open tactic.derive_fintype

/-- Proves `|- fintype T` where `T` is a non-recursive inductive type with no indices,
where all arguments to all constructors are fintypes. -/
meta def mk_fintype_instance : tactic unit :=
do
  intros,
  `(fintype %%e) ← target >>= whnf,
  (const I ls, args) ← pure (get_app_fn_args e),
  env ← get_env,
  let cs := env.constructors_of I,
  guard (env.inductive_num_indices I = 0) <|>
    fail "@[derive fintype]: inductive indices are not supported",
  guard (¬ env.is_recursive I) <|>
    fail ("@[derive fintype]: recursive inductive types are " ++
          "not supported (they are also usually infinite)"),
  applyc ``mk_fintype {new_goals := new_goals.all},
  intro1 >>= cases >>= (λ gs,
    gs.enum.mmap' $ λ ⟨i, _⟩, exact (reflect i)),
  mk_finset ls args 0 cs,
  intro1 >>= cases >>= mk_finset_total skip

/--
Tries to derive a `fintype` instance for inductives and structures.

For example:
```
@[derive fintype]
inductive foo (n m : ℕ)
| zero : foo
| one : bool → foo
| two : fin n → fin m → foo
```
Here, `@[derive fintype]` adds the instance `foo.fintype`. The underlying finset
definitionally unfolds to a list that enumerates the elements of the inductive in
lexicographic order.

If the structure/inductive has a type parameter `α`, then the generated instance will have an
argument `fintype α`, even if it is not used.  (This is due to the implementation using
`instance_derive_handler`.)
-/
@[derive_handler] meta def fintype_instance : derive_handler :=
instance_derive_handler ``fintype mk_fintype_instance

end tactic


import system.io
import query_api
import parse

open widget tactic
section json
meta def list.lookup_prod {α β} : (list (α × β)) → (α → bool) → option β
| [] _ := none
| (⟨a,b⟩::xs) p := if p a then pure b else xs.lookup_prod p

open except

meta def json.lookup : json → string → except string json
| (json.object kvs) str :=
  match kvs.lookup_prod $ λ k, k = str with
  | some v := except.ok v
  | none := except.error ("no key " ++ str)
  end
| _ _ := except.error "not an object"

meta def json.as_string : json → except string string
  | (json.of_string s) := except.ok s
  | _ := except.error "not a string"

meta def json.as_array : json → except string (list json)
  | (json.array xs) := ok xs
  | _ := error "not an array"

meta def except.liftOption {α}: option α → except string α
  | none := except.error "option was none"
  | (some a ) := except.ok a

end json

meta def text_of_return_json (parsed : json) : except string string := do
  choices_json ← json.lookup parsed "choices",
  head :: _ ← json.as_array choices_json | except.error "empty array",
  t ← json.lookup head "text",
  s ← json.as_string t,
  return s

meta def run_except {α} : except string α → io α
  | (except.ok a) := pure a
  | (except.error e) := io.fail e

@[derive inhabited]
meta structure bubble :=
  (body : string) -- [todo] add formatting etc
  (user : string)

meta def chat_props := unit

meta structure chat_state : Type :=
  (bubbles : list bubble)
  (current_text : string)

/- @zhangir: write your code for getting response from codex here :-) -/
meta def get_response : chat_state → io string
| state := do {
    bubbles@(head :: tail) ← pure $ state.bubbles | io.fail "no chat yet",
    let statement :=
      match tail with
      | [] := head.body
      | tail := (list.reverse tail).head.body
      end,
    let rest_of_context :=
      match tail with
      | [] := ""
      | tail := (string.intercalate "\n" $  list.map bubble.body $ list.tail $ list.reverse $ bubbles) ++ " Try again:\ntheorem"
        -- want this python command: "\n".join([x.body for x in state.bubbles].reverse[1:])
      end,
    let prompt := prompt_of_nl_statement statement few_shot_prompt ++ rest_of_context,
    --io.put_str_ln (prompt ++ "<endoftext>"),
    return_json ← get_completion_of_request {prompt:=prompt},
    (some maybe_return_parsed) ← pure (json.parse return_json) | io.fail "not json",
    t : string ← run_except $ text_of_return_json maybe_return_parsed,
    return (t ++ " :=")
  }

/-- Use when testing formatting etc to avoid having to call api. -/
meta def get_response_dummy : chat_state → io string
  | state := do {
      return "meta def unsafe_perform_io {α} (m : io α) : except io.error α :=
match (cast undefined m : unit → sum α io.error) () with
| sum.inl a := except.ok a
| sum.inr err := except.error err
end"
    }

-- use some @gebner magic here
meta def unsafe_perform_io {α} (m : io α) : except io.error α :=
match (cast undefined m : unit → sum α io.error) () with
| sum.inl a := except.ok a
| sum.inr err := except.error err
end

meta def unsafe_get_response (input : chat_state) : string :=
  match unsafe_perform_io (get_response input) with
  | except.ok a := a
  | except.error e := "error"
  end

meta inductive chat_action
  | submit
  | demo_text
  | text_change (s : string)
  | copy_to_comment (s : string)
  | copy_to_script (s : string)
  | clear

meta def code_content (code : string) : html chat_action :=
  h "div" [] [
    h "code" [className "font-code", attr.style [("white-space", "break-spaces")]] [
      code
    ],
    h "div" [] [
      button "paste" (chat_action.copy_to_script ("theorem " ++ code))
    ]
  ]

meta inductive nlrun
  | math (s : string) : nlrun
  | text (s : string) : nlrun

meta def nlrun.to_html {α : Type} : nlrun → html α
  | (nlrun.math s) := html.element "InlineMath" [attr.val "math" s] []
  | (nlrun.text s) := html.of_string s

meta def to_nlrun_aux : list string → list nlrun
  | (text :: latex :: rest) := (nlrun.text text) :: (nlrun.math latex) :: to_nlrun_aux rest
  | [] := []
  | [text] := [nlrun.text text]

meta def to_nlrun (s : string) : list nlrun :=
  to_nlrun_aux $ string.split_on s '$'

meta def nl_content (s : string) : html chat_action :=
  let nlruns := to_nlrun s in
  h "div" [className "f6"] $ nlruns.map nlrun.to_html

meta def chat_view (props : chat_props) (state : chat_state) : list (html chat_action) :=
  [h "div" [className "f6"] [
    h "div" [className "flex flex-column"] (
      state.bubbles.reverse.map (λ bubble,
        h "div" [className "pa2 ma2 bg-lightest-blue"] [
          h "div" [className "mr2"] [bubble.user, ": "],
          if bubble.user ≠ "self" then
            code_content bubble.body
          else
            nl_content bubble.body
        ]
      )
    ),
    h "div" [] [
      textbox state.current_text chat_action.text_change,
      button "submit" chat_action.submit, -- [todo] how to get it to trigger on enter?
      button "demo" chat_action.demo_text,
      button "clear" chat_action.clear
    ]
  ]]

meta def push_bubble (b : bubble) (s : chat_state) : chat_state := {bubbles := b :: s.bubbles, ..s}

#check tactic.unsafe_run_io
/-- runs the lean parser on the response, if it's good then returns an error message, otherwise returns none.
  I am not proud of this method, this should be considered bad Lean practice.
-/
meta def unsafe_parse_result (response : string) : option string := do
  empty_tactic_state ← except.to_option $ unsafe_perform_io (io.run_tactic tactic.read),
  let t := lean.parser.run_with_input parse_decl response empty_tactic_state,
  match t with
  | result.success a s := none
  | result.exception (some msg) pos s := some (to_string $ msg())
  | result.exception _ _ _ := some "exception with no message!"
  end

meta def chat_update (props : chat_props)  : chat_state → chat_action → (chat_state × option effect)
  | state (chat_action.submit) :=
    let text := state.current_text,
        state := {current_text := "", ..state},
        state := push_bubble {body := text, user := "self"} state,
        response := unsafe_get_response state,
        state := push_bubble {body := response, user := "codex"} state,
        state := push_bubble {
          body := match unsafe_parse_result ("def " ++ response) with
                  | (some msg) := "error:\n" ++ msg
                  | none := "check passes!"
                  end,
          user := "lean"} state
        in
    (state, none)
  | state (chat_action.text_change s) := ({current_text := s, ..state}, none)
  | state (chat_action.copy_to_comment str) := (state, some $ effect.insert_text $ "/--\n" ++ str ++ "\n-/")
  | state (chat_action.copy_to_script str) := (state, some $ effect.insert_text str)
  | state chat_action.demo_text :=
    -- @zhangir add your example here to avoid having to paste it in every time!
    let state := {current_text := "If $x$ is an element of infinite order in $G$, prove that the elements $x^n$, $n\\in\\mathbb{Z}$ are all distinct.", ..state} in
    chat_update state chat_action.submit
  | state (chat_action.clear) := ({current_text:="", bubbles:=[]}, none)

meta def chat_init (props : chat_props) (old_state : option chat_state) : chat_state :=
  let s : chat_state := ({bubbles := [], current_text := ""} <| old_state) in
  s

meta def chat_widget {π α : Type} : component π α :=
component.ignore_props
$ component.ignore_action
$ component.with_effects (λ _ action, [action])
$ component.stateful chat_action _ chat_init  chat_update chat_view

-- @zhangir put your cursor on the #html token!
#html chat_widget ()