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


/-
Copyright (c) 2017 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.meta.smt.congruence_closure
import Mathlib.Lean3Lib.init.meta.attribute
import Mathlib.Lean3Lib.init.meta.simp_tactic
import Mathlib.Lean3Lib.init.meta.interactive_base
import Mathlib.Lean3Lib.init.meta.derive

universes l

namespace Mathlib

/-- Heuristic instantiation lemma -/
/-- `mk_core m e as_simp`, m is used to decide which definitions will be unfolded in patterns.
   If as_simp is tt, then this tactic will try to use the left-hand-side of the conclusion
   as a pattern. -/
/--
Create a new "cached" attribute (attr_name : user_attribute hinst_lemmas).
It also creates "cached" attributes for each attr_names and simp_attr_names if they have not been defined
yet. Moreover, the hinst_lemmas for attr_name will be the union of the lemmas tagged with
    attr_name, attrs_name, and simp_attr_names.
For the ones in simp_attr_names, we use the left-hand-side of the conclusion as the pattern.
-/
structure ematch_config
where
  max_instances : ℕ
  max_generation : ℕ


-- 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