import system.io
import utils
open tactic.unsafe

universes u v w

-- tools and help functions

-- def option.mmap {m : Type u → Type v} [monad m] {α : Type w} {β : Type u} (f : α → m β) : option α → m (option β)
-- | none       := return none
-- | (some x) := do x' ← f x, return (some x')

def list.last_option {α : Type u}: list α → option α
| []        := none
| [a]       := some a
| (a::b::l) := list.last_option (b::l)

meta def expr.local_uniq_name_option : expr → option name
| (expr.local_const n m bi t) := some n
| e                      := none

meta def expr.mvar_uniq_name_option : expr → option name
| (expr.mvar n ppn t) := some n
| e                   := none

-- set the tactic state
meta def set_state (new_state: tactic_state): tactic unit :=
  -- this is in mathlib but easier to recreate
  λ _, interaction_monad.result.success () new_state

-- types for encoding tactic state information

meta structure mvar_decl :=
(unique_name : name)
(pp_name : name)
(expr_type : expr)
(local_cxt : list expr)
(type : option expr)
(assignment : option expr)

/- There already is a local_decl type, but
this is some more informtion for understanding
the type better.-/
meta structure local_decl2 :=
(unique_name : name)
(pp_name : name)
(expr_type : expr)
(bi : binder_info)
(type : option expr)
(prev : option name)
(frozen : bool)
(ld : option local_decl)

meta structure univ_mvar_decl :=
(unique_name : name)
(assignment : option level)

meta inductive context.decl
| mvar_decl (mv : mvar_decl) : context.decl
| univ_mvar_decl (mv : univ_mvar_decl) : context.decl
| local_decl (loc : local_decl2) : context.decl

meta structure context :=
-- in order of dependecies
(decls : list context.decl)
(names : name_set)

meta structure tactic_state_data :=
(decls : list context.decl)
(goals : list (name × tactic.tag))

-- meta instance : has_to_format tactic.tag := sorry
-- meta instance : has_to_format context.decl := sorry
-- -- meta instance : has_to_tactic_format context.decl := sorry
-- meta instance foo' : has_to_format $ list (name × tactic.tag) := by apply_instance
-- meta instance foo : has_to_format (list decl) := by apply_instance

-- meta instance : has_to_tactic_format tactic_state_data :=
-- ⟨λ ⟨decls, goals⟩, pure $ format!"tactic_state_data.mk \n\t decls := {decls} \n\t goals := {goals}"⟩

attribute [derive [has_to_format]] mvar_decl univ_mvar_decl binder_info
attribute [derive [has_to_format]] local_decl
attribute [derive [has_to_format]] local_decl2
attribute [derive [has_to_format]] context.decl
meta instance : has_to_format tactic.tag := (by apply_instance : has_to_format (list name))

attribute [derive [has_to_format]] tactic_state_data

section instances

attribute [derive [has_to_tactic_json]] mvar_decl
attribute [derive [has_to_tactic_json]] univ_mvar_decl
attribute [derive [has_to_tactic_json]] local_decl
attribute [derive [has_to_tactic_json]] local_decl
attribute [derive [has_to_tactic_json]] local_decl2

attribute [derive has_to_tactic_json] context.decl

meta instance : has_to_tactic_json tactic.tag :=
⟨by mk_to_tactic_json name.anonymous⟩

meta instance : has_from_json tactic.tag :=
has_from_json_list

attribute [derive has_to_tactic_json] tactic_state_data

meta instance : has_from_json mvar_decl :=
⟨λ msg, match msg with
  | (json.array $ [c, json.array [un, pp, ety, cxt, ty, assn]]) := do
    (c_nm : name) ← has_from_json_name_aux c,
    if c_nm = `mvar_decl.mk then do
      mvar_decl.mk <$> (has_from_json.from_json un) <*> (has_from_json.from_json pp)
        <*> has_from_json.from_json ety <*> has_from_json.from_json cxt
          <*> has_from_json.from_json ty <*> has_from_json.from_json assn
    else
    tactic.fail format!"[has_from_json_mvar_decl] unexpected: {msg}"
  | exc := tactic.fail format!"[has_from_json_mvar_decl] unexpected: {exc}"
  end
⟩

meta instance : has_from_json univ_mvar_decl :=
⟨λ msg, match msg with
  | (json.array $ [c, json.array [un, pp]]) := do
    (c_nm : name) ← has_from_json_name_aux c,
    if c_nm = `univ_mvar_decl.mk then do
      univ_mvar_decl.mk <$> (has_from_json.from_json un) <*> (has_from_json.from_json pp)
    else
    tactic.fail format!"[has_from_json_univ_mvar_decl] unexpected: {msg}"
  | exc := tactic.fail format!"[has_from_json_univ_mvar_decl] unexpected: {exc}"
  end
⟩

-- meta instance : has_from_json univ_mvar_decl := sorry

-- attribute [derive [has_to_tactic_json]] bool

run_cmd (has_to_tactic_json.to_tactic_json tt >>= (has_from_json.from_json : json → tactic bool))

meta instance : has_from_json local_decl :=
⟨λ msg, match msg with
  | (json.array $ [c, json.array [un, pp, ty, val, bi, idx]]) := do
    (c_nm : name) ← has_from_json_name_aux c,
    if c_nm = `local_decl.mk then do
      local_decl.mk <$>
        (has_from_json.from_json un) <*>
          (has_from_json.from_json pp) <*>
                has_from_json.from_json ty <*>
                  has_from_json.from_json val <*>
                    has_from_json.from_json bi <*>
                      has_from_json.from_json idx
    else
    tactic.fail format!"[has_from_json_local_decl] unexpected: {msg}"
  | exc := tactic.fail format!"[has_from_json_local_decl] unexpected: {exc}"
  end
⟩

meta instance : has_from_json local_decl2 :=
⟨λ msg, match msg with
  | (json.array $ [c, json.array [un, pp, ety, bi, ty, prev, frozen, ld]]) := do
    (c_nm : name) ← has_from_json_name_aux c,
    if c_nm = `local_decl2.mk then do
      local_decl2.mk <$>
        (has_from_json.from_json un) <*>
          (has_from_json.from_json pp) <*>
            has_from_json.from_json ety <*>
              has_from_json.from_json bi <*>
                has_from_json.from_json ty <*>
                  has_from_json.from_json prev <*>
                    has_from_json.from_json frozen <*>
                      has_from_json.from_json ld
    else
    tactic.fail format!"[has_from_json_local_decl2] unexpected: {msg}"
  | exc := tactic.fail format!"[has_from_json_local_decl2] unexpected: {exc}"
  end
⟩

meta instance : has_from_json context.decl :=
let ⟨fn₁⟩ := (by apply_instance : has_from_json mvar_decl) in
let ⟨fn₂⟩ := (by apply_instance : has_from_json univ_mvar_decl) in
let ⟨fn₃⟩ := (by apply_instance : has_from_json local_decl2) in
⟨λ msg, match msg with
  | (json.array $ [c, json.array args]) := do
    (c_nm : name) ← has_from_json_name_aux c,
    tactic.trace format!"[has_from_json_context.decl] c_nm: {c_nm}",
    if c_nm = `context.decl.mvar_decl then context.decl.mvar_decl <$> fn₁ args.head else
    if c_nm = `context.decl.univ_mvar_decl then context.decl.univ_mvar_decl <$> fn₂ args.head else
    if c_nm = `context.decl.local_decl then context.decl.local_decl <$> fn₃ args.head else
    tactic.fail format!"[has_from_json_context.decl] unexpected: {msg}"
  | exc := tactic.fail format!"[has_from_json_context.decl] unexpected: {exc}"
  end
⟩

meta instance : has_from_json (list context.decl) := has_from_json_list
meta instance has_from_json_list_name_tactic_tag : has_from_json (list (name × tactic.tag)) := has_from_json_list

meta instance : has_from_json tactic_state_data :=
let ⟨fn₁⟩ := (by apply_instance : has_from_json (list context.decl)) in
let ⟨fn₂⟩ := (by apply_instance : has_from_json (list (name × tactic.tag))) in
⟨λ msg, match msg with
  | (json.array $ [c, json.array [decls_msg, goals_msg]]) := do
    (c_nm : name) ← has_from_json_name_aux c,
    tactic.trace format!"[has_from_json_tactic_state_data] c_nm: {c_nm}",
    if c_nm = `tactic_state_data.mk then tactic_state_data.mk <$> fn₁ decls_msg <*> fn₂ goals_msg else
    tactic.fail format!"[has_from_json_tactic_state_data] unexpected: {msg}"
  | exc := tactic.fail format!"[has_from_json_tactic_state_data] unexpected: {exc}"
  end
⟩

end instances

-- convience functions and instances

meta instance mvar_decl_has_to_string : has_to_format mvar_decl :=
⟨ λ d, format! "{{mvar_decl .\nunique_name := {d.unique_name},\npp_name := {d.pp_name},\nexpr_type := {d.expr_type},\nlocal_cxt := {d.local_cxt},\ntype := {d.type},\nassignment := {d.assignment},\n}" ⟩

meta instance univ_mvar_decl_has_to_string : has_to_format univ_mvar_decl :=
⟨ λ d, format! "{{univ_mvar_decl .\nunique_name := {d.unique_name},\nassignment := {d.assignment},\n}" ⟩

meta instance local_decl_has_to_string : has_to_format local_decl :=
⟨ λ d, format! "{{local_decl .\nunique_name := {d.unique_name},\npp_name := {d.pp_name},\ntype := {d.type},\nvalue := {d.value},\nbi := {repr d.bi},\nidx := {d.idx},\n}" ⟩

meta instance local_decl2_has_to_string : has_to_format local_decl2 :=
⟨ λ d, format! "{{local_decl2 .\nunique_name := {d.unique_name},\npp_name := {d.pp_name},\nexpr_type := {d.expr_type},\nbi := {repr d.bi},\ntype := {d.type},\nprev := {d.prev}\n},\nfrozen := {d.frozen},\nld := {d.ld}" ⟩

meta def context.decl.unique_name : context.decl -> name
| (context.decl.mvar_decl d) := d.unique_name
| (context.decl.univ_mvar_decl d) := d.unique_name
| (context.decl.local_decl d) := d.unique_name

meta instance context_decl_has_to_string : has_to_format context.decl :=
⟨ λ d, match d with
| context.decl.mvar_decl d := format! "{d}"
| context.decl.univ_mvar_decl d := format! "{d}"
| context.decl.local_decl d := format! "{d}"
end ⟩

meta instance context_has_to_string : has_to_format context :=
⟨ λ cxt, format! "{cxt.decls}" ⟩

-- constructors

meta def context.empty : context :=
{ decls := [], names := mk_name_set }

meta def context.mk1 (d : context.decl) : context :=
{ decls := [d], names := name_set.of_list [d.unique_name]}

meta def context.append (cxt1 : context) (cxt2 : context) : context :=
{ decls := cxt1.decls ++ (cxt2.decls.filter (λ d, ¬ (cxt1.names.contains d.unique_name))),
  names := cxt1.names.fold cxt2.names $ λ n ns, ns.insert n
}

meta instance context.has_append : has_append context := ⟨ context.append ⟩

/- Get univ metavariables level expression tree.-/
meta def context.process_level : level -> tactic context
| level.zero := return context.empty
| (level.succ lvl) := context.process_level lvl
| (level.max lvl1 lvl2) := do
  cxt1 <- context.process_level lvl1,
  cxt2 <- context.process_level lvl2,
  return (cxt1 ++ cxt2)
| (level.imax lvl1 lvl2) := do
  cxt1 <- context.process_level lvl1,
  cxt2 <- context.process_level lvl2,
  return (cxt1 ++ cxt2)
| (level.param _) := return context.empty
| lvl@(level.mvar nm) := do
  ass <- optional (tactic.get_univ_assignment lvl),
  let univ_decl := context.decl.univ_mvar_decl {
    unique_name := nm,
    assignment := ass
  },
  return (context.mk1 univ_decl)

def find_prev {α : Type} [decidable_eq α] (a : α) : list α -> option α
| [] := none
| [b] := none
| (b :: c :: ls) := if c = a then some b else find_prev (c :: ls)

/- Get metavariables and local constants inside an expression tree, follow recursively. -/
meta def context.process_expr : expr -> local_context -> tactic context
| (expr.var _) _ := return context.empty
| (expr.sort lvl) _ := context.process_level lvl
| (expr.const _ lvls) _ := do
  cxts <- lvls.mmap context.process_level,
  let cxt := cxts.foldl context.append context.empty,
  return cxt
| mv@(expr.mvar unique_nm pp_nm tp) _ := do
  lcxt <- type_context.run $ type_context.get_context mv,
  let local_cxt := lcxt.to_list,
  cxts <- local_cxt.mmap (λ e, e.unfold_macros >>= flip context.process_expr lcxt),
  let cxt := cxts.foldl context.append context.empty,
  tp_cxt <- tp.unfold_macros >>= flip context.process_expr lcxt,
  mv_type <- optional (tactic.infer_type mv),
  tp_cxt2 <- match mv_type with
  | (some e) := e.unfold_macros >>= flip context.process_expr lcxt
  | none := return context.empty
  end,
  assignment <- optional (tactic.get_assignment mv),
  ass_cxt <- match assignment with
  | (some e) := e.unfold_macros >>= flip context.process_expr lcxt
  | none := return context.empty
  end,
  let mv_dec := context.decl.mvar_decl {
    unique_name := unique_nm,
    pp_name := pp_nm,
    expr_type := tp,
    local_cxt := local_cxt,
    type := mv_type,
    assignment := assignment
  },
  return $ cxt ++ tp_cxt ++ tp_cxt2 ++ ass_cxt ++ (context.mk1 mv_dec)
