/-
Copyright (c) 2017 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Mario Carneiro
-/
import data.seq.seq data.seq.computation data.list.basic data.dlist
universes u v w

/-
coinductive wseq (α : Type u) : Type u
| nil : wseq α
| cons : α → wseq α → wseq α
| think : wseq α → wseq α
-/

/-- Weak sequences.

  While the `seq` structure allows for lists which may not be finite,
  a weak sequence also allows the computation of each element to
  involve an indeterminate amount of computation, including possibly
  an infinite loop. This is represented as a regular `seq` interspersed
  with `none` elements to indicate that computation is ongoing.

  This model is appropriate for Haskell style lazy lists, and is closed
  under most interesting computation patterns on infinite lists,
  but conversely it is difficult to extract elements from it. -/
def wseq (α) := seq (option α)

namespace wseq
variables {α : Type u} {β : Type v} {γ : Type w}

/-- Turn a sequence into a weak sequence -/
def of_seq : seq α → wseq α := (<$>) some

/-- Turn a list into a weak sequence -/
def of_list (l : list α) : wseq α := of_seq l

/-- Turn a stream into a weak sequence -/
def of_stream (l : stream α) : wseq α := of_seq l

instance coe_seq : has_coe (seq α) (wseq α) := ⟨of_seq⟩
instance coe_list : has_coe (list α) (wseq α) := ⟨of_list⟩
instance coe_stream : has_coe (stream α) (wseq α) := ⟨of_stream⟩

/-- The empty weak sequence -/
def nil : wseq α := seq.nil

/-- Prepend an element to a weak sequence -/
def cons (a : α) : wseq α → wseq α := seq.cons (some a)

/-- Compute for one tick, without producing any elements -/
def think : wseq α → wseq α := seq.cons none

/-- Destruct a weak sequence, to (eventually possibly) produce either
  `none` for `nil` or `some (a, s)` if an element is produced. -/
def destruct : wseq α → computation (option (α × wseq α)) :=
computation.corec (λs, match seq.destruct s with
  | none              := sum.inl none
  | some (none, s')   := sum.inr s'
  | some (some a, s') := sum.inl (some (a, s'))
  end)

def cases_on {C : wseq α → Sort v} (s : wseq α) (h1 : C nil)
  (h2 : ∀ x s, C (cons x s)) (h3 : ∀ s, C (think s)) : C s :=
seq.cases_on s h1 (λ o, option.cases_on o h3 h2)

protected def mem (a : α) (s : wseq α) := seq.mem (some a) s

instance : has_mem α (wseq α) :=
⟨wseq.mem⟩

theorem not_mem_nil (a : α) : a ∉ @nil α := seq.not_mem_nil a

/-- Get the head of a weak sequence. This involves a possibly
  infinite computation. -/
def head (s : wseq α) : computation (option α) :=
computation.map ((<$>) prod.fst) (destruct s)

/-- Encode a computation yielding a weak sequence into additional
  `think` constructors in a weak sequence -/
def flatten : computation (wseq α) → wseq α :=
seq.corec (λc, match computation.destruct c with
  | sum.inl s := seq.omap return (seq.destruct s)
  | sum.inr c' := some (none, c')
  end)

/-- Get the tail of a weak sequence. This doesn't need a `computation`
  wrapper, unlike `head`, because `flatten` allows us to hide this
  in the construction of the weak sequence itself. -/
def tail (s : wseq α) : wseq α :=
flatten $ (λo, option.rec_on o nil prod.snd) <$> destruct s

/-- drop the first `n` elements from `s`. -/
def drop (s : wseq α) : ℕ → wseq α
| 0     := s
| (n+1) := tail (drop n)
attribute [simp] drop

/-- Get the nth element of `s`. -/
def nth (s : wseq α) (n : ℕ) : computation (option α) := head (drop s n)

/-- Convert `s` to a list (if it is finite and completes in finite time). -/
def to_list (s : wseq α) : computation (list α) :=
@computation.corec (list α) (list α × wseq α) (λ⟨l, s⟩,
  match seq.destruct s with
  | none              := sum.inl l.reverse
  | some (none, s')   := sum.inr (l, s')
  | some (some a, s') := sum.inr (a::l, s')
  end) ([], s)

/-- Get the length of `s` (if it is finite and completes in finite time). -/
def length (s : wseq α) : computation ℕ :=
@computation.corec ℕ (ℕ × wseq α) (λ⟨n, s⟩,
  match seq.destruct s with
  | none              := sum.inl n
  | some (none, s')   := sum.inr (n, s')
  | some (some a, s') := sum.inr (n+1, s')
  end) (0, s)

/-- A weak sequence is finite if `to_list s` terminates. Equivalently,
  it is a finite number of `think` and `cons` applied to `nil`. -/
@[class] def is_finite (s : wseq α) : Prop := (to_list s).terminates

instance to_list_terminates (s : wseq α) [h : is_finite s] : (to_list s).terminates := h

/-- Get the list corresponding to a finite weak sequence. -/
def get (s : wseq α) [is_finite s] : list α := (to_list s).get

/-- A weak sequence is *productive* if it never stalls forever - there are
 always a finite number of `think`s between `cons` constructors.
 The sequence itself is allowed to be infinite though. -/
@[class] def productive (s : wseq α) : Prop := ∀ n, (nth s n).terminates

instance nth_terminates (s : wseq α) [h : productive s] : ∀ n, (nth s n).terminates := h

instance head_terminates (s : wseq α) [h : productive s] : (head s).terminates := h 0

