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

Monad encapsulating continuation passing programming style, similar to
Haskell's `Cont`, `ContT` and `MonadCont`:
<http://hackage.haskell.org/package/mtl-2.2.2/docs/Control-Monad-Cont.html>
-/
import control.monad.writer

universes u v w u₀ u₁ v₀ v₁

structure monad_cont.label (α : Type w) (m : Type u → Type v) (β : Type u) :=
(apply : α → m β)

def monad_cont.goto {α β} {m : Type u → Type v} (f : monad_cont.label α m β) (x : α) := f.apply x

class monad_cont (m : Type u → Type v) :=
(call_cc : Π {α β}, ((monad_cont.label α m β) → m α) → m α)

open monad_cont

class is_lawful_monad_cont (m : Type u → Type v) [monad m] [monad_cont m]
extends is_lawful_monad m :=
(call_cc_bind_right {α ω γ} (cmd : m α) (next : (label ω m γ) → α → m ω) :
  call_cc (λ f, cmd >>= next f) = cmd >>= λ x, call_cc (λ f, next f x))
(call_cc_bind_left {α} (β) (x : α) (dead : label α m β → β → m α) :
  call_cc (λ f : label α m β, goto f x >>= dead f) = pure x)
(call_cc_dummy {α β} (dummy : m α) :
  call_cc (λ f : label α m β, dummy) = dummy)

export is_lawful_monad_cont

def cont_t (r : Type u) (m : Type u → Type v) (α : Type w) := (α → m r) → m r

@[reducible] def cont (r : Type u) (α : Type w) := cont_t r id α

namespace cont_t

export monad_cont (label goto)

variables {r : Type u} {m : Type u → Type v} {α β γ ω : Type w}

def run : cont_t r m α → (α → m r) → m r := id

def map (f : m r → m r) (x : cont_t r m α) : cont_t r m α := f ∘ x

lemma run_cont_t_map_cont_t (f : m r → m r) (x : cont_t r m α) :
  run (map f x) = f ∘ run x := rfl

def with_cont_t (f : (β → m r) → α → m r) (x : cont_t r m α) : cont_t r m β :=
λ g, x $ f g

lemma run_with_cont_t (f : (β → m r) → α → m r) (x : cont_t r m α) :
  run (with_cont_t f x) = run x ∘ f := rfl

@[ext]
protected lemma ext {x y : cont_t r m α}
  (h : ∀ f, x.run f = y.run f) :
  x = y := by { ext; apply h }

instance : monad (cont_t r m) :=
{ pure := λ α x f, f x,
  bind := λ α β x f g, x $ λ i, f i g }

instance : is_lawful_monad (cont_t r m) :=
{ id_map := by { intros, refl },
  pure_bind := by { intros, ext, refl },
  bind_assoc := by { intros, ext, refl } }

def monad_lift [monad m] {α} : m α → cont_t r m α :=
λ x f, x >>= f

instance [monad m] : has_monad_lift m (cont_t r m) :=
{ monad_lift := λ α, cont_t.monad_lift }

lemma monad_lift_bind [monad m] [is_lawful_monad m] {α β} (x : m α) (f : α → m β) :
  (monad_lift (x >>= f) : cont_t r m β) = monad_lift x >>= monad_lift ∘ f :=
begin
  ext,
  simp only [monad_lift,has_monad_lift.monad_lift,(∘),(>>=),bind_assoc,id.def,run,cont_t.monad_lift]
end

instance : monad_cont (cont_t r m) :=
{ call_cc := λ α β f g, f ⟨λ x h, g x⟩ g }

instance : is_lawful_monad_cont (cont_t r m) :=
{ call_cc_bind_right := by intros; ext; refl,
  call_cc_bind_left := by intros; ext; refl,
  call_cc_dummy := by intros; ext; refl }

instance (ε) [monad_except ε m] : monad_except ε (cont_t r m) :=
{ throw := λ x e f, throw e,
  catch := λ α act h f, catch (act f) (λ e, h e f) }

instance : monad_run (λ α, (α → m r) → ulift.{u v} (m r)) (cont_t.{u v u} r m) :=
{ run := λ α f x, ⟨ f x ⟩ }

end cont_t

variables {m : Type u → Type v} [monad m]

def except_t.mk_label {α β ε} : label (except.{u u} ε α) m β → label α (except_t ε m) β
| ⟨ f ⟩ := ⟨ λ a, monad_lift $ f (except.ok a) ⟩

lemma except_t.goto_mk_label {α β ε : Type*} (x : label (except.{u u} ε α) m β) (i : α) :
  goto (except_t.mk_label x) i = ⟨ except.ok <$> goto x (except.ok i) ⟩ := by cases x; refl

def except_t.call_cc
  {ε} [monad_cont m] {α β : Type*} (f : label α (except_t ε m) β → except_t ε m α) :
  except_t ε m α :=
except_t.mk (call_cc $ λ x : label _ m β, except_t.run $ f (except_t.mk_label x) : m (except ε α))

instance {ε} [monad_cont m] : monad_cont (except_t ε m) :=
{ call_cc := λ α β, except_t.call_cc }