| lconst@(expr.local_const unique_nm pp_nm bi tp) lcxt := do
  tp_cxt <- tp.unfold_macros >>= flip context.process_expr lcxt,
  loc_type <- optional (tactic.infer_type lconst),
  tp_cxt2 <- match loc_type with
  | (some e) := e.unfold_macros >>= flip context.process_expr lcxt
  | none := return context.empty
  end,
  let ld := lcxt.get_local_decl unique_nm,
  tp_cxt3 <- match ld with
  | (some ld) := ld.type.unfold_macros >>= flip context.process_expr  lcxt
  | none := return context.empty
  end,
  value_cxt <- match ld with
  | (some ld) := match ld.value with
    | (some e) := e.unfold_macros >>= flip context.process_expr lcxt
    | none := return context.empty
    end
  | none := return context.empty
  end,
  let (prev : option expr) := find_prev lconst lcxt.to_list,
  let prev_id := match prev with
  | some (expr.local_const id _ _ _) := some id
  | _ := none
  end,
  frozen_instances_opt <- tactic.frozen_local_instances,
  let frozen := match frozen_instances_opt with
  | none := ff
  | some frozen_instances := frozen_instances.any (λ e, e.local_uniq_name_option = some unique_nm)
  end,
  let loc_dec := context.decl.local_decl {
    unique_name := unique_nm,
    pp_name := pp_nm,
    expr_type := tp,
    bi := bi,
    type := loc_type,
    prev := prev_id,
    frozen := frozen,
    ld := lcxt.get_local_decl unique_nm,
  },
  return $ tp_cxt ++ tp_cxt2 ++ tp_cxt3 ++ value_cxt ++ (context.mk1 loc_dec)
| (expr.app expr1 expr2) lcxt := do
  cxt1 <- expr1.unfold_macros >>= flip context.process_expr lcxt,
  cxt2 <- expr2.unfold_macros >>= flip context.process_expr lcxt,
  return (cxt1 ++ cxt2)
| (expr.lam _ _ expr1 expr2) lcxt := do
  cxt1 <- expr1.unfold_macros >>= flip context.process_expr lcxt,
  cxt2 <- expr2.unfold_macros >>= flip context.process_expr lcxt,
  return (cxt1 ++ cxt2)
| (expr.pi _ _ expr1 expr2) lcxt := do
  cxt1 <- expr1.unfold_macros >>= flip context.process_expr lcxt,
  cxt2 <- expr2.unfold_macros >>= flip context.process_expr lcxt,
  return (cxt1 ++ cxt2)
| (expr.elet _ expr1 expr2 expr3) lcxt := do
  cxt1 <- expr1.unfold_macros >>= flip context.process_expr lcxt,
  cxt2 <- expr2.unfold_macros >>= flip context.process_expr lcxt,
  cxt3 <- expr3.unfold_macros >>= flip context.process_expr lcxt,
  return (cxt1 ++ cxt2 ++ cxt3)
| (expr.macro md deps) _ := tactic.fail format!"[process_expr] can't handle macro {expr.macro_def_name md}"

meta def context.get : tactic context := do
  lcxt <- type_context.run $ type_context.get_local_context,
  mvs <- tactic.get_goals,
  cxts <- mvs.mmap (λ e, e.unfold_macros >>= flip context.process_expr lcxt),
  let cxt := cxts.foldl context.append context.empty,
  return cxt

meta def tactic_state_data.get : tactic tactic_state_data := do
  cxt <- context.get,
  gs <- tactic.get_goals,
  goals <- gs.mmap $ λ g, do {
    nm <- g.mvar_uniq_name_option,
    tag <- tactic.get_tag g,
    return (nm, tag)
  },
  return {
    decls := cxt.decls,
    goals := goals
  }

-- tracing code for debugging

meta def trace_context : tactic unit := do
  -- cxt <- context.get,
  cxt ← tactic_state_data.get,
  has_to_tactic_json.to_tactic_json cxt >>= tactic.trace

-- rebuilding the context

meta def swap_univ_mvs (nm_map : name_map context.decl) : level → tactic level
| (level.mvar nm) := do {
  d <- nm_map.find nm,
  nm' <- match d with
  | (context.decl.univ_mvar_decl dd) := return dd.unique_name
  | _ := tactic.failed
  end,
  return $ level.mvar nm'
}
| (level.max lvl1 lvl2) := do {
  lvl1' <- swap_univ_mvs lvl1,
  lvl2' <- swap_univ_mvs lvl2,
  return $ level.max lvl1' lvl2'
}
| (level.imax lvl1 lvl2) := do {
  lvl1' <- swap_univ_mvs lvl1,
  lvl2' <- swap_univ_mvs lvl2,
  return $ level.imax lvl1' lvl2'
}
| (level.succ lvl) := do {
  lvl' <- swap_univ_mvs lvl,
  return $ level.succ lvl'
}
| lvl := return lvl  --level.zero and level.param

meta def swap_mvs (nm_map : name_map context.decl) : expr -> tactic expr
| (expr.mvar unique_nm pp_nm tp) := do {
  d <- nm_map.find unique_nm,
  (unique_nm', tp') <- match d with
  | (context.decl.mvar_decl dd) := return (dd.unique_name, dd.expr_type)
  | _ := tactic.failed
  end,
  return $ expr.mvar unique_nm pp_nm tp'
}
| (expr.local_const unique_nm pp_nm bi tp) := do {
  d <- nm_map.find unique_nm,
  (unique_nm', tp') <- match d with
  | (context.decl.local_decl dd) := return (dd.unique_name, dd.expr_type)
  | _ := tactic.failed
  end,
  return $ expr.local_const unique_nm' pp_nm bi tp'
}
| e@(expr.var _) := return e
| (expr.sort lvl) := do {
  lvl' <- swap_univ_mvs nm_map lvl,
  return $ expr.sort lvl'
}
| (expr.const nm lvls) := do {
  lvls' <- lvls.mmap (swap_univ_mvs nm_map),
  return $ expr.const nm lvls'
}
| (expr.app expr1 expr2) := do {
  expr1' <- swap_mvs expr1.unfold_string_macros.erase_annotations,
  expr2' <- swap_mvs expr2.unfold_string_macros.erase_annotations,
  return $ expr.app expr1' expr2'
}
| (expr.lam nm bi expr1 expr2) := do {
  expr1' <- swap_mvs expr1.unfold_string_macros.erase_annotations,
  expr2' <- swap_mvs expr2.unfold_string_macros.erase_annotations,
  return $ expr.lam nm bi expr1' expr2'
}
| (expr.pi nm bi expr1 expr2) := do {
  expr1' <- swap_mvs expr1.unfold_string_macros.erase_annotations,
  expr2' <- swap_mvs expr2.unfold_string_macros.erase_annotations,
  return $ expr.pi nm bi expr1' expr2'
}
| (expr.elet nm expr1 expr2 expr3) := do {
  expr1' <- swap_mvs expr1.unfold_string_macros.erase_annotations,
  expr2' <- swap_mvs expr2.unfold_string_macros.erase_annotations,
  expr3' <- swap_mvs expr3.unfold_string_macros.erase_annotations,
  return $ expr.elet nm expr1' expr2' expr3'
}
| (expr.macro _ _) := tactic.fail "[swap_mvs] can't handle macros yet"

/- A better constructor for locals which covers
frozen status and assignments. -/
meta def local_context.mk_local2 (pretty_name : name) (type : expr) (bi : binder_info) (frozen : bool) (assignment : option expr) (lcxt : local_context) : tactic (expr × local_context) := do
-- capture state
s <- tactic.read,
-- there are a few ways to add to local context,
-- the most direct being local_context.mk_local
-- however that doesn't handle assignments or frozen locals,
-- so we are setting the local context as the context of a goal
-- and using intro to push a new hypothesis onto the stack
target <- match (assignment, bi) with
| (none, bi) :=
pure $ expr.pi pretty_name bi type `(true)
| (some ass, binder_info.default) :=
pure $ expr.elet pretty_name type ass `(true)
| _ := tactic.fail "Unreachable state reached"
end,
goal_mv <- type_context.run $ type_context.mk_mvar "tmp_goal" target lcxt,
tactic.set_goals [goal_mv],
new_local <- tactic.intro_core pretty_name,
if frozen then
  tactic.freeze_local_instances
else
  pure (),
new_lcxt <- type_context.run $ type_context.get_local_context,
-- reset the state back to the beginning
_root_.set_state s,

return (new_local, new_lcxt)

meta def build_context_aux (nm_map : name_map context.decl) (loc_map : name_map local_context): context.decl -> tactic ((name_map context.decl) × (name_map local_context) × context.decl)
| (context.decl.univ_mvar_decl d) := do
  -- update dependencies
  new_assignment <- d.assignment.mmap (swap_univ_mvs nm_map),
  -- create mvar
  new_univ_mvar <- tactic.mk_meta_univ,
  new_uid <- match new_univ_mvar with
  | level.mvar nm := return nm
  | _ := tactic.failed
  end,
  -- assign mvar
  match new_assignment with
  | some lvl := type_context.run $ type_context.level.assign new_univ_mvar lvl
  | none := return ()
  end,
  -- return new decl
  let new_decl := context.decl.univ_mvar_decl {
    unique_name := new_uid,
    assignment := new_assignment
  },
  let new_nmap := nm_map.insert d.unique_name new_decl,
  return (new_nmap, loc_map, new_decl)
| (context.decl.local_decl d) := do
  -- update dependencies
  let pp_name := d.pp_name,
  new_type <- swap_mvs nm_map (d.type.get_or_else d.expr_type).unfold_string_macros.erase_annotations,
  let (new_lcxt_option : option local_context) := do {
    unique_name <- d.prev,
    loc_map.find unique_name
  },
  let new_lcxt := new_lcxt_option.get_or_else local_context.empty,
  ld <- d.ld,
  new_assignment <- ld.value.mmap ((swap_mvs nm_map) ∘ expr.unfold_string_macros ∘ expr.erase_annotations),
  -- create local
  (new_loc, new_lcxt) <- new_lcxt.mk_local2 pp_name new_type d.bi d.frozen new_assignment,
  (new_uid, new_tp) <- match new_loc with
  | expr.local_const nm _ bi tp := return (nm, tp)
  | _ := tactic.failed
  end,
  let (new_prev : option expr) := new_lcxt.fold (λ prev e, if e = new_loc then prev else some e) none,
  let new_prev_id := match new_prev with
  | some (expr.local_const id _ _ _) := some id
  | _ := none
  end,
  let new_decl := context.decl.local_decl {
    unique_name := new_uid,
    pp_name := pp_name,
    expr_type := new_tp,
    bi := d.bi,
    prev := new_prev_id,
    type := new_type,
    frozen := d.frozen,
    ld := new_lcxt.get_local_decl new_uid
  },
  let new_nmap := nm_map.insert d.unique_name new_decl,
  let new_loc_map := loc_map.insert d.unique_name new_lcxt,
  return (new_nmap, new_loc_map, new_decl)

| (context.decl.mvar_decl d) := do
  -- update dependencies
  let pp_name := d.pp_name,
  new_type <- swap_mvs nm_map (d.type.get_or_else d.expr_type).unfold_string_macros.erase_annotations,
  let (new_lcxt_option : option local_context) := do {
    last <- d.local_cxt.last_option,
    unique_name <- last.local_uniq_name_option,
    loc_map.find unique_name
  },
  let new_lcxt := new_lcxt_option.get_or_else local_context.empty,
  new_assignment <- d.assignment.mmap ((swap_mvs nm_map) ∘ expr.unfold_string_macros ∘ expr.erase_annotations),
  -- create mvar
  new_mvar <- type_context.run $ type_context.mk_mvar pp_name new_type new_lcxt,
  (new_uid, new_tp) <- match new_mvar with
  | expr.mvar nm _ tp := return (nm, tp)
  | _ := tactic.failed
  end,
  -- assign mvar
  match new_assignment with
  | some e := type_context.run $ type_context.assign new_mvar e
  | none := return ()
  end,
  let new_decl := context.decl.mvar_decl {
    unique_name := new_uid,
    pp_name := pp_name,
    expr_type := new_tp,
    local_cxt := new_lcxt.to_list,
    type := new_type,
    assignment := new_assignment
  },
  let new_nmap := nm_map.insert d.unique_name new_decl,
  return (new_nmap, loc_map, new_decl)

meta def rebuild_context : (list context.decl) -> tactic ((name_map context.decl) × (name_map local_context))
| [] := return (mk_name_map, mk_name_map)
| (d :: ds) := do
  (nm_map, loc_map) <- rebuild_context ds,
  (nm_map, loc_map, _) <- build_context_aux nm_map loc_map d,
  return (nm_map, loc_map)

meta def mvar_id : expr -> tactic name
| (expr.mvar uid _ _) := return uid
| _ := tactic.fail "Expecting mvar"

meta def rebuild_tactic_state (ts : tactic_state_data) : tactic unit := do
  (nm_map, _) <- rebuild_context ts.decls.reverse,
  goals_and_tags <- ts.goals.mmap $ λ ⟨nm, tag⟩, do {
    d <- nm_map.find nm,
    nm' <- match d with
    | context.decl.mvar_decl dd := return dd.unique_name
    | _ := tactic.fail "Expecting mvar_decl"
    end,
    mvars <- type_context.run type_context.list_mvars,
    mv <- mvars.mfirst $ λ e, do {
      nm2 <- mvar_id e,
      if nm' = nm2 then return e else failure
    },
    return (mv, tag)
  },
  let goals := goals_and_tags.map prod.fst,
  tactic.set_goals goals,
  goals_and_tags.mmap $ λ ⟨g, tag⟩, do {
    tactic.enable_tags tt,
    tactic.set_tag g tag,
    tactic.enable_tags ff -- seems to be off by default
  },
  return ()

