import Lean
open Lean

-- @bollu: I removed the substring in `list` as we don't really use it for anything and
-- it just complicates things
inductive Sexp
| atom: String → Sexp
| list: List Sexp → Sexp
deriving BEq, Inhabited, Repr

def Sexp.fromString : String → Sexp
| s => Sexp.atom s

instance : Coe String Sexp where
  coe s := Sexp.fromString s

def Sexp.fromList : List Sexp → Sexp
| xs => Sexp.list xs

instance : Coe (List Sexp) Sexp where
  coe := Sexp.fromList

partial def Sexp.toString : Sexp → String
| .atom s => s
| .list xs => "(" ++ " ".intercalate (xs.map Sexp.toString) ++ ")"

instance : ToString Sexp := ⟨Sexp.toString⟩

def Sexp.toList? : Sexp → Option (List Sexp)
| .atom _ => .none
| .list xs => .some xs

def Sexp.toAtom! : Sexp → String
| .atom s => s
| .list xs => panic! s!"expected atom, found list at {List.toString xs}"

inductive SexpTok
| sexp: Sexp →  SexpTok
| opening: String.Pos → SexpTok
deriving BEq, Inhabited, Repr

structure SexpState where
  it: String.Iterator
  stack: List SexpTok := []
  sexps: List Sexp := []
  depth : Nat := 0
deriving BEq, Repr

def SexpState.fromString (s: String): SexpState :=
  { it := s.iter : SexpState }

instance : Inhabited SexpState where
  default := SexpState.fromString ""

inductive SexpError
| unmatchedOpenParen (ix: String.Iterator): SexpError
| unmatchedCloseParen (ix: String.Iterator): SexpError
| notSingleSexp (s: String) (xs: List Sexp): SexpError
deriving BEq, Repr

instance : ToString SexpError where toString := λ err => match err with
  | .unmatchedOpenParen ix => s!"Unmatched open parenthesis at {ix}"
  | .unmatchedCloseParen ix => s!"Unmatched close parenthesis at {ix}"
  | .notSingleSexp s xs => s!"not a single sexp '{s}', parsed as: '{xs}'"

abbrev SexpM := EStateM SexpError SexpState

def SexpM.peek: SexpM (Option (Char × String.Pos)) := do
  let state ← get
  return if state.it.atEnd then .none else .some (state.it.curr, state.it.i)

def SexpM.curPos: SexpM String.Pos := do
  let state ← get
  return state.it.i