instance {ε} [monad_cont m] [is_lawful_monad_cont m] : is_lawful_monad_cont (except_t ε m) :=
{ call_cc_bind_right := by { intros, simp [call_cc,except_t.call_cc,call_cc_bind_right], ext, dsimp,
    congr' with ⟨ ⟩; simp [except_t.bind_cont,@call_cc_dummy m _], },
  call_cc_bind_left  := by { intros,
    simp [call_cc,except_t.call_cc,call_cc_bind_right,except_t.goto_mk_label,map_eq_bind_pure_comp,
      bind_assoc,@call_cc_bind_left m _], ext, refl },
  call_cc_dummy := by { intros, simp [call_cc,except_t.call_cc,@call_cc_dummy m _], ext, refl }, }

def option_t.mk_label {α β} : label (option.{u} α) m β → label α (option_t m) β
| ⟨ f ⟩ := ⟨ λ a, monad_lift $ f (some a) ⟩

lemma option_t.goto_mk_label {α β : Type*} (x : label (option.{u} α) m β) (i : α) :
  goto (option_t.mk_label x) i = ⟨ some <$> goto x (some i) ⟩ := by cases x; refl

def option_t.call_cc [monad_cont m] {α β : Type*} (f : label α (option_t m) β → option_t m α) :
  option_t m α :=
option_t.mk (call_cc $ λ x : label _ m β, option_t.run $ f (option_t.mk_label x) : m (option α))

instance [monad_cont m] : monad_cont (option_t m) :=
{ call_cc := λ α β, option_t.call_cc }

instance [monad_cont m] [is_lawful_monad_cont m] : is_lawful_monad_cont (option_t m) :=
{ call_cc_bind_right := by { intros, simp [call_cc,option_t.call_cc,call_cc_bind_right], ext, dsimp,
    congr' with ⟨ ⟩; simp [option_t.bind_cont,@call_cc_dummy m _], },
  call_cc_bind_left  := by { intros, simp [call_cc,option_t.call_cc,call_cc_bind_right,
    option_t.goto_mk_label,map_eq_bind_pure_comp,bind_assoc,@call_cc_bind_left m _], ext, refl },
  call_cc_dummy := by { intros, simp [call_cc,option_t.call_cc,@call_cc_dummy m _], ext, refl }, }

def writer_t.mk_label {α β ω} [has_one ω] : label (α × ω) m β → label α (writer_t ω m) β
| ⟨ f ⟩ := ⟨ λ a, monad_lift $ f (a,1) ⟩

lemma writer_t.goto_mk_label {α β ω : Type*} [has_one ω] (x : label (α × ω) m β) (i : α) :
  goto (writer_t.mk_label x) i = monad_lift (goto x (i,1)) := by cases x; refl

def writer_t.call_cc [monad_cont m] {α β ω : Type*} [has_one ω]
  (f : label α (writer_t ω m) β → writer_t ω m α) : writer_t ω m α :=
⟨ call_cc (writer_t.run ∘ f ∘ writer_t.mk_label : label (α × ω) m β → m (α × ω)) ⟩

instance (ω) [monad m] [has_one ω] [monad_cont m] : monad_cont (writer_t ω m) :=
{ call_cc := λ α β, writer_t.call_cc }

def state_t.mk_label {α β σ : Type u} : label (α × σ) m (β × σ) → label α (state_t σ m) β
| ⟨ f ⟩ := ⟨ λ a, ⟨ λ s, f (a,s) ⟩ ⟩

lemma state_t.goto_mk_label {α β σ : Type u} (x : label (α × σ) m (β × σ)) (i : α) :
  goto (state_t.mk_label x) i = ⟨ λ s, (goto x (i,s)) ⟩ := by cases x; refl

def state_t.call_cc {σ}  [monad_cont m] {α β : Type*}
  (f : label α (state_t σ m) β → state_t σ m α) : state_t σ m α :=
⟨ λ r, call_cc (λ f', (f $ state_t.mk_label f').run r) ⟩

instance {σ} [monad_cont m] : monad_cont (state_t σ m) :=
{ call_cc := λ α β, state_t.call_cc }

instance {σ} [monad_cont m] [is_lawful_monad_cont m] : is_lawful_monad_cont (state_t σ m) :=
{ call_cc_bind_right := by { intros,
    simp [call_cc,state_t.call_cc,call_cc_bind_right,(>>=),state_t.bind], ext, dsimp,
    congr' with ⟨x₀,x₁⟩, refl },
  call_cc_bind_left  := by { intros, simp [call_cc,state_t.call_cc,call_cc_bind_left,(>>=),
    state_t.bind,state_t.goto_mk_label], ext, refl },
  call_cc_dummy := by { intros, simp [call_cc,state_t.call_cc,call_cc_bind_right,(>>=),
    state_t.bind,@call_cc_dummy m _], ext, refl }, }

def reader_t.mk_label {α β} (ρ) : label α m β → label α (reader_t ρ m) β
| ⟨ f ⟩ := ⟨ monad_lift ∘ f ⟩

lemma reader_t.goto_mk_label {α ρ β} (x : label α m β) (i : α) :
  goto (reader_t.mk_label ρ x) i = monad_lift (goto x i) := by cases x; refl

def reader_t.call_cc {ε}  [monad_cont m] {α β : Type*}
  (f : label α (reader_t ε m) β → reader_t ε m α) : reader_t ε m α :=
⟨ λ r, call_cc (λ f', (f $ reader_t.mk_label _ f').run r) ⟩

instance {ρ} [monad_cont m] : monad_cont (reader_t ρ m) :=
{ call_cc := λ α β, reader_t.call_cc }

instance {ρ} [monad_cont m] [is_lawful_monad_cont m] : is_lawful_monad_cont (reader_t ρ m) :=
{ call_cc_bind_right :=
    by { intros, simp [call_cc,reader_t.call_cc,call_cc_bind_right], ext, refl },
  call_cc_bind_left  := by { intros, simp [call_cc,reader_t.call_cc,call_cc_bind_left,
    reader_t.goto_mk_label], ext, refl },
  call_cc_dummy := by { intros, simp [call_cc,reader_t.call_cc,@call_cc_dummy m _], ext, refl } }

/-- reduce the equivalence between two continuation passing monads to the equivalence between
their underlying monad -/
def cont_t.equiv {m₁ : Type u₀ → Type v₀} {m₂ : Type u₁ → Type v₁}
  {α₁ r₁ : Type u₀} {α₂ r₂ : Type u₁} (F : m₁ r₁ ≃ m₂ r₂) (G : α₁ ≃ α₂) :
  cont_t r₁ m₁ α₁ ≃ cont_t r₂ m₂ α₂ :=
{ to_fun := λ f r, F $ f $ λ x, F.symm $ r $ G x,
  inv_fun := λ f r, F.symm $ f $ λ x, F $ r $ G.symm x,
  left_inv := λ f, by funext r; simp,
  right_inv := λ f, by funext r; simp }


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

def List.asString (s : List Char) : String :=
  ⟨s⟩

namespace String

instance : OfNat String.Pos (nat_lit 0) where
  ofNat := {}

instance : LT String :=
  ⟨fun s₁ s₂ => s₁.data < s₂.data⟩

@[extern "lean_string_dec_lt"]
instance decLt (s₁ s₂ : @& String) : Decidable (s₁ < s₂) :=
  List.hasDecidableLt s₁.data s₂.data

@[extern "lean_string_length"]
def length : (@& String) → Nat
  | ⟨s⟩ => s.length

/-- The internal implementation uses dynamic arrays and will perform destructive updates
   if the String is not shared. -/
@[extern "lean_string_push"]
def push : String → Char → String
  | ⟨s⟩, c => ⟨s ++ [c]⟩

/-- The internal implementation uses dynamic arrays and will perform destructive updates
   if the String is not shared. -/
@[extern "lean_string_append"]
def append : String → (@& String) → String
  | ⟨a⟩, ⟨b⟩ => ⟨a ++ b⟩

/-- O(n) in the runtime, where n is the length of the String -/
def toList (s : String) : List Char :=
  s.data

private def utf8GetAux : List Char → Pos → Pos → Char
  | [],    _, _ => default
  | c::cs, i, p => if i = p then c else utf8GetAux cs (i + c) p

/--
  Return character at position `p`. If `p` is not a valid position
  returns `(default : Char)`.
  See `utf8GetAux` for the reference implementation.
-/
@[extern "lean_string_utf8_get"]
def get (s : @& String) (p : @& Pos) : Char :=
  match s with
  | ⟨s⟩ => utf8GetAux s 0 p

private def utf8GetAux? : List Char → Pos → Pos → Option Char
  | [],    _, _ => none
  | c::cs, i, p => if i = p then c else utf8GetAux cs (i + c) p

@[extern "lean_string_utf8_get_opt"]
def get? : (@& String) → (@& Pos) → Option Char
  | ⟨s⟩, p => utf8GetAux? s 0 p

/--
  Similar to `get`, but produces a panic error message if `p` is not a valid `String.Pos`.
-/
@[extern "lean_string_utf8_get_bang"]
def get! (s : @& String) (p : @& Pos) : Char :=
  match s with
  | ⟨s⟩ => utf8GetAux s 0 p

private def utf8SetAux (c' : Char) : List Char → Pos → Pos → List Char
  | [],    _, _ => []
  | c::cs, i, p =>
    if i = p then (c'::cs) else c::(utf8SetAux c' cs (i + c) p)

@[extern "lean_string_utf8_set"]
def set : String → (@& Pos) → Char → String
  | ⟨s⟩, i, c => ⟨utf8SetAux c s 0 i⟩

def modify (s : String) (i : Pos) (f : Char → Char) : String :=
  s.set i <| f <| s.get i

@[extern "lean_string_utf8_next"]
def next (s : @& String) (p : @& Pos) : Pos :=
  let c := get s p
  p + c

private def utf8PrevAux : List Char → Pos → Pos → Pos
  | [],    _, _ => 0
  | c::cs, i, p =>
    let i' := i + c
    if i' = p then i else utf8PrevAux cs i' p

@[extern "lean_string_utf8_prev"]
def prev : (@& String) → (@& Pos) → Pos
  | ⟨s⟩, p => if p = 0 then 0 else utf8PrevAux s 0 p

def front (s : String) : Char :=
  get s 0

def back (s : String) : Char :=
  get s (prev s s.endPos)

@[extern "lean_string_utf8_at_end"]
def atEnd : (@& String) → (@& Pos) → Bool
  | s, p => p.byteIdx ≥ utf8ByteSize s

/--
Similar to `get` but runtime does not perform bounds check.
-/
@[extern "lean_string_utf8_get_fast"]
def get' (s : @& String) (p : @& Pos) (h : ¬ s.atEnd p) : Char :=
  match s with
  | ⟨s⟩ => utf8GetAux s 0 p

@[extern "lean_string_utf8_next_fast"]
def next' (s : @& String) (p : @& Pos) (h : ¬ s.atEnd p) : Pos :=
  let c := get s p
  p + c

/- TODO: remove `partial` keywords after we restore the tactic
  framework and wellfounded recursion support -/

partial def posOfAux (s : String) (c : Char) (stopPos : Pos) (pos : Pos) : Pos :=
  if pos >= stopPos then pos
  else if s.get pos == c then pos
       else posOfAux s c stopPos (s.next pos)

@[inline] def posOf (s : String) (c : Char) : Pos :=
  posOfAux s c s.endPos 0

partial def revPosOfAux (s : String) (c : Char) (pos : Pos) : Option Pos :=
 if s.get pos == c then some pos
 else if pos == 0 then none
 else revPosOfAux s c (s.prev pos)

def revPosOf (s : String) (c : Char) : Option Pos :=
  if s.endPos == 0 then none
  else revPosOfAux s c (s.prev s.endPos)

partial def findAux (s : String) (p : Char → Bool) (stopPos : Pos) (pos : Pos) : Pos :=
  if pos >= stopPos then pos
  else if p (s.get pos) then pos
       else findAux s p stopPos (s.next pos)

@[inline] def find (s : String) (p : Char → Bool) : Pos :=
  findAux s p s.endPos 0

partial def revFindAux (s : String) (p : Char → Bool) (pos : Pos) : Option Pos :=
 if p (s.get pos) then some pos
 else if pos == 0 then none
 else revFindAux s p (s.prev pos)

def revFind (s : String) (p : Char → Bool) : Option Pos :=
  if s.endPos == 0 then none
  else revFindAux s p (s.prev s.endPos)

abbrev Pos.min (p₁ p₂ : Pos) : Pos :=
  { byteIdx := p₁.byteIdx.min p₂.byteIdx }

/-- Returns the first position where the two strings differ. -/
partial def firstDiffPos (a b : String) : Pos :=
  let stopPos := a.endPos.min b.endPos
  let rec loop (i : Pos) : Pos :=
    if i >= stopPos || a.get i != b.get i then i
    else loop (a.next i)
  loop 0

@[extern "lean_string_utf8_extract"]
def extract : (@& String) → (@& Pos) → (@& Pos) → String
  | ⟨s⟩, b, e => if b.byteIdx ≥ e.byteIdx then ⟨[]⟩ else ⟨go₁ s 0 b e⟩
where
  go₁ : List Char → Pos → Pos → Pos → List Char
    | [],        _, _, _ => []
    | s@(c::cs), i, b, e => if i = b then go₂ s i e else go₁ cs (i + c) b e

  go₂ : List Char → Pos → Pos → List Char
    | [],    _, _ => []
    | c::cs, i, e => if i = e then [] else c :: go₂ cs (i + c) e

@[specialize] partial def splitAux (s : String) (p : Char → Bool) (b : Pos) (i : Pos) (r : List String) : List String :=
  if s.atEnd i then
    let r := (s.extract b i)::r
    r.reverse
  else if p (s.get i) then
    let i := s.next i
    splitAux s p i i (s.extract b { byteIdx := i.byteIdx - 1 } :: r)
  else
    splitAux s p b (s.next i) r

@[specialize] def split (s : String) (p : Char → Bool) : List String :=
  splitAux s p 0 0 []

partial def splitOnAux (s sep : String) (b : Pos) (i : Pos) (j : Pos) (r : List String) : List String :=
  if s.atEnd i then
    let r := if sep.atEnd j then ""::(s.extract b (i - j))::r else (s.extract b i)::r
    r.reverse
  else if s.get i == sep.get j then
    let i := s.next i
    let j := sep.next j
    if sep.atEnd j then
      splitOnAux s sep i i 0 (s.extract b (i - j)::r)
    else
      splitOnAux s sep b i j r
  else
    splitOnAux s sep b (s.next i) 0 r

def splitOn (s : String) (sep : String := " ") : List String :=
  if sep == "" then [s] else splitOnAux s sep 0 0 0 []

instance : Inhabited String := ⟨""⟩

instance : Append String := ⟨String.append⟩

def str : String → Char → String := push

def pushn (s : String) (c : Char) (n : Nat) : String :=
  n.repeat (fun s => s.push c) s

def isEmpty (s : String) : Bool :=
  s.endPos == 0

def join (l : List String) : String :=
  l.foldl (fun r s => r ++ s) ""

def singleton (c : Char) : String :=
  "".push c

def intercalate (s : String) : List String → String
  | []      => ""
  | a :: as => go a s as
where go (acc : String) (s : String) : List String → String
  | a :: as => go (acc ++ s ++ a) s as
  | []      => acc

/-- Iterator for `String`. That is, a `String` and a position in that string. -/
structure Iterator where
  s : String
  i : Pos
  deriving DecidableEq

def mkIterator (s : String) : Iterator :=
  ⟨s, 0⟩

abbrev iter := mkIterator

instance : SizeOf String.Iterator where
  sizeOf i := i.1.utf8ByteSize - i.2.byteIdx

theorem Iterator.sizeOf_eq (i : String.Iterator) : sizeOf i = i.1.utf8ByteSize - i.2.byteIdx :=
  rfl

namespace Iterator
def toString : Iterator → String
  | ⟨s, _⟩ => s

def remainingBytes : Iterator → Nat
  | ⟨s, i⟩ => s.endPos.byteIdx - i.byteIdx

def pos : Iterator → Pos
  | ⟨_, i⟩ => i

def curr : Iterator → Char
  | ⟨s, i⟩ => get s i

def next : Iterator → Iterator
  | ⟨s, i⟩ => ⟨s, s.next i⟩

def prev : Iterator → Iterator
  | ⟨s, i⟩ => ⟨s, s.prev i⟩

def atEnd : Iterator → Bool
  | ⟨s, i⟩ => i.byteIdx ≥ s.endPos.byteIdx

def hasNext : Iterator → Bool
  | ⟨s, i⟩ => i.byteIdx < s.endPos.byteIdx

def hasPrev : Iterator → Bool
  | ⟨_, i⟩ => i.byteIdx > 0

def setCurr : Iterator → Char → Iterator
  | ⟨s, i⟩, c => ⟨s.set i c, i⟩

def toEnd : Iterator → Iterator
  | ⟨s, _⟩ => ⟨s, s.endPos⟩

def extract : Iterator → Iterator → String
  | ⟨s₁, b⟩, ⟨s₂, e⟩ =>
    if s₁ ≠ s₂ || b > e then ""
    else s₁.extract b e

def forward : Iterator → Nat → Iterator
  | it, 0   => it
  | it, n+1 => forward it.next n

def remainingToString : Iterator → String
  | ⟨s, i⟩ => s.extract i s.endPos

def nextn : Iterator → Nat → Iterator
  | it, 0   => it
  | it, i+1 => nextn it.next i

def prevn : Iterator → Nat → Iterator
  | it, 0   => it
  | it, i+1 => prevn it.prev i
end Iterator

partial def offsetOfPosAux (s : String) (pos : Pos) (i : Pos) (offset : Nat) : Nat :=
  if i >= pos || s.atEnd i then
    offset
  else
    offsetOfPosAux s pos (s.next i) (offset+1)

def offsetOfPos (s : String) (pos : Pos) : Nat :=
  offsetOfPosAux s pos 0 0

@[specialize] partial def foldlAux {α : Type u} (f : α → Char → α) (s : String) (stopPos : Pos) (i : Pos) (a : α) : α :=
  let rec loop (i : Pos) (a : α) :=
    if i >= stopPos then a
    else loop (s.next i) (f a (s.get i))
  loop i a

@[inline] def foldl {α : Type u} (f : α → Char → α) (init : α) (s : String) : α :=
  foldlAux f s s.endPos 0 init

@[specialize] partial def foldrAux {α : Type u} (f : Char → α → α) (a : α) (s : String) (stopPos : Pos) (i : Pos) : α :=
  let rec loop (i : Pos) :=
    if i >= stopPos then a
    else f (s.get i) (loop (s.next i))
  loop i

@[inline] def foldr {α : Type u} (f : Char → α → α) (init : α) (s : String) : α :=
  foldrAux f init s s.endPos 0

@[specialize] partial def anyAux (s : String) (stopPos : Pos) (p : Char → Bool) (i : Pos) : Bool :=
  let rec loop (i : Pos) :=
    if i >= stopPos then false
    else if p (s.get i) then true
    else loop (s.next i)
  loop i

@[inline] def any (s : String) (p : Char → Bool) : Bool :=
anyAux s s.endPos p 0

@[inline] def all (s : String) (p : Char → Bool) : Bool :=
!s.any (fun c => !p c)

def contains (s : String) (c : Char) : Bool :=
s.any (fun a => a == c)

@[specialize] partial def mapAux (f : Char → Char) (i : Pos) (s : String) : String :=
  if s.atEnd i then s
  else
    let c := f (s.get i)
    let s := s.set i c
    mapAux f (s.next i) s

@[inline] def map (f : Char → Char) (s : String) : String :=
  mapAux f 0 s

def isNat (s : String) : Bool :=
  !s.isEmpty && s.all (·.isDigit)

def toNat? (s : String) : Option Nat :=
  if s.isNat then
    some <| s.foldl (fun n c => n*10 + (c.toNat - '0'.toNat)) 0
  else
    none

/--
Return `true` iff the substring of byte size `sz` starting at position `off1` in `s1` is equal to that starting at `off2` in `s2.`.
False if either substring of that byte size does not exist. -/
partial def substrEq (s1 : String) (off1 : String.Pos) (s2 : String) (off2 : String.Pos) (sz : Nat) : Bool :=
  off1.byteIdx + sz ≤ s1.endPos.byteIdx && off2.byteIdx + sz ≤ s2.endPos.byteIdx && loop off1 off2 { byteIdx := off1.byteIdx + sz }
where
  loop (off1 off2 stop1 : Pos) :=
    if off1.byteIdx >= stop1.byteIdx then
      true
    else
      let c₁ := s1.get off1
      let c₂ := s2.get off2
      c₁ == c₂ && loop (off1 + c₁) (off2 + c₂) stop1

/-- Return true iff `p` is a prefix of `s` -/
def isPrefixOf (p : String) (s : String) : Bool :=
  substrEq p 0 s 0 p.endPos.byteIdx

/-- Replace all occurrences of `pattern` in `s` with `replacment`. -/
partial def replace (s pattern replacement : String) : String :=
  loop "" 0 0
where
  loop (acc : String) (accStop pos : String.Pos) :=
    if pos.byteIdx + pattern.endPos.byteIdx > s.endPos.byteIdx then
      acc ++ s.extract accStop s.endPos
    else if s.substrEq pos pattern 0 pattern.endPos.byteIdx then
      loop (acc ++ s.extract accStop pos ++ replacement) (pos + pattern) (pos + pattern)
    else
      loop acc accStop (s.next pos)

end String

namespace Substring

@[inline] def isEmpty (ss : Substring) : Bool :=
  ss.bsize == 0

@[inline] def toString : Substring → String
  | ⟨s, b, e⟩ => s.extract b e

@[inline] def toIterator : Substring → String.Iterator
  | ⟨s, b, _⟩ => ⟨s, b⟩

/-- Return the codepoint at the given offset into the substring. -/
@[inline] def get : Substring → String.Pos → Char
  | ⟨s, b, _⟩, p => s.get (b+p)

/-- Given an offset of a codepoint into the substring,
return the offset there of the next codepoint. -/
@[inline] def next : Substring → String.Pos → String.Pos
  | ⟨s, b, e⟩, p =>
    let absP := b+p
    if absP = e then p else { byteIdx := (s.next absP).byteIdx - b.byteIdx }

/-- Given an offset of a codepoint into the substring,
return the offset there of the previous codepoint. -/
@[inline] def prev : Substring → String.Pos → String.Pos
  | ⟨s, b, _⟩, p =>
    let absP := b+p
    if absP = b then p else { byteIdx := (s.prev absP).byteIdx - b.byteIdx }

def nextn : Substring → Nat → String.Pos → String.Pos
  | _,  0,   p => p
  | ss, i+1, p => ss.nextn i (ss.next p)

def prevn : Substring → Nat → String.Pos → String.Pos
  | _,  0,   p => p
  | ss, i+1, p => ss.prevn i (ss.prev p)

@[inline] def front (s : Substring) : Char :=
  s.get 0

/-- Return the offset into `s` of the first occurence of `c` in `s`,
or `s.bsize` if `c` doesn't occur. -/
@[inline] def posOf (s : Substring) (c : Char) : String.Pos :=
  match s with
  | ⟨s, b, e⟩ => { byteIdx := (String.posOfAux s c e b).byteIdx - b.byteIdx }

@[inline] def drop : Substring → Nat → Substring
  | ss@⟨s, b, e⟩, n => ⟨s, b + ss.nextn n 0, e⟩

@[inline] def dropRight : Substring → Nat → Substring
  | ss@⟨s, b, _⟩, n => ⟨s, b, b + ss.prevn n ⟨ss.bsize⟩⟩

@[inline] def take : Substring → Nat → Substring
  | ss@⟨s, b, _⟩, n => ⟨s, b, b + ss.nextn n 0⟩

@[inline] def takeRight : Substring → Nat → Substring
  | ss@⟨s, b, e⟩, n => ⟨s, b + ss.prevn n ⟨ss.bsize⟩, e⟩

@[inline] def atEnd : Substring → String.Pos → Bool
  | ⟨_, b, e⟩, p => b + p == e

@[inline] def extract : Substring → String.Pos → String.Pos → Substring
  | ⟨s, b, e⟩, b', e' => if b' ≥ e' then ⟨"", 0, 0⟩ else ⟨s, e.min (b+b'), e.min (b+e')⟩

partial def splitOn (s : Substring) (sep : String := " ") : List Substring :=
  if sep == "" then
    [s]
  else
    let rec loop (b i j : String.Pos) (r : List Substring) : List Substring :=
      if i.byteIdx == s.bsize then
        let r := if sep.atEnd j then
          "".toSubstring :: s.extract b (i-j) :: r
        else
          s.extract b i :: r
        r.reverse
      else if s.get i == sep.get j then
        let i := s.next i
        let j := sep.next j
        if sep.atEnd j then
          loop i i 0 (s.extract b (i-j) :: r)
        else
          loop b i j r
      else
        loop b (s.next i) 0 r
    loop 0 0 0 []

@[inline] def foldl {α : Type u} (f : α → Char → α) (init : α) (s : Substring) : α :=
  match s with
  | ⟨s, b, e⟩ => String.foldlAux f s e b init

@[inline] def foldr {α : Type u} (f : Char → α → α) (init : α) (s : Substring) : α :=
  match s with
  | ⟨s, b, e⟩ => String.foldrAux f init s e b

@[inline] def any (s : Substring) (p : Char → Bool) : Bool :=
  match s with
  | ⟨s, b, e⟩ => String.anyAux s e p b

@[inline] def all (s : Substring) (p : Char → Bool) : Bool :=
  !s.any (fun c => !p c)

def contains (s : Substring) (c : Char) : Bool :=
  s.any (fun a => a == c)

@[specialize] private partial def takeWhileAux (s : String) (stopPos : String.Pos) (p : Char → Bool) (i : String.Pos) : String.Pos :=
  if i >= stopPos then i
  else if p (s.get i) then takeWhileAux s stopPos p (s.next i)
  else i

@[inline] def takeWhile : Substring → (Char → Bool) → Substring
  | ⟨s, b, e⟩, p =>
    let e := takeWhileAux s e p b;
    ⟨s, b, e⟩

@[inline] def dropWhile : Substring → (Char → Bool) → Substring
  | ⟨s, b, e⟩, p =>
    let b := takeWhileAux s e p b;
    ⟨s, b, e⟩

@[specialize] private partial def takeRightWhileAux (s : String) (begPos : String.Pos) (p : Char → Bool) (i : String.Pos) : String.Pos :=
  if i == begPos then i
  else
    let i' := s.prev i
    let c  := s.get i'
    if !p c then i
    else takeRightWhileAux s begPos p i'

@[inline] def takeRightWhile : Substring → (Char → Bool) → Substring
  | ⟨s, b, e⟩, p =>
    let b := takeRightWhileAux s b p e
    ⟨s, b, e⟩

@[inline] def dropRightWhile : Substring → (Char → Bool) → Substring
  | ⟨s, b, e⟩, p =>
    let e := takeRightWhileAux s b p e
    ⟨s, b, e⟩

@[inline] def trimLeft (s : Substring) : Substring :=
  s.dropWhile Char.isWhitespace

@[inline] def trimRight (s : Substring) : Substring :=
  s.dropRightWhile Char.isWhitespace

@[inline] def trim : Substring → Substring
  | ⟨s, b, e⟩ =>
    let b := takeWhileAux s e Char.isWhitespace b
    let e := takeRightWhileAux s b Char.isWhitespace e
    ⟨s, b, e⟩

def isNat (s : Substring) : Bool :=
  s.all fun c => c.isDigit

def toNat? (s : Substring) : Option Nat :=
  if s.isNat then
    some <| s.foldl (fun n c => n*10 + (c.toNat - '0'.toNat)) 0
  else
    none

def beq (ss1 ss2 : Substring) : Bool :=
  ss1.bsize == ss2.bsize && ss1.str.substrEq ss1.startPos ss2.str ss2.startPos ss1.bsize

instance hasBeq : BEq Substring := ⟨beq⟩

end Substring

namespace String

def drop (s : String) (n : Nat) : String :=
  (s.toSubstring.drop n).toString

def dropRight (s : String) (n : Nat) : String :=
  (s.toSubstring.dropRight n).toString

def take (s : String) (n : Nat) : String :=
  (s.toSubstring.take n).toString

def takeRight (s : String) (n : Nat) : String :=
  (s.toSubstring.takeRight n).toString

def takeWhile (s : String) (p : Char → Bool) : String :=
  (s.toSubstring.takeWhile p).toString

def dropWhile (s : String) (p : Char → Bool) : String :=
  (s.toSubstring.dropWhile p).toString

def takeRightWhile (s : String) (p : Char → Bool) : String :=
  (s.toSubstring.takeRightWhile p).toString

def dropRightWhile (s : String) (p : Char → Bool) : String :=
  (s.toSubstring.dropRightWhile p).toString

def startsWith (s pre : String) : Bool :=
  s.toSubstring.take pre.length == pre.toSubstring

def endsWith (s post : String) : Bool :=
  s.toSubstring.takeRight post.length == post.toSubstring

def trimRight (s : String) : String :=
  s.toSubstring.trimRight.toString

def trimLeft (s : String) : String :=
  s.toSubstring.trimLeft.toString

def trim (s : String) : String :=
  s.toSubstring.trim.toString

@[inline] def nextWhile (s : String) (p : Char → Bool) (i : String.Pos) : String.Pos :=
  Substring.takeWhileAux s s.endPos p i

@[inline] def nextUntil (s : String) (p : Char → Bool) (i : String.Pos) : String.Pos :=
  nextWhile s (fun c => !p c) i

def toUpper (s : String) : String :=
  s.map Char.toUpper

def toLower (s : String) : String :=
  s.map Char.toLower

def capitalize (s : String) :=
  s.set 0 <| s.get 0 |>.toUpper

def decapitalize (s : String) :=
  s.set 0 <| s.get 0 |>.toLower

end String

protected def Char.toString (c : Char) : String :=
  String.singleton c


import Lean
import PremiseSelection.StatementFeatures

open Lean System

namespace String

/-- Check if a string is contained in another one. -/
partial def isSubstrOf (target str : String) : Bool :=
  loop "" 0
where
  loop (acc : String) (pos : String.Pos) : Bool :=
    if pos.byteIdx + target.endPos.byteIdx > str.endPos.byteIdx then
      false
    else if str.substrEq pos target 0 target.endPos.byteIdx then
      true
    else
      loop acc (str.next pos)

end String

/-- Find file path from module imported from Mathbin. -/
def pathFromMathbinImport (mod : Name) : MetaM (Option FilePath) := do
  let mathbinPath : System.FilePath := "." / "lake-packages" / "mathlib3port"
  SearchPath.findModuleWithExt [mathbinPath] "lean" mod

/-- Find file path from module imported from Mathbin. -/
def pathFromMathlibImport (mod : Name) : MetaM (Option FilePath) := do
  let mathbinPath : System.FilePath := "." / "lake-packages" / "mathlib"
  SearchPath.findModuleWithExt [mathbinPath] "lean" mod

/-- Find file path of JSON with proof sources. -/
def proofSourcePath (mod : Name) : MetaM (Option FilePath) := do
  let mathbinPath : System.FilePath := "." / "data" / "proof_sources"
  SearchPath.findModuleWithExt [mathbinPath] "json" mod

/-- Given a theorem name and a file path, extract the proof text. -/
def proofSource (thm : Name) (json : Json) : MetaM (Option String) := do
  if let Name.str _ thmStr := thm then
    match json.getObjVal? thmStr  with
    | Except.ok (Json.str s) =>
        return some s
    | _ => return none
  return none

namespace PremiseSelection

/-- Given a list of premises and proof text, get rid of the ones that do not
appear. We take into account `ToAdditive` name translations. -/
def filterUserPremises (premises : Multiset Name) (proofSource : String)
  : Multiset Name := Id.run <| do
  let appearsInProof (s : String) : Bool := s.isSubstrOf proofSource
  let mut result := Std.RBMap.empty
  for (p, c) in premises do
    let pLast := (Syntax.splitNameLit p.toString.toSubstring).reverse.head!.toString
    if appearsInProof pLast then
      result := result.insert p c
  return result

/-- Like `filterUserPremises` but simply checks that the premise appears
somewhere in the file, instead of looking for the proof source. -/
def filterUserPremisesFromFile
  (premises : Multiset Name) (referencePath : FilePath)
  : IO (Multiset Name) := do
  let appearsInFile (s : String) : IO Bool := do
    let args := #[s, referencePath.toString]
    let output ← IO.Process.output { cmd := "grep", args := args }
    if output.exitCode != 0 then
      return false
    if output.stdout.isEmpty then
      return false
    return true
  let mut result := Std.RBMap.empty
  for (p, c) in premises do
    if ← appearsInFile p.toString then
      result := result.insert p c
  return result

end PremiseSelection


import complexity_class.lemmas
import data.multiset.functor
import data.finset.functor

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
-/
variables {α : Type*} {β β₁ β₂ γ δ : α → Type*}

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

instance has_uncurry_dep_aux_base : has_uncurry_dep_aux (λ x, (β₁ x → β₂ x)) β₁ β₂ := ⟨λ x y, y⟩

instance has_uncurry_dep_aux_ind [has_uncurry_dep_aux γ β₁ β₂] : has_uncurry_dep_aux (λ x, δ x → γ x) (λ x, δ x × β₁ x) β₂ :=
⟨λ x f p, has_uncurry_dep_aux.uncurry _ (f p.1) p.2⟩

class has_uncurry_dep {α : Type*} (β : α → Type*) {α' : out_param 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) β := ⟨λ f x, f x⟩

instance has_uncurry_dep_base₂ [has_uncurry_dep_aux γ β₁ β₂] :
  has_uncurry_dep (λ x, γ x) (λ z : (Σ x, β₁ x), β₂ z.1) :=
⟨λ f x, (has_uncurry_dep_aux.uncurry _ (f x.1)) x.2⟩

end function

namespace complexity_class
open_locale complexity_class
open tencodable function

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

def mem1_dep (f : ∀ x, β x) : Prop :=
∃ (f' : tree unit → tree unit), C.prop f' ∧ ∀ x : α, f' (encode x) = encode (f x)

def mem_dep [has_uncurry_dep γ' γ] (f : ∀ x, γ' x) (C : complexity_class) : Prop :=
C.mem1_dep ↿ₚf

@[simp] lemma mem_dep_iff₁ (f : α → α₁) :
  mem_dep f C ↔ mem f C := iff.rfl

@[simp] lemma mem_dep_iff₂ (f : α → α₁ → α₂) :
  mem_dep f C ↔ mem f C :=
by split; rintro ⟨f, pf, hf⟩; refine ⟨f, pf, _⟩; rintro ⟨a, b⟩; exact hf ⟨a, b⟩

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

lemma mem_iff_comp_encode_dep {f : ∀ x, β x} :
  f ∈ₐ C ↔ (λ x, encode (f x)) ∈ₑ C := iff.rfl

lemma mem_dep_iff_comp_eq_encode {f : ∀ x, β x} (g : ∀ x, β x → α₁)
  (hg : ∀ (x : α) (y : β x), encode (g x y) = encode y) :
  f ∈ₐ C ↔ (λ x, g x (f x)) ∈ₑ C :=
by { rw [mem_iff_comp_encode, mem_iff_comp_encode_dep], simp only [hg], }

section functor

class is_encodable_functor (F : Type → Type) [functor F] :=
(inst : ∀ x, tencodable x → tencodable (F x))
(map_encode : ∀ {α} [I : tencodable α] (x : F α), (@encode _ (inst _ infer_instance) $ @encode _ I <$> x) = @encode _ (inst _ I) x)

instance : is_encodable_functor list :=
⟨@list.tencodable, λ α _ x, by simp [encode]⟩

lemma list_map_encode (x : list α) : encode (x.map encode) = encode x := is_encodable_functor.map_encode x

instance : is_encodable_functor option :=
⟨@option.tencodable, λ α _ x, by cases x; simp [encode]⟩

lemma option_map_encode (x : option α) : encode (x.map encode) = encode x := is_encodable_functor.map_encode x

lemma _root_.list.sorted_map {α β : Type*} {r : β → β → Prop} (f : α → β) {l : list α} :
  (l.map f).sorted r ↔ l.sorted (λ x y, r (f x) (f y)) := list.pairwise_map f

lemma _root_.list.map_merge_sort {α β : Type*} (r : β → β → Prop) [decidable_rel r]
  [is_total β r] [is_trans β r] [is_antisymm β r] (f : α → β) (l : list α)  :
  (l.map f).merge_sort r = (l.merge_sort (λ x y, r (f x) (f y))).map f :=
begin
  refine list.eq_of_perm_of_sorted (trans ((l.map f).perm_merge_sort r) ((l.perm_merge_sort (λ x y, r (f x) (f y))).map f).symm)
    (list.sorted_merge_sort _ _) ((list.sorted_map f).mpr $ _),
  refine @list.sorted_merge_sort _ _ _ ⟨_⟩ ⟨_⟩ _,
  { intros a b, exact is_total.total (f a) (f b), },
  { intros a b c, exact is_trans.trans (f a) (f b) (f c), }
end

instance : is_encodable_functor multiset :=
⟨@multiset.tencodable, λ α _ x, quotient.induction_on x $ λ x, begin
  dsimp [encode_multiset, lift_le],
  rw [← @list_map_encode α, list.map_merge_sort],
end⟩

lemma multiset_map_encode (x : multiset α) : encode (x.map encode) = encode x := is_encodable_functor.map_encode x

section
open_locale classical

lemma _root_.finset.fmap_def' {α β : Type*} [decidable_eq β] (f : α → β) (x : finset α) :
  f <$> x = x.image f := by { dsimp, congr, }

noncomputable instance : is_encodable_functor finset :=
⟨λ α I, @finset.tencodable _ I _, λ α _ x, begin
  resetI,
  simp only [finset.fmap_def'],
  rw (show x.image encode = _, from (finset.map_eq_image ⟨encode, encode_injective⟩ x).symm),
  cases x with v hv,
  simp only [finset.map, encode_finset],
  exact multiset_map_encode _,
end⟩

lemma finset_image_encode [decidable_eq α] (x : finset α) : encode (x.image encode) = encode x := by convert is_encodable_functor.map_encode x

end

lemma functor_map_dep (F : Type → Type) [functor F] [is_lawful_functor F] [is_encodable_functor F]
  {l : ∀ x, F (β x)} {g : ∀ x, β x → γ x} (hl : by { haveI : ∀ x, tencodable (F (β x)) := λ x, is_encodable_functor.inst _ infer_instance, exact l ∈ₐ C, })
  (hg : g ∈ₐ C) (hm : by { haveI : tencodable (F (tree unit)) := is_encodable_functor.inst _ infer_instance,
    exact ∀ {l' : α → F (tree unit)} {g' : α → tree unit → tree unit}, l' ∈ₑ C → g' ∈ₑ C → (λ z, (g' z) <$> (l' z)) ∈ₑ C, }) :
  by { haveI : ∀ x, tencodable (F (γ x)) := λ x, is_encodable_functor.inst _ infer_instance, exact (λ x, (g x) <$> (l x)) ∈ₐ C } :=
begin
  letI : ∀ x, tencodable (F (β x)) := λ x, is_encodable_functor.inst _ infer_instance,
  letI : tencodable (F (tree unit)) := is_encodable_functor.inst _ infer_instance,
  letI : ∀ x, tencodable (F (γ x)) := λ x, is_encodable_functor.inst _ infer_instance,
  rcases hg with ⟨g', pg, hg⟩,
  simp only [sigma.forall] at hg,
  dsimp [encode_sigma, has_uncurry_dep.uncurry, has_uncurry_dep_aux.uncurry] at hg,
  rw C.mem_dep_iff_comp_eq_encode (λ x (y : F (γ x)), encode <$> y) (λ x y, is_encodable_functor.map_encode y),
  rw C.mem_dep_iff_comp_eq_encode (λ x (y : F (β x)), encode <$> y) (λ x y, is_encodable_functor.map_encode y) at hl,
  have pg' : (λ (x : α) (t : tree unit), g' (tree.node () (encode x) t)) ∈ₑ C := by complexity,
  convert hm hl pg',
  funext x,
  simp [functor.map_map, function.comp, hg],
end

end functor

-- lemma mem_iff_comp_list_map_dep {f : ∀ x, list (β x)} :
--   f ∈ₐ C ↔ (λ x, (f x).map encode) ∈ₑ C :=
-- C.mem_dep_iff_comp_eq_encode (λ x (y : list (β x)), y.map encode) (λ (x : α) y, list_map_encode y)

end dep

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

namespace table_fun
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 to_finmap [fintype α] (f : α [→] β) : @finmap α (λ _, β) := finmap.of_fun f

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 f.to_finmap := rfl

end table_fun

open_locale complexity_class
variables {C : complexity_class} {α : Type} {β γ : α → Type} [tencodable α]
  [∀ i, tencodable (β i)] [∀ i, tencodable (γ i)] [∀ i, decidable_eq (β i)] [∀ i, fintype (β i)]

lemma to_finmap_iff {f : ∀ x, (β x) [→] (γ x)} :
  f ∈ₐ C ↔ (λ x, (f x).to_finmap) ∈ₐ C := iff.rfl

@[complexity] lemma to_finmap {f : ∀ x, (β x) [→] (γ x)} (hf : f ∈ₐ C) : (λ x, (f x).to_finmap) ∈ₐ C := hf

end complexity_class


-- This example used to produce
-- `error: (kernel) declaration has metavariables`
-- Because we were not registering postponed instance metavariables
def f {α : Type u} (a : α) : α :=
  let g a b := a + b
  g a a


import ReactorModel.Determinism.State

open Classical ReactorType Indexable

namespace Execution
namespace Instantaneous
namespace Step

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

theorem rcn_not_mem_progress (e : s₁ ⇓ᵢ s₂) : e.rcn ∉ s₁.progress :=
  sorry -- e.allows.unprocessed

theorem preserves_tag (e : s₁ ⇓ᵢ s₂) : s₁.tag = s₂.tag :=
  sorry -- e.exec.preserves_tag

theorem mem_progress_iff :
  (e : s₁ ⇓ᵢ s₂) → (rcn' ∈ s₂.progress ↔ rcn' = e.rcn ∨ rcn' ∈ s₁.progress) := by
  intro h
  constructor
  case mp =>
    intro ho
    by_cases hc : rcn' = h.rcn
    case pos => exact .inl hc
    case neg => sorry
      -- rw [State.progress, h.exec.ctx_adds_rcn, ←rcn] at ho
      -- simp [Context.mem_record_progress_iff _ _ _ |>.mp ho]
  case mpr =>
    intro ho
    by_cases hc : rcn' = h.rcn
    case pos =>
      simp [hc]
      sorry
      -- exact Context.mem_record_progress_iff _ _ _ |>.mpr (.inl rfl)
    case neg =>
      sorry
      -- simp [State.progress, h.exec.ctx_adds_rcn, Context.mem_record_progress_iff _ _ _ |>.mpr (.inr $ ho.resolve_left hc)]

-- Corollary of `InstStep.mem_progress_iff`.
theorem not_mem_progress :
  (e : s₁ ⇓ᵢ s₂) → (rcn' ≠ e.rcn) → rcn' ∉ s₁.progress → rcn' ∉ s₂.progress :=
  sorry -- λ h hn hm => (mt h.mem_progress.mp) $ not_or.mpr ⟨hn, hm⟩

-- Corollary of `InstStep.mem_progress`.
theorem monotonic_progress : (s₁ ⇓ᵢ s₂) → rcn' ∈ s₁.progress → rcn' ∈ s₂.progress :=
  sorry -- (·.mem_progress_iff.mpr $ .inr ·)

-- Corollary of `InstStep.mem_progress`.
theorem rcn_mem_progress : (e : s₁ ⇓ᵢ s₂) → e.rcn ∈ s₂.progress :=
  (·.mem_progress_iff.mpr $ .inl rfl)

theorem Skip.equiv : (s₁ ⇓ₛ s₂) → s₁.rtr ≈ s₂.rtr
  | mk .. => .refl

theorem Skip.progress_eq : (e : s₁ ⇓ₛ s₂) → s₂.progress = s₁.progress.insert e.rcn
  | mk .. => rfl

theorem Skip.progress_mono (e : s₁ ⇓ₛ s₂) : s₁.progress ⊆ s₂.progress := by
  simp [e.progress_eq]
  apply Set.subset_insert

theorem Skip.triggers_iff (e : s₁ ⇓ₛ s₂) : (s₁.Triggers i) ↔ (s₂.Triggers i) := by
  cases e
  case mk rcn _ _ =>
    constructor
    all_goals
      intro h
      sorry -- exact h.progress_agnostic (s₁.record_preserves_rtr rcn) (s₁.record_preserves_tag rcn) |>.choose_spec

theorem Skip.mem_progress_iff (e : s₁ ⇓ₛ s₂) :
    (rcn' ∈ s₂.progress) ↔ (rcn' = e.rcn ∨ rcn' ∈ s₁.progress) := by
  sorry

theorem Skip.preserves_allows_indep (e₁ : s₁ ⇓ₛ s₂) (e₂ : s₂ ⇓ₛ s₃) (h : e₁.rcn ≮[s₁.rtr] e₂.rcn) :
    s₁.Allows e₂.rcn where
  mem := sorry
  unprocessed := Set.not_mem_subset e₁.progress_mono $ e₂.allows_rcn.unprocessed
  deps := by
    intro i hi
    have h' := equiv_eq_dependencies e₁.equiv |>.symm ▸ e₂.allows_rcn.deps
    refine e₁.mem_progress_iff.mp (h' hi) |>.resolve_left ?_
    intro hc; subst hc; contradiction

set_option pp.proofs.withType false
theorem Skip.swap_indep_skip (e₁ : s₁ ⇓ₛ s₂) (e₂ : s₂ ⇓ₛ s₃) (h : e₁.rcn ≮[s₁.rtr] e₂.rcn) :
    ∃ (s₂' : _) (f₁ : s₁ ⇓ₛ s₂') (f₂ : s₂' ⇓ₛ s₃), (f₁.rcn = e₂.rcn) ∧ (f₂.rcn = e₁.rcn) := by
  have ha := e₁.preserves_allows_indep e₂ h
  have ht := e₁.triggers_iff.not.mpr e₂.not_triggers
  have e₁' := Skip.mk ha ht
  simp at e₁'
  exists _, e₁'
  sorry -- TODO: This is super annoying. Is there a better way to approach this?
        --       Do we need more preservation theorems for `Allows` and `Triggers` first?

theorem Skip.swap_indep_exec (e₁ : s₁ ⇓ₛ s₂) (e₂ : s₂ ⇓ₑ s₃) (h : e₁.rcn ≮[s₁.rtr] e₂.rcn) :
    ∃ (s₂' : _) (f₁ : s₁ ⇓ₑ s₂') (f₂ : s₂' ⇓ₛ s₃), (f₁.rcn = e₂.rcn) ∧ (f₂.rcn = e₁.rcn) := by
  sorry

theorem Skip.swap_indep (e₁ : s₁ ⇓ₛ s₂) (e₂ : s₂ ⇓ᵢ s₃) (h : e₁.rcn ≮[s₁.rtr] e₂.rcn) :
    ∃ (s₂' : _) (f₁ : s₁ ⇓ᵢ s₂') (f₂ : s₂' ⇓ₛ s₃), (f₁.rcn = e₂.rcn) ∧ (f₂.rcn = e₁.rcn) := by
  cases e₂
  case skip e₂ =>
    have ⟨_, e₁', e₂', _⟩ := e₁.swap_indep_skip e₂ h
    exists _, skip e₁', e₂'
  case exec e₂ =>
    have ⟨_, e₁', e₂', _⟩ := e₁.swap_indep_exec e₂ h
    exists _, exec e₁', e₂'

theorem Exec.equiv : (s₁ ⇓ₑ s₂) → s₁.rtr ≈ s₂.rtr
  | mk .. => by simp [State.record_preserves_rtr, s₁.exec_equiv]

theorem Exec.progress_eq : (e : s₁ ⇓ₑ s₂) → s₂.progress = s₁.progress.insert e.rcn
  | mk .. => rfl

theorem not_Closed (e : s₁ ⇓ᵢ s₂) : ¬s₁.Closed := by
  intro c
  have h := c ▸ e.allows_rcn.unprocessed
  simp [Partial.mem_iff] at h
  sorry -- have := h e.allows.mem.choose
  -- contradiction

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

theorem deterministic (e₁ : s ⇓ᵢ s₁) (e₂ : s ⇓ᵢ s₂) (h : e₁.rcn = e₂.rcn) : s₁ = s₂ := by
  cases e₁ <;> cases e₂
  all_goals
    case _ e₁ e₂ =>
      cases e₁; cases e₂
      simp [rcn] at h
      subst h
      first | rfl | contradiction

theorem acyclic (e : s₁ ⇓ᵢ s₂) : e.rcn ≮[s₁.rtr] e.rcn :=
  e.allows_rcn.acyclic

theorem progress_eq : (e : s₁ ⇓ᵢ s₂) → s₂.progress = s₁.progress.insert e.rcn
  | skip e | exec e => e.progress_eq

/-
By cases on e₁ and e₂:
(1) skip.skip:
    then both didn't trigger and nothing about the reactors changed, so it is easy to show that
    switching the order preserves non-triggering.

(2)&(3) exec.skip and skip.exec:
        using `preserves_triggers` it should be easy

(4) exec.exec:
    reduces to a theorem on State.exec?
-/
theorem prepend_indep (e₁ : s₁ ⇓ᵢ s₂) (e₂ : s₂ ⇓ᵢ s₃) (h : e₁.rcn ≮[s₁.rtr] e₂.rcn) :
    ∃ (s₂' : _) (e₁' : s₁ ⇓ᵢ s₂') (e₂' : s₂' ⇓ᵢ s₃), e₁'.rcn = e₂.rcn ∧ e₂'.rcn = e₁.rcn := by
  cases e₁
  case skip e₁ =>
    have ⟨_, e₁', e₂', _⟩ := e₁.swap_indep e₂ h
    exists _, e₁', skip e₂'
  case exec e₁ =>
    sorry

end Step
end Instantaneous
end Execution


import Alloy.C
open scoped Alloy.C

/-!
# `S.lean`

An adaption of Lean 4's ['foreign'][1] example for Alloy.

[1]: https://github.com/leanprover/lean4/tree/b278a20ac22adcbfde11db386f2dc874d4a215ad/tests/compiler/foreign
-/

alloy c include <stdint.h> <stdlib.h> <string.h> <lean/lean.h>

--------------------------------------------------------------------------------
/-! ## C Definition of S                                                      -/
--------------------------------------------------------------------------------

alloy c section

typedef struct {
  uint32_t      m_x;
  uint32_t      m_y;
  lean_object * m_s;
} S;

static void S_finalize(void* ptr) {
  lean_dec(((S*)ptr)->m_s);
  free(ptr);
}

static void S_foreach(void* ptr, b_lean_obj_arg f) {
  lean_apply_1(f, ((S*)ptr)->m_s);
}

static lean_external_class * g_S_class = NULL;

static inline lean_object * S_to_lean(S* s) {
  if (g_S_class == NULL) {
    g_S_class = lean_register_external_class(S_finalize, S_foreach);
  }
  return lean_alloc_external(g_S_class, s);
}

static inline S const * to_S(b_lean_obj_arg s) {
  return (S*)(lean_get_external_data(s));
}

static S g_s = {0, 0, NULL};

end

--------------------------------------------------------------------------------
/-! ## Lean Interface                                                         -/
--------------------------------------------------------------------------------

opaque S.nonemptyType : NonemptyType
def S : Type := S.nonemptyType.type
instance : Nonempty S := S.nonemptyType.property

alloy c extern "lean_mk_S"
def mkS (x y : UInt32) (string : String) : S := {
  S* s = malloc(sizeof(S));
  s->m_x = x;
  s->m_y = y;
  s->m_s = string;
  return S_to_lean(s);
}

alloy c extern "lean_S_add_x_y"
def S.addXY (s : @& S) : UInt32 := {
  return to_S(s)->m_x + to_S(s)->m_y;
}

alloy c extern "lean_S_string"
def S.string (s : @& S) : String := {
  lean_inc(to_S(s)->m_s);
  return to_S(s)->m_s;
}

alloy c extern "lean_S_global_append"
def appendToGlobalS (string : String) : BaseIO PUnit := {
  if (g_s.m_s == NULL) {
    g_s.m_s = string;
  } else {
    g_s.m_s = lean_string_append(g_s.m_s, string);
  }
  return lean_io_result_mk_ok(lean_box(0));
}

alloy c extern "lean_S_global_string"
def getGlobalString : BaseIO String := {
  if (g_s.m_s == NULL) {
    g_s.m_s = lean_mk_string("");
  }
  lean_inc(g_s.m_s);
  return lean_io_result_mk_ok(g_s.m_s);
}

alloy c extern "lean_S_update_global"
def updateGlobalS (s : @& S) : BaseIO Unit := {
  if (g_s.m_s != NULL) {
    lean_dec(g_s.m_s);
  }
  lean_inc(to_S(s)->m_s);
  g_s.m_x = to_S(s)->m_x;
  g_s.m_y = to_S(s)->m_y;
  g_s.m_s = to_S(s)->m_s;
  return lean_io_result_mk_ok(lean_box(0));
}


-- Copyright (c) 2018 Scott Morrison. All rights reserved.
-- Released under Apache 2.0 license as described in the file LICENSE.
-- Authors: Keeley Hoek, Scott Morrison

import .init
import .discovery

import tidy.lib.list
import tidy.lib.expr

open tactic

variables {α β γ δ : Type}

namespace tidy.rewrite_search

meta def try_search (cfg : rewrite_search_config α β γ δ) (prog : discovery.progress) (rs : list (expr × bool)) (eqn : sided_pair expr) : tactic (option string) := do
  i ← try_mk_search_instance cfg prog rs eqn,
  match i with
  | none := return none
  | some i := do
    (i, result) ← i.search_until_solved,
    match result with
    | search_result.failure reason := fail reason
    | search_result.success proof steps := do
      exact proof,
      some <$> i.explain proof steps
    end
  end

meta def rewrite_search_pair (cfg : rewrite_search_config α β γ δ) (prog : discovery.progress) (rs : list (expr × bool)) (eqn : sided_pair expr) := do
  result ← try_search cfg prog rs eqn,
  match result with
  | some str := return str
  | none := do
    trace "\nError initialising rewrite_search instance, falling back to emergency config!",
    result ← try_search (mk_fallback_config cfg) prog rs eqn,
    match result with
    | some str := return str
    | none := fail "Could not initialise emergency rewrite_search instance!"
    end
  end

-- TODO If try_search fails due to a failure to init any of the tracer, metric, or strategy we try again
-- using the "fallback" default versions of all three of these. Instead we could be more thoughtful,
-- and try again only replacing the failing one of these with its respective fallback module version.

meta def collect_rw_lemmas (cfg : rewrite_search_config α β γ δ) (use_suggest_annotations : bool) (per : discovery.persistence) (extra_names : list name) (extra_rws : list (expr × bool)) : tactic (discovery.progress × list (expr × bool)) := do
  let per := if cfg.help_me then discovery.persistence.try_everything else per,
  (prog, rws) ← discovery.collect use_suggest_annotations per cfg.suggest extra_names,
  hyp_rws ← discovery.rewrite_list_from_hyps,
  let rws := rws ++ extra_rws ++ hyp_rws,

  locs ← local_context,
  rws ← rws.mmap $ discovery.inflate_rw locs,
  return (prog, rws.join)

meta def rewrite_search_target (cfg : rewrite_search_config α β γ δ) (use_suggest_annotations : bool) (per : discovery.persistence) (extra_names : list name) (extra_rws : list (expr × bool)) : tactic string := do
  t ← target,
  if t.has_meta_var then
    fail "rewrite_search is not suitable for goals containing metavariables"
  else skip,

  (prog, rws) ← collect_rw_lemmas cfg use_suggest_annotations per extra_names extra_rws,

  if cfg.trace_rules then
    do rs_strings ← pp_rules rws,
      trace ("rewrite_search using:\n---\n" ++ (string.intercalate "\n" rs_strings) ++ "\n---")
  else skip,

  match t with
  | `(%%lhs = %%rhs) := rewrite_search_pair cfg prog rws ⟨lhs, rhs⟩
  | `(%%lhs ↔ %%rhs) := rewrite_search_pair cfg prog rws ⟨lhs, rhs⟩
  | _                := fail "target is not an equation or iff"
  end

private meta def add_simps : simp_lemmas → list name → tactic simp_lemmas
| s []      := return s
| s (n::ns) := do s' ← s.add_simp n, add_simps s' ns

private meta def add_expr (s : simp_lemmas) (u : list name) (e : expr) : tactic (simp_lemmas × list name) :=
do
  let e := e.erase_annotations,
  match e with
  | expr.const n _           :=
    (do b ← is_valid_simp_lemma_cnst n, guard b, s ← s.add_simp n, return (s, u))
    <|>
    (do eqns ← get_eqn_lemmas_for tt n, guard (eqns.length > 0), s ← add_simps s eqns, return (s, u))
    <|>
    (do env ← get_env, guard (env.is_projection n).is_some, return (s, n::u))
    <|>
    fail n
  | _ :=
    (do b ← is_valid_simp_lemma e, guard b, s ← s.add e, return (s, u))
    <|>
    fail e
  end

meta def simp_search_target (cfg : rewrite_search_config α β γ δ) (use_suggest_annotations : bool) (per : discovery.persistence) (extra_names : list name) (extra_rws : list (expr × bool)) : tactic unit := do
  t ← target,

  (prog, rws) ← collect_rw_lemmas cfg use_suggest_annotations per extra_names extra_rws,

  if cfg.trace_rules then
    do rs_strings ← pp_rules rws,
      trace ("simp_search using:\n---\n" ++ (string.intercalate "\n" rs_strings) ++ "\n---")
  else skip,

  (s, to_unfold) ← mk_simp_set ff [] [] >>= λ sset, rws.mfoldl (λ c e, add_expr c.1 c.2 e.1 <|> return c) sset,
  (n, pf) ← simplify s to_unfold t {contextual := tt} `eq failed,
  replace_target n pf >> try tactic.triv >> try (tactic.reflexivity reducible)

end tidy.rewrite_search

namespace tactic.interactive

open interactive
open tidy.rewrite_search tidy.rewrite_search.discovery.persistence

meta def rewrite_search (cfg : rewrite_search_config α β γ δ . pick_default_config) : tactic string :=
  rewrite_search_target cfg tt try_everything [] []

meta def rewrite_search_with (rs : parse rw_rules) (cfg : rewrite_search_config α β γ δ . pick_default_config) : tactic string := do
  extra_rws ← discovery.rewrite_list_from_rw_rules rs.rules,
  rewrite_search_target cfg tt try_everything [] extra_rws

meta def rewrite_search_using (as : list name) (cfg : rewrite_search_config α β γ δ . pick_default_config) : tactic string := do
  extra_names ← discovery.load_attr_list as,
  rewrite_search_target cfg ff try_bundles extra_names []

-- @Scott should we still do this?
--  exprs ← close_under_apps exprs, -- TODO don't do this for everything, it's too expensive: only for specially marked lemmas

meta def simp_search (cfg : rewrite_search_config α β γ δ . pick_default_config) : tactic unit := do
  simp_search_target cfg tt try_everything [] []

meta def simp_search_with (rs : parse rw_rules) (cfg : rewrite_search_config α β γ δ . pick_default_config) : tactic unit := do
  extra_rws ← discovery.rewrite_list_from_rw_rules rs.rules,
  simp_search_target cfg tt try_everything [] extra_rws

end tactic.interactive


-- Copyright (c) 2017 Scott Morrison. All rights reserved.
-- Released under Apache 2.0 license as described in the file LICENSE.
-- Authors: Scott Morrison

import tidy.tidy

open tactic
set_option trace.tidy true

@[reducible] def {u} auto_cast {α β : Sort u} {h : α = β} (a : α) := cast h a
@[simp] lemma {u} auto_cast_identity {α : Sort u} (p : α = α) (a : α) : @auto_cast α α p a = a :=
begin unfold auto_cast, unfold cast, end
notation `⟬` p `⟭` := @auto_cast _ _ (by obviously) p


import polytime.lemmas
import complexity_class.dependent

namespace polytime

open_locale complexity_class
open tencodable function polysize

variables {α β γ δ ε η : Type} [tencodable α] [tencodable β] [tencodable γ]
 [tencodable δ] [tencodable ε] [tencodable η]

@[complexity] lemma list_len : (@list.length α) ∈ₑ PTIME :=
begin
  complexity using λ x, (encode x).stack_rec (λ _ : unit, 0) (λ _ _ _, ()) (λ _ _ _, ())
    (λ _ ih _ _ _, ih + 1) (),
  induction x with hd tl ih, { refl, },
  simpa [encode_cons],
end

@[simp] def scanl_step {α β : Type*} (f : β → α → β) : list α × list β → list α × list β
| ((x :: xs), (y :: ys)) := (xs, f y x :: y :: ys)
| x := x

@[complexity] theorem scanl_step_polytime {f : γ → β → α → β} {st : γ → list α × list β} (hf : f ∈ₑ PTIME) (hst : st ∈ₑ PTIME) :
  (λ x, scanl_step (f x) (st x)) ∈ₑ PTIME :=
by { delta scanl_step, clean_target, complexity, }

theorem scanl_step_iterate' {α : Type*} (f : β → α → β) (l : list α) (x : β) (n : ℕ) :
  (scanl_step f)^[n] (l, [x]) = (l.drop n, ((l.scanl f x).take (n + 1)).reverse) :=
begin
  suffices : ∀ xs, (scanl_step f)^[n] (l, x :: xs) = (l.drop n, ((l.scanl f x).take (n + 1)).reverse ++ xs),
  { simpa using this [], },
  intro xs,
  induction l with hd tl ih generalizing x xs n, { rw iterate_fixed; simp, },
  cases n, { simp, }, { simp [ih], },
end

theorem scanl_step_iterate {α : Type*} (f : β → α → β) (l : list α) (x : β) :
  (scanl_step f)^[l.length] (l, [x]) = ([], (l.scanl f x).reverse) :=
by { rw [scanl_step_iterate', ← @list.length_scanl _ _ f x l], simp, }

lemma list_scanl_rev' [polysize α] [polysize β] [polysize γ] {lst : γ → list α} {acc : γ → β} {f : γ → β → α → β}
  (hlst : lst ∈ₑ PTIME) (hacc : acc ∈ₑ PTIME) (hf : f ∈ₑ PTIME)
  (hf' : polysize_fun (λ (ls : list α) (x), ls.foldl (f x) (acc x))) :
  polytime.mem (λ x : γ, ((lst x).scanl (f x) (acc x)).reverse) :=
begin
  convert_to polytime.mem (λ x, ((scanl_step (f x))^[(lst x).length] (lst x, [acc x])).2),
  { simp [scanl_step_iterate], },
  refine complexity_class.mem.snd.comp _,
  apply iterate, complexity,
  cases hf' with pf hpf, cases polytime.size_le hlst with pl hpl,
  use pl + (pl + 1) * (pf.comp (pl + polynomial.X) + 1),
  intros x m _,
  simp [scanl_step_iterate'], apply add_le_add,
  { exact (list.size_le_of_sublist ((lst x).drop_sublist _)).trans (hpl _), },
  refine (list.size_le_of_sublist (list.take_sublist _ _)).trans _,
  apply list.size_le_mul_of_le,
  { rw [list.length_scanl, add_le_add_iff_right], exact (lst x).length_le_size.trans (hpl _), },
  simp_rw list.mem_iff_nth_le,
  rintros e ⟨n, hn, rfl⟩,
  rw list.scanl_nth_le_eq_foldl,
  exact (hpf ((lst x).take n, x)).trans (pf.eval_mono $ add_le_add_right ((list.size_le_of_sublist (list.take_sublist _ _)).trans (hpl _)) _),
end

lemma list_scanl_rev [polysize α] [polysize β] [polysize γ] {lst : γ → list α} {acc : γ → β} {f : γ → β → α → β}
  (hlst : lst ∈ₑ PTIME) (hacc : acc ∈ₑ PTIME) (hf : f ∈ₑ PTIME)
  (hf' : polysize_safe (λ (usf : γ × α) (sf : β), f usf.1 sf usf.2)) :
  polytime.mem (λ x : γ, ((lst x).scanl (f x) (acc x)).reverse) :=
list_scanl_rev' hlst hacc hf
  (polysize_safe.foldl polysize_fun.fst ((polytime.size_le hacc).comp polysize_fun.snd)
    (by { apply hf'.comp₃_1, complexity, }))

lemma list_foldl' [polysize α] [polysize β] [polysize γ] {lst : γ → list α} {acc : γ → β} {f : γ → β → α → β}
  (hlst : lst ∈ₑ PTIME) (hacc : acc ∈ₑ PTIME) (hf : f ∈ₑ PTIME)
  (hf' : polysize_fun (λ (ls : list α) (x), ls.foldl (f x) (acc x))) :
  polytime.mem (λ x : γ, (lst x).foldl (f x) (acc x)) :=
begin
  rw complexity_class.of_some,
  convert complexity_class.head'.comp (list_scanl_rev' hlst hacc hf hf'),
  ext x : 1, simp [list.scanl_last_eq_foldl],
end

@[complexity] lemma list_foldl [polysize α] [polysize β] [polysize γ] {lst : γ → list α} {acc : γ → β} {f : γ → β → α → β}
  (hlst : lst ∈ₑ PTIME) (hacc : acc ∈ₑ PTIME) (hf : f ∈ₑ PTIME)
  (hf' : polysize_safe (λ (usf : γ × α) (sf : β), f usf.1 sf usf.2)) :
  polytime.mem (λ x : γ, (lst x).foldl (f x) (acc x)) :=
begin
  rw complexity_class.of_some,
  convert complexity_class.head'.comp (list_scanl_rev hlst hacc hf hf'),
  ext x : 1, simp [list.scanl_last_eq_foldl],
end

@[complexity] theorem list_reverse : (@list.reverse α) ∈ₑ PTIME :=
begin
  complexity using (λ l : list α, l.foldl (λ (acc : list α) (hd : α), hd :: acc) []),
  rw [← list.foldr_reverse, list.foldr_eta],
end

@[complexity] lemma list_scanl [polysize α] [polysize β] [polysize γ] {lst : γ → list α} {acc : γ → β} {f : γ → β → α → β}
  (hlst : lst ∈ₑ PTIME) (hacc : acc ∈ₑ PTIME) (hf : f ∈ₑ PTIME)
  (hf' : polysize_safe (λ (usf : γ × α) (sf : β), f usf.1 sf usf.2)) :
  polytime.mem (λ x : γ, ((lst x).scanl (f x) (acc x))) :=
by simpa using list_reverse.comp (list_scanl_rev hlst hacc hf hf')

@[complexity] theorem list_foldr [polysize α] [polysize β] [polysize γ] {lst : γ → list α} {acc : γ → β} {f : γ → α → β → β}
  (hlst : lst ∈ₑ PTIME) (hacc : acc ∈ₑ PTIME) (hf : f ∈ₑ PTIME)
  (hf' : polysize_safe (λ (usf : γ × α) (sf : β), f usf.1 usf.2 sf)) :
  polytime.mem (λ x : γ, (lst x).foldr (f x) (acc x)) :=
by { simp_rw ← list.foldl_reverse, complexity, }

@[complexity] theorem list_map {lst : γ → list α} {f : γ → α → β} (hlst : lst ∈ₑ PTIME) (hf : f ∈ₑ PTIME) :
 (λ x, (lst x).map (f x)) ∈ₑ PTIME :=
by { complexity using (λ x, (lst x).foldr (λ hd acc, (f x hd) :: acc) []), induction lst x; simp [*], }

@[complexity] theorem list_all_some : (@list.all_some α) ∈ₑ PTIME :=
begin
  complexity using λ l, l.foldr (λ hd' acc', hd'.bind (λ hd, acc'.map (λ acc, hd :: acc))) (some []),
  induction l with hd, { simp, }, cases hd; simp [*],
end

instance {α : Type*} [polycodable α] : polycodable (list α) :=
{ poly := by { dunfold decode, complexity, } }

@[complexity] theorem list_tails : (@list.tails α) ∈ₑ PTIME :=
by { complexity using λ ls, ls.scanl (λ l _, l.tail) ls, induction ls; simp [*], }

@[complexity] theorem list_init : (@list.init α) ∈ₑ PTIME :=
by { complexity using λ ls, ls.reverse.tail.reverse, induction ls using list.reverse_rec_on; simp [list.init], }

theorem stack_rec_eq [inhabited β] [inhabited γ] (ls : list γ) (base : α → β) (pre : γ → list γ → α → α)
  (post : β → γ → list γ → α → β) (arg : α) : ls.stack_rec base pre post arg =
  (ls.tails.init.scanl (λ (acc : list γ × α) (x : list γ), (x.tail, pre x.head x.tail acc.2)) (ls, arg))
    .foldr (λ ls_arg ih, if ls_arg.1.empty then base ls_arg.2 else post ih ls_arg.1.head ls_arg.1.tail ls_arg.2) (arbitrary _) :=
begin
  induction ls with hd tl ih generalizing arg, { simp [list.init], },
  simp [list.init_cons_of_ne_nil (list.ne_nil_of_length_eq_succ $ list.length_tails _), ih],
end

@[complexity] theorem list_stack_rec [polysize α] [polysize β] [polysize γ] [polysize δ] {ls : δ → list γ} {base : δ → α → β} {pre : δ → γ → list γ → α → α}
  {post : δ → β → γ → list γ → α → β} {arg : δ → α} (hls : ls ∈ₑ PTIME) (hb : base ∈ₑ PTIME)
  (hpre : pre ∈ₑ PTIME) (hpost : post ∈ₑ PTIME) (harg : arg ∈ₑ PTIME)
  (hpre' : polysize_safe (λ (usf : δ × γ × list γ) (sf : α), pre usf.1 usf.2.1 usf.2.2 sf))
  (hpost' : polysize_safe (λ (usf : δ × γ × list γ × α) (sf : β), post usf.1 sf usf.2.1 usf.2.2.1 usf.2.2.2)) :
  (λ x, (ls x).stack_rec (base x) (pre x) (post x) (arg x)) ∈ₑ PTIME :=
begin
  casesI is_empty_or_nonempty δ, { exact complexity_class.of_from_fintype' _, },
  casesI is_empty_or_nonempty γ, { complexity using (λ x, base x (arg x)), cases ls x with hd, { refl, }, exact is_empty.elim' infer_instance hd, },
  inhabit δ, inhabit γ, haveI : inhabited α := ⟨arg default⟩, haveI : inhabited β := ⟨base default default⟩,
  simp_rw stack_rec_eq,
  complexity, { apply polysize_safe.comp₄_3, complexity, },
  { apply polysize_safe.comp₅_1, complexity, },
end

@[complexity] lemma repeat : (@list.repeat α) ∈ₑ PTIME :=
by { complexity using λ x n, (list.cons x)^[n] [], induction n; simp [iterate_succ', *], }

@[complexity] lemma nat_stack_rec {n : γ → ℕ} {base : γ → α → β} {pre : γ → ℕ → α → α} {post : γ → β → ℕ → α → β}
  {arg : γ → α}  (hn : n ∈ₑ PTIME) (hb : base ∈ₑ PTIME) (hpr : pre ∈ₑ PTIME) (hpo : post ∈ₑ PTIME)
  (harg : arg ∈ₑ PTIME) (hpr : polysize_safe (λ (usf : γ × ℕ) (sf : α), pre usf.1 usf.2 sf))
  (hpo' : polysize_safe (λ (usf : γ × ℕ × α) (sf : β), post usf.1 sf usf.2.1 usf.2.2)) : (λ x, (n x).stack_rec (base x) (pre x) (post x) (arg x)) ∈ₑ PTIME :=
begin
  complexity using λ x, (list.repeat () $ n x).stack_rec (base x) (λ _ tl y, pre x tl.length y)
    (λ ih _ tl y, post x ih tl.length y) (arg x),
  { apply polysize_safe.comp₃_2, complexity, },
  { apply polysize_safe.comp₄_1, complexity, },
  generalize : arg x = y, induction n x with n ih generalizing y,
  { simp, }, { simp [ih], },
end

@[complexity] lemma unary_nat_sum : (@list.sum ℕ _ _) ∈ₑ PTIME :=
by { delta list.sum, complexity, }

@[complexity] lemma list_ordered_insert {r : γ → α → α → Prop} [∀ x, decidable_rel (r x)] {a : γ → α} {ls : γ → list α} (hr : r ∈ₚ PTIME)
  (he : a ∈ₑ PTIME) (hls : ls ∈ₑ PTIME) : (λ x, (ls x).ordered_insert (r x) (a x)) ∈ₑ PTIME :=
begin
  complexity using λ x, (ls x).stack_rec (λ _ : unit, [a x]) (λ _ _ _, ())
    (λ ih b l _, if r x (a x) b then a x :: b :: l else b :: ih) (),
  induction ls x; simp [*],
end

@[complexity] lemma list_insertion_sort {r : γ → α → α → Prop} [∀ x, decidable_rel (r x)] {ls : γ → list α} (hr : r ∈ₚ PTIME)
  (hls : ls ∈ₑ PTIME) : (λ x, (ls x).insertion_sort (r x)) ∈ₑ PTIME :=
by { complexity using λ x, (ls x).foldr (λ b ih, list.ordered_insert (r x) b ih) [], induction ls x; simp [*], }

@[complexity] lemma list_append : ((++) : list α → list α → list α) ∈ₑ PTIME :=
by { complexity using λ l₁ l₂, l₁.foldr (λ hd acc, hd :: acc) l₂, induction l₁; simp [*], }

@[complexity] lemma list_drop : @list.drop α ∈ₑ PTIME :=
by { complexity using λ n l, list.tail^[n] l, simp, }

@[complexity] lemma list_any {l : α → list β} {p : α → β → bool} (hl : l ∈ₑ PTIME) (hp : p ∈ₑ PTIME) :
  (λ x, (l x).any (p x)) ∈ₑ PTIME :=
by { delta list.any, complexity, use 0, simp, }

@[complexity] lemma list_all {l : α → list β} {p : α → β → bool} (hl : l ∈ₑ PTIME) (hp : p ∈ₑ PTIME) :
  (λ x, (l x).all (p x)) ∈ₑ PTIME :=
by { delta list.all, complexity, use 0, simp, }

@[complexity] lemma list_bex {l : α → list β} {p : α → β → Prop} [∀ x y, decidable (p x y)] (hl : l ∈ₑ PTIME) (hp : p ∈ₚ PTIME) :
  (λ x, ∃ e ∈ l x, p x e) ∈ₚ PTIME :=
by { simp_rw [← list.any_iff_exists_prop], complexity, }

@[complexity] lemma list_ball {l : α → list β} {p : α → β → Prop} [∀ x y, decidable (p x y)] (hl : l ∈ₑ PTIME) (hp : p ∈ₚ PTIME) :
  (λ x, ∀ e ∈ l x, p x e) ∈ₚ PTIME :=
by { simp_rw ← list.all_iff_forall_prop, complexity, }

@[complexity] lemma list_mem : ((∈) : α → list α → Prop) ∈ₚ PTIME :=
by { haveI : decidable_eq α := decidable_eq_of_encodable _, complexity using λ x l, ∃ z ∈ l, z = x, simp, }

@[complexity] lemma pairwise {l : α → list β} {r : α → β → β → Prop} (hl : l ∈ₑ PTIME)
  (hr : r ∈ₚ PTIME) : (λ x, (l x).pairwise (r x)) ∈ₚ PTIME :=
begin
  classical, rw ← complexity_class.mem_iff_mem_pred,
  complexity using λ x, (l x).stack_rec (λ _ : unit, tt) (λ _ _ _, ())
    (λ ih hd tl _, (∀ a' ∈ tl, r x hd a') && ih) (), { use 0, simp, },
  induction l x; simp [*],
end

@[complexity] lemma list_nodup : (@list.nodup α) ∈ₚ PTIME :=
by { dunfold list.nodup, complexity, }

@[complexity] lemma list_nth {α : Type} [tencodable α] : @list.nth α ∈ₑ PTIME :=
by { complexity using λ l n, (l.drop n).head', rw [← list.nth_zero, list.nth_drop], refl, }

section dep
open_locale tree
variables {σ τ : α → Type} [∀ i, tencodable (σ i)] [∀ i, tencodable (τ i)]

lemma list_map_dep {f : ∀ x, list (σ x)} {g : ∀ x, σ x → τ x} (hf : f ∈ₐ PTIME) (hg : g ∈ₐ PTIME) :
  (λ x, @list.map _ (τ x) (g x) (f x)) ∈ₐ PTIME :=
polytime.functor_map_dep list hf hg (@list_map _ _ _ _ _ _)

end dep

end polytime


import defs.dynamics
import defs.fv
import defs.statics
import lemmas.cx
import lemmas.fv

theorem if_subst (ex: exp) (x: var) (fv_ex: fv ex = []) (e1 e2 e3: exp):
  subst ex x fv_ex (exp.if_ e1 e2 e3) =
    exp.if_
    (subst ex x fv_ex e1)
    (subst ex x fv_ex e2)
    (subst ex x fv_ex e3)
  := by simp [subst]

theorem app_subst (ex: exp) (x: var) (fv_ex: fv ex = []) (e1 e2: exp):
  subst ex x fv_ex (exp.app e1 e2) =
    exp.app
    (subst ex x fv_ex e1)
    (subst ex x fv_ex e2)
  := by simp [subst]

theorem fn_subst (ex: exp) (x: var) (fv_ex: fv ex = [])
  (y: var) (τ: typ) (e: exp):
  subst ex x fv_ex (exp.fn y τ e) =
  exp.fn y τ (if x = y then e else subst ex x fv_ex e) :=
begin
  simp [subst],
  -- super weird... at this point the lhs and rhs are equal so we should be
  -- done by refl but it doesn't work. but this does...?
  by_cases x = y,
  simp [h],
  simp [h],
end

theorem pair_subst (ex: exp) (x: var) (fv_ex: fv ex = []) (e1 e2: exp):
  subst ex x fv_ex (exp.pair e1 e2) =
    exp.pair
    (subst ex x fv_ex e1)
    (subst ex x fv_ex e2)
  := by simp [subst]

theorem pair_left_subst (ex: exp) (x: var) (fv_ex: fv ex = []) (e: exp):
  subst ex x fv_ex (exp.pair_left e) = exp.pair_left (subst ex x fv_ex e)
  := by simp [subst]

theorem pair_right_subst (ex: exp) (x: var) (fv_ex: fv ex = []) (e: exp):
  subst ex x fv_ex (exp.pair_right e) = exp.pair_right (subst ex x fv_ex e)
  := by simp [subst]

theorem either_left_subst {τ: typ} (ex: exp) (x: var) (fv_ex: fv ex = []) (e: exp):
  subst ex x fv_ex (exp.either_left τ e) = exp.either_left τ (subst ex x fv_ex e)
  := by simp [subst]

theorem either_right_subst {τ: typ} (ex: exp) (x: var) (fv_ex: fv ex = []) (e: exp):
  subst ex x fv_ex (exp.either_right τ e) = exp.either_right τ (subst ex x fv_ex e)
  := by simp [subst]

theorem case_never_subst {τ: typ} (ex: exp) (x: var) (fv_ex: fv ex = []) (e: exp):
  subst ex x fv_ex (exp.case_never τ e) = exp.case_never τ (subst ex x fv_ex e)
  := by simp [subst]

theorem case_subst (ex: exp) (x: var) (fv_ex: fv ex = [])
  (eh e1 e2: exp) (x1 x2: var):
  subst ex x fv_ex (exp.case eh x1 e1 x2 e2) =
  exp.case
    (subst ex x fv_ex eh)
    x1 (if x = x1 then e1 else subst ex x fv_ex e1)
    x2 (if x = x2 then e2 else subst ex x fv_ex e2)
   :=
begin
  simp [subst],
  -- same as in fn_subst...?
  by_cases h1: x = x1,
  simp [h1],
  by_cases h2: x1 = x2,
  simp [h2],
  simp [h2],
  simp [h1],
  by_cases h2: x = x2,
  simp [h2],
  simp [h2],
end

theorem weakening_var_helper {Γ: env} {x x': var} {τ1 τ2 τ: typ} {e: exp}:
  has_typ (env.insert_exp Γ x' τ1) e τ2 ->
  (x ∉ fv e -> has_typ (env.insert_exp (env.insert_exp Γ x' τ1) x τ) e τ2) ->
  x ∉ list.filter (ne x') (fv e) ->
  has_typ (env.insert_exp (env.insert_exp Γ x τ) x' τ1) e τ2 :=
begin
  intros et ih fv_e,
  by_cases x = x',
  let a := useless_insert_twice Γ.exps x τ1 τ,
  rw symm h at et ⊢,
  simp [env.insert_exp] at et,
  rw symm a at et,
  exact et,
  let b := ih (not_filter fv_e (fun a, h (symm a))),
  simp [env.insert_exp] at b,
  rw insert_comm Γ.exps x x' τ τ1 h at b,
  exact b,
end

theorem weakening {Γ: env} {e: exp} {τ: typ} (x: var) (τx: typ):
  x ∉ fv e ->
  has_typ Γ e τ ->
  has_typ (env.insert_exp Γ x τx) e τ :=
begin
  intros fv_e et,
  induction et,
  exact has_typ.int,
  exact has_typ.true,
  exact has_typ.false,
  rw if_fv at fv_e,
  simp [list.append] at fv_e,
  let t_e1 := et_ih_a (fun a, fv_e (or.inl a)),
  let t_e2 := et_ih_a_1 (fun a, fv_e (or.inr (or.inl a))),
  let t_e3 := et_ih_a_2 (fun a, fv_e (or.inr (or.inr a))),
  exact has_typ.if_ t_e1 t_e2 t_e3,
  simp [fv] at fv_e,
  let var_ne := fun a, fv_e (symm a),
  exact has_typ.var (iff.elim_right (useless_insert_ne var_ne) et_a),
  rw fn_fv at fv_e,
  exact has_typ.fn (weakening_var_helper et_a et_ih fv_e),
  rw app_fv at fv_e,
  simp [list.append] at fv_e,
  let t_e1 := et_ih_a (fun a, fv_e (or.inl a)),
  let t_e2 := et_ih_a_1 (fun a, fv_e (or.inr a)),
  exact has_typ.app t_e1 t_e2,
  exact has_typ.unit,
  rw pair_fv at fv_e,
  simp [list.append] at fv_e,
  let t_e1 := et_ih_a (fun a, fv_e (or.inl a)),
  let t_e2 := et_ih_a_1 (fun a, fv_e (or.inr a)),
  exact has_typ.pair t_e1 t_e2,
  rw pair_left_fv at fv_e,
  exact has_typ.pair_left (et_ih fv_e),
  rw pair_right_fv at fv_e,
  exact has_typ.pair_right (et_ih fv_e),
  rw either_left_fv at fv_e,
  exact has_typ.either_left (et_ih fv_e),
  rw either_right_fv at fv_e,
  exact has_typ.either_right (et_ih fv_e),
  rw case_never_fv at fv_e,
  exact has_typ.case_never (et_ih fv_e),
  rw case_fv at fv_e,
  simp [list.append] at fv_e,
  let t_eh := et_ih_a (fun a, fv_e (or.inl a)),
  let fv_et_e1 := fun a, fv_e (or.inr (or.inl a)),
  let t_e1 := weakening_var_helper et_a_1 et_ih_a_1 fv_et_e1,
  let fv_et_e2 := fun a, fv_e (or.inr (or.inr a)),
  let t_e2 := weakening_var_helper et_a_2 et_ih_a_2 fv_et_e2,
  exact has_typ.case t_eh t_e1 t_e2,
end

theorem subst_preservation_var_helper
  {Γ Γ': env} {x x': var} {τ1 τ2 τx: typ}
  {ex e: exp} (fv_ex: fv ex = []):
  Γ' = env.insert_exp Γ x τx ->
  has_typ (env.insert_exp Γ' x' τ1) e τ2 ->
  has_typ Γ ex τx ->
  (∀ {Γ : env},
    env.insert_exp Γ' x' τ1 = env.insert_exp Γ x τx →
    has_typ Γ ex τx → has_typ Γ (subst ex x fv_ex e) τ2) ->
  has_typ (env.insert_exp Γ x' τ1) (ite (x = x') e (subst ex x fv_ex e)) τ2 :=
begin
  intros Γ'_is et ext ih,
  by_cases x = x',
  simp [h],
  rw Γ'_is at et,
  rw h at et,
  simp [env.insert_exp] at et,
  rw useless_insert_twice Γ.exps x' τ1 τx at et,
  exact et,
  simp [h],
  rw Γ'_is at ih,
  let notin_fv_ex := list.not_mem_nil x',
  rw symm fv_ex at notin_fv_ex,
  let ext' := weakening x' τ1 notin_fv_ex ext,
  simp [env.insert_exp] at ih,
  let a := insert_comm Γ.exps x' x τ1 τx (fun a, h (symm a)),
  exact @ih (env.insert_exp Γ x' τ1) a ext',
end

theorem subst_preservation
  {Γ Γ': env}
  {e ex: exp}
  {x: var}
  {τ τx: typ}
  (Γ'_is: Γ' = env.insert_exp Γ x τx)
  (fv_ex: fv ex = [])
  (et: has_typ Γ' e τ)
  (ext: has_typ Γ ex τx)
  : has_typ Γ (subst ex x fv_ex e) τ :=
begin
  induction et generalizing Γ,
  exact has_typ.int,
  exact has_typ.true,
  exact has_typ.false,
  let a := et_ih_a Γ'_is ext,
  let b := et_ih_a_1 Γ'_is ext,
  let c := et_ih_a_2 Γ'_is ext,
  exact has_typ.if_ a b c,
  rw Γ'_is at et_a,
  by_cases x = et_x,
  rw h at et_a ⊢,
  rw lookup_uniq (lookup_insert Γ.exps et_x τx) et_a at ext,
  simp [subst],
  exact ext,
  simp [subst, h],
  let h' := fun a, h (symm a),
  let hm := iff.elim_left (useless_insert_ne h') et_a,
  exact has_typ.var hm,
  rw fn_subst,
  let a := subst_preservation_var_helper fv_ex Γ'_is et_a ext @et_ih,
  exact has_typ.fn a,
  let a := et_ih_a Γ'_is ext,
  let b := et_ih_a_1 Γ'_is ext,
  exact has_typ.app a b,
  exact has_typ.unit,
  let a := et_ih_a Γ'_is ext,
  let b := et_ih_a_1 Γ'_is ext,
  exact has_typ.pair a b,
  exact has_typ.pair_left (et_ih Γ'_is ext),
  exact has_typ.pair_right (et_ih Γ'_is ext),
  exact has_typ.either_left (et_ih Γ'_is ext),
  exact has_typ.either_right (et_ih Γ'_is ext),
  exact has_typ.case_never (et_ih Γ'_is ext),
  rw case_subst,
  let t_e1 := subst_preservation_var_helper fv_ex Γ'_is et_a_1 ext @et_ih_a_1,
  let t_e2 := subst_preservation_var_helper fv_ex Γ'_is et_a_2 ext @et_ih_a_2,
  exact has_typ.case (et_ih_a Γ'_is ext) t_e1 t_e2,
end

theorem subst_fv_var_helper {x x': var} {ex e: exp} (fv_ex: fv ex = []):
  fv (subst ex x fv_ex e) = list.filter (ne x) (fv e) ->
  list.filter (ne x') (fv (ite (x = x') e (subst ex x fv_ex e))) =
  list.filter (ne x) (list.filter (ne x') (fv e)) :=
begin
  intro ih,
  by_cases x = x',
  -- can't just `simp [h]` or else weird stuff happens with mismatched types
  rw h,
  simp,
  exact symm (filter_idempotent (ne x') (fv e)),
  simp [h],
  simp [ih],
  exact filter_comm (ne x') (ne x) (fv e),
end

theorem subst_fv (ex: exp) (x: var) (fv_ex: fv ex = []) (e: exp):
  fv (subst ex x fv_ex e) = list.filter (ne x) (fv e) :=
begin
  let s := subst ex x fv_ex,
  induction e,
  simp [subst, fv],
  simp [subst, fv],
  simp [subst, fv],
  rw if_subst ex x fv_ex e_a e_a_1 e_a_2,
  rw if_fv (s e_a) (s e_a_1) (s e_a_2),
  rw e_ih_a,
  rw e_ih_a_1,
  rw e_ih_a_2,
  rw symm (list.filter_append (fv e_a_1) (fv e_a_2)),
  rw symm (list.filter_append (fv e_a) (fv e_a_1 ++ fv e_a_2)),
  simp [subst, fv],
  by_cases x = e,
  rw h,
  simp [subst, var_fv],
  exact fv_ex,
  simp [fv, subst, h],
  simp [fn_subst, fn_fv],
  exact subst_fv_var_helper fv_ex e_ih,
  rw app_subst ex x fv_ex e_a e_a_1,
  rw app_fv (s e_a) (s e_a_1),
  rw e_ih_a,
  rw e_ih_a_1,
  rw symm (list.filter_append (fv e_a) (fv e_a_1)),
  rw symm (app_fv e_a e_a_1),
  simp [subst, fv],
  simp [pair_subst],
  simp [pair_fv],
  rw e_ih_a,
  rw e_ih_a_1,
  simp [pair_left_subst, pair_left_fv],
  rw e_ih,
  simp [pair_right_subst, pair_right_fv],
  rw e_ih,
  simp [either_left_subst, either_left_fv],
  rw e_ih,
  simp [either_right_subst, either_right_fv],
  rw e_ih,
  simp [case_never_subst, case_never_fv],
  rw e_ih,
  simp [case_subst, case_fv],
  rw e_ih_a,
  rw subst_fv_var_helper fv_ex e_ih_a_1,
  rw subst_fv_var_helper fv_ex e_ih_a_2,
end


import Lean

macro "foo!" x:term:max : term => `($x + 1)

#check foo! 0

theorem ex1 : foo! 2 = 3 :=
  rfl

macro "foo!" x:term:max : term => `($x * 2)

#check foo! 1 -- ambiguous

-- macro with higher priority
macro (priority := high) "foo!" x:term:max : term => `($x - 2)

#check foo! 2

theorem ex2 : foo! 2 = 0 :=
  rfl

-- Define elaborator with even higher priority
elab (priority := high+1) "foo!" x:term:max : term <= expectedType =>
  Lean.Elab.Term.elabTerm x expectedType

theorem ex3 : foo! 3 = 3 :=
  rfl


import breen_deligne.constants
import breen_deligne.suitable
import pseudo_normed_group.FP
import system_of_complexes.rescale

noncomputable theory

open_locale classical nnreal big_operators
local attribute [instance] type_pow

universe variable u

namespace category_theory
namespace FreeAb

def of_functor (C : Type*) [category C] : C ⥤ FreeAb C :=
{ obj := of,
  map := λ X Y f, free_abelian_group.of f,
  map_id' := λ X, rfl,
  map_comp' := λ X Y Z f g, rfl }

end FreeAb
end category_theory

open category_theory breen_deligne

namespace breen_deligne

variables (r' : ℝ≥0)
variables (BD : breen_deligne.data)
variables (M : ProFiltPseuNormGrpWithTinv r')
variables (c c₁ c₂ c₃ c₄ : ℝ≥0) (l m n : ℕ)

open category_theory breen_deligne
open Profinite pseudo_normed_group profinitely_filtered_pseudo_normed_group

/-- The "functor" that sends `M` and `c` to `(filtration M c)^n` -/
def FP2 (r' : ℝ≥0) (c : ℝ≥0) (n : ℕ) :
  ProFiltPseuNormGrpWithTinv r' ⥤ FreeAb Profinite :=
FiltrationPow r' c n ⋙ FreeAb.of_functor _

theorem FP2_def (r' : ℝ≥0) (c : ℝ≥0) (n : ℕ) :
  FP2 r' c n = FiltrationPow r' c n ⋙ FreeAb.of_functor _ := rfl

namespace FP2

@[simps {fully_applied := ff}]
def res (r' : ℝ≥0) (c₁ c₂ : ℝ≥0) [fact (c₁ ≤ c₂)] (n : ℕ) : FP2 r' c₁ n ⟶ FP2 r' c₂ n :=
whisker_right (FiltrationPow.cast_le r' c₁ c₂ n) _

@[simp] lemma res_refl : res r' c c n = 𝟙 _ :=
by { simp [res, FiltrationPow.cast_le_refl], refl }

lemma res_comp_res [h₁ : fact (c₁ ≤ c₂)] [h₂ : fact (c₂ ≤ c₃)] :
  res r' c₁ c₂ n ≫ res r' c₂ c₃ n = @res r' c₁ c₃ ⟨le_trans h₁.1 h₂.1⟩ n :=
by simp only [res, ← whisker_right_comp, FiltrationPow.cast_le_comp]

section Tinv
open profinitely_filtered_pseudo_normed_group_with_Tinv
variables [fact (0 < r')]

@[simps {fully_applied := ff}]
def Tinv [fact (c₁ ≤ r' * c₂)] : FP2 r' c₁ n ⟶ FP2 r' c₂ n :=
whisker_right (FiltrationPow.Tinv r' c₁ c₂ n) _

lemma Tinv_def [fact (c₁ ≤ r' * c₂)] :
  Tinv r' c₁ c₂ n = whisker_right (FiltrationPow.Tinv r' c₁ c₂ n) _ := rfl

lemma res_comp_Tinv
  [fact (c₁ ≤ c₂)] [fact (c₂ ≤ c₃)] [fact (c₁ ≤ r' * c₂)] [fact (c₂ ≤ r' * c₃)] :
  res r' c₁ c₂ n ≫ Tinv r' c₂ c₃ n = Tinv r' c₁ c₂ n ≫ res r' c₂ c₃ n :=
by { simp only [Tinv, res, ← whisker_right_comp], refl }

end Tinv

end FP2

open FP2

variables {l m n}

namespace basic_universal_map

open basic_universal_map

variables (ϕ : basic_universal_map m n)

def eval_FP2 (c₁ c₂ : ℝ≥0) [ϕ.suitable c₁ c₂] : FP2 r' c₁ m ⟶ FP2 r' c₂ n :=
whisker_right (ϕ.eval_FP r' c₁ c₂) _

def eval_FP2' (c₁ c₂ : ℝ≥0) : FP2 r' c₁ m ⟶ FP2 r' c₂ n :=
if H : ϕ.suitable c₁ c₂
then by exactI whisker_right (ϕ.eval_FP r' c₁ c₂) _
else 0

lemma eval_FP2_eq_eval_FP2' (h : ϕ.suitable c₁ c₂) :
  eval_FP2 r' ϕ c₁ c₂ = eval_FP2' r' ϕ c₁ c₂ :=
by { delta eval_FP2 eval_FP2', rw dif_pos h }

lemma eval_FP2'_def [h : ϕ.suitable c₁ c₂] :
  eval_FP2' r' ϕ c₁ c₂ = whisker_right (ϕ.eval_FP r' c₁ c₂) _ :=
dif_pos h

lemma eval_FP2'_not_suitable (h : ¬ ϕ.suitable c₁ c₂) :
  eval_FP2' r' ϕ c₁ c₂ = 0 :=
dif_neg h

lemma eval_FP2'_comp (f : basic_universal_map l m) (g : basic_universal_map m n)
  [hf : f.suitable c₁ c₂] [hg : g.suitable c₂ c₃] :
  eval_FP2' r' (comp g f) c₁ c₃ = eval_FP2' r' f c₁ c₂ ≫ eval_FP2' r' g c₂ c₃ :=
begin
  haveI : (comp g f).suitable c₁ c₃ := suitable_comp c₂,
  simp only [eval_FP2'_def, eval_FP_comp r' _ c₂, whisker_right_comp]
end

lemma eval_FP2_comp (f : basic_universal_map l m) (g : basic_universal_map m n)
  [hf : f.suitable c₁ c₂] [hg : g.suitable c₂ c₃] :
  @eval_FP2 r' _ _ (comp g f) c₁ c₃ (suitable_comp c₂) =
    eval_FP2 r' f c₁ c₂ ≫ eval_FP2 r' g c₂ c₃ :=
by { simp only [eval_FP2_eq_eval_FP2'], apply eval_FP2'_comp }

lemma res_comp_eval_FP2
  [fact (c₁ ≤ c₂)] [fact (c₃ ≤ c₄)] [ϕ.suitable c₂ c₄] [ϕ.suitable c₁ c₃] :
  res r' c₁ c₂ m ≫ eval_FP2 r' ϕ c₂ c₄ = eval_FP2 r' ϕ c₁ c₃ ≫ res r' c₃ c₄ n :=
by simp only [res, eval_FP2, ← whisker_right_comp,
  cast_le_comp_eval_FP _ c₁ c₂ c₃ c₄]

lemma Tinv_comp_eval_FP2 [fact (0 < r')] [fact (c₁ ≤ r' * c₂)] [fact (c₃ ≤ r' * c₄)]
  [ϕ.suitable c₁ c₃] [ϕ.suitable c₂ c₄] :
  Tinv r' c₁ c₂ m ≫ eval_FP2 r' ϕ c₂ c₄ = eval_FP2 r' ϕ c₁ c₃ ≫ Tinv r' c₃ c₄ n :=
by simp only [Tinv, eval_FP2, ← whisker_right_comp,
  Tinv_comp_eval_FP _ _ c₁ c₂ c₃ c₄]

end basic_universal_map

namespace universal_map

open free_abelian_group

variables (ϕ : universal_map m n)

def eval_FP2 [ϕ.suitable c₁ c₂] : FP2 r' c₁ m ⟶ FP2 r' c₂ n :=
∑ g : {g : basic_universal_map m n // g ∈ ϕ.support},
  begin
    haveI := suitable_of_mem_support ϕ c₁ c₂ g g.2,
    exact coeff (g : basic_universal_map m n) ϕ • (basic_universal_map.eval_FP2 r' g c₁ c₂)
  end

def eval_FP2' : FP2 r' c₁ m ⟶ FP2 r' c₂ n :=
∑ g in ϕ.support, coeff g ϕ • (g.eval_FP2' r' c₁ c₂)

lemma eval_FP2_eq_eval_FP2' (h : ϕ.suitable c₁ c₂) :
  eval_FP2 r' c₁ c₂ ϕ = eval_FP2' r' c₁ c₂ ϕ :=
begin
  simp only [eval_FP2, eval_FP2', basic_universal_map.eval_FP2_eq_eval_FP2',
    subtype.val_eq_coe],
  symmetry,
  apply finset.sum_subtype ϕ.support (λ _, iff.rfl),
end

@[simp] lemma eval_FP2'_of (f : basic_universal_map m n) :
  eval_FP2' r' c₁ c₂ (of f) = f.eval_FP2' r' c₁ c₂ :=
by simp only [eval_FP2', support_of, coeff_of_self, one_smul, finset.sum_singleton]

@[simp] lemma eval_FP2_of (f : basic_universal_map m n) [f.suitable c₁ c₂] :
  eval_FP2 r' c₁ c₂ (of f) = f.eval_FP2 r' c₁ c₂ :=
by rw [eval_FP2_eq_eval_FP2', eval_FP2'_of, basic_universal_map.eval_FP2_eq_eval_FP2']

@[simp] lemma eval_FP2'_zero :
  eval_FP2' r' c₁ c₂ (0 : universal_map m n) = 0 :=
by rw [eval_FP2', support_zero, finset.sum_empty]

@[simp] lemma eval_FP2_zero :
  eval_FP2 r' c₁ c₂ (0 : universal_map m n) = 0 :=
by rw [eval_FP2_eq_eval_FP2', eval_FP2'_zero]

@[simp] lemma eval_FP2'_neg (f : universal_map m n) :
  eval_FP2' r' c₁ c₂ (-f) = -eval_FP2' r' c₁ c₂ f :=
by simp only [eval_FP2', add_monoid_hom.map_neg, finset.sum_neg_distrib, neg_smul, support_neg]

@[simp] lemma eval_FP2_neg (f : universal_map m n) [f.suitable c₁ c₂] :
  eval_FP2 r' c₁ c₂ (-f) = -eval_FP2 r' c₁ c₂ f :=
by simp only [eval_FP2_eq_eval_FP2', eval_FP2'_neg]

lemma eval_FP2'_add (f g : universal_map m n) :
  eval_FP2' r' c₁ c₂ (f + g) = eval_FP2' r' c₁ c₂ f + eval_FP2' r' c₁ c₂ g :=
begin
  simp only [eval_FP2'],
  rw finset.sum_subset (support_add f g), -- two goals
  simp only [add_monoid_hom.map_add _ f g, add_smul],
  convert finset.sum_add_distrib using 2, -- three goals
  apply finset.sum_subset (finset.subset_union_left _ _), swap,
  apply finset.sum_subset (finset.subset_union_right _ _),
  all_goals { rintros x - h, rw not_mem_support_iff at h, simp [h] },
end

lemma eval_FP2_add (f g : universal_map m n) [f.suitable c₁ c₂] [g.suitable c₁ c₂] :
  eval_FP2 r' c₁ c₂ (f + g) = eval_FP2 r' c₁ c₂ f + eval_FP2 r' c₁ c₂ g :=
by simp only [eval_FP2_eq_eval_FP2', eval_FP2'_add]

lemma eval_FP2_sub (f g : universal_map m n) [f.suitable c₁ c₂] [g.suitable c₁ c₂] :
  eval_FP2 r' c₁ c₂ (f - g) = eval_FP2 r' c₁ c₂ f - eval_FP2 r' c₁ c₂ g :=
by simp only [sub_eq_add_neg, eval_FP2_add, eval_FP2_neg]

lemma eval_FP2'_comp_of (g : basic_universal_map m n) (f : basic_universal_map l m)
  [hf : f.suitable c₁ c₂] [hg : g.suitable c₂ c₃] :
  eval_FP2' r' c₁ c₃ ((universal_map.comp (of g)) (of f)) =
    eval_FP2' r' c₁ c₂ (of f) ≫ eval_FP2' r' c₂ c₃ (of g) :=
begin
  simp only [universal_map.comp_of, eval_FP2'_of],
  haveI hfg : (basic_universal_map.comp g f).suitable c₁ c₃ := basic_universal_map.suitable_comp c₂,
  rw ← basic_universal_map.eval_FP2'_comp,
end

open category_theory category_theory.limits category_theory.preadditive

lemma eval_FP2'_comp (g : universal_map m n) (f : universal_map l m)
  [hf : f.suitable c₁ c₂] [hg : g.suitable c₂ c₃] :
  eval_FP2' r' c₁ c₃ (universal_map.comp g f) = eval_FP2' r' c₁ c₂ f ≫ eval_FP2' r' c₂ c₃ g :=
begin
  unfreezingI { revert hg },
  apply free_abelian_group.induction_on_free_predicate
    (universal_map.suitable c₁ c₂) (universal_map.suitable_free_predicate c₁ c₂) f hf;
      unfreezingI { clear_dependent f },
  { intros h₂,
    simp only [eval_FP2'_zero, zero_comp, pi.zero_apply,
      add_monoid_hom.zero_apply, add_monoid_hom.map_zero] },
  { intros f hf hg,
    -- now do another nested induction on `f`
    apply free_abelian_group.induction_on_free_predicate
      (universal_map.suitable c₂ c₃) (universal_map.suitable_free_predicate c₂ c₃) g hg;
        unfreezingI { clear_dependent g },
    { simp only [universal_map.eval_FP2'_zero, comp_zero, add_monoid_hom.map_zero,
        add_monoid_hom.zero_apply] },
    { intros g hg,
      rw suitable_of_iff at hf hg,
      resetI,
      apply eval_FP2'_comp_of },
    { intros g hg IH,
      simp only [IH, eval_FP2'_neg, add_monoid_hom.map_neg, comp_neg,
        add_monoid_hom.neg_apply] },
    { rintros (g₁ : universal_map m n) (g₂ : universal_map m n) hg₁ hg₂ IH₁ IH₂, resetI,
      haveI Hg₁f : (universal_map.comp g₁ (of f)).suitable c₁ c₃ := suitable.comp c₂,
      haveI Hg₂f : (universal_map.comp g₂ (of f)).suitable c₁ c₃ := suitable.comp c₂,
      simp only [add_monoid_hom.map_add, eval_FP2'_add, IH₁, IH₂, comp_add,
        add_monoid_hom.add_apply] } },
  { intros f hf IH hg, resetI, specialize IH,
    simp only [IH, add_monoid_hom.map_neg, eval_FP2'_neg,
      add_monoid_hom.neg_apply, neg_inj, neg_comp] },
  { rintros (f₁ : universal_map l m) (f₂ : universal_map l m) hf₁ hf₂ IH₁ IH₂ hf, resetI,
    haveI Hgf₁ : (universal_map.comp g f₁).suitable c₁ c₃ := suitable.comp c₂,
    haveI Hgf₂ : (universal_map.comp g f₂).suitable c₁ c₃ := suitable.comp c₂,
    simp only [add_monoid_hom.map_add, add_monoid_hom.add_apply, eval_FP2'_add, IH₁, IH₂, add_comp] }
end

lemma eval_FP2_comp (g : universal_map m n) (f : universal_map l m)
  [hf : f.suitable c₁ c₂] [hg : g.suitable c₂ c₃] :
  @eval_FP2 r' c₁ c₃ _ _ (universal_map.comp g f) (universal_map.suitable.comp c₂) =
    eval_FP2 r' c₁ c₂ f ≫ eval_FP2 r' c₂ c₃ g :=
by { simp only [eval_FP2_eq_eval_FP2'], apply eval_FP2'_comp }

lemma res_comp_eval_FP2 [fact (c₁ ≤ c₂)] [fact (c₃ ≤ c₄)] [ϕ.suitable c₁ c₃] [ϕ.suitable c₂ c₄] :
  res r' c₁ c₂ m ≫ eval_FP2 r' c₂ c₄ ϕ = eval_FP2 r' c₁ c₃ ϕ ≫ res r' c₃ c₄ n :=
begin
  simp only [eval_FP2, comp_sum, sum_comp, comp_zsmul, zsmul_comp],
  apply finset.sum_congr rfl,
  rintros ⟨g, hg⟩ -,
  haveI : g.suitable c₁ c₃ := suitable_of_mem_support ϕ _ _ g hg,
  haveI : g.suitable c₂ c₄ := suitable_of_mem_support ϕ _ _ g hg,
  simp only [subtype.coe_mk, g.res_comp_eval_FP2 r' c₁ c₂ c₃ c₄],
end

lemma Tinv_comp_eval_FP2 [fact (0 < r')] [fact (c₁ ≤ r' * c₂)] [fact (c₃ ≤ r' * c₄)]
  [ϕ.suitable c₁ c₃] [ϕ.suitable c₂ c₄] :
  Tinv r' c₁ c₂ m ≫ eval_FP2 r' c₂ c₄ ϕ = eval_FP2 r' c₁ c₃ ϕ ≫ Tinv r' c₃ c₄ n :=
begin
  simp only [eval_FP2, comp_sum, sum_comp, comp_zsmul, zsmul_comp],
  apply finset.sum_congr rfl,
  rintros ⟨g, hg⟩ -,
  haveI : g.suitable c₁ c₃ := suitable_of_mem_support ϕ _ _ g hg,
  haveI : g.suitable c₂ c₄ := suitable_of_mem_support ϕ _ _ g hg,
  congr' 1, apply basic_universal_map.Tinv_comp_eval_FP2 r',
end

end universal_map

variables (κ : ℝ≥0 → ℕ → ℝ≥0) [∀ c, BD.suitable (κ c)]

def FPsystem.X (c : ℝ≥0) (n : ℕ) : FreeAb Profinite :=
FreeAb.of $ (FiltrationPow r' (κ c n) $ BD.X n).obj M

def FPsystem.d (c : ℝ≥0) (n : ℕ) :
  FPsystem.X r' BD M κ c (n + 1) ⟶ FPsystem.X r' BD M κ c n :=
(universal_map.eval_FP2 r' (κ c (n+1)) (κ c n) (BD.d (n+1) n)).app M

lemma FPsystem.d_comp_d (c : ℝ≥0) (n : ℕ) :
  FPsystem.d r' BD M κ c (n + 1) ≫ FPsystem.d r' BD M κ c n = 0 :=
begin
  delta FPsystem.d,
  rw [← nat_trans.comp_app, ← universal_map.eval_FP2_comp],
  convert nat_trans.app_zero _, refl, refl,
  convert universal_map.eval_FP2_zero _ _ _,
  show BD.d _ _ ≫ BD.d _ _ = 0,
  rw homological_complex.d_comp_d,
end

open opposite

def FPsystem [hκ : ∀ n, fact (monotone (function.swap κ n))] :
  ℝ≥0 ⥤ chain_complex (FreeAb Profinite) ℕ :=
{ obj := λ c, chain_complex.of (FPsystem.X r' BD M κ c) (FPsystem.d r' BD M κ _) (FPsystem.d_comp_d _ _ _ _ _),
  map := λ c₁ c₂ h,
  { f := λ n, by { refine (@FP2.res r' _ _ (id _) (BD.X n)).app M,
      have := (hκ n).out, refine ⟨this h.le⟩, },
    comm' := begin
      rintro i j (rfl : j + 1 = i),
      rw [chain_complex.of_d, chain_complex.of_d],
      delta FPsystem.d, rw [← nat_trans.comp_app, ← nat_trans.comp_app],
      congr' 1,
      apply universal_map.res_comp_eval_FP2,
    end },
  map_id' := λ c, begin
    ext n, dsimp, rw [Filtration.cast_le_refl, (FreeAb.of_functor _).map_id], refl,
  end,
  map_comp' := λ c₁ c₂ c₃ h₁₂ h₂₃, begin
    ext n, dsimp, rw [← (FreeAb.of_functor _).map_comp, Filtration.cast_le_comp],
  end }
.

def FPsystem.Tinv [fact (0 < r')]
  (κ₁ κ₂ : ℝ≥0 → ℕ → ℝ≥0)
  [∀ c, BD.suitable (κ₁ c)] [∀ c, BD.suitable (κ₂ c)]
  [hκ₁ : ∀ n, fact (monotone (function.swap κ₁ n))]
  [hκ₂ : ∀ n, fact (monotone (function.swap κ₂ n))]
  [∀ c n, fact (κ₁ c n ≤ r' * κ₂ c n)] :
  FPsystem r' BD M κ₁ ⟶ FPsystem r' BD M κ₂ :=
{ app := λ c,
  { f := λ n, (FP2.Tinv r' _ _ _).app M,
    comm' := begin
      rintro i j (rfl : j + 1 = i),
      dsimp only [functor.comp_obj, FPsystem],
      rw [chain_complex.of_d, chain_complex.of_d],
      delta FPsystem.d,
      rw [← nat_trans.comp_app, ← nat_trans.comp_app],
      congr' 1,
      apply universal_map.Tinv_comp_eval_FP2
    end },
  naturality' := begin
    intros c₁ c₂ h,
    ext n,
    dsimp only [FPsystem, Tinv_app, homological_complex.comp_f, functor.comp_map, res_app],
    rw [← functor.map_comp, ← functor.map_comp],
    refl,
  end }

def FPsystem.res [fact (r' ≤ 1)]
  (κ₁ κ₂ : ℝ≥0 → ℕ → ℝ≥0)
  [∀ c, BD.suitable (κ₁ c)] [∀ c, BD.suitable (κ₂ c)]
  [hκ₁ : ∀ n, fact (monotone (function.swap κ₁ n))]
  [hκ₂ : ∀ n, fact (monotone (function.swap κ₂ n))]
  [∀ c n, fact (κ₁ c n ≤ κ₂ c n)] :
  FPsystem r' BD M κ₁ ⟶ FPsystem r' BD M κ₂ :=
{ app := λ c,
  { f := λ n, (FP2.res r' _ _ _).app M,
    comm' := begin
      rintro i j (rfl : j + 1 = i),
      dsimp only [functor.comp_obj, FPsystem],
      rw [chain_complex.of_d, chain_complex.of_d],
      delta FPsystem.d,
      rw [← nat_trans.comp_app, ← nat_trans.comp_app],
      congr' 1,
      apply universal_map.res_comp_eval_FP2
    end },
  naturality' := begin
    intros c₁ c₂ h,
    ext n,
    dsimp only [FPsystem, res_app, homological_complex.comp_f, functor.comp_map],
    rw [← functor.map_comp, ← functor.map_comp],
    refl,
  end }

end breen_deligne