import category_theory.adjunction.limits
import category_theory.closed.cartesian
import category_theory.conj

universes v u₁ u₂

namespace category_theory
open category limits
variables {C : Type u₁} [category.{v} C] [has_finite_products C] [cartesian_closed C]
variables {D : Type u₂} [category.{v} D] [has_finite_products D] [cartesian_closed D]
variables (F : C ⥤ D)

class cartesian_closed_functor :=
[preserves_bin_prods : preserves_limits_of_shape (discrete walking_pair) F]
(comparison_iso : ∀ A B, is_iso (exp_comparison F A B))

attribute [instance] cartesian_closed_functor.comparison_iso

variables {F} {L : D ⥤ C}

noncomputable def frobenius_map (A : C) (B : D) (adj : L ⊣ F) : L.obj (F.obj A ⨯ B) ⟶ A ⨯ L.obj B :=
prod_comparison _ _ _ ≫ limits.prod.map (adj.counit.app A) (𝟙 _)

@[reassoc]
lemma frob_naturality {A A' : C} {B B' : D} (adj : L ⊣ F) (f : A ⟶ A') (g : B ⟶ B') :
  frobenius_map A B adj ≫ limits.prod.map f (L.map g) = L.map (limits.prod.map (F.map f) g) ≫ frobenius_map A' B' adj :=
begin
  rw [frobenius_map, frobenius_map, assoc, prod_comparison_natural_assoc],
  apply prod.hom_ext,
  simp only [assoc, limits.prod.map_fst, limits.prod.map_fst_assoc, adjunction.counit_naturality],
  simp only [assoc, limits.prod.map_snd, limits.prod.map_snd_assoc, comp_id, id_comp],
end