/-- Replace the `n`th element of `s` with `a`. -/
def update_nth (s : wseq α) (n : ℕ) (a : α) : wseq α :=
@seq.corec (option α) (ℕ × wseq α) (λ⟨n, s⟩,
  match seq.destruct s, n with
  | none,               n     := none
  | some (none, s'),    n     := some (none, n, s')
  | some (some a', s'), 0     := some (some a', 0, s')
  | some (some a', s'), 1     := some (some a, 0, s')
  | some (some a', s'), (n+2) := some (some a', n+1, s')
  end) (n+1, s)

/-- Remove the `n`th element of `s`. -/
def remove_nth (s : wseq α) (n : ℕ) : wseq α :=
@seq.corec (option α) (ℕ × wseq α) (λ⟨n, s⟩,
  match seq.destruct s, n with
  | none,               n     := none
  | some (none, s'),    n     := some (none, n, s')
  | some (some a', s'), 0     := some (some a', 0, s')
  | some (some a', s'), 1     := some (none, 0, s')
  | some (some a', s'), (n+2) := some (some a', n+1, s')
  end) (n+1, s)

/-- Map the elements of `s` over `f`, removing any values that yield `none`. -/
def filter_map (f : α → option β) : wseq α → wseq β :=
seq.corec (λs, match seq.destruct s with
  | none              := none
  | some (none, s')   := some (none, s')
  | some (some a, s') := some (f a, s')
  end)

/-- Select the elements of `s` that satisfy `p`. -/
def filter (p : α → Prop) [decidable_pred p] : wseq α → wseq α :=
filter_map (λa, if p a then some a else none)

-- example of infinite list manipulations
/-- Get the first element of `s` satisfying `p`. -/
def find (p : α → Prop) [decidable_pred p] (s : wseq α) : computation (option α) :=
head $ filter p s

/-- Zip a function over two weak sequences -/
def zip_with (f : α → β → γ) (s1 : wseq α) (s2 : wseq β) : wseq γ :=
@seq.corec (option γ) (wseq α × wseq β) (λ⟨s1, s2⟩,
  match seq.destruct s1, seq.destruct s2 with
  | some (none, s1'),    some (none, s2')    := some (none, s1', s2')
  | some (some a1, s1'), some (none, s2')    := some (none, s1, s2')
  | some (none, s1'),    some (some a2, s2') := some (none, s1', s2)
  | some (some a1, s1'), some (some a2, s2') := some (some (f a1 a2), s1', s2')
  | _,                   _                   := none
  end) (s1, s2)

/-- Zip two weak sequences into a single sequence of pairs -/
def zip : wseq α → wseq β → wseq (α × β) := zip_with prod.mk

/-- Get the list of indexes of elements of `s` satisfying `p` -/
def find_indexes (p : α → Prop) [decidable_pred p] (s : wseq α) : wseq ℕ :=
(zip s (stream.nats : wseq ℕ)).filter_map
  (λ ⟨a, n⟩, if p a then some n else none)

/-- Get the index of the first element of `s` satisfying `p` -/
def find_index (p : α → Prop) [decidable_pred p] (s : wseq α) : computation ℕ :=
(λ o, option.get_or_else o 0) <$> head (find_indexes p s)

/-- Get the index of the first occurrence of `a` in `s` -/
def index_of [decidable_eq α] (a : α) : wseq α → computation ℕ := find_index (eq a)

/-- Get the indexes of occurrences of `a` in `s` -/
def indexes_of [decidable_eq α] (a : α) : wseq α → wseq ℕ := find_indexes (eq a)

/-- `union s1 s2` is a weak sequence which interleaves `s1` and `s2` in
  some order (nondeterministically). -/
def union (s1 s2 : wseq α) : wseq α :=
@seq.corec (option α) (wseq α × wseq α) (λ⟨s1, s2⟩,
  match seq.destruct s1, seq.destruct s2 with
  | none,                none                := none
  | some (a1, s1'),      none                := some (a1, s1', nil)
  | none,                some (a2, s2')      := some (a2, nil, s2')
  | some (none, s1'),    some (none, s2')    := some (none, s1', s2')
  | some (some a1, s1'), some (none, s2')    := some (some a1, s1', s2')
  | some (none, s1'),    some (some a2, s2') := some (some a2, s1', s2')
  | some (some a1, s1'), some (some a2, s2') := some (some a1, cons a2 s1', s2')
  end) (s1, s2)

/-- Returns `tt` if `s` is `nil` and `ff` if `s` has an element -/
def is_empty (s : wseq α) : computation bool :=
computation.map option.is_none $ head s

/-- Calculate one step of computation -/
def compute (s : wseq α) : wseq α :=
match seq.destruct s with
| some (none, s') := s'
| _               := s
end

/-- Get the first `n` elements of a weak sequence -/
def take (s : wseq α) (n : ℕ) : wseq α :=
@seq.corec (option α) (ℕ × wseq α) (λ⟨n, s⟩,
  match n, seq.destruct s with
  | 0,   _                 := none
  | m+1, none              := none
  | m+1, some (none, s')   := some (none, m+1, s')
  | m+1, some (some a, s') := some (some a, m, s')
  end) (n, s)

/-- Split the sequence at position `n` into a finite initial segment
  and the weak sequence tail -/
def split_at (s : wseq α) (n : ℕ) : computation (list α × wseq α) :=
@computation.corec (list α × wseq α) (ℕ × list α × wseq α) (λ⟨n, l, s⟩,
  match n, seq.destruct s with
  | 0,   _                 := sum.inl (l.reverse, s)
  | m+1, none              := sum.inl (l.reverse, s)
  | m+1, some (none, s')   := sum.inr (n, l, s')
  | m+1, some (some a, s') := sum.inr (m, a::l, s')
  end) (n, [], s)

/-- Returns `tt` if any element of `s` satisfies `p` -/
def any (s : wseq α) (p : α → bool) : computation bool :=
computation.corec (λs : wseq α,
  match seq.destruct s with
  | none              := sum.inl ff
  | some (none, s')   := sum.inr s'
  | some (some a, s') := if p a then sum.inl tt else sum.inr s'
  end) s

/-- Returns `tt` if every element of `s` satisfies `p` -/
def all (s : wseq α) (p : α → bool) : computation bool :=
computation.corec (λs : wseq α,
  match seq.destruct s with
  | none              := sum.inl tt
  | some (none, s')   := sum.inr s'
  | some (some a, s') := if p a then sum.inr s' else sum.inl ff
  end) s

/-- Apply a function to the elements of the sequence to produce a sequence
  of partial results. (There is no `scanr` because this would require
  working from the end of the sequence, which may not exist.) -/
def scanl (f : α → β → α) (a : α) (s : wseq β) : wseq α :=
cons a $ @seq.corec (option α) (α × wseq β) (λ⟨a, s⟩,
  match seq.destruct s with
  | none              := none
  | some (none, s')   := some (none, a, s')
  | some (some b, s') := let a' := f a b in some (some a', a', s')
  end) (a, s)

/-- Get the weak sequence of initial segments of the input sequence -/
def inits (s : wseq α) : wseq (list α) :=
cons [] $ @seq.corec (option (list α)) (dlist α × wseq α) (λ ⟨l, s⟩,
  match seq.destruct s with
  | none              := none
  | some (none, s')   := some (none, l, s')
  | some (some a, s') := let l' := l.concat a in
                         some (some l'.to_list, l', s')
  end) (dlist.empty, s)

/-- Like take, but does not wait for a result. Calculates `n` steps of
  computation and returns the sequence computed so far -/
def collect (s : wseq α) (n : ℕ) : list α :=
(seq.take n s).filter_map id

/-- Append two weak sequences. As with `seq.append`, this may not use
  the second sequence if the first one takes forever to compute -/
def append : wseq α → wseq α → wseq α := seq.append

/-- Map a function over a weak sequence -/
def map (f : α → β) : wseq α → wseq β := seq.map (option.map f)

/-- Flatten a sequence of weak sequences. (Note that this allows
  empty sequences, unlike `seq.join`.) -/
def join (S : wseq (wseq α)) : wseq α :=
seq.join ((λo : option (wseq α), match o with
  | none := seq1.ret none
  | some s := (none, s)
  end) <$> S)

/-- Monadic bind operator for weak sequences -/
def bind (s : wseq α) (f : α → wseq β) : wseq β :=
join (map f s)

@[simp] def lift_rel_o (R : α → β → Prop) (C : wseq α → wseq β → Prop) :
  option (α × wseq α) → option (β × wseq β) → Prop
| none          none          := true
| (some (a, s)) (some (b, t)) := R a b ∧ C s t
| _             _             := false

theorem lift_rel_o.imp {R S : α → β → Prop} {C D : wseq α → wseq β → Prop}
  (H1 : ∀ a b, R a b → S a b) (H2 : ∀ s t, C s t → D s t) :
  ∀ {o p}, lift_rel_o R C o p → lift_rel_o S D o p
| none          none          h := trivial
| (some (a, s)) (some (b, t)) h := and.imp (H1 _ _) (H2 _ _) h
| none          (some _)      h := false.elim h
| (some (_, _)) none          h := false.elim h

theorem lift_rel_o.imp_right (R : α → β → Prop) {C D : wseq α → wseq β → Prop}
  (H : ∀ s t, C s t → D s t) {o p} : lift_rel_o R C o p → lift_rel_o R D o p :=
lift_rel_o.imp (λ _ _, id) H

@[simp] def bisim_o (R : wseq α → wseq α → Prop) :
  option (α × wseq α) → option (α × wseq α) → Prop := lift_rel_o (=) R

theorem bisim_o.imp {R S : wseq α → wseq α → Prop} (H : ∀ s t, R s t → S s t) {o p} :
  bisim_o R o p → bisim_o S o p :=
lift_rel_o.imp_right _ H

/-- Two weak sequences are `lift_rel R` related if they are either both empty,
  or they are both nonempty and the heads are `R` related and the tails are
  `lift_rel R` related. (This is a coinductive definition.) -/
def lift_rel (R : α → β → Prop) (s : wseq α) (t : wseq β) : Prop :=
∃ C : wseq α → wseq β → Prop, C s t ∧
∀ {s t}, C s t → computation.lift_rel (lift_rel_o R C) (destruct s) (destruct t)

/-- If two sequences are equivalent, then they have the same values and
  the same computational behavior (i.e. if one loops forever then so does
  the other), although they may differ in the number of `think`s needed to
  arrive at the answer. -/
def equiv : wseq α → wseq α → Prop := lift_rel (=)

theorem lift_rel_destruct {R : α → β → Prop} {s : wseq α} {t : wseq β} :
  lift_rel R s t →
    computation.lift_rel (lift_rel_o R (lift_rel R)) (destruct s) (destruct t)
| ⟨R, h1, h2⟩ :=
  by refine computation.lift_rel.imp _ _ _ (h2 h1);
     apply lift_rel_o.imp_right; exact λ s' t' h', ⟨R, h', @h2⟩

theorem lift_rel_destruct_iff {R : α → β → Prop} {s : wseq α} {t : wseq β} :
  lift_rel R s t ↔
    computation.lift_rel (lift_rel_o R (lift_rel R)) (destruct s) (destruct t) :=
⟨lift_rel_destruct, λ h, ⟨λ s t, lift_rel R s t ∨
  computation.lift_rel (lift_rel_o R (lift_rel R)) (destruct s) (destruct t),
  or.inr h, λ s t h, begin
    have h : computation.lift_rel (lift_rel_o R (lift_rel R)) (destruct s) (destruct t),
    { cases h with h h, exact lift_rel_destruct h, assumption },
    apply computation.lift_rel.imp _ _ _ h,
    intros a b, apply lift_rel_o.imp_right,
    intros s t, apply or.inl
  end⟩⟩

infix ~ := equiv

theorem destruct_congr {s t : wseq α} :
  s ~ t → computation.lift_rel (bisim_o (~)) (destruct s) (destruct t) :=
lift_rel_destruct

theorem destruct_congr_iff {s t : wseq α} :
  s ~ t ↔ computation.lift_rel (bisim_o (~)) (destruct s) (destruct t) :=
lift_rel_destruct_iff

theorem lift_rel.refl (R : α → α → Prop) (H : reflexive R) : reflexive (lift_rel R) :=
λ s, begin
  refine ⟨(=), rfl, λ s t (h : s = t), _⟩,
  rw ←h, apply computation.lift_rel.refl,
  intro a, cases a with a, simp, cases a; simp, apply H
end

theorem lift_rel_o.swap (R : α → β → Prop) (C) :
  function.swap (lift_rel_o R C) = lift_rel_o (function.swap R) (function.swap C) :=
by funext x y; cases x with x; [skip, cases x]; { cases y with y; [skip, cases y]; refl }

theorem lift_rel.swap_lem {R : α → β → Prop} {s1 s2} (h : lift_rel R s1 s2) :
  lift_rel (function.swap R) s2 s1 :=
begin
  refine ⟨function.swap (lift_rel R), h, λ s t (h : lift_rel R t s), _⟩,
  rw [←lift_rel_o.swap, computation.lift_rel.swap],
  apply lift_rel_destruct h
end

theorem lift_rel.swap (R : α → β → Prop) :
  function.swap (lift_rel R) = lift_rel (function.swap R) :=
funext $ λ x, funext $ λ y, propext ⟨lift_rel.swap_lem, lift_rel.swap_lem⟩

theorem lift_rel.symm (R : α → α → Prop) (H : symmetric R) : symmetric (lift_rel R) :=
λ s1 s2 (h : function.swap (lift_rel R) s2 s1),
by rwa [lift_rel.swap, show function.swap R = R, from
        funext $ λ a, funext $ λ b, propext $ by constructor; apply H] at h

theorem lift_rel.trans (R : α → α → Prop) (H : transitive R) : transitive (lift_rel R) :=
λ s t u h1 h2, begin
  refine ⟨λ s u, ∃ t, lift_rel R s t ∧ lift_rel R t u, ⟨t, h1, h2⟩, λ s u h, _⟩,
  rcases h with ⟨t, h1, h2⟩,
  have h1 := lift_rel_destruct h1,
  have h2 := lift_rel_destruct h2,
  refine computation.lift_rel_def.2
    ⟨(computation.terminates_of_lift_rel h1).trans
     (computation.terminates_of_lift_rel h2), λ a c ha hc, _⟩,
  rcases h1.left ha with ⟨b, hb, t1⟩,
  have t2 := computation.rel_of_lift_rel h2 hb hc,
  cases a with a; cases c with c,
  { trivial },
  { cases b, {cases t2}, {cases t1} },
  { cases a, cases b with b, {cases t1}, {cases b, cases t2} },
  { cases a with a s, cases b with b, {cases t1},
    cases b with b t, cases c with c u,
    cases t1 with ab st, cases t2 with bc tu,
    exact ⟨H ab bc, t, st, tu⟩ }
end

theorem lift_rel.equiv (R : α → α → Prop) : equivalence R → equivalence (lift_rel R)
| ⟨refl, symm, trans⟩ :=
  ⟨lift_rel.refl R refl, lift_rel.symm R symm, lift_rel.trans R trans⟩

@[refl] theorem equiv.refl : ∀ (s : wseq α), s ~ s :=
lift_rel.refl (=) eq.refl

@[symm] theorem equiv.symm : ∀ {s t : wseq α}, s ~ t → t ~ s :=
lift_rel.symm (=) (@eq.symm _)

@[trans] theorem equiv.trans : ∀ {s t u : wseq α}, s ~ t → t ~ u → s ~ u :=
lift_rel.trans (=) (@eq.trans _)

theorem equiv.equivalence : equivalence (@equiv α) :=
⟨@equiv.refl _, @equiv.symm _, @equiv.trans _⟩

open computation
local notation `return` := computation.return

@[simp] theorem destruct_nil : destruct (nil : wseq α) = return none :=
computation.destruct_eq_ret rfl

@[simp] theorem destruct_cons (a : α) (s) : destruct (cons a s) = return (some (a, s)) :=
computation.destruct_eq_ret $ by simp [destruct, cons, computation.rmap]

@[simp] theorem destruct_think (s : wseq α) : destruct (think s) = (destruct s).think :=
computation.destruct_eq_think $ by simp [destruct, think, computation.rmap]

@[simp] theorem seq_destruct_nil : seq.destruct (nil : wseq α) = none :=
seq.destruct_nil

@[simp] theorem seq_destruct_cons (a : α) (s) : seq.destruct (cons a s) = some (some a, s) :=
seq.destruct_cons _ _

@[simp] theorem seq_destruct_think (s : wseq α) : seq.destruct (think s) = some (none, s) :=
seq.destruct_cons _ _

@[simp] theorem head_nil : head (nil : wseq α) = return none := by simp [head]; refl
@[simp] theorem head_cons (a : α) (s) : head (cons a s) = return (some a) := by simp [head]; refl
@[simp] theorem head_think (s : wseq α) : head (think s) = (head s).think := by simp [head]; refl

@[simp] theorem flatten_ret (s : wseq α) : flatten (return s) = s :=
begin
  refine seq.eq_of_bisim (λs1 s2, flatten (return s2) = s1) _ rfl,
  intros s' s h, rw ←h, simp [flatten],
  cases seq.destruct s, { simp },
  { cases val with o s', simp }
end

@[simp] theorem flatten_think (c : computation (wseq α)) : flatten c.think = think (flatten c) :=
seq.destruct_eq_cons $ by simp [flatten, think]

@[simp] theorem destruct_flatten (c : computation (wseq α)) : destruct (flatten c) = c >>= destruct :=
begin
  refine computation.eq_of_bisim (λc1 c2, c1 = c2 ∨
    ∃ c, c1 = destruct (flatten c) ∧ c2 = computation.bind c destruct) _ (or.inr ⟨c, rfl, rfl⟩),
  intros c1 c2 h, exact match c1, c2, h with
  | _, _, (or.inl $ eq.refl c) := by cases c.destruct; simp
  | _, _, (or.inr ⟨c, rfl, rfl⟩) := begin
    apply c.cases_on (λa, _) (λc', _); repeat {simp},
    { cases (destruct a).destruct; simp },
    { exact or.inr ⟨c', rfl, rfl⟩ }
  end end
end

theorem head_terminates_iff (s : wseq α) : terminates (head s) ↔ terminates (destruct s) :=
terminates_map_iff _ (destruct s)

@[simp] theorem tail_nil : tail (nil : wseq α) = nil := by simp [tail]
@[simp] theorem tail_cons (a : α) (s) : tail (cons a s) = s := by simp [tail]
@[simp] theorem tail_think (s : wseq α) : tail (think s) = (tail s).think := by simp [tail]

@[simp] theorem dropn_nil (n) :
  drop (nil : wseq α) n = nil := by induction n; simp [*, drop]
@[simp] theorem dropn_cons (a : α) (s) (n) :
  drop (cons a s) (n+1) = drop s n := by induction n; simp [*, drop]
@[simp] theorem dropn_think (s : wseq α) (n) :
  drop (think s) n = (drop s n).think := by induction n; simp [*, drop]

theorem dropn_add (s : wseq α) (m) : ∀ n, drop s (m + n) = drop (drop s m) n
| 0     := rfl
| (n+1) := congr_arg tail (dropn_add n)

theorem dropn_tail (s : wseq α) (n) : drop (tail s) n = drop s (n + 1) :=
by rw add_comm; symmetry; apply dropn_add

theorem nth_add (s : wseq α) (m n) : nth s (m + n) = nth (drop s m) n :=
congr_arg head (dropn_add _ _ _)

theorem nth_tail (s : wseq α) (n) : nth (tail s) n = nth s (n + 1) :=
congr_arg head (dropn_tail _ _)

@[simp] theorem join_nil : join nil = (nil : wseq α) := seq.join_nil

@[simp] theorem join_think (S : wseq (wseq α)) :
  join (think S) = think (join S) :=
by { simp [think, join], unfold functor.map, simp [join, seq1.ret] }

@[simp] theorem join_cons (s : wseq α) (S) :
  join (cons s S) = think (append s (join S)) :=
by { simp [think, join], unfold functor.map, simp [join, cons, append] }

@[simp] theorem nil_append (s : wseq α) : append nil s = s := seq.nil_append _

@[simp] theorem cons_append (a : α) (s t) :
  append (cons a s) t = cons a (append s t) := seq.cons_append _ _ _

@[simp] theorem think_append (s t : wseq α) :
  append (think s) t = think (append s t) := seq.cons_append _ _ _

@[simp] theorem append_nil (s : wseq α) : append s nil = s := seq.append_nil _

@[simp] theorem append_assoc (s t u : wseq α) :
  append (append s t) u = append s (append t u) := seq.append_assoc _ _ _

@[simp] def tail.aux : option (α × wseq α) → computation (option (α × wseq α))
| none          := return none
| (some (a, s)) := destruct s

theorem destruct_tail (s : wseq α) :
  destruct (tail s) = destruct s >>= tail.aux :=
begin
  dsimp [tail], simp, rw [← bind_pure_comp_eq_map, is_lawful_monad.bind_assoc],
  apply congr_arg, funext o,
  rcases o with _|⟨a, s⟩;
  apply (@pure_bind computation _ _ _ _ _ _).trans _; simp
end

@[simp] def drop.aux : ℕ → option (α × wseq α) → computation (option (α × wseq α))
| 0     := return
| (n+1) := λ a, tail.aux a >>= drop.aux n

theorem drop.aux_none : ∀ n, @drop.aux α n none = return none
| 0     := rfl
| (n+1) := show computation.bind (return none) (drop.aux n) = return none,
           by rw [ret_bind, drop.aux_none]

theorem destruct_dropn :
  ∀ (s : wseq α) n, destruct (drop s n) = destruct s >>= drop.aux n
| s 0     := (bind_ret' _).symm
| s (n+1) := by rw [← dropn_tail, destruct_dropn _ n,
  destruct_tail, is_lawful_monad.bind_assoc]; refl

theorem head_terminates_of_head_tail_terminates (s : wseq α) [T : terminates (head (tail s))] :
  terminates (head s) :=
(head_terminates_iff _).2 $ begin
  cases (head_terminates_iff _).1 T with a h,
  simp [tail] at h,
  rcases exists_of_mem_bind h with ⟨s', h1, h2⟩,
  unfold functor.map at h1,
  exact let ⟨t, h3, h4⟩ := exists_of_mem_map h1 in terminates_of_mem h3
end

theorem destruct_some_of_destruct_tail_some {s : wseq α} {a}
  (h : some a ∈ destruct (tail s)) : ∃ a', some a' ∈ destruct s :=
begin
  unfold tail functor.map at h, simp at h,
  rcases exists_of_mem_bind h with ⟨t, tm, td⟩, clear h,
  rcases exists_of_mem_map tm with ⟨t', ht', ht2⟩, clear tm,
  cases t' with t'; rw ←ht2 at td; simp at td,
  { have := mem_unique td (ret_mem _), contradiction },
  { exact ⟨_, ht'⟩ }
end

theorem head_some_of_head_tail_some {s : wseq α} {a}
  (h : some a ∈ head (tail s)) : ∃ a', some a' ∈ head s :=
begin
  unfold head at h,
  rcases exists_of_mem_map h with ⟨o, md, e⟩, clear h,
  cases o with o; injection e with h', clear e h',
  cases destruct_some_of_destruct_tail_some md with a am,
  exact ⟨_, mem_map ((<$>) (@prod.fst α (wseq α))) am⟩
end

theorem head_some_of_nth_some {s : wseq α} {a n}
  (h : some a ∈ nth s n) : ∃ a', some a' ∈ head s :=
begin
  revert a, induction n with n IH; intros,
  exacts [⟨_, h⟩, let ⟨a', h'⟩ := head_some_of_head_tail_some h in IH h']
end

instance productive_tail (s : wseq α) [productive s] : productive (tail s) :=
λ n, by rw [nth_tail]; apply_instance

instance productive_dropn (s : wseq α) [productive s] (n) : productive (drop s n) :=
λ m, by rw [←nth_add]; apply_instance

/-- Given a productive weak sequence, we can collapse all the `think`s to
  produce a sequence. -/
def to_seq (s : wseq α) [productive s] : seq α :=
⟨λ n, (nth s n).get, λn h,
begin
  induction e : computation.get (nth s (n + 1)), {trivial},
  have := mem_of_get_eq _ e,
  simp [nth] at this h, cases head_some_of_head_tail_some this with a' h',
  have := mem_unique h' (@mem_of_get_eq _ _ _ _ h),
  contradiction
end⟩

theorem nth_terminates_le {s : wseq α} {m n} (h : m ≤ n) : terminates (nth s n) → terminates (nth s m) :=
by induction h with m' h IH; [exact id,
  exact λ T, IH (@head_terminates_of_head_tail_terminates _ _ T)]

theorem head_terminates_of_nth_terminates {s : wseq α} {n} : terminates (nth s n) → terminates (head s) :=
nth_terminates_le (nat.zero_le n)

theorem destruct_terminates_of_nth_terminates {s : wseq α} {n} (T : terminates (nth s n)) : terminates (destruct s) :=
(head_terminates_iff _).1 $ head_terminates_of_nth_terminates T

theorem mem_rec_on {C : wseq α → Prop} {a s} (M : a ∈ s)
  (h1 : ∀ b s', (a = b ∨ C s') → C (cons b s'))
  (h2 : ∀ s, C s → C (think s)) : C s :=
begin
  apply seq.mem_rec_on M,
  intros o s' h, cases o with b,
  { apply h2, cases h, {contradiction}, {assumption} },
  { apply h1, apply or.imp_left _ h, intro h, injection h }
end

@[simp] theorem mem_think (s : wseq α) (a) : a ∈ think s ↔ a ∈ s :=
begin
  cases s with f al,
  change some (some a) ∈ some none :: f ↔ some (some a) ∈ f,
  constructor; intro h,
  { apply (stream.eq_or_mem_of_mem_cons h).resolve_left,
    intro, injections },
  { apply stream.mem_cons_of_mem _ h }
end

theorem eq_or_mem_iff_mem {s : wseq α} {a a' s'} :
  some (a', s') ∈ destruct s → (a ∈ s ↔ a = a' ∨ a ∈ s') :=
begin
  generalize e : destruct s = c, intro h,
  revert s, apply computation.mem_rec_on h _ (λ c IH, _); intro s;
  apply s.cases_on _ (λ x s, _) (λ s, _); intros m;
  have := congr_arg computation.destruct m; simp at this;
  cases this with i1 i2,
  { rw [i1, i2],
    cases s' with f al,
    unfold cons has_mem.mem wseq.mem seq.mem seq.cons, simp,
    have h_a_eq_a' : a = a' ↔ some (some a) = some (some a'), {simp},
    rw [h_a_eq_a'],
    refine ⟨stream.eq_or_mem_of_mem_cons, λo, _⟩,
    { cases o with e m,
      { rw e, apply stream.mem_cons },
      { exact stream.mem_cons_of_mem _ m } } },
  { simp, exact IH this }
end

@[simp] theorem mem_cons_iff (s : wseq α) (b) {a} : a ∈ cons b s ↔ a = b ∨ a ∈ s :=
eq_or_mem_iff_mem $ by simp [ret_mem]

theorem mem_cons_of_mem {s : wseq α} (b) {a} (h : a ∈ s) : a ∈ cons b s :=
(mem_cons_iff _ _).2 (or.inr h)

theorem mem_cons (s : wseq α) (a) : a ∈ cons a s :=
(mem_cons_iff _ _).2 (or.inl rfl)

theorem mem_of_mem_tail {s : wseq α} {a} : a ∈ tail s → a ∈ s :=
begin
  intro h, have := h, cases h with n e, revert s, simp [stream.nth],
  induction n with n IH; intro s; apply s.cases_on _ (λx s, _) (λ s, _);
    repeat{simp}; intros m e; injections,
  { exact or.inr m },
  { exact or.inr m },
  { apply IH m, rw e, cases tail s, refl }
end

theorem mem_of_mem_dropn {s : wseq α} {a} : ∀ {n}, a ∈ drop s n → a ∈ s
| 0     h := h
| (n+1) h := @mem_of_mem_dropn n (mem_of_mem_tail h)

theorem nth_mem {s : wseq α} {a n} : some a ∈ nth s n → a ∈ s :=
begin
  revert s, induction n with n IH; intros s h,
  { rcases exists_of_mem_map h with ⟨o, h1, h2⟩,
    cases o with o; injection h2 with h',
    cases o with a' s',
    exact (eq_or_mem_iff_mem h1).2 (or.inl h'.symm) },
  { have := @IH (tail s), rw nth_tail at this,
    exact mem_of_mem_tail (this h) }
end

theorem exists_nth_of_mem {s : wseq α} {a} (h : a ∈ s) : ∃ n, some a ∈ nth s n :=
begin
  apply mem_rec_on h,
  { intros a' s' h, cases h with h h,
    { existsi 0, simp [nth], rw h, apply ret_mem },
    { cases h with n h, existsi n+1,
      simp [nth], exact h } },
  { intros s' h, cases h with n h,
    existsi n, simp [nth], apply think_mem h }
end

theorem exists_dropn_of_mem {s : wseq α} {a} (h : a ∈ s) :
  ∃ n s', some (a, s') ∈ destruct (drop s n) :=
let ⟨n, h⟩ := exists_nth_of_mem h in ⟨n, begin
  cases (head_terminates_iff _).1 ⟨_, h⟩ with o om,
  have := mem_unique (mem_map _ om) h,
  cases o with o; injection this with i,
  cases o with a' s', dsimp at i,
  rw i at om, exact ⟨_, om⟩
end⟩

theorem lift_rel_dropn_destruct {R : α → β → Prop} {s t} (H : lift_rel R s t) :
  ∀ n, computation.lift_rel (lift_rel_o R (lift_rel R))
    (destruct (drop s n)) (destruct (drop t n))
| 0     := lift_rel_destruct H
| (n+1) := begin
  simp [destruct_tail],
  apply lift_rel_bind,
  apply lift_rel_dropn_destruct n,
  exact λ a b o, match a, b, o with
  | none,       none,         _        := by simp
  | some (a, s), some (b, t), ⟨h1, h2⟩ := by simp [tail.aux]; apply lift_rel_destruct h2
  end
end

theorem exists_of_lift_rel_left {R : α → β → Prop} {s t}
  (H : lift_rel R s t) {a} (h : a ∈ s) : ∃ {b}, b ∈ t ∧ R a b :=
let ⟨n, h⟩ := exists_nth_of_mem h,
    ⟨some (._, s'), sd, rfl⟩ := exists_of_mem_map h,
    ⟨some (b, t'), td, ⟨ab, _⟩⟩ := (lift_rel_dropn_destruct H n).left sd in
⟨b, nth_mem (mem_map ((<$>) prod.fst.{v v}) td), ab⟩

theorem exists_of_lift_rel_right {R : α → β → Prop} {s t}
  (H : lift_rel R s t) {b} (h : b ∈ t) : ∃ {a}, a ∈ s ∧ R a b :=
by rw ←lift_rel.swap at H; exact exists_of_lift_rel_left H h

theorem head_terminates_of_mem {s : wseq α} {a} (h : a ∈ s) : terminates (head s) :=
let ⟨n, h⟩ := exists_nth_of_mem h in head_terminates_of_nth_terminates ⟨_, h⟩

theorem of_mem_append {s₁ s₂ : wseq α} {a : α} : a ∈ append s₁ s₂ → a ∈ s₁ ∨ a ∈ s₂ :=
seq.of_mem_append

theorem mem_append_left {s₁ s₂ : wseq α} {a : α} : a ∈ s₁ → a ∈ append s₁ s₂ :=
seq.mem_append_left

theorem exists_of_mem_map {f} {b : β} : ∀ {s : wseq α}, b ∈ map f s → ∃ a, a ∈ s ∧ f a = b
| ⟨g, al⟩ h := let ⟨o, om, oe⟩ := seq.exists_of_mem_map h in
  by cases o with a; injection oe with h'; exact ⟨a, om, h'⟩

@[simp] theorem lift_rel_nil (R : α → β → Prop) : lift_rel R nil nil :=
by rw [lift_rel_destruct_iff]; simp

@[simp] theorem lift_rel_cons (R : α → β → Prop) (a b s t) :
  lift_rel R (cons a s) (cons b t) ↔ R a b ∧ lift_rel R s t :=
by rw [lift_rel_destruct_iff]; simp

@[simp] theorem lift_rel_think_left (R : α → β → Prop) (s t) :
  lift_rel R (think s) t ↔ lift_rel R s t :=
by rw [lift_rel_destruct_iff, lift_rel_destruct_iff]; simp

@[simp] theorem lift_rel_think_right (R : α → β → Prop) (s t) :
  lift_rel R s (think t) ↔ lift_rel R s t :=
by rw [lift_rel_destruct_iff, lift_rel_destruct_iff]; simp

theorem cons_congr {s t : wseq α} (a : α) (h : s ~ t) : cons a s ~ cons a t :=
by unfold equiv; simp; exact h

theorem think_equiv (s : wseq α) : think s ~ s :=
by unfold equiv; simp; apply equiv.refl

theorem think_congr {s t : wseq α} (a : α) (h : s ~ t) : think s ~ think t :=
by unfold equiv; simp; exact h

theorem head_congr : ∀ {s t : wseq α}, s ~ t → head s ~ head t :=
suffices ∀ {s t : wseq α}, s ~ t → ∀ {o}, o ∈ head s → o ∈ head t, from
λ s t h o, ⟨this h, this h.symm⟩,
begin
  intros s t h o ho,
  rcases @computation.exists_of_mem_map _ _ _ _ (destruct s) ho with ⟨ds, dsm, dse⟩,
  rw ←dse,
  cases destruct_congr h with l r,
  rcases l dsm with ⟨dt, dtm, dst⟩,
  cases ds with a; cases dt with b,
  { apply mem_map _ dtm },
  { cases b, cases dst },
  { cases a, cases dst },
  { cases a with a s', cases b with b t', rw dst.left,
    exact @mem_map _ _ (@functor.map _ _ (α × wseq α) _ prod.fst)
      _ (destruct t) dtm }
end

theorem flatten_equiv {c : computation (wseq α)} {s} (h : s ∈ c) : flatten c ~ s :=
begin
  apply computation.mem_rec_on h, { simp },
  { intro s', apply equiv.trans, simp [think_equiv] }
end

theorem lift_rel_flatten {R : α → β → Prop} {c1 : computation (wseq α)} {c2 : computation (wseq β)}
  (h : c1.lift_rel (lift_rel R) c2) : lift_rel R (flatten c1) (flatten c2) :=
let S := λ s t,
  ∃ c1 c2, s = flatten c1 ∧ t = flatten c2 ∧ computation.lift_rel (lift_rel R) c1 c2 in
⟨S, ⟨c1, c2, rfl, rfl, h⟩, λ s t h,
  match s, t, h with ._, ._, ⟨c1, c2, rfl, rfl, h⟩ := begin
    simp, apply lift_rel_bind _ _ h,
    intros a b ab, apply computation.lift_rel.imp _ _ _ (lift_rel_destruct ab),
    intros a b, apply lift_rel_o.imp_right,
    intros s t h, refine ⟨return s, return t, _, _, _⟩; simp [h]
  end end⟩

theorem flatten_congr {c1 c2 : computation (wseq α)} :
  computation.lift_rel equiv c1 c2 → flatten c1 ~ flatten c2 := lift_rel_flatten

theorem tail_congr {s t : wseq α} (h : s ~ t) : tail s ~ tail t :=
begin
  apply flatten_congr,
  unfold functor.map, rw [←bind_ret, ←bind_ret],
  apply lift_rel_bind _ _ (destruct_congr h),
  intros a b h, simp,
  cases a with a; cases b with b,
  { trivial },
  { cases h },
  { cases a, cases h },
  { cases a with a s', cases b with b t', exact h.right }
end

theorem dropn_congr {s t : wseq α} (h : s ~ t) (n) : drop s n ~ drop t n :=
by induction n; simp [*, tail_congr]

theorem nth_congr {s t : wseq α} (h : s ~ t) (n) : nth s n ~ nth t n :=
head_congr (dropn_congr h _)

theorem mem_congr {s t : wseq α} (h : s ~ t) (a) : a ∈ s ↔ a ∈ t :=
suffices ∀ {s t : wseq α}, s ~ t → a ∈ s → a ∈ t, from ⟨this h, this h.symm⟩,
λ s t h as, let ⟨n, hn⟩ := exists_nth_of_mem as in
nth_mem ((nth_congr h _ _).1 hn)

theorem productive_congr {s t : wseq α} (h : s ~ t) : productive s ↔ productive t :=
forall_congr $ λn, terminates_congr $ nth_congr h _

theorem equiv.ext {s t : wseq α} (h : ∀ n, nth s n ~ nth t n) : s ~ t :=
⟨λ s t, ∀ n, nth s n ~ nth t n, h, λs t h, begin
  refine lift_rel_def.2 ⟨_, _⟩,
  { rw [←head_terminates_iff, ←head_terminates_iff],
    exact terminates_congr (h 0) },
  { intros a b ma mb,
    cases a with a; cases b with b,
    { trivial },
    { injection mem_unique (mem_map _ ma) ((h 0 _).2 (mem_map _ mb)) },
    { injection mem_unique (mem_map _ ma) ((h 0 _).2 (mem_map _ mb)) },
    { cases a with a s', cases b with b t',
      injection mem_unique (mem_map _ ma) ((h 0 _).2 (mem_map _ mb)) with ab,
      refine ⟨ab, λ n, _⟩,
      refine (nth_congr (flatten_equiv (mem_map _ ma)) n).symm.trans
        ((_ : nth (tail s) n ~ nth (tail t) n).trans
        (nth_congr (flatten_equiv (mem_map _ mb)) n)),
      rw [nth_tail, nth_tail], apply h } }
end⟩

theorem length_eq_map (s : wseq α) : length s = computation.map list.length (to_list s) :=
begin
  refine eq_of_bisim
    (λ c1 c2, ∃ (l : list α) (s : wseq α),
      c1 = corec length._match_2 (l.length, s) ∧
      c2 = computation.map list.length (corec to_list._match_2 (l, s)))
    _ ⟨[], s, rfl, rfl⟩,
  intros s1 s2 h, rcases h with ⟨l, s, h⟩, rw [h.left, h.right],
  apply s.cases_on _ (λ a s, _) (λ s, _);
    repeat {simp [to_list, nil, cons, think, length]},
  { refine ⟨a::l, s, _, _⟩; simp },
  { refine ⟨l, s, _, _⟩; simp }
end

@[simp] theorem of_list_nil : of_list [] = (nil : wseq α) := rfl

@[simp] theorem of_list_cons (a : α) (l) :
  of_list (a :: l) = cons a (of_list l) :=
show seq.map some (seq.of_list (a :: l)) =
     seq.cons (some a) (seq.map some (seq.of_list l)), by simp

@[simp] theorem to_list'_nil (l : list α) :
  corec to_list._match_2 (l, nil) = return l.reverse :=
destruct_eq_ret rfl

@[simp] theorem to_list'_cons (l : list α) (s : wseq α) (a : α) :
  corec to_list._match_2 (l, cons a s) =
  (corec to_list._match_2 (a::l, s)).think :=
destruct_eq_think $ by simp [to_list, cons]

@[simp] theorem to_list'_think (l : list α) (s : wseq α) :
  corec to_list._match_2 (l, think s) =
  (corec to_list._match_2 (l, s)).think :=
destruct_eq_think $ by simp [to_list, think]

theorem to_list'_map (l : list α) (s : wseq α) :
  corec to_list._match_2 (l, s) =
  ((++) l.reverse) <$> to_list s :=
begin
  refine eq_of_bisim
    (λ c1 c2, ∃ (l' : list α) (s : wseq α),
      c1 = corec to_list._match_2 (l' ++ l, s) ∧
      c2 = computation.map ((++) l.reverse) (corec to_list._match_2 (l', s)))
    _ ⟨[], s, rfl, rfl⟩,
  intros s1 s2 h, rcases h with ⟨l', s, h⟩, rw [h.left, h.right],
  apply s.cases_on _ (λ a s, _) (λ s, _);
    repeat {simp [to_list, nil, cons, think, length]},
  { refine ⟨a::l', s, _, _⟩; simp },
  { refine ⟨l', s, _, _⟩; simp }
end

@[simp] theorem to_list_cons (a : α) (s) :
  to_list (cons a s) = (list.cons a <$> to_list s).think :=
destruct_eq_think $ by unfold to_list; simp; rw to_list'_map; simp; refl

@[simp] theorem to_list_nil : to_list (nil : wseq α) = return [] :=
destruct_eq_ret rfl

theorem to_list_of_list (l : list α) : l ∈ to_list (of_list l) :=
by induction l with a l IH; simp [ret_mem]; exact think_mem (mem_map _ IH)

@[simp] theorem destruct_of_seq (s : seq α) :
  destruct (of_seq s) = return (s.head.map $ λ a, (a, of_seq s.tail)) :=
destruct_eq_ret $ begin
  simp [of_seq, head, destruct, seq.destruct, seq.head],
  rw [show seq.nth (some <$> s) 0 = some <$> seq.nth s 0, by apply seq.map_nth],
  cases seq.nth s 0 with a, { refl },
  unfold functor.map,
  simp [destruct]
end

@[simp] theorem head_of_seq (s : seq α) : head (of_seq s) = return s.head :=
by simp [head]; cases seq.head s; refl

@[simp] theorem tail_of_seq (s : seq α) : tail (of_seq s) = of_seq s.tail :=
begin
  simp [tail], apply s.cases_on _ (λ x s, _); simp [of_seq], {refl},
  rw [seq.head_cons, seq.tail_cons], refl
end

@[simp] theorem dropn_of_seq (s : seq α) : ∀ n, drop (of_seq s) n = of_seq (s.drop n)
| 0 := rfl
| (n+1) := by dsimp [drop]; rw [dropn_of_seq, tail_of_seq]

theorem nth_of_seq (s : seq α) (n) : nth (of_seq s) n = return (seq.nth s n) :=
by dsimp [nth]; rw [dropn_of_seq, head_of_seq, seq.head_dropn]

instance productive_of_seq (s : seq α) : productive (of_seq s) :=
λ n, by rw nth_of_seq; apply_instance

theorem to_seq_of_seq (s : seq α) : to_seq (of_seq s) = s :=
begin
  apply subtype.eq, funext n,
  dsimp [to_seq], apply get_eq_of_mem,
  rw nth_of_seq, apply ret_mem
end

/-- The monadic `return a` is a singleton list containing `a`. -/
def ret (a : α) : wseq α := of_list [a]

@[simp] theorem map_nil (f : α → β) : map f nil = nil := rfl

@[simp] theorem map_cons (f : α → β) (a s) :
  map f (cons a s) = cons (f a) (map f s) := seq.map_cons _ _ _

@[simp] theorem map_think (f : α → β) (s) :
  map f (think s) = think (map f s) := seq.map_cons _ _ _

@[simp] theorem map_id (s : wseq α) : map id s = s := by simp [map]

@[simp] theorem map_ret (f : α → β) (a) : map f (ret a) = ret (f a) := by simp [ret]

@[simp] theorem map_append (f : α → β) (s t) : map f (append s t) = append (map f s) (map f t) :=
seq.map_append _ _ _

theorem map_comp (f : α → β) (g : β → γ) (s : wseq α) :
  map (g ∘ f) s = map g (map f s) :=
begin
  dsimp [map], rw ←seq.map_comp,
  apply congr_fun, apply congr_arg,
  funext o, cases o; refl
end

theorem mem_map (f : α → β) {a : α} {s : wseq α} : a ∈ s → f a ∈ map f s :=
seq.mem_map (option.map f)

-- The converse is not true without additional assumptions
theorem exists_of_mem_join {a : α} : ∀ {S : wseq (wseq α)}, a ∈ join S → ∃ s, s ∈ S ∧ a ∈ s :=
suffices ∀ ss : wseq α, a ∈ ss → ∀ s S, append s (join S) = ss →
  a ∈ append s (join S) → a ∈ s ∨ ∃ s, s ∈ S ∧ a ∈ s, from λ S h,
  (this _ h nil S (by simp) (by simp [h])).resolve_left (not_mem_nil _),
begin
  intros ss h, apply mem_rec_on h (λ b ss o, _) (λ ss IH, _); intros s S,
  { refine s.cases_on (S.cases_on _ (λ s S, _) (λ S, _)) (λ b' s, _) (λ s, _);
    intros ej m; simp at ej;
    have := congr_arg seq.destruct ej; simp at this;
    try {cases this}; try {contradiction},
    substs b' ss,
    simp at m ⊢,
    cases o with e IH, { simp [e] },
    cases m with e m, { simp [e] },
    exact or.imp_left or.inr (IH _ _ rfl m) },
  { refine s.cases_on (S.cases_on _ (λ s S, _) (λ S, _)) (λ b' s, _) (λ s, _);
    intros ej m; simp at ej;
    have := congr_arg seq.destruct ej; simp at this;
    try { try {have := this.1}, contradiction }; subst ss,
    { apply or.inr, simp at m ⊢,
      cases IH s S rfl m with as ex,
      { exact ⟨s, or.inl rfl, as⟩ },
      { rcases ex with ⟨s', sS, as⟩,
        exact ⟨s', or.inr sS, as⟩ } },
    { apply or.inr, simp at m,
      rcases (IH nil S (by simp) (by simp [m])).resolve_left (not_mem_nil _) with ⟨s, sS, as⟩,
      exact ⟨s, by simp [sS], as⟩ },
    { simp at m IH ⊢, apply IH _ _ rfl m } }
end

theorem exists_of_mem_bind {s : wseq α} {f : α → wseq β} {b}
  (h : b ∈ bind s f) : ∃ a ∈ s, b ∈ f a :=
let ⟨t, tm, bt⟩ := exists_of_mem_join h,
    ⟨a, as, e⟩ := exists_of_mem_map tm in ⟨a, as, by rwa e⟩

theorem destruct_map (f : α → β) (s : wseq α) :
  destruct (map f s) = computation.map (option.map (prod.map f (map f))) (destruct s) :=
begin
  apply eq_of_bisim (λ c1 c2, ∃ s, c1 = destruct (map f s) ∧
    c2 = computation.map (option.map (prod.map f (map f))) (destruct s)),
  { intros c1 c2 h, cases h with s h, rw [h.left, h.right],
    apply s.cases_on _ (λ a s, _) (λ s, _); simp; simp,
    exact ⟨s, rfl, rfl⟩ },
  { exact ⟨s, rfl, rfl⟩ }
end

theorem lift_rel_map {δ} (R : α → β → Prop) (S : γ → δ → Prop)
  {s1 : wseq α} {s2 : wseq β}
  {f1 : α → γ} {f2 : β → δ}
  (h1 : lift_rel R s1 s2) (h2 : ∀ {a b}, R a b → S (f1 a) (f2 b))
  : lift_rel S (map f1 s1) (map f2 s2) :=
⟨λ s1 s2, ∃ s t, s1 = map f1 s ∧ s2 = map f2 t ∧ lift_rel R s t,
⟨s1, s2, rfl, rfl, h1⟩,
λ s1 s2 h, match s1, s2, h with ._, ._, ⟨s, t, rfl, rfl, h⟩ := begin
  simp [destruct_map], apply computation.lift_rel_map _ _ (lift_rel_destruct h),
  intros o p h,
  cases o with a; cases p with b; simp,
  { cases b; cases h },
  { cases a; cases h },
  { cases a with a s; cases b with b t, cases h with r h,
    exact ⟨h2 r, s, rfl, t, rfl, h⟩ }
end end⟩

theorem map_congr (f : α → β) {s t : wseq α} (h : s ~ t) : map f s ~ map f t :=
lift_rel_map _ _ h (λ _ _, congr_arg _)

@[simp] def destruct_append.aux (t : wseq α) :
  option (α × wseq α) → computation (option (α × wseq α))
| none          := destruct t
| (some (a, s)) := return (some (a, append s t))

theorem destruct_append (s t : wseq α) :
  destruct (append s t) = (destruct s).bind (destruct_append.aux t) :=
begin
  apply eq_of_bisim (λ c1 c2, ∃ s t, c1 = destruct (append s t) ∧
    c2 = (destruct s).bind (destruct_append.aux t)) _ ⟨s, t, rfl, rfl⟩,
  intros c1 c2 h, rcases h with ⟨s, t, h⟩, rw [h.left, h.right],
  apply s.cases_on _ (λ a s, _) (λ s, _); simp; simp,
  { apply t.cases_on _ (λ b t, _) (λ t, _); simp; simp,
    { refine ⟨nil, t, _, _⟩; simp } },
  { exact ⟨s, t, rfl, rfl⟩ }
end

@[simp] def destruct_join.aux : option (wseq α × wseq (wseq α)) → computation (option (α × wseq α))
| none          := return none
| (some (s, S)) := (destruct (append s (join S))).think

theorem destruct_join (S : wseq (wseq α)) :
  destruct (join S) = (destruct S).bind destruct_join.aux :=
begin
  apply eq_of_bisim (λ c1 c2, c1 = c2 ∨ ∃ S, c1 = destruct (join S) ∧
    c2 = (destruct S).bind destruct_join.aux) _ (or.inr ⟨S, rfl, rfl⟩),
  intros c1 c2 h, exact match c1, c2, h with
  | _, _, (or.inl $ eq.refl c) := by cases c.destruct; simp
  | _, _, or.inr ⟨S, rfl, rfl⟩ := begin
    apply S.cases_on _ (λ s S, _) (λ S, _); simp; simp,
    { refine or.inr ⟨S, rfl, rfl⟩ }
  end end
end

theorem lift_rel_append (R : α → β → Prop) {s1 s2 : wseq α} {t1 t2 : wseq β}
  (h1 : lift_rel R s1 t1) (h2 : lift_rel R s2 t2) :
  lift_rel R (append s1 s2) (append t1 t2) :=
⟨λ s t, lift_rel R s t ∨ ∃ s1 t1, s = append s1 s2 ∧ t = append t1 t2 ∧ lift_rel R s1 t1,
or.inr ⟨s1, t1, rfl, rfl, h1⟩,
λ s t h, match s, t, h with
| s, t, or.inl h := begin
    apply computation.lift_rel.imp _ _ _ (lift_rel_destruct h),
    intros a b, apply lift_rel_o.imp_right,
    intros s t, apply or.inl
  end
| ._, ._, or.inr ⟨s1, t1, rfl, rfl, h⟩ := begin
    simp [destruct_append],
    apply computation.lift_rel_bind _ _ (lift_rel_destruct h),
    intros o p h,
    cases o with a; cases p with b,
    { simp, apply computation.lift_rel.imp _ _ _ (lift_rel_destruct h2),
      intros a b, apply lift_rel_o.imp_right,
      intros s t, apply or.inl },
    { cases b; cases h },
    { cases a; cases h },
    { cases a with a s; cases b with b t, cases h with r h,
      simp, exact ⟨r, or.inr ⟨s, rfl, t, rfl, h⟩⟩ }
  end
end⟩

theorem lift_rel_join.lem (R : α → β → Prop) {S T} {U : wseq α → wseq β → Prop}
  (ST : lift_rel (lift_rel R) S T) (HU : ∀ s1 s2, (∃ s t S T,
      s1 = append s (join S) ∧ s2 = append t (join T) ∧
      lift_rel R s t ∧ lift_rel (lift_rel R) S T) → U s1 s2) {a} (ma : a ∈ destruct (join S)) :
  ∃ {b}, b ∈ destruct (join T) ∧ lift_rel_o R U a b :=
begin
  cases exists_results_of_mem ma with n h, clear ma, revert a S T,
  apply nat.strong_induction_on n _,
  intros n IH a S T ST ra, simp [destruct_join] at ra, exact
  let ⟨o, m, k, rs1, rs2, en⟩ := of_results_bind ra,
      ⟨p, mT, rop⟩ := computation.exists_of_lift_rel_left (lift_rel_destruct ST) rs1.mem in
  by exact match o, p, rop, rs1, rs2, mT with
  | none, none, _, rs1, rs2, mT := by simp [destruct_join]; exact
    ⟨none, mem_bind mT (ret_mem _), by rw eq_of_ret_mem rs2.mem; trivial⟩
  | some (s, S'), some (t, T'), ⟨st, ST'⟩, rs1, rs2, mT :=
    by simp [destruct_append] at rs2; exact
    let ⟨k1, rs3, ek⟩ := of_results_think rs2,
        ⟨o', m1, n1, rs4, rs5, ek1⟩ := of_results_bind rs3,
        ⟨p', mt, rop'⟩ := computation.exists_of_lift_rel_left (lift_rel_destruct st) rs4.mem in
    by exact match o', p', rop', rs4, rs5, mt with
    | none, none, _, rs4, rs5', mt :=
      have n1 < n, begin
        rw [en, ek, ek1],
        apply lt_of_lt_of_le _ (nat.le_add_right _ _),
        apply nat.lt_succ_of_le (nat.le_add_right _ _)
      end,
      let ⟨ob, mb, rob⟩ := IH _ this ST' rs5' in by refine ⟨ob, _, rob⟩;
      { simp [destruct_join], apply mem_bind mT, simp [destruct_append],
        apply think_mem, apply mem_bind mt, exact mb }
    | some (a, s'), some (b, t'), ⟨ab, st'⟩, rs4, rs5, mt := begin
      simp at rs5,
      refine ⟨some (b, append t' (join T')), _, _⟩,
      { simp [destruct_join], apply mem_bind mT, simp [destruct_append],
        apply think_mem, apply mem_bind mt, apply ret_mem },
      rw eq_of_ret_mem rs5.mem,
      exact ⟨ab, HU _ _ ⟨s', t', S', T', rfl, rfl, st', ST'⟩⟩
    end end
  end
end

theorem lift_rel_join (R : α → β → Prop) {S : wseq (wseq α)} {T : wseq (wseq β)}
  (h : lift_rel (lift_rel R) S T) : lift_rel R (join S) (join T) :=
⟨λ s1 s2, ∃ s t S T,
  s1 = append s (join S) ∧ s2 = append t (join T) ∧
  lift_rel R s t ∧ lift_rel (lift_rel R) S T,
  ⟨nil, nil, S, T, by simp, by simp, by simp, h⟩,
λs1 s2 ⟨s, t, S, T, h1, h2, st, ST⟩, begin
  clear _fun_match _x,
  rw [h1, h2], rw [destruct_append, destruct_append],
  apply computation.lift_rel_bind _ _ (lift_rel_destruct st),
  exact λ o p h, match o, p, h with
  | some (a, s), some (b, t), ⟨h1, h2⟩ :=
    by simp; exact ⟨h1, s, t, S, rfl, T, rfl, h2, ST⟩
  | none, none, _ := begin
    dsimp [destruct_append.aux, computation.lift_rel], constructor,
    { intro, apply lift_rel_join.lem _ ST (λ _ _, id) },
    { intros b mb,
      rw [←lift_rel_o.swap], apply lift_rel_join.lem (function.swap R),
      { rw [←lift_rel.swap R, ←lift_rel.swap], apply ST },
      { rw [←lift_rel.swap R, ←lift_rel.swap (lift_rel R)],
        exact λ s1 s2 ⟨s, t, S, T, h1, h2, st, ST⟩,
                      ⟨t, s, T, S, h2, h1, st, ST⟩ },
      { exact mb } }
  end end
end⟩

theorem join_congr {S T : wseq (wseq α)} (h : lift_rel equiv S T) : join S ~ join T :=
lift_rel_join _ h

theorem lift_rel_bind {δ} (R : α → β → Prop) (S : γ → δ → Prop)
  {s1 : wseq α} {s2 : wseq β}
  {f1 : α → wseq γ} {f2 : β → wseq δ}
  (h1 : lift_rel R s1 s2) (h2 : ∀ {a b}, R a b → lift_rel S (f1 a) (f2 b))
  : lift_rel S (bind s1 f1) (bind s2 f2) :=
lift_rel_join _ (lift_rel_map _ _ h1 @h2)

theorem bind_congr {s1 s2 : wseq α} {f1 f2 : α → wseq β}
  (h1 : s1 ~ s2) (h2 : ∀ a, f1 a ~ f2 a) : bind s1 f1 ~ bind s2 f2 :=
lift_rel_bind _ _ h1 (λ a b h, by rw h; apply h2)

@[simp] theorem join_ret (s : wseq α) : join (ret s) ~ s :=
by simp [ret]; apply think_equiv

@[simp] theorem join_map_ret (s : wseq α) : join (map ret s) ~ s :=
begin
  refine ⟨λ s1 s2, join (map ret s2) = s1, rfl, _⟩,
  intros s' s h, rw ←h,
  apply lift_rel_rec
    (λ c1 c2, ∃ s,
      c1 = destruct (join (map ret s)) ∧ c2 = destruct s),
  { exact λ c1 c2 h, match c1, c2, h with
    | ._, ._, ⟨s, rfl, rfl⟩ := begin
      clear h _match,
      apply s.cases_on _ (λ a s, _) (λ s, _); simp [ret]; simp [ret],
      { refine ⟨_, ret_mem _, _⟩, simp },
      { exact ⟨s, rfl, rfl⟩ }
    end end },
  { exact ⟨s, rfl, rfl⟩ }
end

@[simp] theorem join_append (S T : wseq (wseq α)) :
  join (append S T) ~ append (join S) (join T) :=
begin
  refine ⟨λ s1 s2, ∃ s S T,
    s1 = append s (join (append S T)) ∧
    s2 = append s (append (join S) (join T)), ⟨nil, S, T, by simp, by simp⟩, _⟩,
  intros s1 s2 h,
  apply lift_rel_rec (λ c1 c2, ∃ (s : wseq α) S T,
    c1 = destruct (append s (join (append S T))) ∧
    c2 = destruct (append s (append (join S) (join T)))) _ _ _
    (let ⟨s, S, T, h1, h2⟩ := h in
         ⟨s, S, T, congr_arg destruct h1, congr_arg destruct h2⟩),
  intros c1 c2 h,
  exact match c1, c2, h with ._, ._, ⟨s, S, T, rfl, rfl⟩ := begin
    clear _match h h,
    apply wseq.cases_on s _ (λ a s, _) (λ s, _); simp; simp,
    { apply wseq.cases_on S _ (λ s S, _) (λ S, _); simp; simp,
      { apply wseq.cases_on T _ (λ s T, _) (λ T, _); simp; simp,
        { refine ⟨s, nil, T, _, _⟩; simp },
        { refine ⟨nil, nil, T, _, _⟩; simp } },
      { exact ⟨s, S, T, rfl, rfl⟩ },
      { refine ⟨nil, S, T, _, _⟩; simp } },
    { exact ⟨s, S, T, rfl, rfl⟩ },
    { exact ⟨s, S, T, rfl, rfl⟩ }
  end end
end

@[simp] theorem bind_ret (f : α → β) (s) : bind s (ret ∘ f) ~ map f s :=
begin
  dsimp [bind], change (λx, ret (f x)) with (ret ∘ f),
  rw [map_comp], apply join_map_ret
end

@[simp] theorem ret_bind (a : α) (f : α → wseq β) :
  bind (ret a) f ~ f a := by simp [bind]

@[simp] theorem map_join (f : α → β) (S) :
  map f (join S) = join (map (map f) S) :=
begin
  apply seq.eq_of_bisim (λs1 s2,
    ∃ s S, s1 = append s (map f (join S)) ∧
      s2 = append s (join (map (map f) S))),
  { intros s1 s2 h,
    exact match s1, s2, h with ._, ._, ⟨s, S, rfl, rfl⟩ := begin
      apply wseq.cases_on s _ (λ a s, _) (λ s, _); simp; simp,
      { apply wseq.cases_on S _ (λ s S, _) (λ S, _); simp; simp,
        { exact ⟨map f s, S, rfl, rfl⟩ },
        { refine ⟨nil, S, _, _⟩; simp } },
      { exact ⟨_, _, rfl, rfl⟩ },
      { exact ⟨_, _, rfl, rfl⟩ }
    end end },
  { refine ⟨nil, S, _, _⟩; simp }
end

@[simp] theorem join_join (SS : wseq (wseq (wseq α))) :
  join (join SS) ~ join (map join SS) :=
begin
  refine ⟨λ s1 s2, ∃ s S SS,
    s1 = append s (join (append S (join SS))) ∧
    s2 = append s (append (join S) (join (map join SS))),
    ⟨nil, nil, SS, by simp, by simp⟩, _⟩,
  intros s1 s2 h,
  apply lift_rel_rec (λ c1 c2, ∃ s S SS,
      c1 = destruct (append s (join (append S (join SS)))) ∧
      c2 = destruct (append s (append (join S) (join (map join SS)))))
    _ (destruct s1) (destruct s2)
    (let ⟨s, S, SS, h1, h2⟩ := h in ⟨s, S, SS, by simp [h1], by simp [h2]⟩),
  intros c1 c2 h,
  exact match c1, c2, h with ._, ._, ⟨s, S, SS, rfl, rfl⟩ := begin
    clear _match h h,
    apply wseq.cases_on s _ (λ a s, _) (λ s, _); simp; simp,
    { apply wseq.cases_on S _ (λ s S, _) (λ S, _); simp; simp,
      { apply wseq.cases_on SS _ (λ S SS, _) (λ SS, _); simp; simp,
        { refine ⟨nil, S, SS, _, _⟩; simp },
        { refine ⟨nil, nil, SS, _, _⟩; simp } },
      { exact ⟨s, S, SS, rfl, rfl⟩ },
      { refine ⟨nil, S, SS, _, _⟩; simp } },
    { exact ⟨s, S, SS, rfl, rfl⟩ },
    { exact ⟨s, S, SS, rfl, rfl⟩ }
  end end
end

@[simp] theorem bind_assoc (s : wseq α) (f : α → wseq β) (g : β → wseq γ) :
  bind (bind s f) g ~ bind s (λ (x : α), bind (f x) g) :=
begin
  simp [bind], rw [← map_comp f (map g), map_comp (map g ∘ f) join],
  apply join_join
end

instance : monad wseq :=
{ map  := @map,
  pure := @ret,
  bind := @bind }

/-
  Unfortunately, wseq is not a lawful monad, because it does not satisfy
  the monad laws exactly, only up to sequence equivalence.
  Furthermore, even quotienting by the equivalence is not sufficient,
  because the join operation involves lists of quotient elements,
  with a lifted equivalence relation, and pure quotients cannot handle
  this type of construction.

instance : is_lawful_monad wseq :=
{ id_map := @map_id,
  bind_pure_comp_eq_map := @bind_ret,
  pure_bind := @ret_bind,
  bind_assoc := @bind_assoc }
-/

end wseq


import FOL.completeness

universes u v

namespace fol
open_locale logic_symbol

variables {L : language}

namespace formula

def bfal_le [has_preceq (term L) (formula L)] (t : term L) (p : formula L) : formula L :=
∀.((#0 ≼ (t^1)) ⟶ p)

notation `∀.{≼ ` t `} `:64 p := bfal_le t p

def bex_le [has_preceq (term L) (formula L)] (t : term L) (p : formula L) : formula L :=
∼∀.{≼ t}∼p

notation `∃.{≼ ` t `} `:64 p := bex_le t p

def bfal_lt [has_prec (term L) (formula L)] (t : term L) (p : formula L) : formula L :=
∀.((#0 ≺ t^1) ⟶ p)

notation `∀.{≺ ` t `} `:64 p := bfal_lt t p

def bex_lt [has_prec (term L) (formula L)] (t : term L) (p : formula L) : formula L :=
∼∀.{≺ t}∼p

notation `∃.{≺ ` t `} `:64 p := bex_lt t p

def bfal_mem [has_elem (term L) (formula L)] (t : term L) (p : formula L) : formula L :=
∀.((#0 ∊ t^1) ⟶ p)

notation `∀.{∊ ` t `} `:64 p := bfal_mem t p

def bex_mem [has_elem (term L) (formula L)] (t : term L) (p : formula L) : formula L :=
∃.((#0 ∊ t^1) ⊓ p)

notation `∃.{∊ ` t `} `:64 p := bex_mem t p

@[simp] lemma bfal_le_rew [has_le_symbol L] (t : term L) (p : formula L) (s) :
  (∀.{≼ t} p).rew s = ∀.{≼ t.rew s} p.rew (s^1) :=
by simp[bfal_le, term.pow_rew_distrib]

@[simp] lemma bex_le_rew [has_le_symbol L] (t : term L) (p : formula L) (s) :
  (∃.{≼ t} p).rew s = ∃.{≼ t.rew s} p.rew (s^1) :=
by simp[bex_le, term.pow_rew_distrib]

@[simp] def fal_rank : formula L → ℕ
| (p ⟶ q) := max (fal_rank p) (fal_rank q)
| (∼p)     := fal_rank p
| (∀.p)  := fal_rank p + 1
| _        := 0

@[simp] lemma fal_rank_and (p q : formula L) :
  fal_rank (p ⊓ q) = max (fal_rank p) (fal_rank q) := rfl

@[simp] lemma fal_rank_or (p q : formula L) :
  fal_rank (p ⊔ q) = max (fal_rank p) (fal_rank q) := rfl

@[simp] lemma fal_rank_ex (p : formula L) :
  fal_rank (∃.p) = fal_rank p + 1 := rfl

@[simp] lemma fal_rank_eq (t u : term L) :
  fal_rank (t =' u : formula L) = 0 := rfl

@[simp] lemma fal_rank_le [has_le_symbol L] (t u : term L) :
  fal_rank (t ≼ u : formula L) = 0 := rfl

@[simp] lemma fal_rank_mem [has_mem_symbol L] (t u : term L) :
  fal_rank (t ∊ u : formula L) = 0 := rfl

@[simp] def binary_inv : formula L → option (formula L × formula L)
| (p ⟶ q) := some (p, q)
| _        := none

@[simp] def unary_inv : formula L → option (formula L)
| ∼p := some p
| ∀.p := some p
| _  := none

@[simp] def quantifier_fn_aux : ℕ → (term L → formula L) → formula L → formula L
| s f ⊤        := ⊤
| s f (p ⟶ q) := quantifier_fn_aux s (λ t, (f t).binary_inv.iget.1) p ⟶ quantifier_fn_aux s (λ t, (f t).binary_inv.iget.2) q
| s f ∼p       := ∼quantifier_fn_aux s (λ t, (f t).unary_inv.iget) p
| s f (∀.p)  := ∀.quantifier_fn_aux (s + 1) (λ t, (f t).unary_inv.iget) p
| s f _        := f #s

@[simp] lemma quantifier_fn_aux_imply (s) (f g : term L → formula L) (p q : formula L) :
  quantifier_fn_aux s (λ x, f x ⟶ g x) (p ⟶ q) = quantifier_fn_aux s f p ⟶ quantifier_fn_aux s g q := rfl

@[simp] lemma quantifier_fn_aux_neg (s) (f : term L → formula L) (p : formula L) :
  quantifier_fn_aux s (λ x, ∼f x) (∼p) = ∼quantifier_fn_aux s f p := rfl

@[simp] lemma quantifier_fn_aux_fal (s) (f : term L → formula L) (p : formula L) :
  quantifier_fn_aux s (λ x, ∀.(f x)) (∀.p) = ∀.quantifier_fn_aux (s + 1) f p := rfl

@[simp] lemma quantifier_fn_aux_and (s) (f g : term L → formula L) (p q : formula L) :
  quantifier_fn_aux s (λ x, f x ⊓ g x) (p ⊓ q) = quantifier_fn_aux s f p ⊓ quantifier_fn_aux s g q := rfl

@[simp] lemma quantifier_fn_aux_or (s) (f g : term L → formula L) (p q : formula L) :
  quantifier_fn_aux s (λ x, f x ⊔ g x) (p ⊔ q) = quantifier_fn_aux s f p ⊔ quantifier_fn_aux s g q := rfl

@[simp] lemma quantifier_fn_aux_iff (s) (f g : term L → formula L) (p q : formula L) :
  quantifier_fn_aux s (λ x, f x ⟷ g x) (p ⟷ q) = quantifier_fn_aux s f p ⟷ quantifier_fn_aux s g q :=
by simp[lrarrow_def]

@[simp] lemma quantifier_fn_aux_ex (s) (f : term L → formula L) (p : formula L) :
  quantifier_fn_aux s (λ x, ∃.(f x)) (∃.p) = ∃.quantifier_fn_aux (s + 1) f p := rfl

@[simp] lemma quantifier_fn_aux_eq (s) (f g : term L → term L) (t u : term L) :
  quantifier_fn_aux s (λ x, f x =' g x) (t =' u) = (f #s =' g #s) := rfl

@[simp] lemma quantifier_fn_aux_le [has_le_symbol L] (s) (f g : term L → term L) (t u : term L) :
  quantifier_fn_aux s (λ x, f x ≼ g x) (t ≼ u) = (f #s ≼ g #s) := rfl

@[simp] lemma quantifier_fn_aux_mem [has_mem_symbol L] (s) (f g : term L → term L) (t u : term L) :
  quantifier_fn_aux s (λ x, f x ∊ g x) (t ∊ u) = (f #s ∊ g #s) := rfl

@[simp] lemma fal_fn_constant (s) (p : formula L) :
  quantifier_fn_aux s (λ x, p) p = p :=
by induction p generalizing s; simp*

end formula

def fal_fn (p : term L → formula L) : formula L := ∀.formula.quantifier_fn_aux 0 p (p #0)

def ex_fn (p : term L → formula L) : formula L := ∃.formula.quantifier_fn_aux 0 p (p #0)

notation `∀₁` binders `, ` r:(scoped p, fal_fn p) := r

notation `∃₁` binders `, ` r:(scoped p, ex_fn p) := r

def bfal_le_fn [has_preceq (term L) (formula L)] (t : term L) (p : term L → formula L) : formula L :=
∀₁ x, ((x ≼ t) ⟶ p x)

notation `∀₁` binders ` ≼ᵇ ` t `, ` r:(scoped p, bfal_le_fn t p) := r

def bex_le_fn [has_preceq (term L) (formula L)] (t : term L) (p : term L → formula L) : formula L :=
∃₁ x, ((x ≼ t) ⊓ p x)

notation `∃₁` binders ` ≼ᵇ ` t `, ` r:(scoped p, bex_le_fn t p) := r

def bfal_lt_fn [has_prec (term L) (formula L)] (t : term L) (p : term L → formula L) : formula L :=
∀₁ x, ((x ≺ t) ⟶ p x)

notation `∀₁` binders ` ≺ᵇ ` t `, ` r:(scoped p, bfal_le_fn t p) := r

def bex_lt_fn [has_prec (term L) (formula L)] (t : term L) (p : term L → formula L) : formula L :=
∃₁ x, ((x ≺ t) ⊓ p x)

notation `∃₁` binders ` ≺ᵇ ` t `, ` r:(scoped p, bex_le_fn t p) := r

def bfal_mem_fn [has_elem (term L) (formula L)] (t : term L) (p : term L → formula L) : formula L :=
∀₁ x, (x ∊ t) ⟶ p x

notation `∀₁` binders ` ∊ᵇ ` t `, ` r:(scoped p, bfal_mem_fn t p) := r

def bex_mem_fn [has_elem (term L) (formula L)] (t : term L) (p : term L → formula L) : formula L :=
∃₁ x, (x ∊ t) ⊓ p x

notation `∃₁` binders ` ∊ᵇ ` t `, ` r:(scoped p, bex_mem_fn t p) := r

#check ∀₁ x ≺ᵇ #4, x =' x

#check ∃.{≼ #3} #4 =' #9

variables [has_le_symbol L]

namespace formula

section
variables {L₁ L₂ : language.{u}} [L₁.language_translation_coe L₂]

end

inductive bounded : Theory L
| verum : bounded ⊤
| predicate {n} {p : L.pr n} {v} : bounded (❴p❵ v)
| equal {t u : term L} : bounded (t =' u)
| imply {p q} : bounded p → bounded q → bounded (p ⟶ q)
| neg {p} : bounded p → bounded (∼p)
| bfal {t} {p} : bounded p → bounded ∀.{≼ t} p

attribute [simp] bounded.verum bounded.predicate bounded.equal bounded.neg

@[simp] lemma bounded_preceq (t u : term L) : bounded (t ≼ u : formula L) :=
bounded.predicate

@[simp] lemma bounded_imply_iff (p q : formula L) : bounded (p ⟶ q) ↔ bounded p ∧ bounded q :=
⟨λ h, by { cases h, simp* }, λ ⟨hp, hq⟩, bounded.imply hp hq⟩

@[simp] lemma bounded_neg_iff (p : formula L) : bounded (∼p) ↔ bounded p :=
⟨λ h, by { cases h, simp* }, bounded.neg⟩

lemma bounded_of_open {p : formula L} (h : p.is_open) : bounded p :=
begin
  induction p,
  case verum { simp },
  case app { simp },
  case equal { simp },
  case imply : p q IHp IHq { simp at h, exact bounded.imply (IHp h.1) (IHq h.2) },
  case neg : p IH { simp at h, exact bounded.neg (IH h) },
  case fal { simp at h, contradiction }
end

lemma bounded_of_lt {p q : formula L} (hq : bounded q) (h : p < q) : bounded p :=
begin
  induction q; try { simp at h, contradiction },
  case imply : q₁ q₂ IH₁ IH₂
  { simp at h hq, rcases h with (le | le),
    { rcases lt_or_eq_of_le le with (lt | rfl), { exact IH₁ hq.1 lt }, { exact hq.1 } },
    { rcases lt_or_eq_of_le le with (lt | rfl), { exact IH₂ hq.2 lt }, { exact hq.2 } } },
  case neg : q IH
  { simp at h hq, rcases lt_or_eq_of_le h with (lt | rfl), { refine IH hq lt }, { exact hq } },
  case fal : p IH
  { rcases hq with ⟨lo, loo⟩, simp at h,
    rcases lt_or_eq_of_le h with (lt | rfl), { refine IH (by { simp, exact hq_ᾰ}) lt }, { simp, exact hq_ᾰ } }
end

@[simp] lemma bounded_bfal_iff {p : formula L} {t : term L} : bounded (∀.{≼ t} p) ↔ bounded p :=
⟨λ h, bounded_of_lt h (by { simp[bfal_le], refine le_of_lt (by simp) }), bounded.bfal⟩

lemma bounded_bex_iff {p : formula L} {t : term L} : bounded (∃.{≼ t} p) ↔ bounded p :=
by simp[bex_le]

lemma bounded_rew (s : ℕ → term L) (p : formula L) (h : bounded p) : bounded (rew s p) :=
by induction h generalizing s; simp*

instance bounded_proper : proper_Theory (bounded : Theory L) :=
⟨λ p s mem, by { simp[set.mem_def] at mem ⊢, induction mem generalizing s; try {simp*} }⟩

mutual inductive is_sigma, is_pi
with is_sigma : ℕ → formula L → Prop
| zero : ∀ {p : formula L}, p.bounded → is_sigma 0 p
| succ : ∀ {p} {n}, is_pi n p → is_sigma (n+1) ∃.p
with is_pi : ℕ → formula L → Prop
| zero : ∀ {p : formula L}, p.bounded → is_pi 0 p
| succ : ∀ {p} {n}, is_sigma n p → is_pi (n+1) ∀.p

notation `𝛴₀` := is_sigma 0

notation `𝛴₁` := is_sigma 1

notation `𝛱₀` := is_pi 0

notation `𝛱₁` := is_pi 1

@[simp] lemma sigma_0_iff_bounded : (𝛴₀ : Theory L) = bounded :=
by funext p; simp; exact ⟨λ h, by { cases h, simp* }, is_sigma.zero⟩

@[simp] lemma pi_0_iff_bounded : (𝛱₀ : Theory L) = bounded :=
by funext p; simp; exact ⟨λ h, by { cases h, simp* }, is_pi.zero⟩

private lemma sigma_pi_proper (n : ℕ) : proper_at 0 (is_sigma n : Theory L) ∧ proper_at 0 (is_pi n : Theory L) :=
begin
  induction n with n IH,
  { simp, refine proper_Theory.proper },
  { refine ⟨λ p s hs, _, λ p s hp, _⟩,
    { cases hs with _ _ p _ hp, simp,
      have : is_pi n (p.rew (s^1)),
      { have := IH.2 p (s^1) hp, simp at this, exact this },
      refine is_sigma.succ this },
    { cases hp with _ _ p _ hs, simp,
      have : is_sigma n (p.rew (s^1)),
      { have := IH.1 p (s^1) hs, simp at this, exact this },
      refine is_pi.succ this } }
end

instance is_sigma_proper (n) : proper_Theory (is_sigma n : Theory L) := ⟨(sigma_pi_proper n).1⟩

instance is_pi_proper (n) : proper_Theory (is_pi n : Theory L) := ⟨(sigma_pi_proper n).2⟩

end formula

def arithmetical_sigma (T : Theory L) (n : ℕ) : Theory L :=
λ p, ∃ q, T ⊢ p ⟷ q ∧ q.is_sigma n

notation `𝜮`:60 n ` in ` T :60 := arithmetical_sigma T n

def arithmetical_pi (T : Theory L) (n : ℕ) : Theory L :=
λ p, ∃ q, T ⊢ p ⟷ q ∧ q.is_pi n

notation `𝜫`:60 n ` in ` T :60 := arithmetical_pi T n

def arithmetical_delta (T : Theory L) (n : ℕ) : Theory L :=
λ p, p ∈ 𝜮n in T ∧ p ∈ 𝜫n in T

notation `𝜟`:60 n ` in ` T :60 := arithmetical_delta T n

end fol


/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Meta.WHNF
import Lean.Meta.Transform
import Lean.Meta.SynthInstance
import Lean.Meta.AppBuilder

namespace Lean.Meta

builtin_initialize coeDeclAttr : TagAttribute ←
  registerTagAttribute `coe_decl "auxiliary definition used to implement coercion (unfolded during elaboration)"

/--
  Return true iff `declName` is one of the auxiliary definitions/projections
  used to implement coercions.
-/
def isCoeDecl (env : Environment) (declName : Name) : Bool :=
  coeDeclAttr.hasTag env declName

/-- Expand coercions occurring in `e` -/
partial def expandCoe (e : Expr) : MetaM Expr :=
  withReducibleAndInstances do
    transform e fun e => do
      let f := e.getAppFn
      if f.isConst then
        let declName := f.constName!
        if isCoeDecl (← getEnv) declName then
          if let some e ← unfoldDefinition? e then
            return .visit e.headBeta
      return .continue

register_builtin_option autoLift : Bool := {
  defValue := true
  descr    := "insert monadic lifts (i.e., `liftM` and coercions) when needed"
}

/-- Coerces `expr` to `expectedType` using `CoeT`. -/
def coerceSimple? (expr expectedType : Expr) : MetaM (LOption Expr) := do
  let eType ← inferType expr
  let u ← getLevel eType
  let v ← getLevel expectedType
  let coeTInstType := mkAppN (mkConst ``CoeT [u, v]) #[eType, expr, expectedType]
  match ← trySynthInstance coeTInstType with
  | .some inst =>
    let result ← expandCoe (mkAppN (mkConst ``CoeT.coe [u, v]) #[eType, expr, expectedType, inst])
    unless ← isDefEq (← inferType result) expectedType do
      throwError "could not coerce{indentExpr expr}\nto{indentExpr expectedType}\ncoerced expression has wrong type:{indentExpr result}"
    return .some result
  | .undef => return .undef
  | .none => return .none

/-- Coerces `expr` to a function type. -/
def coerceToFunction? (expr : Expr) : MetaM (Option Expr) := do
  -- constructing expression manually because mkAppM wouldn't assign universe mvars
  let α ← inferType expr
  let u ← getLevel α
  let v ← mkFreshLevelMVar
  let γ ← mkFreshExprMVar (← mkArrow α (mkSort v))
  let .some inst ← trySynthInstance (mkApp2 (.const ``CoeFun [u,v]) α γ) | return none
  let expanded ← expandCoe (mkApp4 (.const ``CoeFun.coe [u,v]) α γ inst expr)
  unless (← whnf (← inferType expanded)).isForall do
    throwError "failed to coerce{indentExpr expr}\nto a function, after applying `CoeFun.coe`, result is still not a function{indentExpr expanded}\nthis is often due to incorrect `CoeFun` instances, the synthesized instance was{indentExpr inst}"
  return expanded

/-- Coerces `expr` to a type. -/
def coerceToSort? (expr : Expr) : MetaM (Option Expr) := do
  -- constructing expression manually because mkAppM wouldn't assign universe mvars
  let α ← inferType expr
  let u ← getLevel α
  let v ← mkFreshLevelMVar
  let β ← mkFreshExprMVar (mkSort v)
  let .some inst ← trySynthInstance (mkApp2 (.const ``CoeSort [u,v]) α β) | return none
  let expanded ← expandCoe (mkApp4 (.const ``CoeSort.coe [u,v]) α β inst expr)
  unless (← whnf (← inferType expanded)).isSort do
    throwError "failed to coerce{indentExpr expr}\nto a type, after applying `CoeSort.coe`, result is still not a type{indentExpr expanded}\nthis is often due to incorrect `CoeSort` instances, the synthesized instance was{indentExpr inst}"
  return expanded

/-- Return `some (m, α)` if `type` can be reduced to an application of the form `m α` using `[reducible]` transparency. -/
def isTypeApp? (type : Expr) : MetaM (Option (Expr × Expr)) := do
  let type ← withReducible <| whnf type
  match type with
  | .app m α => return some ((← instantiateMVars m), (← instantiateMVars α))
  | _        => return none

/--
Return `true` if `type` is of the form `m α` where `m` is a `Monad`.
Note that we reduce `type` using transparency `[reducible]`.
-/
def isMonadApp (type : Expr) : MetaM Bool := do
  let some (m, _) ← isTypeApp? type | return false
  return (← isMonad? m).isSome

/--
Try coercions and monad lifts to make sure `e` has type `expectedType`.

If `expectedType` is of the form `n β`, we try monad lifts and other extensions.

Extensions for monads.

1. Try to unify `n` and `m`. If it succeeds, then we use
  ```
  coeM {m : Type u → Type v} {α β : Type u} [∀ a, CoeT α a β] [Monad m] (x : m α) : m β
  ```
  `n` must be a `Monad` to use this one.

2. If there is monad lift from `m` to `n` and we can unify `α` and `β`, we use
  ```
  liftM : ∀ {m : Type u_1 → Type u_2} {n : Type u_1 → Type u_3} [self : MonadLiftT m n] {α : Type u_1}, m α → n α
  ```
  Note that `n` may not be a `Monad` in this case. This happens quite a bit in code such as
  ```
  def g (x : Nat) : IO Nat := do
    IO.println x
    pure x

  def f {m} [MonadLiftT IO m] : m Nat :=
    g 10

  ```

3. If there is a monad lift from `m` to `n` and a coercion from `α` to `β`, we use
  ```
  liftCoeM {m : Type u → Type v} {n : Type u → Type w} {α β : Type u} [MonadLiftT m n] [∀ a, CoeT α a β] [Monad n] (x : m α) : n β
  ```

Note that approach 3 does not subsume 1 because it is only applicable if there is a coercion from `α` to `β` for all values in `α`.
This is not the case for example for `pure $ x > 0` when the expected type is `IO Bool`. The given type is `IO Prop`, and
we only have a coercion from decidable propositions.  Approach 1 works because it constructs the coercion `CoeT (m Prop) (pure $ x > 0) (m Bool)`
using the instance `pureCoeDepProp`.

Note that, approach 2 is more powerful than `tryCoe`.
Recall that type class resolution never assigns metavariables created by other modules.
Now, consider the following scenario
```lean
def g (x : Nat) : IO Nat := ...
deg h (x : Nat) : StateT Nat IO Nat := do
v ← g x;
IO.Println v;
...
```
Let's assume there is no other occurrence of `v` in `h`.
Thus, we have that the expected of `g x` is `StateT Nat IO ?α`,
and the given type is `IO Nat`. So, even if we add a coercion.
```
instance {α m n} [MonadLiftT m n] {α} : Coe (m α) (n α) := ...
```
It is not applicable because TC would have to assign `?α := Nat`.
On the other hand, TC can easily solve `[MonadLiftT IO (StateT Nat IO)]`
since this goal does not contain any metavariables. And then, we
convert `g x` into `liftM $ g x`.
-/
def coerceMonadLift? (e expectedType : Expr) : MetaM (Option Expr) := do
  let expectedType ← instantiateMVars expectedType
  let eType ← instantiateMVars (← inferType e)
  let some (n, β) ← isTypeApp? expectedType | return none
  let some (m, α) ← isTypeApp? eType | return none
  if (← isDefEq m n) then
    let some monadInst ← isMonad? n | return none
    try expandCoe (← mkAppOptM ``Lean.Internal.coeM #[m, α, β, none, monadInst, e]) catch _ => return none
  else if autoLift.get (← getOptions) then
    try
      -- Construct lift from `m` to `n`
      -- Note: we cannot use mkAppM here because mkAppM does not assign universe metavariables,
      -- but we need to make sure that the domains of `m` and `n` have the same level.
      let .forallE _ (.sort um₁) (.sort um₂) _ ← whnf (← inferType m) | return none
      let .forallE _ (.sort un₁) (.sort un₂) _ ← whnf (← inferType n) | return none
      let u ← decLevel um₁
      let .true ← isLevelDefEq u (← decLevel un₁) | return none
      let v ← decLevel um₂
      let w ← decLevel un₂
      let monadLiftType := mkAppN (.const ``MonadLiftT [u, v, w]) #[m, n]
      let .some monadLiftVal ← trySynthInstance monadLiftType | return none
      let u_1 ← getDecLevel α
      let u_2 ← getDecLevel eType
      let u_3 ← getDecLevel expectedType
      let eNew := mkAppN (Lean.mkConst ``liftM [u_1, u_2, u_3]) #[m, n, monadLiftVal, α, e]
      let eNewType ← inferType eNew
      if (← isDefEq expectedType eNewType) then
        return some eNew -- approach 2 worked
      else
        let some monadInst ← isMonad? n | return none
        let u ← getLevel α
        let v ← getLevel β
        let coeTInstType := Lean.mkForall `a BinderInfo.default α <| mkAppN (mkConst ``CoeT [u, v]) #[α, mkBVar 0, β]
        let .some coeTInstVal ← trySynthInstance coeTInstType | return none
        let eNew ← expandCoe (mkAppN (Lean.mkConst ``Lean.Internal.liftCoeM [u_1, u_2, u_3]) #[m, n, α, β, monadLiftVal, coeTInstVal, monadInst, e])
        let eNewType ← inferType eNew
        unless (← isDefEq expectedType eNewType) do return none
        return some eNew -- approach 3 worked
    catch _ =>
      /- If `m` is not a monad, then we try to use `tryCoe?`. -/
      return none
  else
    return none

/-- Coerces `expr` to the type `expectedType`.
Returns `.some coerced` on successful coercion,
`.none` if the expression cannot by coerced to that type,
or `.undef` if we need more metavariable assignments. -/
def coerce? (expr expectedType : Expr) : MetaM (LOption Expr) := do
  if let some lifted ← coerceMonadLift? expr expectedType then
    return .some lifted
  if (← whnfR expectedType).isForall then
    if let some fn ← coerceToFunction? expr then
      if ← isDefEq (← inferType fn) expectedType then
        return .some fn
  coerceSimple? expr expectedType

end Lean.Meta


import Lib.Data.MonoFunctor
import Lib.Data.MonoFoldable
import Lib.Data.MonoTraversable
import Lib.Data.List.Instances

namespace String

@[implementedBy map]
def map' (f : Char → Char) : String → String
| ⟨ data ⟩ => ⟨ Functor.map f data ⟩

instance : MonoFunctor String Char where
  map := map'

instance : LawfulMonoFunctor String Char where
  id_map := by
    intros; next x =>
    cases x; simp [MonoFunctor.map, map']
  comp_map := by
    intros; next x =>
    cases x; simp [MonoFunctor.map, map']

variable {β : Type}

@[inline]
def foldlImpl (f : β → Char → β) (x₀ : β) : String → β :=
String.foldl f x₀

@[implementedBy foldlImpl]
def foldl' (f : β → Char → β) (x₀ : β) : String → β
| ⟨ data ⟩ => Foldable.foldl f x₀ data

@[inline]
def foldrImpl (f : Char → β → β) (x₀ : β) : String → β :=
String.foldr f x₀

@[implementedBy foldrImpl]
def foldr' (f : Char → β → β) (x₀ : β) : String → β
| ⟨ data ⟩ => Foldable.foldr f x₀ data

def toArray (s : String) : Array Char :=
String.foldl' Array.push (Array.mkEmpty s.length) s

instance : MonoFoldable String Char where
  foldl := String.foldl'
  foldr := String.foldr'
  toList := String.toList
  toArray := String.toArray
  length := String.length

open MonoFoldable

instance : LawfulMonoFoldable String Char where
  foldl_sim :=
    by intros _ _ _ f SIM x₀ y₀ t a ih
       cases t; simp [foldl, MonoFoldable.foldl]
       auto [LawfulFoldable.foldl_sim]
  foldr_eq_foldMap :=
    by intros _ _ t _
       cases t; simp [foldr, MonoFoldable.foldr]
       auto [LawfulFoldable.foldr_eq_foldMap]
  toArray_toList :=
    by intros t
       cases t with | mk data =>
       simp [toList, toArray, String.foldl']
       simp [String.toList,
             String.toArray, String.foldl']
       show Foldable.toArray data = _
       rw [LawfulFoldable.toArray_eq]
       refl
  length_toList := by
    intros t; cases t; refl
  foldl_toList := by
    intros β f x₀ t; cases t; refl

section mapM

variable {M} [Monad M]
@[specialize] partial def mapMAux (f : Char → M Char) (i : Pos) (s : String) : M String :=
  if s.atEnd i then return s
  else do
    let c ← f (s.get i)
    let s := s.set i c
    mapMAux f (s.next i) s

@[inline] def mapMImpl (f : Char → M Char) (s : String) : M String :=
  mapMAux f 0 s

@[implementedBy mapMImpl] def mapM (f : Char → M Char) (s : String) : M String :=
  String.mk <$> s.data.mapM f

end mapM

section mapM

variable {F} [Applicative F]

def mapA (f : Char → F Char) (s : String) : F String :=
String.mk <$> Traversable.traverse f s.data

end mapM

instance : MonoTraversable String Char where
  traverse := mapA
  mapM := mapM

open LawfulMonoTraversable
instance : LawfulMonoTraversable String Char where
  map_eq_traverse := by
    intros; next x => cases x with | mk data =>
    simp [MonoTraversable.traverse, MonoFunctor.map, map']
    simp [LawfulTraversable.map_eq_traverse]; refl
  foldl_eq_traverse := by
    intros; next x y => cases x with | mk data =>
    simp [MonoFoldable.foldl, foldl']
    simp [LawfulTraversable.foldl_eq_traverse]; refl
  traverse_eq_mapM := by
    intros; next x => cases x with | mk data =>
    simp [MonoTraversable.traverse, mapA]
    simp [LawfulTraversable.traverse_eq_mapM]; refl
  comp_traverse := by
    intros; next x y z => cases x with | mk data =>
    simp [MonoTraversable.traverse, mapA]
    simp [LawfulTraversable.comp_traverse]; refl
  traverse_sim := by
    intros; next x SIM _ _ _ => cases x with | mk data =>
    simp [MonoTraversable.traverse, mapA]
    auto [LawfulTraversable.traverse_sim, ApplicativeRel.naturality]

theorem toList_inj {s₀ s₁ : String} :
  s₀.toList = s₁.toList → s₀ = s₁ :=
by intro h; cases s₀; cases s₁; cases h; rfl

attribute [auto] one_le_csize

@[simp]
theorem toList_mk {x : List Char} :
  String.toList ⟨x⟩ = x := rfl

@[simp]
theorem toList_append {s₀ s₁ : String} :
  (s₀ ++ s₁).toList = s₀.toList ++ s₁.toList :=
rfl

@[simp]
theorem length_toList {s : String} :
  s.toList.length = s.length := rfl