-- for testing

meta def refresh_context : tactic unit := do
  cxt <- context.get,
  (nm_map, _) <- rebuild_context cxt.decls.reverse,
  gs <- tactic.get_goals,
  new_goals <- gs.mmap $ λ g, do {
    nm <- mvar_id g,
    d <- nm_map.find nm,
    tactic.trace (nm, d),
    nm' <- match d with
    | context.decl.mvar_decl dd := return dd.unique_name
    | _ := tactic.fail "Expecting mvar_decl"
    end,
    mvars <- type_context.run type_context.list_mvars,
    mv <- mvars.mfirst $ λ e, do {
      nm2 <- mvar_id e,
      if nm' = nm2 then return e else failure
    },
    return mv
  },
  tactic.set_goals new_goals

meta def refresh_tactic_state : tactic unit := do
  ts_data <- tactic_state_data.get,
  -- go into a clean tactic environment and build the tactic state
  ts <- tactic.unsafe_run_io $ io.run_tactic' $ do {
    rebuild_tactic_state ts_data,
    tactic.read  -- return tactic state
  },
  -- set tactic state to new one
  _root_.set_state ts

-- examples

section examples

-- example (α : Type) (a : nat): a=a := begin
-- trace_context,
-- refresh_tactic_state,
-- trace_context,
-- induction a,
-- trace_context,
-- refresh_tactic_state,  -- check that tags are still there
-- trace_context,
-- refl,
-- refl,
-- done,
-- trace_context,
-- end

-- -- Debug
-- example {α β γ : Type} --(f : α → β) (g : β → γ)
--   : α :=
-- begin
--   trace_context,
--   refresh_tactic_state,
--   trace_context,
-- end

-- -- Frozen local instances
-- def fish {α β γ} {m : Type → Type} [monad m] (f : m α → m β) (g : m β → m γ)
--   : m α → m γ :=
-- begin
--   trace_context,
--   refresh_tactic_state,
--   trace_context,
--   revert m, -- fails, `monad m` is a frozen instance
--   revert f, -- succeeds
--   revert α  -- succeeds
-- end

-- example : let x := 0 in x=0 := begin
-- intro,
-- trace_context,
-- refresh_tactic_state,
-- trace_context,
-- simp,
-- done,
-- end

end examples


/-
Copyright (c) 2017 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro

Parallel computation of a computable sequence of computations by
a diagonal enumeration.
The important theorems of this operation are proven as
terminates_parallel and exists_of_mem_parallel.
(This operation is nondeterministic in the sense that it does not
honor sequence equivalence (irrelevance of computation time).)

! This file was ported from Lean 3 source module data.seq.parallel
! leanprover-community/mathlib commit a7e36e48519ab281320c4d192da6a7b348ce40ad
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
import Mathbin.Data.Seq.Wseq

universe u v

namespace Computation

/- ./././Mathport/Syntax/Translate/Command.lean:224:11: unsupported: unusual advanced open style -/
/- ./././Mathport/Syntax/Translate/Command.lean:224:11: unsupported: unusual advanced open style -/
variable {α : Type u} {β : Type v}

/- ./././Mathport/Syntax/Translate/Expr.lean:177:8: unsupported: ambiguous notation -/
def Parallel.aux2 : List (Computation α) → Sum α (List (Computation α)) :=
  List.foldr
    (fun c o =>
      match o with
      | Sum.inl a => Sum.inl a
      | Sum.inr ls => rmap (fun c' => c'::ls) (destruct c))
    (Sum.inr [])
#align computation.parallel.aux2 Computation.Parallel.aux2

