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

The writer monad transformer for passing immutable state.
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.control.monad.basic
import Mathlib.algebra.group.basic
import Mathlib.PostPort

universes u v l u_1 u_2 u_3 u₀ u₁ v₀ v₁

namespace Mathlib

structure writer_t (ω : Type u) (m : Type u → Type v) (α : Type u)
where
  run : m (α × ω)

def writer (ω : Type u) (α : Type u) :=
  writer_t ω id

namespace writer_t

protected theorem ext {ω : Type u} {m : Type u → Type v} [Monad m] {α : Type u} (x : writer_t ω m α) (x' : writer_t ω m α) (h : run x = run x') : x = x' := sorry

protected def tell {ω : Type u} {m : Type u → Type v} [Monad m] (w : ω) : writer_t ω m PUnit :=
  mk (pure (PUnit.unit, w))

protected def listen {ω : Type u} {m : Type u → Type v} [Monad m] {α : Type u} : writer_t ω m α → writer_t ω m (α × ω) :=
  sorry

protected def pass {ω : Type u} {m : Type u → Type v} [Monad m] {α : Type u} : writer_t ω m (α × (ω → ω)) → writer_t ω m α :=
  sorry

protected def pure {ω : Type u} {m : Type u → Type v} [Monad m] {α : Type u} [HasOne ω] (a : α) : writer_t ω m α :=
  mk (pure (a, 1))

protected def bind {ω : Type u} {m : Type u → Type v} [Monad m] {α : Type u} {β : Type u} [Mul ω] (x : writer_t ω m α) (f : α → writer_t ω m β) : writer_t ω m β :=
  mk
    (do
      let x ← run x
      let x' ← run (f (prod.fst x))
      pure (prod.fst x', prod.snd x * prod.snd x'))

protected instance monad {ω : Type u} {m : Type u → Type v} [Monad m] [HasOne ω] [Mul ω] : Monad (writer_t ω m) := sorry

protected instance is_lawful_monad {ω : Type u} {m : Type u → Type v} [Monad m] [monoid ω] [is_lawful_monad m] : is_lawful_monad (writer_t ω m) := sorry

protected def lift {ω : Type u} {m : Type u → Type v} [Monad m] {α : Type u} [HasOne ω] (a : m α) : writer_t ω m α :=
  mk (flip Prod.mk 1 <$> a)

protected instance has_monad_lift {ω : Type u} (m : Type u → Type u_1) [Monad m] [HasOne ω] : has_monad_lift m (writer_t ω m) :=
  has_monad_lift.mk fun (α : Type u) => writer_t.lift

protected def monad_map {ω : Type u} {m : Type u → Type u_1} {m' : Type u → Type u_2} [Monad m] [Monad m'] {α : Type u} (f : {α : Type u} → m α → m' α) : writer_t ω m α → writer_t ω m' α :=
  fun (x : writer_t ω m α) => mk (f (run x))

protected instance monad_functor {ω : Type u} (m : Type u → Type u_1) (m' : Type u → Type u_1) [Monad m] [Monad m'] : monad_functor m m' (writer_t ω m) (writer_t ω m') :=
  monad_functor.mk writer_t.monad_map

protected def adapt {ω : Type u} {m : Type u → Type v} [Monad m] {ω' : Type u} {α : Type u} (f : ω → ω') : writer_t ω m α → writer_t ω' m α :=
  fun (x : writer_t ω m α) => mk (prod.map id f <$> run x)

protected instance monad_except {ω : Type u} {m : Type u → Type v} [Monad m] (ε : outParam (Type u_1)) [HasOne ω] [Monad m] [monad_except ε m] : monad_except ε (writer_t ω m) :=
  monad_except.mk (fun (α : Type u) => writer_t.lift ∘ throw)
    fun (α : Type u) (x : writer_t ω m α) (c : ε → writer_t ω m α) => mk (catch (run x) fun (e : ε) => run (c e))

end writer_t