-- Stop is a good name, because it indicates that it's exclusive
-- (AG: I don't read it as being exclusive from the name 'stop')
def SexpM.mkSubstring (l: String.Pos) (r: String.Pos): SexpM Substring := do
  let state ← get
  return { str := state.it.s, startPos := l, stopPos := r}

def SexpM.advance: SexpM Unit := do
  modify (fun state => { state with it := state.it.next })

def SexpM.pushTok (tok: SexpTok): SexpM Unit := do
  modify (fun state => { state with stack := tok :: state.stack })

def SexpM.pushSexp (sexp: Sexp): SexpM Unit := do
  let state ← get
  if state.stack.length == 0
  then set { state with stack := [], sexps := sexp :: state.sexps }
  else set { state with stack := (SexpTok.sexp sexp) :: state.stack }

def SexpM.incrementDepth: SexpM Unit :=
  modify (fun state => { state with depth := state.depth + 1 })

def SexpM.decrementDepth: SexpM Unit :=
  modify (fun state => { state with depth := state.depth - 1 })

instance [Inhabited α] : Inhabited (SexpM α) where
  default := do return default

def SexpM.pop: SexpM SexpTok := do
  let state ← get
  match state.stack with
  | [] => panic! "empty stack"
  | x::xs => do
      set { state with stack := xs }
      return x

-- abbrev SexpTokStack := List SexpTok

-- Remove elements from the stack of tokens `List SexpToken` till we find a `SexpToken.opening`.
-- When we do, return (1) the position of the open paren, (2) the list of SexpTokens left on the stack, and (3) the list of Sexps
-- Until then, accumulate the `SexpToken.sexp`s into `sexps`.
def stackPopTillOpen (stk: List SexpTok) (sexps: List Sexp := []): Option (String.Pos × (List SexpTok) × (List Sexp)) :=
  match stk with
  | [] => .none
  | SexpTok.opening openPos :: rest => (.some (openPos, rest, sexps))
  | SexpTok.sexp s :: rest => stackPopTillOpen rest (s :: sexps)

-- collapse the current stack till the last ( into a single Sexp.list
def SexpM.matchClosingParen: SexpM Unit := do
  let state ← get
  match stackPopTillOpen state.stack with
  | (.some (_, stk, sexps)) =>
    let sexp := Sexp.list sexps
    modify (fun state => { state with stack := stk })
    SexpM.pushSexp sexp
  | (.none) => throw (SexpError.unmatchedCloseParen state.it)

partial def SexpM.takeString (startPos: String.Pos): SexpM Substring := do
  match (← SexpM.peek) with
  | .none => SexpM.mkSubstring startPos (← SexpM.curPos)
  | .some (' ', _) => SexpM.mkSubstring startPos (← SexpM.curPos)
  | .some ('(', _) => SexpM.mkSubstring startPos (← SexpM.curPos)
  | .some (')', _) => SexpM.mkSubstring startPos (← SexpM.curPos)
  | .some _ => do
     SexpM.advance
     SexpM.takeString startPos

partial def SexpM.parse: SexpM Unit := do
  match (← SexpM.peek) with
  | .some  ('(', i) => do
     SexpM.advance
     SexpM.pushTok (SexpTok.opening i)
     SexpM.incrementDepth
     SexpM.parse
  | .some (')', _) => do
     SexpM.advance
     SexpM.matchClosingParen
     SexpM.parse
     -- return cur ++ rest
  | .some (' ', _) => do
      SexpM.advance
      SexpM.parse
  | .some (_, i) => do
      let s ← SexpM.takeString i
      SexpM.pushSexp ((Sexp.atom s.toString))
      SexpM.parse
  | .none => do
      let state ← get
      match stackPopTillOpen state.stack with
      | (.some (openPos, _, _)) =>
          throw <| SexpError.unmatchedOpenParen   ({ s := state.it.s, i := openPos : String.Iterator })
      | (.none) => return ()

-- | Parse a list of (possibly empty) sexps.
def parseSexpList (s: String):  Except SexpError (List Sexp) :=
  let initState := SexpState.fromString s
  match EStateM.run SexpM.parse initState with
  | .ok () state => .ok state.sexps.reverse
  | .error e _ => .error e

-- | Parse a single s-expression, and error if found no sexp or multiple sexps
def parseSingleSexp (s: String): Except SexpError Sexp := do
  match (← parseSexpList s) with
  | [x] => .ok x
  | xs => .error (.notSingleSexp s xs)

-- To simplify Sexps, we want to replace some subterms in an Sexp:
-- Have to mark this as partial since the termination checker doesn't like
-- these higher-order functions like map. See: https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/.E2.9C.94.20using.20higher-order.20functions.20on.20inductive.20types.3A.20termin.2E.2E.2E
partial def replaceTerm (toReplace : Sexp) (replaceWith : Sexp) (atSexp : Sexp) : Sexp :=
  if toReplace == atSexp then replaceWith
  else match atSexp with
  | .atom _ => atSexp
  | .list sexps => sexps.map $ replaceTerm toReplace replaceWith

-- The idea of this simplification is to do substitutions that replace a term,
-- such that the replaced term never appears as a proper subterm elsewhere in
-- the request, i.e. neither as the goal/starting point nor in any of the
-- rewrites. Ideally maximal subterms with this property

-- For this, we start by finding wether a subexpression is contained in an Sexp
-- Yes, this is not optimal because of the checks, feel free to rewrite.

-- Need this order of arguments (awkward for recursive call) for `.` notation
partial def Sexp.containsSubexpr (mainExpr : Sexp) (subExpr : Sexp) : Bool :=
  if subExpr == mainExpr then true
  else match mainExpr with
    | .atom _ => false
    | .list sexps => sexps.any (containsSubexpr · subExpr)

partial def Sexp.vars : Sexp → List String
  | .atom s => [s]
  | .list sexps => List.join $ sexps.map vars

partial def Sexp.fvarsConstsVars : Sexp → List Sexp × List Sexp × List String
  | .atom s => ([],[],[s])
  | c@(.list ("const"::_)) => ([],[c],[])
  | fvar@(.list ("fvar"::_)) => ([fvar],[],[])
  | .list sexps => sexps.foldl (init := ([],[],[]))
    λ (consts,fvars,vars) sexp =>
      let res := sexp.fvarsConstsVars
      (consts.append res.1, fvars.append res.2.1, vars.append res.2.2)

-- We could maybe replace this with `Std.HashMap`, but this should do it for now.
abbrev VariableMapping := List (String × Sexp)

-- Some generic helper functions
def _root_.List.revLookup? {α β : Type 0} [BEq β] : List (α × β) → β → Option α
  | [], _ => none
  | (a,b)::rest, b' => if b == b' then some a else rest.revLookup? b'

def _root_.List.unique {α : Type 0} [BEq α] : List α → List α
  | [] => []
  | a :: as => if as.contains a then as.unique else (a :: as.unique)

def _root_.List.unzip3 {α β γ : Type 0} : List (α × β × γ) → List α × List β × List γ
  | abc => let (a,bc) := abc.unzip
    (a,bc.unzip)

def freshVar (vars : List String) : String := Id.run do
  let mut idx := vars.length
  let mut fresh := s!"v{idx}"
  while vars.contains fresh do
    idx := idx + 1
    fresh := s!"v{idx}"
  return fresh

def Sexp.head : Sexp → String
  | .atom s => s
  | .list [] => ""
  | .list (hd::_) => head hd

def Sexp.uncurry : Sexp → Sexp
  | a@(.atom _) => a
  | .list ["ap", (.list ["ap", (.list ["ap", (.list ["ap", (.atom fname), args4]), args3]), args2]), args1] => .list [(.atom s!"ap4-{fname}"), args4.uncurry, args3.uncurry, args2.uncurry, args1.uncurry]
  | .list ["ap", (.list ["ap", (.list ["ap", (.atom fname), args3]), args2]), args1] => .list [s!"ap3-{fname}", args3.uncurry, args2.uncurry, args1.uncurry]
  | .list ["ap", (.list ["ap", (.atom fname), args2]), args1] => .list [s!"ap2-{fname}", args2.uncurry, args1.uncurry]
  | .list ["ap", (.atom fname), args] => .list [s!"ap-{fname}", args.uncurry]
  | l@(.list _) => l

#eval Sexp.uncurry (parseSingleSexp "(ap (ap mul (ap inv y)) (ap inv x))" |>.toOption |>.get!) |>.toString

-- partial because of map..
partial def Sexp.curry : Sexp → Sexp
  | a@(.atom _) => a
  | .list [(.atom (.mk ('a'::'p'::'4'::'-'::fname))), args4, args3, args2, args1] => .list ["ap", (.list ["ap", (.list ["ap", (.list ["ap", (.atom (.mk fname)), args4.curry]), args3.curry]), args2.curry]), args1.curry]
  | .list [(.atom (.mk ('a'::'p'::'3'::'-'::fname))), args3, args2, args1] => .list ["ap", (.list ["ap", (.list ["ap", (.atom (.mk fname)), args3.curry]), args2.curry]), args1.curry]
  | .list [(.atom (.mk ('a'::'p'::'2'::'-'::fname))), args2, args1] => .list ["ap", (.list ["ap", (.atom (.mk fname)), args2.curry]), args1.curry]
  | .list [(.atom (.mk ('a'::'p'::'-'::fname))), args] => .list ["ap", (.atom (.mk fname)), args.curry]
  | l@(.list _) => l

def simplifySexps : List Sexp → List Sexp × VariableMapping
  | sexps =>
    let fvarsConstsVars := sexps.foldl (init := ([],[],[]))
      λ (fvs,cs,vs) exp =>
        let res := exp.fvarsConstsVars
        ((fvs ++ res.1).unique, (cs ++ res.2.1).unique, (vs ++ res.2.2).unique)
    let fvars := fvarsConstsVars.1
    let consts := fvarsConstsVars.2.1
    Id.run do
      let mut allVars := fvarsConstsVars.2.2
      let mut mapping := []
      let mut exps := sexps
      for fvar in fvars do
        let vname := freshVar allVars
        mapping := (vname,fvar)::mapping
        allVars := vname::allVars
        exps := exps.map λ exp => replaceTerm fvar (Sexp.atom vname) exp
      for c in consts do
        let vname := freshVar allVars
        mapping := (vname,c)::mapping
        allVars := vname::allVars
        exps := exps.map λ exp => replaceTerm c (Sexp.atom vname) exp
      return (exps, mapping)

def Sexp.unsimplify : Sexp →  VariableMapping → Sexp
  | sexp, mapping => sexp.vars.foldl (init := sexp)
    λ e var => match mapping.lookup var with
      | none => e
      | some subexp => replaceTerm (Sexp.atom var) subexp e

def unsimplifySExps : List Sexp →  VariableMapping → List Sexp
  | sexps, mapping => sexps.map
    λ exp => exp.vars.foldl (init := exp)
      λ e var => match mapping.lookup var with
        | none => e
        | some subexp => replaceTerm (Sexp.atom var) subexp e

def ab := parseSingleSexp "(a b)" |>.toOption |>.get!
def aab := parseSingleSexp "(a (a b))" |>.toOption |>.get!
def c := Sexp.atom "c"
def a := Sexp.atom "a"
#eval ab.toString
#eval replaceTerm ab c ab |>.toString
#eval replaceTerm ab c aab |>.toString
#eval replaceTerm a c aab |>.toString
#eval replaceTerm ab c aab |> replaceTerm c ab |>.toString

def realexample := parseSexpList "(ap (fvar (num (str anonymous _uniq) 547)) (ap (fvar (num (str anonymous _uniq) 547)) (fvar (num (str anonymous _uniq) 550)))) (fvar (num (str anonymous _uniq) 550)) (fvar (num (str anonymous _uniq) 549)) (ap (ap (fvar (num (str anonymous _uniq) 548)) (fvar (num (str anonymous _uniq) 550))) (ap (fvar (num (str anonymous _uniq) 547)) (fvar (num (str anonymous _uniq) 550)))) ?_uniq.562 (ap (ap (fvar (num (str anonymous _uniq) 548)) ?_uniq.562) (fvar (num (str anonymous _uniq) 549))) (ap (ap (fvar (num (str anonymous _uniq) 548)) (ap (fvar (num (str anonymous _uniq) 547)) ?_uniq.561)) ?_uniq.561) (fvar (num (str anonymous _uniq) 549)) (ap (ap (fvar (num (str anonymous _uniq) 548)) ?_uniq.558) (ap (ap (fvar (num (str anonymous _uniq) 548)) ?_uniq.559) ?_uniq.560)) (ap (ap (fvar (num (str anonymous _uniq) 548)) (ap (ap (fvar (num (str anonymous _uniq) 548)) ?_uniq.558) ?_uniq.559)) ?_uniq.560)" |>.toOption.get!
def realexampleSimplified := simplifySexps realexample
#eval realexampleSimplified.1.toString
#eval realexampleSimplified.1.map (λ e => e.uncurry ) |>.map toString
#eval (realexampleSimplified.1.map (λ e => e.uncurry.curry) |>.zip realexampleSimplified.1).map λ (a,b) => a == b

#eval realexampleSimplified.2.map λ (s,sexp) => (s,sexp.toString)
#eval realexampleSimplified.1.map (Sexp.unsimplify · realexampleSimplified.2) |>.zip realexample |>.map λ (a,b) => a == b
def exp1 := parseSexpList "(a (a b)) (b (a b))" |>.toOption |>.get!
def exp2 := parseSexpList "(c (a b)) ((a b) c)" |>.toOption |>.get!
def exp3 := parseSexpList "(d ((a b) c)) (((a b) c) d)" |>.toOption |>.get!
def simp1 := simplifySexps exp1
def simp2 := simplifySexps exp2
def simp3 := simplifySexps exp3

#eval simp1.1 |>.map Sexp.toString
#eval simp2.1 |>.map Sexp.toString
#eval simp3.1 |>.map Sexp.toString

#eval unsimplifySExps simp1.1 simp1.2 == exp1
#eval unsimplifySExps simp2.1 simp2.2 == exp2
#eval unsimplifySExps simp3.1 simp3.2 == exp3

#eval aab.containsSubexpr a
#eval aab.containsSubexpr ab
#eval aab.containsSubexpr c

#eval parseSexpList ""
#eval parseSexpList "(a, b)"
#eval parseSexpList "(a, "
#eval parseSexpList "a)"
#eval parseSexpList "a b c"
#eval parseSexpList "(a b) (c d)"
#eval parseSingleSexp "(a b)"
#eval parseSingleSexp "(a (b c) d)"


import Structure.Generic.Axioms.Universes
import Structure.Generic.Axioms.AbstractFunctors

set_option autoBoundImplicitLocal false
--set_option pp.universes true

namespace HasLinearFunOp

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

  -- The "swap" functor swaps the arguments of a nested functor. Its plain version `swapFun` actually
  -- just fixes the second argument.

  def swapIsFun {α β γ : U} (F : α ⟶' β ⟶ γ) (b : β) : HasExternalFunctors.IsFun (λ a : α => F a b) :=
  h.compIsFun F (appFun' b γ)

  def swapFun' {α β γ : U} (F : α ⟶' β ⟶ γ) (b : β) : α ⟶' γ := BundledFunctor.mkFun (swapIsFun F b)
  def swapFun  {α β γ : U} (F : α ⟶  β ⟶ γ) (b : β) : α ⟶  γ := HasInternalFunctors.fromBundled (swapFun' (HasInternalFunctors.toBundled F) b)

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

  theorem swapFunFun.def {α β γ : U} (F : α ⟶' β ⟶ γ) (b : β) :
    HasInternalFunctors.fromBundled (swapFun' F b) = HasInternalFunctors.fromBundled (HasInternalFunctors.toBundled (appFun b γ) ⊙' F) :=
  HasInternalFunctors.toFromBundled (appFun' b γ) ▸ rfl

  def swapFunIsFun {α β γ : U} (F : α ⟶' β ⟶ γ) : HasExternalFunctors.IsFun (λ b : β => HasInternalFunctors.fromBundled (swapFun' F b)) :=
  funext (swapFunFun.def F) ▸ h.compIsFun (appFunFun' β γ) (compFunFun' F γ)

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

  @[simp] theorem swapFunFun.eff {α β γ : U} (F : α ⟶ β ⟶ γ) (b : β) : (swapFunFun F) b = swapFun F b :=
  by apply HasInternalFunctors.fromBundled.eff

  @[simp] theorem swapFunFun.effEff {α β γ : U} (F : α ⟶ β ⟶ γ) (b : β) (a : α) : ((swapFunFun F) b) a = F a b :=
  by simp

  theorem swapFunFunFun.def (α β γ : U) (F : α ⟶ β ⟶ γ) :
    HasInternalFunctors.mkFun (swapFunIsFun (HasInternalFunctors.toBundled F)) = HasInternalFunctors.fromBundled (HasInternalFunctors.toBundled (compFunFun F γ) ⊙' appFunFun' β γ) :=
  HasInternalFunctors.toFromBundled (compFunFun' (HasInternalFunctors.toBundled F) γ) ▸ elimRec

  def swapFunFunIsFun (α β γ : U) : HasExternalFunctors.IsFun (λ F : α ⟶ β ⟶ γ => swapFunFun F) :=
  funext (swapFunFunFun.def α β γ) ▸ h.compIsFun (compFunFunFun' α (β ⟶ γ) γ) (compFunFun' (appFunFun' β γ) (α ⟶ γ))

  def swapFunFunFun' (α β γ : U) : (α ⟶ β ⟶ γ) ⟶' (β ⟶ α ⟶ γ) := BundledFunctor.mkFun (swapFunFunIsFun α β γ)
  def swapFunFunFun  (α β γ : U) : (α ⟶ β ⟶ γ) ⟶  (β ⟶ α ⟶ γ) := HasInternalFunctors.fromBundled (swapFunFunFun' α β γ)

  @[simp] theorem swapFunFunFun.eff (α β γ : U) (F : α ⟶ β ⟶ γ) : (swapFunFunFun α β γ) F = swapFunFun F :=
  by apply HasInternalFunctors.fromBundled.eff

  @[simp] theorem swapFunFunFun.effEff (α β γ : U) (F : α ⟶ β ⟶ γ) (b : β) : ((swapFunFunFun α β γ) F) b = swapFun F b :=
  by simp

  @[simp] theorem swapFunFunFun.effEffEff (α β γ : U) (F : α ⟶ β ⟶ γ) (b : β) (a : α) : (((swapFunFunFun α β γ) F) b) a = F a b :=
  by simp

  -- In particular, reverse composition is also functorial.

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

  @[simp] theorem revCompFun.eff {α β γ : U} (F : α ⟶ β) (G : β ⟶ γ) (a : α) : (G ⊙ F) a = G (F a) :=
  compFun.eff F G a

  theorem revCompFunFun.def (α : U) {β γ : U} (G : β ⟶' γ) (F : α ⟶ β) :
    HasInternalFunctors.fromBundled (G ⊙' HasInternalFunctors.toBundled F) = (compFunFun F γ) (HasInternalFunctors.fromBundled G) :=
  Eq.subst (motive := λ H => HasInternalFunctors.fromBundled (H ⊙' HasInternalFunctors.toBundled F) = (compFunFun F γ) (HasInternalFunctors.fromBundled G))
           (HasInternalFunctors.toFromBundled G)
           (Eq.symm (compFunFun.eff F γ (HasInternalFunctors.fromBundled G)))

  def revCompFunIsFun (α : U) {β γ : U} (G : β ⟶' γ) :
    HasExternalFunctors.IsFun (λ F : α ⟶ β => HasInternalFunctors.fromBundled (G ⊙' HasInternalFunctors.toBundled F)) :=
  funext (revCompFunFun.def α G) ▸ swapIsFun (compFunFunFun' α β γ) (HasInternalFunctors.fromBundled G)

  def revCompFunFun' (α : U) {β γ : U} (G : β ⟶' γ) : (α ⟶ β) ⟶' (α ⟶ γ) := BundledFunctor.mkFun (revCompFunIsFun α G)
  def revCompFunFun  (α : U) {β γ : U} (G : β ⟶  γ) : (α ⟶ β) ⟶  (α ⟶ γ) := HasInternalFunctors.fromBundled (revCompFunFun' α (HasInternalFunctors.toBundled G))

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

  @[simp] theorem revCompFunFun.effEff (α : U) {β γ : U} (G : β ⟶ γ) (F : α ⟶ β) (a : α) : ((revCompFunFun α G) F) a = G (F a) :=
  by simp

  theorem revCompFunFunFun.def (α β γ : U) (G : β ⟶ γ) :
    HasInternalFunctors.mkFun (revCompFunIsFun α (HasInternalFunctors.toBundled G)) = HasInternalFunctors.fromBundled (swapFun' (compFunFunFun' α β γ) G) :=
  congrArg (λ H => HasInternalFunctors.fromBundled (swapFun' (compFunFunFun' α β γ) H)) (HasInternalFunctors.fromToBundled G) ▸ elimRec

  def revCompFunFunIsFun (α β γ : U) : HasExternalFunctors.IsFun (λ G : β ⟶ γ => revCompFunFun α G) :=
  funext (revCompFunFunFun.def α β γ) ▸ swapFunIsFun (compFunFunFun' α β γ)

  def revCompFunFunFun' (α β γ : U) : (β ⟶ γ) ⟶' (α ⟶ β) ⟶ (α ⟶ γ) := BundledFunctor.mkFun (revCompFunFunIsFun α β γ)
  def revCompFunFunFun  (α β γ : U) : (β ⟶ γ) ⟶  (α ⟶ β) ⟶ (α ⟶ γ) := HasInternalFunctors.fromBundled (revCompFunFunFun' α β γ)

  @[simp] theorem revCompFunFunFun.eff (α β γ : U) (G : β ⟶ γ) : (revCompFunFunFun α β γ) G = revCompFunFun α G :=
  by apply HasInternalFunctors.fromBundled.eff

  @[simp] theorem revCompFunFunFun.effEff (α β γ : U) (G : β ⟶ γ) (F : α ⟶ β) : ((revCompFunFunFun α β γ) G) F = G ⊙ F :=
  by simp

  @[simp] theorem revCompFunFunFun.effEffEff (α β γ : U) (G : β ⟶ γ) (F : α ⟶ β) (a : α) : (((revCompFunFunFun α β γ) G) F) a = G (F a) :=
  by simp

  -- Composition of a function with two arguments.

  def compFun₂ {α β γ δ : U} (F : α ⟶ β ⟶ γ) (G : γ ⟶ δ) : α ⟶ β ⟶ δ := swapFunFun (revCompFunFun α G ⊙ swapFunFun F)

end HasLinearFunOp

namespace HasFullFunOp

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

  -- The S combinator (see https://en.wikipedia.org/wiki/SKI_combinator_calculus), which in our case says
  -- that if we can functorially construct a functor `H : β ⟶ γ` and an argument `b : β`, then the
  -- construction of `H b` is also functorial.

  theorem substFun.def {α β γ : U} (F : α ⟶' β ⟶ γ) (G : α ⟶' β) (a : α) :
    F a (G a) = HasInternalFunctors.fromBundled (HasLinearFunOp.swapFun' F (G a)) a :=
  Eq.symm (HasInternalFunctors.fromBundled.eff (HasLinearFunOp.swapFun' F (G a)) a)

  def substIsFun {α β γ : U} (F : α ⟶' β ⟶ γ) (G : α ⟶' β) : HasExternalFunctors.IsFun (λ a : α => F a (G a)) :=
  funext (substFun.def F G) ▸ h.dupIsFun (HasLinearFunOp.swapFunFun' F ⊙' G)

  def substFun' {α β γ : U} (F : α ⟶' β ⟶ γ) (G : α ⟶' β) : α ⟶' γ := BundledFunctor.mkFun (substIsFun F G)
  def substFun  {α β γ : U} (F : α ⟶  β ⟶ γ) (G : α ⟶  β) : α ⟶  γ := HasInternalFunctors.fromBundled (substFun' (HasInternalFunctors.toBundled F) (HasInternalFunctors.toBundled G))

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

  theorem substFunFun.def {α β γ : U} (F : α ⟶' β ⟶ γ) (G : α ⟶ β) :
    HasInternalFunctors.mkFun (substIsFun F (HasInternalFunctors.toBundled G)) =
    HasInternalFunctors.fromBundled (HasNonLinearFunOp.dupFun' (HasInternalFunctors.toBundled (HasInternalFunctors.fromBundled (HasLinearFunOp.swapFunFun' F ⊙' HasInternalFunctors.toBundled G)))) :=
  HasInternalFunctors.toFromBundled _ ▸ elimRec

  def substFunIsFun {α β γ : U} (F : α ⟶' β ⟶ γ) :
    HasExternalFunctors.IsFun (λ G : α ⟶ β => HasInternalFunctors.fromBundled (substFun' F (HasInternalFunctors.toBundled G))) :=
  funext (substFunFun.def F) ▸ h.compIsFun (HasLinearFunOp.revCompFunFun' α (HasLinearFunOp.swapFunFun' F)) (HasNonLinearFunOp.dupFunFun' α γ)

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

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

  @[simp] theorem substFunFun.effEff {α β γ : U} (F : α ⟶ β ⟶ γ) (G : α ⟶ β) (a : α) : ((substFunFun F) G) a = F a (G a) :=
  by simp

  theorem substFunFunFun.def (α β γ : U) (F : α ⟶ β ⟶ γ) :
    HasInternalFunctors.mkFun (substFunIsFun (HasInternalFunctors.toBundled F)) = HasInternalFunctors.fromBundled (HasNonLinearFunOp.dupFunFun' α γ ⊙' HasInternalFunctors.toBundled (HasLinearFunOp.revCompFunFun α (HasLinearFunOp.swapFunFun F))) :=
  HasInternalFunctors.toFromBundled _ ▸ HasInternalFunctors.toFromBundled _ ▸ elimRec

  def substFunFunIsFun (α β γ : U) : HasExternalFunctors.IsFun (λ F : α ⟶ β ⟶ γ => substFunFun F) :=
  funext (substFunFunFun.def α β γ) ▸ h.compIsFun (HasLinearFunOp.revCompFunFunFun' α β (α ⟶ γ) ⊙' HasLinearFunOp.swapFunFunFun' α β γ)
                                                  (HasLinearFunOp.revCompFunFun' (α ⟶ β) (HasNonLinearFunOp.dupFunFun' α γ))

  def substFunFunFun' (α β γ : U) : (α ⟶ β ⟶ γ) ⟶' (α ⟶ β) ⟶ (α ⟶ γ) := BundledFunctor.mkFun (substFunFunIsFun α β γ)
  def substFunFunFun  (α β γ : U) : (α ⟶ β ⟶ γ) ⟶  (α ⟶ β) ⟶ (α ⟶ γ) := HasInternalFunctors.fromBundled (substFunFunFun' α β γ)

  @[simp] theorem substFunFunFun.eff (α β γ : U) (F : α ⟶ β ⟶ γ) : (substFunFunFun α β γ) F = substFunFun F :=
  by apply HasInternalFunctors.fromBundled.eff

  @[simp] theorem substFunFunFun.effEff (α β γ : U) (F : α ⟶ β ⟶ γ) (G : α ⟶ β) : ((substFunFunFun α β γ) F) G = substFun F G :=
  by simp

  @[simp] theorem substFunFunFun.effEffEff (α β γ : U) (F : α ⟶ β ⟶ γ) (G : α ⟶ β) (a : α) : (((substFunFunFun α β γ) F) G) a = F a (G a) :=
  by simp

end HasFullFunOp

-- Using the functoriality axioms and the constructions above, we can algorithmically prove
-- functoriality of lambda terms. The algorithm to prove `HasExternalFunctors.IsFun (λ a : α => t)`
-- is as follows:
--
--  Case                           | Proof
-- --------------------------------+--------------------------------------------------------------
--  `t` does not contain `a`       | `constIsFun α t`
--  `t` is `a`                     | `idIsFun α`
--  `t` is `G b` with `G : β ⟶ γ`: |
--    `a` appears only in `b`      | Prove that `λ a => b` is functorial, yielding a functor
--                                 | `F : α ⟶ β`. Then the proof is `compIsFun F G`.
--      `b` is `a`                 | Optimization: `HasInternalFunctors.isFun G`
--    `a` appears only in `G`      | Prove that `λ a => G` is functorial, yielding a functor
--                                 | `F : α ⟶ β ⟶ γ`. Then the proof is `swapIsFun F b`.
--      `G` is `a`                 | Optimization: `appIsFun b γ`
--    `a` appears in both          | Prove that `λ a => G` is functorial, yielding a functor
--                                 | `F₁ : α ⟶ β ⟶ γ`. Prove that `λ a => b` is functorial,
--                                 | yielding a functor `F₂ : α ⟶ β`. Then the proof is
--                                 | `substIsFun F₁ F₂`.
--  `t` is `mkFun (λ b : β => c)`  | Prove that `λ a => c` is functorial when regarding `b` as
--                                 | a constant, yielding a functor `F : α ⟶ γ` for every `b`.
--                                 | Prove that  `λ b => F` is functorial, yielding a functor
--                                 | `G : β ⟶ α ⟶ γ`. Then the proof is `swapFunIsFun G`.
--
-- (This list does not contain all possible optimizations.)


import .nonneg_rat
import .mask

namespace simplify

structure O := (n : nat)

instance dec_eq_o (o o' : O) [nat_dec : ∀n n' : nat, decidable (n = n')] : decidable (o = o') := begin
  induction o, induction o',
  rw O.mk.inj_eq,
  exact nat_dec o o',
end

inductive X
| div : X
| null : X
| ass : X
| perm : X

inductive Ref
| o : O → Ref
| null : Ref

instance dec_eq_ref (r r' : Ref) [o_dec : ∀o o' : O, decidable (o = o')] [f_dec : decidable false] [t_dec : decidable true] : decidable (r = r') := begin
  have ineq : ∀o : O, (Ref.null = Ref.o o) = false, cc,
  have ineq' : ∀o : O, (Ref.o o = Ref.null) = false, cc,

  induction r,

  induction r',
  rw Ref.o.inj_eq,
  exact o_dec r r',
  rw ineq' r,
  exact f_dec,

  induction r',
  rw ineq r',
  exact f_dec,
  rw Ref.null.inj_eq,
  exact t_dec,
end

inductive V
| b : bool → V
| z : ℤ → V
| q : ℚ* → V
| o : Ref → V

structure var := (n : nat)
structure field := (n : nat)
structure func := (n : nat)
structure pred := (n : nat)

inductive T
| B : T
| Z : T
| O : T
| Q : T

def V.typ : V → T
| (V.b _) := T.B
| (V.z _) := T.Z
| (V.o _) := T.O
| (V.q _) := T.Q

def H := field × O → V
def S := var → V
def FieldPermMask := mask (field × O)
def PredPermMask := mask (pred × list V)

structure conf :=
(h : H)
(s : S)
(fieldPerm : FieldPermMask)
(predPerm : PredPermMask)

inductive E
| null : E
| const_b : bool → E
| const_n : ℤ → E
| x : var → E
| deref : E → field → E
| negate : E → E
| add : E → E → E
| div : E → E → E
| not : E → E
| eq : E → E → E
| apply : func → list E → E
| unfolding : pred → list E → E → E → E
| fieldPerm : E → field → E
| predPerm : pred → list E → E

inductive A
| e : E → A
| fieldAcc : E → field → E → A
| predAcc : pred → list E → E → A
| imp : E → A → A
| conj : A → A → A
| all : T → (E → A) → A
| ex : T → (E → A) → A
| forperm : field → E → A

inductive Eval
| type_err : Eval
| err : X → Eval
| ok : V → Eval

def Eval.is_x : Eval → Prop
| Eval.type_err := false
| (Eval.err _) := true
| (Eval.ok _) := false

def Eval.is_v : Eval → Prop
| Eval.type_err := false
| (Eval.err _) := false
| (Eval.ok _) := true

def v_of_is_v {e : Eval} (h : e.is_v) : V := begin
  cases e,
  case ok { exact e, },
  repeat { simp [Eval.is_v] at h, apply false.elim, exact h, },
end

def Eval.is_t : Eval → T → Prop
| Eval.type_err _ := false
| (Eval.err _) _ := false
| (Eval.ok v) t := v.typ = t

structure defs :=
(var : var → T)
(field : field → T)
(func : func → (list T) × T × (list V → Eval)) -- shamelessly shallowly embed lean functions, to not deal with termination.
(pred : pred → (list T))

def E.typ (d : defs) : E → T
| (E.null) := T.O
| (E.const_b b) := T.B
| (E.const_n z) := T.Z
| (E.x x) := d.var x
| (E.deref o f) := d.field f
| (E.negate z) := z.typ
| (E.add z z') := z.typ
| (E.div z z') := z.typ
| (E.not b) := T.B
| (E.eq v v') := T.B
| (E.apply f args) := (d.func f).snd.fst
| (E.unfolding p args q e) := e.typ
| (E.fieldPerm o f) := T.Q
| (E.predPerm p args) := T.Q

def all_typed (d : defs) : list E → list T → Prop
| [] [] := true
| [] (_ :: _) := false
| (_ :: _) [] := false
| (e :: es) (t :: ts) := e.typ d = t ∧ all_typed es ts

mutual def all_wt, wt (d : defs)
with all_wt : list E → Prop
| [] := true
| (e :: es) := (wt e ∧ all_wt es)

with wt : E → Prop
| E.null := true
| (E.const_b b) := true
| (E.const_n z) := true
| (E.x x) := true
| (E.deref o f) := o.typ d = T.O ∧ wt o
| (E.negate z) := z.typ d = T.Z ∧ wt z
| (E.add z z') := z.typ d = z'.typ d ∧ (z.typ d = T.Z ∨ z.typ d = T.Q) ∧ wt z ∧ wt z'
| (E.div z z') := z.typ d = z'.typ d ∧ (z.typ d = T.Z ∨ z.typ d = T.Q) ∧ wt z ∧ wt z'
| (E.not b) := b.typ d = T.B ∧ wt b
| (E.eq v v') := v.typ d = v'.typ d ∧ wt v ∧ wt v'
| (E.apply f args) := all_typed d args (d.func f).fst ∧ all_wt args
| (E.unfolding p args q e) := all_typed d args (d.pred p) ∧ all_wt args ∧ wt q ∧ wt e
| (E.fieldPerm o f) := o.typ d = T.O ∧ wt o
| (E.predPerm p args) := all_typed d args (d.pred p) ∧ all_wt args

def conf.wt (d : defs) (c : conf) : Prop :=
  (∀x, (c.s x).typ = d.var x)
  ∧ (∀f o, (c.h (f, o)).typ = d.field f)

def Eval.map_v (e : Eval) (f : V → Eval) : Eval := match e with
| Eval.type_err := Eval.type_err
| Eval.err x := Eval.err x
| Eval.ok v := f v
end

#print Eval.rec

lemma map_v_is_x_or_t
  : ∀ {t t' : T} {f : V → Eval} (e : Eval),
    e.is_x ∨ e.is_t t' →
    (∀ (v : V) (h : e = Eval.ok v), (f v).is_x ∨ (f v).is_t t) →
    (e.map_v f).is_x ∨ (e.map_v f).is_t t
  :=
begin
  intros t t' f e,
  cases e,

  case type_err {
    intros e_x_or_t f_all_v,
    simp [Eval.is_x, Eval.is_v] at e_x_or_t,
    apply false.elim, exact e_x_or_t,
  },

  case err {
    intros e_x_or_t v,
    simp [Eval.map_v, Eval.is_x],
  },

  case ok {
    intros e_x_or_t v,
    simp [Eval.map_v],
    exact v e rfl,
  },
end

lemma eval_ok_t {e : Eval} {t : T} {v : V} :
  (Eval.ok v).is_t t ↔ v.typ = t := by simp [Eval.is_t]

def Eval.cases_v (e : Eval) (f_b : bool → Eval) (f_z : ℤ → Eval) (f_q : ℚ* → Eval) (f_o : Ref → Eval) :=
  e.map_v $ λv, match v with
  | V.b b := f_b b
  | V.z z := f_z z
  | V.q q := f_q q
  | V.o o := f_o o
  end

def evals_map' : list V → (list V → Eval) → list Eval → Eval
| acc f [] := f acc
| acc f (e :: es) := e.map_v $ λv, evals_map' (acc ++ [v]) f es

def evals_map (es : list Eval) (f : list V → Eval) : Eval :=
  evals_map' [] f es

lemma evals_map_x_or_t :
  ∀ {t  : T} (f : list V → Eval) (es : list Eval),
    (∀e : Eval, ∃t : T, e ∈ es → e.is_x ∨ e.is_t t) →
    (∀ (vs : list V) (h : es = vs.map Eval.ok), (f vs).is_x ∨ (f vs).is_t t) →
    (evals_map es f).is_x ∨ (evals_map es f).is_t t :=
begin
  intros t f es,
  intros no_type_err,
  intros ok_f_no_type_err,

  induction es,

  case nil {
    simp only [evals_map, evals_map'],
    apply ok_f_no_type_err list.nil,
    simp only [list.map],
  },

  case cons : e es ih {
    simp only [evals_map, evals_map'],
    apply map_v_is_x_or_t,
    apply exists.elim,
    exact no_type_err e,
    intro t_before,
    simp,
    intro e_ok,
    /-exact e_ok,-/ sorry,

    sorry,
    sorry,
  },
end

mutual def eval, evals (d : defs)
with eval : conf → E → Eval
| c (E.null) := Eval.ok (V.o Ref.null)
| c (E.const_b b) := Eval.ok (V.b b)
| c (E.const_n z) := Eval.ok (V.z z)
| c (E.x x) := Eval.ok (c.s x)
| c (E.deref o f) := (eval c o).map_v $ λv,
  match v with
  | (V.o Ref.null) := Eval.err X.null
  | (V.o (Ref.o o)) := if c.fieldPerm (f, o) = 0 then Eval.err X.perm else Eval.ok $ c.h (f, o)
  | _ := Eval.type_err
  end
| c (E.negate z) := (eval c z).map_v $ λv,
  match v with
  | (V.z z) := Eval.ok $ V.z (-z)
  | _ := Eval.type_err
  end
| c (E.add z z') := (eval c z).map_v $ λz, (eval c z').map_v $ λz',
  match z, z' with
  | (V.z z), (V.z z') := Eval.ok $ V.z (z + z')
  | (V.q q), (V.q q') := Eval.ok $ V.q (q + q')
  | _, _ := Eval.type_err
  end
| c (E.div z z') := (eval c z).map_v $ λz, (eval c z').map_v $ λz',
  match z, z' with
  | (V.z z), (V.z z') := if z' = 0 then Eval.err X.div else Eval.ok $ V.z $ int.div /- or flooring, round to zero? -/ z z'
  | (V.q q), (V.q q') := if q' = 0 then Eval.err X.div else Eval.ok $ V.q (q / q')
  | _, _ := Eval.type_err
  end
| c (E.not b) := (eval c b).map_v $ λv,
  match v with
  | (V.b b) := Eval.ok $ V.b ¬b
  | _ := Eval.type_err
  end
| c (E.eq v v') := (eval c v).map_v $ λv, (eval c v').map_v $ λv',
  match v, v' with
  | (V.b b), (V.b b') := Eval.ok $ V.b $ if b = b' then tt else ff
  | (V.z z), (V.z z') := Eval.ok $ V.b $ if z = z' then tt else ff
  | (V.q q), (V.q q') := Eval.ok $ V.b $ if q = q' then tt else ff
  | (V.o o), (V.o o') := Eval.ok $ V.b $ if o = o' then tt else ff
  | _, _ := Eval.type_err
  end
| c (E.apply f args) := evals_map (evals c args) $ λargs, ((d.func f).snd.snd args)
| c (E.unfolding p args q e) := evals_map (evals c args) $
    λargs, if c.predPerm (p, args) < 1
      then Eval.err X.perm
      else (eval c e).map_v $ λbody, Eval.ok (V.b ff)
| c (E.fieldPerm o f) := (eval c o).map_v $ λo,
  match o with
  | (V.o Ref.null) := Eval.err X.null
  | (V.o (Ref.o o)) := Eval.ok $ V.q $ c.fieldPerm (f, o)
  | _ := Eval.type_err
  end
| c (E.predPerm p args) := evals_map (evals c args)$ λargs,
    Eval.ok $ V.q $ c.predPerm (p, args)

with evals : conf → list E → list Eval
| c [] := []
| c (e :: es) := (eval c e) :: (evals c es)

#print decidable.rec

lemma x {α : Type} (c : Prop) [decidable c] (a a' : α) : (ite c a a') = (ite (¬c) a' a) := begin
  exact (ite_not c a' a).symm
end

lemma ite_prop {α : Type} {prop : α → Prop} {c : Prop} [h : decidable c] {a a' : α} :
  (c → prop a) ∧ (¬c → prop a') ↔ prop (ite c a a') := begin
  apply iff.intro,
  intro props,
  rw ite,
  apply h.rec_on,

  intro not_c,
  simp [(is_false not_c).rec_on], exact props.elim_right not_c,
  intro c,
  simp [(is_true c).rec_on], exact props.elim_left c,

  rw ite,
  apply h.rec_on,
  intro not_c,
  simp [(is_false not_c).rec_on],
  intro prop_a', finish,
  intro c,
  simp [(is_true c).rec_on],
  intro prop_a, finish,
end

theorem wt_sufficient {c : conf} {d : defs} {e : E} (h : wt d e) (ch : c.wt d)
  : (eval d c e).is_x ∨ (eval d c e).is_t (e.typ d) :=
begin
  rw conf.wt at ch,
  have var_ok : ∀x, (c.s x).typ = d.var x, exact and.elim_left ch,
  have field_ok : ∀f o, (c.h (f,o)).typ = d.field f, exact and.elim_right ch,
  clear ch,

  induction e,
  repeat { simp only [eval, wt, Eval.is_t, E.typ, V.typ, Eval.is_x, Eval.is_t, false_or, rfl] at *, },

  -- case x : x { exact var_ok x, },

  -- case deref : o f ih {
  --   apply map_v_is_x_or_t,
  --   exact ih (and.elim_right h),
  --   intros v v_ok,
  --   rw v_ok at *,
  --   cases v,
  --   repeat { simp [eval, Eval.is_x, Eval.is_t, V.typ, h.elim_left] at *, exact ih h, },
  --   cases v,
  --   simp [eval],
  --   by_cases c.fieldPerm (f, v) = 0,
  --   simp [h, Eval.is_x, Eval.is_t],
  --   rw ←ite_not,
  --   simp at h,
  --   simp [h, Eval.is_x, Eval.is_t, field_ok f v],
  --   simp [eval, Eval.is_t, Eval.is_x],
  -- },

  -- case negate : z ih {
  --   apply map_v_is_x_or_t,
  --   exact ih h.elim_right,
  --   intros v v_ok,
  --   rw v_ok at *,
  --   cases v,
  --   repeat { simp [eval] at *, },

  --   repeat { simp [eval, Eval.is_x, Eval.is_t, h.elim_left, V.typ] at *, },
  --   repeat {  exact ih h, },
  -- },

  -- case add : z z' ih ih' {
  --   apply map_v_is_x_or_t,
  --   exact ih h.elim_right.elim_right.elim_left,
  --   intros v v_ok,
  --   rw v_ok at *,
  --   apply map_v_is_x_or_t,
  --   exact ih' h.elim_right.elim_right.elim_right,
  --   intros v' v_ok',
  --   rw v_ok' at *,
  --   by_cases is_t_z : (z.typ d) = T.Z,

  --   cases v,
  --   repeat { simp [is_t_z] at *, },
  --   repeat { simp [eval, Eval.is_x, Eval.is_t, V.typ] at ih, apply false.elim, exact ih h.elim_right.elim_left, },
  --   cases v',
  --   repeat { simp [eval, Eval.is_x, Eval.is_t, V.typ, ←h.elim_left] at ih', apply false.elim, exact ih' h.elim_right.elim_right, },

  --   simp [eval, Eval.is_t, Eval.is_x, V.typ],

  --   cases v,
  --   repeat { simp [eval, Eval.is_x, Eval.is_t, V.typ, h.elim_right.elim_left] at ih, apply false.elim, exact ih h.elim_right.elim_right.elim_left, },
  --   cases v',
  --   repeat { simp [eval, Eval.is_x, Eval.is_t, V.typ, h.elim_right.elim_left, ←h.elim_left] at ih', apply false.elim, exact ih' h.elim_right.elim_right.elim_right, },

  --   simp [eval, Eval.is_t, Eval.is_x, V.typ, h.elim_right.elim_left],
  -- },

  -- case div : z z' ih ih' {
  --   apply map_v_is_x_or_t,
  --   exact ih h.elim_right.elim_right.elim_left,
  --   intros v v_ok,
  --   rw v_ok at *,
  --   apply map_v_is_x_or_t,
  --   exact ih' h.elim_right.elim_right.elim_right,
  --   intros v' v_ok',
  --   rw v_ok' at *,
  --   by_cases is_t_z : (z.typ d) = T.Z,

  --   cases v,
  --   repeat { simp [is_t_z] at *, },
  --   repeat { simp [eval, Eval.is_x, Eval.is_t, V.typ] at ih, apply false.elim, exact ih h.elim_right.elim_left, },
  --   cases v',
  --   repeat { simp [eval, Eval.is_x, Eval.is_t, V.typ, ←h.elim_left] at ih', apply false.elim, exact ih' h.elim_right.elim_right, },

  --   rw eval,
  --   by_cases div_zero : v' = 0,
  --   simp [div_zero, Eval.is_x],
  --   rw ←ite_not,
  --   simp [div_zero, Eval.is_t, V.typ],

  --   cases v,
  --   repeat { simp [eval, Eval.is_x, Eval.is_t, V.typ, h.elim_right.elim_left] at ih, apply false.elim, exact ih h.elim_right.elim_right.elim_left, },
  --   cases v',
  --   repeat { simp [eval, Eval.is_x, Eval.is_t, V.typ, ←h.elim_left, h.elim_right.elim_left] at ih', apply false.elim, exact ih' h.elim_right.elim_right.elim_right, },

  --   rw eval,
  --   by_cases div_zero : v' = 0,
  --   simp [div_zero, Eval.is_x],
  --   rw ←ite_not,
  --   simp [div_zero, Eval.is_t, V.typ, h.elim_right.elim_left],
  -- },

  -- case not : b ih {
  --   apply map_v_is_x_or_t,
  --   exact ih h.elim_right,
  --   intros v v_ok,
  --   rw v_ok at *,
  --   cases v,
  --   repeat { simp only [eval, Eval.is_x, Eval.is_t, V.typ, false_or], },
  --   repeat { simp only [h.elim_left, Eval.is_x, Eval.is_t, V.typ, false_or] at ih, },
  --   repeat { exact ih h.elim_right },
  -- },

  -- case eq : l r ih ih' {
  --   have t_eq : l.typ d = r.typ d, from h.elim_left,
  --   have wt_l : wt d l, from h.elim_right.elim_left,
  --   have wt_r : wt d r, from h.elim_right.elim_right,
  --   clear h,

  --   apply map_v_is_x_or_t,
  --   exact ih wt_l,
  --   intros v v_ok,
  --   apply map_v_is_x_or_t,
  --   exact ih' wt_r,
  --   intros v' v_ok',
  --   rw v_ok at *, rw v_ok' at *,

  --   cases v,
  --   repeat {
  --     cases v',
  --     repeat {
  --       simp only [eval],
  --       simp only [Eval.is_t, V.typ],
  --       apply or.intro_right,
  --       refl,
  --     },
  --     repeat {
  --       simp only [eval], simp only [Eval.is_t, Eval.is_x, V.typ, or_false],
  --       simp only [Eval.is_x, Eval.is_t, V.typ, false_or] at ih,
  --       simp only [Eval.is_x, Eval.is_t, V.typ, false_or, ←t_eq, ←ih wt_l] at ih',
  --       exact ih' wt_r,
  --     },
  --   },
  -- },

  case apply : f args {
    apply map_v_is_x_or_t,
  },

  repeat { sorry, },
end

end simplify


-- import circuits.circuit_encoding
-- import polytime.data_structures.finset

-- namespace function
-- /- Note that Lean fails to infer something like
-- (infer_instance : ∀ x : α, has_uncurry (list (β x) → list (β x) → list (β x)) _ _)
-- unless all the underscores are made explicit for some reason.
-- (infer_instance : ∀ x : α, has_uncurry (list (β x) → list (β x) → list (β x)) (list (β x) × list (β x)) (list $ β x))

-- This is why we make the following instances
-- -/
-- universes u v
-- variables {α : Type} {β β₁ β₂ γ δ : α → Type}

-- class has_uncurry_dep {α : Type} (β : α → Type) (β₁ β₂ : out_param (α → Type)) :=
-- (uncurry : ∀ {x}, (β x) → ((β₁ x) → (β₂ x)))

-- notation (name := uncurry_dep) `↾`:max x:max := has_uncurry_dep.uncurry _ x

-- instance has_uncurry_dep_base : has_uncurry_dep (λ x, (β₁ x → β₂ x)) β₁ β₂ := ⟨λ x y, y⟩

-- instance has_uncurry_dep_induction [has_uncurry_dep γ β₁ β₂] : has_uncurry_dep (λ x, δ x → γ x) (λ x, δ x × β₁ x) β₂ :=
-- ⟨λ x f p, ↾(f p.1) p.2⟩

-- end function

-- open_locale complexity_class
-- open function tencodable

-- namespace complexity_class

-- variables {α : Type} {β : α → Type} {β₁ : α → Type} {β₂ : α → Type} {γ δ : Type}
--   [tencodable α] [∀ x, tencodable (β x)] [∀ x, tencodable (β₁ x)] [∀ x, tencodable (β₂ x)]
--   [tencodable γ] [tencodable δ] {C : complexity_class}

-- /-- Membership of a dependent function in a complexity class;
--   Note: we only ever need *one* dependent argument for almost everything we do

--   TODO: can we unify `mem_dep` ("base case") and `mem_dep₂`?
--   Can we and should we use has_uncurry_dep so that `mem_dep₂` is automatically
--   generalized for >2 arguments?

--   Should we generate composition lemmas for `mem_dep` and `mem_dep₂`? If so, what are they?  -/
-- def mem_dep (f : ∀ x, β x) (C : complexity_class) : Prop :=
-- ∃ (f' : tree unit → tree unit), f' ∈ₑ C ∧
--   ∀ x : α, f' (encode x) = encode (f x)

-- localized "infix ` ∈ₐ `:50 := complexity_class.mem_dep" in complexity_class

-- /-- `fintype` but "C"-constructible -/
-- def mem_types (β : α → Type) [∀ x, tencodable (β x)] [∀ x, decidable_eq (β x)]
--   [∀ x, fintype (β x)] (C : complexity_class) : Prop :=
-- (λ x : α, @finset.univ (β x) _) ∈ₐ C

-- def mem_dep₂ (f : ∀ x, β₁ x → β₂ x) (C : complexity_class) : Prop :=
-- (λ x : sigma β₁, f x.1 x.2) ∈ₐ C

-- -- "t" for "two" ??
-- localized "infix ` ∈ₜ `:50 := complexity_class.mem_dep₂" in complexity_class
-- open_locale tree

-- @[simp] lemma mem_dep_iff {f : α → γ} : f ∈ₐ C ↔ f ∈ₑ C :=
-- by { simp_rw [mem_dep, ← prop_iff_mem], refl, }

-- lemma mem_dep₂_iff {f : α → γ → δ} : f ∈ₜ C ↔ f ∈ₑ C :=
-- by { dunfold mem_dep₂, rw mem_dep_iff, split; { rintro ⟨f', pf, hf⟩, refine ⟨f', pf, _⟩, rintro ⟨a, b⟩, exact hf ⟨a, b⟩, }, }

-- @[complexity] lemma mem_dep_of_mem {f : α → γ} (h : f ∈ₑ C) : f ∈ₐ C := by rwa mem_dep_iff
-- @[complexity] lemma mem_dep₂_of_mem {f : α → γ → δ} (h : f ∈ₑ C) : f ∈ₜ C := by rwa mem_dep₂_iff

-- lemma mem_iff_comp_encode_dep {f : ∀ x, β x} :
--   f ∈ₐ C ↔ (λ x, encode (f x)) ∈ₑ C :=
-- by { rw ← mem_dep_iff, refl, }

-- lemma _root_.list.encode_map_encode (l : list α) :
--   encode (l.map encode) = encode l := by simp only [encode, list.map_id]

-- lemma mem_iff_comp_list_encode_dep {f : ∀ x, list (β x)} :
--   f ∈ₐ C ↔ (λ x, (f x).map encode) ∈ₑ C :=
-- by { rw [mem_iff_comp_encode, mem_iff_comp_encode_dep], simp only [list.encode_map_encode], }

-- end complexity_class

-- /-- A function which is encoded as a table -/
-- structure complexity_class.table_fun (α β : Type*) :=
-- (to_fun : α → β)

-- namespace complexity_class.table_fun
-- open_locale complexity_class
-- variables {α β γ : Type*}

-- localized "infixr ` [→] `:25 := complexity_class.table_fun" in complexity_class

-- instance : has_coe_to_fun (α [→] β) (λ _, α → β) := ⟨table_fun.to_fun⟩

-- @[ext]
-- protected lemma ext : ∀ (f g : α [→] β), ⇑f = (by exact ⇑g) → f = g
-- | ⟨f⟩ ⟨g⟩ rfl := rfl

-- @[simp] lemma to_fun_eq_coe (f : α [→] β) : f.to_fun = ⇑f := rfl

-- @[simps]
-- def equiv_fun : (α [→] β) ≃ (α → β) := ⟨λ f, ⇑f, λ f, ⟨f⟩, λ f, by ext; refl, λ f, rfl⟩

-- @[simps]
-- def sum (f : α [→] γ) (g : β [→] γ) : α ⊕ β [→] γ := ⟨sum.elim ⇑f ⇑g⟩

-- @[simps]
-- def map (f : α [→] β) (g : β → γ) : α [→] γ := ⟨λ x, g (f x)⟩

-- @[simps]
-- def comp (f : α [→] β) (g : γ [→] α) : γ [→] β := ⟨λ x, f (g x)⟩

-- def finmap_equiv_fun [fintype α] [decidable_eq α] :
--   {x : @finmap α (λ _, β) // ∀ k : α, k ∈ x} ≃ (α → β) :=
-- { to_fun := λ f x, @option.get _ ((↑f : finmap _).lookup x) (finmap.lookup_is_some.mpr $ f.prop x),
--   inv_fun := λ f, ⟨finmap.of_fun f, λ k, finmap.mem_iff.mpr ⟨_, finmap.of_fun_lookup k⟩⟩,
--   left_inv := λ f, by { ext : 1, apply finmap.ext_lookup, simp, },
--   right_inv := λ f, by { ext, simp, } }

-- variables [tencodable α] [fintype α] [decidable_eq α] [tencodable β] [tencodable γ]

-- instance : tencodable (α [→] β) :=
-- tencodable.of_equiv _ (equiv_fun.trans finmap_equiv_fun.symm)

-- lemma encode_table_fun (f : α [→] β) : encode f = encode (finmap_equiv_fun.symm ⇑f) :=
-- rfl

-- lemma table_fun_mk_of {γ : Type} {ψ₁ ψ₂ : γ → Type} [tencodable γ] [∀ x, tencodable (ψ₁ x)]
--   [∀ x, fintype (ψ₁ x)] [∀ x, decidable_eq (ψ₁ x)] [∀ x, tencodable (ψ₂ x)] {f : ∀ x, (ψ₁ x → ψ₂ x)}
--   (hψ : polytime.mem_types ψ₁) (hf : f ∈ₜ PTIME) : (λ x, table_fun.mk (f x) : ∀ x, ψ₁ x [→] ψ₂ x) ∈ₐ PTIME :=
-- sorry

-- end complexity_class.table_fun

-- namespace polytime
-- open_locale complexity_class
-- open_locale tree
-- open complexity_class

-- variables {α : Type} {β : α → Type} {β₁ : α → Type} {β₂ : α → Type} {γ δ : Type}
--   [tencodable α] [∀ x, tencodable (β x)] [∀ x, tencodable (β₁ x)] [∀ x, tencodable (β₂ x)]
--   [tencodable γ] [tencodable δ] {C : complexity_class}

-- @[complexity] lemma list_map_dep {l : ∀ x, list (β₁ x)} {f : ∀ x, β₁ x → β₂ x} (hl : l ∈ₐ PTIME)
--   (hf : f ∈ₜ PTIME) : (λ x : α, (l x).map (f x) : ∀ x, list (β₂ x)) ∈ₐ PTIME :=
-- begin
--   rcases hf with ⟨f', pf, hf⟩,
--   rw mem_iff_comp_list_encode_dep at ⊢ hl,
--   complexity using λ x, ((l x).map encode).map (λ y, f' (encode x △ y)),
--   simp at hf, dsimp [tencodable.encode_sigma] at hf, simp [function.comp, hf],
-- end

-- @[complexity] lemma list_append_dep {l₁ l₂ : ∀ x, list (β x)} (hl₁ : l₁ ∈ₐ PTIME) (hl₂ : l₂ ∈ₐ PTIME) :
--   (λ x, (l₁ x) ++ (l₂ x)) ∈ₐ PTIME :=
-- by { rw [mem_iff_comp_list_encode_dep] at *, simp only [list.map_append], complexity, }

-- end polytime


/-
Copyright (c) 2016 Gabriel Ebner. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Gabriel Ebner, Sebastian Ullrich

Classy functions for lifting monadic actions of different shapes.

This theory is roughly modeled after the Haskell 'layers' package https://hackage.haskell.org/package/layers-0.1.
Please see https://hackage.haskell.org/package/layers-0.1/docs/Documentation-Layers-Overview.html for an exhaustive discussion of the different approaches to lift functions.
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.function
import Mathlib.Lean3Lib.init.coe
import Mathlib.Lean3Lib.init.control.monad

universes u v w l u_1 u_2 u_3 u_4

namespace Mathlib

/-- A function for lifting a computation from an inner monad to an outer monad.
    Like [MonadTrans](https://hackage.haskell.org/package/transformers-0.5.5.0/docs/Control-Monad-Trans-Class.html),
    but `n` does not have to be a monad transformer.
    Alternatively, an implementation of [MonadLayer](https://hackage.haskell.org/package/layers-0.1/docs/Control-Monad-Layer.html#t:MonadLayer) without `layerInvmap` (so far). -/
class has_monad_lift (m : Type u → Type v) (n : Type u → Type w) where
  monad_lift : {α : Type u} → m α → n α

/-- The reflexive-transitive closure of `has_monad_lift`.
    `monad_lift` is used to transitively lift monadic computations such as `state_t.get` or `state_t.put s`.
    Corresponds to [MonadLift](https://hackage.haskell.org/package/layers-0.1/docs/Control-Monad-Layer.html#t:MonadLift). -/
class has_monad_lift_t (m : Type u → Type v) (n : Type u → Type w) where
  monad_lift : {α : Type u} → m α → n α

/-- A coercion that may reduce the need for explicit lifting.
    Because of [limitations of the current coercion resolution](https://github.com/leanprover/lean/issues/1402), this definition is not marked as a global instance and should be marked locally instead. -/
def has_monad_lift_to_has_coe {m : Type u_1 → Type u_2} {n : Type u_1 → Type u_3}
    [has_monad_lift_t m n] {α : Type u_1} : has_coe (m α) (n α) :=
  has_coe.mk monad_lift

protected instance has_monad_lift_t_trans (m : Type u_1 → Type u_2) (n : Type u_1 → Type u_3)
    (o : Type u_1 → Type u_4) [has_monad_lift_t m n] [has_monad_lift n o] : has_monad_lift_t m o :=
  has_monad_lift_t.mk fun (α : Type u_1) (ma : m α) => has_monad_lift.monad_lift (monad_lift ma)

protected instance has_monad_lift_t_refl (m : Type u_1 → Type u_2) : has_monad_lift_t m m :=
  has_monad_lift_t.mk fun (α : Type u_1) => id

@[simp] theorem monad_lift_refl {m : Type u → Type v} {α : Type u} : monad_lift = id := rfl

/-- A functor in the category of monads. Can be used to lift monad-transforming functions.
    Based on pipes' [MFunctor](https://hackage.haskell.org/package/pipes-2.4.0/docs/Control-MFunctor.html),
    but not restricted to monad transformers.
    Alternatively, an implementation of [MonadTransFunctor](http://duairc.netsoc.ie/layers-docs/Control-Monad-Layer.html#t:MonadTransFunctor). -/
class monad_functor (m : Type u → Type v) (m' : Type u → Type v) (n : Type u → Type w)
    (n' : Type u → Type w)
    where
  monad_map : {α : Type u} → ({α : Type u} → m α → m' α) → n α → n' α

/-- The reflexive-transitive closure of `monad_functor`.
    `monad_map` is used to transitively lift monad morphisms such as `state_t.zoom`.
    A generalization of [MonadLiftFunctor](http://duairc.netsoc.ie/layers-docs/Control-Monad-Layer.html#t:MonadLiftFunctor), which can only lift endomorphisms (i.e. m = m', n = n'). -/
class monad_functor_t (m : Type u → Type v) (m' : Type u → Type v) (n : Type u → Type w)
    (n' : Type u → Type w)
    where
  monad_map : {α : Type u} → ({α : Type u} → m α → m' α) → n α → n' α

protected instance monad_functor_t_trans (m : Type u_1 → Type u_2) (m' : Type u_1 → Type u_2)
    (n : Type u_1 → Type u_3) (n' : Type u_1 → Type u_3) (o : Type u_1 → Type u_4)
    (o' : Type u_1 → Type u_4) [monad_functor_t m m' n n'] [monad_functor n n' o o'] :
    monad_functor_t m m' o o' :=
  monad_functor_t.mk
    fun (α : Type u_1) (f : {α : Type u_1} → m α → m' α) =>
      monad_functor.monad_map fun (α : Type u_1) => monad_map f

protected instance monad_functor_t_refl (m : Type u_1 → Type u_2) (m' : Type u_1 → Type u_2) :
    monad_functor_t m m' m m' :=
  monad_functor_t.mk fun (α : Type u_1) (f : {α : Type u_1} → m α → m' α) => f

@[simp] theorem monad_map_refl {m : Type u → Type v} {m' : Type u → Type v}
    (f : {α : Type u} → m α → m' α) {α : Type u} : monad_map f = f :=
  rfl

/-- Run a monad stack to completion.
    `run` should be the composition of the transformers' individual `run` functions.
    This class mostly saves some typing when using highly nested monad stacks:
    ```
    @[reducible] def my_monad := reader_t my_cfg $ state_t my_state $ except_t my_err id
    -- def my_monad.run {α : Type} (x : my_monad α) (cfg : my_cfg) (st : my_state) := ((x.run cfg).run st).run
    def my_monad.run {α : Type} (x : my_monad α) := monad_run.run x
    ```
    -/
class monad_run (out : outParam (Type u → Type v)) (m : Type u → Type v) where
  run : {α : Type u} → m α → out α

end Mathlib


/-
Copyright (c) 2021 Mario Carneiro. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro
-/
import Lean
import Std
import Mathlib.Tactic.Cases

namespace Mathlib.Tactic
open Lean Parser.Tactic Elab Command Elab.Tactic Meta

syntax (name := «variables») "variables" (bracketedBinder)* : command

@[command_elab «variables»] def elabVariables : CommandElab
  | `(variables%$pos $binders*) => do
    logWarningAt pos "'variables' has been replaced by 'variable' in lean 4"
    elabVariable (← `(variable%$pos $binders*))
  | _ => throwUnsupportedSyntax

/-- `lemma` means the same as `theorem`. It is used to denote "less important" theorems -/
syntax (name := lemma)
  declModifiers group("lemma" declId declSig declVal Parser.Command.terminationSuffix) : command

/-- Implementation of the `lemma` command, by macro expansion to `theorem`. -/
@[macro «lemma»] def expandLemma : Macro := fun stx =>
  -- FIXME: this should be a macro match, but terminationSuffix is not easy to bind correctly.
  -- This implementation ensures that any future changes to `theorem` are reflected in `lemma`
  let stx := stx.modifyArg 1 fun stx =>
    let stx := stx.modifyArg 0 (mkAtomFrom · "theorem" (canonical := true))
    stx.setKind ``Parser.Command.theorem
  pure <| stx.setKind ``Parser.Command.declaration

/-- Given two arrays of `FVarId`s, one from an old local context and the other from a new local
context, pushes `FVarAliasInfo`s into the info tree for corresponding pairs of `FVarId`s.
Recall that variables linked this way should be considered to be semantically identical.

The effect of this is, for example, the unused variable linter will see that variables
from the first array are used if corresponding variables in the second array are used. -/
def pushFVarAliasInfo [Monad m] [MonadInfoTree m]
    (oldFVars newFVars : Array FVarId) (newLCtx : LocalContext) : m Unit := do
  for old in oldFVars, new in newFVars do
    if old != new then
      let decl := newLCtx.get! new
      pushInfoLeaf (.ofFVarAliasInfo { id := new, baseId := old, userName := decl.userName })

/-- Function to help do the revert/intro pattern, running some code inside a context
where certain variables have been reverted before re-introing them.
It will push `FVarId` alias information into info trees for you according to a simple protocol.

- `fvarIds` is an array of `fvarIds` to revert. These are passed to
  `Lean.MVarId.revert` with `preserveOrder := true`, hence the function
  raises an error if they cannot be reverted in the provided order.
- `k` is given the goal with all the variables reverted and
  the array of reverted `FVarId`s, with the requested `FVarId`s at the beginning.
  It must return a tuple of a value, an array describing which `FVarIds` to link,
  and a mutated `MVarId`.

The `a : Array (Option FVarId)` array returned by `k` is interpreted in the following way.
The function will intro `a.size` variables, and then for each non-`none` entry we
create an FVar alias between it and the corresponding `intro`ed variable.
For example, having `k` return `fvars.map .some` causes all reverted variables to be
`intro`ed and linked.

Returns the value returned by `k` along with the resulting goal.
 -/
def _root_.Lean.MVarId.withReverted (mvarId : MVarId) (fvarIds : Array FVarId)
    (k : MVarId → Array FVarId → MetaM (α × Array (Option FVarId) × MVarId))
    (clearAuxDeclsInsteadOfRevert := false) : MetaM (α × MVarId) := do
  let (xs, mvarId) ← mvarId.revert fvarIds true clearAuxDeclsInsteadOfRevert
  let (r, xs', mvarId) ← k mvarId xs
  let (ys, mvarId) ← mvarId.introNP xs'.size
  mvarId.withContext do
    for x? in xs', y in ys do
      if let some x := x? then
        pushInfoLeaf (.ofFVarAliasInfo { id := y, baseId := x, userName := ← y.getUserName })
  return (r, mvarId)

/--
Replace the type of the free variable `fvarId` with `typeNew`.

If `checkDefEq = true` then throws an error if `typeNew` is not definitionally
equal to the type of `fvarId`. Otherwise this function assumes `typeNew` and the type
of `fvarId` are definitionally equal.

This function is the same as `Lean.MVarId.changeLocalDecl` but makes sure to push substitution
information into the infotree.
-/
def _root_.Lean.MVarId.changeLocalDecl' (mvarId : MVarId) (fvarId : FVarId) (typeNew : Expr)
    (checkDefEq := true) : MetaM MVarId := do
  mvarId.checkNotAssigned `changeLocalDecl
  let (_, mvarId) ← mvarId.withReverted #[fvarId] fun mvarId fvars => mvarId.withContext do
    let check (typeOld : Expr) : MetaM Unit := do
      if checkDefEq then
        unless ← isDefEq typeNew typeOld do
          throwTacticEx `changeLocalDecl mvarId
            m!"given type{indentExpr typeNew}\nis not definitionally equal to{indentExpr typeOld}"
    let finalize (targetNew : Expr) := do
      return ((), fvars.map .some, ← mvarId.replaceTargetDefEq targetNew)
    match ← mvarId.getType with
    | .forallE n d b bi => do check d; finalize (.forallE n typeNew b bi)
    | .letE n t v b ndep  => do check t; finalize (.letE n typeNew v b ndep)
    | _ => throwTacticEx `changeLocalDecl mvarId "unexpected auxiliary target"
  return mvarId

/-- `change` can be used to replace the main goal or its local
variables with definitionally equal ones.

For example, if `n : ℕ` and the current goal is `⊢ n + 2 = 2`, then
```lean
change _ + 1 = _
```
changes the goal to `⊢ n + 1 + 1 = 2`. The tactic also applies to the local context.
If `h : n + 2 = 2` and `h' : n + 3 = 4` are in the local context, then
```lean
change _ + 1 = _ at h h'
```
changes their types to be `h : n + 1 + 1 = 2` and `h' : n + 2 + 1 = 4`.

Change is like `refine` in that every placeholder needs to be solved for by unification,
but you can use named placeholders and `?_` where you want `change` to create new goals.

The the tactic `show e` is interchangeable with `change e`, where the pattern `e` is applied to
the main goal. -/
elab_rules : tactic
  | `(tactic| change $newType:term $[$loc:location]?) => do
    withLocation (expandOptLocation (Lean.mkOptionalNode loc))
      (atLocal := fun h ↦ do
        let hTy ← h.getType
        -- This is a hack to get the new type to elaborate in the same sort of way that
        -- it would for a `show` expression for the goal.
        let mvar ← mkFreshExprMVar none
        let (_, mvars) ← elabTermWithHoles
                          (← `(term | show $newType from $(← Term.exprToSyntax mvar))) hTy `change
        liftMetaTactic fun mvarId ↦ do
          return (← mvarId.changeLocalDecl' h (← inferType mvar)) :: mvars)
      (atTarget := evalTactic <| ← `(tactic| show $newType))
      (failed := fun _ ↦ throwError "change tactic failed")

/--
`by_cases p` makes a case distinction on `p`,
resulting in two subgoals `h : p ⊢` and `h : ¬ p ⊢`.
-/
macro "by_cases " e:term : tactic =>
  `(tactic| by_cases $(mkIdent `h) : $e)

syntax "transitivity" (colGt term)? : tactic
set_option hygiene false in
macro_rules
  | `(tactic| transitivity) => `(tactic| apply Nat.le_trans)
  | `(tactic| transitivity $e) => `(tactic| apply Nat.le_trans (m := $e))
set_option hygiene false in
macro_rules
  | `(tactic| transitivity) => `(tactic| apply Nat.lt_trans)
  | `(tactic| transitivity $e) => `(tactic| apply Nat.lt_trans (m := $e))

/--
The tactic `introv` allows the user to automatically introduce the variables of a theorem and
explicitly name the non-dependent hypotheses.
Any dependent hypotheses are assigned their default names.

Examples:
```
example : ∀ a b : Nat, a = b → b = a := by
  introv h,
  exact h.symm
```
The state after `introv h` is
```
a b : ℕ,
h : a = b
⊢ b = a
```

```
example : ∀ a b : Nat, a = b → ∀ c, b = c → a = c := by
  introv h₁ h₂,
  exact h₁.trans h₂
```
The state after `introv h₁ h₂` is
```
a b : ℕ,
h₁ : a = b,
c : ℕ,
h₂ : b = c
⊢ a = c
```
-/
syntax (name := introv) "introv " (colGt binderIdent)* : tactic
@[tactic introv] partial def evalIntrov : Tactic := fun stx ↦ do
  match stx with
  | `(tactic| introv)                     => introsDep
  | `(tactic| introv $h:ident $hs:binderIdent*) =>
    evalTactic (← `(tactic| introv; intro $h:ident; introv $hs:binderIdent*))
  | `(tactic| introv _%$tk $hs:binderIdent*) =>
    evalTactic (← `(tactic| introv; intro _%$tk; introv $hs:binderIdent*))
  | _ => throwUnsupportedSyntax
where
  introsDep : TacticM Unit := do
    let t ← getMainTarget
    match t with
    | Expr.forallE _ _ e _ =>
      if e.hasLooseBVars then
        intro1PStep
        introsDep
    | _ => pure ()
  intro1PStep : TacticM Unit :=
    liftMetaTactic fun goal ↦ do
      let (_, goal) ← goal.intro1P
      pure [goal]

/-- Try calling `assumption` on all goals; succeeds if it closes at least one goal. -/
macro "assumption'" : tactic => `(tactic| any_goals assumption)

elab "match_target" t:term : tactic  => do
  withMainContext do
    let (val) ← elabTerm t (← inferType (← getMainTarget))
    if not (← isDefEq val (← getMainTarget)) then
      throwError "failed"

/-- This tactic clears all auxiliary declarations from the context. -/
elab (name := clearAuxDecl) "clear_aux_decl" : tactic => withMainContext do
  let mut g ← getMainGoal
  for ldec in ← getLCtx do
    if ldec.isAuxDecl then
      g ← g.tryClear ldec.fvarId
  replaceMainGoal [g]

/-- Clears the value of the local definition `fvarId`. Ensures that the resulting goal state
is still type correct. Throws an error if it is a local hypothesis without a value. -/
def _root_.Lean.MVarId.clearValue (mvarId : MVarId) (fvarId : FVarId) : MetaM MVarId := do
  mvarId.checkNotAssigned `clear_value
  let tag ← mvarId.getTag
  let (_, mvarId) ← mvarId.withReverted #[fvarId] fun mvarId' fvars => mvarId'.withContext do
    let tgt ← mvarId'.getType
    unless tgt.isLet do
      mvarId.withContext <|
        throwTacticEx `clear_value mvarId m!"{Expr.fvar fvarId} is not a local definition"
    let tgt' := Expr.forallE tgt.letName! tgt.letType! tgt.letBody! .default
    unless ← isTypeCorrect tgt' do
      mvarId.withContext <|
        throwTacticEx `clear_value mvarId
          m!"cannot clear {Expr.fvar fvarId}, the resulting context is not type correct"
    let mvarId'' ← mkFreshExprSyntheticOpaqueMVar tgt' tag
    mvarId'.assign <| .app mvarId'' tgt.letValue!
    return ((), fvars.map .some, mvarId''.mvarId!)
  return mvarId

/-- `clear_value n₁ n₂ ...` clears the bodies of the local definitions `n₁, n₂ ...`, changing them
into regular hypotheses. A hypothesis `n : α := t` is changed to `n : α`.

The order of `n₁ n₂ ...` does not matter, and values will be cleared in reverse order of
where they appear in the context. -/
elab (name := clearValue) "clear_value" hs:(colGt term:max)+ : tactic => do
  let fvarIds ← getFVarIds hs
  let fvarIds ← withMainContext <| sortFVarIds fvarIds
  for fvarId in fvarIds.reverse do
    withMainContext do
      let mvarId ← (← getMainGoal).clearValue fvarId
      replaceMainGoal [mvarId]


import Do.Return

/-! # Iteration -/

open Lean

/- Disable the automatic monadic lifting feature described in the paper.
   We want to make it clear that we do not depend on it. -/
set_option autoLift false

syntax "for" ident "in" term "do'" stmt:1 : stmt
syntax "break " : stmt
syntax "continue " : stmt

syntax "break" : expander
syntax "continue" : expander
syntax "lift" : expander

macro_rules
  | `(stmt| expand! $_   in break) => `(stmt| break)                                              -- subsumes (S7, R7, B2, L1)
  | `(stmt| expand! $_   in continue) => `(stmt| continue)                                        -- subsumes (S8, R8, L2)
  | `(stmt| expand! $exp in for $x in $e do' $s) => `(stmt| for $x in $e do' expand! $exp in $s)  -- subsumes (L8, R9)

macro_rules
  | `(d! for $x in $e do' $s) => do  -- (D5), optimized like (1')
    let mut s := s
    let sb ← expandStmt (← `(stmt| expand! break in $s))
    let hasBreak := sb.raw.count (· matches `(stmt| break)) < s.raw.count (· matches `(stmt| break))
    if hasBreak then
      s := sb
    let sc ← expandStmt (← `(stmt| expand! continue in $s))
    let hasContinue := sc.raw.count (· matches `(stmt| continue)) < s.raw.count (· matches `(stmt| continue))
    if hasContinue then
      s := sc
    let mut body ← `(d! $s)
    if hasContinue then
      body ← `(ExceptCpsT.runCatch $body)
    let mut loop ← `(forM $e (fun $x => $body))
    if hasBreak then
      loop ← `(ExceptCpsT.runCatch $loop)
    pure loop
  | `(d! break%$b) =>
    throw <| Macro.Exception.error b "unexpected 'break' outside loop"
  | `(d! continue%$c) =>
    throw <| Macro.Exception.error c "unexpected 'continue' outside loop"

macro_rules
  | `(stmt| expand! break in break) => `(stmt| throw ())                                           -- (B1)
  | `(stmt| expand! break in $e:term) => `(stmt| ExceptCpsT.lift $e)                               -- (B3)
  | `(stmt| expand! break in for $x in $e do' $s) => `(stmt| for $x in $e do' expand! lift in $s)  -- (B8)
  | `(stmt| expand! continue in continue) => `(stmt| throw ())
  | `(stmt| expand! continue in $e:term) => `(stmt| ExceptCpsT.lift $e)
  | `(stmt| expand! continue in for $x in $e do' $s) => `(stmt| for $x in $e do' expand! lift in $s)

macro_rules
  | `(stmt| expand! lift in $e:term) => `(stmt| ExceptCpsT.lift $e)  -- (L3)

macro_rules
  | `(stmt| expand! mut $y in for $x in $e do' $s) => `(stmt| for $x in $e do' { let $y ← get; expand! mut $y in $s })  -- (S9)

variable [Monad m]
variable (ma ma' : m α)
variable (b : Bool)
variable (xs : List α) (act : α → m Unit)

attribute [local simp] map_eq_pure_bind

example [LawfulMonad m] :
    (do' for x in xs do' {
           act x
         })
    =
    xs.forM act
:= by induction xs <;> simp_all!

def ex2 (f : β → α → m β) (init : β) (xs : List α) : m β := do'
  let mut y := init;
  for x in xs do' {
    y ← f y x
  };
  return y

example [LawfulMonad m] (f : β → α → m β) :
    ex2 f init xs = xs.foldlM f init := by
  unfold ex2; induction xs generalizing init <;> simp_all!

@[simp] theorem List.find?_cons {xs : List α} : (x::xs).find? p = if p x then some x else xs.find? p := by
  cases h : p x <;> simp_all!

example (p : α → Bool) : Id.run
    (do' for x in xs do' {
           if p x then {
             return some x
           }
         };
         pure none)
    =
    xs.find? p
:= by induction xs with
      | nil => simp [Id.run, List.find?]
      | cons x => cases h : p x <;> simp_all [Id.run]

variable (p : α → m Bool)

theorem byCases_Bool_bind (x : m Bool) (f g : Bool → m β) (isTrue : f true = g true) (isFalse : f false = g false) : (x >>= f) = (x >>= g) := by
  have : f = g := by
    funext b
    cases b with
    | true  => exact isTrue
    | false => exact isFalse
  rw [this]

theorem eq_findM [LawfulMonad m] :
    (do' for x in xs do' {
           let b ← p x;
           if b then {
             return some x
           }
         };
         pure none)
    =
    xs.findM? p
:= by induction xs with
      | nil => simp!
      | cons x xs ih =>
        rw [List.findM?, ← ih]; simp
        apply byCases_Bool_bind <;> simp

def ex3 [Monad m] (p : α → m Bool) (xss : List (List α)) : m (Option α) := do'
  for xs in xss do' {
    for x in xs do' {
      let b ← p x;
      if b then {
        return some x
      }
    }
  };
  pure none

theorem eq_findSomeM_findM [LawfulMonad m] (xss : List (List α)) :
    ex3 p xss = xss.findSomeM? (fun xs => xs.findM? p) := by
  unfold ex3
  induction xss with
  | nil => simp!
  | cons xs xss ih =>
    simp [List.findSomeM?]
    rw [← ih, ← eq_findM]
    induction xs with
    | nil => simp
    | cons x xs ih => simp; apply byCases_Bool_bind <;> simp [ih]

def List.untilM (p : α → m Bool) : List α → m Unit
  | []    => pure ()
  | a::as => p a >>= fun | true => pure () | false => as.untilM p

theorem eq_untilM [LawfulMonad m] :
  (do' for x in xs do' {
         let b ← p x;
         if b then {
           break
         }
       })
  =
  xs.untilM p
:= by induction xs with
      | nil => simp!
      | cons x xs ih =>
        simp [List.untilM]; rw [← ih]; clear ih
        apply byCases_Bool_bind <;> simp

/-
The notation `[0:10]` is a range from 0 to 10 (exclusively).
-/

#eval do'
  for x in [0:10] do' {
    if x > 5 then {
      break
    };
    for y in [0:x] do' {
      IO.println y;
      break
    };
    IO.println x
  }

#eval do'
  for x in [0:10] do' {
    if x > 5 then {
      break
    };
    IO.println x
  }

#eval do'
  for x in [0:10] do' {
    if x % 2 == 0 then {
      continue
    };
    if x > 5 then {
      break
    };
    IO.println x
  }

#eval do'
  for x in [0:10] do' {
    if x % 2 == 0 then {
      continue
    };
    if x > 5 then {
      return ()
    };
    IO.println x
  }

-- set_option trace.compiler.ir.init true
def ex1 (xs : List Nat) (z : Nat) : Id Nat := do'
  let mut s1 := 0;
  let mut s2 := 0;
  for x in xs do' {
    if x % 2 == 0 then {
      continue
    };
    if x == z then {
      return s1
    };
    s1 := s1 + x;
    s2 := s2 + s1
  };
  return (s1 + s2)

/-
Adding `repeat` and `while` statements
-/

-- The "partial" keyword allows users to define non-terminating functions in Lean.
-- However, we currently cannot reason about them.
@[specialize] partial def loopForever [Monad m] (f : Unit → m Unit) : m Unit :=
  f () *> loopForever f

-- `Loop'` is a "helper" type. It is similar to `Unit`.
-- Its `ForM` instance produces an "infinite" sequence of units.
inductive Loop' where
  | mk : Loop'

instance : ForM m Loop' Unit where
  forM _ f := loopForever f

macro:0 "repeat" s:stmt:1 : stmt => `(stmt| for u in Loop'.mk do' $s)

#eval do'
  let mut i := 0;
  repeat {
    if i > 10 then {
      break
    };
    IO.println i;
    i := i + 1
  };
  return i

macro:0 "while" c:term "do'" s:stmt:1 : stmt => `(stmt| repeat { unless $c do' break; { $s } })

#eval do'
  let mut i := 0;
  while (i < 10) do' {
    IO.println i;
    i := i + 1
  };
  return i


import category_theory.preadditive
import category_theory.abelian.projective
import tactic.interval_cases

noncomputable theory

open category_theory
open category_theory.limits

universe variables v u

namespace category_theory

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

namespace fin3_functor_mk

variables (F : fin 3 → C) (a : F 0 ⟶ F 1) (b : F 1 ⟶ F 2)

def map' : Π (i j : fin 3) (hij : i ≤ j), F i ⟶ F j
| ⟨0,hi⟩ ⟨0,hj⟩ _ := 𝟙 _
| ⟨1,hi⟩ ⟨1,hj⟩ _ := 𝟙 _
| ⟨2,hi⟩ ⟨2,hj⟩ _ := 𝟙 _
| ⟨0,hi⟩ ⟨1,hj⟩ _ := a
| ⟨1,hi⟩ ⟨2,hj⟩ _ := b
| ⟨0,hi⟩ ⟨2,hj⟩ _ := a ≫ b
| ⟨i+3,hi⟩ _ _ := by { exfalso, revert hi, dec_trivial }
| _ ⟨j+3,hj⟩ _ := by { exfalso, revert hj, dec_trivial }
| ⟨i+1,hi⟩ ⟨0,hj⟩ H := by { exfalso, revert H, dec_trivial }
| ⟨i+2,hi⟩ ⟨1,hj⟩ H := by { exfalso, revert H, dec_trivial }
.

lemma map'_id : ∀ (i : fin 3), map' F a b i i le_rfl = 𝟙 _
| ⟨0,hi⟩ := rfl
| ⟨1,hi⟩ := rfl
| ⟨2,hi⟩ := rfl
| ⟨i+3,hi⟩ := by { exfalso, revert hi, dec_trivial }

lemma map'_comp : Π (i j k : fin 3) (hij : i ≤ j) (hjk : j ≤ k),
  map' F a b i j hij ≫ map' F a b j k hjk = map' F a b i k (hij.trans hjk)
| ⟨0, _⟩ ⟨0, _⟩ k _ _ := category.id_comp _
| ⟨1, _⟩ ⟨1, _⟩ k _ _ := category.id_comp _
| i ⟨1, _⟩ ⟨1, _⟩ _ _ := category.comp_id _
| i ⟨2, _⟩ ⟨2, _⟩ _ _ := category.comp_id _
| ⟨0, _⟩ ⟨1, _⟩ ⟨2, _⟩ _ _ := rfl
| ⟨i+3,hi⟩ _ _ _ _ := by { exfalso, revert hi, dec_trivial }
| _ ⟨j+3,hj⟩ _ _ _ := by { exfalso, revert hj, dec_trivial }
| _ _ ⟨k+3,hk⟩ _ _ := by { exfalso, revert hk, dec_trivial }
| ⟨i+1,hi⟩ ⟨0,hj⟩ _ H _ := by { exfalso, revert H, dec_trivial }
| ⟨i+2,hi⟩ ⟨1,hj⟩ _ H _ := by { exfalso, revert H, dec_trivial }
| _ ⟨i+1,hi⟩ ⟨0,hj⟩ _ H := by { exfalso, revert H, dec_trivial }
| _ ⟨i+2,hi⟩ ⟨1,hj⟩ _ H := by { exfalso, revert H, dec_trivial }

end fin3_functor_mk

def fin3_functor_mk (F : fin 3 → C) (a : F 0 ⟶ F 1) (b : F 1 ⟶ F 2) : fin 3 ⥤ C :=
{ obj := F,
  map := λ i j hij, fin3_functor_mk.map' F a b i j hij.le,
  map_id' := λ i, fin3_functor_mk.map'_id F a b i,
  map_comp' := λ i j k hij hjk, by rw fin3_functor_mk.map'_comp F a b i j k hij.le hjk.le }

namespace fin4_functor_mk

variables (F : fin 4 → C) (a : F 0 ⟶ F 1) (b : F 1 ⟶ F 2) (c : F 2 ⟶ F 3)

def map' : Π (i j : fin 4) (hij : i ≤ j), F i ⟶ F j
| ⟨0,hi⟩ ⟨0,hj⟩ _ := 𝟙 _
| ⟨1,hi⟩ ⟨1,hj⟩ _ := 𝟙 _
| ⟨2,hi⟩ ⟨2,hj⟩ _ := 𝟙 _
| ⟨3,hi⟩ ⟨3,hj⟩ _ := 𝟙 _
| ⟨0,hi⟩ ⟨1,hj⟩ _ := a
| ⟨1,hi⟩ ⟨2,hj⟩ _ := b
| ⟨2,hi⟩ ⟨3,hj⟩ _ := c
| ⟨0,hi⟩ ⟨2,hj⟩ _ := a ≫ b
| ⟨1,hi⟩ ⟨3,hj⟩ _ := b ≫ c
| ⟨0,hi⟩ ⟨3,hj⟩ _ := a ≫ b ≫ c
| ⟨i+4,hi⟩ _ _ := by { exfalso, revert hi, dec_trivial }
| _ ⟨j+4,hj⟩ _ := by { exfalso, revert hj, dec_trivial }
| ⟨i+1,hi⟩ ⟨0,hj⟩ H := by { exfalso, revert H, dec_trivial }
| ⟨i+2,hi⟩ ⟨1,hj⟩ H := by { exfalso, revert H, dec_trivial }
| ⟨3,hi⟩ ⟨2,hj⟩ H := by { exfalso, revert H, dec_trivial }
.

lemma map'_id : ∀ (i : fin 4), map' F a b c i i le_rfl = 𝟙 _
| ⟨0,hi⟩ := rfl
| ⟨1,hi⟩ := rfl
| ⟨2,hi⟩ := rfl
| ⟨3,hi⟩ := rfl
| ⟨i+4,hi⟩ := by { exfalso, revert hi, dec_trivial }

lemma map'_comp : Π (i j k : fin 4) (hij : i ≤ j) (hjk : j ≤ k),
  map' F a b c i j hij ≫ map' F a b c j k hjk = map' F a b c i k (hij.trans hjk)
| ⟨0, _⟩ ⟨0, _⟩ k _ _ := category.id_comp _
| ⟨1, _⟩ ⟨1, _⟩ k _ _ := category.id_comp _
| ⟨2, _⟩ ⟨2, _⟩ k _ _ := category.id_comp _
| i ⟨1, _⟩ ⟨1, _⟩ _ _ := category.comp_id _
| i ⟨2, _⟩ ⟨2, _⟩ _ _ := category.comp_id _
| i ⟨3, _⟩ ⟨3, _⟩ _ _ := category.comp_id _
| ⟨0, _⟩ ⟨1, _⟩ ⟨2, _⟩ _ _ := rfl
| ⟨0, _⟩ ⟨1, _⟩ ⟨3, _⟩ _ _ := rfl
| ⟨0, _⟩ ⟨2, _⟩ ⟨3, _⟩ _ _ := category.assoc a b c
| ⟨1, _⟩ ⟨2, _⟩ ⟨3, _⟩ _ _ := rfl
| ⟨i+4,hi⟩ _ _ _ _ := by { exfalso, revert hi, dec_trivial }
| _ ⟨j+4,hj⟩ _ _ _ := by { exfalso, revert hj, dec_trivial }
| _ _ ⟨k+4,hk⟩ _ _ := by { exfalso, revert hk, dec_trivial }
| ⟨i+1,hi⟩ ⟨0,hj⟩ _ H _ := by { exfalso, revert H, dec_trivial }
| ⟨i+2,hi⟩ ⟨1,hj⟩ _ H _ := by { exfalso, revert H, dec_trivial }
| ⟨3,hi⟩ ⟨2,hj⟩ _ H _ := by { exfalso, revert H, dec_trivial }
| _ ⟨i+1,hi⟩ ⟨0,hj⟩ _ H := by { exfalso, revert H, dec_trivial }
| _ ⟨i+2,hi⟩ ⟨1,hj⟩ _ H := by { exfalso, revert H, dec_trivial }
| _ ⟨3,hi⟩ ⟨2,hj⟩ _ H := by { exfalso, revert H, dec_trivial }

end fin4_functor_mk

def fin4_functor_mk (F : fin 4 → C) (a : F 0 ⟶ F 1) (b : F 1 ⟶ F 2) (c : F 2 ⟶ F 3) : fin 4 ⥤ C :=
{ obj := F,
  map := λ i j hij, fin4_functor_mk.map' F a b c i j hij.le,
  map_id' := λ i, fin4_functor_mk.map'_id F a b c i,
  map_comp' := λ i j k hij hjk, by rw fin4_functor_mk.map'_comp F a b c i j k hij.le hjk.le }

end category_theory


import system.io
import .solvers.z3
import .syntax
import .builder
import .tactic
import .attributes
import .lol
import init.data.option.basic

declare_trace smt2

open tactic
open smt2.builder
open native

meta structure smt2_state : Type :=
(ctxt : lol.context)
(type_map : rb_map expr lol.type)

meta def smt2_state.initial : smt2_state :=
⟨ lol.context.empty, rb_map.mk _ _ ⟩

@[reducible] meta def smt2_m (α : Type) :=
state_t smt2_state tactic α

meta instance tactic_to_smt2_m (α : Type) : has_coe (tactic α) (smt2_m α) :=
⟨ fun tc, state_t.mk (fun s, do res ← tc, return (res, s)) ⟩

namespace smt2

meta def trace_smt2 (msg : string) : smt2_m unit :=
  tactic.when_tracing `smt2 (tactic.trace msg)

meta def fail {α : Type} (msg : string) : smt2_m α :=
tactic.fail $ "smt2_tactic: " ++ msg

meta def mangle_name (n : name) : string :=
"lean_" ++ n^.to_string_with_sep "-"

meta def insert_type (n : string) (ty : expr) (lty : lol.type) : smt2_m unit :=
do st ← get,
   put ⟨
     st.ctxt.declare_type n lty,
     st.type_map.insert ty lty
   ⟩

meta def fn_type : expr → (list expr × expr)
| (expr.pi _ _ ty rest) :=
    let (args, rt) := fn_type rest
    in (ty :: args, rt)
| rt := ([], rt)

-- Currently we only support first order fn types
meta def compile_arrow_type (ty : expr) (cb : expr → smt2_m lol.type) : smt2_m lol.type :=
let (args, rt) := fn_type ty
in lol.type.fn <$> monad.mapm cb args <*> cb rt

meta def compile_type : expr → smt2_m lol.type :=
fun ty,
do st ← get,
   match st.type_map.find ty with
   | some lty := return lty
   | none := do
     lty ← match ty with
     | `(int) := pure $ lol.type.int
     | `(nat) := pure $ lol.type.refinement lol.type.int (fun x, lol.term.lte (lol.term.int 0) (lol.term.var x))
     | `(Prop) := pure $ lol.type.bool
     | _ := if ty.is_arrow
            then compile_arrow_type ty compile_type
            else if ty.is_constant
            then do insert_type (mangle_name ty.const_name) ty (lol.type.fn [] (lol.type.var $ mangle_name ty.const_name)),
                 return $ (lol.type.fn [] (lol.type.var $ mangle_name ty.const_name))
            else fail $ "unsupported type: " ++ to_string ty
     end,
     -- insert_type ty lty,
     return lty
   end

meta def add_decl (n : name) (ty : expr) : smt2_m unit :=
  do st ← get,
     ct ← compile_type ty,
     let d := lol.decl.fn (mangle_name n) ct none,
     put { st with ctxt := st.ctxt.declare d }

-- meta def ensure_constant (e : expr) (n : name) : smt2_m lol.decl :=
--   do ty ← infer_type e,
--    let (arg_tys, ret_ty) := fn_type ty,
--    let mangled_name := mangle_name n,
--    arg_sorts ← monad.mapm compile_type arg_tys,
--    ret_sort ← compile_type ret_ty,
--    -- ensure_constant_core e (return $ (mangled_name, arg_sorts, ret_sort)),
--    return $ lol.decl.fn mangled_name arg_sorts ret_sort

-- meta def formula_type_from_arrow (n : name) (e : expr) : smt2_m formula_type :=
-- do (lol.decl.fn _ arg_sorts ret_sort) ← ensure_constant e n,
--    return $ formula_type.fn n arg_sorts ret_sort

-- /-- The goal of this function is to categorize the set of formulas in the hypotheses,
--     and goal. We want to narrow down from the full term language of Lean to a fragment
--     of formula's we suppose. The below code makes some assumptions:

--    A local constant of the form `(P : Prop)`, must be reflected as declaration
--    in SMT2 that is `(declare-const P Bool)`.

--    An occurence of a proof of `P`, `(p : P)`, must be transformed into
--    `(assert P)`. If P is a formula, not an atom, we must transform P into a corresponding
--    SMT2 formula and `(assert P)`.
-- -/

meta def extract_coe_args (args : list expr) : smt2_m (expr × expr × expr) :=
match args with
| (source :: target :: inst :: e :: []) := return (source, target, e)
| _ := fail "internal tactic error expected `coe` to have exactly 4 arguments"
end

meta def reflect_coercion (source target e : expr) (callback : expr → smt2_m lol.term) : smt2_m lol.term :=
if source = `(nat) ∧ target = `(int)
then callback e
else fail $ "unsupported coercion between " ++ "`" ++ to_string source ++ "` and `" ++ to_string target ++ "`"

meta def reflect_application (fn : expr) (args : list expr) (callback : expr → smt2_m lol.term) : smt2_m lol.term :=
    if fn.is_constant
    then if fn.const_name = `coe
          then do (source, target, e) ← extract_coe_args args,
                   reflect_coercion source target e callback
          else do ty ← infer_type fn,
                  let mangled := (mangle_name fn.const_name),
                  add_decl fn.const_name ty,
                  lol.term.apply mangled <$> monad.mapm callback args
    else if fn.is_local_constant
    then lol.term.apply (mangle_name fn.local_uniq_name) <$> monad.mapm callback args
    else fail $ "unsupported head symbol `" ++ to_string fn ++ "`"

-- meta def is_supported_head_symbol (e : expr) : bool := true

meta def is_supported_numeric_ty (ty : expr) : bool :=
(ty = `(int) ∨ ty = `(nat))

-- /-- This function is the meat of the tactic, it takes a propositional formula in Lean, and transforms
--    it into a corresponding term in SMT2. -/
meta def reflect_arith_formula (reflect_base : expr → smt2_m lol.term) : expr → smt2_m lol.term
| `(%%a + %%b) := lol.term.add <$> reflect_arith_formula a <*> reflect_arith_formula b
| `(%%a - %%b) := lol.term.sub <$> reflect_arith_formula a <*> reflect_arith_formula b
| `(%%a * %%b) := lol.term.mul <$> reflect_arith_formula a <*> reflect_arith_formula b
| `(%%a / %%b) := lol.term.div <$> reflect_arith_formula a <*> reflect_arith_formula b
| `(%%a % %%b) := lol.term.mod <$> reflect_arith_formula a <*> reflect_arith_formula b
| `(- %%a) := lol.term.neg <$> reflect_arith_formula a
-- /- Constants -/
| `(has_zero.zero _) := lol.term.int <$> eval_expr int `(has_zero.zero int)
| `(has_one.one _) := lol.term.int <$> eval_expr int `(has_one.one int)
| `(bit0 %%Bits) :=
  do ty ← infer_type Bits,
     if is_supported_numeric_ty ty
     then lol.term.int <$> eval_expr int `(bit0 %%Bits : int)
     else if (ty = `(nat))
     then lol.term.int <$> int.of_nat <$> eval_expr nat `(bit0 %%Bits : nat)
     else fail $ "unknown numeric literal: " ++ (to_string ```(bit0 %%Bits : int))
| `(bit1 %%Bits) :=
  do ty ← infer_type Bits,
     if is_supported_numeric_ty ty
     then lol.term.int <$> eval_expr int `(bit1 %%Bits : int)
     else if (ty = `(nat))
     then lol.term.int <$> (int.of_nat <$> eval_expr nat `(bit1 %%Bits : nat))
     else fail $ "unknown numeric literal: " ++ (to_string `(bit1 %%Bits : int))
| a :=
    if a.is_local_constant
    then return $ lol.term.var (mangle_name a.local_uniq_name)
    else if a.is_constant
    then return $ lol.term.var (mangle_name a.const_name)
    else if a.is_app
    then reflect_application (a.get_app_fn) (a.get_app_args) reflect_base
    else fail $ "unsupported arithmetic formula: " ++ to_string a

-- /-- Check if the type is an `int` or logically a subtype of an `int` like nat. -/
meta def is_int (e : expr) : tactic bool :=
do ty ← infer_type e,
   return $ (ty = `(int)) || (ty = `(nat))

meta def unsupported_ordering_on {α : Type} (elem : expr) : tactic α :=
do ty ← infer_type elem,
   tactic.fail $ "unable to translate orderings for values of type: " ++ to_string ty

meta def reflect_ordering (reflect_arith : expr → smt2_m lol.term) (R : lol.term → lol.term → lol.term) (P Q : expr) : smt2_m lol.term :=
do is ← is_int P, -- NB: P and Q should have the same type.
   if is
   then R <$> (reflect_arith P) <*> (reflect_arith Q)
   else unsupported_ordering_on P

meta def supported_pi_binder (ty : expr) : bool :=
match ty with
| `(int) := tt
| `(nat) := tt
| `(Prop) := tt
| _ := if ty.is_constant
       then tt
       else ff
end

meta def add_assertion (t : lol.term) : smt2_m unit :=
  do st ← get,
     put { st with ctxt := st.ctxt.assert t }

meta def compile_pi (e : expr) (cb : expr → smt2_m lol.term) : smt2_m lol.term :=
if supported_pi_binder e.binding_domain
then do loc ← tactic.mk_local' e.binding_name e.binding_info e.binding_domain,
        lol.term.forallq
          (mangle_name $ loc.local_uniq_name) <$>
          (compile_type $ e.binding_domain) <*>
          (cb (expr.instantiate_var (e.binding_body) loc))
else fail $ "arbitrary Π types are not supported, unable to translate term: `" ++ to_string e ++ "`"

meta def reflect_prop_formula' : expr → smt2_m lol.term
| `(¬ %%P) := lol.term.not <$> (reflect_prop_formula' P)
| `(%%P = %% Q) := lol.term.equals <$> (reflect_prop_formula' P) <*> (reflect_prop_formula' Q)
| `(%%P ∧ %%Q) := lol.term.and <$> (reflect_prop_formula' P) <*> (reflect_prop_formula' Q)
| `(%%P ∨ %%Q) := lol.term.or <$> (reflect_prop_formula' P) <*> (reflect_prop_formula' Q)
| `(%%P ↔ %%Q) := lol.term.iff <$> (reflect_prop_formula' P) <*> (reflect_prop_formula' Q)
| `(%%P < %%Q) := reflect_ordering (reflect_arith_formula reflect_prop_formula') lol.term.lt P Q
| `(%%P <= %%Q) := reflect_ordering (reflect_arith_formula reflect_prop_formula') lol.term.lte P Q
| `(%%P > %%Q) := reflect_ordering (reflect_arith_formula reflect_prop_formula') lol.term.gt P Q
| `(%%P >= %%Q) := reflect_ordering (reflect_arith_formula reflect_prop_formula') lol.term.gte P Q
| `(true) := return $ lol.term.true
| `(false) := return $ lol.term.false
| e := do ty ← infer_type e,
       if e.is_local_constant
       then pure $ lol.term.var (mangle_name e.local_uniq_name)
       else if e.is_arrow
       then lol.term.implies <$> (reflect_prop_formula' e.binding_domain) <*> (reflect_prop_formula' e.binding_body )
       else if e.is_pi
       then compile_pi e reflect_prop_formula'
       else if is_supported_numeric_ty ty
       then reflect_arith_formula reflect_prop_formula' e
       else if e.is_app
       then reflect_application (e.get_app_fn) (e.get_app_args) reflect_prop_formula'
       else tactic.fail $ "unsupported propositional formula : " ++ to_string e

meta def reflect_prop_formula (e : expr) : smt2_m unit :=
reflect_prop_formula' e >>= add_assertion

-- meta def warn_unable_to_trans_local (e : expr) : smt2_m (builder unit) := do
--   trace_smt2 $ "unable to translate local variable: " ++ to_string e,
--   return $ return ()

meta def is_builtin_type : expr → bool
| `(int) := tt
| `(Prop) := tt
| `(nat) := tt
| _ := ff

meta def unsupported_formula (e : expr) : smt2_m unit :=
fail $ "unsupported formula: " ++ to_string e

meta def compile_local (e : expr) : smt2_m unit :=
do ty ← infer_type e,
   prop_sorted ← is_prop ty,
   if e.is_local_constant
   then if is_builtin_type ty
        then add_decl e.local_uniq_name ty
        else if ty.is_arrow
        then add_decl e.local_uniq_name ty
        else if prop_sorted
        then reflect_prop_formula ty
        else unsupported_formula ty
   else if e.is_constant
   then if is_builtin_type ty ∨ ty.is_arrow
        then add_decl e.const_name ty
        else if prop_sorted
        then reflect_prop_formula ty
        else unsupported_formula ty
   else if (ty = `(Prop))
   then reflect_prop_formula e
   else unsupported_formula e

meta def reflect_attr_decl (n : name) : smt2_m unit :=
do exp ← mk_const n,
   compile_local exp

/- Reflect the environment consisting of declarations with the `smt2` attribute. -/
meta def reflect_environment : smt2_m unit :=
do decls ← attribute.get_instances `smt2,
   bs ← monad.mapm reflect_attr_decl decls.reverse,
   return ()

meta def reflect_context : smt2_m unit :=
 do ls ← local_context,
    bs ← monad.mapm (fun e, compile_local e) ls,
    return ()

meta def reflect_goal : smt2_m unit :=
  do tgt ← target,
     -- SMT solvers are looking for satisfiabiltiy, so we must negate to check validity.
     reflect_prop_formula `(_root_.not %%tgt),
     return ()

meta def reflect : smt2_m (builder unit) :=
do reflect_environment,
   reflect_context,
   reflect_goal,
   st ← get,
   return $ (lol.to_builder (lol.smt2_compiler_state.mk (rb_map.mk _ _) st.ctxt []) lol.compile >> check_sat)

end smt2

universe u

@[smt2] lemma int_of_nat_is_pos :
  forall (n : nat), 0 <= int.of_nat n :=
begin
  intros, trivial
end

axiom proof_by_z3 (A : Sort u) : A

meta def z3 (log_file : option string := none) : tactic unit :=
do (builder, _) ← smt2.reflect.run smt2_state.initial,
   resp ← unsafe_run_io (smt2 builder log_file),
   match resp with
   | smt2.response.sat := fail "z3 was unable to prove the goal"
   | smt2.response.unknown := fail "z3 was unable to prove the goal"
   | smt2.response.other str := fail $ "z3 communication error, unexpected response:\n\n" ++ str ++ "\n"
   | smt2.response.unsat := do
        tgt ← target,
        sry ← to_expr $ ``(proof_by_z3 %%tgt),
        exact sry
   end