/- ./././Mathport/Syntax/Translate/Expr.lean:177:8: unsupported: ambiguous notation -/
def Parallel.aux1 :
    List (Computation α) × Wseq (Computation α) →
      Sum α (List (Computation α) × Wseq (Computation α))
  | (l, S) =>
    rmap
      (fun l' =>
        match Seq.destruct S with
        | none => (l', Seq.nil)
        | some (none, S') => (l', S')
        | some (some c, S') => (c::l', S'))
      (Parallel.aux2 l)
#align computation.parallel.aux1 Computation.Parallel.aux1

/-- Parallel computation of an infinite stream of computations,
  taking the first result -/
def parallel (S : Wseq (Computation α)) : Computation α :=
  corec Parallel.aux1 ([], S)
#align computation.parallel Computation.parallel

theorem TerminatesParallel.aux :
    ∀ {l : List (Computation α)} {S c},
      c ∈ l → Terminates c → Terminates (corec Parallel.aux1 (l, S)) :=
  by
  have lem1 :
    ∀ l S, (∃ a : α, parallel.aux2 l = Sum.inl a) → terminates (corec parallel.aux1 (l, S)) :=
    by
    intro l S e
    cases' e with a e
    have : corec parallel.aux1 (l, S) = return a :=
      by
      apply destruct_eq_ret
      simp [parallel.aux1]
      rw [e]
      simp [rmap]
    rw [this]
    infer_instance
  intro l S c m T
  revert l S
  apply @terminates_rec_on _ _ c T _ _
  · intro a l S m
    apply lem1
    induction' l with c l IH generalizing m <;> simp at m
    · contradiction
    cases' m with e m
    · rw [← e]
      simp [parallel.aux2]
      cases' List.foldr parallel.aux2._match_1 (Sum.inr List.nil) l with a' ls
      exacts[⟨a', rfl⟩, ⟨a, rfl⟩]
    · cases' IH m with a' e
      simp [parallel.aux2]
      simp [parallel.aux2] at e
      rw [e]
      exact ⟨a', rfl⟩
  · intro s IH l S m
    have H1 : ∀ l', parallel.aux2 l = Sum.inr l' → s ∈ l' :=
      by
      induction' l with c l IH' generalizing m <;> intro l' e' <;> simp at m
      · contradiction
      cases' m with e m <;> simp [parallel.aux2] at e'
      · rw [← e] at e'
        cases' List.foldr parallel.aux2._match_1 (Sum.inr List.nil) l with a' ls <;>
          injection e' with e'
        rw [← e']
        simp
      · induction' e : List.foldr parallel.aux2._match_1 (Sum.inr List.nil) l with a' ls <;>
          rw [e] at e'
        · contradiction
        have := IH' m _ e
        simp [parallel.aux2] at e'
        cases destruct c <;> injection e' with h'
        rw [← h']
        simp [this]
    induction' h : parallel.aux2 l with a l'
    · exact lem1 _ _ ⟨a, h⟩
    · have H2 : corec parallel.aux1 (l, S) = think _ :=
        by
        apply destruct_eq_think
        simp [parallel.aux1]
        rw [h]
        simp [rmap]
      rw [H2]
      apply @Computation.think_terminates _ _ _
      have := H1 _ h
      rcases seq.destruct S with (_ | ⟨_ | c, S'⟩) <;> simp [parallel.aux1] <;> apply IH <;>
        simp [this]
#align computation.terminates_parallel.aux Computation.TerminatesParallel.aux

theorem terminates_parallel {S : Wseq (Computation α)} {c} (h : c ∈ S) [T : Terminates c] :
    Terminates (parallel S) :=
  by
  suffices
    ∀ (n) (l : List (Computation α)) (S c),
      c ∈ l ∨ some (some c) = Seq.nth S n → Terminates c → Terminates (corec Parallel.aux1 (l, S))
    from
    let ⟨n, h⟩ := h
    this n [] S c (Or.inr h) T
  intro n; induction' n with n IH <;> intro l S c o T
  · cases' o with a a
    · exact terminates_parallel.aux a T
    have H : seq.destruct S = some (some c, _) :=
      by
      unfold seq.destruct Functor.map
      rw [← a]
      simp
    induction' h : parallel.aux2 l with a l' <;> have C : corec parallel.aux1 (l, S) = _
    · apply destruct_eq_ret
      simp [parallel.aux1]
      rw [h]
      simp [rmap]
    · rw [C]
      skip
      infer_instance
    · apply destruct_eq_think
      simp [parallel.aux1]
      rw [h, H]
      simp [rmap]
    · rw [C]
      apply @Computation.think_terminates _ _ _
      apply terminates_parallel.aux _ T
      simp
  · cases' o with a a
    · exact terminates_parallel.aux a T
    induction' h : parallel.aux2 l with a l' <;> have C : corec parallel.aux1 (l, S) = _
    · apply destruct_eq_ret
      simp [parallel.aux1]
      rw [h]
      simp [rmap]
    · rw [C]
      skip
      infer_instance
    · apply destruct_eq_think
      simp [parallel.aux1]
      rw [h]
      simp [rmap]
    · rw [C]
      apply @Computation.think_terminates _ _ _
      have TT : ∀ l', terminates (corec parallel.aux1 (l', S.tail)) :=
        by
        intro
        apply IH _ _ _ (Or.inr _) T
        rw [a]
        cases' S with f al
        rfl
      induction' e : seq.nth S 0 with o
      · have D : seq.destruct S = none := by
          dsimp [seq.destruct]
          rw [e]
          rfl
        rw [D]
        simp [parallel.aux1]
        have TT := TT l'
        rwa [seq.destruct_eq_nil D, seq.tail_nil] at TT
      · have D : seq.destruct S = some (o, S.tail) :=
          by
          dsimp [seq.destruct]
          rw [e]
          rfl
        rw [D]
        cases' o with c <;> simp [parallel.aux1, TT]
#align computation.terminates_parallel Computation.terminates_parallel

theorem exists_of_mem_parallel {S : Wseq (Computation α)} {a} (h : a ∈ parallel S) :
    ∃ c ∈ S, a ∈ c :=
  by
  suffices
    ∀ C,
      a ∈ C →
        ∀ (l : List (Computation α)) (S),
          corec Parallel.aux1 (l, S) = C → ∃ c, (c ∈ l ∨ c ∈ S) ∧ a ∈ c
    from
    let ⟨c, h1, h2⟩ := this _ h [] S rfl
    ⟨c, h1.resolve_left id, h2⟩
  let F : List (Computation α) → Sum α (List (Computation α)) → Prop :=
    by
    intro l a
    cases' a with a l'
    exact ∃ c ∈ l, a ∈ c
    exact ∀ a', (∃ c ∈ l', a' ∈ c) → ∃ c ∈ l, a' ∈ c
  have lem1 : ∀ l : List (Computation α), F l (parallel.aux2 l) :=
    by
    intro l
    induction' l with c l IH <;> simp [parallel.aux2]
    · intro a h
      rcases h with ⟨c, hn, _⟩
      exact False.elim hn
    · simp [parallel.aux2] at IH
      cases' List.foldr parallel.aux2._match_1 (Sum.inr List.nil) l with a ls <;>
        simp [parallel.aux2]
      · rcases IH with ⟨c', cl, ac⟩
        refine' ⟨c', Or.inr cl, ac⟩
      · induction' h : destruct c with a c' <;> simp [rmap]
        · refine' ⟨c, List.mem_cons_self _ _, _⟩
          rw [destruct_eq_ret h]
          apply ret_mem
        · intro a' h
          rcases h with ⟨d, dm, ad⟩
          simp at dm
          cases' dm with e dl
          · rw [e] at ad
            refine' ⟨c, List.mem_cons_self _ _, _⟩
            rw [destruct_eq_think h]
            exact think_mem ad
          · cases' IH a' ⟨d, dl, ad⟩ with d dm
            cases' dm with dm ad
            exact ⟨d, Or.inr dm, ad⟩
  intro C aC
  refine' mem_rec_on aC _ fun C' IH => _ <;> intro l S e <;> have e' := congr_arg destruct e <;>
          have := lem1 l <;>
        simp [parallel.aux1] at e' <;>
      cases' parallel.aux2 l with a' l' <;>
    injection e' with h'
  · rw [h'] at this
    rcases this with ⟨c, cl, ac⟩
    exact ⟨c, Or.inl cl, ac⟩
  · induction' e : seq.destruct S with a <;> rw [e] at h'
    ·
      exact
        let ⟨d, o, ad⟩ := IH _ _ h'
        let ⟨c, cl, ac⟩ := this a ⟨d, o.resolve_right (wseq.not_mem_nil _), ad⟩
        ⟨c, Or.inl cl, ac⟩
    · cases' a with o S'
      cases' o with c <;> simp [parallel.aux1] at h' <;> rcases IH _ _ h' with ⟨d, dl | dS', ad⟩
      ·
        exact
          let ⟨c, cl, ac⟩ := this a ⟨d, dl, ad⟩
          ⟨c, Or.inl cl, ac⟩
      · refine' ⟨d, Or.inr _, ad⟩
        rw [seq.destruct_eq_cons e]
        exact seq.mem_cons_of_mem _ dS'
      · simp at dl
        cases' dl with dc dl
        · rw [dc] at ad
          refine' ⟨c, Or.inr _, ad⟩
          rw [seq.destruct_eq_cons e]
          apply seq.mem_cons
        ·
          exact
            let ⟨c, cl, ac⟩ := this a ⟨d, dl, ad⟩
            ⟨c, Or.inl cl, ac⟩
      · refine' ⟨d, Or.inr _, ad⟩
        rw [seq.destruct_eq_cons e]
        exact seq.mem_cons_of_mem _ dS'
#align computation.exists_of_mem_parallel Computation.exists_of_mem_parallel

theorem map_parallel (f : α → β) (S) : map f (parallel S) = parallel (S.map (map f)) :=
  by
  refine'
    eq_of_bisim
      (fun c1 c2 =>
        ∃ l S,
          c1 = map f (corec parallel.aux1 (l, S)) ∧
            c2 = corec parallel.aux1 (l.map (map f), S.map (map f)))
      _ ⟨[], S, rfl, rfl⟩
  intro c1 c2 h;
  exact
    match c1, c2, h with
    | _, _, ⟨l, S, rfl, rfl⟩ => by
      clear _match
      have : parallel.aux2 (l.map (map f)) = lmap f (rmap (List.map (map f)) (parallel.aux2 l)) :=
        by
        simp [parallel.aux2]
        induction' l with c l IH <;> simp
        rw [IH]
        cases List.foldr parallel.aux2._match_1 (Sum.inr List.nil) l <;> simp [parallel.aux2]
        cases destruct c <;> simp
      simp [parallel.aux1]
      rw [this]
      cases' parallel.aux2 l with a l' <;> simp
      apply S.rec_on _ (fun c S => _) fun S => _ <;> simp <;> simp [parallel.aux1] <;>
        exact ⟨_, _, rfl, rfl⟩
#align computation.map_parallel Computation.map_parallel

theorem parallel_empty (S : Wseq (Computation α)) (h : S.headI ~> none) : parallel S = empty _ :=
  eq_empty_of_not_terminates fun ⟨⟨a, m⟩⟩ =>
    by
    let ⟨c, cs, ac⟩ := exists_of_mem_parallel m
    let ⟨n, nm⟩ := Wseq.exists_nth_of_mem cs
    let ⟨c', h'⟩ := Wseq.head_some_of_nth_some nm
    injection h h'
#align computation.parallel_empty Computation.parallel_empty

-- The reason this isn't trivial from exists_of_mem_parallel is because it eliminates to Sort
def parallelRec {S : Wseq (Computation α)} (C : α → Sort v) (H : ∀ s ∈ S, ∀ a ∈ s, C a) {a}
    (h : a ∈ parallel S) : C a :=
  by
  let T : wseq (Computation (α × Computation α)) := S.map fun c => c.map fun a => (a, c)
  have : S = T.map (map fun c => c.1) :=
    by
    rw [← wseq.map_comp]
    refine' (wseq.map_id _).symm.trans (congr_arg (fun f => wseq.map f S) _)
    funext c
    dsimp [id, Function.comp]
    rw [← map_comp]
    exact (map_id _).symm
  have pe := congr_arg parallel this
  rw [← map_parallel] at pe
  have h' := h
  rw [pe] at h'
  haveI : terminates (parallel T) := (terminates_map_iff _ _).1 ⟨⟨_, h'⟩⟩
  induction' e : get (parallel T) with a' c
  have : a ∈ c ∧ c ∈ S := by
    rcases exists_of_mem_map h' with ⟨d, dT, cd⟩
    rw [get_eq_of_mem _ dT] at e
    cases e
    dsimp at cd
    cases cd
    rcases exists_of_mem_parallel dT with ⟨d', dT', ad'⟩
    rcases wseq.exists_of_mem_map dT' with ⟨c', cs', e'⟩
    rw [← e'] at ad'
    rcases exists_of_mem_map ad' with ⟨a', ac', e'⟩
    injection e' with i1 i2
    constructor
    rwa [i1, i2] at ac'
    rwa [i2] at cs'
  cases' this with ac cs
  apply H _ cs _ ac
#align computation.parallel_rec Computation.parallelRec

theorem parallel_promises {S : Wseq (Computation α)} {a} (H : ∀ s ∈ S, s ~> a) : parallel S ~> a :=
  fun a' ma' =>
  let ⟨c, cs, ac⟩ := exists_of_mem_parallel ma'
  H _ cs ac
#align computation.parallel_promises Computation.parallel_promises

theorem mem_parallel {S : Wseq (Computation α)} {a} (H : ∀ s ∈ S, s ~> a) {c} (cs : c ∈ S)
    (ac : a ∈ c) : a ∈ parallel S := by
  haveI := terminates_of_mem ac <;> haveI := terminates_parallel cs <;>
    exact mem_of_promises _ (parallel_promises H)
#align computation.mem_parallel Computation.mem_parallel

theorem parallel_congr_lem {S T : Wseq (Computation α)} {a} (H : S.LiftRel Equiv T) :
    (∀ s ∈ S, s ~> a) ↔ ∀ t ∈ T, t ~> a :=
  ⟨fun h1 t tT =>
    let ⟨s, sS, se⟩ := Wseq.exists_of_liftRel_right H tT
    (promises_congr se _).1 (h1 _ sS),
    fun h2 s sS =>
    let ⟨t, tT, se⟩ := Wseq.exists_of_liftRel_left H sS
    (promises_congr se _).2 (h2 _ tT)⟩
#align computation.parallel_congr_lem Computation.parallel_congr_lem

/- ./././Mathport/Syntax/Translate/Expr.lean:177:8: unsupported: ambiguous notation -/
-- The parallel operation is only deterministic when all computation paths lead to the same value
theorem parallel_congr_left {S T : Wseq (Computation α)} {a} (h1 : ∀ s ∈ S, s ~> a)
    (H : S.LiftRel Equiv T) : parallel S ~ parallel T :=
  let h2 := (parallel_congr_lem H).1 h1
  fun a' =>
  ⟨fun h => by
    have aa := parallel_promises h1 h <;> rw [← aa] <;> rw [← aa] at h <;>
      exact
        let ⟨s, sS, as⟩ := exists_of_mem_parallel h
        let ⟨t, tT, st⟩ := wseq.exists_of_lift_rel_left H sS
        let aT := (st _).1 as
        mem_parallel h2 tT aT,
    fun h => by
    have aa := parallel_promises h2 h <;> rw [← aa] <;> rw [← aa] at h <;>
      exact
        let ⟨s, sS, as⟩ := exists_of_mem_parallel h
        let ⟨t, tT, st⟩ := wseq.exists_of_lift_rel_right H sS
        let aT := (st _).2 as
        mem_parallel h1 tT aT⟩
#align computation.parallel_congr_left Computation.parallel_congr_left

/- ./././Mathport/Syntax/Translate/Expr.lean:177:8: unsupported: ambiguous notation -/
theorem parallel_congr_right {S T : Wseq (Computation α)} {a} (h2 : ∀ t ∈ T, t ~> a)
    (H : S.LiftRel Equiv T) : parallel S ~ parallel T :=
  parallel_congr_left ((parallel_congr_lem H).2 h2) H
#align computation.parallel_congr_right Computation.parallel_congr_right

end Computation


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

Monotone version of 1-point elimination for ∃ (and ∀).

  `∃x, t[x, R x a] ~> t[a, rfl]`

with the (reflexive) relation `R`, and `t` is monotone in `R`:
  `∀x y, R x y → t x → t y`.

Then
  (∃x, R x a ∧ p x) ↔ p a
then
  assume h : p a, ⟨a, h, R.refl a⟩
and
  assume ⟨x, hx, hxa⟩, mono x a hxa hx

Or with the following dependent monotonicity:
  `∀x y (h : R x y), t x h → t y R.refl`.

Then
  (∃x, ∃h:R x a, p x h) ↔ p a R.refl
then
  assume h : p a R.refl, ⟨a, h, R.refl a⟩
and
  assume ⟨x, hx, hxa⟩, mono x a hxa hx

-/
import simp_loop.conv

inductive bintree (α : Type*)
| leaf (a : α) : bintree
| node (l r : bintree) : bintree

def list.dedup {α : Type*} [decidable_eq α] : list α → list α
| []        := []
| (x :: xs) := (if x ∈ xs then xs.dedup else x :: xs.dedup)

namespace bintree
variables {α : Type*} {β : Type*}

def pos := list bool

def left : bintree α → bintree α
| (leaf a)   := leaf a
| (node l r) := l

def right : bintree α → bintree α
| (leaf a)   := leaf a
| (node l r) := r

def at_pos : pos → bintree α → bintree α
| []        t := t
| (ff :: p) t := at_pos p (left t)
| (tt :: p) t := at_pos p (right t)

def map (f : α → β) : bintree α → bintree β
| (leaf a) := leaf (f a)
| (node l r) := node (map l) (map r)

def mmap {m : Type* → Type*} [monad m] (f : α → m β) : bintree α → m (bintree β)
| (leaf a)   := leaf <$> f a
| (node l r) := node <$> mmap l <*> mmap r

end bintree

def monotone {α : Sort*} (r : α → α → Prop) (p : α → Prop) : Prop :=
∀x y, r x y → p x → p y

lemma monotone_eq {α : Type*} (p : α → Prop) : monotone (=) p :=
assume x y h, h ▸ id

-- do we/want need a dependent version?
lemma exists_elim_rel {α : Sort*} {r : α → α → Prop} {p : α → Prop} {a : α} [is_refl α r]
  (h : monotone r p) :
  (∃x, r x a ∧ p x) ↔ p a :=
⟨assume ⟨x, hxa, hx⟩, h x a hxa hx, assume ha, ⟨a, is_refl.refl r a, ha⟩⟩

lemma ex_extend {α : Type*} {p : α → Prop} : (∃a, p a) ↔ (∃a, p a ∧ true) :=
⟨assume ⟨a, ha⟩, ⟨a, ha, ⟨⟩⟩, assume ⟨a, ha, ⟨⟩⟩, ⟨a, ha⟩⟩

lemma l_eq {p E : Prop} : (p ∧ E) ↔ (E ∧ p) :=
and_comm _ _
lemma l_ex {p E : Prop} {r : E → Prop} : (p ∧ (∃h:E, r h)) ↔ (∃h:E, p ∧ r h) :=
⟨λ⟨h, e, r⟩, ⟨e, h, r⟩, λ⟨e, h, r⟩, ⟨h, e, r⟩⟩
lemma l_cn {p q E : Prop} : (p ∧ (E ∧ q)) ↔ (E ∧ (p ∧ q)) :=
⟨λ⟨h, e, r⟩, ⟨e, h, r⟩, λ⟨e, h, r⟩, ⟨h, e, r⟩⟩
lemma r_eq {p E : Prop} : (E ∧ p) ↔ (E ∧ p) :=
iff.refl _
lemma r_ex {p E : Prop} {r : E → Prop} : ((∃h:E, r h) ∧ p) ↔ (∃h:E, r h ∧ p) :=
⟨λ⟨⟨e, r⟩, h⟩, ⟨e, r, h⟩, λ⟨e, r, h⟩, ⟨⟨e, r⟩, h⟩⟩
lemma r_cn {p q E : Prop} : ((E ∧ q) ∧ p) ↔ (E ∧ (q ∧ p)) :=
⟨λ⟨⟨e, r⟩, h⟩, ⟨e, r, h⟩, λ⟨e, r, h⟩, ⟨⟨e, r⟩, h⟩⟩
lemma e_eq_Prop {p E : Prop} : (∃a:p, E) ↔ (E ∧ p) :=
⟨λ⟨a, h⟩, ⟨h, a⟩, λ⟨h, a⟩, ⟨a, h⟩⟩
lemma e_eq {α : Sort*} {E : Prop} : (∃a:α, E) ↔ (E ∧ ∃a:α, true) :=
⟨λ⟨a, h⟩, ⟨h, ⟨a, ⟨⟩⟩⟩, λ⟨h, ⟨a, _⟩⟩, ⟨a, h⟩⟩
lemma e_ex {α : Sort*} {E : Prop} {s : E → α → Prop} : (∃a:α, ∃h:E, s h a) ↔ (∃h:E, ∃a:α, s h a) :=
⟨λ⟨a, h, t⟩, ⟨h, a, t⟩, λ⟨a, h, t⟩, ⟨h, a, t⟩⟩
lemma e_cn {α : Sort*} {E : Prop} {t : α → Prop} : (∃a:α, E ∧ t a) ↔ (E ∧ (∃a:α, t a)) :=
⟨λ⟨a, e, t⟩, ⟨e, ⟨a, t⟩⟩, λ⟨e, ⟨a, t⟩⟩, ⟨a, e, t⟩⟩

lemma l_congr {p q r : Prop} (h : q ↔ r) : p ∧ q ↔ p ∧ r :=
and_congr (iff.refl p) h
lemma r_congr {p q r : Prop} (h : q ↔ r) : q ∧ p ↔ r ∧ p :=
and_congr h (iff.refl p)
lemma ex_congr {α : Sort*} {p q : α → Prop} (h : ∀a, p a ↔ q a) : (∃a, p a) ↔ (∃a, q a) :=
exists_congr h

lemma comm_l {α : Sort*} {p : Prop} {q : α → Prop} : p ∧ (∃a, q a) ↔ ∃a, p ∧ q a :=
⟨λ⟨a, hq, hp⟩, ⟨hq, ⟨a, hp⟩⟩, λ ⟨hq, ⟨a, hp⟩⟩, ⟨a, hq, hp⟩⟩
lemma comm_r {α : Sort*} {p : Prop} {q : α → Prop} : (∃a, q a) ∧ p ↔ ∃a, q a ∧ p :=
⟨λ ⟨⟨a, hp⟩, hq⟩, ⟨a, hp, hq⟩, λ⟨a, hp, hq⟩, ⟨⟨a, hp⟩, hq⟩⟩
lemma comm_ex {α : Sort*} {β : Sort*} {p : α → β → Prop} : (∃a b, p a b) ↔ (∃b a, p a b) :=
⟨λ⟨a, ⟨b, h⟩⟩, ⟨b, ⟨a, h⟩⟩, λ⟨a, ⟨b, h⟩⟩, ⟨b, ⟨a, h⟩⟩⟩

namespace simp_loop
open conv_t tactic expr

meta inductive info
| binder (v : expr) (dependent : bool) | operator (side : bool)

namespace info
open format
meta instance : has_to_format info :=
⟨λi, match i with
| info.binder v d := to_fmt "binder " ++ to_fmt v ++ " " ++ to_fmt d
| info.operator s := to_fmt "operator " ++ to_fmt s
end⟩
end info

meta inductive norm_form
| eq | ex | cn

namespace norm_form
open format
meta instance : has_to_format norm_form :=
⟨λi, match i with
| norm_form.eq := to_fmt "eq"
| norm_form.ex := to_fmt "ex"
| norm_form.cn := to_fmt "cn"
end⟩
end norm_form

meta def congr_ex {α} (c : conv α) : conv α := congr_binder ``ex_congr (λ_, c)

meta def congr {α} (c : conv α) : info → conv α
| (info.binder _ _) := congr_ex c
| (info.operator tt) := congr_simple ``r_congr c
| (info.operator ff) := congr_simple ``l_congr c

meta def apply_norm (n_eq n_ex n_cn : list name) : norm_form → conv norm_form
| norm_form.eq := do n_eq.mfirst apply_const, return norm_form.cn
| norm_form.ex := do n_ex.mfirst apply_const, return norm_form.ex
| norm_form.cn := do n_cn.mfirst apply_const, return norm_form.cn

meta def analyse (chk : expr → list expr → expr → bool)
  (v : expr) (deps : list expr) : expr → tactic (option $ list $ info × expr)
| `(@Exists %%α %%p) := if chk v deps α then
    return none
  else do
    (lam pp_n bi domain body) ← return p | return (some []),
    x ← mk_local' pp_n bi domain,
    return [(info.binder x (deps.any $ λv, v.occurs α), body.instantiate_var x)]
| `(%%p ∧ %%q) := return [(info.operator tt, p), (info.operator ff, q)]
| t := return $ if chk v deps t then none else some []

meta def find (chk : expr → list expr → expr → bool)
  (v : expr) : list expr → expr → tactic (list $ list info) | deps e := do
some is ← analyse chk v deps e | return [[]],
iss ← is.mmap (λ⟨i, t⟩, do
  deps ← return $ match i with (info.binder v tt) := v :: deps | _ := deps end,
  iss ← find deps t,
  return $ iss.map $ λis, i :: is),
return iss.join

section reorder

private meta def reorder_dependent : list info → conv norm_form
| [] := (do -- trace "reorder_equality []", trace_lhs,
  `(_ = _) ← lhs, return norm_form.eq) <|> return norm_form.ex
| (i::xs)  := do
  n ← congr (reorder_dependent xs) i,
  match i with
  | info.binder _ _  := apply_norm [``e_eq_Prop, ``e_eq] [``e_ex] [``e_cn] n
  | info.operator tt := apply_norm [``r_eq] [``r_ex] [``r_cn] n
  | info.operator ff := apply_norm [``l_eq] [``l_ex] [``l_cn] n
  end

private meta def reorder_non_dependent : list info → conv (option (list info))
| []      := return none
| (i::is) := (do info.binder _ ff ← return i, return is) <|> (do
  some is' ← congr (reorder_non_dependent is) i | return none,
  r ← return $ match i with
  | info.binder _ _  := ``comm_ex
  | info.operator tt := ``comm_r
  | info.operator ff := ``comm_l
  end,
  apply_const r,
  return (i :: is'))

meta def reorder {α} (elim : norm_form → conv α) : list info → conv α | l := do
-- trace (to_fmt "reorder_and_elim " ++ to_fmt l),
some l' ← congr_ex (reorder_non_dependent l) |
  (congr_ex (reorder_dependent l) >>= elim),
apply_const ``comm_ex,
congr_ex (reorder l')

end reorder

section term_focus

/-

Move a term on one side of a relation using Galois connections:

Symmetric rules:
  f x R y ↔ x Q g z

    ⟹ f x R t ↔ x Q t'
    ⟹ t Q g x ↔ t' R x

Injectivity rules:
  f x R g y ↔ x Q y

    ⟹ f x R t ↔ x Q t'
    ⟹ t R g x ↔ t' Q x

Splitting rules:
  f x R t ↔ (C₁ ∧ x Q₁ t₁) ∨ (C₂ ∧ x Q₂ t₂)

Setup:

* allow symmetric relations
* apply the rules symmetrically
* should we add dischargers for conditional rules, ala:
    0 < R → x / R ≤ S ↔ x ≤ S * R
  or:
    x / R ≤ S ↔ (0 < R ∧ x ≤ S * R) ∨ (R < 0 ∧ S * R ≤ x) ∨ (R = 0 ∧ 0 ≤ S)
    This one doesn't work for monotone elimination as we have x ≤ S * R and S * R ≤ x case.
* conditionals:
    ∀x n i : ℕ, x + n = i ↔ (n ≤ i ∧ x = i - n)
    ∀x n i : ℕ, x - n = i ↔ ((i = 0 ∧ x ≤ n) ∨ x = i + n)
    ∀x n i : ℕ, n - x = i ↔ ((i = 0 ∧ n ≤ x) ∨ (i ≤ n ∧ x = i - n))

-/

meta def connection_iff : user_attribute :=
{ name := `connection_iff,
  descr := "Connection rules of the form f x R y ↔ x Q g z, used for term focusing" }

section analyse

meta def conjs : expr → tactic (list expr)
| `(%%a ∧ %%b) := do
  as ← conjs a,
  bs ← conjs b,
  return (as ++ bs)
| e := return [e]

meta def disjs_of_conjs : expr → tactic (list $ list expr)
| `(%%a ∨ %%b) := do
  as ← disjs_of_conjs a,
  bs ← disjs_of_conjs b,
  return (as ++ bs)
| e := do
  d ← conjs e,
  return [d]

-- better `parse_rel`: this doesn't work with heq!
-- idea: use relation manager
meta def parse_rel : expr → tactic (expr × expr × expr)
| (expr.app (expr.app f a) b) := return (f, a, b)
| _ := fail "term is not a relation application"

def peep {α} : list α → list (list α × α × list α)
| []        := []
| (a :: xs) := ([], a, xs) :: (peep xs).map (λ⟨p, a', s⟩, (a::p, a', s))

private meta def analyse_connection_aux (ls : list level) (vs : list expr) (lhs rhs : expr) :
  tactic $ list pattern := do
disjs ← disjs_of_conjs rhs,
candidates ← disjs.mmap (λconjs, do
  candidates ← (peep conjs).mmap (λxs, (do
    (ps, e, ss) ← return xs,
    (rel, l, r) ← parse_rel e,
    return $
      (if l ∈ vs ∧ (r :: rel :: ps ++ ss).all (λe, ¬ l.occurs e) then [l] else []) ++
      (if r ∈ vs ∧ (l :: rel :: ps ++ ss).all (λe, ¬ r.occurs e) then [r] else []))
      <|> return []),
  return candidates.join),
let candidates := vs.filter $ λv, candidates.all $ λcs, v ∈ cs,

(rel, l, r) ← parse_rel lhs,
let candidates := candidates.filter $ λc,
  ¬ c.occurs rel ∧ ((c.occurs l ∧ ¬ c.occurs r) ∨ c.occurs r ∧ ¬ c.occurs l),
let candidates := candidates.dedup,
let vs' := vs.filter (λv, v.occurs lhs),
candidates.mmap (λc, mk_pattern ls vs' lhs [] [c])

/-- `analyse_connection ls r` a list of symm-flag and pattern. The pattern matches if the rule is
applicable. In this case `tactic.match_pattern` returns one expression where the focused variable
should occur. The symm-flag indicates if the iff-rule should be applied in its symmetric variant. -/
meta def analyse_connection (n : name) : tactic $ list $ bool × pattern := do
e ← get_env,
d ← e.get n,
let ls := d.univ_params.map level.param,
(vs, `(%%lhs ↔ %%rhs)) ← mk_local_pis d.type,
l ← analyse_connection_aux ls vs lhs rhs,
r ← analyse_connection_aux ls vs rhs lhs,
return (l.map (λp, (ff, p)) ++ r.map (λp, (tt, p)))

meta def ors (c : conv unit) : conv unit := do
`(_ ∨ _) ← lhs | c,
congr_core (congr_core skip ors) ors,
skip

meta def ands (c : conv unit) : conv unit := do
`(_ ∧ _) ← lhs | c,
congr_core (congr_core skip ands) ands,
skip

meta def local_eq (v₀ v₁ : expr) : bool :=
v₀.is_local_constant ∧ v₁.is_local_constant ∧ v₀.local_uniq_name = v₁.local_uniq_name

meta def term_focus (l : list (name × bool × pattern)) (v : expr) : conv unit :=
l.mfirst $ assume ⟨n, symm, pat⟩, do
  ([], [t]) ← match_pattern pat,
  guard $ v.occurs t,
  (if symm then apply_const n else fail "symmetric apply not supported yet"),
  ors (ands $ (do
    l ← lhs,
    (_, l, r) ← lift_tactic $ parse_rel l,
    guard ((v.occurs l ∧ ¬ local_eq l v) ∨ (v.occurs r ∧ ¬ local_eq r v)),
    term_focus) <|> skip)

end analyse

/-

Lattices

  x ≤ y ⊓ z ↔ x ≤ y ∧ x ≤ z

  y ⊓ z ≤ x ↔ y \ z ≤ x (Heyting algebra)

  a ⊔ b ≤ c ↔ a ≤ c ∧ b ≤ c

Case: Focus on a
  (∃ x, x ⊔ b ≤ c ∧ p x) ↔ (∃x, x ≤ c ∧ b ≤ c ∧ p x)

-/

end term_focus

#exit

section equality_elim

meta def check_eq (x : expr) (deps : list expr) : expr → bool
| `(%%l = %%r) := (l = x ∧ deps.all (λx, ¬ x.occurs r)) ∨ (r = x ∧ deps.all (λx, ¬ x.occurs l))
| _ := ff

meta def elim_equality (n : norm_form) : conv unit := do
-- trace (to_fmt "elim_equality [reordered] " ++ to_fmt n), trace_lhs,
apply_norm [``elim_eq_left, ``elim_eq_right] [``elim_ex_left, ``elim_ex_right] [``elim_cn_left, ``elim_cn_right] n,
skip

meta def run : conv unit := do
pss ← congr_binder ``ex_congr (λv, do t ← lhs, find check_eq v [v] t),
pss.mfirst (reorder elim_equality)

end equality_elim

end simp_loop


import ReactorModel.Determinism.InstantaneousExecution

namespace Execution

open ReactorType
open State (Closed)

variable [Indexable α] {s s₁ s₂ : State α} [State.Nontrivial s] [State.Nontrivial s₁]

namespace AdvanceTag

theorem not_Closed (a : s₁ ⇓- s₂) : ¬(Closed s₂) :=
  have := a.advance.preserves_Nontrivial -- TODO: Make this work via type class inference.
  (absurd a.advance.progress_empty ·.progress_Nonempty.ne_empty)

theorem nonrepeatable (a₁ : s₁ ⇓- s₂) (a₂ : s₂ ⇓- s₃) : False :=
  absurd a₂.closed a₁.not_Closed

theorem tag_lt (a : s₁ ⇓- s₂) : s₁.tag < s₂.tag :=
  a.advance.tag_lt

theorem tag_ne (a : s₁ ⇓- s₂) : s₁.tag ≠ s₂.tag :=
  ne_of_lt a.tag_lt

theorem determinisic (a₁ : s ⇓- s₁) (a₂ : s ⇓- s₂) : s₁ = s₂ :=
  a₁.advance.determinisic a₂.advance

instance preserves_Nontrivial [State.Nontrivial s₁] {e : s₁ ⇓- s₂} : State.Nontrivial s₂ :=
  e.advance.preserves_Nontrivial

end AdvanceTag

namespace Instantaneous
namespace ClosedExecution

theorem not_Closed (e : s₁ ⇓| s₂) : ¬(Closed s₁) := by
  simp [Closed]
  have h := Partial.Nonempty.iff_ids_nonempty.mp $ State.Nontrivial.nontrivial (s := s₁)
  exact e.fresh ▸ h.ne_empty.symm

theorem preserves_tag (e : s₁ ⇓| s₂) : s₁.tag = s₂.tag :=
  e.exec.preserves_tag

theorem equiv (e : s₁ ⇓| s₂) : s₁.rtr ≈ s₂.rtr :=
  e.exec.equiv

theorem rcns_Nodup (e : s₁ ⇓| s₂) : e.rcns.Nodup :=
  e.exec.rcns_nodup

theorem progress_def (e : s₁ ⇓| s₂) : s₂.progress = s₁.rtr[.rcn].ids :=
  Equivalent.obj?_rcn_eq e.equiv ▸ e.closed

theorem mem_rcns_iff (e : s₁ ⇓| s₂) : rcn ∈ e.rcns ↔ (rcn ∈ s₁.rtr[.rcn] ∧ rcn ∉ s₁.progress) := by
  simp [Partial.mem_def, e.progress_def ▸ e.exec.mem_rcns_iff (rcn := rcn)]

theorem rcns_perm (e₁ : s ⇓| s₁) (e₂ : s ⇓| s₂) : e₁.rcns ~ e₂.rcns := by
  simp [List.perm_ext e₁.rcns_Nodup e₂.rcns_Nodup, e₁.mem_rcns_iff, e₂.mem_rcns_iff]

theorem tag_eq (e₁ : s ⇓| s₁) (e₂ : s ⇓| s₂) : s₁.tag = s₂.tag :=
  e₁.exec.preserves_tag ▸ e₂.exec.preserves_tag

theorem progress_eq (e₁ : s ⇓| s₁) (e₂ : s ⇓| s₂) : s₁.progress = s₂.progress := by
  simp [e₁.progress_def, e₂.progress_def]

theorem deterministic (e₁ : s ⇓| s₁) (e₂ : s ⇓| s₂) : s₁ = s₂ :=
  e₁.exec.deterministic e₂.exec (e₁.tag_eq e₂) (e₁.progress_eq e₂)

theorem step_determined (e : s ⇓| s₁) (a : s ⇓- s₂) : False :=
  absurd a.closed e.not_Closed

instance preserves_Nontrivial [h : State.Nontrivial s₁] {e : s₁ ⇓| s₂} : State.Nontrivial s₂ where
  nontrivial := Equivalent.obj?_rcn_eq e.equiv ▸ h.nontrivial

theorem nonrepeatable (e₁ : s₁ ⇓| s₂) (e₂ : s₂ ⇓| s₃) : False :=
  have := e₁.preserves_Nontrivial -- TODO: Make this work via type class inference.
  absurd e₁.closed $ e₂.not_Closed

theorem progress_ssubset (e : s₁ ⇓| s₂) : s₁.progress ⊂ s₂.progress := by
  have := e.preserves_Nontrivial -- TODO: Make this work via type class inference.
  rw [e.fresh]
  exact e.closed.progress_Nonempty.empty_ssubset

end ClosedExecution
end Instantaneous

namespace Step

theorem tag_le : (s₁ ⇓ s₂) → s₁.tag ≤ s₂.tag
  | close e   => le_of_eq e.preserves_tag
  | advance a => le_of_lt a.tag_lt

theorem deterministic : (s ⇓ s₁) → (s ⇓ s₂) → s₁ = s₂
  | close e₁, close e₂                      => e₁.deterministic e₂
  | advance a₁, advance a₂                  => a₁.determinisic a₂
  | close e, advance a | advance a, close e => e.step_determined a |>.elim

theorem seq_tag_lt : (s₁ ⇓ s₂) → (s₂ ⇓ s₃) → s₁.tag < s₃.tag
  | close e₁,   close e₂   => e₁.nonrepeatable e₂ |>.elim
  | advance a₁, advance a₂ => a₁.nonrepeatable a₂ |>.elim
  | close e,    advance a  => e.preserves_tag ▸ a.tag_lt
  | advance a,  close e    => e.preserves_tag ▸ a.tag_lt

instance preserves_Nontrivial [State.Nontrivial s₁] : (s₁ ⇓ s₂) → State.Nontrivial s₂
  | close e   => e.preserves_Nontrivial
  | advance a => a.preserves_Nontrivial

end Step

end Execution


import Qq
open Qq

set_option linter.unusedVariables false in
def typeClassArgument (α : Q(Sort u)) (inst : Q(Inhabited $α)) : Q($α) :=
  q(Inhabited.default)

example : Q(Nat) :=
  typeClassArgument q(Nat) q(inferInstance)

open Lean in
#eval show MetaM Q(Nat) from do
  let _ ← synthInstanceQ q(Inhabited Nat)
  return typeClassArgument (u := levelOne) q(Nat) q(inferInstance)


/-
Copyright (c) 2018 Mario Carneiro. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro, Kenny Lau
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.data.list.basic
import Mathlib.PostPort

universes u v w z u_1 u_2 u_3

namespace Mathlib

namespace list

/- zip & unzip -/

@[simp] theorem zip_with_cons_cons {α : Type u} {β : Type v} {γ : Type w} (f : α → β → γ) (a : α)
    (b : β) (l₁ : List α) (l₂ : List β) :
    zip_with f (a :: l₁) (b :: l₂) = f a b :: zip_with f l₁ l₂ :=
  rfl

@[simp] theorem zip_cons_cons {α : Type u} {β : Type v} (a : α) (b : β) (l₁ : List α)
    (l₂ : List β) : zip (a :: l₁) (b :: l₂) = (a, b) :: zip l₁ l₂ :=
  rfl

@[simp] theorem zip_with_nil_left {α : Type u} {β : Type v} {γ : Type w} (f : α → β → γ)
    (l : List β) : zip_with f [] l = [] :=
  rfl

@[simp] theorem zip_with_nil_right {α : Type u} {β : Type v} {γ : Type w} (f : α → β → γ)
    (l : List α) : zip_with f l [] = [] :=
  list.cases_on l (Eq.refl (zip_with f [] []))
    fun (l_hd : α) (l_tl : List α) => Eq.refl (zip_with f (l_hd :: l_tl) [])

@[simp] theorem zip_nil_left {α : Type u} {β : Type v} (l : List α) : zip [] l = [] := rfl

@[simp] theorem zip_nil_right {α : Type u} {β : Type v} (l : List α) : zip l [] = [] :=
  zip_with_nil_right Prod.mk l

@[simp] theorem zip_swap {α : Type u} {β : Type v} (l₁ : List α) (l₂ : List β) :
    map prod.swap (zip l₁ l₂) = zip l₂ l₁ :=
  sorry

@[simp] theorem length_zip_with {α : Type u} {β : Type v} {γ : Type w} (f : α → β → γ) (l₁ : List α)
    (l₂ : List β) : length (zip_with f l₁ l₂) = min (length l₁) (length l₂) :=
  sorry

@[simp] theorem length_zip {α : Type u} {β : Type v} (l₁ : List α) (l₂ : List β) :
    length (zip l₁ l₂) = min (length l₁) (length l₂) :=
  length_zip_with Prod.mk

theorem lt_length_left_of_zip_with {α : Type u} {β : Type v} {γ : Type w} {f : α → β → γ} {i : ℕ}
    {l : List α} {l' : List β} (h : i < length (zip_with f l l')) : i < length l :=
  and.left
    (eq.mp (Eq._oldrec (Eq.refl (i < min (length l) (length l'))) (propext lt_min_iff))
      (eq.mp (Eq._oldrec (Eq.refl (i < length (zip_with f l l'))) (length_zip_with f l l')) h))

theorem lt_length_right_of_zip_with {α : Type u} {β : Type v} {γ : Type w} {f : α → β → γ} {i : ℕ}
    {l : List α} {l' : List β} (h : i < length (zip_with f l l')) : i < length l' :=
  and.right
    (eq.mp (Eq._oldrec (Eq.refl (i < min (length l) (length l'))) (propext lt_min_iff))
      (eq.mp (Eq._oldrec (Eq.refl (i < length (zip_with f l l'))) (length_zip_with f l l')) h))

theorem lt_length_left_of_zip {α : Type u} {β : Type v} {i : ℕ} {l : List α} {l' : List β}
    (h : i < length (zip l l')) : i < length l :=
  lt_length_left_of_zip_with h

theorem lt_length_right_of_zip {α : Type u} {β : Type v} {i : ℕ} {l : List α} {l' : List β}
    (h : i < length (zip l l')) : i < length l' :=
  lt_length_right_of_zip_with h

theorem zip_append {α : Type u} {β : Type v} {l₁ : List α} {r₁ : List α} {l₂ : List β} {r₂ : List β}
    (h : length l₁ = length l₂) : zip (l₁ ++ r₁) (l₂ ++ r₂) = zip l₁ l₂ ++ zip r₁ r₂ :=
  sorry

theorem zip_map {α : Type u} {β : Type v} {γ : Type w} {δ : Type z} (f : α → γ) (g : β → δ)
    (l₁ : List α) (l₂ : List β) : zip (map f l₁) (map g l₂) = map (prod.map f g) (zip l₁ l₂) :=
  sorry

theorem zip_map_left {α : Type u} {β : Type v} {γ : Type w} (f : α → γ) (l₁ : List α)
    (l₂ : List β) : zip (map f l₁) l₂ = map (prod.map f id) (zip l₁ l₂) :=
  eq.mpr
    (id
      (Eq._oldrec (Eq.refl (zip (map f l₁) l₂ = map (prod.map f id) (zip l₁ l₂)))
        (Eq.symm (zip_map f id l₁ l₂))))
    (eq.mpr (id (Eq._oldrec (Eq.refl (zip (map f l₁) l₂ = zip (map f l₁) (map id l₂))) (map_id l₂)))
      (Eq.refl (zip (map f l₁) l₂)))

theorem zip_map_right {α : Type u} {β : Type v} {γ : Type w} (f : β → γ) (l₁ : List α)
    (l₂ : List β) : zip l₁ (map f l₂) = map (prod.map id f) (zip l₁ l₂) :=
  eq.mpr
    (id
      (Eq._oldrec (Eq.refl (zip l₁ (map f l₂) = map (prod.map id f) (zip l₁ l₂)))
        (Eq.symm (zip_map id f l₁ l₂))))
    (eq.mpr (id (Eq._oldrec (Eq.refl (zip l₁ (map f l₂) = zip (map id l₁) (map f l₂))) (map_id l₁)))
      (Eq.refl (zip l₁ (map f l₂))))

theorem zip_map' {α : Type u} {β : Type v} {γ : Type w} (f : α → β) (g : α → γ) (l : List α) :
    zip (map f l) (map g l) = map (fun (a : α) => (f a, g a)) l :=
  sorry

theorem mem_zip {α : Type u} {β : Type v} {a : α} {b : β} {l₁ : List α} {l₂ : List β} :
    (a, b) ∈ zip l₁ l₂ → a ∈ l₁ ∧ b ∈ l₂ :=
  sorry

theorem map_fst_zip {α : Type u} {β : Type v} (l₁ : List α) (l₂ : List β) :
    length l₁ ≤ length l₂ → map prod.fst (zip l₁ l₂) = l₁ :=
  sorry

theorem map_snd_zip {α : Type u} {β : Type v} (l₁ : List α) (l₂ : List β) :
    length l₂ ≤ length l₁ → map prod.snd (zip l₁ l₂) = l₂ :=
  sorry

@[simp] theorem unzip_nil {α : Type u} {β : Type v} : unzip [] = ([], []) := rfl

@[simp] theorem unzip_cons {α : Type u} {β : Type v} (a : α) (b : β) (l : List (α × β)) :
    unzip ((a, b) :: l) = (a :: prod.fst (unzip l), b :: prod.snd (unzip l)) :=
  sorry

theorem unzip_eq_map {α : Type u} {β : Type v} (l : List (α × β)) :
    unzip l = (map prod.fst l, map prod.snd l) :=
  sorry

theorem unzip_left {α : Type u} {β : Type v} (l : List (α × β)) :
    prod.fst (unzip l) = map prod.fst l :=
  sorry

theorem unzip_right {α : Type u} {β : Type v} (l : List (α × β)) :
    prod.snd (unzip l) = map prod.snd l :=
  sorry

theorem unzip_swap {α : Type u} {β : Type v} (l : List (α × β)) :
    unzip (map prod.swap l) = prod.swap (unzip l) :=
  sorry

theorem zip_unzip {α : Type u} {β : Type v} (l : List (α × β)) :
    zip (prod.fst (unzip l)) (prod.snd (unzip l)) = l :=
  sorry

theorem unzip_zip_left {α : Type u} {β : Type v} {l₁ : List α} {l₂ : List β} :
    length l₁ ≤ length l₂ → prod.fst (unzip (zip l₁ l₂)) = l₁ :=
  sorry

theorem unzip_zip_right {α : Type u} {β : Type v} {l₁ : List α} {l₂ : List β}
    (h : length l₂ ≤ length l₁) : prod.snd (unzip (zip l₁ l₂)) = l₂ :=
  eq.mpr (id (Eq._oldrec (Eq.refl (prod.snd (unzip (zip l₁ l₂)) = l₂)) (Eq.symm (zip_swap l₂ l₁))))
    (eq.mpr
      (id
        (Eq._oldrec (Eq.refl (prod.snd (unzip (map prod.swap (zip l₂ l₁))) = l₂))
          (unzip_swap (zip l₂ l₁))))
      (unzip_zip_left h))

theorem unzip_zip {α : Type u} {β : Type v} {l₁ : List α} {l₂ : List β}
    (h : length l₁ = length l₂) : unzip (zip l₁ l₂) = (l₁, l₂) :=
  sorry

theorem zip_of_prod {α : Type u} {β : Type v} {l : List α} {l' : List β} {lp : List (α × β)}
    (hl : map prod.fst lp = l) (hr : map prod.snd lp = l') : lp = zip l l' :=
  sorry

theorem map_prod_left_eq_zip {α : Type u} {β : Type v} {l : List α} (f : α → β) :
    map (fun (x : α) => (x, f x)) l = zip l (map f l) :=
  sorry

theorem map_prod_right_eq_zip {α : Type u} {β : Type v} {l : List α} (f : α → β) :
    map (fun (x : α) => (f x, x)) l = zip (map f l) l :=
  sorry

@[simp] theorem length_revzip {α : Type u} (l : List α) : length (revzip l) = length l := sorry

@[simp] theorem unzip_revzip {α : Type u} (l : List α) : unzip (revzip l) = (l, reverse l) :=
  unzip_zip (Eq.symm (length_reverse l))

@[simp] theorem revzip_map_fst {α : Type u} (l : List α) : map prod.fst (revzip l) = l :=
  eq.mpr (id (Eq._oldrec (Eq.refl (map prod.fst (revzip l) = l)) (Eq.symm (unzip_left (revzip l)))))
    (eq.mpr (id (Eq._oldrec (Eq.refl (prod.fst (unzip (revzip l)) = l)) (unzip_revzip l)))
      (Eq.refl (prod.fst (l, reverse l))))

@[simp] theorem revzip_map_snd {α : Type u} (l : List α) : map prod.snd (revzip l) = reverse l :=
  eq.mpr
    (id
      (Eq._oldrec (Eq.refl (map prod.snd (revzip l) = reverse l))
        (Eq.symm (unzip_right (revzip l)))))
    (eq.mpr (id (Eq._oldrec (Eq.refl (prod.snd (unzip (revzip l)) = reverse l)) (unzip_revzip l)))
      (Eq.refl (prod.snd (l, reverse l))))

theorem reverse_revzip {α : Type u} (l : List α) : reverse (revzip l) = revzip (reverse l) := sorry

theorem revzip_swap {α : Type u} (l : List α) : map prod.swap (revzip l) = revzip (reverse l) :=
  sorry

theorem nth_zip_with {α : Type u_1} {β : Type u_1} {γ : Type u_1} (f : α → β → γ) (l₁ : List α)
    (l₂ : List β) (i : ℕ) : nth (zip_with f l₁ l₂) i = f <$> nth l₁ i <*> nth l₂ i :=
  sorry

theorem nth_zip_with_eq_some {α : Type u_1} {β : Type u_2} {γ : Type u_3} (f : α → β → γ)
    (l₁ : List α) (l₂ : List β) (z : γ) (i : ℕ) :
    nth (zip_with f l₁ l₂) i = some z ↔
        ∃ (x : α), ∃ (y : β), nth l₁ i = some x ∧ nth l₂ i = some y ∧ f x y = z :=
  sorry

theorem nth_zip_eq_some {α : Type u} {β : Type v} (l₁ : List α) (l₂ : List β) (z : α × β) (i : ℕ) :
    nth (zip l₁ l₂) i = some z ↔ nth l₁ i = some (prod.fst z) ∧ nth l₂ i = some (prod.snd z) :=
  sorry

@[simp] theorem nth_le_zip_with {α : Type u} {β : Type v} {γ : Type w} {f : α → β → γ} {l : List α}
    {l' : List β} {i : ℕ} {h : i < length (zip_with f l l')} :
    nth_le (zip_with f l l') i h =
        f (nth_le l i (lt_length_left_of_zip_with h))
          (nth_le l' i (lt_length_right_of_zip_with h)) :=
  sorry

@[simp] theorem nth_le_zip {α : Type u} {β : Type v} {l : List α} {l' : List β} {i : ℕ}
    {h : i < length (zip l l')} :
    nth_le (zip l l') i h =
        (nth_le l i (lt_length_left_of_zip h), nth_le l' i (lt_length_right_of_zip h)) :=
  nth_le_zip_with

theorem mem_zip_inits_tails {α : Type u} {l : List α} {init : List α} {tail : List α} :
    (init, tail) ∈ zip (inits l) (tails l) ↔ init ++ tail = l :=
  sorry

end Mathlib


import Lean
import Lean.Meta
import Lean.Parser
import LeanCodePrompts.CheckParse
import LeanCodePrompts.ParseJson
import LeanCodePrompts.Autocorrect
import LeanCodePrompts.KeywordSummary.KeywordExtraction
import LeanCodePrompts.EgsTranslate
open Lean Meta

open Lean Elab Parser Command

def fileName := "data/safe_prompts.json"

/-- extract prompt pairs from JSON response to local server -/
def sentenceSimPairs
  (s: String)
  (theoremField : String := "theorem")
   : MetaM  <| Except String (Array (String × String)) := do
  let json ← readJson s
  return do
    (← json.getArr?).mapM <| fun j => do
      let docstring ← j.getObjValAs? String "doc_string"
      let typeField :=
        if fileName ∈ ["data/mathlib4-prompts.json"] then "type"
        else theoremField
      let thm ← j.getObjValAs? String typeField
      pure (docstring, thm)

-- #eval sentenceSimPairs egSen

namespace GPT

def message (role content : String) : Json :=
  Json.mkObj [("role", role), ("content", content)]

def prompt (sys: String) (egs : List <| String × String)(query : String) : Json :=
  let head := message "system" sys
  let egArr :=
    egs.bind (fun (ds, thm) => [message "user" ds, message "assistant" thm])
  Json.arr <| head :: egArr ++ [message "user" query] |>.toArray

def sysPrompt := "You are a coding assistant who translates from natural language to Lean Theorem Prover code following examples. Follow EXACTLY the examples given."

def makePrompt(query : String)(pairs: Array (String × String)) : Json:= prompt sysPrompt pairs.toList query

def makeFlipPrompt(query : String)(pairs: Array (String × String)) : Json:= prompt sysPrompt (pairs.toList.map (fun (x, y) => (y, x))) query

def jsonToExprStrArray (json: Json) : TermElabM (Array String) := do
  let outArr : Array String ←
    match json.getArr? with
    | Except.ok arr =>
        let parsedArr : Array String ←
          arr.filterMapM <| fun js =>
            match js.getObjVal? "message" with
            | Except.ok jsobj =>
                match jsobj.getObjVal? "content" with
                | Except.ok jsstr =>
                  match jsstr.getStr? with
                  | Except.ok str => pure (some str)
                  | Except.error e =>
                    throwError m!"json string expected but got {js}, error: {e}"
                | Except.error _ =>
                  throwError m!"no content field in {jsobj}"
            | Except.error _ =>
                throwError m!"no message field in {js}"

        pure parsedArr
    | Except.error e => throwError m!"json parsing error: {e}"

end GPT

/-- make prompt from prompt pairs -/
@[deprecated GPT.makePrompt]
def makePrompt(prompt : String)(pairs: Array (String × String)) : String :=
      pairs.foldr (fun  (ds, thm) acc =>
        -- acc ++ "/-- " ++ ds ++" -/\ntheorem" ++ thm ++ "\n" ++ "\n"
s!"/-- {ds} -/
theorem {thm} :=

{acc}"
          ) s!"/-- {prompt} -/
theorem "

/-- make prompt for reverse translation from prompt pairs -/
@[deprecated GPT.makeFlipPrompt]
def makeFlipPrompt(statement : String)(pairs: Array (String × String)) : String :=
      pairs.foldr (fun  (ds, thm) acc =>
s!"theorem {thm} :=
/- {ds} -/

{acc}"
          ) s!"theorem {statement} :=
/- "

/-- make prompt for reverse translation from prompt pairs -/
def makeFlipStatementsPrompt(statement : String)(pairs: Array (String × String)) : String :=
      pairs.foldr (fun  (ds, thm) acc =>
s!"{thm} :=
/- {ds} -/

{acc}"
          ) s!"{statement} :=
/- "

def openAIKey : IO (Option String) := IO.getEnv "OPENAI_API_KEY"

/--query OpenAI Codex with given prompt and parameters -/
def codexQuery(prompt: String)(n: Nat := 1)
  (temp : JsonNumber := ⟨2, 1⟩)(stopTokens: Array String :=  #[":=", "-/"]) : MetaM Json := do
  let key? ← openAIKey
  let key :=
    match key? with
    | some k => k
    | none => panic! "OPENAI_API_KEY not set"
  let dataJs := Json.mkObj [("model", "code-davinci-002"), ("prompt", prompt), ("temperature", Json.num temp), ("n", n), ("max_tokens", 150), ("stop", Json.arr <| stopTokens |>.map Json.str)]
  let data := dataJs.pretty
  trace[Translate.info] "OpenAI query: {data}"
  let out ←  IO.Process.output {
        cmd:= "curl",
        args:= #["https://api.openai.com/v1/completions",
        "-X", "POST",
        "-H", "Authorization: Bearer " ++ key,
        "-H", "Content-Type: application/json",
        "--data", data]}
  readJson out.stdout

def gptQuery(messages: Json)(n: Nat := 1)
  (temp : JsonNumber := ⟨2, 1⟩)(stopTokens: Array String :=  #[":=", "-/"]) : MetaM Json := do
  let key? ← openAIKey
  let key :=
    match key? with
    | some k => k
    | none => panic! "OPENAI_API_KEY not set"
  let dataJs := Json.mkObj [("model", "gpt-3.5-turbo"), ("messages", messages)
  , ("temperature", Json.num temp), ("n", n), ("max_tokens", 150), ("stop", Json.arr <| stopTokens |>.map Json.str)
  ]
  let data := dataJs.pretty
  trace[Translate.info] "OpenAI query: {data}"
  let out ←  IO.Process.output {
        cmd:= "curl",
        args:= #["https://api.openai.com/v1/chat/completions",
        "-X", "POST",
        "-H", "Authorization: Bearer " ++ key,
        "-H", "Content-Type: application/json",
        "--data", data]}
  trace[Translate.info] "OpenAI response: {out.stdout} (stderr: {out.stderr})"
  readJson out.stdout

def openAIQuery(prompt: String)(n: Nat := 1)
  (temp : JsonNumber := ⟨2, 1⟩)(stopTokens: Array String :=  #[":=", "-/"]) : MetaM Json :=
  codexQuery prompt n temp stopTokens

/-!
Caching, polling etc to avoid repeatedly calling servers
-/

initialize webCacheJson : IO.Ref (HashMap String Json) ← IO.mkRef (HashMap.empty)

initialize pendingJsonQueries : IO.Ref (HashSet String)
    ← IO.mkRef (HashSet.empty)

initialize logCache : IO.Ref (Array String) ← IO.mkRef (#[])

def mkLog{α : Type _}[ToString α](msg: α) : IO Unit := do
  let cache ← logCache.get
  logCache.set (cache.push (toString msg))

def logs (num: Nat) : IO (List String) := do
  let cache ← logCache.get
  return cache.reverse.toList.take num

def showLogs (num: Nat) : IO Unit := do
  let cache ← logCache.get
  let ls := cache.reverse.toList.take num
  for lines in ls do
  for l in lines.splitOn "\n" do
    IO.println l

def getCachedJson? (s: String) : IO (Option Json) := do
  let cache ← webCacheJson.get
  return cache.find? s

def cacheJson (s: String)(js: Json)  : IO Unit := do
  let cache ← webCacheJson.get
  webCacheJson.set (cache.insert s js)
  return ()

partial def pollCacheJson (s : String) : IO Json := do
  let cache ← webCacheJson.get
  match cache.find? s with
  | some jsBlob => return jsBlob
  | none => do
    IO.sleep 200
    pollCacheJson s

/-- check if there is a valid elaboration after translation, autocorrection -/
def hasElab (s: String)(limit : Option Nat := none) : TermElabM Bool := do
    -- (elabThmTrans s).map (fun e => e.toBool)
  let elab? ← polyElabThmTrans s limit
  match elab? with
  | Except.error _ => return Bool.false
  | Except.ok els => return !els.isEmpty

/-- log to file -/
def elabLog (s: String) : IO Unit := do
  let logFile := System.mkFilePath ["results/elab_logs.txt"]
  let h ← IO.FS.Handle.mk logFile IO.FS.Mode.append Bool.false
  h.putStrLn s
  h.putStrLn ""

def fixedPrompts:= #[("If $z_1, \\dots, z_n$ are complex, then $|z_1 + z_2 + \\dots + z_n|\\leq |z_1| + |z_2| + \\dots + |z_n|$.", "(n : ℕ) (f : ℕ → ℂ) :\n abs (∑ i in finset.range n, f i) ≤ ∑ i in finset.range n, abs (f i) :="), ("If x and y are in $\\mathbb{R}^n$, then $|x+y|^2 + |x-y|^2 = 2|x|^2 + 2|y|^2$.", "(n : ℕ) (x y : euclidean_space ℝ (fin n)) :\n ∥x + y∥^2 + ∥x - y∥^2 = 2*∥x∥^2 + 2*∥y∥^2 :="), ("If $x$ is an element of infinite order in $G$, prove that the elements $x^n$, $n\\in\\mathbb{Z}$ are all distinct.", "(G : Type*) [group G] (x : G) (hx : x ≠ 1) (hx_inf : ∀ n : ℕ, x ^ n ≠ 1) : ∀ m n : ℤ, m ≠ n → x ^ m ≠ x ^ n :="), ("Let $X$ be a topological space; let $A$ be a subset of $X$. Suppose that for each $x\\in A$ there is an open set $U$ containing $x$ such that $U\\subset A$. Show that $A$ is open in $X$.", "(X : Type*) [topological_space X]\n (A : set X) (hA : ∀ x ∈ A, ∃ U : set X, is_open U ∧ x ∈ U ∧ U ⊆ A):\n is_open A :=")]

/-- choosing pairs to build a prompt -/
def getPromptPairs(s: String)(numSim : Nat)(numKW: Nat)
    (scoreBound: Float)(matchBound: Nat)
   : TermElabM (Array (String × String) × IO.Process.Output) := do
      let jsData := Json.mkObj [
        ("filename", fileName),
        ("field", "doc_string"),
        ("doc_string", s),
        ("n", numSim),
        ("model_name", "all-mpnet-base-v2")
      ]
      let simJsonOut ←
        IO.Process.output {cmd:= "curl", args:=
          #["-X", "POST", "-H", "Content-type: application/json", "-d", jsData.pretty, s!"{← leanAideIP}/nearest_prompts"]}
      let pairs? ← sentenceSimPairs simJsonOut.stdout "theorem"
      -- IO.println s!"obtained sentence similarity; time : {← IO.monoMsNow}"
      let allPairs : Array (String × String) ←
        match pairs? with
        | Except.error e =>
            throwError e
        | Except.ok pairs => pure pairs
      -- logInfo m!"all pairs: {allPairs}"
      let kwPairs :=
        if numKW >0
        then ←  keywordBasedPrompts docPair s numKW scoreBound matchBound
        else #[]
      -- IO.println s!"obtained keyword pairs; time : {← IO.monoMsNow}"
      let allPairs := (allPairs ++ kwPairs).toList.eraseDups.toArray
      let pairs -- := allPairs --
        ←  allPairs.filterM (fun (_, s) => do
            isElabPrompt s )
      let kwPairs ←  keywordBasedPrompts docPair s
      return (
          (pairs ++ kwPairs).toList.eraseDups.toArray, simJsonOut)

/-- choosing pairs to build a prompt -/
def getPromptPairsGeneral(s: String)(numSim : Nat)(field: String := "doc_string")
    (theoremField : String := "theorem")
   : TermElabM (Array (String × String) × IO.Process.Output) := do
      let jsData := Json.mkObj [
        ("filename", fileName),
        ("field", field),
        (field, s),
        ("n", numSim),
        ("model_name", "all-mpnet-base-v2")
      ]
      let simJsonOut ←
        IO.Process.output {cmd:= "curl", args:=
          #["-X", "POST", "-H", "Content-type: application/json", "-d", jsData.pretty, s!"{← leanAideIP}/nearest_prompts"]}
      let pairs? ← sentenceSimPairs simJsonOut.stdout theoremField
      -- IO.println s!"obtained sentence similarity; time : {← IO.monoMsNow}"
      let allPairs : Array (String × String) ←
        match pairs? with
        | Except.error e =>
            throwError e

        | Except.ok pairs => pure pairs
      -- logInfo m!"all pairs: {allPairs}"
      return (
          allPairs.toList.eraseDups.toArray, simJsonOut)

/-- given string to translate, build prompt and query OpenAI; returns JSON response
-/
def getCodeJson (s: String)(numSim : Nat:= 8)(numKW: Nat := 0)(includeFixed: Bool := Bool.false)(queryNum: Nat := 5)(temp : JsonNumber := ⟨2, 1⟩)(scoreBound: Float := 0.2)(matchBound: Nat := 15) : TermElabM Json := do
  match ← getCachedJson? s with
  | some js => return js
  | none =>
    let pending ←  pendingJsonQueries.get
    if pending.contains s then pollCacheJson s
    else
      let pending ←  pendingJsonQueries.get
      pendingJsonQueries.set (pending.insert s)
      -- work starts here; before this was caching, polling etc
      let (pairs, IOOut) ←
        if numSim > 0 then
          getPromptPairs s numSim numKW scoreBound matchBound
        else pure (#[], ⟨0, "", ""⟩)
      let pairs := if includeFixed then pairs ++ fixedPrompts else pairs
      let pairs  := pairs.filter (fun (s, _) => s.length < 100)
      let prompt := GPT.makePrompt s pairs
      trace[Translate.info] m!"prompt: \n{prompt.pretty}"
      -- mkLog prompt
      let fullJson ←
        gptQuery prompt queryNum temp
      let outJson :=
        (fullJson.getObjVal? "choices").toOption.getD (Json.arr #[])
      let pending ←  pendingJsonQueries.get
      pendingJsonQueries.set (pending.erase s)
      if IOOut.exitCode = 0 then cacheJson s outJson
        else throwError m!"Web query error: {IOOut.stderr}"
      return outJson

/-- Given an array of outputs, tries to elaborate them with translation and autocorrection and returns the best choice, throwing an error if nothing elaborates.  -/
def arrayToExpr (output: Array String) : TermElabM Expr := do
  let output := output.toList.eraseDups.toArray
  trace[Translate.info] m!"output:\n{output}"
  -- mkLog output
  let mut elaborated : Array String := Array.empty
  -- translation, autocorrection and filtering by elaboration
  for out in output do
    let ployElab? ← polyElabThmTrans out
    match ployElab? with
      | Except.error _ => pure ()
      | Except.ok es =>
        for (_ , _, s) in es do
            elaborated := elaborated.push s
  if elaborated.isEmpty then do
    -- information with failed logs
    logWarning m!"No valid output from Codex; outputs below"
    for out in output do
      let polyOut ←  polyStrThmTrans out
      for str in polyOut do
        logWarning m!"{str}"
    mkSyntheticSorry (mkSort levelZero)
  else
    -- grouping by trying to prove equality and selecting
    let groupSorted ← groupFuncStrs elaborated
    let topStr := groupSorted[0]![0]!
    let thmExc ← elabFuncTyp topStr
    match thmExc with
    | Except.ok (_, thm) => return thm
    | Except.error s => throwError s

/-- Given an array of outputs, tries to elaborate them with translation and autocorrection and returns the best choice, throwing an error if nothing elaborates.  -/
def arrayToStx (output: Array String) : TermElabM Syntax := do
  let output := output.toList.eraseDups.toArray
  trace[Translate.info] m!"output:\n{output}"
  -- mkLog output
  let mut elaborated : Array String := Array.empty
  -- translation, autocorrection and filtering by elaboration
  for out in output do
    let ployElab? ← polyElabThmTrans out
    match ployElab? with
      | Except.error _ => pure ()
      | Except.ok es =>
        for (_ , _, s) in es do
            elaborated := elaborated.push s
  if elaborated.isEmpty then do
    -- information with failed logs
    logWarning m!"No valid output from Codex; outputs below"
    for out in output do
      let polyOut ←  polyStrThmTrans out
      for str in polyOut do
        logWarning m!"{str}"
    pure Syntax.missing
  else
    -- grouping by trying to prove equality and selecting
    let groupSorted ← groupFuncStrs elaborated
    let topStr := groupSorted[0]![0]!
    let thmExc ← elabFuncTyp topStr
    match thmExc with
    | Except.ok (stx, _) => return stx
    | Except.error s => throwError s

/-- Given an array of outputs, tries to elaborate them with translation and autocorrection and optionally returns the best choice as well as all elaborated terms (used for batch processing, interactive code uses `arrayToExpr` instead)  -/
def arrayToExpr? (output: Array String) : TermElabM (Option (Expr× (Array String))) := do
  -- IO.println s!"arrayToExpr? called with {output.size} outputs"
  let mut elaborated : Array String := Array.empty
  let mut fullElaborated : Array String := Array.empty
  for out in output do
    -- IO.println s!"elaboration called: {out}"
    let ployElab? ← polyElabThmTrans out
    match ployElab? with
      | Except.error _ => pure ()
      | Except.ok es =>
        for (expr, _, s) in es do
          elaborated := elaborated.push s
          if !expr.hasExprMVar then
            fullElaborated := fullElaborated.push s
  if elaborated.isEmpty then
    elabLog "No valid output from Codex; outputs below"
    for out in output do
      let polyOut ←  polyStrThmTrans out
      for str in polyOut do
        elabLog s!"{str}"
    return none
  else
    let priority :=
        if fullElaborated.isEmpty then elaborated else fullElaborated
    let groupSorted ← groupFuncStrs priority
    let topStr := groupSorted[0]![0]!
    let thmExc ← elabFuncTyp topStr
    match thmExc with
    | Except.ok (_, thm) => return some (thm, elaborated)
    | Except.error s =>
        elabLog s!"Second round error : {s}"
        return none

def greedyArrayToExpr? (output: Array String) : TermElabM (Option Expr) := do
    output.findSomeM? <| fun out => do
      let t? ← elabThmTrans? out
      return t?.map fun (expr, _, _) => expr

/-- reverse translation from `Lean` to natural language -/
def leanToPrompt (thm: String)(numSim : Nat:= 5)(numKW: Nat := 1)(temp : JsonNumber := 0)(scoreBound: Float := 0.2)(matchBound: Nat := 15)(textField : String := "text") : TermElabM String := do
    let (pairs, _) ← getPromptPairs thm numSim numKW scoreBound matchBound
    let prompt := GPT.makeFlipPrompt thm pairs
    -- elabLog prompt
    let fullJson ← gptQuery prompt 1 temp
    let outJson :=
      (fullJson.getObjVal? "choices").toOption.getD (Json.arr #[])
    let out? := (outJson.getArrVal? 0).bind fun js => js.getObjVal? textField
    let outJson :=
        match (out?) with
        | Except.error s => Json.str s!"query for translation failed: {s}"
        | Except.ok js => js
    return outJson.getStr!

/-- reverse translation from `Lean` to natural language -/
@[deprecated leanToPrompt]
def statementToDoc (thm: String)(numSim : Nat:= 5)(temp : JsonNumber := 0) : TermElabM String := do
    let (pairs, _) ← getPromptPairsGeneral thm numSim "statement"
    let prompt := makeFlipStatementsPrompt thm pairs
    -- elabLog prompt
    let fullJson ← openAIQuery prompt 1 temp
    let outJson :=
      (fullJson.getObjVal? "choices").toOption.getD (Json.arr #[])
    let out? := (outJson.getArrVal? 0).bind fun js => js.getObjVal? "text"
    let outJson :=
        match (out?) with
        | Except.error s => Json.str s!"query for translation failed: {s}"
        | Except.ok js => js
    return outJson.getStr!

def egThm := "theorem eg_thm : ∀ n: Nat, ∃ m : Nat, m > n ∧ m % 2 = 0"

def egPairs := getPromptPairsGeneral egThm 5 "statement" "statement"

def egPrompt := do
  let (pairs, _) ← egPairs
  return makeFlipStatementsPrompt egThm pairs

-- #eval egPrompt

-- #eval statementToDoc egThm 5 0

-- #eval leanToPrompt "∀ {p : ℕ} [inst : Fact (Nat.Prime p)], p = 2 ∨ p % 2 = 1"

-- #eval leanToPrompt "∀ {α : Type u} {x : FreeGroup α}, x ≠ 1 → ¬IsOfFinOrder x"

-- #eval leanToPrompt "{  n :  ℕ } ->  Even   (    (   n +  1  ) * n  )"

/-- array of outputs extracted from OpenAI Json -/
def jsonToExprStrArray (json: Json) : TermElabM (Array String) := do
  let outArr : Array String ←
    match json.getArr? with
    | Except.ok arr =>
        let parsedArr : Array String ←
          arr.filterMapM <| fun js =>
            match js.getObjVal? "text" with
              | Except.ok jsstr =>
                match jsstr.getStr? with
                | Except.ok str => pure (some str)
                | Except.error e =>
                  throwError m!"json string expected but got {js}, error: {e}"
              | Except.error _ =>
                throwError m!"no text field"
        pure parsedArr
    | Except.error e => throwError m!"json parsing error: {e}"
  return outArr

/-- array of outputs extracted from Json Array -/
def jsonStringToExprStrArray (jsString: String) : TermElabM (Array String) := do
  try
  let json ← readJson jsString
  let outArr : Array String ←
    match json.getArr? with
    | Except.ok arr =>
        let parsedArr : Array String ←
          arr.filterMapM <| fun js =>
            match js.getStr? with
            | Except.ok str => pure (some str)
            | Except.error e =>
              throwError m!"json string expected but got {js}, error: {e}"
        pure parsedArr
    | Except.error _ => pure #[jsString]
  return outArr
  catch _ =>
    pure #[jsString]

-- #eval jsonStringToExprStrArray "simple"
-- #eval jsonStringToExprStrArray "[\"simple\", \"simple2\"]"

/-- given json returned by open-ai obtain the best translation -/
def jsonToExpr' (json: Json) : TermElabM Expr := do
  let output ← GPT.jsonToExprStrArray json
  arrayToExpr output

/-- translation from a comment-like syntax to a theorem statement -/
elab "//-" cb:commentBody  : term => do
  let s := cb.raw.getAtomVal
  let s := (s.dropRight 2).trim
  -- querying codex
  let js ← getCodeJson  s
  -- filtering, autocorrection and selection
  let e ← jsonToExpr' js
  trace[Translate.info] m!"{e}"
  return e

def uncurriedView(numArgs: Nat)(e: Expr) : MetaM String :=
  match numArgs with
  | 0 => do return " : " ++ (← e.view)
  | k +1 =>
    match e with
    | Expr.forallE n t _ bi => do
      let core := s!"{n.eraseMacroScopes} : {← t.view}"
      let typeString :=s!"{← t.view}"
      let argString := match bi with
      | BinderInfo.implicit => "{"++ core ++ "}"
      | BinderInfo.strictImplicit => "{{ "++ core ++ "}}"
      | BinderInfo.instImplicit =>
        if (`inst).isPrefixOf n then s!"[{typeString}]"
          else s!"[{core}]"
      | BinderInfo.default => s!"({core})"
      let tail : String ←
        withLocalDecl `func BinderInfo.default e fun func =>
          withLocalDecl n bi t fun arg => do
            let fx := mkAppN func #[arg]
            let newType ← inferType fx
            uncurriedView k newType
      return " " ++ argString ++ tail
    | _ => do return " : " ++ (← e.view)

elab "uncurry2" e:term : term => do
  let e ← Term.elabTerm e none
  let e ← uncurriedView 2 e
  return mkStrLit e

universe u

def translateViewM (s: String) : TermElabM String := do
  let js ← getCodeJson  s
  let output ← GPT.jsonToExprStrArray js
  trace[Translate.info] m!"{output}"
  let e? ← arrayToExpr? output
  match e? with
  | some (e, _) => do
    e.view
  | none => do
    let stx ← output.findSomeM? <| fun s => do
      let exp ←  identMappedFunStx s
      return exp.toOption
    return stx.getD "False"

/-- view of string in core; to be run with Snapshot.runCore
-/
def translateViewCore (s: String) : CoreM String :=
  (translateViewM s).run'.run'