/-
## Rewriting MLIR fragments with PDL

This file implements some framework support for PDL rewrites, which is used to
formalize and prove rewrite operations. We interpret the matching section of
PDL rewrites with the framework's basic pattern-matching tools, which allows us
to unify (enrich) them with constraints specified by the dialect. This is
important because we later need to reason on the semantics of matched programs,
which are only defined if the constraints are met.

The end goal is to output a template theorem that looks like this:

```lean4
-- (unification variables)
∀ (value_1: Value) (type_1: Type) (op_1: Operation),
  -- (non-unification-based constraints)
  is_float_type type_1 →

  semantics_of [mlir_bb|
    -- (symbolic pattern, almost verbatim from the `pdl.pattern` description)
    op_1 (value_1: type_1) 2
  ] =
  semantics_of [mlir_bb|
    -- (rewritten version of above)
    output_op_1 ...
  ]
```

For readability/accessibility reasons, this statement would probably be
synthetized into Lean code rather than being a general theorem instantiated
with that particular PDL rewrite. The second option would require considering
"any matching basic block" and then painstakingly proving that such a match has
the structure underlined by the PDL rewrite, which seems annoying.

Operations of PDL that we support right now:

* `pdl.operand`:
  Used as unification variable. Supports type requirement.
* `pdl.operation`:
  Used as unification term.
* `pdl.pattern`:
  Input operation for the whole rewriting workflow.
* `pdl.result`:
  Utility to reference results from `pdl.operation` pointers
* `pdl.type`:
  Used as unification variable.

Operations of PDL that we should partially or totally express:

* `pdl.attribute`:
  We don't have attributes yet
* `pdl.erase`:
  We don't compute the target yet
* `pdl.operation`:
  We should support attributes. And we don't handle the output either
* `pdl.replace`:
  We don't compute the target yet
* `pdl.rewrite`:
  One of the main components of the specification. We might not care about the
  root selection since we don't implement the pattern matching/search.

Operations or features of PDL that we don't try to express (yet):

* `pdl.operands`, `pdl.results`, `pdl.types`
  Anything related to ranges
* `pdl.apply_native_constraint`, `pdl.apply_native_rewrite`:
  This is possible in principle, but requires the native constraints/rewrites
  to be rewritten in Lean, which is slightly inconvenient.
* Use of native rewrite function in `pdl.rewrite`:
  Same as above
* Using variadic arguments or matching operands with variadic arguments
-/

import MLIR.Semantics.Types
import MLIR.Semantics.Matching
import MLIR.Semantics.Unification
import MLIR.AST
import MLIR.EDSL
import Lean.Exception
import Lean.Elab.Term
open Lean

open MLIR.AST
open MLIR.EDSL

-- TODO: We updated to new matching but did not add regular type checks

/- Utilities -/

private def MLIR.AST.SSAVal.str: MLIR.AST.SSAVal → String
  | SSAVal.SSAVal s => s

-- Get the n-th category of arguments from the specified variadic list by using
-- the operand segment size attribute
private def operandSegment (args: List SSAVal) (oss: TensorElem) (n: Nat) :=
  -- Flatten segment sizes
  let oss := match oss with
  | TensorElem.nested l => l.map (
      match · with | TensorElem.int i => i.toNat | _ => 0)
  | _ => []
  -- Accumulate
  let (oss, _) := oss.foldl
    (fun (segments, acc) size => (segments ++ [(acc, size)], acc + size))
    ([], 0)
  -- Get the segment for that argument
  let (start, size) := oss.getD n (0,0)
  (List.range size).filterMap (fun i => args.get? (i + start))

/-
### Monad for the analysis of PDL programs

While PDL programs are mostly similar to match terms, the translation involves
some bookkeeping and fresh name generation. The Translation structure keeps
track of this information, which is carried around in a state monad called
TranslationM.
-/

structure Translation (δ: Dialect α σ ε) where mk ::
  -- Unification problem
  u: Unification δ
  -- All variables defined so far (for fresh name generation)
  allvars: List String
  -- Typing judgements (to be inserted into operation terms)
  judgements: List (String × MTerm δ)
  -- List of operations that we want to collect after unification
  operations: List (MTerm δ)
  -- Names of results for each operation, and their types
  opresults: List (String × List (String × MTerm δ))
  -- Whether translation completed successfully
  success: Bool

def Translation.empty: Translation δ :=
  { u           := Unification.empty,
    allvars     := [],
    judgements  := [],
    operations  := [],
    opresults   := [],
    success     := false }

instance {δ: Dialect α σ ε} : Inhabited (Translation δ) := ⟨Translation.empty⟩

def Translation.str: Translation δ → String := fun tr =>
  "Unification problem:\n" ++
    (toString tr.u) ++
  "\nAll variables:\n " ++
    (String.join $ tr.allvars.map (s!" %{·}")) ++
  "\nTyping judgements:\n" ++
    (String.join $ tr.judgements.map (fun (v,t) => s!"  %{v}: {t}\n")) ++
  "Operation results:\n" ++
    (String.join $ tr.opresults.map (fun (n,l) => s!"  %{n}: {l}\n"))

abbrev TranslationM (δ: Dialect α σ ε) := StateT (Translation δ) IO

def TranslationM.toIO {α} (tr: Translation δ) (x: TranslationM δ α): IO α :=
  Prod.fst <$> StateT.run x tr

def TranslationM.error {α} (s: String): TranslationM δ α :=
  throw <| IO.userError (s ++ " o(x_x)o")

/-
#### Name recording and name generation

The following monad functions are concerned with tracking variables,
guaranteeing uniqueness, and generating fresh names.
-/

-- Record that [name] is now used in the problem, and check uniqueness
def TranslationM.addName (name: String): TranslationM δ Unit := do
  let tr ← get
  if name ∈ tr.allvars then
    error s!"addName: {name} is already used!"
  else
    set { tr with allvars := tr.allvars ++ [name] }

-- Check that [name] is known
def TranslationM.checkNameDefined (name: String): TranslationM δ Unit := do
  let tr ← get
  if ! name ∈ tr.allvars then
    error s!"checkNameDefined: {name} is undefined!"