-- @[reassoc]
-- lemma frob_inv_naturality_left {A A' : C} {B : D} (adj : L ⊣ F) (f : A ⟶ A')
--   [is_iso (frobenius_map A B adj)]
--   [is_iso (frobenius_map A' B adj)] :
-- inv (frobenius_map A B adj) ≫ L.map (limits.prod.map (F.map f) (𝟙 _)) =
-- limits.prod.map f (𝟙 _) ≫ inv (frobenius_map A' B adj) :=
-- begin
--   erw [(as_iso]
--   -- erw [(as_iso (frobenius_map A B adj)).inv_comp_eq, assoc, ← prod_map_comp_id_assoc, ← assoc,
--   --      (as_iso (frobenius_map A' B adj)).eq_comp_inv, prod_comparison_natural_assoc, L.map_id,
--   --      ← prod_map_comp_id, adj.counit.naturality], refl,
-- end

@[reassoc]
lemma frob_inv_naturality {A A' : C} {B B' : D} (adj : L ⊣ F) (f : A ⟶ A') (g : B ⟶ B')
  [is_iso (frobenius_map A B adj)]
  [is_iso (frobenius_map A' B' adj)] :
inv (frobenius_map A B adj) ≫ L.map (limits.prod.map (F.map f) g) =
limits.prod.map f (L.map g) ≫ inv (frobenius_map A' B' adj) :=
begin
  erw [(as_iso (frobenius_map A B adj)).inv_comp_eq],
  rw ← assoc,
  erw [(as_iso (frobenius_map A' B' adj)).eq_comp_inv],
  symmetry,
  apply frob_naturality,
end

variables (adj : L ⊣ F) [∀ A B, is_iso (frobenius_map A B adj)]

noncomputable def biject (adj : L ⊣ F) [∀ A B, is_iso (frobenius_map A B adj)]
  {A B : C} {c : D} : (c ⟶ F.obj (B ^^ A)) ≃ (c ⟶ F.obj B ^^ F.obj A) :=
calc (c ⟶ F.obj (B ^^ A)) ≃ (L.obj c ⟶ B ^^ A) : (adj.hom_equiv _ _).symm
     ... ≃ (A ⨯ L.obj c ⟶ B) : ((exp.adjunction A).hom_equiv _ _).symm
     ... ≃ (L.obj (F.obj A ⨯ c) ⟶ B) : iso.hom_congr (as_iso (frobenius_map _ _ adj)).symm (iso.refl _)
     ... ≃ (F.obj A ⨯ c ⟶ F.obj B) : adj.hom_equiv _ _
     ... ≃ (c ⟶ F.obj B ^^ F.obj A) : ((exp.adjunction _).hom_equiv _ _)

-- tidy this up
lemma biject_id {A B : C} [preserves_limits_of_shape (discrete walking_pair) F] :
  biject adj (𝟙 _) = exp_comparison F A B :=
begin
  dsimp [biject, iso.hom_congr, frobenius_map, exp_comparison],
  rw [comp_id],
  change cartesian_closed.curry _ = _,
  congr' 1,
  erw (as_iso (prod_comparison F A ((exp A).obj B))).eq_inv_comp,
  rw adj.hom_equiv_unit,
  erw adj.unit.naturality_assoc,
  dsimp [prod_comparison],
  erw ← F.map_comp,
  rw adjunction.hom_equiv_counit,
  rw adjunction.hom_equiv_counit,
  rw L.map_id,
  rw id_comp,
  dsimp,
  rw prod.lift_map,
  rw prod.lift_map_assoc,
  rw comp_id,
  rw comp_id,
  simp only [functor.map_comp],
  rw ← F.map_comp,
  rw ← F.map_comp,
  rw ← adj.hom_equiv_unit,
  rw adj.hom_equiv_apply_eq,
  rw adj.hom_equiv_counit,
  erw adj.counit.naturality,
  rw ← assoc,
  congr' 1,
  apply prod.hom_ext; simp [← functor.map_comp_assoc],
end

lemma biject_id' {A B : C} [preserves_limits_of_shape (discrete walking_pair) F] :
  𝟙 (F.obj ((exp A).obj B)) = (biject adj).symm (exp_comparison F A B) :=
begin
  rw equiv.eq_symm_apply,
  rw biject_id,
end

lemma biject_natural
  {A B B' : C} {c c' : D} (f : c' ⟶ c) (g : B ⟶ B') (q : c ⟶ F.obj (B ^^ A)) :
f ≫ biject adj q ≫ (exp _).map (F.map g) = biject adj (f ≫ q ≫ F.map ((exp _).map g)) :=
begin
  dsimp [biject, iso.hom_congr],
  rw [comp_id, comp_id],
  rw ← adjunction.hom_equiv_naturality_right,
  rw ← adjunction.hom_equiv_naturality_left,
  erw ← adjunction.hom_equiv_naturality_right,
  rw ← adjunction.hom_equiv_naturality_left,
  rw assoc,
  rw adj.hom_equiv_naturality_left_symm,
  rw adjunction.hom_equiv_naturality_left_symm,
  erw frob_naturality_assoc,
  rw adj.hom_equiv_naturality_right_symm,
  rw adjunction.hom_equiv_naturality_right_symm,
  congr' 4,
  dsimp, simp,
end

lemma biject_natural_left
  {A B : C} {c c' : D} (f : c' ⟶ c) (q : c ⟶ F.obj (B ^^ A)) :
f ≫ biject adj q = biject adj (f ≫ q) :=
by simpa using biject_natural adj f (𝟙 _) q

lemma biject_natural_right
  {A B B' : C} {c : D} (g : B ⟶ B') (q : c ⟶ F.obj (B ^^ A)) :
biject adj q ≫ (exp _).map (F.map g) = biject adj (q ≫ F.map ((exp _).map g)) :=
by simpa using biject_natural adj (𝟙 _) g q

noncomputable def cartesian_closed_of_frobenius_iso : cartesian_closed_functor F :=
{ preserves_bin_prods :=
  begin
    letI := adj.right_adjoint_preserves_limits,
    apply_instance,
  end,
  comparison_iso := λ A B,
  { inv := (biject adj).symm (𝟙 _),
    hom_inv_id' :=
    begin
      rw ← (biject adj).apply_eq_iff_eq,
      rw biject_id _,
      rw ← biject_natural_left,
      rw equiv.apply_symm_apply,
      rw comp_id,
    end,
    inv_hom_id' :=
    begin
      letI := adj.right_adjoint_preserves_limits,
      rw ← biject_id adj,
      rw biject_natural_left,
      rw comp_id,
      rw equiv.apply_symm_apply,
    end } }

-- /--
-- The exponential comparison map.
-- `F` is a cartesian closed functor if this is an iso for all `A,B`.
-- -/
-- def exp_comparison (A B : C) :
--   F.obj (A ⟹ B) ⟶ F.obj A ⟹ F.obj B :=
-- curry (inv (prod_comparison F A _) ≫ F.map ((ev _).app _))

-- /-- The exponential comparison map is natural in its left argument. -/
-- lemma exp_comparison_natural_left (A A' B : C) (f : A' ⟶ A) :
--   exp_comparison F A B ≫ pre (F.obj B) (F.map f) = F.map (pre B f) ≫ exp_comparison F A' B :=
-- begin
--   rw [exp_comparison, exp_comparison, ← curry_natural_left, eq_curry_iff, uncurry_natural_left,
--        pre, uncurry_curry, prod_map_map_assoc, curry_eq, prod_map_id_comp, assoc],
--   erw [(ev _).naturality, ev_coev_assoc, ← F.map_id, ← prod_comparison_inv_natural_assoc,
--        ← F.map_id, ← prod_comparison_inv_natural_assoc, ← F.map_comp, ← F.map_comp, pre, curry_eq,
--        prod_map_id_comp, assoc, (ev _).naturality, ev_coev_assoc], refl,
-- end

-- /-- The exponential comparison map is natural in its right argument. -/
-- lemma exp_comparison_natural_right (A B B' : C) (f : B ⟶ B') :
--   exp_comparison F A B ≫ (exp (F.obj A)).map (F.map f) =
--     F.map ((exp A).map f) ≫ exp_comparison F A B' :=
-- by
--   erw [exp_comparison, ← curry_natural_right, curry_eq_iff, exp_comparison, uncurry_natural_left,
--        uncurry_curry, assoc, ← F.map_comp, ← (ev _).naturality, F.map_comp,
--        prod_comparison_inv_natural_assoc, F.map_id]

-- -- TODO: If F has a left adjoint L, then F is cartesian closed if and only if
-- -- L (B ⨯ F A) ⟶ L B ⨯ L F A ⟶ L B ⨯ A
-- -- is an iso for all A ∈ D, B ∈ C.
-- -- Corollary: If F has a left adjoint L which preserves finite products, F is cartesian closed iff
-- -- F is full and faithful.
end category_theory


import Lean.Meta
import Lean.Elab
import Scratch.TermSeq
open Lean.Core
open Lean.Meta
open Lean.Elab.Term
open Lean
open Nat

-- testing decomposition, used later (and better forms)
-- also logging info in tactics and elaborators
-- also adding a definition in a tactic context
-- with a tactic doing this
-- experiments with thunks, lazylists

partial def decomposeSeq : Expr → MetaM (List Expr) :=
  fun expr =>
  do
    let mvar ←  mkFreshExprMVar none
    let tmvar ← mkFreshExprMVar (mkConst `TermSeq)
    let sExp ←  mkAppM ``TermSeq.cons #[mvar, tmvar]
    if ← isDefEq sExp expr then
      let prev ← decomposeSeq tmvar
      return mvar :: prev
    else
      return []

partial def seqLength : Expr → MetaM Nat := fun expr =>
  do
    let l ← decomposeSeq expr
    return l.length

syntax (name:= tsl) "tsl!" term : term

@[termElab tsl] def tslImpl : TermElab :=
  fun stx expectedType? =>
    match stx with
    | `(tsl!%$tk $s) =>
      do
        let e ← elabTerm s none
        let l ← seqLength e
        Lean.Elab.logInfoAt tk m!"expr: {(← e)}"
        Lean.Elab.logInfoAt tk m!"whnf: {(← whnf e)}"
        Lean.Elab.logInfoAt tk m!"length: {(← l)}"
        return ← ToExpr.toExpr l
    | _ => Elab.throwIllFormedSyntax

def ts := TermSeq.cons 3 TermSeq.empty

#eval tsl! ts
#eval tsl! #⟨3, 4, "this"⟩

#check  #⟨3, 4, "this"⟩
#check fun x: Nat => #⟨x, x, "this", 4⟩

example : (fun x: Nat => #⟨x, x, "this", 4⟩) 3 = #⟨3, 3, "this", 4⟩ := rfl

open Lean.Elab.Tactic

syntax (name:= blah) "blah" : tactic
@[tactic blah] def blahImpl : Tactic :=
  fun stx =>
    do
      Lean.Elab.logInfo "blah say I"
      return ()

syntax (name := tsltac) "tslength" term : tactic
@[tactic tsltac] def tslTacImpl : Tactic :=
  fun stx =>
    match stx with
    | `(tactic|tslength $s) =>
      do
        let mvar ← getMainGoal
        let e ← liftM (Elab.Term.elabTerm s none true true)
        let l ← seqLength e
        let n := ToExpr.toExpr l
        replaceMainGoal []
        assignExprMVar mvar n
    | _ => Elab.throwIllFormedSyntax

def tstacEg : Nat := by
        blah
        tslength #⟨3, 4, "this"⟩

#eval tstacEg

universe u v

def factorThroughEg(α : Sort u) (β : Sort v)(b : β ) : (β  → α ) → α   :=
    fun g => g b

def addToContextM3(name: Name) (type : Expr)(value: Expr) :
     MVarId → MetaM (List MVarId) :=
  fun m =>
      do
        let target ← getMVarType m
        let exp ← mkAppM `factorThroughEg #[target, type, value]
        let appGoalList ←  apply m exp
        let appGoal := appGoalList.head!
        let ⟨_, introGoal⟩ ←  intro appGoal name
        return [introGoal]

syntax (name:= useterm) "use" term ("with" term)? "as" ident : tactic
@[tactic useterm] def usetermImpl : Tactic :=
   fun stx =>
    match stx with
    | `(tactic|use $s with $t as $n) =>
    withMainContext $
    do
      let name ← n.getId
      let typ ← elabType t
      let value ← Elab.Tactic.elabTerm s (some typ)
      liftMetaTactic $ addToContextM3 name typ value
    | `(tactic|use $s as $n) =>
    withMainContext $
    do
      let name ← n.getId
      let value ← Elab.Tactic.elabTerm s none
      let typ ← inferType value
      liftMetaTactic $ addToContextM3 name typ value
    | _ => Elab.throwIllFormedSyntax

example : Nat := by
        use 3 with Nat as n
        use "this" as s
        let x := 3
        use (succ x) as y
        have b := "d"
        exact y

syntax (name:= dupllet) "assign" ident "::" term "as" term : tactic
@[tactic dupllet] def assignImpl : Tactic :=
  fun stx =>
    match stx with
    | `(tactic|assign $n:ident :: $t as  $i) =>
      do
        let name ← n.getId
        let typ ← liftM (Elab.Term.elabTerm t none true true)
        let value ← liftM (Elab.Term.elabTerm i (some typ) true true)
        let mvar ← getMainGoal
        withMVarContext mvar do
          let mvar2 ← mkFreshExprMVar (← getMVarType mvar)
          let mvarId2 := mvar2.mvarId!
          replaceMainGoal [mvarId2]
          withLetDecl  name  typ value $ fun x =>
            do
            assignExprMVar mvar (← mkLetFVars #[x] mvar2)
            return ()
    | _ => Elab.throwIllFormedSyntax

def fl := 4.5
def three := 3

def letTac : Nat   := by
        let p := 3
        assign y :: Nat as 3
        skip
        exact p
        done

#eval tsl! #⟨three, fl, "this"⟩

def getFloat (s: String) : Option Float :=
  (Syntax.decodeScientificLitVal? s).map (fun ⟨m, s, e⟩ => OfScientific.ofScientific m s e)

#eval getFloat "3.1415"

#eval Syntax.decodeScientificLitVal? (Float.toString (3.145))

-- copied from @Mac's code below and modified
def parseOpt (s : String) : Option Float :=
  match Json.Parser.num s.mkIterator with
  | Parsec.ParseResult.success _ res =>
    some <| Float.ofInt res.mantissa * 10 ^ - Float.ofNat res.exponent
  | Parsec.ParseResult.error it err  => none

def parseGet (s : String) : Nat ×  Nat :=
  match Json.Parser.num s.mkIterator with
  | Parsec.ParseResult.success _ res =>
      (res.mantissa.toNat, res.exponent)
  | Parsec.ParseResult.error it err  => (0, 0)

syntax (name:= floatlit) "float!" term : term

@[termElab floatlit] def floatlitImpl : TermElab :=
    fun stx expectedType? =>
      match stx with
      | `(float! $s) =>
        do
           let fl ← Elab.Term.elabTerm s (some (Lean.mkConst `Float))
           let strRaw ← mkAppM ``Float.toString #[fl]
           let str ← whnf strRaw
          --  let ⟨n, _, _⟩ := (Syntax.isScientificLit? s).get!
           return ← mkAppM ``parseGet #[str]
          --  return  ToExpr.toExpr n
      | _ => Elab.throwIllFormedSyntax

def pi := 3.1415

#eval float! 3.14
#eval float! pi
#eval float! (3.1 * 6)

syntax ident ("⌣" ident)? "↦" term : term

macro_rules
  | `( $x:ident ↦ $y:term ) => `(fun $x => $y)
  | `( $x:ident ⌣ $t: ident ↦ $y:term ) => `(fun ($x : $t) => $y)

def mapfn : Nat → Nat := n ↦ n + 1

def mapFn2 := k ⌣ Nat ↦ (k + 1)

-- copied from Zulip chat due to @mac

-- import Lean.Data.Json.Parser

def parse (s : String) : Except String Float :=
  match (Json.Parser.num <* Parsec.eof) s.mkIterator with
  | Parsec.ParseResult.success _ res =>
    Except.ok <| Float.ofInt res.mantissa * 10 ^ - Float.ofNat res.exponent
  | Parsec.ParseResult.error it err  => Except.error s!"offset {it.i.repr}: {err}"

#eval parse "1.3" -- Except.ok 1.300000

def th3:= Thunk.mk (fun () => 3)

def addN (n : Nat) : Thunk Nat →  Thunk Nat :=
      let succT : Nat → Thunk Nat := fun n => Thunk.mk (fun () => n + 1)
      match n with
      | 0 => id
      | k + 1 => fun tn => (addN k tn).bind succT

#reduce (addN 20) th3

-- credit:
-- https://www.classes.cs.uchicago.edu/archive/2019/spring/22300-1/lectures/LazyLists/index.html

mutual
  inductive LazyList (α : Type) where
  | mk : Thunk (LazyListCell α) → LazyList α

  inductive LazyListCell (α : Type) where
  | nil : LazyListCell α
  | cons : α → LazyList α → LazyListCell α
end


import data.int.parity
import tactic.apply_fun
import data.real.basic
import algebra.group.pi

import lib.tactiques

namespace tactic.interactive
setup_tactic_parser

open tactic

meta def verifie : tactic unit :=
`[ { repeat { unfold limite_suite},
   repeat { unfold continue_en },
   push_neg,
   try { simp only [exists_prop] },
   try { exact iff.rfl },
   done } <|> fail "Ce n'est pas cela. Essayez encore." ]

end tactic.interactive

notation `|`:1024 x:1 `|`:1 := abs x

namespace m154

lemma inferieur_ssi {x y : ℝ} : x ≤ y ↔ 0 ≤ y - x :=
sub_nonneg.symm

lemma pos_pos {x y : ℝ} (hx : 0 ≤ x) (hy : 0 ≤ y) : 0 ≤ x*y :=
mul_nonneg hx hy

lemma neg_neg {x y : ℝ} (hx : x ≤ 0) (hy : y ≤ 0) : 0 ≤ x*y :=
mul_nonneg_of_nonpos_of_nonpos hx hy

lemma inferieur_ssi' {x y : ℝ} : x ≤ y ↔ x - y ≤ 0 :=
by rw [show x-y = -(y-x), by ring, inferieur_ssi, neg_le, neg_zero]
end m154

open nat
def pgcd := nat.gcd

lemma divise_refl (a : ℕ) : a ∣ a :=
dvd_refl a

lemma divise_pgcd_ssi {a b c : ℕ} : c ∣ pgcd a b ↔ c ∣ a ∧ c ∣ b :=
dvd_gcd_iff

lemma divise_antisym {a b : ℕ} : a ∣ b → b ∣ a → a = b :=
dvd_antisymm

lemma divise_def (a b : ℤ) : a ∣ b ↔ ∃ k, b = a*k :=
iff.rfl

def pair (n : ℤ) := ∃ k, n = 2*k
def impair (n : ℤ) := ∃ k, n = 2*k + 1

lemma pair_ou_impair (n : ℤ) : pair n ∨ impair n :=
by by_cases h : n % 2 = 0 ; [left, {right ; rw int.mod_two_ne_zero at h}] ;
  rw [← int.mod_add_div n 2, h] ; use n/2 ; ring

lemma non_pair_et_impair (n : ℤ) : ¬ (pair n ∧ impair n) :=
begin
  rintro ⟨h, h'⟩,
  change even n at h,
  rw int.even_iff at h,
  rcases h' with ⟨k, rfl⟩,
  simp only [int.add_mul_mod_self_left, add_comm, euclidean_domain.mod_eq_zero] at h,
  cases h with l hl,
  rw eq_comm at hl,
  have := int.eq_one_of_mul_eq_one_right (by linarith) hl,
  linarith
end

lemma abs_inferieur_ssi (x y : ℝ) : |x| ≤ y ↔ -y ≤ x ∧ x ≤ y :=
abs_le

lemma abs_diff (x y : ℝ) : |x - y| = |y - x| :=
abs_sub_comm x y

lemma pos_abs (x : ℝ) : |x| > 0 ↔ x ≠ 0 :=
abs_pos

variables {α : Type*} [linear_order α]
lemma superieur_max_ssi (p q r : α) : r ≥ max p q  ↔ r ≥ p ∧ r ≥ q :=
max_le_iff

lemma inferieur_max_gauche (p q : α) : p ≤ max p q :=
le_max_left _ _

lemma inferieur_max_droite (p q : α) : q ≤ max p q :=
le_max_right _ _

lemma egal_si_abs_diff_neg {a b : ℝ} : |a - b| ≤ 0 → a = b :=
eq_of_abs_sub_nonpos

lemma egal_si_abs_eps (x y : ℝ) : (∀ ε > 0, |x - y| ≤ ε) → x = y :=
begin
  intro h,
  apply egal_si_abs_diff_neg,
  by_contradiction H,
  push_neg at H,
  specialize h ( |x-y|/2) (by linarith),
  linarith
end

lemma ineg_triangle (x y : ℝ) : |x + y| ≤ |x| + |y| :=
abs_add x y

namespace m154
def limite_suite (u : ℕ → ℝ) (l : ℝ) : Prop :=
∀ ε > 0, ∃ N, ∀ n ≥ N, |u n - l| ≤ ε

lemma unicite_limite {u l l'}: limite_suite u l → limite_suite u l' → l = l' :=
begin
  -- sorry
  intros hl hl',
  apply egal_si_abs_eps,
  intros ε ε_pos,
  specialize hl (ε/2) (by linarith),
  cases hl with N hN,
  specialize hl' (ε/2) (by linarith),
  cases hl' with N' hN',
  specialize hN (max N N') (inferieur_max_gauche _ _),
  specialize hN' (max N N') (inferieur_max_droite _ _),
  calc |l - l'| = |(l-u (max N N')) + (u (max N N') -l')| : by ring_nf
  ... ≤ |l - u (max N N')| + |u (max N N') - l'| : by apply ineg_triangle
  ... =  |u (max N N') - l| + |u (max N N') - l'| : by rw abs_diff
  ... ≤ ε/2 + ε/2 : by linarith
  ... = ε : by ring,
end

end m154

open m154

def extraction (φ : ℕ → ℕ) := ∀ n m, n < m → φ n < φ m

-- Dans la suite, φ désignera toujours une fonction de ℕ dans ℕ
variable { φ : ℕ → ℕ}

/-- Un réel `a` est valeur d'adhérence d'une suite `u` s'il
existe une suite extraite de `u` qui tend vers `a`. -/
def valeur_adherence (u : ℕ → ℝ) (a : ℝ) :=
∃ φ, extraction φ ∧ limite_suite (u ∘ φ) a

/-- Toute extraction est supérieure à l'identité. -/
lemma extraction_superieur_id : extraction φ → ∀ n, n ≤ φ n :=
begin
  intros hyp n,
  induction n with n hn,
    exact nat.zero_le _,
  exact nat.succ_le_of_lt (by linarith [hyp n (n+1) (by linarith)]),
end

open filter

lemma extraction_machine (ψ : ℕ → ℕ) (hψ : ∀ n, ψ n ≥ n) : ∃ f : ℕ → ℕ, extraction (ψ ∘ f) ∧ ∀ n, f n ≥ n :=
begin
  refine ⟨λ n, nat.rec_on n 0 (λ n ih, ψ ih + 1), λ m n h, _, λ n, _⟩,
  { induction h; dsimp [(∘)],
    { exact hψ _ },
    { exact lt_trans h_ih (hψ _) } },
  { induction n, {apply le_refl},
    exact nat.succ_le_succ (le_trans n_ih (hψ _)) }
end


/-
Copyright (c) 2020 Bhavik Mehta. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Bhavik Mehta
-/
import category_theory.limits.shapes.finite_products
import category_theory.limits.shapes.binary_products
import category_theory.limits.preserves.shapes.products
import category_theory.limits.preserves.shapes.binary_products
import category_theory.limits.shapes.pullbacks
import category_theory.pempty
import data.equiv.fin

/-!
# Constructing finite products from binary products and terminal.

If a category has binary products and a terminal object then it has finite products.
If a functor preserves binary products and the terminal object then it preserves finite products.

# TODO

Provide the dual results.
Show the analogous results for functors which reflect or create (co)limits.
-/

universes v u u'

noncomputable theory
open category_theory category_theory.category category_theory.limits
namespace category_theory

variables {J : Type v} [small_category J]
variables {C : Type u} [category.{v} C]
variables {D : Type u'} [category.{v} D]

/--
Given `n+1` objects of `C`, a fan for the last `n` with point `c₁.X` and a binary fan on `c₁.X` and
`f 0`, we can build a fan for all `n+1`.

In `extend_fan_is_limit` we show that if the two given fans are limits, then this fan is also a
limit.
-/
@[simps {rhs_md := semireducible}]
def extend_fan {n : ℕ} {f : ulift (fin (n+1)) → C}
  (c₁ : fan (λ (i : ulift (fin n)), f ⟨i.down.succ⟩))
  (c₂ : binary_fan (f ⟨0⟩) c₁.X) :
  fan f :=
fan.mk c₂.X
begin
  rintro ⟨i⟩,
  revert i,
  refine fin.cases _ _,
  { apply c₂.fst },
  { intro i,
    apply c₂.snd ≫ c₁.π.app (ulift.up i) },
end

/--
Show that if the two given fans in `extend_fan` are limits, then the constructed fan is also a
limit.
-/
def extend_fan_is_limit {n : ℕ} (f : ulift (fin (n+1)) → C)
  {c₁ : fan (λ (i : ulift (fin n)), f ⟨i.down.succ⟩)} {c₂ : binary_fan (f ⟨0⟩) c₁.X}
  (t₁ : is_limit c₁) (t₂ : is_limit c₂) :
  is_limit (extend_fan c₁ c₂) :=
{ lift := λ s,
  begin
    apply (binary_fan.is_limit.lift' t₂ (s.π.app ⟨0⟩) _).1,
    apply t₁.lift ⟨_, discrete.nat_trans (λ i, s.π.app ⟨i.down.succ⟩)⟩
  end,
  fac' := λ s,
  begin
    rintro ⟨j⟩,
    apply fin.induction_on j,
    { apply (binary_fan.is_limit.lift' t₂ _ _).2.1 },
    { rintro i -,
      dsimp only [extend_fan_π_app],
      rw [fin.cases_succ, ← assoc, (binary_fan.is_limit.lift' t₂ _ _).2.2, t₁.fac],
      refl }
  end,
  uniq' := λ s m w,
  begin
    apply binary_fan.is_limit.hom_ext t₂,
    { rw (binary_fan.is_limit.lift' t₂ _ _).2.1,
      apply w ⟨0⟩ },
    { rw (binary_fan.is_limit.lift' t₂ _ _).2.2,
      apply t₁.uniq ⟨_, _⟩,
      rintro ⟨j⟩,
      rw assoc,
      dsimp only [discrete.nat_trans_app],
      rw ← w ⟨j.succ⟩,
      dsimp only [extend_fan_π_app],
      rw fin.cases_succ }
  end }

section
variables [has_binary_products.{v} C] [has_terminal C]

/--
If `C` has a terminal object and binary products, then it has a product for objects indexed by
`ulift (fin n)`.
This is a helper lemma for `has_finite_products_of_has_binary_and_terminal`, which is more general
than this.
-/
private lemma has_product_ulift_fin :
  Π (n : ℕ) (f : ulift (fin n) → C), has_product f
| 0 := λ f,
  begin
    letI : has_limits_of_shape (discrete (ulift (fin 0))) C :=
      has_limits_of_shape_of_equivalence
        (discrete.equivalence (equiv.ulift.trans fin_zero_equiv').symm),
    apply_instance,
  end
| (n+1) := λ f,
  begin
    haveI := has_product_ulift_fin n,
    apply has_limit.mk ⟨_, extend_fan_is_limit f (limit.is_limit _) (limit.is_limit _)⟩,
  end

/--
If `C` has a terminal object and binary products, then it has limits of shape
`discrete (ulift (fin n))` for any `n : ℕ`.
This is a helper lemma for `has_finite_products_of_has_binary_and_terminal`, which is more general
than this.
-/
private lemma has_limits_of_shape_ulift_fin (n : ℕ) :
  has_limits_of_shape (discrete (ulift (fin n))) C :=
{ has_limit := λ K,
begin
  letI := has_product_ulift_fin n K.obj,
  let : discrete.functor K.obj ≅ K := discrete.nat_iso (λ i, iso.refl _),
  apply has_limit_of_iso this,
end }

/-- If `C` has a terminal object and binary products, then it has finite products. -/
lemma has_finite_products_of_has_binary_and_terminal : has_finite_products C :=
⟨λ J 𝒥₁ 𝒥₂, begin
  resetI,
  let e := fintype.equiv_fin J,
  apply has_limits_of_shape_of_equivalence (discrete.equivalence (e.trans equiv.ulift.symm)).symm,
  refine has_limits_of_shape_ulift_fin (fintype.card J),
end⟩

end

section preserves
variables (F : C ⥤ D)
variables [preserves_limits_of_shape (discrete walking_pair) F]
variables [preserves_limits_of_shape (discrete pempty) F]
variables [has_finite_products.{v} C]

/--
If `F` preserves the terminal object and binary products, then it preserves products indexed by
`ulift (fin n)` for any `n`.
-/
noncomputable def preserves_fin_of_preserves_binary_and_terminal  :
  Π (n : ℕ) (f : ulift (fin n) → C), preserves_limit (discrete.functor f) F
| 0 := λ f,
  begin
    letI : preserves_limits_of_shape (discrete (ulift (fin 0))) F :=
      preserves_limits_of_shape_of_equiv
        (discrete.equivalence (equiv.ulift.trans fin_zero_equiv').symm) _,
    apply_instance,
  end
| (n+1) :=
  begin
    haveI := preserves_fin_of_preserves_binary_and_terminal n,
    intro f,
    refine preserves_limit_of_preserves_limit_cone
      (extend_fan_is_limit f (limit.is_limit _) (limit.is_limit _)) _,
    apply (is_limit_map_cone_fan_mk_equiv _ _ _).symm _,
    let := extend_fan_is_limit (λ i, F.obj (f i))
              (is_limit_of_has_product_of_preserves_limit F _)
              (is_limit_of_has_binary_product_of_preserves_limit F _ _),
    refine is_limit.of_iso_limit this _,
    apply cones.ext _ _,
    apply iso.refl _,
    rintro ⟨j⟩,
    apply fin.induction_on j,
    { apply (category.id_comp _).symm },
    { rintro i -,
      dsimp only [extend_fan_π_app, iso.refl_hom, fan.mk_π_app],
      rw [fin.cases_succ, fin.cases_succ],
      change F.map _ ≫ _ = 𝟙 _ ≫ _,
      rw [id_comp, ←F.map_comp],
      refl }
  end

/--
If `F` preserves the terminal object and binary products, then it preserves limits of shape
`discrete (ulift (fin n))`.
-/
def preserves_ulift_fin_of_preserves_binary_and_terminal (n : ℕ) :
  preserves_limits_of_shape (discrete (ulift (fin n))) F :=
{ preserves_limit := λ K,
  begin
    let : discrete.functor K.obj ≅ K := discrete.nat_iso (λ i, iso.refl _),
    haveI := preserves_fin_of_preserves_binary_and_terminal F n K.obj,
    apply preserves_limit_of_iso_diagram F this,
  end }

/-- If `F` preserves the terminal object and binary products then it preserves finite products. -/
def preserves_finite_products_of_preserves_binary_and_terminal
  (J : Type v) [fintype J] :
  preserves_limits_of_shape.{v} (discrete J) F :=
begin
  classical,
  let e := fintype.equiv_fin J,
  haveI := preserves_ulift_fin_of_preserves_binary_and_terminal F (fintype.card J),
  apply preserves_limits_of_shape_of_equiv (discrete.equivalence (e.trans equiv.ulift.symm)).symm,
end

end preserves

/--
Given `n+1` objects of `C`, a cofan for the last `n` with point `c₁.X`
and a binary cofan on `c₁.X` and `f 0`, we can build a cofan for all `n+1`.

In `extend_cofan_is_colimit` we show that if the two given cofans are colimits,
then this cofan is also a colimit.
-/
@[simps {rhs_md := semireducible}]
def extend_cofan {n : ℕ} {f : ulift (fin (n+1)) → C}
  (c₁ : cofan (λ (i : ulift (fin n)), f ⟨i.down.succ⟩))
  (c₂ : binary_cofan (f ⟨0⟩) c₁.X) :
  cofan f :=
cofan.mk c₂.X
begin
  rintro ⟨i⟩,
  revert i,
  refine fin.cases _ _,
  { apply c₂.inl },
  { intro i,
    apply c₁.ι.app (ulift.up i) ≫ c₂.inr },
end

/--
Show that if the two given cofans in `extend_cofan` are colimits,
then the constructed cofan is also a colimit.
-/
def extend_cofan_is_colimit {n : ℕ} (f : ulift (fin (n+1)) → C)
  {c₁ : cofan (λ (i : ulift (fin n)), f ⟨i.down.succ⟩)} {c₂ : binary_cofan (f ⟨0⟩) c₁.X}
  (t₁ : is_colimit c₁) (t₂ : is_colimit c₂) :
  is_colimit (extend_cofan c₁ c₂) :=
{ desc := λ s,
  begin
    apply (binary_cofan.is_colimit.desc' t₂ (s.ι.app ⟨0⟩) _).1,
    apply t₁.desc ⟨_, discrete.nat_trans (λ i, s.ι.app ⟨i.down.succ⟩)⟩
  end,
  fac' := λ s,
  begin
    rintro ⟨j⟩,
    apply fin.induction_on j,
    { apply (binary_cofan.is_colimit.desc' t₂ _ _).2.1 },
    { rintro i -,
      dsimp only [extend_cofan_ι_app],
      rw [fin.cases_succ, assoc, (binary_cofan.is_colimit.desc' t₂ _ _).2.2, t₁.fac],
      refl }
  end,
  uniq' := λ s m w,
  begin
    apply binary_cofan.is_colimit.hom_ext t₂,
    { rw (binary_cofan.is_colimit.desc' t₂ _ _).2.1,
      apply w ⟨0⟩ },
    { rw (binary_cofan.is_colimit.desc' t₂ _ _).2.2,
      apply t₁.uniq ⟨_, _⟩,
      rintro ⟨j⟩,
      dsimp only [discrete.nat_trans_app],
      rw ← w ⟨j.succ⟩,
      dsimp only [extend_cofan_ι_app],
      rw [fin.cases_succ, assoc], }
  end }

section
variables [has_binary_coproducts.{v} C] [has_initial C]

/--
If `C` has an initial object and binary coproducts, then it has a coproduct for objects indexed by
`ulift (fin n)`.
This is a helper lemma for `has_cofinite_products_of_has_binary_and_terminal`, which is more general
than this.
-/
private lemma has_coproduct_ulift_fin :
  Π (n : ℕ) (f : ulift (fin n) → C), has_coproduct f
| 0 := λ f,
  begin
    letI : has_colimits_of_shape (discrete (ulift (fin 0))) C :=
      has_colimits_of_shape_of_equivalence
        (discrete.equivalence (equiv.ulift.trans fin_zero_equiv').symm),
    apply_instance,
  end
| (n+1) := λ f,
  begin
    haveI := has_coproduct_ulift_fin n,
    apply has_colimit.mk
      ⟨_, extend_cofan_is_colimit f (colimit.is_colimit _) (colimit.is_colimit _)⟩,
  end

/--
If `C` has an initial object and binary coproducts, then it has colimits of shape
`discrete (ulift (fin n))` for any `n : ℕ`.
This is a helper lemma for `has_cofinite_products_of_has_binary_and_terminal`, which is more general
than this.
-/
private lemma has_colimits_of_shape_ulift_fin (n : ℕ) :
  has_colimits_of_shape (discrete (ulift (fin n))) C :=
{ has_colimit := λ K,
begin
  letI := has_coproduct_ulift_fin n K.obj,
  let : K ≅ discrete.functor K.obj := discrete.nat_iso (λ i, iso.refl _),
  apply has_colimit_of_iso this,
end }

/-- If `C` has an initial object and binary coproducts, then it has finite coproducts. -/
lemma has_finite_coproducts_of_has_binary_and_terminal : has_finite_coproducts C :=
⟨λ J 𝒥₁ 𝒥₂, begin
  resetI,
  let e := fintype.equiv_fin J,
  apply has_colimits_of_shape_of_equivalence (discrete.equivalence (e.trans equiv.ulift.symm)).symm,
  refine has_colimits_of_shape_ulift_fin (fintype.card J),
end⟩

end

section preserves
variables (F : C ⥤ D)
variables [preserves_colimits_of_shape (discrete walking_pair) F]
variables [preserves_colimits_of_shape (discrete pempty) F]
variables [has_finite_coproducts.{v} C]

/--
If `F` preserves the initial object and binary coproducts, then it preserves products indexed by
`ulift (fin n)` for any `n`.
-/
noncomputable def preserves_fin_of_preserves_binary_and_initial  :
  Π (n : ℕ) (f : ulift (fin n) → C), preserves_colimit (discrete.functor f) F
| 0 := λ f,
  begin
    letI : preserves_colimits_of_shape (discrete (ulift (fin 0))) F :=
      preserves_colimits_of_shape_of_equiv
        (discrete.equivalence (equiv.ulift.trans fin_zero_equiv').symm) _,
    apply_instance,
  end
| (n+1) :=
  begin
    haveI := preserves_fin_of_preserves_binary_and_initial n,
    intro f,
    refine preserves_colimit_of_preserves_colimit_cocone
      (extend_cofan_is_colimit f (colimit.is_colimit _) (colimit.is_colimit _)) _,
    apply (is_colimit_map_cocone_cofan_mk_equiv _ _ _).symm _,
    let := extend_cofan_is_colimit (λ i, F.obj (f i))
              (is_colimit_of_has_coproduct_of_preserves_colimit F _)
              (is_colimit_of_has_binary_coproduct_of_preserves_colimit F _ _),
    refine is_colimit.of_iso_colimit this _,
    apply cocones.ext _ _,
    apply iso.refl _,
    rintro ⟨j⟩,
    apply fin.induction_on j,
    { apply category.comp_id },
    { rintro i -,
      dsimp only [extend_cofan_ι_app, iso.refl_hom, cofan.mk_ι_app],
      rw [fin.cases_succ, fin.cases_succ],
      erw [comp_id, ←F.map_comp],
      refl, }
  end

/--
If `F` preserves the initial object and binary coproducts, then it preserves colimits of shape
`discrete (ulift (fin n))`.
-/
def preserves_ulift_fin_of_preserves_binary_and_initial (n : ℕ) :
  preserves_colimits_of_shape (discrete (ulift (fin n))) F :=
{ preserves_colimit := λ K,
  begin
    let : discrete.functor K.obj ≅ K := discrete.nat_iso (λ i, iso.refl _),
    haveI := preserves_fin_of_preserves_binary_and_initial F n K.obj,
    apply preserves_colimit_of_iso_diagram F this,
  end }

/-- If `F` preserves the initial object and binary coproducts then it preserves finite products. -/
def preserves_finite_coproducts_of_preserves_binary_and_initial
  (J : Type v) [fintype J] :
  preserves_colimits_of_shape.{v} (discrete J) F :=
begin
  classical,
  let e := fintype.equiv_fin J,
  haveI := preserves_ulift_fin_of_preserves_binary_and_initial F (fintype.card J),
  apply preserves_colimits_of_shape_of_equiv (discrete.equivalence (e.trans equiv.ulift.symm)).symm,
end

end preserves

end category_theory


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

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

instance : inhabited (wseq α) := ⟨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 is_finite (s : wseq α) : Prop := (out : (to_list s).terminates)

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

/-- 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 productive (s : wseq α) : Prop := (nth_terminates : ∀ n, (nth s n).terminates)

theorem productive_iff (s : wseq α) : productive s ↔ ∀ n, (nth s n).terminates :=
⟨λ h, h.1, λ h, ⟨h⟩⟩

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

instance head_terminates (s : wseq α) [productive s] :
  (head s).terminates := s.nth_terminates 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 ` ~ `:50 := 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) :
  swap (lift_rel_o R C) = lift_rel_o (swap R) (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 (swap R) s2 s1 :=
begin
  refine ⟨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) :
  swap (lift_rel R) = lift_rel (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 : swap (lift_rel R) s2 s1),
by rwa [lift_rel.swap, show 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.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
  simp [tail], rw [← bind_pure_comp_eq_map, is_lawful_monad.bind_assoc],
  apply congr_arg, ext1 (_|⟨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
  rcases (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
  cases e : computation.get (nth s (n + 1)), {assumption},
  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
  rcases (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 :=
by simp only [productive_iff]; exact
  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,
  ext ⟨⟩; 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,
    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,
  { apply t.cases_on _ (λ b t, _) (λ t, _); 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,
    { 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 only [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 (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,
      have : ∀ s, ∃ s' : wseq α, (map ret s).join.destruct = (map ret s').join.destruct ∧
        destruct s = s'.destruct, from λ s, ⟨s, rfl, rfl⟩,
      apply s.cases_on _ (λ a s, _) (λ s, _); simp [ret, ret_mem, this, option.exists]
    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,
    { apply wseq.cases_on S _ (λ s S, _) (λ S, _); simp,
      { apply wseq.cases_on T _ (λ s T, _) (λ T, _); 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,
      { apply wseq.cases_on S _ (λ s S, _) (λ S, _); 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,
    { apply wseq.cases_on S _ (λ s S, _) (λ S, _); simp,
      { apply wseq.cases_on SS _ (λ S SS, _) (λ SS, _); 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


inductive Mem' (a : α) : List α → Prop where
  | intro (as bs) : Mem' a (as ++ (a :: bs))

example {x : α} (h : Mem' x l) : True :=
  match h with
  | ⟨as', bs'⟩ => True.intro

example {x : α} (h : Mem' x l ∧ Mem' x l') : True :=
  match h with
  | ⟨⟨as', bs'⟩, _⟩ => True.intro


import data.fin.tuple.basic
import data.finsupp.basic
import data.finsupp.pointwise
import data.list.basic
import data.list.range
import data.list.fin_range
import data.option.basic
import data.pfun
import data.prod.lex

lemma bool.coe_iff_eq_tt (b : bool) : b ↔ b = tt := iff.rfl
@[simp] lemma option.bind_const_none {α β} (x : option α) :
  x.bind (λ _, none) = (none : option β) :=
by cases x; simp
@[simp] lemma option.is_none_ff_iff_is_some {α} (x : option α) :
  x.is_none = ff ↔ x.is_some = tt :=
by cases x; simp

@[simp] lemma fin.tuple_eval_one {n : ℕ} {α : fin (n + 2) → Type*}
  (x₀ : α 0) (x₁ : α 1) (x₂ : Π i : fin n, α i.succ.succ) :
  fin.cons x₀ (fin.cons x₁ x₂) 1 = x₁ := rfl

@[simp] lemma fin.tuple_eval_two {n : ℕ} {α : fin (n + 3) → Type*}
  (x₀ : α 0) (x₁ : α 1) (x₂ : α 2) (x₃ : Π i : fin n, α i.succ.succ.succ) :
  fin.cons x₀ (fin.cons x₁ (fin.cons x₂ x₃)) 2 = x₂ := rfl

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

/- Technically `m` only needs to be applicative but whatever -/
def fin.tuple_sequence : ∀ {n : ℕ} {α : fin n → Type u},
  (Π i, m (α i)) → m (Π i, α i)
| 0 _ _ := pure default
| (n+1) α x := (x 0) >>= λ r, (fin.tuple_sequence (fin.tail x)) >>= λ x', return (fin.cons r x')

@[simp] lemma fin.tuple_sequence₁ [is_lawful_monad m] {α : fin 1 → Type u} (x : Π i, m (α i)) :
  fin.tuple_sequence x = (x 0) >>= λ r₀, pure (fin.cons r₀ default) :=
by simp [fin.tuple_sequence] with functor_norm

@[simp] lemma fin.tuple_sequence₂ [is_lawful_monad m] {α : fin 2 → Type u} (x : Π i, m (α i)) :
  fin.tuple_sequence x = (x 0) >>= λ r₀, (x 1) >>= λ r₁, pure (fin.cons r₀ $ fin.cons r₁ $ default) :=
by { simp [fin.tuple_sequence] with functor_norm, refl, }

def fin.tuple_some {n : ℕ} {α : fin n → Type u} (x : Π i, option (α i)) : option (Π i, α i) :=
fin.tuple_sequence x

@[simp] lemma option.bind_is_some {α β} (x : option α) (y : α → option β):
  (x >>= y).is_some ↔ (∃ (h : x.is_some), (y (option.get h)).is_some) :=
by cases x; simp

@[simp] lemma option.map_is_some {α β} (x : option α) (y : α → β) :
  (y <$> x).is_some = x.is_some := by cases x; simp

@[simp] lemma option.not_is_some' {α} (x : option α) : !x.is_some = x.is_none :=
by { cases x; simp, }

def option.guard_prop {α} (p : Prop) [decidable p] (x : α) : option α :=
  if p then some x else none
@[simp] lemma option.guard_prop_is_some {α} {p : Prop} [decidable p] {x : α} :
  (option.guard_prop p x).is_some ↔ p :=
by { dsimp only [option.guard_prop], split_ifs; simpa }

@[simp] lemma option.coe_part_dom {α} (x : option α) :
  (x : part α).dom ↔ x.is_some := by cases x; simp

@[simp] lemma option.coe_part_eq_some {α} (x : option α) (y : α) :
  (x : part α) = part.some y ↔ x = some y :=
by simp [part.eq_some_iff]

@[simp] lemma list.nth_is_some_iff {α} {x : list α} {n : ℕ} :
  (x.nth n).is_some ↔ n < x.length :=
by { rw ← not_iff_not, simp [option.is_none_iff_eq_none], }

@[simp] lemma option.map_is_some' {α β} (x : option α) (f : α → β) :
  (x.map f).is_some = x.is_some := by cases x; simp

lemma list.zip_with_fst {α β} {l₁ : list α} {l₂ : list β} (hl : l₁.length ≤ l₂.length) :
  list.zip_with (λ a b, a) l₁ l₂ = l₁ :=
by { erw [← list.map_uncurry_zip_eq_zip_with, list.map_fst_zip], exact hl, }

lemma list.zip_with_snd {α β} {l₁ : list α} {l₂ : list β} (hl : l₂.length ≤ l₁.length) :
  list.zip_with (λ a b, b) l₁ l₂ = l₂ :=
by { erw [← list.map_uncurry_zip_eq_zip_with, list.map_snd_zip], exact hl, }

@[simp] lemma multiset.map_nth_le {α} {n : ℕ} {l : list α} (hn : l.length = n) :
    (finset.univ.val : multiset (fin n)).map (λ i, l.nth_le i (by rw hn; exact i.prop)) = l :=
by { subst hn, simp [finset.univ, fintype.elems], erw list.map_nth_le, }

@[simp] lemma le_ff_iff {b : bool} : b ≤ ff ↔ b = ff :=
by cases b; simp

lemma ne_min_of_ne_and_ne {ι : Type*} [linear_order ι] {a x y : ι} (hx : a ≠ x) (hy : a ≠ y) :
  a ≠ min x y := by cases min_choice x y with h; rw h; assumption

@[simp] lemma max_ne_self_iff {ι : Type*} [linear_order ι] (a b : ι) :
  ¬(a = max a b) ↔ a < b :=
by { rw max_def, split_ifs, { simpa using h }, { simpa using le_of_not_ge h } }

@[simp] lemma max_ne_self_iff' {ι : Type*} [linear_order ι] (a b : ι) :
  ¬(b = max a b) ↔ b < a :=
by { rw max_comm, simp, }

section with_top
variables {ι : Type} [partial_order ι]

@[simp] lemma with_top.is_some_iff_lt_top {x : with_top ι} :
  x.is_some ↔ x < ⊤ :=
by { rw [← not_iff_not, eq_ff_eq_not_eq_tt, option.not_is_some, option.is_none_iff_eq_none, lt_top_iff_ne_top, ne, not_not], refl, }

lemma prod.lex.le_iff' {α β : Type} [has_lt α] [has_le β] {x y : α ×ₗ β} :
  x ≤ y ↔ x.1 < y.1 ∨ (x.1 = y.1 ∧ x.2 ≤ y.2) := prod.lex_def _ _

lemma prod.lex.le_iff'' {α β : Type} [partial_order α] [preorder β] {x y : α ×ₗ β} :
  x ≤ y ↔ x.1 ≤ y.1 ∧ (x.1 = y.1 → x.2 ≤ y.2) :=
by { rw [prod.lex.le_iff', le_iff_lt_or_eq], have := @ne_of_lt _ _ x.1 y.1, tauto!, }

lemma prod.lex.lt_iff' {α β : Type} [has_lt α] [has_lt β] {x y : α ×ₗ β} :
  x < y ↔ x.1 < y.1 ∨ (x.1 = y.1 ∧ x.2 < y.2) := prod.lex_def _ _

lemma prod.lex.fst_le_of_le {α β : Type} [preorder α] [preorder β] {x y : α ×ₗ β} (h : x ≤ y) : x.1 ≤ y.1 :=
by { rw prod.lex.le_iff' at h, cases h, { exact h.le, }, { exact h.1.le, }, }

lemma prod.lex.fst_lt_of_lt_of_le {α β : Type} [preorder α] [partial_order β] {x y : α ×ₗ β}
  (h : x < y) (h' : y.2 ≤ x.2) : x.1 < y.1 :=
by { rw prod.lex.lt_iff' at h, cases h, { exact h, }, cases h.2.not_le h', }

@[simp, norm_cast] lemma prod.with_top.coe_inj {α : Type} (x y : α) : (x : with_top α) = y ↔ x = y :=
option.some_inj

end with_top

/- NOTE: This stuff is already in mathlib (`data.finsupp.pointwise`) -/

-- variables {ι α : Type}

-- noncomputable instance finsupp.has_mul [mul_zero_class α] : has_mul (ι →₀ α) :=
-- ⟨λ a b, finsupp.zip_with (*) (zero_mul _) a b⟩

-- lemma finsupp.mul_apply [mul_zero_class α] (g₁ g₂ : ι →₀ α) (a : ι) : (g₁ * g₂) a = g₁ a * g₂ a := rfl

-- -- #check pi.distrib -- todo, tactic like this?
-- noncomputable instance finsupp.non_unital_semiring [non_unital_semiring α] : non_unital_semiring (ι →₀ α) :=
-- {
--   zero := 0,
--   add_assoc := λ a b c, fun_like.ext _ _ (by simp [finsupp.add_apply, add_assoc]),
--   zero_add  := λ a,     fun_like.ext _ _ (by simp [finsupp.add_apply]),
--   add_zero  := λ a,     fun_like.ext _ _ (by simp [finsupp.add_apply]),
--   add_comm  := λ a b,   fun_like.ext _ _ (by simp [finsupp.add_apply, add_comm] ),
--   zero_mul  := λ a,     fun_like.ext _ _ (by simp [finsupp.mul_apply]),
--   mul_zero  := λ a,     fun_like.ext _ _ (by simp [finsupp.mul_apply]),

--   left_distrib  := λ a b c, by simp [fun_like.ext_iff, finsupp.mul_apply, finsupp.add_apply, left_distrib],
--   right_distrib := λ a b c, by simp [fun_like.ext_iff, finsupp.mul_apply, finsupp.add_apply, right_distrib],

--   mul_assoc     := λ a b c, by simp [fun_like.ext_iff, finsupp.mul_apply, mul_assoc],

--   ..finsupp.has_mul, ..finsupp.has_add, }

@[simp] lemma finsupp.mul_single {ι β : Type*} [mul_zero_class β] (i : ι) (x y : β) :
  (finsupp.single i x) * (finsupp.single i y) = finsupp.single i (x * y) :=
by { ext a, by_cases i = a; simp [h], }

lemma finsupp.mul_eq_zero_of_disjoint_support {ι β : Type*} [decidable_eq ι] [mul_zero_class β] (f g : ι →₀ β) (h : disjoint f.support g.support) :
  f * g = 0 :=
begin
  rw [← finsupp.support_eq_empty, ← finset.subset_empty],
  refine trans finsupp.support_mul _,
  rwa [finset.subset_empty, ← finset.disjoint_iff_inter_eq_empty],
end

lemma finsupp.mul_single_eq_zero {ι β : Type*} [mul_zero_class β] (i₁ i₂ : ι) (hi : i₁ ≠ i₂) (x y : β) :
  (finsupp.single i₁ x) * (finsupp.single i₂ y) = 0 :=
by { classical, rw [finsupp.mul_eq_zero_of_disjoint_support], simp [finset.disjoint_iff_ne, finsupp.mem_support_single], intros, exact hi, }

set_option pp.implicit true