/--
An implementation of [MonadReader](
https://hackage.haskell.org/package/mtl-2.2.2/docs/Control-Monad-Reader-Class.html#t:MonadReader).
It does not contain `local` because this function cannot be lifted using `monad_lift`.
Instead, the `monad_reader_adapter` class provides the more general `adapt_reader` function.

Note: This class can be seen as a simplification of the more "principled" definition
```
class monad_reader (ρ : out_param (Type u)) (n : Type u → Type u) :=
(lift {α : Type u} : (∀ {m : Type u → Type u} [monad m], reader_t ρ m α) → n α)
```
-/
class monad_writer (ω : outParam (Type u)) (m : Type u → Type v)
where
  tell : ω → m PUnit
  listen : {α : Type u} → m α → m (α × ω)
  pass : {α : Type u} → m (α × (ω → ω)) → m α

protected instance writer_t.monad_writer {ω : Type u} {m : Type u → Type v} [Monad m] : monad_writer ω (writer_t ω m) :=
  monad_writer.mk writer_t.tell (fun (α : Type u) => writer_t.listen) fun (α : Type u) => writer_t.pass

protected instance reader_t.monad_writer {ω : Type u} {ρ : Type u} {m : Type u → Type v} [Monad m] [monad_writer ω m] : monad_writer ω (reader_t ρ m) :=
  monad_writer.mk (fun (x : ω) => monad_lift (monad_writer.tell x)) (fun (α : Type u) (_x : reader_t ρ m α) => sorry)
    fun (α : Type u) (_x : reader_t ρ m (α × (ω → ω))) => sorry

def swap_right {α : Type u_1} {β : Type u_2} {γ : Type u_3} : (α × β) × γ → (α × γ) × β :=
  sorry

protected instance state_t.monad_writer {ω : Type u} {σ : Type u} {m : Type u → Type v} [Monad m] [monad_writer ω m] : monad_writer ω (state_t σ m) :=
  monad_writer.mk (fun (x : ω) => monad_lift (monad_writer.tell x)) (fun (α : Type u) (_x : state_t σ m α) => sorry)
    fun (α : Type u) (_x : state_t σ m (α × (ω → ω))) => sorry

def except_t.pass_aux {ε : Type u_1} {α : Type u_2} {ω : Type u_3} : except ε (α × (ω → ω)) → except ε α × (ω → ω) :=
  sorry

protected instance except_t.monad_writer {ω : Type u} {ε : Type u} {m : Type u → Type v} [Monad m] [monad_writer ω m] : monad_writer ω (except_t ε m) :=
  monad_writer.mk (fun (x : ω) => monad_lift (monad_writer.tell x)) (fun (α : Type u) (_x : except_t ε m α) => sorry)
    fun (α : Type u) (_x : except_t ε m (α × (ω → ω))) => sorry

def option_t.pass_aux {α : Type u_1} {ω : Type u_2} : Option (α × (ω → ω)) → Option α × (ω → ω) :=
  sorry

protected instance option_t.monad_writer {ω : Type u} {m : Type u → Type v} [Monad m] [monad_writer ω m] : monad_writer ω (option_t m) :=
  monad_writer.mk (fun (x : ω) => monad_lift (monad_writer.tell x)) (fun (α : Type u) (_x : option_t m α) => sorry)
    fun (α : Type u) (_x : option_t m (α × (ω → ω))) => sorry

/-- Adapt a monad stack, changing the type of its top-most environment.

This class is comparable to
[Control.Lens.Magnify](https://hackage.haskell.org/package/lens-4.15.4/docs/Control-Lens-Zoom.html#t:Magnify),
but does not use lenses (why would it), and is derived automatically for any transformer
implementing `monad_functor`.

Note: This class can be seen as a simplification of the more "principled" definition
```
class monad_reader_functor (ρ ρ' : out_param (Type u)) (n n' : Type u → Type u) :=
(map {α : Type u} : (∀ {m : Type u → Type u} [monad m], reader_t ρ m α → reader_t ρ' m α) → n α → n' α)
```
-/
class monad_writer_adapter (ω : outParam (Type u)) (ω' : outParam (Type u)) (m : Type u → Type v) (m' : Type u → Type v)
where
  adapt_writer : {α : Type u} → (ω → ω') → m α → m' α

/-- Transitivity.

This instance generates the type-class problem with a metavariable argument (which is why this
is marked as `[nolint dangerous_instance]`).
Currently that is not a problem, as there are almost no instances of `monad_functor` or
`monad_writer_adapter`.

see Note [lower instance priority] -/
protected instance monad_writer_adapter_trans {ω : Type u} {ω' : Type u} {m : Type u → Type v} {m' : Type u → Type v} {n : Type u → Type v} {n' : Type u → Type v} [monad_writer_adapter ω ω' m m'] [monad_functor m m' n n'] : monad_writer_adapter ω ω' n n' :=
  monad_writer_adapter.mk fun (α : Type u) (f : ω → ω') => monad_map fun (α : Type u) => adapt_writer f

protected instance writer_t.monad_writer_adapter {ω : Type u} {ω' : Type u} {m : Type u → Type v} [Monad m] : monad_writer_adapter ω ω' (writer_t ω m) (writer_t ω' m) :=
  monad_writer_adapter.mk fun (α : Type u) => writer_t.adapt

protected instance writer_t.monad_run (ω : Type u) (m : Type u → Type (max u u_1)) (out : outParam (Type u → Type (max u u_1))) [monad_run out m] : monad_run (fun (α : Type u) => out (α × ω)) (writer_t ω m) :=
  monad_run.mk fun (α : Type u) (x : writer_t ω m α) => run (writer_t.run x)

/-- reduce the equivalence between two writer monads to the equivalence between
their underlying monad -/
def writer_t.equiv {m₁ : Type u₀ → Type v₀} {m₂ : Type u₁ → Type v₁} {α₁ : Type u₀} {ω₁ : Type u₀} {α₂ : Type u₁} {ω₂ : Type u₁} (F : m₁ (α₁ × ω₁) ≃ m₂ (α₂ × ω₂)) : writer_t ω₁ m₁ α₁ ≃ writer_t ω₂ m₂ α₂ :=
  equiv.mk (fun (_x : writer_t ω₁ m₁ α₁) => sorry) (fun (_x : writer_t ω₂ m₂ α₂) => sorry) sorry sorry


/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Leonardo de Moura
-/
prelude
import Init.Control.Basic
import Init.Data.List.Basic

namespace List
universes u v w u₁ u₂

/-
Remark: we can define `mapM`, `mapM₂` and `forM` using `Applicative` instead of `Monad`.
Example:
```
def mapM {m : Type u → Type v} [Applicative m] {α : Type w} {β : Type u} (f : α → m β) : List α → m (List β)
  | []    => pure []
  | a::as => List.cons <$> (f a) <*> mapM as
```

However, we consider `f <$> a <*> b` an anti-idiom because the generated code
may produce unnecessary closure allocations.
Suppose `m` is a `Monad`, and it uses the default implementation for `Applicative.seq`.
Then, the compiler expands `f <$> a <*> b <*> c` into something equivalent to
```
(Functor.map f a >>= fun g_1 => Functor.map g_1 b) >>= fun g_2 => Functor.map g_2 c
```
In an ideal world, the compiler may eliminate the temporary closures `g_1` and `g_2` after it inlines
`Functor.map` and `Monad.bind`. However, this can easily fail. For example, suppose
`Functor.map f a >>= fun g_1 => Functor.map g_1 b` expanded into a match-expression.
This is not unreasonable and can happen in many different ways, e.g., we are using a monad that
may throw exceptions. Then, the compiler has to decide whether it will create a join-point for
the continuation of the match or float it. If the compiler decides to float, then it will
be able to eliminate the closures, but it may not be feasible since floating match expressions
may produce exponential blowup in the code size.

Finally, we rarely use `mapM` with something that is not a `Monad`.

Users that want to use `mapM` with `Applicative` should use `mapA` instead.
-/

@[specialize]
def mapM {m : Type u → Type v} [Monad m] {α : Type w} {β : Type u} (f : α → m β) : List α → m (List β)
  | []    => pure []
  | a::as => return (← f a) :: (← mapM f as)

@[specialize]
def mapA {m : Type u → Type v} [Applicative m] {α : Type w} {β : Type u} (f : α → m β) : List α → m (List β)
  | []    => pure []
  | a::as => List.cons <$> f a <*> mapA f as

@[specialize]
protected def forM {m : Type u → Type v} [Monad m] {α : Type w} (as : List α) (f : α → m PUnit) : m PUnit :=
  match as with
  | []      => pure ⟨⟩
  | a :: as => do f a; List.forM as f

@[specialize]
def forA {m : Type u → Type v} [Applicative m] {α : Type w} (as : List α) (f : α → m PUnit) : m PUnit :=
  match as with
  | []      => pure ⟨⟩
  | a :: as => f a *> forA as f

@[specialize]
def filterAuxM {m : Type → Type v} [Monad m] {α : Type} (f : α → m Bool) : List α → List α → m (List α)
  | [],     acc => pure acc
  | h :: t, acc => do
    let b ← f h
    filterAuxM f t (cond b (h :: acc) acc)

@[inline]
def filterM {m : Type → Type v} [Monad m] {α : Type} (f : α → m Bool) (as : List α) : m (List α) := do
  let as ← filterAuxM f as []
  pure as.reverse

@[inline]
def filterRevM {m : Type → Type v} [Monad m] {α : Type} (f : α → m Bool) (as : List α) : m (List α) :=
  filterAuxM f as.reverse []

@[inline]
def filterMapM {m : Type u → Type v} [Monad m] {α β : Type u} (f : α → m (Option β)) (as : List α) : m (List β) :=
  let rec @[specialize] loop
    | [],     bs => pure bs
    | a :: as, bs => do
      match (← f a) with
      | none   => loop as bs
      | some b => loop as (b::bs)
  loop as.reverse []

@[specialize]
protected def foldlM {m : Type u → Type v} [Monad m] {s : Type u} {α : Type w} : (f : s → α → m s) → (init : s) → List α → m s
  | f, s, []      => pure s
  | f, s, a :: as => do
    let s' ← f s a
    List.foldlM f s' as

@[specialize]
def foldrM {m : Type u → Type v} [Monad m] {s : Type u} {α : Type w} : (f : α → s → m s) → (init : s) → List α → m s
  | f, s, []      => pure s
  | f, s, a :: as => do
    let s' ← foldrM f s as
    f a s'

@[specialize]
def firstM {m : Type u → Type v} [Monad m] [Alternative m] {α : Type w} {β : Type u} (f : α → m β) : List α → m β
  | []    => failure
  | a::as => f a <|> firstM f as

@[specialize]
def anyM {m : Type → Type u} [Monad m] {α : Type v} (f : α → m Bool) : List α → m Bool
  | []    => pure false
  | a::as => do
    match (← f a) with
    | true  => pure true
    | false => anyM f as

@[specialize]
def allM {m : Type → Type u} [Monad m] {α : Type v} (f : α → m Bool) : List α → m Bool
  | []    => pure true
  | a::as => do
    match (← f a) with
    | true  => allM f as
    | false => pure false

@[specialize]
def findM? {m : Type → Type u} [Monad m] {α : Type} (p : α → m Bool) : List α → m (Option α)
  | []    => pure none
  | a::as => do
    match (← p a) with
    | true  => pure (some a)
    | false => findM? p as

@[specialize]
def findSomeM? {m : Type u → Type v} [Monad m] {α : Type w} {β : Type u} (f : α → m (Option β)) : List α → m (Option β)
  | []    => pure none
  | a::as => do
    match (← f a) with
    | some b => pure (some b)
    | none   => findSomeM? f as

@[inline] protected def forIn {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : List α) (init : β) (f : α → β → m (ForInStep β)) : m β :=
  let rec @[specialize] loop
    | [], b    => pure b
    | a::as, b => do
      match (← f a b) with
      | ForInStep.done b  => pure b
      | ForInStep.yield b => loop as b
  loop as init

instance : ForIn m (List α) α where
  forIn := List.forIn

@[simp] theorem forIn_nil [Monad m] (f : α → β → m (ForInStep β)) (b : β) : forIn [] b f = pure b :=
  rfl

@[simp] theorem forIn_cons [Monad m] (f : α → β → m (ForInStep β)) (a : α) (as : List α) (b : β)
    : forIn (a::as) b f = f a b >>= fun | ForInStep.done b => pure b | ForInStep.yield b => forIn as b f :=
  rfl

instance : ForM m (List α) α where
  forM := List.forM

@[simp] theorem forM_nil  [Monad m] (f : α → m PUnit) : forM [] f = pure ⟨⟩ :=
  rfl
@[simp] theorem forM_cons [Monad m] (f : α → m PUnit) (a : α) (as : List α) : forM (a::as) f = f a >>= fun _ => forM as f :=
  rfl

end List


import Lean
import Lean.Meta
import Lean.Elab
import Lean.Parser
import Lean.Parser.Extension
import LeanCodePrompts.Utils
open Lean Meta Elab Parser  Tactic

def depsPrompt : IO (Array String) := do
  let file ← reroutePath <| System.mkFilePath ["data/types.txt"]
  IO.FS.lines file

declare_syntax_cat typed_ident
syntax "(" ident ":" term ")" : typed_ident
syntax "{" ident ":" term "}" : typed_ident

#check Array.foldrM
#check TSyntaxArray.rawImpl
#check TSyntax.mk

instance : Coe (Syntax) (TSyntax n) where
  coe := TSyntax.mk

instance : Coe (Array Syntax) (Array (TSyntax n)) where
  coe := Array.map Coe.coe

/-- check whether a string parses as a term -/
def checkTerm (s : String) : MetaM Bool := do
  let env ← getEnv
  let chk := Lean.Parser.runParserCategory env `term  s
  match chk with
  | Except.ok _  => pure true
  | Except.error _  => pure false

/-- split prompts into those that parse -/
def promptsSplit : MetaM ((Array String) × (Array String)) := do
  let deps ← depsPrompt
  let mut succ: Array String := Array.empty
  let mut fail: Array String := Array.empty
  for type in deps do
    let chk ←  checkTerm type
    if chk then
      succ := succ.push type
    else
      fail := fail.push type
  return (succ, fail)

declare_syntax_cat argument
syntax "(" ident+ " : " term ")" : argument
syntax "{" ident+ " : " term "}" : argument
syntax "[" ident " : " term "]" : argument
syntax "[" term "]" : argument

declare_syntax_cat thmStat
syntax argument* docComment "theorem"  argument*  ":" term : thmStat
syntax "theorem" (ident)? argument*  ":" term : thmStat
syntax "def" ident argument*  ":" term : thmStat
syntax argument*  ":" term : thmStat

def thmsPrompt : IO (Array String) := do
  let file ← reroutePath <| System.mkFilePath ["data/thms.txt"]
  IO.FS.lines file

/-- check whether a string parses as a theorem -/
def checkThm (s : String) : MetaM Bool := do
  let env ← getEnv
  let chk := Lean.Parser.runParserCategory env `thmStat  s
  match chk with
  | Except.ok stx  =>
      IO.println stx
      pure true
  | Except.error _  => pure false

#check Syntax
partial def tokens (s : Syntax) : Array String :=
match s with
| .missing => Array.empty
| .node _ _ args => args.foldl (fun acc x => acc ++ tokens x) Array.empty
| .atom _  val => #[val]
| .ident _ val .. => #[val.toString]

def getTokens (s: String) : MetaM <| Array String := do
  let env ← getEnv
  let chk := Lean.Parser.runParserCategory env `thmStat  s
  match chk with
  | Except.ok stx  =>
      pure <| tokens stx
  | Except.error _  => pure Array.empty

-- #eval getTokens "{α : Type u} [group α] [has_lt α] [covariant_class α α (function.swap has_mul.mul) has_lt.lt] {a : α} : 1 < a⁻¹ ↔ a < 1"

/-- split prompts into those that parse -/
def promptsThmSplit : MetaM ((Array String) × (Array String)) := do
  let deps ← thmsPrompt
  let mut succ: Array String := Array.empty
  let mut fail: Array String := Array.empty
  for type in deps do
    let chk ←  checkThm type
    if chk then
      succ := succ.push type
    else
      fail := fail.push type
  return (succ, fail)

def promptsThmSplitCore : CoreM ((Array String) × (Array String)) :=
  promptsThmSplit.run'

def levelNames :=
  [`u, `v, `u_1, `u_2, `u_3, `u_4, `u_5, `u_6, `u_7, `u_8, `u_9, `u_10, `u_11, `u₁, `u₂, `W₁, `W₂, `w₁, `w₂, `u', `v', `uu, `w, `wE]

partial def idents : Syntax → List String
| Syntax.ident _ s .. => [s.toString]
| Syntax.node _ _ ss => ss.toList.bind idents
| _ => []

def elabThm (s : String)(opens: List String := [])
  (levelNames : List Lean.Name := levelNames)
  : TermElabM <| Except String Expr := do
  let env ← getEnv
  let chk := Lean.Parser.runParserCategory env `thmStat  s
  match chk with
  | Except.ok stx  =>
      match stx with
      | `(thmStat| $_:docComment theorem  $args:argument* : $type:term) =>
        elabAux type args
      | `(thmStat|theorem $_ $args:argument* : $type:term) =>
        elabAux type args
      | `(thmStat|theorem $args:argument* : $type:term) =>
        elabAux type args
      | `(thmStat|$vars:argument* $_:docComment  theorem $args:argument* : $type:term ) =>
        elabAux type (vars ++ args)
      | `(thmStat|def $_ $args:argument* : $type:term) =>
        elabAux type args
      | `(thmStat|$args:argument* : $type:term) =>
        elabAux type args
      | _ => return Except.error s!"parsed incorrectly to {stx}"
  | Except.error e  => return Except.error e
  where elabAux (type: Syntax)(args: Array Syntax) :
        TermElabM <| Except String Expr := do
        let header := if opens.isEmpty then "" else
          (opens.foldl (fun acc s => acc ++ " " ++ s) "open ") ++ " in "
        let mut argS := ""
        for arg in args do
          argS := argS ++ (showSyntax arg) ++ " -> "
        let funStx := s!"{header}{argS}{showSyntax type}"
        match Lean.Parser.runParserCategory (← getEnv) `term funStx with
        | Except.ok termStx => Term.withLevelNames levelNames <|
          try
            let expr ← Term.withoutErrToSorry <|
                Term.elabTerm termStx none
            return Except.ok expr
          catch e =>
            return Except.error s!"{← e.toMessageData.toString} ; identifiers {idents termStx} (during elaboration)"
        | Except.error e =>
            return Except.error s!"parsed to {funStx}; error while parsing as theorem: {e}"

def elabThmCore (s : String)(opens: List String := [])
  (levelNames : List Lean.Name := levelNames)
  : CoreM <| Except String Expr :=
    (elabThm s opens levelNames).run'.run'

theorem true_true_iff_True : true = true ↔ True := by
    apply Iff.intro
    intros
    exact True.intro
    intros
    rfl

theorem true_false_iff_false : false = true ↔ False := by
    apply Iff.intro
    intro hyp
    simp at hyp
    intro hyp
    contradiction

syntax "lynx" ("at" ident)? : tactic
syntax "lynx" "at" "*" : tactic
macro_rules
| `(tactic| lynx) =>
  `(tactic|try(repeat rw [true_true_iff_True]);try (repeat (rw [true_false_iff_false])))
| `(tactic| lynx at $t:ident) =>
  `(tactic| try(repeat rw [true_true_iff_True] at $t:ident);try (repeat (rw [true_false_iff_false] at $t:ident)))
| `(tactic| lynx at *) =>
  `(tactic|try(repeat rw [true_true_iff_True] at *);try (repeat (rw [true_false_iff_false] at *)))

def provedEqual (e₁ e₂ : Expr) : TermElabM Bool := do
  let type ← mkEq e₁ e₂
  let mvar ← mkFreshExprMVar <| some type
  let mvarId := mvar.mvarId!
  let stx ← `(tactic| lynx;  try (rfl))
  let res ←  runTactic mvarId stx
  let (remaining, _) := res
  return remaining.isEmpty

def provedEquiv (e₁ e₂ : Expr) : TermElabM Bool := do
  try
  let type ← mkAppM ``Iff #[e₁, e₂]
  let mvar ← mkFreshExprMVar <| some type
  let mvarId := mvar.mvarId!
  let stx ← `(tactic| intros; lynx at *<;> apply Iff.intro <;> intro hyp  <;> (lynx at *) <;> (try assumption) <;> try (intros; apply Eq.symm; apply hyp))
  let res ←  runTactic mvarId stx
  let (remaining, _) := res
  return remaining.isEmpty
  catch _ => pure false

def compareThms(s₁ s₂ : String)(opens: List String := [])
  (levelNames : List Lean.Name := levelNames)
  : TermElabM <| Except String Bool := do
  let e₁ ← elabThm s₁ opens levelNames
  let e₂ ← elabThm s₂ opens levelNames
  match e₁ with
  | Except.ok e₁ => match e₂ with
    | Except.ok e₂ =>
        let p := (← provedEqual e₁ e₂) ||
          (← provedEquiv e₁ e₂)
        return Except.ok p
    | Except.error e₂ => return Except.error e₂
  | Except.error e₁ => return Except.error e₁

def compareThmsCore(s₁ s₂ : String)(opens: List String := [])
  (levelNames : List Lean.Name := levelNames)
  : CoreM <| Except String Bool :=
    (compareThms s₁ s₂ opens levelNames).run'.run'

def compareThmExps(e₁ e₂: Expr)
  : TermElabM <| Except String Bool := do
      let p := (← provedEqual e₁ e₂) ||
        (← provedEquiv e₁ e₂)
      return Except.ok p

def compareThmExpsCore(e₁ e₂: Expr)
  : CoreM <| Except String Bool := do
      (compareThmExps e₁ e₂).run'.run'

def equalThms(s₁ s₂ : String)(opens: List String := [])
  (levelNames : List Lean.Name := levelNames)
  : TermElabM Bool := do
  match ← compareThms s₁ s₂ opens levelNames with
  | Except.ok p => return p
  | Except.error _ => return false

def groupThms(ss: Array String)(opens: List String := [])
  (levelNames : List Lean.Name := levelNames)
  : TermElabM (Array (Array String)) := do
    let mut groups: Array (Array String) := Array.empty
    for s in ss do
      match ← groups.findIdxM? (fun g =>
          equalThms s g[0]! opens levelNames) with
      |none  =>
        groups := groups.push #[s]
      | some j =>
        groups := groups.set! j (groups[j]!.push s)
    return groups

def groupTheoremsCore(ss: Array String)(opens: List String := [])
  (levelNames : List Lean.Name := levelNames)
  : CoreM (Array (Array String)) :=
    (groupThms ss opens levelNames).run'.run'

def groupThmsSort(ss: Array String)(opens: List String := [])
  (levelNames : List Lean.Name := levelNames)
  : TermElabM (Array (Array String)) := do
  let gps ← groupThms ss opens levelNames
  return gps.qsort (fun xs ys => xs.size > ys.size)

def groupThmsSortCore(ss: Array String)(opens: List String := [])
  (levelNames : List Lean.Name := levelNames)
  : CoreM (Array (Array String)) :=
    (groupThmsSort ss opens levelNames).run'.run'

-- Tests

-- #eval checkTerm "(fun x : Nat => x + 1)"

-- #eval checkTerm "a • s"

-- #eval checkTerm "λ x : Nat, x + 1"

-- #eval checkTerm "a - t = 0"

def checkStatements : MetaM (List (String × Bool)) := do
  let prompts ← depsPrompt
  (prompts.toList.take 50).mapM fun s =>
    do return (s, ← checkTerm s)

def tryParseThm (s : String) : MetaM String := do
  let env ← getEnv
  let chk := Lean.Parser.runParserCategory env `thmStat  s
  match chk with
  | Except.ok stx  =>
      match stx with
      | `(thmStat|theorem $_ $args:argument* : $type:term) =>
        let mut argS := ""
        for arg in args do
          argS := argS ++ (showSyntax arg) ++ " -> "
        let funStx := s!"{argS}{showSyntax type}"
        pure s!"match: {funStx}"
      | `(thmStat|$args:argument* : $type:term) =>
        let mut argS := ""
        for arg in args do
          argS := argS ++ (showSyntax arg) ++ " -> "
        let funStx := s!"{argS}{showSyntax type}"
        pure s!"match: {funStx}"
      | _ => pure s!"parsed to mysterious {stx}"
  | Except.error e  => pure s!"error: {e}"

-- #eval tryParseThm "theorem blah (n : Nat) {m: Type} : n  = n"

-- #eval elabThm "(p: Nat)/-- blah test -/ theorem  (n : Nat) {m: Type} : n  = p"

def eg :=
"section
variable (α : Type) {n : Nat}
/-- A doc that should be ignored -/
theorem blah (m: Nat) : n  = m "

-- #eval checkThm eg

-- #eval checkThm "(n : Nat) {m: Type} : n  = n"

-- #eval tryParseThm "theorem subfield.list_sum_mem {K : Type u} [field K] (s : subfield K) {l : list K} : (∀ (x : K), x ∈ l → x ∈ s) → l.sum ∈ s"

def checkElabThm (s : String) : TermElabM String := do
  let env ← getEnv
  let chk := Lean.Parser.runParserCategory env `thmStat  s
  match chk with
  | Except.ok stx  =>
      match stx with
      | `(thmStat|theorem $_ $args:argument* : $type:term) =>
        let mut argS := ""
        for arg in args do
          argS := argS ++ (showSyntax arg) ++ " -> "
        let funStx := s!"{argS}{showSyntax type}"
        match Lean.Parser.runParserCategory env `term funStx with
        | Except.ok termStx => Term.withLevelNames levelNames <|
          try
            let expr ← Term.withoutErrToSorry <|
                Term.elabTerm termStx none
            pure s!"elaborated: {← expr.view} from {funStx}"
          catch e =>
            pure s!"{← e.toMessageData.toString} during elaboration"
        | Except.error e =>
            pure s!"parsed to {funStx}; error while parsing: {e}"
      | `(thmStat|$vars:argument* $_:docComment theorem $args:argument* : $type:term ) =>
        let mut argS := ""
        for arg in vars ++ args do
          argS := argS ++ (showSyntax arg) ++ " -> "
        let funStx := s!"{argS}{showSyntax type}"
        match Lean.Parser.runParserCategory env `term funStx with
        | Except.ok termStx => Term.withLevelNames levelNames <|
          try
            let expr ← Term.withoutErrToSorry <|
                Term.elabTerm termStx none
            pure s!"elaborated: {← expr.view} from {funStx}"
          catch e =>
            pure s!"{← e.toMessageData.toString} during elaboration"
        | Except.error e =>
            pure s!"parsed to {funStx}; error while parsing: {e}"
      | `(thmStat|$args:argument* : $type:term) =>
        let mut argS := ""
        for arg in args do
          argS := argS ++ (showSyntax arg) ++ " -> "
        let funStx := s!"{argS}{showSyntax type}"
        match Lean.Parser.runParserCategory env `term funStx with
        | Except.ok termStx => Term.withLevelNames levelNames <|
          try
            let expr ← Term.withoutErrToSorry <|
                Term.elabTerm termStx none
            pure s!"elaborated: {← expr.view} from {funStx}"
          catch e =>
            pure s!"{← e.toMessageData.toString} during elaboration"
        | Except.error e =>
            pure s!"parsed to {funStx}; error while parsing: {e}"
      | _ => pure s!"parsed to mysterious {stx}"
  | Except.error e  => pure s!"error: {e}"

-- #eval checkElabThm "theorem blah (n : Nat) {m : Nat} : n  = m"

-- #eval checkElabThm eg

-- #eval checkElabThm "theorem subfield.list_sum_mem {K : Type u} [field K] (s : subfield K) {l : list K} : (∀ (x : K), x ∈ l → x ∈ s) → l.sum ∈ s"

-- #eval elabThm "theorem blah (n : Nat) {m : Nat} : n  = m"

-- #eval elabThm "theorem (n : Nat) {m : Nat} : n  = m"

-- #eval elabThm "theorem blah (n : Nat) {m : Nat} : n  = succ n" ["Nat"]

-- #eval elabThm "theorem blah (n : Nat) {m : Nat} : n  = succ n" ["Nat"]

-- #eval elabThm "(n : Nat) {m : Nat} : n  = succ n" ["Nat"]

-- #eval elabThmCore "(n : Nat) {m : Nat} : n  = succ n" ["Nat"]

-- #eval elabThm "theorem subfield.list_sum_mem {K : Type u} [field K] (s : subfield K) {l : list K} : (∀ (x : K), x ∈ l → x ∈ s) → l.sum ∈ s"

-- #eval compareThms "theorem nonsense(n : Nat) (m : Nat) : n = m" "(p : Nat)(q: Nat) : p = q"

-- #eval compareThms ": True" ": true = true"

-- #eval compareThms "{A: Type} : A →  True" "{A: Type}: A →  true"

-- #eval compareThms ": False" ": false = true"

-- #eval compareThms "{A: Sort} : False →  A" "{A: Sort} : false = true →  A"

example : (∀ {A: Sort}, False → A) ↔ (∀ {A: Sort}, false = true → A) := by
  intros; lynx at *<;> apply Iff.intro <;> intro hyp  <;> (lynx at *) <;> (try assumption) <;> try (intros; apply Eq.symm; apply hyp)

example : (∀ (a b c: Nat),
  a + (b + c) = (a + b) + c) ↔ (∀ (a b c: Nat), (a + b) + c = a + (b + c)) := by
  intros; apply Iff.intro <;> intro hyp  <;> (try assumption) <;> try (intros; apply Eq.symm; apply hyp)

-- #eval compareThms "(a b c: Nat): a + (b + c) = (a + b) + c" "(a b c: Nat): (a + b) + c = a + (b + c)"


/-
Copyright (c) 2019 Simon Hudon. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Simon Hudon
-/
import control.traversable.derive
import control.traversable.lemmas
import data.dlist
import tactic.monotonicity.basic

variables {a b c p : Prop}

namespace tactic.interactive

open lean lean.parser  interactive
open interactive.types
open tactic

local postfix `?`:9001 := optional
local postfix *:9001 := many

meta inductive mono_function (elab : bool := tt)
 | non_assoc : expr elab → list (expr elab) → list (expr elab) → mono_function
 | assoc : expr elab → option (expr elab) → option (expr elab) → mono_function
 | assoc_comm : expr elab → expr elab → mono_function

meta instance : decidable_eq mono_function :=
by mk_dec_eq_instance

meta def mono_function.to_tactic_format : mono_function → tactic format
 | (mono_function.non_assoc fn xs ys) := do
  fn' ← pp fn,
  xs' ← mmap pp xs,
  ys' ← mmap pp ys,
  return format!"{fn'} {xs'} _ {ys'}"
 | (mono_function.assoc fn xs ys) := do
  fn' ← pp fn,
  xs' ← pp xs,
  ys' ← pp ys,
  return format!"{fn'} {xs'} _ {ys'}"
 | (mono_function.assoc_comm fn xs) := do
  fn' ← pp fn,
  xs' ← pp xs,
  return format!"{fn'} _ {xs'}"

meta instance has_to_tactic_format_mono_function : has_to_tactic_format mono_function :=
{ to_tactic_format := mono_function.to_tactic_format }

@[derive traversable]
meta structure ac_mono_ctx' (rel : Type) :=
  (to_rel : rel)
  (function : mono_function)
  (left right rel_def : expr)

@[reducible]
meta def ac_mono_ctx := ac_mono_ctx' (option (expr → expr → expr))
@[reducible]
meta def ac_mono_ctx_ne := ac_mono_ctx' (expr → expr → expr)

meta def ac_mono_ctx.to_tactic_format (ctx : ac_mono_ctx) : tactic format :=
do fn  ← pp ctx.function,
   l   ← pp ctx.left,
   r   ← pp ctx.right,
   rel ← pp ctx.rel_def,
   return format!"{{ function := {fn}\n, left  := {l}\n, right := {r}\n, rel_def := {rel} }}"

meta instance has_to_tactic_format_mono_ctx : has_to_tactic_format ac_mono_ctx :=
{ to_tactic_format := ac_mono_ctx.to_tactic_format }

meta def as_goal (e : expr) (tac : tactic unit) : tactic unit :=
do gs ← get_goals,
   set_goals [e],
   tac,
   set_goals gs

open list (hiding map) functor dlist

section config

parameter opt : mono_cfg
parameter asms : list expr

meta def unify_with_instance (e : expr) : tactic unit :=
as_goal e $
apply_instance
<|>
apply_opt_param
<|>
apply_auto_param
<|>
tactic.solve_by_elim { lemmas := some asms }
<|>
reflexivity
<|>
applyc ``id
<|>
return ()

private meta def match_rule_head  (p : expr)
: list expr → expr → expr → tactic expr
 | vs e t :=
(unify t p >> mmap' unify_with_instance vs >> instantiate_mvars e)
<|>
do (expr.pi _ _ d b) ← return t | failed,
   v ← mk_meta_var d,
   match_rule_head (v::vs) (expr.app e v) (b.instantiate_var v)

meta def pi_head : expr → tactic expr
| (expr.pi n _ t b) :=
do v ← mk_meta_var t,
   pi_head (b.instantiate_var v)
| e := return e

meta def delete_expr (e : expr)
: list expr → tactic (option (list expr))
 | [] := return none
 | (x :: xs) :=
(compare opt e x >> return (some xs))
<|>
(map (cons x) <$> delete_expr xs)

meta def match_ac'
: list expr → list expr → tactic (list expr × list expr × list expr)
 | es (x :: xs) := do
    es' ← delete_expr x es,
    match es' with
     | (some es') := do
       (c,l,r) ← match_ac' es' xs, return (x::c,l,r)
     | none := do
       (c,l,r) ← match_ac' es xs, return (c,l,x::r)
    end
 | es [] := do
return ([],es,[])

meta def match_ac (l : list expr) (r : list expr)
: tactic (list expr × list expr × list expr) :=
do (s',l',r') ← match_ac' l r,
   s' ← mmap instantiate_mvars s',
   l' ← mmap instantiate_mvars l',
   r' ← mmap instantiate_mvars r',
   return (s',l',r')

meta def match_prefix
: list expr → list expr → tactic (list expr × list expr × list expr)
| (x :: xs) (y :: ys) :=
  (do compare opt x y,
      prod.map ((::) x) id <$> match_prefix xs ys)
<|> return ([],x :: xs,y :: ys)
| xs ys := return ([],xs,ys)

/--
`(prefix,left,right,suffix) ← match_assoc unif l r` finds the
longest prefix and suffix common to `l` and `r` and
returns them along with the differences  -/
meta def match_assoc (l : list expr) (r : list expr)
: tactic (list expr × list expr × list expr × list expr) :=
do (pre,l₁,r₁) ← match_prefix l r,
   (suf,l₂,r₂) ← match_prefix (reverse l₁) (reverse r₁),
   return (pre,reverse l₂,reverse r₂,reverse suf)

meta def check_ac : expr → tactic (bool × bool × option (expr × expr × expr) × expr)
 | (expr.app (expr.app f x) y) :=
   do t ← infer_type x,
      a ← try_core $ to_expr ``(is_associative %%t %%f) >>= mk_instance,
      c ← try_core $ to_expr ``(is_commutative %%t %%f) >>= mk_instance,
      i ← try_core (do
          v ← mk_meta_var t,
          l_inst_p ← to_expr ``(is_left_id %%t %%f %%v),
          r_inst_p ← to_expr ``(is_right_id %%t %%f %%v),
          l_v ← mk_meta_var l_inst_p,
          r_v ← mk_meta_var r_inst_p ,
          l_id ← mk_mapp `is_left_id.left_id [some t,f,v,some l_v],
          mk_instance l_inst_p >>= unify l_v,
          r_id ← mk_mapp `is_right_id.right_id [none,f,v,some r_v],
          mk_instance r_inst_p >>= unify r_v,
          v' ← instantiate_mvars v,
          return (l_id,r_id,v')),
      return (a.is_some,c.is_some,i,f)
 | _ := return (ff,ff,none,expr.var 1)

meta def parse_assoc_chain' (f : expr) : expr → tactic (dlist expr)
 | e :=
 (do (expr.app (expr.app f' x) y) ← return e,
     is_def_eq f f',
     (++) <$> parse_assoc_chain' x <*> parse_assoc_chain' y)
<|> return (singleton e)

meta def parse_assoc_chain (f : expr) : expr → tactic (list expr) :=
map dlist.to_list ∘ parse_assoc_chain' f

meta def fold_assoc (op : expr) :
  option (expr × expr × expr) → list expr → option (expr × list expr)
| _ (x::xs) := some (foldl (expr.app ∘ expr.app op) x xs, [])
| none []   := none
| (some (l_id,r_id,x₀)) [] := some (x₀,[l_id,r_id])

meta def fold_assoc1 (op : expr) : list expr → option expr
| (x::xs) := some $ foldl (expr.app ∘ expr.app op) x xs
| []   := none

meta def same_function_aux
: list expr → list expr → expr → expr → tactic (expr × list expr × list expr)
 | xs₀ xs₁ (expr.app f₀ a₀) (expr.app f₁ a₁) :=
   same_function_aux (a₀ :: xs₀) (a₁ :: xs₁) f₀ f₁
 | xs₀ xs₁ e₀ e₁ := is_def_eq e₀ e₁ >> return (e₀,xs₀,xs₁)

meta def same_function : expr → expr → tactic (expr × list expr × list expr) :=
same_function_aux [] []

meta def parse_ac_mono_function (l r : expr)
: tactic (expr × expr × list expr × mono_function) :=
do (full_f,ls,rs) ← same_function l r,
   (a,c,i,f) ← check_ac l,
   if a
   then if c
   then do
     (s,ls,rs) ← monad.join (match_ac
                   <$> parse_assoc_chain f l
                   <*> parse_assoc_chain f r),
     (l',l_id) ← fold_assoc f i ls,
     (r',r_id) ← fold_assoc f i rs,
     s' ← fold_assoc1 f s,
     return (l',r',l_id ++ r_id,mono_function.assoc_comm f s')
   else do -- a ∧ ¬ c
     (pre,ls,rs,suff) ← monad.join (match_assoc
                   <$> parse_assoc_chain f l
                   <*> parse_assoc_chain f r),
     (l',l_id) ← fold_assoc f i ls,
     (r',r_id) ← fold_assoc f i rs,
     let pre'  := fold_assoc1 f pre,
     let suff' := fold_assoc1 f suff,
     return (l',r',l_id ++ r_id,mono_function.assoc f pre' suff')
   else do -- ¬ a
     (xs₀,x₀,x₁,xs₁) ← find_one_difference opt ls rs,
     return (x₀,x₁,[],mono_function.non_assoc full_f xs₀ xs₁)

meta def parse_ac_mono_function' (l r : pexpr) :=
do l' ← to_expr l,
   r' ← to_expr r,
   parse_ac_mono_function l' r'

meta def ac_monotonicity_goal : expr → tactic (expr × expr × list expr × ac_mono_ctx)
 | `(%%e₀ → %%e₁) :=
  do (l,r,id_rs,f) ← parse_ac_mono_function e₀ e₁,
     t₀ ← infer_type e₀,
     t₁ ← infer_type e₁,
     rel_def ← to_expr ``(λ x₀ x₁, (x₀ : %%t₀) → (x₁ : %%t₁)),
     return (e₀, e₁, id_rs,
            { function := f
            , left := l, right := r
            , to_rel := some $ expr.pi `x binder_info.default
            , rel_def := rel_def })
 | `(%%e₀ = %%e₁) :=
  do (l,r,id_rs,f) ← parse_ac_mono_function e₀ e₁,
     t₀ ← infer_type e₀,
     t₁ ← infer_type e₁,
     rel_def ← to_expr ``(λ x₀ x₁, (x₀ : %%t₀) = (x₁ : %%t₁)),
     return (e₀, e₁, id_rs,
            { function := f
            , left := l, right := r
            , to_rel := none
            , rel_def := rel_def })
 | (expr.app (expr.app rel e₀) e₁) :=
  do (l,r,id_rs,f) ← parse_ac_mono_function e₀ e₁,
     return (e₀, e₁, id_rs,
            { function := f
            , left := l, right := r
            , to_rel := expr.app ∘ expr.app rel
            , rel_def := rel })
 | _ := fail "invalid monotonicity goal"

meta def bin_op_left (f : expr)  : option expr → expr → expr
| none e := e
| (some e₀) e₁ := f.mk_app [e₀,e₁]

meta def bin_op (f a b : expr) : expr :=
f.mk_app [a,b]

meta def bin_op_right (f : expr) : expr → option expr → expr
| e none := e
| e₀ (some e₁) := f.mk_app [e₀,e₁]

meta def mk_fun_app : mono_function → expr → expr
 | (mono_function.non_assoc f x y) z := f.mk_app (x ++ z :: y)
 | (mono_function.assoc f x y) z := bin_op_left f x (bin_op_right f z y)
 | (mono_function.assoc_comm f x) z := f.mk_app [z,x]

meta inductive mono_law
   /- `assoc (l₀,r₀) (r₁,l₁)` gives first how to find rules to prove
      x+(y₀+z) R x+(y₁+z);
      if that fails, helps prove (x+y₀)+z R (x+y₁)+z -/
 | assoc : expr × expr → expr × expr → mono_law
   /- `congr r` gives the rule to prove `x = y → f x = f y` -/
 | congr : expr → mono_law
 | other : expr → mono_law

meta def mono_law.to_tactic_format : mono_law → tactic format
 | (mono_law.other e) := do e ← pp e, return format!"other {e}"
 | (mono_law.congr r) := do e ← pp r, return format!"congr {e}"
 | (mono_law.assoc (x₀,x₁) (y₀,y₁)) :=
do x₀ ← pp x₀,
   x₁ ← pp x₁,
   y₀ ← pp y₀,
   y₁ ← pp y₁,
   return format!"assoc {x₀}; {x₁} | {y₀}; {y₁}"

meta instance has_to_tactic_format_mono_law : has_to_tactic_format mono_law :=
{ to_tactic_format := mono_law.to_tactic_format }

meta def mk_rel (ctx : ac_mono_ctx_ne) (f : expr → expr) : expr :=
ctx.to_rel (f ctx.left) (f ctx.right)

meta def mk_congr_args (fn : expr) (xs₀ xs₁ : list expr) (l r : expr) : tactic expr :=
do p ← mk_app `eq [fn.mk_app $ xs₀ ++ l :: xs₁,fn.mk_app $ xs₀ ++ r :: xs₁],
   prod.snd <$> solve_aux p
     (do iterate_exactly (xs₁.length) (applyc `congr_fun),
         applyc `congr_arg)

meta def mk_congr_law (ctx : ac_mono_ctx) : tactic expr :=
match ctx.function with
 | (mono_function.assoc f x₀ x₁) :=
    if (x₀ <|> x₁).is_some
       then mk_congr_args f x₀.to_monad x₁.to_monad ctx.left ctx.right
       else failed
 | (mono_function.assoc_comm f x₀) := mk_congr_args f [x₀] [] ctx.left ctx.right
 | (mono_function.non_assoc f x₀ x₁) := mk_congr_args f x₀ x₁ ctx.left ctx.right
end

meta def mk_pattern (ctx : ac_mono_ctx) : tactic mono_law :=
match (sequence ctx : option (ac_mono_ctx' _)) with
 | (some ctx) :=
   match ctx.function with
    | (mono_function.assoc f (some x) (some y)) :=
      return $ mono_law.assoc
       ( mk_rel ctx (λ i, bin_op f x (bin_op f i y))
       , mk_rel ctx (λ i, bin_op f i y))
       ( mk_rel ctx (λ i, bin_op f (bin_op f x i) y)
       , mk_rel ctx (λ i, bin_op f x i))
    | (mono_function.assoc f (some x) none) :=
      return $ mono_law.other $
        mk_rel ctx (λ e, mk_fun_app ctx.function e)
    | (mono_function.assoc f none (some y)) :=
      return $ mono_law.other $
        mk_rel ctx (λ e, mk_fun_app ctx.function e)
    | (mono_function.assoc f none none) :=
      none
    | _ :=
      return $ mono_law.other $
         mk_rel ctx (λ e, mk_fun_app ctx.function e)
   end
 | none := mono_law.congr <$> mk_congr_law ctx
end

meta def match_rule (pat : expr) (r : name) : tactic expr :=
do  r' ← mk_const r,
    t  ← infer_type r',
    t  ← expr.dsimp t { fail_if_unchanged := ff } tt [] [
      simp_arg_type.expr ``(monotone), simp_arg_type.expr ``(strict_mono)],
    match_rule_head pat [] r' t

meta def find_lemma (pat : expr) : list name → tactic (list expr)
 | [] := return []
 | (r :: rs) :=
 do (cons <$> match_rule pat r <|> pure id) <*> find_lemma rs

meta def match_chaining_rules (ls : list name) (x₀ x₁ : expr) : tactic (list expr) :=
do x' ← to_expr ``(%%x₁ → %%x₀),
   r₀ ← find_lemma x' ls,
   r₁ ← find_lemma x₁ ls,
   return (expr.app <$> r₀ <*> r₁)

meta def find_rule (ls : list name) : mono_law → tactic (list expr)
 | (mono_law.assoc (x₀,x₁) (y₀,y₁)) :=
(match_chaining_rules ls x₀ x₁)
<|> (match_chaining_rules ls y₀ y₁)
 | (mono_law.congr r) := return [r]
 | (mono_law.other p) := find_lemma p ls

universes u v

def apply_rel {α : Sort u} (R : α → α → Sort v) {x y : α}
  (x' y' : α)
  (h : R x y)
  (hx : x = x')
  (hy : y = y')
: R x' y' :=
by { rw [← hx,← hy], apply h }

meta def ac_refine (e : expr) : tactic unit :=
refine ``(eq.mp _ %%e) ; ac_refl

meta def one_line (e : expr) : tactic format :=
do lbl ← pp e,
   asm ← infer_type e >>= pp,
   return format!"\t{asm}\n"

meta def side_conditions (e : expr) : tactic format :=
do let vs := e.list_meta_vars,
   ts ← mmap one_line vs.tail,
   let r := e.get_app_fn.const_name,
   return format!"{r}:\n{format.join ts}"

open monad

/-- tactic-facing function, similar to `interactive.tactic.generalize` with the
exception that meta variables -/
private meta def monotonicity.generalize' (h : name) (v : expr) (x : name) : tactic (expr × expr) :=
do tgt ← target,
   t ← infer_type v,
   tgt' ← do
   { ⟨tgt', _⟩ ← solve_aux tgt (tactic.generalize v x >> target),
     to_expr ``(λ y : %%t, Π x, y = x → %%(tgt'.binding_body.lift_vars 0 1)) }
   <|> to_expr ``(λ y : %%t, Π x, %%v = x → %%tgt),
   t ← head_beta (tgt' v) >>= assert h,
   swap,
   r ← mk_eq_refl v,
   solve1 $ tactic.exact (t v r),
   prod.mk <$> tactic.intro x <*> tactic.intro h

private meta def hide_meta_vars (tac : list expr → tactic unit) : tactic unit :=
focus1 $
do tgt ← target >>= instantiate_mvars,
   tactic.change tgt,
   ctx ← local_context,
   let vs := tgt.list_meta_vars,
   vs' ← mmap (λ v,
             do h ← get_unused_name `h,
                x ← get_unused_name `x,
                prod.snd <$> monotonicity.generalize' h v x) vs,
     tac ctx;
     vs'.mmap' (try ∘ tactic.subst)

meta def hide_meta_vars' (tac : itactic) : itactic :=
hide_meta_vars $ λ _, tac

end config

meta def solve_mvar (v : expr) (tac : tactic unit) : tactic unit :=
do gs ← get_goals,
   set_goals [v],
   target >>= instantiate_mvars >>= tactic.change,
   tac, done,
   set_goals $ gs

def list.minimum_on {α β} [linear_order β] (f : α → β) : list α → list α
| [] := []
| (x :: xs) := prod.snd $ xs.foldl (λ ⟨k,a⟩ b,
     let k' := f b in
     if k < k' then (k,a)
     else if k' < k then (k', [b])
     else (k,b :: a)) (f x, [x])

open format mono_selection

meta def best_match {β} (xs : list expr) (tac : expr → tactic β) : tactic unit :=
do t ← target,
   xs ← xs.mmap (λ x,
     try_core $ prod.mk x <$> solve_aux t (tac x >> get_goals)),
   let xs := xs.filter_map id,
   let r := list.minimum_on (list.length ∘ prod.fst ∘ prod.snd) xs,
   match r with
   | [(_,gs,pr)] :=  tactic.exact pr >> set_goals gs
   | [] := fail "no good match found"
   | _ :=
     do lmms ← r.mmap (λ ⟨l,gs,_⟩,
          do ts ← gs.mmap infer_type,
             msg ← ts.mmap pp,
             pure $ foldl compose "\n\n" $
               list.intersperse "\n" $ to_fmt l.get_app_fn.const_name :: msg),
        let msg := foldl compose "" lmms,
        fail format!("ambiguous match: {msg}\n\n" ++
          "Tip: try asserting a side condition to distinguish between the lemmas")
   end

meta def mono_aux (dir : parse side) :
  tactic unit :=
do t ← target >>= instantiate_mvars,
   ns ← get_monotonicity_lemmas t dir,
   asms ← local_context,
   rs ← find_lemma asms t ns,
   focus1 $ () <$ best_match rs (λ law, tactic.refine $ to_pexpr law)

/--
- `mono` applies a monotonicity rule.
- `mono*` applies monotonicity rules repetitively.
- `mono with x ≤ y` or `mono with [0 ≤ x,0 ≤ y]` creates an assertion for the listed
  propositions. Those help to select the right monotonicity rule.
- `mono left` or `mono right` is useful when proving strict orderings:
   for `x + y < w + z` could be broken down into either
    - left:  `x ≤ w` and `y < z` or
    - right: `x < w` and `y ≤ z`
- `mono using [rule1,rule2]` calls `simp [rule1,rule2]` before applying mono.
- The general syntax is
  `mono '*'? ('with' hyp | 'with' [hyp1,hyp2])? ('using' [hyp1,hyp2])? mono_cfg?`

To use it, first import `tactic.monotonicity`.

Here is an example of mono:

```lean
example (x y z k : ℤ)
  (h : 3 ≤ (4 : ℤ))
  (h' : z ≤ y) :
  (k + 3 + x) - y ≤ (k + 4 + x) - z :=
begin
  mono, -- unfold `(-)`, apply add_le_add
  { -- ⊢ k + 3 + x ≤ k + 4 + x
    mono, -- apply add_le_add, refl
    -- ⊢ k + 3 ≤ k + 4
    mono },
  { -- ⊢ -y ≤ -z
    mono /- apply neg_le_neg -/ }
end
```

More succinctly, we can prove the same goal as:

```lean
example (x y z k : ℤ)
  (h : 3 ≤ (4 : ℤ))
  (h' : z ≤ y) :
  (k + 3 + x) - y ≤ (k + 4 + x) - z :=
by mono*
```

-/
meta def mono (many : parse (tk "*")?)
  (dir : parse side)
  (hyps : parse $ tk "with" *> pexpr_list_or_texpr <|> pure [])
  (simp_rules : parse $ tk "using" *> simp_arg_list <|> pure []) :
  tactic unit :=
do hyps ← hyps.mmap (λ p, to_expr p >>= mk_meta_var),
   hyps.mmap' (λ pr, do h ← get_unused_name `h, note h none pr),
   when (¬ simp_rules.empty) (simp_core { } failed tt simp_rules [] (loc.ns [none]) >> skip),
   if many.is_some
     then repeat $ mono_aux dir
     else mono_aux dir,
   gs ← get_goals,
   set_goals $ hyps ++ gs

add_tactic_doc
{ name       := "mono",
  category   := doc_category.tactic,
  decl_names := [`tactic.interactive.mono],
  tags       := ["monotonicity"] }

/--
transforms a goal of the form `f x ≼ f y` into `x ≤ y` using lemmas
marked as `monotonic`.

Special care is taken when `f` is the repeated application of an
associative operator and if the operator is commutative
-/
meta def ac_mono_aux (cfg : mono_cfg := { mono_cfg . }) :
  tactic unit :=
hide_meta_vars $ λ asms,
do try `[simp only [sub_eq_add_neg]],
   tgt ← target >>= instantiate_mvars,
   (l,r,id_rs,g) ← ac_monotonicity_goal cfg tgt
             <|> fail "monotonic context not found",
   ns ← get_monotonicity_lemmas tgt both,
   p ← mk_pattern g,
   rules ← find_rule asms ns p <|> fail "no applicable rules found",
   when (rules = []) (fail "no applicable rules found"),
   err ← format.join <$> mmap side_conditions rules,
   focus1 $ best_match rules (λ rule, do
     t₀ ← mk_meta_var `(Prop),
     v₀ ← mk_meta_var t₀,
     t₁ ← mk_meta_var `(Prop),
     v₁ ← mk_meta_var t₁,
     tactic.refine $ ``(apply_rel %%(g.rel_def) %%l %%r %%rule %%v₀ %%v₁),
     solve_mvar v₀ (try (any_of id_rs rewrite_target) >>
             ( done <|>
               refl <|>
               ac_refl <|>
               `[simp only [is_associative.assoc]]) ),
     solve_mvar v₁ (try (any_of id_rs rewrite_target) >>
             ( done <|>
               refl <|>
               ac_refl <|>
               `[simp only [is_associative.assoc]]) ),
     n ← num_goals,
     iterate_exactly (n-1) (try $ solve1 $ apply_instance <|>
       tactic.solve_by_elim { lemmas := some asms }))

open sum nat

/-- (repeat_until_or_at_most n t u): repeat tactic `t` at most n times or until u succeeds -/
meta def repeat_until_or_at_most : nat → tactic unit → tactic unit → tactic unit
| 0        t _ := fail "too many applications"
| (succ n) t u := u <|> (t >> repeat_until_or_at_most n t u)

meta def repeat_until : tactic unit → tactic unit → tactic unit :=
repeat_until_or_at_most 100000

@[derive _root_.has_reflect, derive _root_.inhabited]
inductive rep_arity : Type
| one | exactly (n : ℕ) | many

meta def repeat_or_not : rep_arity → tactic unit → option (tactic unit) → tactic unit
 | rep_arity.one  tac none := tac
 | rep_arity.many tac none := repeat tac
 | (rep_arity.exactly n) tac none := iterate_exactly' n tac
 | rep_arity.one  tac (some until) := tac >> until
 | rep_arity.many tac (some until) := repeat_until tac until
 | (rep_arity.exactly n) tac (some until) := iterate_exactly n tac >> until

meta def assert_or_rule : lean.parser (pexpr ⊕ pexpr) :=
(tk ":=" *> inl <$> texpr <|> (tk ":" *> inr <$> texpr))

meta def arity : lean.parser rep_arity :=
rep_arity.many <$ tk "*" <|>
rep_arity.exactly <$> (tk "^" *> small_nat) <|>
pure rep_arity.one

/--

`ac_mono` reduces the `f x ⊑ f y`, for some relation `⊑` and a
monotonic function `f` to `x ≺ y`.

`ac_mono*` unwraps monotonic functions until it can't.

`ac_mono^k`, for some literal number `k` applies monotonicity `k`
times.

`ac_mono := h`, with `h` a hypothesis, unwraps monotonic functions and
uses `h` to solve the remaining goal. Can be combined with `*` or `^k`:
`ac_mono* := h`

`ac_mono : p` asserts `p` and uses it to discharge the goal result
unwrapping a series of monotonic functions. Can be combined with * or
^k: `ac_mono* : p`

In the case where `f` is an associative or commutative operator,
`ac_mono` will consider any possible permutation of its arguments and
use the one the minimizes the difference between the left-hand side
and the right-hand side.

To use it, first import `tactic.monotonicity`.

`ac_mono` can be used as follows:

```lean
example (x y z k m n : ℕ)
  (h₀ : z ≥ 0)
  (h₁ : x ≤ y) :
  (m + x + n) * z + k ≤ z * (y + n + m) + k :=
begin
  ac_mono,
  -- ⊢ (m + x + n) * z ≤ z * (y + n + m)
  ac_mono,
  -- ⊢ m + x + n ≤ y + n + m
  ac_mono,
end
```

As with `mono*`, `ac_mono*` solves the goal in one go and so does
`ac_mono* := h₁`. The latter syntax becomes especially interesting in the
following example:

```lean
example (x y z k m n : ℕ)
  (h₀ : z ≥ 0)
  (h₁ : m + x + n ≤ y + n + m) :
  (m + x + n) * z + k ≤ z * (y + n + m) + k :=
by ac_mono* := h₁.
```

By giving `ac_mono` the assumption `h₁`, we are asking `ac_refl` to
stop earlier than it would normally would.
-/
meta def ac_mono (rep : parse arity) :
         parse assert_or_rule? →
         opt_param mono_cfg { mono_cfg . } →
         tactic unit
 | none opt := focus1 $ repeat_or_not rep (ac_mono_aux opt) none
 | (some (inl h)) opt :=
do focus1 $ repeat_or_not rep (ac_mono_aux opt) (some $ done <|> to_expr h >>= ac_refine)
 | (some (inr t)) opt :=
do h ← i_to_expr t >>= assert `h,
   tactic.swap,
   focus1 $ repeat_or_not rep (ac_mono_aux opt) (some $ done <|> ac_refine h)
/-
TODO(Simon): with `ac_mono := h` and `ac_mono : p` split the remaining
  gaol if the provided rule does not solve it completely.
-/

add_tactic_doc
{ name       := "ac_mono",
  category   := doc_category.tactic,
  decl_names := [`tactic.interactive.ac_mono],
  tags       := ["monotonicity"] }

attribute [mono] and.imp or.imp

end tactic.interactive


import its reducibility

open encodable denumerable

attribute [simp] set.set_of_app_iff

variables {α : Type*} {β : Type*} {γ : Type*} {σ : Type*} {τ : Type*} {μ : Type*}
  [primcodable α] [primcodable β] [primcodable γ] [primcodable σ] [primcodable τ] [primcodable μ]
  {o : σ →. τ}
variables {n : ℕ} (S : strategy n) {k : ℕ}

open strategy encodable

def ancestor.equiv_fin {k : ℕ} (η : Tree k) :
ancestor η ≃ fin η.length :=
{ to_fun := (λ μ, ⟨μ.val.length, list.is_initial_length μ.property⟩),
  inv_fun := (λ n, ⟨η↾*n.val, list.is_initial_initial η n.val n.property⟩),
  left_inv := (λ ⟨μ, ⟨l, x, p⟩⟩,
    by { simp[←p],  simp only [show l ++ x :: μ = (l ++ [x]) ++ μ, by simp, list.initial_append] }),
  right_inv := (λ ⟨n, lt⟩, by { simp, exact list.initial_length lt }) }

instance primcodable.ancestor_arrow {α : Type*} [primcodable α] {k} (η : Tree k) : primcodable (ancestor η → α) :=
primcodable.of_equiv (fin η.length → α) (equiv.arrow_congr (ancestor.equiv_fin η) (by refl))

@[rcomputability]
lemma rcomputable.is_pi {k} : @Tree'.is_pi k computable_in o :=
begin
  induction k with k IH,
  { exact rcomputable.id.of_eq (λ b, by cases b; simp) },
  { let F : Tree' (k + 1) → bool := λ μ, list.cases_on μ ff (λ η _, !@Tree'.is_pi k η),
    have : F computable_in o,
    { refine (rcomputable.list_rec (rcomputable.id') (rcomputable.const ff)
      ((rcomputable.dom_fintype bnot).comp (IH.comp (rcomputable.fst.comp rcomputable.snd)))) },
    exact this.of_eq (λ μ, by { induction μ with ν μ IH; simp[F, Tree'.is_sigma] }) }
end

@[rcomputability]
lemma rcomputable.is_sigma {k} : @Tree'.is_sigma k computable_in o :=
(rcomputable.dom_fintype bnot).comp rcomputable.is_pi

def ancestor.extend_fn'_enc (α : Type*) [inhabited α] [primcodable α] (η : Tree k) (x : Tree' k) (n : ℕ) : ℕ :=
  let f := (decode (ancestor (x :: η) → α) n).iget in
encode (ancestor.extend_fn f η (list.suffix_cons x η))

namespace strategy

namespace approx_enc

def derivative (η : Tree (k + 1)) (μ : Tree k) (υ : list (Tree (k + 1))) : list ℕ :=
(list.range_r μ.length).filter (λ i, υ.nth i = η)

def pi_derivative
  (η : Tree (k + 1)) (μ : Tree k) (υ : list (Tree (k + 1))) : list ℕ :=
(derivative η μ υ).filter (λ i, @Tree'.is_sigma (k + 1) (μ↾*(i + 1)))

def lambda : Π (μ : Tree k) (υ : list (Tree (k + 1))), Tree (k + 1)
| []       _ := []
| (x :: μ) υ := let ih := lambda μ υ in
    option.cases_on (υ.nth μ.length) []
    (λ uμ, if uμ = ih ∨ (x.is_pi ∧ pi_derivative uμ μ υ = [])
    then (x :: μ) :: uμ else ih)

def assignment (μ : Tree k) (υ : list (Tree (k + 1))) : Tree (k + 1) × ℕ :=
(S.priority (k + 1)).Min_le
  ((lambda μ υ, 0) :: ((list.range_r (lambda μ υ).length).filter
    (λ i, @Tree'.is_sigma (k + 2) (lambda μ υ↾*(i + 1)))).map
  (λ i, ((lambda μ υ)↾*i, (derivative (lambda μ υ↾*i) μ υ).length))) (by simp)

def up (μ : Tree k) (υ : list (Tree (k + 1))) : Tree (k + 1) :=
(assignment S μ υ).1

lemma ancestors_eq_range (μ : Tree k) : μ.ancestors.map ancestor.index = list.range_r μ.length :=
by { induction μ with ν μ IH; simp[ancestor.index, (∘)], exact IH }

lemma derivative_eq {μ : Tree k} {υ : ancestor μ → Tree (k + 1)} {υ' : list (Tree (k + 1))}
  (h : ∀ ν : ancestor μ, υ'.nth ν.index = some (υ ν)) (η : Tree (k + 1)) :
  derivative η μ υ' = (approx.derivative η υ).map ancestor.index :=
by { simp[derivative, approx.derivative, ←ancestors_eq_range, list.map_filter, (∘)],
     congr, funext ν, simp [h] }

lemma pi_derivative_eq {μ : Tree k} {υ : ancestor μ → Tree (k + 1)} {υ' : list (Tree (k + 1))}
  (h : ∀ ν : ancestor μ, υ'.nth ν.index = some (υ ν)) (η : Tree (k + 1)) :
  pi_derivative η μ υ' = (approx.pi_derivative η υ).map ancestor.index :=
by { simp[pi_derivative, approx.pi_derivative, derivative_eq h, list.map_filter, (∘)],
     congr, funext ν, simp[ancestor_initial_index_succ] }

lemma lambda_eq {μ : Tree k} {υ : ancestor μ → Tree (k + 1)} {υ' : list (Tree (k + 1))}
  (h : ∀ ν : ancestor μ, υ'.nth ν.index = some (υ ν)) :
  lambda μ υ' = approx.lambda υ :=
begin
  induction μ with ν μ IH; simp[lambda, approx.lambda],
  rw [show υ'.nth μ.length = some (υ ⟨μ, by simp⟩), from h ⟨μ, by simp⟩], simp,
  have la_eq : lambda μ υ' = approx.lambda (ancestor.extend_fn υ μ _),
    from @IH (ancestor.extend_fn υ μ (by simp)) (λ σ, h (σ.extend (by simp))),
  have pider_eq : ∀ η,
    pi_derivative η μ υ' = list.map ancestor.index (approx.pi_derivative η (ancestor.extend_fn υ μ _)),
    from @pi_derivative_eq _ _ (ancestor.extend_fn υ μ (by simp)) υ'
    (λ σ, h (σ.extend (by simp))),
  simp[la_eq, pider_eq]
end

lemma assignment_eq {μ : Tree k} {υ : ancestor μ → Tree (k + 1)} {υ' : list (Tree (k + 1))}
  (h : ∀ ν : ancestor μ, υ'.nth ν.index = some (υ ν)) :
  assignment S μ υ' = approx.assignment S υ :=
begin
  simp[assignment, approx.assignment, lambda_eq h],
  refine omega_ordering.Min_le_eq _ _ _ _,
  simp,
  rw [←ancestors_eq_range],
  simp[list.map_filter, (∘)],
  congr,
  { funext ν, simp[ancestor_initial_index, derivative_eq h] },
  { ext ν, simp[ancestor_initial_index_succ] }
end

lemma up_eq {μ : Tree k} {υ : ancestor μ → Tree (k + 1)} {υ' : list (Tree (k + 1))}
  (h : ∀ ν : ancestor μ, υ'.nth ν.index = some (υ ν)) :
  up S μ υ' = approx.up S υ :=
congr_arg prod.fst (assignment_eq S h)

variables {S}
open rcomputable computable₂

lemma rcomputable.derivative :
  (prod.unpaired3 (derivative : Tree (k + 1) → Tree k → list (Tree (k + 1)) → list ℕ)) computable_in o :=
begin
  refine rcomputable.list_filter _ _,
  { refine (rcomputable₂.to_bool_eq (option (Tree (k + 1)))).comp₂ _ _,
    { exact rcomputable₂.list_nth.comp₂ (rcomputable.to_unary₁ rcomputable.snd.to_unary₂) rcomputable.id'.to_unary₂ },
    { exact rcomputable.to_unary₁ rcomputable.option_some.to_unary₁ } },
  { exact rcomputable.list_range_r.comp (rcomputable.list_length.comp rcomputable.fst.to_unary₂) }
end

lemma rcomputable.pi_derivative :
  (prod.unpaired3 (pi_derivative : Tree (k + 1) → Tree k → list (Tree (k + 1)) → list ℕ)) computable_in o :=
begin
  refine rcomputable.list_filter _ rcomputable.derivative,
  { simp, exact rcomputable.is_sigma.comp₂
    (rcomputable.rcomputable.list_initial.comp₂ (rcomputable.to_unary₁ rcomputable.fst.to_unary₂)
    rcomputable.succ.to_unary₂) }
end

lemma rcomputable.lambda :
  (lambda : Tree k → list (Tree (k + 1)) → Tree (k + 1)) computable₂_in o :=
begin
  let F : Tree k → list (Tree (k + 1)) → Tree (k + 1) :=
    λ μ υ, list.rec_on μ []
      (λ x μ ih, option.cases_on (υ.nth μ.length) []
        (λ uμ, if uμ = ih ∨ (x.is_pi ∧ pi_derivative uμ μ υ = []) then (x :: μ) :: uμ else ih)),
  have : F computable₂_in o,
  { simp[F],
    refine rcomputable.list_rec (rpartrec.some.to_unary₁) (rcomputable.const list.nil)
      (rcomputable.option_rec
        (rcomputable₂.list_nth.comp snd.to_unary₁
          (rcomputable.list_length.comp (fst.comp snd.to_unary₂)))
        (rcomputable.const list.nil)
      (rcomputable.ite _ _ _)),
    { simp, refine rcomputable.bor.comp _
        (rcomputable.band.comp (rcomputable.is_pi.comp (fst.comp snd.to_unary₁)) _),
      { refine (rcomputable₂.to_bool_eq _).comp snd
        (snd.comp (snd.comp snd.to_unary₁)) },
      { refine (rcomputable₂.to_bool_eq _).comp
          (rcomputable.pi_derivative.unpaired3 snd
            (fst.comp (snd.comp snd.to_unary₁))
            (snd.comp fst.to_unary₁)) (rcomputable.const []) } },
    { exact rcomputable₂.list_cons.comp
        (rcomputable₂.list_cons.comp (fst.comp snd.to_unary₁)
        (fst.comp (snd.comp snd.to_unary₁)))
        snd },
    { exact snd.comp (snd.comp snd.to_unary₁) } },
  exact this.of_eq (λ μ υ, by {
    induction μ with x μ IH; simp[F, lambda],
    { cases C : υ.nth μ.length with ν; simp[C],
      { simp[F, lambda] at IH, simp[IH] } } })
end

lemma rcomputable.assignment_enc :
  (assignment S : Tree k → list (Tree (k + 1)) → Tree (k + 1) × ℕ) computable₂_in o :=
(omega_ordering_Min_le (S.priority (k + 1)) _ _ (S.effective _).to_rcomp
  (rcomputable₂.list_cons.comp (pair (rcomputable.lambda.comp fst snd) (const 0))
    (list_map
      (pair (rcomputable.list_initial.comp (rcomputable.lambda.comp fst.to_unary₁ snd.to_unary₁) snd)
        (list_length.comp (rcomputable.derivative.unpaired3
          (rcomputable.list_initial.comp (rcomputable.lambda.comp fst.to_unary₁ snd.to_unary₁) snd)
          (fst.to_unary₁) (snd.to_unary₁))))
      (list_filter
        (by {simp, exact is_sigma.comp₂
          (rcomputable.list_initial.comp₂ (rcomputable.lambda.comp fst snd).to_unary₁ succ.to_unary₂) })
        (list_range_r.comp (list_length.comp (rcomputable.lambda.comp fst snd)))))))

lemma rcomputable.up_enc :
  (up S : Tree k → list (Tree (k + 1)) → Tree (k + 1)) computable₂_in o :=
fst.comp rcomputable.assignment_enc

end approx_enc

def up'_enc : Tree k → list (Tree (k + 1))
| []       := []
| (_ :: η) := up'_enc η ++ [approx_enc.up S η (up'_enc η)]

@[simp] lemma up'_enc_length (μ : Tree k) : (up'_enc S μ).length = μ.length :=
by induction μ with ν μ IH; simp[up'_enc]; exact IH

variables {S}
open rcomputable rcomputable₂

lemma rcomputable.up'_enc : (up'_enc S : Tree k → list (Tree (k + 1))) computable_in o :=
begin
  let F : Tree k → list (Tree (k + 1)) := λ μ, list.rec_on μ [] (λ _ η IH, IH ++ [approx_enc.up S η IH]),
  have : F computable_in o,
  { simp[F],
    refine rcomputable.list_rec rcomputable.id' (const []) _,
    exact list_append.comp (snd.comp snd.to_unary₂)
      (list_cons.comp (approx_enc.rcomputable.up_enc.comp (fst.comp snd.to_unary₂) (snd.comp snd.to_unary₂))
      (const [])) },
  exact this.of_eq (λ μ, by {
    induction μ with ν μ IH; simp[F, up'_enc],
    { simp[F] at IH, simp[IH] }  })
end

lemma up_enc_eq_up {μ : Tree k} : ∀ ν : ancestor μ, (S.up'_enc μ).nth ν.index = some (S.up' μ ν) :=
begin
  induction μ with ν μ IH; simp[up'_enc, up' S, -up'_up_consistent],
    {  rintros ⟨σ, lt⟩,simp at lt, contradiction },
    { rintros ⟨σ, lt⟩, simp[up', -up'_up_consistent],
      have : σ = μ ∨ σ ⊂ᵢ μ, from list.is_initial_cons_iff.mp lt,
      rcases this with (rfl | lt),
      { simp[ancestor.index],
        rw [show σ.length = (S.up'_enc σ).length, by simp, list.nth_concat_length],
        simp, refine approx_enc.up_eq S IH },
      { simp[lt, ancestor.index],
        rw [list.nth_append (show list.length σ < (S.up'_enc μ).length, by simp[list.is_initial_length lt])],
        have := IH ⟨σ, lt⟩, simp[ancestor.index] at this, exact this } }
end

@[rcomputability]
theorem rcomputable.up : (up[S] : Tree k → Tree (k + 1)) computable_in o :=
(approx_enc.rcomputable.up_enc.comp id' rcomputable.up'_enc).of_eq
(λ μ, approx_enc.up_eq S up_enc_eq_up)

@[rcomputability]
theorem rcomputable.lambda : (λ[S] : Tree k → Tree (k + 1)) computable_in o :=
by { have : (λ μ : Tree k, approx_enc.lambda μ (S.up'_enc μ)) computable_in o,
       from approx_enc.rcomputable.lambda.comp id' rcomputable.up'_enc,
     exact this.of_eq (λ μ, approx_enc.lambda_eq up_enc_eq_up) }

lemma rcomputable.Tree'_weight_aux {k : ℕ} :
  (Tree'.weight_aux : Tree' k → ℕ) computable_in o :=
begin
  induction k with k IH,
  { exact (id'.cond (const 1) (const 0)).of_eq (λ μ, by cases μ; simp[Tree'.weight_aux]) },
  { exact list_weight_of IH }
end

@[rcomputability]
lemma rcomputable.Tree_weight {k : ℕ} :
  (Tree.weight : Tree k → ℕ) computable_in o :=
begin
  induction k with k IH,
  { exact rcomputable.Tree'_weight_aux },
  { let F : Tree (k + 1) → ℕ := λ μ, list.rec_on μ 0 (λ ν _ _, ν.weight_aux + 1),
    have : F computable_in o,
    { simp[F], refine rcomputable.list_rec id (const 0)
      (nat_add.comp (rcomputable.Tree'_weight_aux.comp fst.to_unary₂) (const 1)) },
    exact this.of_eq (λ μ, by cases μ; simp[F, Tree.weight]) }
end

end strategy


import Lean
import Crypto.Util

open Lean Elab Tactic Meta

namespace Crypto

/-- Return true iff `constName` is the a non-recursive inductive datatype without indices
  that has only one constructor. If `includeProp` is false, we also check that the datatype is
  not `Prop` or `Sort*` -/
def isVeryStructureLike (env : Environment) (constName : Name) (includeProp := true) : Bool :=
  match env.find? constName with
  | some d@(ConstantInfo.inductInfo { isRec := false, ctors := [ctor], numIndices := 0, .. }) =>
    includeProp || d.toConstantVal.type.getForallBody.level!.isASucc
  | _ => false

/-- Good names for new hypotheses when casing on a variable with name `nm` and type `constName`. -/
def giveNames (env : Environment) (constName nm : Name) : Array Name :=
  match getStructureInfo? env constName with
  | some p => p.fieldNames.map λ s => nm.appendAfter ("_" ++ s.toString)
  | _ => if constName == `Exists then #[nm.appendAfter "_val", nm.appendAfter "_prop"] else #[]

/- the data we record for each field -/
structure FieldData where
(decl : LocalDecl)
(fvar : Expr)
(type : Expr)
(depends : NameSet)
(isProp : Bool)
deriving Inhabited

/- additionally an expression that gives us the translation to another field -/
structure FieldMapping extends FieldData where
(tgt : Option Expr)
deriving Inhabited

structure Context where
(target : Array FieldData)
deriving Inhabited

structure State where
(mapping : Array FieldMapping)
deriving Inhabited

open Std.Format

instance : ToFormat NameSet :=
⟨λ x => format x.toList⟩

instance : ToFormat FieldData :=
⟨λ ⟨a, b, c, d, e⟩ => group (nest 1 (format "⟨"  ++ format a ++ format "," ++ line ++ format b ++
format "," ++ line ++ format c ++ format "," ++ line ++ format d ++ format "," ++ line ++ format e
++ format "⟩"))⟩

namespace FieldData

def name (fields : FieldData) : Name :=
fields.decl.userName.head!

def toFieldMapping (fields : FieldData) : FieldMapping :=
{ toFieldData := fields, tgt := none }

end FieldData

namespace FieldMapping
def update (field : FieldMapping) (newM : Option Expr) : FieldMapping :=
if field.tgt.isNone then { field with tgt := newM } else field

def updateM (field : FieldMapping) (tac : MetaM (Option Expr)) : MetaM FieldMapping :=
if field.tgt.isNone then do
  let tgt ← tac
  return { field with tgt := tgt }
else
  return field

end FieldMapping

def ppFieldData (x : FieldData) : MetaM Format :=
do return group (nest 1 (format "⟨"  ++
format x.name ++ format "," ++ line ++
format (← Meta.ppExpr x.fvar) ++ format "," ++ line ++
format (← Meta.ppExpr x.type) ++ format "," ++ line ++
format (← x.depends.toList.mapM λ e => Meta.ppExpr (mkFVar ⟨e⟩)) ++ format "," ++ line ++
format x.isProp ++ format "⟩"))

/-- make a field mapping with no connections. -/
def mkState (fields1 : Array FieldData) : State :=
{ mapping := fields1.map (·.toFieldMapping) }

namespace State

def done (st : State) : Bool :=
st.mapping.all (·.tgt.isSome)

def update (st : State) (f : FieldData → Option Expr) : State :=
{ mapping := st.mapping.map λ info => info.update (f info.toFieldData) }

def updateM (st : State) (f : FieldData → MetaM (Option Expr)) : MetaM State := do
  let fieldMapping ← st.mapping.mapM λ info => info.updateM (f info.toFieldData)
  return { mapping := fieldMapping }

def getDataMapping (st : State) : Array (Expr × Expr) :=
st.mapping.filterMap λ map => match map.tgt with
  | none => none
  | some e => if map.isProp then none else some (map.fvar, e)

def getMapping (st : State) : Array (Expr × Expr) :=
st.mapping.filterMap λ map => match map.tgt with
  | none => none
  | some e => some (map.fvar, e)

def missing (st : State) : Array FieldData :=
st.mapping.filterMap λ info => if info.tgt.isNone then some info.toFieldData else none

def missingData (st : State) : Array FieldData :=
st.mapping.filterMap λ info =>
  if info.tgt.isNone && !info.isProp then some info.toFieldData else none

def traceMapping (st : State) : MetaM (Array Format) :=
st.mapping.mapM λ info : FieldMapping =>
  if info.tgt.isSome then Meta.ppExpr info.tgt.get! else return Format.nil

end State

namespace Meta

def fieldDataofCaseGoals (l : Array CasesSubgoal) : MetaM (MVarId × Array FieldData) := do
  let casegoal := l.get! 0
  let mvarId := casegoal.mvarId
  withMVarContext mvarId do
  let lctx ← getLCtx
  let fieldExprs := casegoal.fields
  let fields := fieldExprs.map Expr.fvarId!
  let ldecls := fields.map λ e => lctx.fvarIdToDecl.find! e
  let axiom_fields ← fieldExprs.mapM isProof
  let types ← fieldExprs.mapM inferType
  let depends := types.map λ tp => tp.ListFvarIds
  let fieldData := ldecls.zipWith5 FieldData.mk fieldExprs types depends axiom_fields
  return (mvarId, fieldData)

def updateFieldData (l : Array CasesSubgoal) (fields : Array FieldData) :
  MetaM (Array FieldData) := do
  let casegoal := l.get! 0
  let mvarId := casegoal.mvarId
  let fieldData ← withMVarContext mvarId $ fields.mapM λ info =>
    match casegoal.subst.find? info.fvar.fvarId! with
    | (some e) => do
      let lctx ← getLCtx
      let id := e.fvarId!
      let ldecl := lctx.fvarIdToDecl.find! id
      let t ← inferType e
      let depends := NameSet.empty -- todo
      return ⟨ldecl, e, t, depends, info.isProp⟩
    | none => return info
  return fieldData

def AddFieldsToContext (mvarId : MVarId) (nm : Name) (us : List Level) (args : Array Expr) :
  MetaM (MVarId × MVarId × Array FieldData) := do
  let env ← getEnv
  let d ← env.find? nm
  let true ← isStructureLike env nm
  let eStr ← mkAppN (mkConst nm us) args
  let (h, mvarId, m2) ← assertm mvarId `h eStr
  let l ← cases mvarId h
  let (mvarId, fieldData) ← fieldDataofCaseGoals l
  return (mvarId, m2, fieldData)

/-- map the data fields to data fields of the same name. -/
def trivialMapping (st : State) (ctx : Context) : State :=
st.update $ λ info =>
  if info.isProp then none else
    ctx.target.findSome? λ info' => if info.name == info'.name then some info'.fvar else none

/-- map the data fields to data fields with the same type, if a unique such data field exists. -/
def uniqueMapping (st : State) (ctx : Context) : MetaM State :=
st.updateM $ λ info => if info.isProp then return none else do
  let sources ← st.mapping.filterM λ info' => isDefEq info.type info'.type
  let targets ← ctx.target.filterM λ info' => isDefEq info.type info'.type
  return (if sources.size = 1 ∧ targets.size = 1 then some (targets.get! 0).fvar else none)

set_option pp.all true
partial def caseOnStructures (mvarId : MVarId) (st : Array FieldData) (ctx : Context)
  (includeProp := true) :
  MetaM (MVarId × Array FieldData) := do
  let env ← getEnv
  let info? := st.find? λ info =>
    info.type.getAppFn.constName?.any (isVeryStructureLike env · includeProp)
  match info? with
  | some info => do
    let rest := st.filter λ info' => info'.fvar.fvarId! != info.fvar.fvarId!
    let str := info.type.getAppFn.constName!
    -- todo: there seems to be a bug with the given names when transitively extending structures. Do we need to skip fields?
    -- IO.println (giveNames env str info.name)
    let l ← cases mvarId info.fvar.fvarId! #[⟨false, (giveNames env str info.name).toList⟩]
    let (mvarId, fieldData) ← fieldDataofCaseGoals l
    let rest ← updateFieldData l rest
    caseOnStructures mvarId (rest ++ fieldData) ctx includeProp
  | none =>
    return (mvarId, st)

/-- Find which axioms of the first structure that occur in the second structure. -/
def matchingAxioms (st : State) (ctx : Context) : MetaM State := do
  let targetTypes : Array (Expr × Expr) := ctx.target.filterMap λ info =>
    if info.isProp then (info.fvar, info.type) else none
  st.updateM λ info => do
    let e := info.type.instantiateFVars st.getDataMapping
    match (← targetTypes.findM? λ (e', t) => isDefEq e t) with
    | some (e, t) => return (some e)
    | none => return none

/-- some copy-pasted simp code -/
def getPropHyps : MetaM (Array FVarId) := do
  let mut result := #[]
  for localDecl in (← getLCtx) do
    unless localDecl.isAuxDecl do
      if (← isProp localDecl.type) then
        result := result.push localDecl.fvarId
  return result

/-- some mostly copy-pasted simp code -/
def mkSimpContext (simpOnly := false) (starArg := true) : MetaM Simp.Context := do
  let ctx : Simp.Context :=
  { config      := {}
    simpLemmas  := if simpOnly then {} else (← getSimpLemmas)
    congrLemmas := (← getCongrLemmas) }
  if !starArg then
    return ctx
  else
    let hs ← getPropHyps
    let mut ctx := ctx
    for h in hs do
      let localDecl ← getLocalDecl h
      let fvarId := localDecl.fvarId
      let proof  := localDecl.toExpr
      let id     ← mkFreshUserName `h
      let simpLemmas ← ctx.simpLemmas.add #[] proof (name? := id)
      ctx := { ctx with simpLemmas }
    return ctx

/-- The tactic we use to automatically prove axioms. -/
def currentAutomation (mvarId : MVarId) : MetaM Unit := do
let newgoal ← simpTarget mvarId (← mkSimpContext false)

-- match (← simpTarget mvarId (← mkSimpContext false)) with
-- | some x => IO.println "failure!"
-- | none => IO.println "success!"

/-- Tries to prove `e` in the local context, returns the proof if successful. -/
def tryToProve (mvarId : MVarId) (tac : MVarId → MetaM Unit) (e : Expr) : MetaM (Option Expr) :=
withoutModifyingState do
  let (fvar, mvarId, m2) ← assertm mvarId `h e
  try
    tac m2
    -- IO.println s!"is assigned: {← isExprMVarAssigned m2}"
    let some e ← getExprMVarAssignment? m2 | return none
    let e ← instantiateMVars e
    if (← e.hasExprMVar) then return none else return some e
  catch err =>
    return none

/-- Tests whether nm1 is a subclass of nm1. Currently the data fields must have the same Name for
this tactic to work. -/
def isSubclass (mvarId : MVarId) (nm1 nm2 : Name) (trace := false) :
  MetaM (MVarId × MVarId × MVarId × State) := do
  let u := mkLevelParam `u
  let (M, mvarId) ← asserti mvarId `M (mkSort (mkLevelSucc u)) (mkConst `PUnit [mkLevelSucc u])
  let (mvarId, m2, fields2) ← AddFieldsToContext mvarId nm2 [u] #[mkFVar M]
  let (mvarId, fields2) ← caseOnStructures mvarId fields2 ⟨#[]⟩ false
  let ctx : Context := ⟨fields2⟩
  if trace then IO.println s!"cases on fields class 2: {← ctx.target.map (·.name)}"
  let (mvarId, m1, fields1) ← AddFieldsToContext mvarId nm1 [u] #[mkFVar M]
  let (mvarId, fields1) ← caseOnStructures mvarId fields1 ctx
  let st := mkState fields1
  withMVarContext mvarId do
  let st ← uniqueMapping st ctx
  let st := trivialMapping st ctx
  if trace then IO.println s!"map of data: {← st.traceMapping}"
  let (mvarId, fields2) ← caseOnStructures mvarId fields2 ⟨#[]⟩ -- is this dangerous?
  let ctx : Context := ⟨fields2⟩
  withMVarContext mvarId do
  let ctx : Context := ⟨fields2⟩
  let st ← matchingAxioms st ctx
  let st ← st.updateM λ info => if !info.isProp then return none else
    tryToProve mvarId currentAutomation (info.type.instantiateFVars st.getDataMapping)
  if trace then IO.println s!"map: {← st.traceMapping}"
  return (mvarId, m1, m2, st)

end Meta

syntax (name := guardHyp) "guardHyp " (" : " term)? : tactic
@[tactic guardHyp] def evalGuardHyp : Lean.Elab.Tactic.Tactic := fun stx =>
  match stx with
  | `(tactic| guardHyp $[: $ty]?) => do
    return ()
  | _ => throwUnsupportedSyntax

def isSubclassTac (nm1 nm2 : Name) (trace := false) : TacticM Unit := withMainContext do
let mvarId ← getMainGoal
let (mvarId, m1, m2, st) ← Meta.isSubclass mvarId nm1 nm2 trace
if st.done then IO.println s!"{nm1} is a subclass of {nm2}"
else IO.println s!"Cannot construct the following fields of {nm1} from {nm2}:
{st.missing.map (·.name)}."
let l ← getUnsolvedGoals
setGoals (mvarId::m1::m2::l)

def cryptomorphicTac (nm1 nm2 : Name) (trace := false) : TacticM Unit := withMainContext do
let mvarId ← getMainGoal
let (mvarId, m1, m2, st1) ← Meta.isSubclass mvarId nm1 nm2 trace
let (mvarId, m3, m4, st2) ← Meta.isSubclass mvarId nm2 nm1 trace
if st1.done && st2.done then IO.println s!"{nm1} and {nm2} are cryptomorphic"
else do
  IO.println s!"Cannot prove that {nm1} and {nm2} are cryptomorphic"
  if !st1.done then
    IO.println s!"{nm2} → {nm1}: cannot construct {st1.missing.map (·.name)}"
  if !st2.done then
    IO.println s!"{nm1} → {nm2}: Cannot construct {st2.missing.map (·.name)}"
let l ← getUnsolvedGoals
setGoals (mvarId::m1::m2::m3::m4::l)

syntax (name := isSubclass) "isSubclass " ident ident : tactic
syntax (name := isSubclassE) "isSubclass! " ident ident : tactic
@[tactic «isSubclass»] def evalIsSubclass : Tactic := fun stx =>
  match stx with
  | `(tactic| isSubclass $nm1 $nm2) => withoutModifyingState $ isSubclassTac nm1.getId nm2.getId
  | _ => throwUnsupportedSyntax
@[tactic «isSubclassE»] def evalIsSubclassE : Tactic := fun stx =>
  match stx with
  | `(tactic| isSubclass! $nm1 $nm2) => isSubclassTac nm1.getId nm2.getId true
  | _ => throwUnsupportedSyntax

syntax (name := cryptomorphic) "cryptomorphic " ident ident : tactic
@[tactic «cryptomorphic»] def evalCryptomorphic : Tactic := fun stx =>
  match stx with
  | `(tactic| cryptomorphic $nm1 $nm2) => withoutModifyingState $ cryptomorphicTac nm1.getId nm2.getId
  | _ => throwUnsupportedSyntax

end Crypto
open Crypto

/-!
## Demo and tests

We define some notions of commutative Monoids,
(1) right-unital
(2) right-unital, and then has a superfluous axiom `1 * 1 = 1`
(3) both left-unital and right-unital.
(4) denoted additively
(5) with a unit given by an existential quantifier (`∃ one, ...`)
(6) by extending a `Monoid` structure.
-/

class Zero (α : Type u) where
  zero : α

instance [Zero α] : OfNat α (nat_lit 0) where
  ofNat := Zero.zero

class One (α : Type u) where
  one : α

instance [One α] : OfNat α (nat_lit 1) where
  ofNat := One.one

class CommMonoid1 (M : Type _) extends Mul M, One M :=
(mul_assoc : ∀ x y z : M, (x * y) * z = x * (y * z))
(mul_comm : ∀ x y : M, x * y = y * x)
(mul_one : ∀ x : M, x * 1 = x)

class CommMonoid2 (M : Type _) extends Mul M, One M :=
(mul_one : ∀ x : M, x * 1 = x)
(one_mul_one : 1 * 1 = 1)
(mul_assoc : ∀ x y z : M, (x * y) * z = x * (y * z))
(mul_comm : ∀ x y : M, x * y = y * x)

class CommMonoid3 (M : Type _) extends Mul M, One M :=
(one_mul : ∀ x : M, 1 * x = x)
(mul_assoc : ∀ x y z : M, (x * y) * z = x * (y * z))
(mul_comm : ∀ x y : M, x * y = y * x)

class CommMonoid4 (M : Type _) extends Add M, Zero M :=
(add_assoc : ∀ x y z : M, (x + y) + z = x + (y + z))
(add_comm : ∀ x y : M, x + y = y + x)
(add_zero : ∀ x : M, x + 0 = x)

class CommMonoid5 (M : Type _) extends Mul M :=
(mul_axioms : (∀ x y z : M, (x * y) * z = x * (y * z)) ∧ (∀ x y : M, x * y = y * x))
(exists_one : ∃ one : M, ∀ x : M, x * one = x)

class Monoid (M : Type _) extends Mul M, One M :=
(mul_assoc : ∀ x y z : M, (x * y) * z = x * (y * z))
(mul_one : ∀ x : M, x * 1 = x)

class CommMonoid6 (M : Type _) extends Monoid M :=
(mul_comm : ∀ x y : M, x * y = y * x)

open CommMonoid1

-- example (M : Type _) [CommMonoid1 M] (x : M) : 1 * x = 1 := by
--   simp [mul_comm]

example : True := by
  cryptomorphic CommMonoid1 CommMonoid2 -- yes
  cryptomorphic CommMonoid1 CommMonoid3 -- mul_one, one_mul [need better automation]
  cryptomorphic CommMonoid1 CommMonoid4 -- yes
  cryptomorphic CommMonoid1 CommMonoid5 -- one, mul_one [need support for existentials]
  cryptomorphic CommMonoid1 CommMonoid6 -- yes
  trivial

/-! As a sanity check: we cannot prove commutativity on an arbitrary Monoid. -/

class MyMonoid (M : Type _) extends Mul M :=
(mul_assoc : ∀ x y z : M, (x * y) * z = x * (y * z))
(one : M)
(mul_one : ∀ x : M, x * one = x)

example : True := by
  cryptomorphic MyMonoid CommMonoid1 -- missing (expected): mul_comm
  trivial

/-!
If two data fields have the same type, we try to get the one with the same name.
In the future we could look at which choice will make more axioms overlap.
-/

class MyAlmostRing (M : Type _) extends Mul M :=
(add : M → M → M)
(mul_assoc : ∀ x y z : M, (x * y) * z = x * (y * z))
(mul_comm : ∀ x y : M, x * y = y * x)
(one : M)
(mul_one : ∀ x : M, x * one = x)

example : True := by
  cryptomorphic CommMonoid1 MyAlmostRing -- missing (expected): add
  trivial

/-! Test which "fields" are missing when inside nested structures. -/

class CommMonoidBundled1 (M : Type _) extends Mul M :=
(mul_assoc : ∀ x y z : M, (x * y) * z = x * (y * z))
(mul_comm : ∀ x y : M, x * y = y * x)
(one_axioms : ∃ one : M, (∀ x : M, x * one = x) ∧ (∀ x : M, x * x = one))

class CommMonoidBundled2 (M : Type _) :=
(data : (M → M → M) × M)
(mul_assoc : ∀ x y z : M, data.1 (data.1 x y) z = data.1 x (data.1 y z))
(mul_comm : ∀ x y : M, data.1 x y = data.1 y x)
(mul_one : ∀ x : M, data.1 x data.2 = x)

example : True := by
  cryptomorphic CommMonoidBundled1 CommMonoid1 -- missing (expected): one_axioms_prop_right
  -- missing: one, one_mul [need support for existentials]
  cryptomorphic CommMonoidBundled2 CommMonoid1 -- yesss
  trivial