-- Make up to [n] attempts at finding a fresh name by suffixing [s]
private def TranslationM.freshNameAux (s: String) (n p: Nat):
    TranslationM δ String := do
  let tr ← get
  match n with
  | 0 =>
      return s
  | m+1 =>
      let s' := s!"{s}{p}"
      if tr.allvars.all (· != s') then
        set { tr with allvars := tr.allvars ++ [s'] }
        return s'
      else
        freshNameAux s m (p+1)

-- Generate a new fresh name based on [name]; if not available, resort to
-- adding numbered suffixes
def TranslationM.makeFreshName (name: String): TranslationM δ String := do
  let tr ← get
  if tr.allvars.all (· != name) then
    addName name
    return name
  else
    let f ← freshNameAux name (tr.allvars.length+1) 0
    addName f
    return f

-- Generate [n] fresh names based on [name]
def TranslationM.makeFreshNames (name: String) (n: Nat):
    TranslationM δ (List String) :=
  (List.range n).mapM (fun idx => makeFreshName (name ++ toString idx))

-- Generate a copy of the operation with the specified prefix and fresh names.
-- Returns a pair with the new term and the list of all variables involved.
def TranslationM.makeFreshOp (prefix_: String) (op: MTerm δ) (priority: Nat):
    TranslationM δ (MTerm δ) := do
  let renameVars (done: List String) (vars: List (String × MSort)):
      TranslationM δ (List (String × MTerm δ) × List String) :=
    vars.foldlM
      (fun (repl, done) (var, sort) => do
        if var ∈ done then
          return (repl, done)
        else
          let var' ← makeFreshName (prefix_ ++ var)
          return (repl ++ [(var, MTerm.Var priority var' sort)], done ++ [var]))
      ([], done)

  let (repl, done) ← renameVars [] op.varsWithSorts
  return op.substVars repl

/-
#### Access to translation data

The following utilities query data recorded in the translation state.
-/

def TranslationM.addEquation (equation: UEq δ): TranslationM δ Unit := do
  let tr ← get
  set { tr with u := { equations := tr.u.equations ++ [equation] } }

def TranslationM.addOperation (op: MTerm δ): TranslationM δ Unit := do
  let tr ← get
  set { tr with operations := tr.operations ++ [op] }

def TranslationM.addJudgement (s: String) (type: MTerm δ): TranslationM δ Unit := do
  let tr ← get
  set { tr with judgements := tr.judgements ++ [(s, type)] }

def TranslationM.addOpResults (name: String) (results: List (String × MTerm δ)):
    TranslationM δ Unit := do
  let tr ← get
  set { tr with opresults := tr.opresults ++ [(name, results)] }

def TranslationM.findJudgement? (s: String): TranslationM δ (Option (MTerm δ)) := do
  let cmpName := fun (name, type) => if name = s then some type else none
  return (← get).judgements.findSome? cmpName

def TranslationM.findJudgement (s: String): TranslationM δ (MTerm δ) := do
  match ← findJudgement? s with
  | some type   => return type
  | none        => error s!"findJudgement: no type information for {s}!"

def TranslationM.findJudgements (l: List String):
    TranslationM δ (List (String × MTerm δ)) :=
  l.mapM (fun var => do return (var, ← findJudgement var))

def TranslationM.findOpResult? (op: String) (idx: Nat):
    TranslationM δ (Option (String × MTerm δ)) := do
  let cmpName := fun (name, rets) => if name = op then rets.get? idx else none
  return (← get).opresults.findSome? cmpName

def TranslationM.findOpResult (op: String) (idx: Nat):
    TranslationM δ (String × MTerm δ) := do
  match ← findOpResult? op idx with
  | some info   => return info
  | none        => error s!"findOpResultName: no return #{idx} for {op}!"

/-
### Translation of PDL statements to match problems

Interpretation PDL programs as match problems is fairly straightforward; most
PDL operations simply add new names or equations to the. Most of the work is
spent on parsing the input, adding new names and keeping track of information.

PDL has a lot of restrictions on what you can write; you can't use the same
variable at two different places (implicit unification), you can't have a
declaration operation (pdl.value, pdl.type, etc) without binding it to a
variable, etc. This allows us to make assumptions on the shape of the input.
-/

section
open TranslationM

-- TODO: Provide dialect data properly once implemented
private def PDLToMatch.readStatement (operationMatchTerms: List (MTerm builtin))
    (op: Op builtin): TranslationM builtin Unit := do
  let tr ← get

  match op with
  -- %name = pdl.type
  | Op.mk "pdl.type" [⟨SSAVal.SSAVal name, .undefined "pdl.type"⟩]
         [] [] (AttrDict.mk []) => do
      IO.println s!"Found new type variable: {name}"
      addName name

  -- %name = pdl.type: TYPE
  | Op.mk "pdl.type" [⟨SSAVal.SSAVal name, .undefined "pdl.type"⟩]
         [] [] (AttrDict.mk [AttrEntry.mk "type" (AttrValue.type τ)]) => do
      IO.println s!"Found new type variable: {name} (= {τ})"
      addName name
      addEquation (.Var 1 name .MMLIRType, .ConstMLIRType τ)

  -- %name = pdl.operand
  | Op.mk "pdl.operand" [⟨SSAVal.SSAVal name, .undefined "pdl.value"⟩]
         [] [] (AttrDict.mk []) => do
      IO.println s!"Found new variable: {name}"
      addName name
      let typeName ← makeFreshName (name ++ "_T")
      IO.println s!"→ Generated type name: {typeName}"
      addJudgement name (.Var 1 typeName .MMLIRType)

  -- %name = pdl.operand: %typeName
  | Op.mk "pdl.operand" [⟨SSAVal.SSAVal name, .undefined "pdl.value"⟩]
        [⟨SSAVal.SSAVal typeName, .undefined "pdl.type"⟩] [] (AttrDict.mk []) => do
      IO.println s!"Found new variable: {name} of type {typeName}"
      addName name
      checkNameDefined typeName
      addJudgement name (.Var 0 typeName .MMLIRType)

  -- %name = pdl.operation "OPNAME"(ARGS) -> TYPE
  | Op.mk "pdl.operation" [⟨SSAVal.SSAVal name, _⟩] args [] attrs => do
      let (attributeNames, opname, operand_segment_sizes) :=
        (attrs.find "attributeNames",
         attrs.find "name",
         attrs.find "operand_segment_sizes")
      let argsVal := args.map Prod.fst

      match attributeNames, opname, operand_segment_sizes with
      | some (.list _),
        some (.str opname),
        some (builtin.dense_vector_attr oss _ _ _) =>
          let values := (operandSegment argsVal oss 0).map (·.str)
          let types  := (operandSegment argsVal oss 2).map (·.str)
          IO.println s!"Found new operation: {name} matching {opname}"
          addName name

          IO.println s!"→ Arguments: {values}, return types: {types}"
          values.forM checkNameDefined
          types.forM checkNameDefined

          let valuesTypes ← findJudgements values
          let retNames ← makeFreshNames (name ++ "_res") types.length
          IO.println s!"→ Arg types: {valuesTypes}, return names: {retNames}"

          let insPattern := operationMatchTerms.find? fun
            | .App .OP [.ConstString n, _, _] => n = opname
            | _ => false
          if insPattern.isNone then
            error s!"pdl.operation: no pattern known for {opname}!"
          let insPattern := insPattern.get!

          IO.println s!"→ Using match term: {insPattern}"
          let ins ← makeFreshOp (name ++ "_") insPattern 2
          addOperation ins
          IO.println s!"→ Instantiated match term: {ins}"

          let operands_args: List (MTerm _) := valuesTypes.map (fun (v,t) =>
            .App .OPERAND [.Var 0 v .MSSAVal, t])
          let operands_rets: List (MTerm _) :=
            List.zip retNames types |>.map (fun (v, t) =>
              .App .OPERAND [.Var 1 v .MSSAVal, .Var 0 t .MMLIRType])
          let op_mterm :=
            .App .OP [
              .ConstString opname,
              .App (.LIST .MOperand) operands_args,
              .App (.LIST .MOperand) operands_rets
            ]
          addEquation (op_mterm, ins)

          let opResults := List.zip retNames (types.map (.Var 0 · .MMLIRType))
          addOpResults name opResults

      | _, _, _ =>
          error s!"pdl.operation: unexpected attributes: {attrs}"

    -- %name = pdl.result INDEX of %op
  | Op.mk "pdl.result" [⟨SSAVal.SSAVal name, .undefined "pdl.value"⟩]
        [⟨SSAVal.SSAVal opname, .undefined "pdl.operation"⟩] []  attrs => do
      match attrs.find "index" with
      | some (AttrValue.int index (MLIRType.int _ _)) =>
          IO.println
            s!"Found new variable: {name} aliasing result {index} of {opname}"
          checkNameDefined opname
          addName name

          let (resName, resType) ← findOpResult opname index.toNat
          addJudgement name resType
          addEquation (.Var 0 name .MSSAVal, .Var 1 resName .MSSAVal)

      | _ =>
          error s!"pdl.result: unexpected attributes on {opname}: {attrs}"

  | Op.mk "pdl.rewrite" [] args regions attrs =>
      return ()

  | _ => do
      error s!"{op.name}: unrecognized PDL operation"

def PDLToMatch.convert (PDLProgram: Op builtin)
    (operationMatchTerms: List (MTerm builtin)):
    TranslationM builtin Unit :=
  match PDLProgram with
  | Op.mk "pdl.pattern" _ [] [region] attrs =>
      match region with
      | Region.mk name [] stmts => do
          stmts.forM (readStatement operationMatchTerms)
          set { ← get with success := true }
      | Region.mk _ _ _ => do
          error (s!"PDLToMatch.convert: expected only one BB with no " ++
            "arguments in the pattern region")
  | Op.mk "pdl.pattern" _ _ _ _ => do
      error (s!"PDLToMatch.convert: expected operation to have exactly one " ++
        "argument (a region):\n{pattern}")
  | _ => do
      error s!"PDLToMatch.convert: not a PDL program: {PDLProgram}"

def PDLToMatch.unify: TranslationM δ Unit := do
  let tr ← get
  if ! tr.success then
    error "unify: translation did not complete successfully"
  if let some u_unified ← tr.u.solve then
    set { tr with u := u_unified }
  else
    error "unify: unification failed"

def PDLToMatch.getOperationMatchTerms: TranslationM δ (List (MTerm δ)) := do
  let tr ← get
  return tr.operations.map tr.u.applyOnTerm

end

/-
### Example PDL program
-/

private def ex_pdl: Op builtin := [mlir_op|
  "pdl.pattern"() ({
    -- %T0 = pdl.type
    %T0 = "pdl.type"() : () -> !"pdl.type"
    -- %T1 = pdl.type: i32
    %T1 = "pdl.type"() {type = i32} : () -> !"pdl.type"
    -- %v2 = pdl.operand
    %v2 = "pdl.operand"() : () -> !"pdl.value"
    -- %O3 = pdl.operation "foo.op1"(%v2) -> %T0
    %O3 = "pdl.operation"(%v2, %T0) {attributeNames = [], name = "foo.op1", operand_segment_sizes = dense<[1, 0, 1]> : vector<3×i32>} : (!"pdl.value", !"pdl.type") -> !"pdl.operation"
    -- %v4 = pdl.result 0 of %O3
    %v4 = "pdl.result"(%O3) {index = 0 : i32} : (!"pdl.operation") -> !"pdl.value"
    -- %v5 = pdl.operand: %T0
    %v5 = "pdl.operand"(%T0) : (!"pdl.type") -> !"pdl.value"
    -- %O6 = pdl.operation "foo.op2"(%v4, %v5) -> %T1
    %O6 = "pdl.operation"(%v4, %v5, %T1) {attributeNames = [], name = "foo.op2", operand_segment_sizes = dense<[2, 0, 1]> : vector<3×i32>} : (!"pdl.value", !"pdl.value", !"pdl.type") -> !"pdl.operation"

    -- TODO
    "pdl.rewrite"(%O6) ({
      "pdl.replace"(%O6, %v2) {operand_segment_sizes = dense<[1, 0, 1]> : vector<3×i32>} : (!"pdl.operation", !"pdl.value") -> ()
    }) {operand_segment_sizes = dense<[1, 0]> : vector<2×i32>} : (!"pdl.operation") -> ()
  }) {benefit = 1 : i16} : () -> ()
]

-- %res:!T = "foo.op1"(%x:!T)
private def foo_op1_pattern: MTerm builtin :=
  .App .OP [
    .ConstString "foo.op1",
    .App (.LIST .MOperand) [
      .App .OPERAND [.Var 1 "x" .MSSAVal, .Var 1 "T" .MMLIRType]],
    .App (.LIST .MOperand) [
      .App .OPERAND [.Var 1 "ret" .MSSAVal, .Var 1 "T" .MMLIRType]]
  ]
#eval foo_op1_pattern

-- %res:!T = "foo.op2"(%x:!T, %y:i32)
private def foo_op2_pattern: MTerm builtin :=
  .App .OP [
    .ConstString "foo.op2",
    .App (.LIST .MOperand) [
      .App .OPERAND [.Var 1 "x" .MSSAVal, .Var 1 "T" .MMLIRType],
      .App .OPERAND [.Var 1 "y" .MSSAVal, .ConstMLIRType .i32]],
    .App (.LIST .MOperand) [
      .App .OPERAND [.Var 1 "ret" .MSSAVal, .Var 1 "T" .MMLIRType]]
  ]
#eval foo_op2_pattern

private def foo_dialect := [foo_op1_pattern, foo_op2_pattern]

#eval show IO Unit from do
  TranslationM.toIO Translation.empty do
    PDLToMatch.convert ex_pdl foo_dialect
    IO.println $ "\n## Translation result ##\n\n" ++ (← get).str
    PDLToMatch.unify
    let matchTerms ← PDLToMatch.getOperationMatchTerms
    IO.println "\n## Final operation match terms ##\n"
    IO.println $ "\n".intercalate (matchTerms.map toString)


def HList (αs : List (Type u)) : Type u := αs.foldr Prod.{u, u} PUnit

@[match_pattern] def HList.nil : HList [] := ⟨⟩
@[match_pattern] def HList.cons (a : α) (as : HList αs): HList (α :: αs) := (a, as)

def HList.set : {αs : _} → HList αs → (i : Fin αs.length) → αs.get i → HList αs
  | _ :: _, cons a as, ⟨0,          h⟩, b => cons b as
  | _ :: _, cons a as, ⟨Nat.succ n, h⟩, b => cons a (set as ⟨n, Nat.le_of_succ_le_succ h⟩ b)
  | [],     nil,       _,               _ => nil

instance : EmptyCollection (HList ∅) where
  emptyCollection := HList.nil

notation:30 Γ " ⊢ " α => HList Γ → α

-- simplify well-founded recursion proofs by ignoring context sizes
local instance : SizeOf (List α) := ⟨fun _ => 0⟩ in

-- m: base monad
-- ω: `return` type, `m ω` is the type of the entire `do` block
-- Γ: `do`-local immutable context
-- Δ: `do`-local mutable context
-- b: `break` allowed
-- c: `continue` allowed
-- α: local result type, `m α` is the type of the statement
inductive Stmt (m : Type u → Type _) (ω : Type u) : (Γ Δ : List (Type u)) → (b c : Bool) → (α : Type u) → Type _ where
  | expr (e : Γ ⊢ Δ ⊢ m α) : Stmt m ω Γ Δ b c α
  | bind (s₁ : Stmt m ω Γ Δ b c α) (s₂ : Stmt m ω (α :: Γ) Δ b c β) : Stmt m ω Γ Δ b c β
  | letmut (e : Γ ⊢ Δ ⊢ α) (s : Stmt m ω Γ (α :: Δ) b c β) : Stmt m ω Γ Δ b c β
  | ass (x : Fin Δ.length) (e : Γ ⊢ Δ ⊢ Δ.get x) : Stmt m ω Γ Δ b c PUnit
  | ite (e : Γ ⊢ Δ ⊢ Bool) (s₁ s₂ : Stmt m ω Γ Δ b c α) : Stmt m ω Γ Δ b c α
  | ret (e : Γ ⊢ Δ ⊢ ω) : Stmt m ω Γ Δ b c α
  --| sfor [ForM m γ α] (e : Σ Γ → γ) (body : α → Stmt m ω Γ Δ true PUnit) : Stmt m ω Γ Δ b c PUnit
  | sfor (e : Γ ⊢ Δ ⊢ List α) (body : Stmt m ω (α :: Γ) Δ true true PUnit) : Stmt m ω Γ Δ b c PUnit
  | sbreak : Stmt m ω Γ Δ true c α
  | scont : Stmt m ω Γ Δ b true α

-- normal and abnormal result values
inductive Res (ω α : Type _) : (b c : Bool) → Type _ where
  | val (a : α) : Res ω α b c
  | ret (o : ω) : Res ω α b c
  | rbreak : Res ω α true c
  | rcont : Res ω α b true

instance : Coe α (Res ω α b c) := ⟨Res.val⟩
instance : Coe (Id α) (Res ω α b c) := ⟨Res.val⟩

def Ctx.extendBot (x : α) : {Γ : _} → HList Γ → HList (Γ ++ [α])
  | [],     _               => HList.cons x HList.nil
  | _ :: _, HList.cons a as => HList.cons a (extendBot x as)

def Ctx.extend (x : α) : HList Γ → HList (α :: Γ) :=
  fun σ => HList.cons x σ

def Ctx.drop : HList (α :: Γ) → HList Γ
  | HList.cons a as => as

-- custom wf tactic
theorem Nat.le_add_right_of_le (n m : Nat) : n ≤ m → n ≤ m + k :=
  fun h => add_le_add h (Nat.zero_le _)

macro_rules
| `(tactic| decreasing_tactic) =>
 `(tactic|
   (simp_wf
    repeat (first | apply PSigma.Lex.right | apply PSigma.Lex.left)
    simp [Nat.add_comm (n := 1), Nat.succ_add, Nat.mul_succ]
    try apply Nat.lt_succ_of_le
    repeat apply Nat.le_step
    first
    | repeat first | apply Nat.le_add_left | apply Nat.le_add_right_of_le
    | assumption
    all_goals apply Nat.le_refl
))

@[simp]
def Stmt.mapCtx (f : HList Γ' → HList Γ) : Stmt m ω Γ Δ b c β → Stmt m ω Γ' Δ b c β
  | expr e => expr (e ∘ f)
  | bind s₁ s₂ => bind (s₁.mapCtx f) (s₂.mapCtx (fun | HList.cons a as => HList.cons a (f as)))
  | letmut e s => letmut (e ∘ f) (s.mapCtx f)
  | ass x e => ass x (e ∘ f)
  | ite e s₁ s₂ => ite (e ∘ f) (s₁.mapCtx f) (s₂.mapCtx f)
  | ret e => ret (e ∘ f)
  | sfor e body => sfor (e ∘ f) (body.mapCtx (fun | HList.cons a as => HList.cons a (f as)))
  | sbreak => sbreak
  | scont => scont
termination_by mapCtx _ s => sizeOf s


/-
Copyright (c) 2020 Scott Morrison. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Scott Morrison, Bhavik Mehta
-/
import category_theory.limits.shapes.pullbacks
import category_theory.limits.shapes.strong_epi
import category_theory.limits.shapes.equalizers

/-!
# Definitions and basic properties of regular monomorphisms and epimorphisms.

A regular monomorphism is a morphism that is the equalizer of some parallel pair.

We give the constructions
* `split_mono → regular_mono` and
* `regular_mono → mono`
as well as the dual constructions for regular epimorphisms. Additionally, we give the construction
* `regular_epi ⟶ strong_epi`.

We also define classes `regular_mono_category` and `regular_epi_category` for categories in which
every monomorphism or epimorphism is regular, and deduce that these categories are
`strong_mono_category`s resp. `strong_epi_category`s.

-/

noncomputable theory

namespace category_theory
open category_theory.limits

universes v₁ u₁ u₂

variables {C : Type u₁} [category.{v₁} C]

variables {X Y : C}

/-- A regular monomorphism is a morphism which is the equalizer of some parallel pair. -/
class regular_mono (f : X ⟶ Y) :=
(Z : C)
(left right : Y ⟶ Z)
(w : f ≫ left = f ≫ right)
(is_limit : is_limit (fork.of_ι f w))

attribute [reassoc] regular_mono.w

/-- Every regular monomorphism is a monomorphism. -/
@[priority 100]
instance regular_mono.mono (f : X ⟶ Y) [regular_mono f] : mono f :=
mono_of_is_limit_parallel_pair regular_mono.is_limit

instance equalizer_regular (g h : X ⟶ Y) [has_limit (parallel_pair g h)] :
  regular_mono (equalizer.ι g h) :=
{ Z := Y,
  left := g,
  right := h,
  w := equalizer.condition g h,
  is_limit := fork.is_limit.mk _ (λ s, limit.lift _ s) (by simp) (λ s m w, by { ext1, simp [←w] }) }

/-- Every split monomorphism is a regular monomorphism. -/
@[priority 100]
instance regular_mono.of_split_mono (f : X ⟶ Y) [split_mono f] : regular_mono f :=
{ Z     := Y,
  left  := 𝟙 Y,
  right := retraction f ≫ f,
  w     := by tidy,
  is_limit := split_mono_equalizes f }

/-- If `f` is a regular mono, then any map `k : W ⟶ Y` equalizing `regular_mono.left` and
    `regular_mono.right` induces a morphism `l : W ⟶ X` such that `l ≫ f = k`. -/
def regular_mono.lift' {W : C} (f : X ⟶ Y) [regular_mono f] (k : W ⟶ Y)
  (h : k ≫ (regular_mono.left : Y ⟶ @regular_mono.Z _ _ _ _ f _) = k ≫ regular_mono.right) :
  {l : W ⟶ X // l ≫ f = k} :=
fork.is_limit.lift' regular_mono.is_limit _ h

/--
The second leg of a pullback cone is a regular monomorphism if the right component is too.

See also `pullback.snd_of_mono` for the basic monomorphism version, and
`regular_of_is_pullback_fst_of_regular` for the flipped version.
-/
def regular_of_is_pullback_snd_of_regular {P Q R S : C} {f : P ⟶ Q} {g : P ⟶ R} {h : Q ⟶ S}
  {k : R ⟶ S} [hr : regular_mono h] (comm : f ≫ h = g ≫ k)
  (t : is_limit (pullback_cone.mk _ _ comm)) :
regular_mono g :=
{ Z := hr.Z,
  left := k ≫ hr.left,
  right := k ≫ hr.right,
  w := by rw [← reassoc_of comm, ← reassoc_of comm, hr.w],
  is_limit :=
  begin
    apply fork.is_limit.mk' _ _,
    intro s,
    have l₁ : (fork.ι s ≫ k) ≫ regular_mono.left = (fork.ι s ≫ k) ≫ regular_mono.right,
      rw [category.assoc, s.condition, category.assoc],
    obtain ⟨l, hl⟩ := fork.is_limit.lift' hr.is_limit _ l₁,
    obtain ⟨p, hp₁, hp₂⟩ := pullback_cone.is_limit.lift' t _ _ hl,
    refine ⟨p, hp₂, _⟩,
    intros m w,
    have z : m ≫ g = p ≫ g := w.trans hp₂.symm,
    apply t.hom_ext,
    apply (pullback_cone.mk f g comm).equalizer_ext,
    { erw [← cancel_mono h, category.assoc, category.assoc, comm, reassoc_of z] },
    { exact z },
  end }

/--
The first leg of a pullback cone is a regular monomorphism if the left component is too.

See also `pullback.fst_of_mono` for the basic monomorphism version, and
`regular_of_is_pullback_snd_of_regular` for the flipped version.
-/
def regular_of_is_pullback_fst_of_regular {P Q R S : C} {f : P ⟶ Q} {g : P ⟶ R} {h : Q ⟶ S}
  {k : R ⟶ S} [hr : regular_mono k] (comm : f ≫ h = g ≫ k)
  (t : is_limit (pullback_cone.mk _ _ comm)) :
regular_mono f :=
regular_of_is_pullback_snd_of_regular comm.symm (pullback_cone.flip_is_limit t)

@[priority 100]
instance strong_mono_of_regular_mono (f : X ⟶ Y) [regular_mono f] : strong_mono f :=
{ mono := by apply_instance,
  has_lift :=
  begin
    introsI,
    have : v ≫ (regular_mono.left : Y ⟶ regular_mono.Z f) = v ≫ regular_mono.right,
    { apply (cancel_epi z).1,
      simp only [regular_mono.w, ← reassoc_of h] },
    obtain ⟨t, ht⟩ := regular_mono.lift' _ _ this,
    refine arrow.has_lift.mk ⟨t, (cancel_mono f).1 _, ht⟩,
    simp only [arrow.mk_hom, arrow.hom_mk'_left, category.assoc, ht, h]
  end }

/-- A regular monomorphism is an isomorphism if it is an epimorphism. -/
lemma is_iso_of_regular_mono_of_epi (f : X ⟶ Y) [regular_mono f] [e : epi f] : is_iso f :=
is_iso_of_epi_of_strong_mono _

section
variables (C)

/-- A regular mono category is a category in which every monomorphism is regular. -/
class regular_mono_category :=
(regular_mono_of_mono : ∀ {X Y : C} (f : X ⟶ Y) [mono f], regular_mono f)

end

/-- In a category in which every monomorphism is regular, we can express every monomorphism as
    an equalizer. This is not an instance because it would create an instance loop. -/
def regular_mono_of_mono [regular_mono_category C] (f : X ⟶ Y) [mono f] : regular_mono f :=
regular_mono_category.regular_mono_of_mono _

@[priority 100]
instance regular_mono_category_of_split_mono_category [split_mono_category C] :
  regular_mono_category C :=
{ regular_mono_of_mono := λ _ _ f _,
  by { haveI := by exactI split_mono_of_mono f, apply_instance } }

@[priority 100]
instance strong_mono_category_of_regular_mono_category [regular_mono_category C] :
  strong_mono_category C :=
{ strong_mono_of_mono := λ _ _ f _,
    by { haveI := by exactI regular_mono_of_mono f, apply_instance } }

/-- A regular epimorphism is a morphism which is the coequalizer of some parallel pair. -/
class regular_epi (f : X ⟶ Y) :=
(W : C)
(left right : W ⟶ X)
(w : left ≫ f = right ≫ f)
(is_colimit : is_colimit (cofork.of_π f w))

attribute [reassoc] regular_epi.w

/-- Every regular epimorphism is an epimorphism. -/
@[priority 100]
instance regular_epi.epi (f : X ⟶ Y) [regular_epi f] : epi f :=
epi_of_is_colimit_parallel_pair regular_epi.is_colimit

instance coequalizer_regular (g h : X ⟶ Y) [has_colimit (parallel_pair g h)] :
  regular_epi (coequalizer.π g h) :=
{ W := X,
  left := g,
  right := h,
  w := coequalizer.condition g h,
  is_colimit := cofork.is_colimit.mk _ (λ s, colimit.desc _ s) (by simp)
    (λ s m w, by { ext1, simp [←w] }) }

/-- Every split epimorphism is a regular epimorphism. -/
@[priority 100]
instance regular_epi.of_split_epi (f : X ⟶ Y) [split_epi f] : regular_epi f :=
{ W     := X,
  left  := 𝟙 X,
  right := f ≫ section_ f,
  w     := by tidy,
  is_colimit := split_epi_coequalizes f }

/-- If `f` is a regular epi, then every morphism `k : X ⟶ W` coequalizing `regular_epi.left` and
    `regular_epi.right` induces `l : Y ⟶ W` such that `f ≫ l = k`. -/
def regular_epi.desc' {W : C} (f : X ⟶ Y) [regular_epi f] (k : X ⟶ W)
  (h : (regular_epi.left : regular_epi.W f ⟶ X) ≫ k = regular_epi.right ≫ k) :
  {l : Y ⟶ W // f ≫ l = k} :=
cofork.is_colimit.desc' (regular_epi.is_colimit) _ h

/--
The second leg of a pushout cocone is a regular epimorphism if the right component is too.

See also `pushout.snd_of_epi` for the basic epimorphism version, and
`regular_of_is_pushout_fst_of_regular` for the flipped version.
-/
def regular_of_is_pushout_snd_of_regular
  {P Q R S : C} {f : P ⟶ Q} {g : P ⟶ R} {h : Q ⟶ S} {k : R ⟶ S}
  [gr : regular_epi g] (comm : f ≫ h = g ≫ k) (t : is_colimit (pushout_cocone.mk _ _ comm)) :
regular_epi h :=
{ W := gr.W,
  left := gr.left ≫ f,
  right := gr.right ≫ f,
  w := by rw [category.assoc, category.assoc, comm, reassoc_of gr.w],
  is_colimit :=
  begin
    apply cofork.is_colimit.mk' _ _,
    intro s,
    have l₁ : gr.left ≫ f ≫ s.π = gr.right ≫ f ≫ s.π,
      rw [← category.assoc, ← category.assoc, s.condition],
    obtain ⟨l, hl⟩ := cofork.is_colimit.desc' gr.is_colimit (f ≫ cofork.π s) l₁,
    obtain ⟨p, hp₁, hp₂⟩ := pushout_cocone.is_colimit.desc' t _ _ hl.symm,
    refine ⟨p, hp₁, _⟩,
    intros m w,
    have z := w.trans hp₁.symm,
    apply t.hom_ext,
    apply (pushout_cocone.mk _ _ comm).coequalizer_ext,
    { exact z },
    { erw [← cancel_epi g, ← reassoc_of comm, ← reassoc_of comm, z], refl },
  end }

/--
The first leg of a pushout cocone is a regular epimorphism if the left component is too.

See also `pushout.fst_of_epi` for the basic epimorphism version, and
`regular_of_is_pushout_snd_of_regular` for the flipped version.
-/
def regular_of_is_pushout_fst_of_regular
  {P Q R S : C} {f : P ⟶ Q} {g : P ⟶ R} {h : Q ⟶ S} {k : R ⟶ S}
  [fr : regular_epi f] (comm : f ≫ h = g ≫ k) (t : is_colimit (pushout_cocone.mk _ _ comm)) :
regular_epi k :=
regular_of_is_pushout_snd_of_regular comm.symm (pushout_cocone.flip_is_colimit t)

@[priority 100]
instance strong_epi_of_regular_epi (f : X ⟶ Y) [regular_epi f] : strong_epi f :=
{ epi := by apply_instance,
  has_lift :=
  begin
    introsI,
    have : (regular_epi.left : regular_epi.W f ⟶ X) ≫ u = regular_epi.right ≫ u,
    { apply (cancel_mono z).1,
      simp only [category.assoc, h, regular_epi.w_assoc] },
    obtain ⟨t, ht⟩ := regular_epi.desc' f u this,
    exact arrow.has_lift.mk ⟨t, ht, (cancel_epi f).1
      (by simp only [←category.assoc, ht, ←h, arrow.mk_hom, arrow.hom_mk'_right])⟩,
  end }

/-- A regular epimorphism is an isomorphism if it is a monomorphism. -/
lemma is_iso_of_regular_epi_of_mono (f : X ⟶ Y) [regular_epi f] [m : mono f] : is_iso f :=
is_iso_of_mono_of_strong_epi _

section
variables (C)

/-- A regular epi category is a category in which every epimorphism is regular. -/
class regular_epi_category :=
(regular_epi_of_epi : ∀ {X Y : C} (f : X ⟶ Y) [epi f], regular_epi f)

end

/-- In a category in which every epimorphism is regular, we can express every epimorphism as
    a coequalizer. This is not an instance because it would create an instance loop. -/
def regular_epi_of_epi [regular_epi_category C] (f : X ⟶ Y) [epi f] : regular_epi f :=
regular_epi_category.regular_epi_of_epi _

@[priority 100]
instance regular_epi_category_of_split_epi_category [split_epi_category C] :
  regular_epi_category C :=
{ regular_epi_of_epi := λ _ _ f _, by { haveI := by exactI split_epi_of_epi f, apply_instance } }

@[priority 100]
instance strong_epi_category_of_regular_epi_category [regular_epi_category C] :
  strong_epi_category C :=
{ strong_epi_of_epi := λ _ _ f _, by { haveI := by exactI regular_epi_of_epi f, apply_instance } }

end category_theory


import Lean
open List Lean

example (l : List α) (h : length l = length l) : length (a::l) = length (a::l) :=
  congrArg (·+1) h

elab:max "(" tm:term ":)" : term => Elab.Term.elabTerm tm none

example (l : List α) (h : length l = length l) : length (a::l) = length (a::l) :=
  (congrArg (·+1) h :)

example (l : List α) (h : length l = length l) : length (a::l) = length (a::l) :=
  have := congrArg (·+1) h; this

example (l : List α) (h : length l = length l) : length (a::l) = length (a::l) := by
  have := congrArg (·+1) h; exact this


/-
Copyright (c) 2016 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura

! This file was ported from Lean 3 source module init.meta.smt.smt_tactic
! leanprover-community/mathlib commit 4a03bdeb31b3688c31d02d7ff8e0ff2e5d6174db
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
prelude
import Leanbin.Init.Control.Default
import Leanbin.Init.Meta.SimpTactic
import Leanbin.Init.Meta.Smt.CongruenceClosure
import Leanbin.Init.Meta.Smt.Ematch

universe u

run_cmd
  mk_simp_attr `pre_smt

run_cmd
  mk_hinst_lemma_attr_set `ematch [] [`ematch_lhs]

/-- Configuration for the smt tactic preprocessor. The preprocessor
  is applied whenever a new hypothesis is introduced.

  - simp_attr: is the attribute name for the simplification lemmas
    that are used during the preprocessing step.

  - max_steps: it is the maximum number of steps performed by the simplifier.

  - zeta: if tt, then zeta reduction (i.e., unfolding let-expressions)
    is used during preprocessing.
-/
structure SmtPreConfig where
  simpAttr : Name := `pre_smt
  maxSteps : Nat := 1000000
  zeta : Bool := false
#align smt_pre_config SmtPreConfig

/-- Configuration for the smt_state object.

- em_attr: is the attribute name for the hinst_lemmas
  that are used for ematching -/
structure SmtConfig where
  ccCfg : CcConfig := { }
  emCfg : EmatchConfig := { }
  preCfg : SmtPreConfig := { }
  emAttr : Name := `ematch
#align smt_config SmtConfig

unsafe def smt_config.set_classical (c : SmtConfig) (b : Bool) : SmtConfig :=
  { c with ccCfg := { c.ccCfg with em := b } }
#align smt_config.set_classical smt_config.set_classical

unsafe axiom smt_goal : Type
#align smt_goal smt_goal

unsafe def smt_state :=
  List smt_goal
#align smt_state smt_state

unsafe axiom smt_state.mk : SmtConfig → tactic smt_state
#align smt_state.mk smt_state.mk

unsafe axiom smt_state.to_format : smt_state → tactic_state → format
#align smt_state.to_format smt_state.to_format

/-- Return tt iff classical excluded middle was enabled at  smt_state.mk -/
unsafe axiom smt_state.classical : smt_state → Bool
#align smt_state.classical smt_state.classical

unsafe def smt_tactic :=
  StateT smt_state tactic
#align smt_tactic smt_tactic

unsafe instance : Append smt_state :=
  List.hasAppend

section

attribute [local reducible] smt_tactic

unsafe instance : Monad smt_tactic := by infer_instance

unsafe instance : Alternative smt_tactic := by infer_instance

unsafe instance : MonadState smt_state smt_tactic := by infer_instance

end

/- We don't use the default state_t lift operation because only
   tactics that do not change hypotheses can be automatically lifted to smt_tactic. -/
unsafe axiom tactic_to_smt_tactic (α : Type) : tactic α → smt_tactic α
#align tactic_to_smt_tactic tactic_to_smt_tactic

unsafe instance : HasMonadLift tactic smt_tactic :=
  ⟨tactic_to_smt_tactic⟩

unsafe instance (α : Type) : Coe (tactic α) (smt_tactic α) :=
  ⟨monadLift⟩

unsafe instance : MonadFail smt_tactic :=
  { smt_tactic.monad with fail := fun α s => (tactic.fail (to_fmt s) : smt_tactic α) }

namespace SmtTactic

open Tactic (Transparency)

unsafe axiom intros : smt_tactic Unit
#align smt_tactic.intros smt_tactic.intros

unsafe axiom intron : Nat → smt_tactic Unit
#align smt_tactic.intron smt_tactic.intron

unsafe axiom intro_lst : List Name → smt_tactic Unit
#align smt_tactic.intro_lst smt_tactic.intro_lst

/-- Try to close main goal by using equalities implied by the congruence
  closure module.
-/
unsafe axiom close : smt_tactic Unit
#align smt_tactic.close smt_tactic.close

/-- Produce new facts using heuristic lemma instantiation based on E-matching.
  This tactic tries to match patterns from lemmas in the main goal with terms
  in the main goal. The set of lemmas is populated with theorems
  tagged with the attribute specified at smt_config.em_attr, and lemmas
  added using tactics such as `smt_tactic.add_lemmas`.
  The current set of lemmas can be retrieved using the tactic `smt_tactic.get_lemmas`.

  Remark: the given predicate is applied to every new instance. The instance
  is only added to the state if the predicate returns tt.
-/
unsafe axiom ematch_core : (expr → Bool) → smt_tactic Unit
#align smt_tactic.ematch_core smt_tactic.ematch_core

/-- Produce new facts using heuristic lemma instantiation based on E-matching.
  This tactic tries to match patterns from the given lemmas with terms in
  the main goal.
-/
unsafe axiom ematch_using : hinst_lemmas → smt_tactic Unit
#align smt_tactic.ematch_using smt_tactic.ematch_using

unsafe axiom mk_ematch_eqn_lemmas_for_core : Transparency → Name → smt_tactic hinst_lemmas
#align smt_tactic.mk_ematch_eqn_lemmas_for_core smt_tactic.mk_ematch_eqn_lemmas_for_core

unsafe axiom to_cc_state : smt_tactic cc_state
#align smt_tactic.to_cc_state smt_tactic.to_cc_state

unsafe axiom to_em_state : smt_tactic ematch_state
#align smt_tactic.to_em_state smt_tactic.to_em_state

unsafe axiom get_config : smt_tactic SmtConfig
#align smt_tactic.get_config smt_tactic.get_config

/-- Preprocess the given term using the same simplifications rules used when
  we introduce a new hypothesis. The result is pair containing the resulting
  term and a proof that it is equal to the given one.
-/
unsafe axiom preprocess : expr → smt_tactic (expr × expr)
#align smt_tactic.preprocess smt_tactic.preprocess

unsafe axiom get_lemmas : smt_tactic hinst_lemmas
#align smt_tactic.get_lemmas smt_tactic.get_lemmas

unsafe axiom set_lemmas : hinst_lemmas → smt_tactic Unit
#align smt_tactic.set_lemmas smt_tactic.set_lemmas

unsafe axiom add_lemmas : hinst_lemmas → smt_tactic Unit
#align smt_tactic.add_lemmas smt_tactic.add_lemmas

unsafe def add_ematch_lemma_core (md : Transparency) (as_simp : Bool) (e : expr) :
    smt_tactic Unit := do
  let h ← hinst_lemma.mk_core md e as_simp
  add_lemmas (mk_hinst_singleton h)
#align smt_tactic.add_ematch_lemma_core smt_tactic.add_ematch_lemma_core

unsafe def add_ematch_lemma_from_decl_core (md : Transparency) (as_simp : Bool) (n : Name) :
    smt_tactic Unit := do
  let h ← hinst_lemma.mk_from_decl_core md n as_simp
  add_lemmas (mk_hinst_singleton h)
#align smt_tactic.add_ematch_lemma_from_decl_core smt_tactic.add_ematch_lemma_from_decl_core

unsafe def add_ematch_eqn_lemmas_for_core (md : Transparency) (n : Name) : smt_tactic Unit := do
  let hs ← mk_ematch_eqn_lemmas_for_core md n
  add_lemmas hs
#align smt_tactic.add_ematch_eqn_lemmas_for_core smt_tactic.add_ematch_eqn_lemmas_for_core

unsafe def ematch : smt_tactic Unit :=
  ematch_core fun _ => true
#align smt_tactic.ematch smt_tactic.ematch

unsafe def failed {α} : smt_tactic α :=
  tactic.failed
#align smt_tactic.failed smt_tactic.failed

unsafe def fail {α : Type} {β : Type u} [has_to_format β] (msg : β) : smt_tactic α :=
  tactic.fail msg
#align smt_tactic.fail smt_tactic.fail

unsafe def try {α : Type} (t : smt_tactic α) : smt_tactic Unit :=
  ⟨fun ss ts =>
    result.cases_on (t.run ss ts) (fun ⟨a, new_ss⟩ => result.success ((), new_ss)) fun e ref s' =>
      result.success ((), ss) ts⟩
#align smt_tactic.try smt_tactic.try

/-- `iterate_at_most n t`: repeat the given tactic at most n times or until t fails -/
unsafe def iterate_at_most : Nat → smt_tactic Unit → smt_tactic Unit
  | 0, t => return ()
  | n + 1, t =>
    (do
        t
        iterate_at_most n t) <|>
      return ()
#align smt_tactic.iterate_at_most smt_tactic.iterate_at_most

/-- `iterate_exactly n t` : execute t n times -/
unsafe def iterate_exactly : Nat → smt_tactic Unit → smt_tactic Unit
  | 0, t => return ()
  | n + 1, t => do
    t
    iterate_exactly n t
#align smt_tactic.iterate_exactly smt_tactic.iterate_exactly

unsafe def iterate : smt_tactic Unit → smt_tactic Unit :=
  iterate_at_most 100000
#align smt_tactic.iterate smt_tactic.iterate

unsafe def eblast : smt_tactic Unit :=
  iterate (ematch >> try close)
#align smt_tactic.eblast smt_tactic.eblast

open Tactic

protected unsafe def read : smt_tactic (smt_state × tactic_state) := do
  let s₁ ← get
  let s₂ ← tactic.read
  return (s₁, s₂)
#align smt_tactic.read smt_tactic.read

protected unsafe def write : smt_state × tactic_state → smt_tactic Unit := fun ⟨ss, ts⟩ =>
  ⟨fun _ _ => result.success ((), ss) ts⟩
#align smt_tactic.write smt_tactic.write

private unsafe def mk_smt_goals_for (cfg : SmtConfig) :
    List expr → List smt_goal → List expr → tactic (List smt_goal × List expr)
  | [], sr, tr => return (sr.reverse, tr.reverse)
  | tg :: tgs, sr, tr => do
    tactic.set_goals [tg]
    let [new_sg] ← smt_state.mk cfg |
      tactic.failed
    let [new_tg] ← get_goals |
      tactic.failed
    mk_smt_goals_for tgs (new_sg :: sr) (new_tg :: tr)
#align smt_tactic.mk_smt_goals_for smt_tactic.mk_smt_goals_for

/-- See slift -/
unsafe def slift_aux {α : Type} (t : tactic α) (cfg : SmtConfig) : smt_tactic α :=
  ⟨fun ss => do
    let _ :: sgs ← return ss |
      tactic.fail "slift tactic failed, there no smt goals to be solved"
    let tg :: tgs ← tactic.get_goals |
      tactic.failed
    tactic.set_goals [tg]
    let a ← t
    let new_tgs ← tactic.get_goals
    let (new_sgs, new_tgs) ← mk_smt_goals_for cfg new_tgs [] []
    tactic.set_goals (new_tgs ++ tgs)
    return (a, new_sgs ++ sgs)⟩
#align smt_tactic.slift_aux smt_tactic.slift_aux

/-- This lift operation will restart the SMT state.
  It is useful for using tactics that change the set of hypotheses. -/
unsafe def slift {α : Type} (t : tactic α) : smt_tactic α :=
  get_config >>= slift_aux t
#align smt_tactic.slift smt_tactic.slift

unsafe def trace_state : smt_tactic Unit := do
  let (s₁, s₂) ← smt_tactic.read
  trace (smt_state.to_format s₁ s₂)
#align smt_tactic.trace_state smt_tactic.trace_state

unsafe def trace {α : Type} [has_to_tactic_format α] (a : α) : smt_tactic Unit :=
  tactic.trace a
#align smt_tactic.trace smt_tactic.trace

unsafe def to_expr (q : pexpr) (allow_mvars := true) : smt_tactic expr :=
  tactic.to_expr q allow_mvars
#align smt_tactic.to_expr smt_tactic.to_expr

unsafe def classical : smt_tactic Bool := do
  let s ← get
  return s
#align smt_tactic.classical smt_tactic.classical

unsafe def num_goals : smt_tactic Nat :=
  List.length <$> get
#align smt_tactic.num_goals smt_tactic.num_goals

-- Low level primitives for managing set of goals
unsafe def get_goals : smt_tactic (List smt_goal × List expr) := do
  let (g₁, _) ← smt_tactic.read
  let g₂ ← tactic.get_goals
  return (g₁, g₂)
#align smt_tactic.get_goals smt_tactic.get_goals

unsafe def set_goals : List smt_goal → List expr → smt_tactic Unit := fun g₁ g₂ =>
  ⟨fun ss => tactic.set_goals g₂ >> return ((), g₁)⟩
#align smt_tactic.set_goals smt_tactic.set_goals

private unsafe def all_goals_core (tac : smt_tactic Unit) :
    List smt_goal → List expr → List smt_goal → List expr → smt_tactic Unit
  | [], ts, acs, act => set_goals acs (ts ++ act)
  | s :: ss, [], acs, act => fail "ill-formed smt_state"
  | s :: ss, t :: ts, acs, act => do
    set_goals [s] [t]
    tac
    let (new_ss, new_ts) ← get_goals
    all_goals_core ss ts (acs ++ new_ss) (act ++ new_ts)
#align smt_tactic.all_goals_core smt_tactic.all_goals_core

/-- Apply the given tactic to all goals. -/
unsafe def all_goals (tac : smt_tactic Unit) : smt_tactic Unit := do
  let (ss, ts) ← get_goals
  all_goals_core tac ss ts [] []
#align smt_tactic.all_goals smt_tactic.all_goals

/--
LCF-style AND_THEN tactic. It applies tac1, and if succeed applies tac2 to each subgoal produced by tac1 -/
unsafe def seq (tac1 : smt_tactic Unit) (tac2 : smt_tactic Unit) : smt_tactic Unit := do
  let (s :: ss, t :: ts) ← get_goals
  set_goals [s] [t]
  tac1
  all_goals tac2
  let (new_ss, new_ts) ← get_goals
  set_goals (new_ss ++ ss) (new_ts ++ ts)
#align smt_tactic.seq smt_tactic.seq

unsafe instance : AndThen' (smt_tactic Unit) (smt_tactic Unit) (smt_tactic Unit) :=
  ⟨seq⟩

unsafe def focus1 {α} (tac : smt_tactic α) : smt_tactic α := do
  let (s :: ss, t :: ts) ← get_goals
  match ss with
    | [] => tac
    | _ => do
      set_goals [s] [t]
      let a ← tac
      let (ss', ts') ← get_goals
      set_goals (ss' ++ ss) (ts' ++ ts)
      return a
#align smt_tactic.focus1 smt_tactic.focus1

unsafe def solve1 (tac : smt_tactic Unit) : smt_tactic Unit := do
  let (ss, gs) ← get_goals
  match ss, gs with
    | [], _ => fail "solve1 tactic failed, there isn't any goal left to focus"
    | _, [] => fail "solve1 tactic failed, there isn't any smt goal left to focus"
    | s :: ss, g :: gs => do
      set_goals [s] [g]
      tac
      let (ss', gs') ← get_goals
      match ss', gs' with
        | [], [] => set_goals ss gs
        | _, _ => fail "solve1 tactic failed, focused goal has not been solved"
#align smt_tactic.solve1 smt_tactic.solve1

unsafe def swap : smt_tactic Unit := do
  let (ss, ts) ← get_goals
  match ss, ts with
    | s₁ :: s₂ :: ss, t₁ :: t₂ :: ts => set_goals (s₂ :: s₁ :: ss) (t₂ :: t₁ :: ts)
    | _, _ => failed
#align smt_tactic.swap smt_tactic.swap

/-- Add a new goal for t, and the hypothesis (h : t) in the current goal. -/
unsafe def assert (h : Name) (t : expr) : smt_tactic Unit :=
  (((tactic.assert_core h t >> swap) >> intros) >> swap) >> try close
#align smt_tactic.assert smt_tactic.assert

/-- Add the hypothesis (h : t) in the current goal if v has type t. -/
unsafe def assertv (h : Name) (t : expr) (v : expr) : smt_tactic Unit :=
  (tactic.assertv_core h t v >> intros) >> return ()
#align smt_tactic.assertv smt_tactic.assertv

/-- Add a new goal for t, and the hypothesis (h : t := ?M) in the current goal. -/
unsafe def define (h : Name) (t : expr) : smt_tactic Unit :=
  (((tactic.define_core h t >> swap) >> intros) >> swap) >> try close
#align smt_tactic.define smt_tactic.define

/-- Add the hypothesis (h : t := v) in the current goal if v has type t. -/
unsafe def definev (h : Name) (t : expr) (v : expr) : smt_tactic Unit :=
  (tactic.definev_core h t v >> intros) >> return ()
#align smt_tactic.definev smt_tactic.definev

/-- Add (h : t := pr) to the current goal -/
unsafe def pose (h : Name) (t : Option expr := none) (pr : expr) : smt_tactic Unit :=
  match t with
  | none => do
    let t ← infer_type pr
    definev h t pr
  | some t => definev h t pr
#align smt_tactic.pose smt_tactic.pose

/-- Add (h : t) to the current goal, given a proof (pr : t) -/
unsafe def note (h : Name) (t : Option expr := none) (pr : expr) : smt_tactic Unit :=
  match t with
  | none => do
    let t ← infer_type pr
    assertv h t pr
  | some t => assertv h t pr
#align smt_tactic.note smt_tactic.note

unsafe def destruct (e : expr) : smt_tactic Unit :=
  smt_tactic.seq (tactic.destruct e) smt_tactic.intros
#align smt_tactic.destruct smt_tactic.destruct

unsafe def by_cases (e : expr) : smt_tactic Unit := do
  let c ← classical
  if c then destruct (expr.app (expr.const `classical.em []) e)
    else do
      let dec_e ←
        mk_app `decidable [e] <|> fail "by_cases smt_tactic failed, type is not a proposition"
      let inst ←
        mk_instance dec_e <|>
            fail "by_cases smt_tactic failed, type of given expression is not decidable"
      let em ← mk_app `decidable.em [e, inst]
      destruct em
#align smt_tactic.by_cases smt_tactic.by_cases

unsafe def by_contradiction : smt_tactic Unit := do
  let t ← target
  let c ← classical
  if t then skip
    else
      if c then do
        apply (expr.app (expr.const `classical.by_contradiction []) t)
        intros
      else do
        let dec_t ←
          mk_app `decidable [t] <|>
              fail "by_contradiction smt_tactic failed, target is not a proposition"
        let inst ←
          mk_instance dec_t <|> fail "by_contradiction smt_tactic failed, target is not decidable"
        let a ← mk_mapp `decidable.by_contradiction [some t, some inst]
        apply a
        intros
#align smt_tactic.by_contradiction smt_tactic.by_contradiction

/-- Return a proof for e, if 'e' is a known fact in the main goal. -/
unsafe def proof_for (e : expr) : smt_tactic expr := do
  let cc ← to_cc_state
  cc e
#align smt_tactic.proof_for smt_tactic.proof_for

/--
Return a refutation for e (i.e., a proof for (not e)), if 'e' has been refuted in the main goal. -/
unsafe def refutation_for (e : expr) : smt_tactic expr := do
  let cc ← to_cc_state
  cc e
#align smt_tactic.refutation_for smt_tactic.refutation_for

unsafe def get_facts : smt_tactic (List expr) := do
  let cc ← to_cc_state
  return <| cc expr.mk_true
#align smt_tactic.get_facts smt_tactic.get_facts

unsafe def get_refuted_facts : smt_tactic (List expr) := do
  let cc ← to_cc_state
  return <| cc expr.mk_false
#align smt_tactic.get_refuted_facts smt_tactic.get_refuted_facts

unsafe def add_ematch_lemma : expr → smt_tactic Unit :=
  add_ematch_lemma_core reducible false
#align smt_tactic.add_ematch_lemma smt_tactic.add_ematch_lemma

unsafe def add_ematch_lhs_lemma : expr → smt_tactic Unit :=
  add_ematch_lemma_core reducible true
#align smt_tactic.add_ematch_lhs_lemma smt_tactic.add_ematch_lhs_lemma

unsafe def add_ematch_lemma_from_decl : Name → smt_tactic Unit :=
  add_ematch_lemma_from_decl_core reducible false
#align smt_tactic.add_ematch_lemma_from_decl smt_tactic.add_ematch_lemma_from_decl

unsafe def add_ematch_lhs_lemma_from_decl : Name → smt_tactic Unit :=
  add_ematch_lemma_from_decl_core reducible false
#align smt_tactic.add_ematch_lhs_lemma_from_decl smt_tactic.add_ematch_lhs_lemma_from_decl

unsafe def add_ematch_eqn_lemmas_for : Name → smt_tactic Unit :=
  add_ematch_eqn_lemmas_for_core reducible
#align smt_tactic.add_ematch_eqn_lemmas_for smt_tactic.add_ematch_eqn_lemmas_for

/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `f -/
unsafe def add_lemmas_from_facts_core : List expr → smt_tactic Unit
  | [] => return ()
  | f :: fs => do
    try
        ((is_prop f >> guard (f && not (f f.is_arrow))) >> proof_for f >>=
          add_ematch_lemma_core reducible ff)
    add_lemmas_from_facts_core fs
#align smt_tactic.add_lemmas_from_facts_core smt_tactic.add_lemmas_from_facts_core

unsafe def add_lemmas_from_facts : smt_tactic Unit :=
  get_facts >>= add_lemmas_from_facts_core
#align smt_tactic.add_lemmas_from_facts smt_tactic.add_lemmas_from_facts

unsafe def induction (e : expr) (ids : List Name := []) (rec : Option Name := none) :
    smt_tactic Unit :=
  slift (tactic.induction e ids rec >> return ())
#align smt_tactic.induction smt_tactic.induction

-- pass on the information?
unsafe def when (c : Prop) [Decidable c] (tac : smt_tactic Unit) : smt_tactic Unit :=
  if c then tac else skip
#align smt_tactic.when smt_tactic.when

unsafe def when_tracing (n : Name) (tac : smt_tactic Unit) : smt_tactic Unit :=
  when (is_trace_enabled_for n = true) tac
#align smt_tactic.when_tracing smt_tactic.when_tracing

end SmtTactic

open SmtTactic

unsafe def using_smt {α} (t : smt_tactic α) (cfg : SmtConfig := { }) : tactic α := do
  let ss ← smt_state.mk cfg
  let (a, _) ←
    (do
            let a ← t
            iterate close
            return a).run
        ss
  return a
#align using_smt using_smt

unsafe def using_smt_with {α} (cfg : SmtConfig) (t : smt_tactic α) : tactic α :=
  using_smt t cfg
#align using_smt_with using_smt_with


import Lean

namespace SciLean.Meta.CustomSimp

open Lean Meta Elab Elab.Term

-- Maybe add option to write logical formulas in simp guard
-- something like @[simp_guard (n = 0) || (m = n)]

/-- Prevent applying the theorem if the specified argument is equation to the specified value.

Warning: It only works with custom simplifiers! Normal simplifier ignores this.

Example:
---
Adding the following simp guard to the chain rule prevents appying it if `g` is an identity function.

```
  @[simp_guard g (λ x => x)]
  theorem chain_rule {α β γ : Type}
    (f : β → γ) (g : α → β)
    : ∂ (λ x => f (g x)) = λ x dx => ∂ f (g x) (∂ g x dx) := ...
```
-/
syntax (name := simp_guard) "simp_guard" (ident term),+ : attr

initialize simpGuardAttr : ParametricAttribute (Array (Nat × Expr × Nat)) ←
  registerParametricAttribute {
    name := `simp_guard
    descr := "Do not apply this simp theorem if the specified argument has the specified value."
    getParam := fun name => fun
      | `(attr| simp_guard $[$ids $vals],*) =>
        MetaM.run' <| TermElabM.run' <| (ids.zip vals).mapM λ (id, val) => do
          let info ← getConstInfo name

          let nth ← forallTelescope info.type λ args _ => do
            let i? ← args.findIdxM?
              (λ arg => do
                let argDecl ← getFVarLocalDecl arg
                pure (argDecl.userName = id.getId))
            match i? with
            | some i => pure i
            | none => throwError "Theorem does not have an argument with the name `{id.getId}`"

          -- `valueFun` is a function taking all theorem arguments [0,..,nth) and returning guard value
          let (valFun, numMVars) ← forallBoundedTelescope info.type nth λ args _ => do
            let value ← elabTerm val none
            let value ← abstractMVars value

            pure (← mkLambdaFVars args value.expr, value.numMVars)

          pure (nth, valFun, numMVars)
      | _ => Elab.throwUnsupportedSyntax
  }

def hasCustomSimpGuard (env : Environment) (n : Name) : Bool :=
  match simpGuardAttr.getParam? env n with
  | some _ => true
  | none => false


import data.pfun
import phase2.approximation
import phase2.complete_orbit
import phase2.reduction

open cardinal quiver set sum with_bot
open_locale cardinal classical pointwise

universe u

namespace con_nf
variable [params.{u}]

/-!
# Weak approximations
-/

/-- Noncomputably eliminates a disjunction into a (possibly predicative) universe. -/
noncomputable def _root_.or.elim' {α : Sort*} {p q : Prop}
  (h : p ∨ q) (f : p → α) (g : q → α) : α :=
if hp : p then f hp else g (h.resolve_left hp)

lemma _root_.or.elim'_left {α : Sort*} {p q : Prop}
  (h : p ∨ q) (f : p → α) (g : q → α) (hp : p) : h.elim' f g = f hp :=
by rw [or.elim', dif_pos hp]

lemma _root_.or.elim'_right {α : Sort*} {p q : Prop}
  (h : p ∨ q) (f : p → α) (g : q → α) (hp : ¬p) : h.elim' f g = g (h.resolve_left hp) :=
by rw [or.elim', dif_neg hp]

/-- A *weak near-litter approximation* is a partial function from atoms to atoms and a partial
function from litters to near-litters, both of which have small domain.
The image of a litter under the `litter_map` should be interpreted as the intended *precise* image
of this litter under an allowable permutation.
The atom and litter maps should be injective (in suitable senses) and cohere in the sense that
images of atoms in litters are mapped to atoms inside the corresponding near-litters. -/
@[ext] structure weak_near_litter_approx :=
(atom_map : atom →. atom)
(litter_map : litter →. near_litter)
(atom_map_dom_small : small atom_map.dom)
(litter_map_dom_small : small litter_map.dom)
(atom_map_injective : ∀ ⦃a b⦄ ha hb, (atom_map a).get ha = (atom_map b).get hb → a = b)
(litter_map_injective : ∀ ⦃L₁ L₂ : litter⦄ hL₁ hL₂,
  (((litter_map L₁).get hL₁ : set atom) ∩ (litter_map L₂).get hL₂).nonempty → L₁ = L₂)
(atom_mem : ∀ (a : atom) ha L hL, a.1 = L ↔ (atom_map a).get ha ∈ (litter_map L).get hL)

/-- A `β`-weak structural approximation is a product that assigns a weak near-litter approximation
to each `β`-extended index. -/
def weak_struct_approx (β : type_index) := extended_index β → weak_near_litter_approx

namespace weak_near_litter_approx

variable (w : weak_near_litter_approx)

/-- A litter that is not allowed to be used as a sandbox because it appears somewhere that
we need to preserve. -/
@[mk_iff] inductive banned_litter : litter → Prop
| atom_dom (a : atom) : (w.atom_map a).dom → banned_litter a.1
| litter_dom (L : litter) : (w.litter_map L).dom → banned_litter L
| atom_map (a : atom) (h) : banned_litter ((w.atom_map a).get h).1
| litter_map (L : litter) (h) : banned_litter ((w.litter_map L).get h).1
| diff (L : litter) (h) (a : atom) :
    a ∈ ((w.litter_map L).get h : set atom) \ litter_set ((w.litter_map L).get h).1 →
    banned_litter a.1

lemma banned_litter.mem_map (a : atom) (L : litter) (hL)
  (ha : a ∈ ((w.litter_map L).get hL : set atom)) : w.banned_litter a.1 :=
begin
  by_cases a.1 = ((w.litter_map L).get hL).1,
  { rw h,
    exact banned_litter.litter_map L hL, },
  { exact banned_litter.diff L hL a ⟨ha, h⟩, },
end

/-- There are only a small amount of banned litters. -/
lemma banned_litter_small : small {L | w.banned_litter L} :=
begin
  simp only [banned_litter_iff, mem_diff, set_like.mem_coe, mem_litter_set],
  refine small.union _ (small.union _ (small.union _ (small.union _ _))),
  { refine lt_of_le_of_lt _ w.atom_map_dom_small,
    refine ⟨⟨λ a, ⟨_, a.prop.some_spec.1⟩, λ a₁ a₂ h, _⟩⟩,
    simp only [subtype.mk_eq_mk, prod.mk.inj_iff] at h,
    have := a₁.prop.some_spec.2,
    rw h at this,
    exact subtype.coe_injective (this.trans a₂.prop.some_spec.2.symm), },
  { refine lt_of_le_of_lt _ w.litter_map_dom_small,
    refine ⟨⟨λ L, ⟨_, L.prop⟩, λ L₁ L₂ h, _⟩⟩,
    simp only [subtype.mk_eq_mk, prod.mk.inj_iff] at h,
    exact subtype.coe_injective h, },
  { refine lt_of_le_of_lt _ w.atom_map_dom_small,
    refine ⟨⟨λ L, ⟨_, L.prop.some_spec.some⟩, λ L₁ L₂ h, _⟩⟩,
    simp only [subtype.mk_eq_mk, prod.mk.inj_iff] at h,
    have := L₁.prop.some_spec.some_spec,
    simp_rw h at this,
    exact subtype.coe_injective (this.trans L₂.prop.some_spec.some_spec.symm), },
  { refine lt_of_le_of_lt _ w.litter_map_dom_small,
    refine ⟨⟨λ L, ⟨_, L.prop.some_spec.some⟩, λ L₁ L₂ h, _⟩⟩,
    simp only [subtype.mk_eq_mk, prod.mk.inj_iff] at h,
    have := L₁.prop.some_spec.some_spec,
    simp_rw h at this,
    exact subtype.coe_injective (this.trans L₂.prop.some_spec.some_spec.symm), },
  { have : small ⋃ (L : litter) (h : (w.litter_map L).dom),
      ((w.litter_map L).get h : set atom) \ litter_set ((w.litter_map L).get h).1,
    { refine small.bUnion _ _,
      { refine lt_of_le_of_lt _ w.litter_map_dom_small,
        refine ⟨⟨λ N, ⟨_, N.prop⟩, λ N₁ N₂ h, _⟩⟩,
        simp only [subtype.mk_eq_mk, prod.mk.inj_iff] at h,
        exact subtype.coe_inj.mp h, },
      { intros L hL,
        refine small.mono _ ((w.litter_map L).get hL).2.prop,
        exact λ x hx, or.inr hx, }, },
    refine lt_of_le_of_lt _ this,
    refine ⟨⟨λ L, ⟨L.prop.some_spec.some_spec.some, _⟩, λ L₁ L₂ h, _⟩⟩,
    { simp only [mem_Union],
      exact ⟨_, _, L.prop.some_spec.some_spec.some_spec.1⟩, },
    simp only [subtype.mk_eq_mk, prod.mk.inj_iff] at h,
    have := L₁.prop.some_spec.some_spec.some_spec.2,
    rw h at this,
    exact subtype.coe_injective
      (this.trans L₂.prop.some_spec.some_spec.some_spec.2.symm), },
end

lemma mk_not_banned_litter : #{L | ¬w.banned_litter L} = #μ :=
begin
  have := mk_sum_compl {L | w.banned_litter L},
  rw [compl_set_of, mk_litter] at this,
  rw [← this, add_eq_right],
  { by_contra' h,
    have h' := add_le_add (le_of_lt w.banned_litter_small) h.le,
    rw this at h',
    refine not_lt_of_le h' _,
    refine cardinal.add_lt_of_lt μ_strong_limit.is_limit.aleph_0_le κ_lt_μ _,
    exact lt_of_le_of_lt κ_regular.aleph_0_le κ_lt_μ, },
  { by_contra' h,
    have h' := add_le_add (le_of_lt w.banned_litter_small) h.le,
    rw this at h',
    refine not_lt_of_le h' _,
    refine cardinal.add_lt_of_lt μ_strong_limit.is_limit.aleph_0_le κ_lt_μ _,
    exact lt_trans w.banned_litter_small κ_lt_μ, },
end

lemma not_banned_litter_nonempty : nonempty {L | ¬w.banned_litter L} :=
by simp only [← mk_ne_zero_iff, mk_not_banned_litter, ne.def, mk_ne_zero, not_false_iff]

/-- The *sandbox litter* for a weak near-litter approximation is an arbitrarily chosen litter that
isn't banned. -/
noncomputable def sandbox_litter : litter := w.not_banned_litter_nonempty.some

lemma sandbox_litter_not_banned : ¬w.banned_litter w.sandbox_litter :=
w.not_banned_litter_nonempty.some.prop

/-- If `a` is in the domain, this is the atom map. Otherwise, this gives an arbitrary atom. -/
noncomputable def atom_map_or_else (a : atom) : atom := (w.atom_map a).get_or_else (arbitrary atom)

lemma atom_map_or_else_of_dom {a : atom} (ha : (w.atom_map a).dom) :
  w.atom_map_or_else a = (w.atom_map a).get ha :=
by rw [atom_map_or_else, part.get_or_else_of_dom]

lemma mk_atom_map_image_le_mk_sandbox :
  #(w.atom_map.dom ∆ (w.atom_map_or_else '' w.atom_map.dom) : set atom) ≤
    #(litter_set w.sandbox_litter) :=
begin
  rw mk_litter_set,
  refine le_trans (mk_subtype_mono symm_diff_subset_union) (le_trans (mk_union_le _ _) _),
  refine add_le_of_le κ_regular.aleph_0_le _ _,
  exact le_of_lt w.atom_map_dom_small,
  exact le_trans mk_image_le (le_of_lt w.atom_map_dom_small),
end

lemma disjoint_sandbox :
  disjoint (w.atom_map.dom ∪ w.atom_map_or_else '' w.atom_map.dom) (litter_set w.sandbox_litter) :=
begin
  rw [disjoint_iff_inter_eq_empty, eq_empty_iff_forall_not_mem],
  rintros a ⟨ha₁, ha₂⟩,
  rw mem_litter_set at ha₂,
  have hnb := w.sandbox_litter_not_banned,
  rw ← ha₂ at hnb,
  cases ha₁,
  { exact hnb (banned_litter.atom_dom a ha₁), },
  { refine hnb _,
    simp only [mem_image, pfun.mem_dom] at ha₁,
    obtain ⟨b, ⟨_, hb, rfl⟩, rfl⟩ := ha₁,
    rw w.atom_map_or_else_of_dom hb,
    exact banned_litter.atom_map b hb, },
end

lemma atom_map_or_else_injective : inj_on w.atom_map_or_else w.atom_map.dom :=
begin
  intros a ha b hb h,
  rw [w.atom_map_or_else_of_dom ha, w.atom_map_or_else_of_dom hb] at h,
  exact w.atom_map_injective ha hb h,
end

/-- If `L` is in the domain, this is the litter map.
Otherwise, this gives an arbitrary near-litter. -/
noncomputable def litter_map_or_else (L : litter) : near_litter :=
(w.litter_map L).get_or_else (arbitrary near_litter)

lemma litter_map_or_else_of_dom {L : litter} (hL : (w.litter_map L).dom) :
  w.litter_map_or_else L = (w.litter_map L).get hL :=
by rw [litter_map_or_else, part.get_or_else_of_dom]

noncomputable def rough_litter_map_or_else (L : litter) : litter :=
(w.litter_map_or_else L).1

lemma rough_litter_map_or_else_of_dom {L : litter} (hL : (w.litter_map L).dom) :
  w.rough_litter_map_or_else L = ((w.litter_map L).get hL).1 :=
by rw [rough_litter_map_or_else, litter_map_or_else_of_dom]

/-- The induced action of this weak approximation on near-litters. -/
noncomputable def near_litter_map_or_else (N : near_litter) : near_litter :=
⟨(w.litter_map_or_else N.fst).fst,
  w.litter_map_or_else N.fst ∆ (w.atom_map_or_else '' litter_set N.fst ∆ N),
  begin
    rw [is_near_litter, is_near, ← symm_diff_assoc],
    exact (w.litter_map_or_else N.fst).snd.prop.symm_diff (small.image N.2.prop),
  end⟩

/-- A weak approximation is precise at a litter in its domain if all atoms in the symmetric
difference of its image are accounted for. -/
@[mk_iff] structure precise_at {L : litter} (hL : (w.litter_map L).dom) : Prop :=
(diff : ((w.litter_map L).get hL : set atom) ∆ litter_set ((w.litter_map L).get hL).1 ⊆
  w.atom_map.ran)
(fwd : ∀ a ha, (w.atom_map a).get ha ∈ litter_set L → (w.atom_map ((w.atom_map a).get ha)).dom)
(back : w.atom_map.dom ∩ (w.litter_map L).get hL ⊆ w.atom_map.ran)

/-- A weak approximation is precise if it is precise at every litter in its domain. -/
def precise : Prop := ∀ ⦃L⦄ (hL : (w.litter_map L).dom), w.precise_at hL

/-!
## Induced litter permutation
-/

lemma mk_dom_symm_diff_le :
  #↥(w.litter_map.dom ∆ (w.rough_litter_map_or_else '' w.litter_map.dom)) ≤
  #{L : litter | ¬w.banned_litter L} :=
begin
  rw mk_not_banned_litter,
  refine le_trans (le_of_lt _) κ_le_μ,
  exact small.symm_diff w.litter_map_dom_small w.litter_map_dom_small.image,
end

lemma aleph_0_le_not_banned_litter : ℵ₀ ≤ #{L | ¬w.banned_litter L} :=
begin
  rw mk_not_banned_litter,
  exact μ_strong_limit.is_limit.aleph_0_le,
end

lemma disjoint_dom_not_banned_litter :
  disjoint (w.litter_map.dom ∪ w.rough_litter_map_or_else '' w.litter_map.dom)
    {L : litter | ¬w.banned_litter L} :=
begin
  simp only [set.disjoint_left, mem_union, pfun.mem_dom, mem_image, mem_set_of_eq, not_not],
  rintros _ (⟨_, hL, rfl⟩ | ⟨L, ⟨_, hL, rfl⟩, rfl⟩),
  { exact banned_litter.litter_dom _ hL, },
  { rw w.rough_litter_map_or_else_of_dom hL,
    exact banned_litter.litter_map _ hL, },
end

lemma rough_litter_map_or_else_inj_on : inj_on w.rough_litter_map_or_else w.litter_map.dom :=
begin
  intros L₁ hL₁ L₂ hL₂ h,
  rw [w.rough_litter_map_or_else_of_dom hL₁, w.rough_litter_map_or_else_of_dom hL₂] at h,
  exact w.litter_map_injective hL₁ hL₂ (near_litter.inter_nonempty_of_fst_eq_fst h),
end

/-- A local permutation on the set of litters that occur in the domain or range of `w`.
This permutes both flexible and inflexible litters. -/
noncomputable def litter_perm' : local_perm litter :=
local_perm.complete
  w.rough_litter_map_or_else
  w.litter_map.dom
  {L | ¬w.banned_litter L}
  w.mk_dom_symm_diff_le
  w.aleph_0_le_not_banned_litter
  w.disjoint_dom_not_banned_litter
  w.rough_litter_map_or_else_inj_on

def id_on_banned (s : set litter) : local_perm litter := {
  to_fun := id,
  inv_fun := id,
  domain := {L | w.banned_litter L} \ s,
  to_fun_domain' := λ L h, h,
  inv_fun_domain' := λ L h, h,
  left_inv' := λ L h, rfl,
  right_inv' := λ L h, rfl,
}

noncomputable def litter_perm : local_perm litter :=
local_perm.piecewise w.litter_perm' (w.id_on_banned w.litter_perm'.domain)
  (by rw ← set.subset_compl_iff_disjoint_left; exact λ L h, h.2)

lemma litter_perm'_apply_eq (L : litter) (hL : L ∈ w.litter_map.dom) :
  w.litter_perm' L = w.rough_litter_map_or_else L :=
local_perm.complete_apply_eq _ _ _ hL

lemma litter_perm_apply_eq (L : litter) (hL : L ∈ w.litter_map.dom) :
  w.litter_perm L = w.rough_litter_map_or_else L :=
begin
  rw ← w.litter_perm'_apply_eq L hL,
  exact local_perm.piecewise_apply_eq_left (or.inl (or.inl hL)),
end

lemma litter_perm'_domain_small : small w.litter_perm'.domain :=
begin
  refine small.union (small.union w.litter_map_dom_small w.litter_map_dom_small.image) _,
  rw small,
  rw cardinal.mk_congr (local_perm.sandbox_subset_equiv _ _),
  simp only [mk_sum, mk_prod, mk_denumerable, lift_aleph_0, lift_uzero, lift_id],
  refine add_lt_of_lt κ_regular.aleph_0_le _ _;
    refine (mul_lt_of_lt κ_regular.aleph_0_le (lt_of_le_of_lt Λ_limit.aleph_0_le Λ_lt_κ) _);
    refine lt_of_le_of_lt (mk_subtype_mono (diff_subset _ _)) _,
  exact w.litter_map_dom_small,
  exact w.litter_map_dom_small.image,
end

lemma litter_perm_domain_small : small w.litter_perm.domain :=
small.union w.litter_perm'_domain_small (small.mono (diff_subset _ _) w.banned_litter_small)

variables {α : Λ} [position_data.{}] [phase_2_assumptions α] {β : Iio α} {A : extended_index β}

lemma mk_not_banned_litter_and_flexible : #{L | ¬w.banned_litter L ∧ flexible α L A} = #μ :=
begin
  refine le_antisymm ((mk_subtype_le _).trans mk_litter.le) _,
  by_contra,
  rw not_le at h,
  have h₁ := cardinal.le_mk_diff_add_mk {L | flexible α L A} {L | w.banned_litter L},
  rw [mk_flexible, diff_eq, inter_comm] at h₁,
  have h₂ := add_lt_of_lt μ_strong_limit.is_limit.aleph_0_le h
    (lt_trans w.banned_litter_small κ_lt_μ),
  exact h₁.not_lt h₂,
end

lemma mk_dom_inter_flexible_symm_diff_le :
  #↥((w.litter_map.dom ∩ {L | flexible α L A}) ∆
    (w.rough_litter_map_or_else '' (w.litter_map.dom ∩ {L | flexible α L A}))) ≤
  #{L : litter | ¬w.banned_litter L ∧ flexible α L A} :=
begin
  rw mk_not_banned_litter_and_flexible,
  refine le_trans (le_of_lt _) κ_le_μ,
  exact small.symm_diff
    (small.mono (inter_subset_left _ _) w.litter_map_dom_small)
    (small.mono (inter_subset_left _ _) w.litter_map_dom_small).image,
end

lemma aleph_0_le_not_banned_litter_and_flexible : ℵ₀ ≤ #{L | ¬w.banned_litter L ∧ flexible α L A} :=
begin
  rw mk_not_banned_litter_and_flexible,
  exact μ_strong_limit.is_limit.aleph_0_le,
end

lemma disjoint_dom_inter_flexible_not_banned_litter :
  disjoint ((w.litter_map.dom ∩ {L | flexible α L A})
    ∪ w.rough_litter_map_or_else '' (w.litter_map.dom ∩ {L | flexible α L A}))
    {L : litter | ¬w.banned_litter L ∧ flexible α L A} :=
begin
  refine disjoint_of_subset _ (inter_subset_left _ _) w.disjoint_dom_not_banned_litter,
  rintros a (ha | ⟨b, hb, rfl⟩),
  exact or.inl ha.1,
  exact or.inr ⟨b, hb.1, rfl⟩,
end

lemma rough_litter_map_or_else_inj_on_dom_inter_flexible :
  inj_on w.rough_litter_map_or_else (w.litter_map.dom ∩ {L | flexible α L A}) :=
w.rough_litter_map_or_else_inj_on.mono (inter_subset_left _ _)

noncomputable def flexible_litter_perm (A : extended_index β) :
  local_perm litter :=
local_perm.complete
  w.rough_litter_map_or_else
  (w.litter_map.dom ∩ {L | flexible α L A})
  {L | ¬w.banned_litter L ∧ flexible α L A}
  w.mk_dom_inter_flexible_symm_diff_le
  w.aleph_0_le_not_banned_litter_and_flexible
  w.disjoint_dom_inter_flexible_not_banned_litter
  w.rough_litter_map_or_else_inj_on_dom_inter_flexible

lemma flexible_litter_perm_apply_eq (L : litter)
  (hL₁ : L ∈ w.litter_map.dom) (hL₂ : flexible α L A) :
  w.flexible_litter_perm A L = w.rough_litter_map_or_else L :=
local_perm.complete_apply_eq _ _ _ ⟨hL₁, hL₂⟩

lemma flexible_litter_perm_domain_small : small (w.flexible_litter_perm A).domain :=
begin
  refine small.union (small.union _ _) _,
  { exact w.litter_map_dom_small.mono (inter_subset_left _ _) },
  { exact (w.litter_map_dom_small.mono (inter_subset_left _ _)).image, },
  { rw small,
    rw cardinal.mk_congr (local_perm.sandbox_subset_equiv _ _),
    simp only [mk_sum, mk_prod, mk_denumerable, lift_aleph_0, lift_uzero, lift_id],
    refine add_lt_of_lt κ_regular.aleph_0_le _ _;
      refine (mul_lt_of_lt κ_regular.aleph_0_le (lt_of_le_of_lt Λ_limit.aleph_0_le Λ_lt_κ) _);
      refine lt_of_le_of_lt (mk_subtype_mono (diff_subset _ _)) _,
    exact w.litter_map_dom_small.mono (inter_subset_left _ _),
    exact (w.litter_map_dom_small.mono (inter_subset_left _ _)).image, },
end

/-!
# Completed permutations
-/

/-- A local permutation induced by completing the orbits of atoms in a weak near-litter
approximation. This function creates forward and backward images of atoms in the *sandbox litter*,
a litter which is away from the domain and range of the approximation in question, so it should
not interfere with other constructions. -/
noncomputable def complete_atom_perm : local_perm atom :=
local_perm.complete
  w.atom_map_or_else
  w.atom_map.dom
  (litter_set w.sandbox_litter)
  w.mk_atom_map_image_le_mk_sandbox
  (by simpa only [mk_litter_set] using κ_regular.aleph_0_le)
  w.disjoint_sandbox
  w.atom_map_or_else_injective

lemma sandbox_subset_small : small (local_perm.sandbox_subset
  w.mk_atom_map_image_le_mk_sandbox
  (by simpa only [mk_litter_set] using κ_regular.aleph_0_le)) :=
begin
  rw small,
  rw cardinal.mk_congr (local_perm.sandbox_subset_equiv _ _),
  simp only [mk_sum, mk_prod, mk_denumerable, lift_aleph_0, lift_uzero, lift_id],
  refine add_lt_of_lt κ_regular.aleph_0_le _ _;
    refine (mul_lt_of_lt κ_regular.aleph_0_le (lt_of_le_of_lt Λ_limit.aleph_0_le Λ_lt_κ) _);
    refine lt_of_le_of_lt (mk_subtype_mono (diff_subset _ _)) _,
  { exact w.atom_map_dom_small, },
  { exact lt_of_le_of_lt mk_image_le w.atom_map_dom_small, },
end

lemma complete_atom_perm_domain_small : small w.complete_atom_perm.domain :=
small.union (small.union w.atom_map_dom_small
  (lt_of_le_of_lt mk_image_le w.atom_map_dom_small)) w.sandbox_subset_small

/-- A near-litter approximation built from this weak near-litter approximation.
Its action on atoms matches that of the weak approximation, and its rough action on litters
matches the given litter permutation. -/
noncomputable def complete (A : extended_index β) : near_litter_approx := {
  atom_perm := w.complete_atom_perm,
  litter_perm := w.flexible_litter_perm A,
  domain_small := λ L, small.mono (inter_subset_right _ _) w.complete_atom_perm_domain_small,
}

variable {litter_perm : local_perm litter}

lemma complete_atom_perm_apply_eq {a : atom} (ha : (w.atom_map a).dom) :
  w.complete_atom_perm a = (w.atom_map a).get ha :=
by rwa [complete_atom_perm, local_perm.complete_apply_eq, atom_map_or_else_of_dom]

lemma complete_smul_atom_eq {a : atom} (ha : (w.atom_map a).dom) :
  w.complete A • a = (w.atom_map a).get ha := w.complete_atom_perm_apply_eq ha

@[simp] lemma complete_smul_litter_eq (L : litter) :
  w.complete A • L = w.flexible_litter_perm A L := rfl

lemma smul_atom_eq
  {π : near_litter_perm} (hπ : (w.complete A).exactly_approximates π)
  {a : atom} (ha : (w.atom_map a).dom) :
  π • a = (w.atom_map a).get ha :=
by rw [← hπ.map_atom a (or.inl (or.inl ha)), w.complete_smul_atom_eq ha]

lemma smul_to_near_litter_eq_of_precise_at
  {π : near_litter_perm} (hπ : (w.complete A).exactly_approximates π)
  {L : litter} (hL : (w.litter_map L).dom) (hw : w.precise_at hL)
  (hπL : π • L = ((w.litter_map L).get hL).1) :
  π • L.to_near_litter = (w.litter_map L).get hL :=
begin
  refine set_like.coe_injective _,
  ext a : 1,
  simp only [mem_smul_set_iff_inv_smul_mem, near_litter_perm.coe_smul, litter.coe_to_near_litter,
    mem_litter_set, set_like.mem_coe],
  split,
  { intro ha,
    by_cases π.is_exception a,
    { suffices h' : π⁻¹ • a ∈ w.atom_map.dom,
      { rw w.atom_mem _ h' L hL at ha,
        have := hπ.map_atom _ (or.inl (or.inl h')),
        rw w.complete_smul_atom_eq h' at this,
        rw [this, smul_inv_smul] at ha,
        exact ha, },
      rw ← hπ.symm_map_atom a (hπ.exception_mem _ h) at ha ⊢,
      obtain ((hdom | hdom) | hdom) := (w.complete A).atom_perm.symm.map_domain
        (hπ.exception_mem _ h),
      { exact hdom, },
      { obtain ⟨c, hc₁, hc₂⟩ := hdom,
        rw w.atom_map_or_else_of_dom hc₁ at hc₂,
        have := hw.fwd c hc₁ (by rwa hc₂),
        rw hc₂ at this,
        exact this, },
      { cases w.sandbox_litter_not_banned _,
        rw ← eq_of_mem_litter_set_of_mem_litter_set ha
          (local_perm.sandbox_subset_subset _ _ hdom),
        exact banned_litter.litter_dom L hL, }, },
    { by_contradiction h',
      simp only [near_litter_perm.is_exception, mem_litter_set, not_or_distrib, not_not, ha] at h,
      obtain ⟨b, hb, rfl⟩ := hw.diff
        (or.inr ⟨by rw [← hπL, h.2, smul_inv_smul, mem_litter_set], h'⟩),
      refine h' ((w.atom_mem b hb L hL).mp _),
      have := hπ.map_atom b (or.inl (or.inl hb)),
      rw [w.complete_smul_atom_eq hb] at this,
      rw [this, inv_smul_smul] at ha,
      exact ha, }, },
  { intro ha,
    -- TODO: probably possible to clean up `by_cases` into a `suffices`
    by_cases π⁻¹ • a ∈ w.atom_map.dom,
    { rw w.atom_mem _ h L hL,
      have := hπ.map_atom _ (or.inl (or.inl h)),
      rw w.complete_smul_atom_eq h at this,
      rw [this, smul_inv_smul],
      exact ha, },
    have haL : a ∈ litter_set ((w.litter_map L).get hL).fst,
    { by_contradiction h',
      obtain ⟨b, hb, rfl⟩ := hw.diff (or.inl ⟨ha, h'⟩),
      have := hπ.map_atom b (or.inl (or.inl hb)),
      rw [w.complete_smul_atom_eq hb] at this,
      rw [this, inv_smul_smul] at h,
      exact h hb, },
    by_contradiction h',
    have hex : π.is_exception a,
    { refine or.inr (λ h'', h' (h''.trans _)),
      rw [inv_smul_eq_iff, hπL],
      exact haL, },
    obtain ((hdom | ⟨b, hb₁, hb₂⟩) | hdom) := hπ.exception_mem a hex,
    { obtain ⟨b, hb₁, hb₂⟩ := hw.back ⟨hdom, ha⟩,
      have := hπ.map_atom b (or.inl (or.inl hb₁)),
      rw [w.complete_smul_atom_eq hb₁] at this,
      rw [this, smul_eq_iff_eq_inv_smul] at hb₂,
      rw hb₂ at hb₁,
      exact h hb₁, },
    { rw w.atom_map_or_else_of_dom hb₁ at hb₂,
      have := hπ.map_atom b (or.inl (or.inl hb₁)),
      rw [w.complete_smul_atom_eq hb₁, hb₂, ← inv_smul_eq_iff] at this,
      rw this at h,
      exact h hb₁, },
    { refine w.sandbox_litter_not_banned _,
      rw eq_of_mem_litter_set_of_mem_litter_set (local_perm.sandbox_subset_subset _ _ hdom) haL,
      exact banned_litter.litter_map L hL, }, },
end

lemma smul_near_litter_eq_of_precise_at
  {π : near_litter_perm} (hπ : (w.complete A).exactly_approximates π)
  {N : near_litter} (hN : (w.litter_map N.1).dom) (hw : w.precise_at hN)
  (hπL : π • N.1 = ((w.litter_map N.1).get hN).1) :
  ((π • N : near_litter) : set atom) = (w.litter_map N.1).get hN ∆ (π • (litter_set N.1 ∆ N)) :=
begin
  refine (near_litter_perm.smul_near_litter_eq_smul_symm_diff_smul _ _).trans _,
  rw ← w.smul_to_near_litter_eq_of_precise_at hπ hN hw hπL,
  refl,
end

end weak_near_litter_approx

namespace weak_struct_approx

section

def precise {β : type_index} (w : weak_struct_approx β) : Prop := ∀ B, (w B).precise

variables {α : Λ} [position_data.{}] [phase_2_assumptions α] {β : Iio α} (w : weak_struct_approx β)

noncomputable def complete : struct_approx β :=
λ B, (w B).complete B

lemma smul_atom_eq
  {π : struct_perm β} (hπ : w.complete.exactly_approximates π)
  {a : atom} {B : extended_index β} (ha : ((w B).atom_map a).dom) :
  struct_perm.derivative B π • a = ((w B).atom_map a).get ha :=
begin
  have := (w B).smul_atom_eq (hπ B) ha,
  rw struct_perm.of_bot_smul at this,
  exact this,
end

lemma smul_to_near_litter_eq_of_precise (hw : w.precise)
  {π : struct_perm β} (hπ : w.complete.exactly_approximates π)
  {L : litter} {B : extended_index β} (hL : ((w B).litter_map L).dom)
  (hπL : struct_perm.derivative B π • L = (((w B).litter_map L).get hL).1) :
  struct_perm.derivative B π • L.to_near_litter = ((w B).litter_map L).get hL :=
begin
  have := (w B).smul_to_near_litter_eq_of_precise_at (hπ B) hL (hw B hL) _,
  { rw struct_perm.of_bot_smul at this,
    exact this, },
  { rw struct_perm.of_bot_smul,
    exact hπL, },
end

lemma smul_near_litter_eq_of_precise (hw : w.precise)
  {π : struct_perm β} (hπ : w.complete.exactly_approximates π)
  {N : near_litter} {B : extended_index β} (hN : ((w B).litter_map N.1).dom)
  (hπL : struct_perm.derivative B π • N.1 = (((w B).litter_map N.1).get hN).1) :
  ((struct_perm.derivative B π • N : near_litter) : set atom) =
  ((w B).litter_map N.1).get hN ∆ (struct_perm.derivative B π • (litter_set N.1 ∆ N)) :=
begin
  have := (w B).smul_near_litter_eq_of_precise_at (hπ B) hN (hw B hN) _,
  { rw struct_perm.of_bot_smul at this,
    exact this, },
  { rw struct_perm.of_bot_smul,
    exact hπL, },
end

end

variables {α : Λ} [position_data.{}] [phase_2_assumptions α] {β : Iio α}

/-- A weak structural approximation *supports* a tangle if it defines an image for everything
in the reduction of its designated support. -/
structure supports (w : weak_struct_approx β) (t : tangle β) : Prop :=
(atom_mem : ∀ a B, (inl a, B) ∈ reduction α (designated_support t : set (support_condition β)) →
  ((w B).atom_map a).dom)
(litter_mem : ∀ (L : litter) B,
  (inr L.to_near_litter, B) ∈ reduction α (designated_support t : set (support_condition β)) →
  ((w B).litter_map L).dom)

/-- Two weak structural approximations are *compatible* for a tangle if they both support the
tangle and agree on the reduction of its designated support. -/
structure compatible (w v : weak_struct_approx β) (t : tangle β) : Prop :=
(w_supports : w.supports t)
(v_supports : v.supports t)
(atom_map : ∀ a B ha, ((w B).atom_map a).get (w_supports.atom_mem a B ha) =
  ((v B).atom_map a).get (v_supports.atom_mem a B ha))
(litter_map : ∀ L B hL, ((w B).litter_map L).get (w_supports.litter_mem L B hL) =
  ((v B).litter_map L).get (v_supports.litter_mem L B hL))

/-- The action of a weak structural approximation on support conditions. -/
noncomputable def support_condition_map_or_else (w : weak_struct_approx β) :
  support_condition β → support_condition β
| (inl a, B) := (inl ((w B).atom_map_or_else a), B)
| (inr N, B) := (inr ((w B).near_litter_map_or_else N), B)

def coherent_coe (w : weak_struct_approx β) (t : tangle β) : Prop :=
∀ {π : allowable β} (hπ : w.complete.exactly_approximates π.to_struct_perm)
  (γ : Iic α) (δ ε : Iio α) (hδ : (δ : Λ) < γ) (hε : (ε : Λ) < γ) (hδε : δ ≠ ε)
  (C : path (β : type_index) γ) (t' : tangle δ) (hL)
  (hc₁ : ∃ (d : support_condition β), d ∈ (designated_support t).carrier ∧
    relation.refl_trans_gen (constrains α β)
    (inr (f_map (coe_ne_coe.mpr (coe_ne' hδε)) t').to_near_litter,
      (C.cons (coe_lt hε)).cons (bot_lt_coe _)) d)
  (hc₂ : ∀ (c : support_condition δ), c ∈ (designated_support t').carrier →
    π • (show support_condition β, from (c.fst, (C.cons (coe_lt hδ)).comp c.snd)) =
      w.support_condition_map_or_else (c.fst, (C.cons (coe_lt hδ)).comp c.snd)),
  f_map (subtype.coe_injective.ne (Iio.coe_injective.ne hδε))
      (show tangle δ, from
        (show allowable δ, from allowable_derivative (γ : Iic_index α) δ (coe_lt_coe.mpr hδ)
          (allowable.derivative
            (show path ((β : Iic_index α) : type_index) (γ : Iic_index α), from C) π)) • t') =
    (((w ((C.cons (coe_lt hε)).cons (bot_lt_coe _))).litter_map
      (f_map (subtype.coe_injective.ne (Iio.coe_injective.ne hδε)) t')).get hL).fst

def coherent_bot (w : weak_struct_approx β) : Prop :=
∀ {π : allowable β} (hπ : w.complete.exactly_approximates π.to_struct_perm)
  (γ : Iic α) (ε : Iio α) (hε : (ε : Λ) < γ)
  (C : path (β : type_index) γ) (a : tangle ⊥) (hL)
  (hc : struct_perm.derivative (C.cons (bot_lt_coe _)) π.to_struct_perm • a =
    (w (C.cons (bot_lt_coe _))).atom_map_or_else a),
  f_map (show ((⊥ : Iio_index α) : type_index) ≠ (ε : Iio_index α),
    from subtype.coe_injective.ne Iio_index.bot_ne_coe)
      ((struct_perm.derivative (C.cons (bot_lt_coe _))) π.to_struct_perm • a) =
    (((w ((C.cons (coe_lt hε)).cons (bot_lt_coe _))).litter_map
      (f_map (show (⊥ : type_index) ≠ (ε : Λ), from bot_ne_coe) a)).get hL).fst

@[mk_iff] structure coherent (w : weak_struct_approx β) (t : tangle β) : Prop :=
(coe : w.coherent_coe t)
(bot : w.coherent_bot)

lemma smul_litter_eq_of_supports (w : weak_struct_approx β)
  {π : allowable β} (hπ : w.complete.exactly_approximates π.to_struct_perm)
  (t : tangle β) (hwc : w.coherent t) (hws : w.supports t)
  (d : support_condition β) (hd : d ∈ designated_support t)
  (B : extended_index β) (L : litter)
  (ih : ∀ (e : support_condition β),
    relation.trans_gen (constrains α β) e (inr L.to_near_litter, B) →
    π • e = w.support_condition_map_or_else e)
  (hc : relation.refl_trans_gen (constrains α β) (inr L.to_near_litter, B) d) :
  struct_perm.derivative B π.to_struct_perm • L =
  (((w B).litter_map L).get
    (hws.litter_mem L B ⟨⟨d, hd, refl_trans_gen_near_litter hc⟩, reduced.mk_litter _ _⟩)).fst :=
begin
  by_cases hflex : inflexible α L B,
  rw inflexible_iff at hflex,
  obtain (⟨γ, δ, ε, hδ, hε, hδε, C, t', rfl, rfl⟩ | ⟨γ, ε, hε, C, a, rfl, rfl⟩) := hflex,
  { have hc₂ := λ c hc, ih _ (relation.trans_gen.single $ constrains.f_map hδ hε hδε C t' c hc),
    have := smul_f_map (δ : Iio_index α) ε _ _ (Iio.coe_injective.ne hδε)
      (allowable.derivative
        (show path ((β : Iic_index α) : type_index) (γ : Iic_index α), from C) π) t',
    rw [← allowable.derivative_cons_apply, allowable.derivative_smul,
      ← struct_perm.derivative_bot_smul, ← struct_perm.derivative_cons] at this,
    exact this.trans (hwc.coe hπ γ δ ε hδ hε hδε C t' _ ⟨d, hd, hc⟩ hc₂), },
  { have hc : (_, _) = (_, _) := ih _ (relation.trans_gen.single $ constrains.f_map_bot hε C a),
    simp only [smul_inl, prod.mk.inj_iff, eq_self_iff_true, and_true] at hc,
    have := smul_f_map (⊥ : Iio_index α) ε _ _ _
      (allowable.derivative
        (show path ((β : Iic_index α) : type_index) (γ : Iic_index α), from C) π) a,
    rw [← allowable.derivative_cons_apply, allowable.derivative_smul,
      ← struct_perm.derivative_bot_smul, ← struct_perm.derivative_cons] at this,
    rw ← hwc.bot hπ γ ε hε C a _ hc,
    refine this.trans _,
    swap 3,
    refine congr_arg _ _,
    swap 3,
    { rw ← allowable.derivative_cons_apply,
      rw ← allowable.derivative_smul
        (show path ((β : Iic_index α) : type_index) ((⊥ : Iic_index α) : type_index),
          from C.cons (bot_lt_coe _)) π a,
      congr,
      sorry, },
    all_goals { sorry, }, },
  { have := hws.litter_mem L B ⟨⟨d, hd, refl_trans_gen_near_litter hc⟩, reduced.mk_litter _ _⟩,
    rw [← struct_perm.of_bot_smul, ← (hπ B).map_litter _ (or.inl (or.inl ⟨this, hflex⟩))],
    refine ((w B).complete_smul_litter_eq L).trans _,
    rw [(w B).flexible_litter_perm_apply_eq, (w B).rough_litter_map_or_else_of_dom],
    exact this,
    exact hflex, },
end

lemma smul_support_condition_eq (w : weak_struct_approx β) (hw : w.precise)
  {π : allowable β} (hπ : w.complete.exactly_approximates π.to_struct_perm)
  (t : tangle β) (hwc : w.coherent t) (hws : w.supports t)
  (c d : support_condition β)
  (hc : relation.refl_trans_gen (constrains α β) c d)
  (hd : d ∈ designated_support t) :
  π • c = w.support_condition_map_or_else c :=
begin
  revert d,
  refine (constrains_wf α β).trans_gen.induction c _,
  rintros c ih d hc hd,
  obtain ⟨a | N, B⟩ := c,
  { refine prod.ext _ rfl,
    change inl _ = inl _,
    refine congr_arg inl _,
    rw [w.smul_atom_eq hπ (hws.atom_mem a B ⟨⟨d, hd, hc⟩, reduced.mk_atom a B⟩),
      weak_near_litter_approx.atom_map_or_else_of_dom], },
  refine prod.ext _ rfl,
  change inr _ = inr _,
  refine congr_arg inr (set_like.coe_injective _),
  have ih' := λ e he, ih e (relation.trans_gen.single he) d
    (relation.refl_trans_gen.head he hc) hd,
  rw w.smul_near_litter_eq_of_precise hw hπ (hws.litter_mem N.1 B _) _,
  { simp only [weak_near_litter_approx.near_litter_map_or_else,
      near_litter.coe_mk, subtype.coe_mk],
    rw (w B).litter_map_or_else_of_dom (hws.litter_mem N.1 B _),
    congr' 1,
    ext a : 1,
    rw [mem_smul_set, mem_image],
    split,
    { rintro ⟨b, hb₁, hb₂⟩,
      have : (_, _) = (_, _) := ih' _ (constrains.symm_diff N _ hb₁ B),
      simp only [smul_inl, smul_inv_smul, prod.mk.inj_iff] at this,
      rw this.1 at hb₂,
      exact ⟨b, hb₁, hb₂⟩, },
    { rintro ⟨b, hb₁, hb₂⟩,
      have : (_, _) = (_, _) := ih' _ (constrains.symm_diff N _ hb₁ B),
      simp only [smul_inl, smul_inv_smul, prod.mk.inj_iff] at this,
      rw ← this.1 at hb₂,
      exact ⟨b, hb₁, hb₂⟩, },
    { exact ⟨⟨d, hd, refl_trans_gen_near_litter hc⟩, reduced.mk_litter _ _⟩, }, },
  refine w.smul_litter_eq_of_supports hπ t hwc hws d hd B N.1 _ (refl_trans_gen_near_litter hc),
  exact λ e he, ih e (trans_gen_near_litter he) d
    (relation.refl_trans_gen.trans he.to_refl (refl_trans_gen_near_litter hc)) hd,
end

lemma smul_eq_smul_tangle (w v : weak_struct_approx β)
  (hw : w.precise) (hv : v.precise)
  (t : tangle β) (h : compatible w v t)
  (hwc : w.coherent t) (hvc : v.coherent t)
  {πw πv : allowable β} (hπw : w.complete.exactly_approximates πw.to_struct_perm)
  (hπv : v.complete.exactly_approximates πv.to_struct_perm) :
  πw • t = πv • t :=
begin
  rw [smul_eq_iff_eq_inv_smul, smul_smul],
  symmetry,
  refine (designated_support t).supports _ _,
  intros c hc,
  rw [mul_smul, inv_smul_eq_iff],
  symmetry,
  rw smul_support_condition_eq w hw hπw t hwc h.w_supports c c relation.refl_trans_gen.refl hc,
  rw smul_support_condition_eq v hv hπv t hvc h.v_supports c c relation.refl_trans_gen.refl hc,
  obtain ⟨a | N, B⟩ := c,
  { simp only [support_condition_map_or_else, prod.mk.inj_iff, eq_self_iff_true, and_true],
    rw [(w B).atom_map_or_else_of_dom, (v B).atom_map_or_else_of_dom],
    refine h.atom_map a B _,
    exact ⟨⟨_, hc, relation.refl_trans_gen.refl⟩, reduced.mk_atom _ _⟩, },
  { simp only [support_condition_map_or_else, prod.mk.inj_iff, eq_self_iff_true, and_true,
      weak_near_litter_approx.near_litter_map_or_else],
    refine set_like.coe_injective _,
    simp only [near_litter.coe_mk, subtype.coe_mk],
    congr' 1,
    { rw [(w B).litter_map_or_else_of_dom, (v B).litter_map_or_else_of_dom, h.litter_map N.1 B _],
      exact ⟨⟨_, hc, refl_trans_gen_near_litter relation.refl_trans_gen.refl⟩,
        reduced.mk_litter _ _⟩, },
    { ext a : 1,
      rw [mem_image, mem_image],
      split;
      rintro ⟨b, hb₁, hb₂⟩;
      refine ⟨b, hb₁, _⟩;
      rw [← hb₂, (w B).atom_map_or_else_of_dom, (v B).atom_map_or_else_of_dom],
      { refine (h.atom_map b B _).symm,
        exact ⟨⟨_, hc, relation.refl_trans_gen.single (constrains.symm_diff N b hb₁ B)⟩,
          reduced.mk_atom _ _⟩, },
      { refine h.atom_map b B _,
        exact ⟨⟨_, hc, relation.refl_trans_gen.single (constrains.symm_diff N b hb₁ B)⟩,
          reduced.mk_atom _ _⟩, }, }, },
end

end weak_struct_approx

end con_nf


/-
Copyright (c) 2021 Gabriel Ebner. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Gabriel Ebner, Scott Morrison
-/
import Std.Tactic.TryThis
import Mathlib.Lean.Expr.Basic
import Mathlib.Tactic.Cache
import Mathlib.Tactic.Core
import Mathlib.Tactic.SolveByElim

/-!
# Library search

This file defines a tactic `library_search`
and a term elaborator `library_search%`
that tries to find a lemma
solving the current goal
(subgoals are solved using `solveByElim`).

```
example : x < x + 1 := library_search%
example : Nat := by library_search
```
-/

namespace Lean.Meta.DiscrTree

/--
Inserts a new key into a discrimination tree,
but only if it is not of the form `#[*]` or `#[=, *, *, *]`.
-/
def insertIfSpecific {α : Type} {s : Bool} [BEq α] (d : DiscrTree α s)
    (keys : Array (DiscrTree.Key s)) (v : α) : DiscrTree α s :=
  if keys == #[Key.star] || keys == #[Key.const `Eq 3, Key.star, Key.star, Key.star] then
    d
  else
    d.insertCore keys v

end Lean.Meta.DiscrTree

namespace Mathlib.Tactic.LibrarySearch

open Lean Meta Std.Tactic.TryThis

initialize registerTraceClass `Tactic.librarySearch

-- from Lean.Server.Completion
private def isBlackListed (declName : Name) : MetaM Bool := do
  if declName == ``sorryAx then return true
  if declName matches .str _ "inj" then return true
  if declName matches .str _ "noConfusionType" then return true
  let env ← getEnv
  pure $ declName.isInternal'
   || isAuxRecursor env declName
   || isNoConfusion env declName
  <||> isRec declName <||> isMatcher declName

/--
A "modifier" for a declaration.
* `none` indicates the original declaration,
* `symm` indicates that (possibly after binders) the declaration is an `=`,
  and we want to consider the symmetric version,
* `mp` indicates that (possibly after binders) the declaration is an `iff`,
  and we want to consider the forward direction,
* `mpr` similarly, but for the backward direction.
-/
inductive DeclMod
| none | symm | mp | mpr
deriving DecidableEq

initialize librarySearchLemmas : DeclCache (DiscrTree (Name × DeclMod) true) ←
  DeclCache.mk "librarySearch: init cache" {} fun name constInfo lemmas => do
    if constInfo.isUnsafe then return lemmas
    if ← isBlackListed name then return lemmas
    withNewMCtxDepth do withReducible do
      let (_, _, type) ← forallMetaTelescopeReducing constInfo.type
      let keys ← DiscrTree.mkPath type
      let lemmas := lemmas.insertIfSpecific keys (name, .none)
      match type.getAppFnArgs with
      | (``Eq, #[_, lhs, rhs]) => do
        let keys_symm ← DiscrTree.mkPath (← mkEq rhs lhs)
        pure (lemmas.insertIfSpecific keys_symm (name, .symm))
      | (``Iff, #[lhs, rhs]) => do
        let keys_mp ← DiscrTree.mkPath rhs
        let keys_mpr ← DiscrTree.mkPath lhs
        pure <| (lemmas.insertIfSpecific keys_mp (name, .mp)).insertIfSpecific keys_mpr (name, .mpr)
      | _ => pure lemmas

/-- Shortcut for calling `solveByElim`. -/
def solveByElim (goals : List MVarId) (required : List Expr) (depth) := do
  -- There is only a marginal decrease in performance for using the `symm` and `exfalso`
  -- options for `solveByElim`.
  -- (measured via `lake build && time lake env lean test/librarySearch.lean`).
  let cfg : SolveByElim.Config := { maxDepth := depth, exfalso := true, symm := true }
  let cfg := if !required.isEmpty then cfg.requireUsingAll required else cfg
  _ ← SolveByElim.solveByElim.processSyntax cfg false false [] [] #[] goals

/--
Try to solve the goal either by:
* calling `solveByElim`
* or applying a library lemma then calling `solveByElim` on the resulting goals.

If it successfully closes the goal, returns `none`.
Otherwise, it returns `some a`, where `a : Array (MetavarContext × List MVarId)`,
with an entry for each library lemma which was successfully applied,
containing the metavariable context after the application, and a list of the subsidiary goals.

(Always succeeds, and the metavariable context stored in the monad is reverted,
unless the goal was completely solved.)

(Note that if `solveByElim` solves some but not all subsidiary goals,
this is not currently tracked.)
-/
def librarySearch (goal : MVarId) (lemmas : DiscrTree (Name × DeclMod) s) (required : List Expr)
    (solveByElimDepth := 6) : MetaM <| Option (Array <| MetavarContext × List MVarId) := do
  profileitM Exception "librarySearch" (← getOptions) do
  let ty ← goal.getType
  withTraceNode `Tactic.librarySearch (return m!"{exceptOptionEmoji ·} {ty}") do

  let mut suggestions := #[]

  let state0 ← get

  try
    solveByElim [goal] required solveByElimDepth
    return none
  catch _ =>
    set state0

  for (lem, mod) in ← lemmas.getMatch ty do
    trace[Tactic.librarySearch] "{lem}"
    let result ← withTraceNode `Tactic.librarySearch (return m!"{exceptOptionEmoji ·} trying {lem}")
      try
        let lem ← mkConstWithFreshMVarLevels lem
        let lem ← match mod with
        | .none => pure lem
        | .symm => mapForallTelescope (fun e => mkAppM ``Eq.symm #[e]) lem
        | .mp => mapForallTelescope (fun e => mkAppM ``Iff.mp #[e]) lem
        | .mpr => mapForallTelescope (fun e => mkAppM ``Iff.mpr #[e]) lem
        let newGoals ← goal.apply lem
        (try
          for newGoal in newGoals do
            trace[Tactic.librarySearch] "proving {← addMessageContextFull (mkMVar newGoal)}"
          solveByElim newGoals required solveByElimDepth
          pure $ some $ Sum.inr ()
        catch _ =>
          let res := some $ Sum.inl (← getMCtx, newGoals)
          set state0
          return res)
    catch _ =>
      set state0
      pure none
    match result with
    | none => pure ()
    | some (Sum.inr ()) => return none
    | some (Sum.inl suggestion) => suggestions := suggestions.push suggestion

  pure $ some suggestions

def lines (ls : List MessageData) :=
  MessageData.joinSep ls (MessageData.ofFormat Format.line)

open Lean.Parser.Tactic

-- TODO: implement the additional options for `library_search` from Lean 3,
-- in particular including additional lemmas
-- with `library_search [X, Y, Z]` or `library_search with attr`.
syntax (name := librarySearch') "library_search" (config)? (simpArgs)?
  (" using " (colGt term),+)? : tactic
syntax (name := librarySearch!) "library_search!" (config)? (simpArgs)?
  (" using " (colGt term),+)? : tactic

-- For now we only implement the basic functionality.
-- The full syntax is recognized, but will produce a "Tactic has not been implemented" error.

open Elab.Tactic Elab Tactic in
elab_rules : tactic | `(tactic| library_search%$tk $[using $[$required:term],*]?) => do
  let mvar ← getMainGoal
  let (_, goal) ← (← getMainGoal).intros
  goal.withContext do
    let required := (← (required.getD #[]).mapM getFVarId).toList.map .fvar
    if let some suggestions ← librarySearch goal (← librarySearchLemmas.get) required then
      for suggestion in suggestions do
        withMCtx suggestion.1 do
          addExactSuggestion tk (← instantiateMVars (mkMVar mvar)).headBeta
      admitGoal goal
    else
      addExactSuggestion tk (← instantiateMVars (mkMVar mvar)).headBeta

open Elab Term in
elab tk:"library_search%" : term <= expectedType => do
  let goal ← mkFreshExprMVar expectedType
  let (_, introdGoal) ← goal.mvarId!.intros
  introdGoal.withContext do
    if let some suggestions ← librarySearch introdGoal (← librarySearchLemmas.get) [] then
      for suggestion in suggestions do
        withMCtx suggestion.1 do
          addTermSuggestion tk (← instantiateMVars goal).headBeta
      mkSorry expectedType (synthetic := true)
    else
      addTermSuggestion tk (← instantiateMVars goal).headBeta
      instantiateMVars goal