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

import category_theory.whiskering
import category_theory.yoneda
import category_theory.limits.cones

open category_theory category_theory.category category_theory.functor

namespace category_theory.limits

universes v u u' u'' w -- declare the `v`'s first; see `category_theory.category` for an explanation

variables {J : Type v} [small_category J]
variables {C : Type u} [𝒞 : category.{v} C]
include 𝒞

variables {F : J ⥤ C}

/-- A cone `t` on `F` is a limit cone if each cone on `F` admits a unique
  cone morphism to `t`. -/
structure is_limit (t : cone F) :=
(lift  : Π (s : cone F), s.X ⟶ t.X)
(fac'  : ∀ (s : cone F) (j : J), lift s ≫ t.π.app j = s.π.app j . obviously)
(uniq' : ∀ (s : cone F) (m : s.X ⟶ t.X) (w : ∀ j : J, m ≫ t.π.app j = s.π.app j),
  m = lift s . obviously)

restate_axiom is_limit.fac'
attribute [simp] is_limit.fac
restate_axiom is_limit.uniq'

namespace is_limit

instance subsingleton {t : cone F} : subsingleton (is_limit t) :=
⟨by intros P Q; cases P; cases Q; congr; ext; solve_by_elim⟩

/- Repackaging the definition in terms of cone morphisms. -/

def lift_cone_morphism {t : cone F} (h : is_limit t) (s : cone F) : s ⟶ t :=
{ hom := h.lift s }

lemma uniq_cone_morphism {s t : cone F} (h : is_limit t) {f f' : s ⟶ t} :
  f = f' :=
have ∀ {g : s ⟶ t}, g = h.lift_cone_morphism s, by intro g; ext; exact h.uniq _ _ g.w,
this.trans this.symm

def mk_cone_morphism {t : cone F}
  (lift : Π (s : cone F), s ⟶ t)
  (uniq' : ∀ (s : cone F) (m : s ⟶ t), m = lift s) : is_limit t :=
{ lift := λ s, (lift s).hom,
  uniq' := λ s m w,
    have cone_morphism.mk m w = lift s, by apply uniq',
    congr_arg cone_morphism.hom this }

/-- Limit cones on `F` are unique up to isomorphism. -/
def unique {s t : cone F} (P : is_limit s) (Q : is_limit t) : s ≅ t :=
{ hom := Q.lift_cone_morphism s,
  inv := P.lift_cone_morphism t,
  hom_inv_id' := P.uniq_cone_morphism,
  inv_hom_id' := Q.uniq_cone_morphism }

def of_iso_limit {r t : cone F} (P : is_limit r) (i : r ≅ t) : is_limit t :=
is_limit.mk_cone_morphism
  (λ s, P.lift_cone_morphism s ≫ i.hom)
  (λ s m, by rw ←i.comp_inv_eq; apply P.uniq_cone_morphism)

variables {t : cone F}

lemma hom_lift (h : is_limit t) {W : C} (m : W ⟶ t.X) :
  m = h.lift { X := W, π := { app := λ b, m ≫ t.π.app b } } :=
h.uniq { X := W, π := { app := λ b, m ≫ t.π.app b } } m (λ b, rfl)

/-- Two morphisms into a limit are equal if their compositions with
  each cone morphism are equal. -/
lemma hom_ext (h : is_limit t) {W : C} {f f' : W ⟶ t.X}
  (w : ∀ j, f ≫ t.π.app j = f' ≫ t.π.app j) : f = f' :=
by rw [h.hom_lift f, h.hom_lift f']; congr; exact funext w

/-- The universal property of a limit cone: a map `W ⟶ X` is the same as
  a cone on `F` with vertex `W`. -/
def hom_iso (h : is_limit t) (W : C) : (W ⟶ t.X) ≅ ((const J).obj W ⟹ F) :=
{ hom := λ f, (t.extend f).π,
  inv := λ π, h.lift { X := W, π := π },
  hom_inv_id' := by ext f; apply h.hom_ext; intro j; simp; dsimp; refl }

@[simp] lemma hom_iso_hom (h : is_limit t) {W : C} (f : W ⟶ t.X) :
  (is_limit.hom_iso h W).hom f = (t.extend f).π := rfl

/-- The limit of `F` represents the functor taking `W` to
  the set of cones on `F` with vertex `W`. -/
def nat_iso (h : is_limit t) : yoneda.obj t.X ≅ F.cones :=
nat_iso.of_components (λ W, is_limit.hom_iso h (unop W)) (by tidy)

def hom_iso' (h : is_limit t) (W : C) :
  (W ⟶ t.X) ≅ { p : Π j, W ⟶ F.obj j // ∀ {j j'} (f : j ⟶ j'), p j ≫ F.map f = p j' } :=
h.hom_iso W ≪≫
{ hom := λ π,
  ⟨λ j, π.app j, λ j j' f,
   by convert ←(π.naturality f).symm; apply id_comp⟩,
  inv := λ p,
  { app := λ j, p.1 j,
    naturality' := λ j j' f, begin dsimp, rw [id_comp], exact (p.2 f).symm end } }

/-- If G : C → D is a faithful functor which sends t to a limit cone,
  then it suffices to check that the induced maps for the image of t
  can be lifted to maps of C. -/
def of_faithful {t : cone F} {D : Type u'} [category.{v} D] (G : C ⥤ D) [faithful G]
  (ht : is_limit (G.map_cone t)) (lift : Π (s : cone F), s.X ⟶ t.X)
  (h : ∀ s, G.map (lift s) = ht.lift (G.map_cone s)) : is_limit t :=
{ lift := lift,
  fac' := λ s j, by apply G.injectivity; rw [G.map_comp, h]; apply ht.fac,
  uniq' := λ s m w, begin
    apply G.injectivity, rw h,
    refine ht.uniq (G.map_cone s) _ (λ j, _),
    convert ←congr_arg (λ f, G.map f) (w j),
    apply G.map_comp
  end }

end is_limit

def is_limit_iso_unique_cone_morphism {t : cone F} :
  is_limit t ≅ Π s, unique (s ⟶ t) :=
{ hom := λ h s,
  { default := h.lift_cone_morphism s,
    uniq := λ _, h.uniq_cone_morphism },
  inv := λ h,
  { lift := λ s, (h s).default.hom,
    uniq' := λ s f w, congr_arg cone_morphism.hom ((h s).uniq ⟨f, w⟩) } }

/-- A cocone `t` on `F` is a colimit cocone if each cocone on `F` admits a unique
  cocone morphism from `t`. -/
structure is_colimit (t : cocone F) :=
(desc  : Π (s : cocone F), t.X ⟶ s.X)
(fac'  : ∀ (s : cocone F) (j : J), t.ι.app j ≫ desc s = s.ι.app j . obviously)
(uniq' : ∀ (s : cocone F) (m : t.X ⟶ s.X) (w : ∀ j : J, t.ι.app j ≫ m = s.ι.app j),
  m = desc s . obviously)

restate_axiom is_colimit.fac'
attribute [simp] is_colimit.fac
restate_axiom is_colimit.uniq'

namespace is_colimit

instance subsingleton {t : cocone F} : subsingleton (is_colimit t) :=
⟨by intros P Q; cases P; cases Q; congr; ext; solve_by_elim⟩

/- Repackaging the definition in terms of cone morphisms. -/

def desc_cocone_morphism {t : cocone F} (h : is_colimit t) (s : cocone F) : t ⟶ s :=
{ hom := h.desc s }

lemma uniq_cocone_morphism {s t : cocone F} (h : is_colimit t) {f f' : t ⟶ s} :
  f = f' :=
have ∀ {g : t ⟶ s}, g = h.desc_cocone_morphism s, by intro g; ext; exact h.uniq _ _ g.w,
this.trans this.symm

def mk_cocone_morphism {t : cocone F}
  (desc : Π (s : cocone F), t ⟶ s)
  (uniq' : ∀ (s : cocone F) (m : t ⟶ s), m = desc s) : is_colimit t :=
{ desc := λ s, (desc s).hom,
  uniq' := λ s m w,
    have cocone_morphism.mk m w = desc s, by apply uniq',
    congr_arg cocone_morphism.hom this }

/-- Limit cones on `F` are unique up to isomorphism. -/
def unique {s t : cocone F} (P : is_colimit s) (Q : is_colimit t) : s ≅ t :=
{ hom := P.desc_cocone_morphism t,
  inv := Q.desc_cocone_morphism s,
  hom_inv_id' := P.uniq_cocone_morphism,
  inv_hom_id' := Q.uniq_cocone_morphism }

def of_iso_colimit {r t : cocone F} (P : is_colimit r) (i : r ≅ t) : is_colimit t :=
is_colimit.mk_cocone_morphism
  (λ s, i.inv ≫ P.desc_cocone_morphism s)
  (λ s m, by rw i.eq_inv_comp; apply P.uniq_cocone_morphism)

variables {t : cocone F}

lemma hom_desc (h : is_colimit t) {W : C} (m : t.X ⟶ W) :
  m = h.desc { X := W, ι := { app := λ b, t.ι.app b ≫ m,
    naturality' := by intros; erw [←assoc, t.ι.naturality, comp_id, comp_id] } } :=
h.uniq { X := W, ι := { app := λ b, t.ι.app b ≫ m, naturality' := _ } } m (λ b, rfl)

/-- Two morphisms out of a colimit are equal if their compositions with
  each cocone morphism are equal. -/
lemma hom_ext (h : is_colimit t) {W : C} {f f' : t.X ⟶ W}
  (w : ∀ j, t.ι.app j ≫ f = t.ι.app j ≫ f') : f = f' :=
by rw [h.hom_desc f, h.hom_desc f']; congr; exact funext w

/-- The universal property of a colimit cocone: a map `X ⟶ W` is the same as
  a cocone on `F` with vertex `W`. -/
def hom_iso (h : is_colimit t) (W : C) : (t.X ⟶ W) ≅ (F ⟹ (const J).obj W) :=
{ hom := λ f, (t.extend f).ι,
  inv := λ ι, h.desc { X := W, ι := ι },
  hom_inv_id' := by ext f; apply h.hom_ext; intro j; simp; dsimp; refl }

@[simp] lemma hom_iso_hom (h : is_colimit t) {W : C} (f : t.X ⟶ W) :
  (is_colimit.hom_iso h W).hom f = (t.extend f).ι := rfl

/-- The colimit of `F` represents the functor taking `W` to
  the set of cocones on `F` with vertex `W`. -/
def nat_iso (h : is_colimit t) : coyoneda.obj (op t.X) ≅ F.cocones :=
nat_iso.of_components (is_colimit.hom_iso h) (by intros; ext; dsimp; rw ←assoc; refl)

def hom_iso' (h : is_colimit t) (W : C) :
  (t.X ⟶ W) ≅ { p : Π j, F.obj j ⟶ W // ∀ {j j' : J} (f : j ⟶ j'), F.map f ≫ p j' = p j } :=
h.hom_iso W ≪≫
{ hom := λ ι,
  ⟨λ j, ι.app j, λ j j' f,
   by convert ←(ι.naturality f); apply comp_id⟩,
  inv := λ p,
  { app := λ j, p.1 j,
    naturality' := λ j j' f, begin dsimp, rw [comp_id], exact (p.2 f) end } }

/-- If G : C → D is a faithful functor which sends t to a colimit cocone,
  then it suffices to check that the induced maps for the image of t
  can be lifted to maps of C. -/
def of_faithful {t : cocone F} {D : Type u'} [category.{v} D] (G : C ⥤ D) [faithful G]
  (ht : is_colimit (G.map_cocone t)) (desc : Π (s : cocone F), t.X ⟶ s.X)
  (h : ∀ s, G.map (desc s) = ht.desc (G.map_cocone s)) : is_colimit t :=
{ desc := desc,
  fac' := λ s j, by apply G.injectivity; rw [G.map_comp, h]; apply ht.fac,
  uniq' := λ s m w, begin
    apply G.injectivity, rw h,
    refine ht.uniq (G.map_cocone s) _ (λ j, _),
    convert ←congr_arg (λ f, G.map f) (w j),
    apply G.map_comp
  end }

end is_colimit

def is_colimit_iso_unique_cocone_morphism {t : cocone F} :
  is_colimit t ≅ Π s, unique (t ⟶ s) :=
{ hom := λ h s,
  { default := h.desc_cocone_morphism s,
    uniq := λ _, h.uniq_cocone_morphism },
  inv := λ h,
  { desc := λ s, (h s).default.hom,
    uniq' := λ s f w, congr_arg cocone_morphism.hom ((h s).uniq ⟨f, w⟩) } }

section limit

/-- `has_limit F` represents a particular chosen limit of the diagram `F`. -/
class has_limit (F : J ⥤ C) :=
(cone : cone F)
(is_limit : is_limit cone)

variables (J C)

/-- `C` has limits of shape `J` if we have chosen a particular limit of
  every functor `F : J ⥤ C`. -/
@[class] def has_limits_of_shape := Π F : J ⥤ C, has_limit F

/-- `C` has all (small) limits if it has limits of every shape. -/
@[class] def has_limits :=
Π {J : Type v} {𝒥 : small_category J}, by exactI has_limits_of_shape J C

variables {J C}

instance has_limit_of_has_limits_of_shape
  {J : Type v} [small_category J] [H : has_limits_of_shape J C] (F : J ⥤ C) : has_limit F :=
H F

instance has_limits_of_shape_of_has_limits
  {J : Type v} [small_category J] [H : has_limits.{v} C] : has_limits_of_shape J C :=
H

/- Interface to the `has_limit` class. -/

def limit.cone (F : J ⥤ C) [has_limit F] : cone F := has_limit.cone F

def limit (F : J ⥤ C) [has_limit F] := (limit.cone F).X

def limit.π (F : J ⥤ C) [has_limit F] (j : J) : limit F ⟶ F.obj j :=
(limit.cone F).π.app j

@[simp] lemma limit.cone_π {F : J ⥤ C} [has_limit F] (j : J) :
  (limit.cone F).π.app j = limit.π _ j := rfl

@[simp] lemma limit.w (F : J ⥤ C) [has_limit F] {j j' : J} (f : j ⟶ j') :
  limit.π F j ≫ F.map f = limit.π F j' := (limit.cone F).w f

def limit.is_limit (F : J ⥤ C) [has_limit F] : is_limit (limit.cone F) :=
has_limit.is_limit.{v} F

def limit.lift (F : J ⥤ C) [has_limit F] (c : cone F) : c.X ⟶ limit F :=
(limit.is_limit F).lift c

@[simp] lemma limit.is_limit_lift {F : J ⥤ C} [has_limit F] (c : cone F) :
  (limit.is_limit F).lift c = limit.lift F c := rfl

@[simp] lemma limit.lift_π {F : J ⥤ C} [has_limit F] (c : cone F) (j : J) :
  limit.lift F c ≫ limit.π F j = c.π.app j :=
is_limit.fac _ c j

def limit.cone_morphism {F : J ⥤ C} [has_limit F] (c : cone F) :
  cone_morphism c (limit.cone F) :=
(limit.is_limit F).lift_cone_morphism c

@[simp] lemma limit.cone_morphism_hom {F : J ⥤ C} [has_limit F] (c : cone F) :
  (limit.cone_morphism c).hom = limit.lift F c := rfl
@[simp] lemma limit.cone_morphism_π {F : J ⥤ C} [has_limit F] (c : cone F) (j : J) :
  (limit.cone_morphism c).hom ≫ limit.π F j = c.π.app j :=
by erw is_limit.fac

@[extensionality] lemma limit.hom_ext {F : J ⥤ C} [has_limit F] {X : C} {f f' : X ⟶ limit F}
  (w : ∀ j, f ≫ limit.π F j = f' ≫ limit.π F j) : f = f' :=
(limit.is_limit F).hom_ext w

def limit.hom_iso (F : J ⥤ C) [has_limit F] (W : C) : (W ⟶ limit F) ≅ (F.cones.obj (op W)) :=
(limit.is_limit F).hom_iso W

@[simp] lemma limit.hom_iso_hom (F : J ⥤ C) [has_limit F] {W : C} (f : W ⟶ limit F):
  (limit.hom_iso F W).hom f = (const J).map f ≫ (limit.cone F).π :=
(limit.is_limit F).hom_iso_hom f

def limit.hom_iso' (F : J ⥤ C) [has_limit F] (W : C) :
  (W ⟶ limit F) ≅ { p : Π j, W ⟶ F.obj j // ∀ {j j' : J} (f : j ⟶ j'), p j ≫ F.map f = p j' } :=
(limit.is_limit F).hom_iso' W

lemma limit.lift_extend {F : J ⥤ C} [has_limit F] (c : cone F) {X : C} (f : X ⟶ c.X) :
  limit.lift F (c.extend f) = f ≫ limit.lift F c :=
by obviously

section pre
variables {K : Type v} [small_category K]
variables (F) [has_limit F] (E : K ⥤ J) [has_limit (E ⋙ F)]

def limit.pre : limit F ⟶ limit (E ⋙ F) :=
limit.lift (E ⋙ F)
  { X := limit F,
    π := { app := λ k, limit.π F (E.obj k) } }

@[simp] lemma limit.pre_π (k : K) : limit.pre F E ≫ limit.π (E ⋙ F) k = limit.π F (E.obj k) :=
by erw is_limit.fac

@[simp] lemma limit.lift_pre (c : cone F) :
  limit.lift F c ≫ limit.pre F E = limit.lift (E ⋙ F) (c.whisker E) :=
by ext; simp

variables {L : Type v} [small_category L]
variables (D : L ⥤ K) [has_limit (D ⋙ E ⋙ F)]

@[simp] lemma limit.pre_pre : limit.pre F E ≫ limit.pre (E ⋙ F) D = limit.pre F (D ⋙ E) :=
by ext j; erw [assoc, limit.pre_π, limit.pre_π, limit.pre_π]; refl

end pre

section post
variables {D : Type u'} [𝒟 : category.{v} D]
include 𝒟

variables (F) [has_limit F] (G : C ⥤ D) [has_limit (F ⋙ G)]

def limit.post : G.obj (limit F) ⟶ limit (F ⋙ G) :=
limit.lift (F ⋙ G)
{ X := G.obj (limit F),
  π :=
  { app := λ j, G.map (limit.π F j),
    naturality' :=
      by intros j j' f; erw [←G.map_comp, limits.cone.w, id_comp]; refl } }

@[simp] lemma limit.post_π (j : J) : limit.post F G ≫ limit.π (F ⋙ G) j = G.map (limit.π F j) :=
by erw is_limit.fac

@[simp] lemma limit.lift_post (c : cone F) :
  G.map (limit.lift F c) ≫ limit.post F G = limit.lift (F ⋙ G) (G.map_cone c) :=
by ext; rw [assoc, limit.post_π, ←G.map_comp, limit.lift_π, limit.lift_π]; refl

@[simp] lemma limit.post_post
  {E : Type u''} [category.{v} E] (H : D ⥤ E) [has_limit ((F ⋙ G) ⋙ H)] :
/- H G (limit F) ⟶ H (limit (F ⋙ G)) ⟶ limit ((F ⋙ G) ⋙ H) equals -/
/- H G (limit F) ⟶ limit (F ⋙ (G ⋙ H)) -/
  H.map (limit.post F G) ≫ limit.post (F ⋙ G) H = limit.post F (G ⋙ H) :=
by ext; erw [assoc, limit.post_π, ←H.map_comp, limit.post_π, limit.post_π]; refl

end post

lemma limit.pre_post {K : Type v} [small_category K] {D : Type u'} [category.{v} D]
  (E : K ⥤ J) (F : J ⥤ C) (G : C ⥤ D)
  [has_limit F] [has_limit (E ⋙ F)] [has_limit (F ⋙ G)] [has_limit ((E ⋙ F) ⋙ G)] :
/- G (limit F) ⟶ G (limit (E ⋙ F)) ⟶ limit ((E ⋙ F) ⋙ G) vs -/
/- G (limit F) ⟶ limit F ⋙ G ⟶ limit (E ⋙ (F ⋙ G)) or -/
  G.map (limit.pre F E) ≫ limit.post (E ⋙ F) G = limit.post F G ≫ limit.pre (F ⋙ G) E :=
by ext; erw [assoc, limit.post_π, ←G.map_comp, limit.pre_π, assoc, limit.pre_π, limit.post_π]; refl

section lim_functor

variables [has_limits_of_shape J C]

/-- `limit F` is functorial in `F`, when `C` has all limits of shape `J`. -/
def lim : (J ⥤ C) ⥤ C :=
{ obj := λ F, limit F,
  map := λ F G α, limit.lift G
    { X := limit F,
      π :=
      { app := λ j, limit.π F j ≫ α.app j,
        naturality' := λ j j' f,
          by erw [id_comp, assoc, ←α.naturality, ←assoc, limit.w] } },
  map_comp' := λ F G H α β,
    by ext; erw [assoc, is_limit.fac, is_limit.fac, ←assoc, is_limit.fac, assoc]; refl }

variables {F} {G : J ⥤ C} (α : F ⟹ G)

@[simp] lemma lim.map_π (j : J) : lim.map α ≫ limit.π G j = limit.π F j ≫ α.app j :=
by apply is_limit.fac

@[simp] lemma limit.lift_map (c : cone F) :
  limit.lift F c ≫ lim.map α = limit.lift G ((cones.postcompose α).obj c) :=
by ext; rw [assoc, lim.map_π, ←assoc, limit.lift_π, limit.lift_π]; refl

lemma limit.map_pre {K : Type v} [small_category K] [has_limits_of_shape K C] (E : K ⥤ J) :
  lim.map α ≫ limit.pre G E = limit.pre F E ≫ lim.map (whisker_left E α) :=
by ext; rw [assoc, limit.pre_π, lim.map_π, assoc, lim.map_π, ←assoc, limit.pre_π]; refl

lemma limit.map_pre' {K : Type v} [small_category K] [has_limits_of_shape.{v} K C]
  (F : J ⥤ C) {E₁ E₂ : K ⥤ J} (α : E₁ ⟹ E₂) :
  limit.pre F E₂ = limit.pre F E₁ ≫ lim.map (whisker_right α F) :=
by ext1; simp [(category.assoc _ _ _ _).symm]

lemma limit.id_pre (F : J ⥤ C) :
limit.pre F (functor.id _) = lim.map (functor.left_unitor F).inv := by tidy

lemma limit.map_post {D : Type u'} [category.{v} D] [has_limits_of_shape J D] (H : C ⥤ D) :
/- H (limit F) ⟶ H (limit G) ⟶ limit (G ⋙ H) vs
   H (limit F) ⟶ limit (F ⋙ H) ⟶ limit (G ⋙ H) -/
  H.map (lim.map α) ≫ limit.post G H = limit.post F H ≫ lim.map (whisker_right α H) :=
begin
  ext,
  rw [assoc, limit.post_π, ←H.map_comp, lim.map_π, H.map_comp],
  rw [assoc, lim.map_π, ←assoc, limit.post_π],
  refl
end

def lim_yoneda : lim ⋙ yoneda ≅ category_theory.cones J C :=
nat_iso.of_components (λ F, nat_iso.of_components (λ W, limit.hom_iso F (unop W)) (by tidy))
  (by tidy)

end lim_functor

end limit

section colimit

/-- `has_colimit F` represents a particular chosen colimit of the diagram `F`. -/
class has_colimit (F : J ⥤ C) :=
(cocone : cocone F)
(is_colimit : is_colimit cocone)

variables (J C)

/-- `C` has colimits of shape `J` if we have chosen a particular colimit of
  every functor `F : J ⥤ C`. -/
@[class] def has_colimits_of_shape := Π F : J ⥤ C, has_colimit F

/-- `C` has all (small) colimits if it has limits of every shape. -/
@[class] def has_colimits :=
Π {J : Type v} {𝒥 : small_category J}, by exactI has_colimits_of_shape J C

variables {J C}

instance has_colimit_of_has_colimits_of_shape
  {J : Type v} [small_category J] [H : has_colimits_of_shape J C] (F : J ⥤ C) : has_colimit F :=
H F

instance has_colimits_of_shape_of_has_colimits
  {J : Type v} [small_category J] [H : has_colimits.{v} C] : has_colimits_of_shape J C :=
H

/- Interface to the `has_colimit` class. -/

def colimit.cocone (F : J ⥤ C) [has_colimit F] : cocone F := has_colimit.cocone F

def colimit (F : J ⥤ C) [has_colimit F] := (colimit.cocone F).X

def colimit.ι (F : J ⥤ C) [has_colimit F] (j : J) : F.obj j ⟶ colimit F :=
(colimit.cocone F).ι.app j

@[simp] lemma colimit.cocone_ι {F : J ⥤ C} [has_colimit F] (j : J) :
  (colimit.cocone F).ι.app j = colimit.ι _ j := rfl

@[simp] lemma colimit.w (F : J ⥤ C) [has_colimit F] {j j' : J} (f : j ⟶ j') :
  F.map f ≫ colimit.ι F j' = colimit.ι F j := (colimit.cocone F).w f

def colimit.is_colimit (F : J ⥤ C) [has_colimit F] : is_colimit (colimit.cocone F) :=
has_colimit.is_colimit.{v} F

def colimit.desc (F : J ⥤ C) [has_colimit F] (c : cocone F) : colimit F ⟶ c.X :=
(colimit.is_colimit F).desc c

@[simp] lemma colimit.is_colimit_desc {F : J ⥤ C} [has_colimit F] (c : cocone F) :
  (colimit.is_colimit F).desc c = colimit.desc F c := rfl

@[simp] lemma colimit.ι_desc {F : J ⥤ C} [has_colimit F] (c : cocone F) (j : J) :
  colimit.ι F j ≫ colimit.desc F c = c.ι.app j :=
is_colimit.fac _ c j

def colimit.cocone_morphism {F : J ⥤ C} [has_colimit F] (c : cocone F) :
  cocone_morphism (colimit.cocone F) c :=
(colimit.is_colimit F).desc_cocone_morphism c

@[simp] lemma colimit.cocone_morphism_hom {F : J ⥤ C} [has_colimit F] (c : cocone F) :
  (colimit.cocone_morphism c).hom = colimit.desc F c := rfl
@[simp] lemma colimit.ι_cocone_morphism {F : J ⥤ C} [has_colimit F] (c : cocone F) (j : J) :
  colimit.ι F j ≫ (colimit.cocone_morphism c).hom = c.ι.app j :=
by erw is_colimit.fac

@[extensionality] lemma colimit.hom_ext {F : J ⥤ C} [has_colimit F] {X : C} {f f' : colimit F ⟶ X}
  (w : ∀ j, colimit.ι F j ≫ f = colimit.ι F j ≫ f') : f = f' :=
(colimit.is_colimit F).hom_ext w

def colimit.hom_iso (F : J ⥤ C) [has_colimit F] (W : C) : (colimit F ⟶ W) ≅ (F.cocones.obj W) :=
(colimit.is_colimit F).hom_iso W

@[simp] lemma colimit.hom_iso_hom (F : J ⥤ C) [has_colimit F] {W : C} (f : colimit F ⟶ W):
  (colimit.hom_iso F W).hom f = (colimit.cocone F).ι ≫ (const J).map f :=
(colimit.is_colimit F).hom_iso_hom f

def colimit.hom_iso' (F : J ⥤ C) [has_colimit F] (W : C) :
  (colimit F ⟶ W) ≅ { p : Π j, F.obj j ⟶ W // ∀ {j j'} (f : j ⟶ j'), F.map f ≫ p j' = p j } :=
(colimit.is_colimit F).hom_iso' W

lemma colimit.desc_extend (F : J ⥤ C) [has_colimit F] (c : cocone F) {X : C} (f : c.X ⟶ X) :
  colimit.desc F (c.extend f) = colimit.desc F c ≫ f :=
begin
  ext1, simp [category.assoc_symm], refl
end

section pre
variables {K : Type v} [small_category K]
variables (F) [has_colimit F] (E : K ⥤ J) [has_colimit (E ⋙ F)]

def colimit.pre : colimit (E ⋙ F) ⟶ colimit F :=
colimit.desc (E ⋙ F)
  { X := colimit F,
    ι := { app := λ k, colimit.ι F (E.obj k) } }

@[simp] lemma colimit.ι_pre (k : K) : colimit.ι (E ⋙ F) k ≫ colimit.pre F E = colimit.ι F (E.obj k) :=
by erw is_colimit.fac

@[simp] lemma colimit.pre_desc (c : cocone F) :
  colimit.pre F E ≫ colimit.desc F c = colimit.desc (E ⋙ F) (c.whisker E) :=
by ext; rw [←assoc, colimit.ι_pre]; simp

variables {L : Type v} [small_category L]
variables (D : L ⥤ K) [has_colimit (D ⋙ E ⋙ F)]

@[simp] lemma colimit.pre_pre : colimit.pre (E ⋙ F) D ≫ colimit.pre F E = colimit.pre F (D ⋙ E) :=
begin
  ext j,
  rw [←assoc, colimit.ι_pre, colimit.ι_pre],
  letI : has_colimit ((D ⋙ E) ⋙ F) := show has_colimit (D ⋙ E ⋙ F), by apply_instance,
  exact (colimit.ι_pre F (D ⋙ E) j).symm
end

end pre

section post
variables {D : Type u'} [𝒟 : category.{v} D]
include 𝒟

variables (F) [has_colimit F] (G : C ⥤ D) [has_colimit (F ⋙ G)]

def colimit.post : colimit (F ⋙ G) ⟶ G.obj (colimit F) :=
colimit.desc (F ⋙ G)
{ X := G.obj (colimit F),
  ι :=
  { app := λ j, G.map (colimit.ι F j),
    naturality' :=
      by intros j j' f; erw [←G.map_comp, limits.cocone.w, comp_id]; refl } }

@[simp] lemma colimit.ι_post (j : J) : colimit.ι (F ⋙ G) j ≫ colimit.post F G  = G.map (colimit.ι F j) :=
by erw is_colimit.fac

@[simp] lemma colimit.post_desc (c : cocone F) :
  colimit.post F G ≫ G.map (colimit.desc F c) = colimit.desc (F ⋙ G) (G.map_cocone c) :=
by ext; rw [←assoc, colimit.ι_post, ←G.map_comp, colimit.ι_desc, colimit.ι_desc]; refl

@[simp] lemma colimit.post_post
  {E : Type u''} [category.{v} E] (H : D ⥤ E) [has_colimit ((F ⋙ G) ⋙ H)] :
/- H G (colimit F) ⟶ H (colimit (F ⋙ G)) ⟶ colimit ((F ⋙ G) ⋙ H) equals -/
/- H G (colimit F) ⟶ colimit (F ⋙ (G ⋙ H)) -/
  colimit.post (F ⋙ G) H ≫ H.map (colimit.post F G) = colimit.post F (G ⋙ H) :=
begin
  ext,
  rw [←assoc, colimit.ι_post, ←H.map_comp, colimit.ι_post],
  exact (colimit.ι_post F (G ⋙ H) j).symm
end

end post

lemma colimit.pre_post {K : Type v} [small_category K] {D : Type u'} [category.{v} D]
  (E : K ⥤ J) (F : J ⥤ C) (G : C ⥤ D)
  [has_colimit F] [has_colimit (E ⋙ F)] [has_colimit (F ⋙ G)] [has_colimit ((E ⋙ F) ⋙ G)] :
/- G (colimit F) ⟶ G (colimit (E ⋙ F)) ⟶ colimit ((E ⋙ F) ⋙ G) vs -/
/- G (colimit F) ⟶ colimit F ⋙ G ⟶ colimit (E ⋙ (F ⋙ G)) or -/
  colimit.post (E ⋙ F) G ≫ G.map (colimit.pre F E) = colimit.pre (F ⋙ G) E ≫ colimit.post F G :=
begin
  ext,
  rw [←assoc, colimit.ι_post, ←G.map_comp, colimit.ι_pre, ←assoc],
  letI : has_colimit (E ⋙ F ⋙ G) := show has_colimit ((E ⋙ F) ⋙ G), by apply_instance,
  erw [colimit.ι_pre (F ⋙ G) E j, colimit.ι_post]
end

section colim_functor

variables [has_colimits_of_shape J C]

/-- `colimit F` is functorial in `F`, when `C` has all colimits of shape `J`. -/
def colim : (J ⥤ C) ⥤ C :=
{ obj := λ F, colimit F,
  map := λ F G α, colimit.desc F
    { X := colimit G,
      ι :=
      { app := λ j, α.app j ≫ colimit.ι G j,
        naturality' := λ j j' f,
          by erw [comp_id, ←assoc, α.naturality, assoc, colimit.w] } },
  map_comp' := λ F G H α β,
    by ext; erw [←assoc, is_colimit.fac, is_colimit.fac, assoc, is_colimit.fac, ←assoc]; refl }

variables {F} {G : J ⥤ C} (α : F ⟹ G)

@[simp] lemma colim.ι_map (j : J) : colimit.ι F j ≫ colim.map α = α.app j ≫ colimit.ι G j :=
by apply is_colimit.fac

@[simp] lemma colimit.map_desc (c : cocone G) :
  colim.map α ≫ colimit.desc G c = colimit.desc F ((cocones.precompose α).obj c) :=
by ext; rw [←assoc, colim.ι_map, assoc, colimit.ι_desc, colimit.ι_desc]; refl

lemma colimit.pre_map {K : Type v} [small_category K] [has_colimits_of_shape K C] (E : K ⥤ J) :
  colimit.pre F E ≫ colim.map α = colim.map (whisker_left E α) ≫ colimit.pre G E :=
by ext; rw [←assoc, colimit.ι_pre, colim.ι_map, ←assoc, colim.ι_map, assoc, colimit.ι_pre]; refl

lemma colimit.pre_map' {K : Type v} [small_category K] [has_colimits_of_shape.{v} K C]
  (F : J ⥤ C) {E₁ E₂ : K ⥤ J} (α : E₁ ⟹ E₂) :
  colimit.pre F E₁ = colim.map (whisker_right α F) ≫ colimit.pre F E₂ :=
by ext1; simp [(category.assoc _ _ _ _).symm]

lemma colimit.pre_id (F : J ⥤ C) :
colimit.pre F (functor.id _) = colim.map (functor.left_unitor F).hom := by tidy

lemma colimit.map_post {D : Type u'} [category.{v} D] [has_colimits_of_shape J D] (H : C ⥤ D) :
/- H (colimit F) ⟶ H (colimit G) ⟶ colimit (G ⋙ H) vs
   H (colimit F) ⟶ colimit (F ⋙ H) ⟶ colimit (G ⋙ H) -/
  colimit.post F H ≫ H.map (colim.map α) = colim.map (whisker_right α H) ≫ colimit.post G H:=
begin
  ext,
  rw [←assoc, colimit.ι_post, ←H.map_comp, colim.ι_map, H.map_comp],
  rw [←assoc, colim.ι_map, assoc, colimit.ι_post],
  refl
end

def colim_coyoneda : colim.op ⋙ coyoneda ≅ category_theory.cocones J C :=
nat_iso.of_components (λ F, nat_iso.of_components (colimit.hom_iso (unop F)) (by tidy))
  (by {tidy, rw [← category.assoc,← category.assoc], tidy})

end colim_functor

end colimit

end category_theory.limits


import Lean.Elab.Command
import Lean.Elab.Tactic.Basic
import Lean.Elab.Tactic.Location
import Lean.Elab.Tactic.Match
import Lean.Meta.Tactic.Split
import Lean.PrettyPrinter

import Lib.Data.List.Control
import Lib.Data.Array.Control
import Lib.Logic.Relation
import Lib.Meta.Dump

macro "rintro1 " t:term : tactic =>
  `(tactic| intro x; match x with | $t:term => ?_)

syntax "rintro " term,* : tactic

macro_rules
| `(tactic| rintro $([]),*) => `(tactic| skip)
| `(tactic| rintro $(t::ts),*) =>
  `(tactic| rintro1 $t <;> rintro $ts )

macro "obtain " p:term " from " d:term : tactic =>
  `(tactic| match $d:term with | $p:term => ?_)

macro "obtain " p:term " : " q:term " from " d:term : tactic =>
  `(tactic|
    have h : $q := $d;
    match h with | $p:term => ?_)

namespace Lean.Elab.Tactic
open Lean.Meta

elab "obtain" "!" p:term " from " d:term : tactic =>
  withMainContext do
    let (e, mvarIds') ← elabTermWithHoles d none `specialize (allowNaturalHoles := true)
    let h := e.getAppFn
    if h.isFVar then
      let localDecl ← getLocalDecl h.fvarId!
      let mvarId ← assert (← getMainGoal) localDecl.userName (← inferType e).headBeta e
      let (newHyp', mvarId) ← intro1P mvarId
      let mvarId ← tryClear mvarId h.fvarId!
      replaceMainGoal (mvarId :: mvarIds')
      withMainContext do
        let newHyp ← mkIdentFromRef
            <| (← getLCtx).get! newHyp' |>.userName
        let myMatch ←
          `(tactic| match $newHyp:term with | $p => ?_ )
        let stx ← getRef
        withMacroExpansion stx myMatch
          <| Lean.Elab.Tactic.evalMatch myMatch
        liftMetaTactic1 <| (tryClear . newHyp')
    else
      throwError "'specialize' requires a term of the form `h x_1 .. x_n` where `h` appears in the local context"

macro "obtain!" p:term " from " d:term : tactic =>
  `(tactic| obtain ! $p from $d)

end Lean.Elab.Tactic

macro "left" : tactic =>
  `(tactic| apply Or.inl)

macro "right" : tactic =>
  `(tactic| apply Or.inr)

syntax "refl" : tactic

macro "congr" : tactic =>
  `(tactic| repeat first | apply congrArg | apply congrFun | refl)

-- TODO: import mathlib and get rid of this
syntax "extOrSkip" (ident)* : tactic
syntax "ext" (ident)* : tactic

macro_rules
| `(tactic| extOrSkip) => `(tactic| skip)
| `(tactic| extOrSkip $x $xs) =>
  `(tactic| apply funext; intro $x; extOrSkip $xs)

macro_rules
| `(tactic| ext) => `(tactic| repeat (apply funext; intro))
| `(tactic| ext $xs) =>
  `(tactic| extOrSkip $xs)

syntax "change" term "with" term : tactic

macro "change" t:term "with" t':term : tactic =>
  `(tactic| have H : $t = $t' := rfl; rw [H]; clear H )

namespace Lean.Elab.Tactic

initialize registerTraceClass `refl
initialize registerTraceClass `substAll

open Expr Lean.Meta

def applyUnifyAll (g : MVarId) (lmm : Expr) (allowMVars := false) : MetaM (List MVarId) := do
  trace[auto.lemmas]"try {lmm}"
  trace[auto.lemmas]"type: {(← inferType lmm)}"
  trace[auto.goal]"goal {(← inferType (mkMVar g))}"
  let gs ← try apply g lmm
           catch e =>
             trace[auto.apply.failure]"error: {lmm}"
             trace[auto.apply.failure]"error: {e.toMessageData}"
             throw e
  trace[auto.lemmas]"{gs.length} new goals"
  let ls ← gs.mapM λ v => inferType (mkMVar v)
  trace[auto.lemmas]"{ls} all new goals"
  let b ← gs.allM λ v => do
    return (← inferType (← inferType (mkMVar v))).isProp
  if ¬ (allowMVars ∨ b) then
    trace[auto.lemmas]"failed"
    failure
  trace[auto.lemmas]"success"
  return gs

def tacRefl : TacticM Unit := do
  let g := (← instantiateMVars (← getMainTarget)).consumeMData
  match g with
  | (app (app R x _) y _) => liftMetaTactic λ g => do
    let cl ← mkAppOptM ``Reflexive #[none, R]
    let inst ← synthInstance cl
    let reflLmm ← mkAppOptM ``Reflexive.refl #[none, R, inst]
    apply g reflLmm
  | _ =>
    trace[refl]"ctorName: {g.ctorName}"
    let reflLmm ← mkConstWithFreshMVarLevels ``Reflexive.refl
    liftMetaTactic (apply . reflLmm) <|>
      throwError "Expection a reflexive relation: R x y"

def tacSymm : TacticM Unit := do
  let g := (← instantiateMVars (← getMainTarget)).consumeMData
  match g with
  | (app (app R x _) y _) => liftMetaTactic λ g => do
    let cl ← mkAppOptM ``Symmetric #[none, R]
    let inst ← synthInstance cl
    let symmLmm ← mkAppOptM ``Symmetric.symmetry #[none, R, inst]
    apply g symmLmm
  | _ =>
    trace[refl]"ctorName: {g.ctorName}"
    let symmLmm ← mkConstWithFreshMVarLevels ``Symmetric.symmetry
    liftMetaTactic (apply . symmLmm) <|>
      throwError "Expection a symmetric relation: R x y"

def tacSubstAll : TacticM Unit := do
  let lctx ← getLCtx
  trace[substAll] "hyps: {lctx.getFVars}"

  for h in lctx do
    let t := (← instantiateMVars h.type).consumeMData

    if ¬ h.isAuxDecl then
      match t with
      | app (app eq lhs _) rhs _ =>
        let lc ← getLCtx
        trace[substAll] "{h.userName} : {t}"
        if let Option.some _ := lc.find? h.fvarId then
          trace[substAll] "valid hyp"
          if eq.isAppOf `Eq ∧ (lhs.isFVar ∨ rhs.isFVar) then
            trace[substAll] "var eq"
            liftMetaTactic1 (subst .  h.fvarId)
      | _ =>
        trace[substAll] "ignore {h.userName} :  {t.ctorName} {t}"

def SearchT (δ : Type u) (m : Type u → Type v) (α : Type u) := (α → m δ) → m δ

namespace SearchT

variable {σ : Type} {m : Type → Type}
variable {α β : Type}

def pure (x : α) : SearchT δ m α :=
λ f => f x

def bind (x : SearchT δ m α) (f : α → SearchT δ m β) : SearchT δ m β :=
λ g => x λ a => f a g

instance : Monad (SearchT δ m) where
  pure := pure
  bind := bind

section Alternative
variable [Alternative m] [Monad m] [MonadBacktrack σ m]

def failure  : SearchT δ m α :=
λ f => Alternative.failure

def orElse (x : SearchT δ m α) (y : Unit → SearchT δ m α) : SearchT δ m α :=
λ f => do
  let s ← saveState
  Alternative.orElse (x f) λ _ => restoreState s >>= λ _ => y () f

instance : Alternative (SearchT δ m) where
  failure := failure
  orElse := orElse

end Alternative

def pick [Alternative m] [Monad m] [MonadBacktrack σ m] (l : List α) : SearchT δ m α :=
λ g => do
  let s ← saveState
  l.firstM λ a => do
    restoreState s
    g a

def pick' [Alternative m] [Monad m] [MonadBacktrack σ m] (l : Array α) : SearchT δ m α :=
λ g => do
  let s ← saveState
  l.firstM λ a => do
    restoreState s
    g a

instance [Bind m] : MonadLift m (SearchT δ m) where
  monadLift x f := x >>= f

instance [Bind m] [MonadEnv m] : MonadEnv (SearchT δ m) where
  getEnv f := getEnv >>= f
  modifyEnv x f := modifyEnv x >>= f

-- instance [Bind m] [MonadExceptOf ε m] : MonadExceptOf ε (SearchT δ m) where
--   throw e f := throw e
--   tryCatch x f g := _

instance [Monad m] [MonadRef m] : MonadRef (SearchT δ m) where
  getRef f := getRef >>= f
  withRef s x f := withRef s (x f)

instance [Bind m] [AddErrorMessageContext m] : AddErrorMessageContext (SearchT δ m) where
  add s m f := AddErrorMessageContext.add s m >>= f

def run [Pure m] (x : SearchT α m α) : m α := x Pure.pure

end SearchT

abbrev SearchTacticM δ := SearchT δ TacticM

namespace SearchTacticM

-- instance [Bind m] : MonadFunctor m (SearchT m) where
--   monadMap f x δ g := f _

def focus {α : Type} (x : SearchTacticM δ α) : SearchTacticM δ α := do
  let mvarId :: mvarIds ← getUnsolvedGoals | throwNoGoalsToBeSolved
  setGoals [mvarId]
  let a ← x
  let mvarIds' ← getUnsolvedGoals
  setGoals (mvarIds' ++ mvarIds)
  pure a

def focusAndDone {α} (tactic : SearchTacticM δ α) : SearchTacticM δ α :=
  focus do
    let a ← tactic
    done
    pure a

end SearchTacticM

def isDone : TacticM Bool :=
Option.isSome <$> optional done

def allGoals [Monad m] [MonadLiftT TacticM m] (tac : m PUnit) : m PUnit := do
let gs ← getGoals
let gs' ← gs.mapM λ g => do setGoals [g]; tac; getGoals
setGoals gs'.join

def tryTac [Alternative m] [Pure m] (x: m Unit) : m Unit :=
Alternative.orElse x (λ _ => pure ())

def iterate [Monad m] [MonadLiftT TacticM m]: Nat → m PUnit → m PUnit
| 0, _ => pure ()
| Nat.succ n, tac => do
  unless (← isDone) do
    -- tryTac (do
      ( traceM `auto.iterate <| return s!"iterate n = {n}" : TacticM Unit)
      tac
      allGoals $ iterate n tac

def tacMyApply (e : Expr) : TacticM Unit :=
  liftMetaTactic (applyUnifyAll . e)

end Lean.Elab.Tactic

open Lean.Elab.Tactic

elab "apply1" t:term : tactic =>
  withMainContext (do tacMyApply (← elabTerm t none))
elab "refl" : tactic => withMainContext Lean.Elab.Tactic.tacRefl
elab "symmetry" : tactic => withMainContext Lean.Elab.Tactic.tacSymm
elab "substAll" : tactic => withMainContext Lean.Elab.Tactic.tacSubstAll

macro "exfalso" : tactic =>
  `(tactic| apply False.elim)

macro "byContradiction" h: ident : tactic =>
  `(tactic| apply Classical.byContradiction; intro h)

syntax "trans" (term)? : tactic

macro_rules
| `(tactic| trans ) => `(tactic| trans ?middle)
| `(tactic| trans $t:term) =>
  `(tactic|
    -- show (?rel : ?α → ?α → ?t) ?x ?y ;
    focus
      refine' Trans.trans (r := ?rel) (s := ?rel) (t := ?rel)
                         (self := ?inst) (b := $t) ?first ?second;
      case inst => infer_instance
      rotate_right 2 )

open Lean.Elab.Tactic
open Lean

-- syntax (name := auto) "auto" : attr
-- syntax (name := eauto) "eauto" : attr

abbrev AutoExtension := SimpleScopedEnvExtension Name NameSet

def mkAutoAttr (attrName : Name) (attrDescr : String) (ext : AutoExtension) : IO Unit :=
  registerBuiltinAttribute {
    name  := attrName
    descr := attrDescr
    add   := fun declName stx attrKind =>
      let go : MetaM Unit := do
        let info ← getConstInfo declName
        match info with
        | ConstantInfo.inductInfo i =>
          for c in i.ctors do
            ext.add c attrKind
        | _ =>
          ext.add declName attrKind
      discard <| go.run {} {}
    erase := fun declName => do
      let s := ext.getState (← getEnv)
      let s := s.erase declName
      modifyEnv fun env => ext.modifyState env fun _ => s
  }

def mkAutoExt (extName : Name) : IO AutoExtension :=
  registerSimpleScopedEnvExtension {
    name     := extName
    initial  := {}
    addEntry := fun d e => d.insert e
  }

def registerAutoAttr (attrName : Name) (attrDescr : String) (extName : Name := attrName.appendAfter "Ext") : IO AutoExtension := do
  let ext ← mkAutoExt extName
  mkAutoAttr attrName attrDescr ext
  return ext

initialize autoExtension : AutoExtension ← registerAutoAttr `auto "auto closing lemma"

-- initialize autoLemmasAttr : TagAttribute ← registerTagAttribute `auto "auto lemmas"
-- initialize autoLemmasAttr : ParametricAttribute Unit ←
--   registerAutoAttribute {
--     name := `auto,
--     descr := "auto lemmas",
--     getParam := λ _ _ => () }
open Lean
-- open Tactic
open Lean Meta

-- initialize extLemmasCache : DeclCache (DiscrTree Name) ←
--   DeclCache.mk "ext: initialize cache" {} fun decl ci lemmas => do
--     if let some keys := extAttribute.getParam (← getEnv) decl then
--       lemmas.insertCore keys decl
--     else
--       lemmas

def getAutoLemmas [Monad m] [MonadEnv m] : m NameSet := do
  let ns := autoExtension.getState (← getEnv)
  return ns
    |>.insert ``True.intro
    |>.insert ``Iff.intro
    |>.insert ``And.intro

def getAutoList [Monad m] [MonadEnv m] (hyps : Array Name := #[]) : m (Array Name) := do
  return hyps ++ (← getAutoLemmas).toArray

initialize registerTraceClass `auto
initialize registerTraceClass `auto.apply
initialize registerTraceClass `auto.apply.attempts
initialize registerTraceClass `auto.apply.failure
initialize registerTraceClass `auto.destruct_hyp
initialize registerTraceClass `auto.goal
initialize registerTraceClass `auto.iterate
initialize registerTraceClass `auto.lemmas

open Lean

def Meta.applyAuto (ns : Array Name) (allowMVars := false) : SearchTacticM δ Unit :=
SearchTacticM.focus do
  let n ← SearchT.pick' ns
  traceM `auto.lemmas <| return s!"Lemma: {n}"
  let mut lmm ← (mkConstWithFreshMVarLevels n : TacticM _)
  Lean.Elab.Term.synthesizeSyntheticMVars true
  lmm ← instantiateMVars lmm
  liftMetaTactic (applyUnifyAll . lmm allowMVars)

def Meta.applyAssumption (allowMVars := false) : SearchTacticM δ Unit := SearchTacticM.focus $ do
  let x ← SearchT.pick' (← getLCtx).getFVarIds
  let lctx ← getLCtx
  guard (¬ (lctx.get! x).isAuxDecl)
  traceM `auto.lemmas <| return s!"Hyp: {lctx.get! x |>.userName}"
  liftMetaTactic (applyUnifyAll . (mkFVar x) allowMVars)

elab "#print" "auto_db" : command => do
  IO.println (← getAutoLemmas).toList

instance : ToMessageData LocalContext where
  toMessageData lctx := toMessageData
    <| lctx.fvarIdToDecl.toList.map
    <| LocalDecl.userName ∘ Prod.snd

def Meta.destructHyp : TacticM Unit := focus $ do
  let lctx ← getLCtx
  let mut changed := false
  trace[auto.destruct_hyp] "local context: {lctx}"
  for x in lctx do
    trace[auto.destruct_hyp]"local: {x.userName}"
    trace[auto.destruct_hyp]"is type: {x.type}"
    let type ← instantiateMVars x.type
    if ¬ x.isAuxDecl ∧ type.isAppOf ``And then
      let g ← getMainGoal
      let gs ← cases g x.fvarId
      replaceMainGoal <|
        gs.toList.map <|
          InductionSubgoal.mvarId ∘ CasesSubgoal.toInductionSubgoal
      changed := true
  guard changed

section HOrElse
variable  [Alternative m] [Monad m] [MonadBacktrack σ m]
instance : HOrElse (m α) (SearchT δ m α) (SearchT δ m α) where
  hOrElse x y := (Alternative.orElse (liftM x) y)

instance : HOrElse (SearchT δ m α) (m α) (SearchT δ m α) where
  hOrElse x y := (Alternative.orElse x (liftM ∘ y))

end HOrElse

def withMainContext' (x : SearchTacticM δ α) : SearchTacticM δ α :=
λ f => withMainContext (x f)

def Meta.tacAutoStep (ns : Array Name) (allowMVars := false) : SearchTacticM δ Unit :=
withMainContext' $
  Lean.Elab.Tactic.tacRefl <|>
  liftMetaTactic (do Meta.contradiction .; return []) <|>
  Meta.applyAssumption allowMVars <|>
  Meta.destructHyp <|>
  liftMetaTactic1 ((some ∘ Prod.snd) <$> intro1 .) <|>
  Meta.applyAuto (allowMVars := allowMVars) ns
  -- Meta.applyAuto ns allowMVars <|>
  -- liftMetaTactic1 ((some ∘ Prod.snd) <$> intro1 .)

def Meta.tacAuto (ns : Array Name) (bound : Option Nat)
  (allowMVars := false) : SearchTacticM δ Unit :=
let bound := bound.getD 5
SearchTacticM.focusAndDone $
iterate bound <| Meta.tacAutoStep ns allowMVars

def autoTac : TacticM Unit  := do
Meta.tacAuto (← getAutoList) none |>.run

def autoStepTac : TacticM Unit  := do
Meta.tacAutoStep (← getAutoList) |>.run

namespace Parser

elab "destruct_hyp" : tactic => withMainContext Meta.destructHyp

syntax "eauto" "[" ident,* "]" : tactic
syntax "auto" ("[" ident,* "]")? (" with " num)? : tactic

elab "auto" : tactic => do
  withMainContext (Meta.tacAuto (← getAutoList) none).run

elab "auto" " with " n:num : tactic => do
  withMainContext (Meta.tacAuto (← getAutoList) (Syntax.isNatLit? n)).run

elab "auto" "[" ids:ident,* "]": tactic => do
  let ids ← getAutoList (← ids.getElems.mapM resolveGlobalConstNoOverload)
  withMainContext (Meta.tacAuto ids none).run

elab "auto" "[" ids:ident,* "]" " with " n:num : tactic => do
  let ids ← getAutoList (← ids.getElems.mapM resolveGlobalConstNoOverload)
  withMainContext (Meta.tacAuto ids (Syntax.isNatLit? n)).run

elab "eauto" : tactic => do
  withMainContext (Meta.tacAuto (← getAutoList) none true).run
elab "auto_step" : tactic => do
  withMainContext (Meta.tacAutoStep (← getAutoList)).run
elab "eauto_step" : tactic => do
  withMainContext (Meta.tacAutoStep (allowMVars := true) (← getAutoList)).run
elab "apply_auto" : tactic => do
  withMainContext (Meta.applyAuto (← getAutoList)).run

elab "eauto" "[" ids:ident,* "]" : tactic => do
  let ids ← getAutoList (← ids.getElems.mapM resolveGlobalConstNoOverload)
  withMainContext (Meta.tacAuto ids none true).run

elab "eauto" "[" ids:ident,* "]" " with " n:num : tactic => do
  let ids ← getAutoList (← ids.getElems.mapM resolveGlobalConstNoOverload)
  withMainContext (Meta.tacAuto ids (Syntax.isNatLit? n) true).run

syntax "change" term "at" ident : tactic

elab "change" t:term "at" h:ident : tactic =>
  withMainContext do
    let h ← getFVarId h
    liftMetaTactic1 (changeLocalDecl . h (← elabTerm t none))

elab "apply_assumption" : tactic =>
  withMainContext (Meta.applyAssumption true).run

end Parser

-- macro "auto" : tactic =>
--   `(tactic|
--     solve
--     | repeat
--         first
--         | refl
--         | apply True.intro
--         | assumption
--         | destruct_hyp
--         | apply And.intro
--         | apply Iff.intro
--         | intros _
--         | apply_auto_lemma )

-- macro "auto_step" : tactic =>
--   `(tactic|
--         first
--         | refl
--         | apply True.intro
--         | assumption
--         | destruct_hyp
--         | apply And.intro
--         | apply Iff.intro
--         | intros _
--         | apply_auto_lemma )

-- theorem swapHyp {p q : Prop} (h : p) (h' : ¬ p) : q := by
--   cases h' h

-- macro "swapHyp" h:term "as" h':ident : tactic =>
--   `(tactic| apply Classical.byContradiction; intro $h' <;>
--             first
--             | apply $h ; clear $h
--             | apply swapHyp $h <;> clear $h
--               )

open Lean.Elab.Tactic

elab "all_but_first " tac:tacticSeq : tactic => do
  let mvarId :: mvarIds ← getUnsolvedGoals
    | throwNoGoalsToBeSolved
  let mut gs := #[mvarId]
  for g in mvarIds do
    if ← not <$> isExprMVarAssigned g then
      setGoals [g]
      tryTac <| withMainContext <| evalTactic tac
      gs := gs.appendList (← getGoals)
  setGoals gs.toList

macro:1 x:tactic " </> " y:tactic:0 : tactic =>
  `(tactic| focus ($x:tactic; all_but_first ($y:tactic; done)))

syntax "split" "*" : tactic
macro_rules
  | `(tactic| split*) => `(tactic| first | split <;> split* | skip)

macro "have " " ← " " : " p:term " := " proof:term : tactic =>
  `(tactic|
    have h : $p := $proof ;
    rw [← h] <;> clear h )

macro "have " " → " " : " p:term " := " proof:term : tactic =>
  `(tactic|
    have h : $p := $proof ;
    rw [h] <;> clear h )

namespace Classical

theorem contradiction {p q} (hp : p) (hnp : ¬ p) : q :=
by cases hnp hp

end Classical

syntax "apply' " term,* : tactic

macro_rules
| `(tactic| apply' $(xs):term,* ) => do
    if xs.elemsAndSeps.size == 0 then
      `(tactic| skip)
    else if xs.elemsAndSeps.size == 0 then
      `(tactic| apply $(xs.elemsAndSeps[0]):term)
    else
      let xs' := { xs with elemsAndSeps := xs.elemsAndSeps[2:] }
      `(tactic| apply $(xs.elemsAndSeps[0]):term; apply' $(xs'):term,*)

macro "falseHyp" h:ident : tactic =>
  `(first
    | refine' Classical.contradiction _ $h; clear $h
    | apply Classical.contradiction $h; clear $h )

elab "fold" foo:ident : tactic => do
  let mut eqns := #[]
  if let some xs ← Lean.Meta.getEqnsFor? foo.getId then
    eqns := xs
  else if let some x ← Lean.Meta.getUnfoldEqnFor? foo.getId then
    eqns := #[x]
  else
    throwError "{foo.getId} has no equational lemmas"
  liftMetaTactic1 λ mvar => do
    let tgt ← inferType (mkMVar mvar)
    let mut simpLmms : SimpTheorems := {}
    for x in eqns do
      simpLmms ← simpLmms.addConst (inv := true) x
    let r ← Lean.Meta.simp tgt { simpTheorems := #[simpLmms] }
    let newGoal ← mkFreshExprMVar (some r.expr)
    if let some pr := r.proof? then
      assignExprMVar mvar (← mkEqMP pr newGoal)
    else assignExprMVar mvar newGoal
    return some newGoal.mvarId!

macro "assume " h:ident " : " t:term : tactic =>
  `(tactic|
      refine' λ $h:ident : $t => ?_  )


import Do.Basic
import Do.LazyList

/-! # Local Mutation -/

open Lean

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

syntax "let" "mut" ident ":=" term ";" stmt : stmt
syntax ident ":=" term : stmt
syntax "if" term "then" stmt "else" stmt:1 : stmt

declare_syntax_cat expander
-- generic syntax for traversal-like functions S_y/R/B/L
syntax "expand!" expander "in" stmt:1 : stmt
syntax "mut" ident : expander  -- corresponds to `S_y`

-- generic traversal rules
macro_rules
  | `(stmt| expand! $exp in let $x ← $s; $s') => `(stmt| let $x ← expand! $exp in $s; expand! $exp in $s')                -- subsumes (R3, B4, L4)
  | `(stmt| expand! $exp in let mut $x := $e; $s') => `(stmt| let mut $x := $e; expand! $exp in $s')                      -- subsumes (R4, B5, L5)
  | `(stmt| expand! $_ in $x:ident := $e) => `(stmt| $x:ident := $e)                                                      -- subsumes (R5, B6, L6)
  | `(stmt| expand! $exp in if $e then $s₁ else $s₂) => `(stmt| if $e then expand! $exp in $s₁ else expand! $exp in $s₂)  -- subsumes (S6, R6, B7, L7)
  | `(stmt| expand! $exp in $s) => do
    let s' ← expandStmt s
    `(stmt| expand! $exp in $s')

macro_rules
  | `(d! let mut $x := $e; $s) => `(let $x := $e; StateT.run' (d! expand! mut $x in $s) $x) -- (D3)
  | `(d! $x:ident := $_:term) =>
      -- `s!"..."` is an interpolated string. For more information, see https://leanprover.github.io/lean4/doc/stringinterp.html.
      throw <| Macro.Exception.error x s!"variable '{x.getId}' is not reassignable in this scope"
  | `(d! if $e then $s₁ else $s₂) => `(if $e then d! $s₁ else d! $s₂)                       -- (D4)

macro_rules
  | `(stmt| expand! mut $_ in $e:term) => `(stmt| StateT.lift $e)  -- (S1)
  | `(stmt| expand! mut $y in let $x ← $s; $s') =>                -- (S2)
    if x == y then
      throw <| Macro.Exception.error x s!"cannot shadow 'mut' variable '{x.getId}'"
    else
      `(stmt| let $x ← expand! mut $y in $s; let $y ← get; expand! mut $y in $s')
  | `(stmt| expand! mut $y in let mut $x := $e; $s') =>           -- (S3)
    if x == y then
      throw <| Macro.Exception.error x s!"cannot shadow 'mut' variable '{x.getId}'"
    else
      `(stmt| let mut $x := $e; expand! mut $y in $s')
  | `(stmt| expand! mut $y in $x:ident := $e) =>
    if x == y then
      `(stmt| set $e)                                             -- (S5)
    else
      `(stmt| $x:ident := $e)                                     -- (S4)

macro:0 "let" "mut" x:ident "←" s:stmt:1 ";" s':stmt : stmt => `(let y ← $s; let mut $x := y; $s') -- (A3)
macro:0 x:ident "←" s:stmt:1 : stmt => `(let y ← $s; $x:ident := y)                                -- (A4)
-- a variant of (A4) since we technically cannot make the above macro a `stmt`
macro:0 x:ident "←" s:stmt:1 ";" s':stmt : stmt => `(let y ← $s; $x:ident := y; $s')
macro "if" e:term "then" s₁:stmt:1 : stmt => `(if $e then $s₁ else pure ())                        -- (A5)
macro "unless" e:term "do'" s₂:stmt:1 : stmt => `(if $e then pure () else $s₂)                     -- (A6)

/-
The `variable` command instructs Lean to insert the declared variables as bound variables
in definitions that refer to them.
-/

variable [Monad m]
variable (ma ma' : m α)

/-
  Mark `map_eq_pure_bind` as a simplification lemma.
  It is a theorem for
  `f <$> x = x >>= pure (f a)`
-/
attribute [local simp] map_eq_pure_bind

/-
  Remark: an `example` in Lean is like a "nameless" definition, and it does not update the environment.
  It is useful for writing tests.
-/

/-
  The instance `[LawfulMonad m]` contains the monadic laws.
  For more information, see https://github.com/leanprover/lean4/blob/v4.0.0-m4/src/Init/Control/Lawful.lean
-/

example [LawfulMonad m] :
    (do' let mut x ← ma;
         pure x : m α)
    =
    ma
:= by simp

example [LawfulMonad m] :
    (do' let mut x ← ma;
         x ← ma';
         pure x)
    =
    (ma >>= fun _ => ma')
:= by simp

/- The command `#check_failure <term>` succeeds only if `<term>` fails to be elaborated. -/

#check_failure do'
  let mut x ← ma;
  let x ← ma';  -- cannot shadow 'mut' variable 'x'
  pure x

#check_failure do'
  x ← ma;  -- variable 'x' is not reassignable in this scope
  pure ()

variable (b : Bool)

-- The following equivalence is true even if `m` does not satisfy the monadic laws.
example :
    (do' if b then {
           discard ma
         })
    =
    (if b then discard ma else pure ())
:= rfl

theorem simple [LawfulMonad m] :
    (do' let mut x ← ma;
         if b then {
           x ← ma'
         };
         pure x)
    =
    (ma >>= fun x => if b then ma' else pure x)
:= by cases b <;> simp

example [LawfulMonad m] (f : α → α → α) :
    (do' let mut x ← ma;
         let y ←
           if b then {
             x ← ma;
             ma'
           } else {
             ma'
           };
         pure (f x y))
    =
    (ma >>= fun x => if b then ma >>= fun x => ma' >>= fun y => pure (f x y) else ma' >>= fun y => pure (f x y))
:= by cases b <;> simp

/-
Nondeterminism example from Section 2.
-/
def choose := @List.toLazy

def ex : LazyList Nat := do'
  let mut x := 0;
  let y ← choose [0, 1, 2, 3];
  x := x + 1;
  guard (x < 3);
  pure (x + y)

-- Generate all solutions
#eval ex.toList


import operations
import params
import littleendian

import category_theory.category.basic
import category_theory.core

open operations
open params
open littleendian

open category_theory

namespace utils

variables [category (bitvec word_len)]

/-!
  # Utilities
-/

-- A random input matrix to be used as `rowround` and `columnround` inputs.
variable M : matrixType

/--
  Prepare the `rowround` matrix input.

  Any `rowround` input matrix is in the form:

  (y₀, y₁, y₂, y₃)
  (y₄, y₅, y₆, y₇)
  (y₈, y₉, y₁₀, y₁₁)
  (y₁₂, y₁₃, y₁₄, y₁₅)

  But we need this to be converted to:

  (y₀, y₁, y₂, y₃)
  (y₅, y₆, y₇, y₄)
  (y₁₀, y₁₁, y₈, y₉)
  (y₁₅, y₁₂, y₁₃, y₁₄)
-/
@[simp] def rowround_input : matrixType :=
  (
    ((M.fst).fst,                 (M.fst).snd.fst,              (M.fst).snd.snd.fst,      (M.fst).snd.snd.snd),
    ((M.snd.fst).snd.fst,         (M.snd.fst).snd.snd.fst,      (M.snd.fst).snd.snd.snd,  (M.snd.fst).fst),
    ((M.snd.snd.fst).snd.snd.fst, (M.snd.snd.fst).snd.snd.snd,  (M.snd.snd.fst).fst,      (M.snd.snd.fst).snd.fst),
    ((M.snd.snd.snd).snd.snd.snd, (M.snd.snd.snd).fst,          (M.snd.snd.snd).snd.fst,  (M.snd.snd.snd).snd.snd.fst)
  )

/--
  Prepare the `rowround` matrix output.

  Any `rowround` output matrix is of the form:

  (z₀, z₁, z₂, z₃)
  (z₅, z₆, z₇, z₄)
  (z₁₀, z₁₁, z₈, z₉)
  (z₁₅, z₁₂, z₁₃, z₁₄)

  But we need this to be converted to:

  (z₀, z₁, z₂, z₃)
  (z₄, z₅, z₆, z₇)
  (z₈, z₉, z₁₀, z₁₁)
  (z₁₂, z₁₃, z₁₄, z₁₅)
-/
@[simp] def rowround_output : matrixType :=
  (
    ((M.fst).fst,                 (M.fst).snd.fst,              (M.fst).snd.snd.fst,          (M.fst).snd.snd.snd),
    ((M.snd.fst).snd.snd.snd,     (M.snd.fst).fst,              (M.snd.fst).snd.fst,          (M.snd.fst).snd.snd.fst),
    ((M.snd.snd.fst).snd.snd.fst, (M.snd.snd.fst).snd.snd.snd,  (M.snd.snd.fst).fst,          (M.snd.snd.fst).snd.fst),
    ((M.snd.snd.snd).snd.fst,     (M.snd.snd.snd).snd.snd.fst,  (M.snd.snd.snd).snd.snd.snd,  (M.snd.snd.snd).fst)
  )

/-- The `rowround_output` function is the inverse of the `rowround_input` function. -/
lemma rowround_output_is_inv_of_input : rowround_output (rowround_input M) = M :=
begin
  unfold rowround_input,
  unfold rowround_output,
  simp only [prod.mk.eta],
end

/--
  Prepare the `columnround` matrix input.

  Any `columnround` input matrix is in the form:

  (x₀, x₁, x₂, x₃)
  (x₄, x₅, x₆, x₇)
  (x₈, x₉, x₁₀, x₁₁)
  (x₁₂, x₁₃, x₁₄, x₁₅)

  But we need this to be converted to:

  (x₀, x₄, x₈, x₁₂)
  (x₅, x₉, x₁₃, x₁)
  (x₁₀, x₁₄, x₂, x₆)
  (x₁₅, x₃, x₇, x₁₁)
-/
@[simp] def columnround_input : matrixType :=
  (
    ((M.fst).fst,                 (M.snd.fst).fst,              (M.snd.snd.fst).fst,      (M.snd.snd.snd).fst),
    ((M.snd.fst).snd.fst,         (M.snd.snd.fst).snd.fst,      (M.snd.snd.snd).snd.fst,  (M.fst).snd.fst),
    ((M.snd.snd.fst).snd.snd.fst, (M.snd.snd.snd).snd.snd.fst,  (M.fst).snd.snd.fst,      (M.snd.fst).snd.snd.fst),
    ((M.snd.snd.snd).snd.snd.snd, (M.fst).snd.snd.snd,          (M.snd.fst).snd.snd.snd,  (M.snd.snd.fst).snd.snd.snd)
  )

/--
  Prepare the `columnround` matrix output.

  Any `columnround` output matrix is in the form:

  (y₀, y₄, y₈, y₁₂)
  (y₅, y₉, y₁₃, y₁)
  (y₁₀, y₁₄, y₂, y₆)
  (y₁₅, y₃, y₇, y₁₁)

  But we need this to be converted to:

  (y₀, y₁, y₂, y₃)
  (y₄, y₅, y₆, y₇)
  (y₈, y₉, y₁₀, y₁₁)
  (y₁₂, y₁₃, y₁₄, y₁₅)
-/
@[simp] def columnround_output : matrixType :=
  (
    ((M.fst).fst,         (M.snd.fst).snd.snd.snd,  (M.snd.snd.fst).snd.snd.fst,  (M.snd.snd.snd).snd.fst),
    ((M.fst).snd.fst,     (M.snd.fst).fst,          (M.snd.snd.fst).snd.snd.snd,  (M.snd.snd.snd).snd.snd.fst),
    ((M.fst).snd.snd.fst, (M.snd.fst).snd.fst,      (M.snd.snd.fst).fst,          (M.snd.snd.snd).snd.snd.snd),
    ((M.fst).snd.snd.snd, (M.snd.fst).snd.snd.fst,  (M.snd.snd.fst).snd.fst,      (M.snd.snd.snd).fst)
  )

/-- The `columnround_output` function is the inverse of the `columnround_input` function. -/
lemma columnround_output_is_inv_of_input : columnround_output (columnround_input M) = M :=
begin
  unfold columnround_input,
  unfold columnround_output,
  simp only [prod.mk.eta],
end

-- A random input 64 bytes matrix that we can reduce using `littleendian` function.
variable X : matrix64Type

-- A random input 16 bytes matrix that we can aument using the `littleendian_inv` function.
variable Y : matrixType

/-- Reduce the 64 bytes sequence to a 16 bytes one by using little endian. -/
def reduce : matrixType :=
  (
    (
      littleendian (((X.fst).fst).fst,          ((X.fst).fst).snd.fst,          ((X.fst).fst).snd.snd.fst,          ((X.fst).fst).snd.snd.snd),
      littleendian (((X.fst).snd.fst).fst,      ((X.fst).snd.fst).snd.fst,      ((X.fst).snd.fst).snd.snd.fst,      ((X.fst).snd.fst).snd.snd.snd),
      littleendian (((X.fst).snd.snd.fst).fst,  ((X.fst).snd.snd.fst).snd.fst,  ((X.fst).snd.snd.fst).snd.snd.fst,  ((X.fst).snd.snd.fst).snd.snd.snd),
      littleendian (((X.fst).snd.snd.snd).fst,  ((X.fst).snd.snd.snd).snd.fst,  ((X.fst).snd.snd.snd).snd.snd.fst,  ((X.fst).snd.snd.snd).snd.snd.snd)
    ),
    (
      littleendian (((X.snd.fst).fst).fst,          ((X.snd.fst).fst).snd.fst,          ((X.snd.fst).fst).snd.snd.fst,          ((X.snd.fst).fst).snd.snd.snd),
      littleendian (((X.snd.fst).snd.fst).fst,      ((X.snd.fst).snd.fst).snd.fst,      ((X.snd.fst).snd.fst).snd.snd.fst,      ((X.snd.fst).snd.fst).snd.snd.snd),
      littleendian (((X.snd.fst).snd.snd.fst).fst,  ((X.snd.fst).snd.snd.fst).snd.fst,  ((X.snd.fst).snd.snd.fst).snd.snd.fst,  ((X.snd.fst).snd.snd.fst).snd.snd.snd),
      littleendian (((X.snd.fst).snd.snd.snd).fst,  ((X.snd.fst).snd.snd.snd).snd.fst,  ((X.snd.fst).snd.snd.snd).snd.snd.fst,  ((X.snd.fst).snd.snd.snd).snd.snd.snd)
    ),
    (
      littleendian (((X.snd.snd.fst).fst).fst,          ((X.snd.snd.fst).fst).snd.fst,          ((X.snd.snd.fst).fst).snd.snd.fst,          ((X.snd.snd.fst).fst).snd.snd.snd),
      littleendian (((X.snd.snd.fst).snd.fst).fst,      ((X.snd.snd.fst).snd.fst).snd.fst,      ((X.snd.snd.fst).snd.fst).snd.snd.fst,      ((X.snd.snd.fst).snd.fst).snd.snd.snd),
      littleendian (((X.snd.snd.fst).snd.snd.fst).fst,  ((X.snd.snd.fst).snd.snd.fst).snd.fst,  ((X.snd.snd.fst).snd.snd.fst).snd.snd.fst,  ((X.snd.snd.fst).snd.snd.fst).snd.snd.snd),
      littleendian (((X.snd.snd.fst).snd.snd.snd).fst,  ((X.snd.snd.fst).snd.snd.snd).snd.fst,  ((X.snd.snd.fst).snd.snd.snd).snd.snd.fst,  ((X.snd.snd.fst).snd.snd.snd).snd.snd.snd)
    ),
    (
      littleendian (((X.snd.snd.snd).fst).fst,          ((X.snd.snd.snd).fst).snd.fst,          ((X.snd.snd.snd).fst).snd.snd.fst,          ((X.snd.snd.snd).fst).snd.snd.snd),
      littleendian (((X.snd.snd.snd).snd.fst).fst,      ((X.snd.snd.snd).snd.fst).snd.fst,      ((X.snd.snd.snd).snd.fst).snd.snd.fst,      ((X.snd.snd.snd).snd.fst).snd.snd.snd),
      littleendian (((X.snd.snd.snd).snd.snd.fst).fst,  ((X.snd.snd.snd).snd.snd.fst).snd.fst,  ((X.snd.snd.snd).snd.snd.fst).snd.snd.fst,  ((X.snd.snd.snd).snd.snd.fst).snd.snd.snd),
      littleendian (((X.snd.snd.snd).snd.snd.snd).fst,  ((X.snd.snd.snd).snd.snd.snd).snd.fst,  ((X.snd.snd.snd).snd.snd.snd).snd.snd.fst,  ((X.snd.snd.snd).snd.snd.snd).snd.snd.snd)
    )
  )

/-- Aument a given 16 bytes sequence to a 64 bytes one using `littleenedian_inv`. -/
def aument : matrix64Type := (
  (
    littleendian_inv Y.fst.fst,
    littleendian_inv Y.fst.snd.fst,
    littleendian_inv Y.fst.snd.snd.fst,
    littleendian_inv Y.fst.snd.snd.snd
  ),
  (
    littleendian_inv Y.snd.fst.fst,
    littleendian_inv Y.snd.fst.snd.fst,
    littleendian_inv Y.snd.fst.snd.snd.fst,
    littleendian_inv Y.snd.fst.snd.snd.snd
  ),
  (
    littleendian_inv Y.snd.snd.fst.fst,
    littleendian_inv Y.snd.snd.fst.snd.fst,
    littleendian_inv Y.snd.snd.fst.snd.snd.fst,
    littleendian_inv Y.snd.snd.fst.snd.snd.snd
  ),
  (
    littleendian_inv Y.snd.snd.snd.fst,
    littleendian_inv Y.snd.snd.snd.snd.fst,
    littleendian_inv Y.snd.snd.snd.snd.snd.fst,
    littleendian_inv Y.snd.snd.snd.snd.snd.snd
  )
)

/--
Modular 2^32 addition of 4x4 matrices by doing Aᵢⱼ + Bᵢⱼ

The `MOD` operation (modulo 2^32 addition) is the key to make the salsa20 hash function irreversible.
Everything is reversible except for this addition.
-/
@[simp] def mod_matrix (A B : matrixType) : matrixType := (
  (
    A.fst.fst          MOD B.fst.fst,
    A.fst.snd.fst      MOD B.fst.snd.fst,
    A.fst.snd.snd.fst  MOD B.fst.snd.snd.fst,
    A.fst.snd.snd.snd  MOD B.fst.snd.snd.snd
  ),
  (
    A.snd.fst.fst          MOD B.snd.fst.fst,
    A.snd.fst.snd.fst      MOD B.snd.fst.snd.fst,
    A.snd.fst.snd.snd.fst  MOD B.snd.fst.snd.snd.fst,
    A.snd.fst.snd.snd.snd  MOD B.snd.fst.snd.snd.snd
  ),
  (
    A.snd.snd.fst.fst          MOD B.snd.snd.fst.fst,
    A.snd.snd.fst.snd.fst      MOD B.snd.snd.fst.snd.fst,
    A.snd.snd.fst.snd.snd.fst  MOD B.snd.snd.fst.snd.snd.fst,
    A.snd.snd.fst.snd.snd.snd  MOD B.snd.snd.fst.snd.snd.snd
  ),
  (
    A.snd.snd.snd.fst          MOD B.snd.snd.snd.fst,
    A.snd.snd.snd.snd.fst      MOD B.snd.snd.snd.snd.fst,
    A.snd.snd.snd.snd.snd.fst  MOD B.snd.snd.snd.snd.snd.fst,
    A.snd.snd.snd.snd.snd.snd  MOD B.snd.snd.snd.snd.snd.snd
  )
)

/-
/-- The inverse of a `mod_matrix` operation is not a function. -/
@[simp] lemma inv_of_mod_matrix_is_not_a_function : ∃ (A B C D : matrixType), mod_matrix A B = mod_matrix C D :=
begin
  simp only [mod_matrix, prod.mk.inj_iff, prod.exists, exists_and_distrib_left, exists_and_distrib_right,
  inv_of_mod_is_not_a_function, and_true],
end
-/

/-- We define the xor of a matrix to be the xor of each individual bitvector of matrix A and matrix B. -/
def xor_matrix (A B : matrixType) : matrixType := (
  (
    A.fst.fst          XOR B.fst.fst,
    A.fst.snd.fst      XOR B.fst.snd.fst,
    A.fst.snd.snd.fst  XOR B.fst.snd.snd.fst,
    A.fst.snd.snd.snd  XOR B.fst.snd.snd.snd
  ),
  (
    A.snd.fst.fst          XOR B.snd.fst.fst,
    A.snd.fst.snd.fst      XOR B.snd.fst.snd.fst,
    A.snd.fst.snd.snd.fst  XOR B.snd.fst.snd.snd.fst,
    A.snd.fst.snd.snd.snd  XOR B.snd.fst.snd.snd.snd
  ),
  (
    A.snd.snd.fst.fst          XOR B.snd.snd.fst.fst,
    A.snd.snd.fst.snd.fst      XOR B.snd.snd.fst.snd.fst,
    A.snd.snd.fst.snd.snd.fst  XOR B.snd.snd.fst.snd.snd.fst,
    A.snd.snd.fst.snd.snd.snd  XOR B.snd.snd.fst.snd.snd.snd
  ),
  (
    A.snd.snd.snd.fst          XOR B.snd.snd.snd.fst,
    A.snd.snd.snd.snd.fst      XOR B.snd.snd.snd.snd.fst,
    A.snd.snd.snd.snd.snd.fst  XOR B.snd.snd.snd.snd.snd.fst,
    A.snd.snd.snd.snd.snd.snd  XOR B.snd.snd.snd.snd.snd.snd
  )
)

-- Have 16 random numbers.
variables a₀ a₁ a₂ a₃ a₄ a₅ a₆ a₇ a₈ a₉ a₁₀ a₁₁ a₁₂ a₁₃ a₁₄ a₁₅ : bitvec word_len

/-- Distribute 2 * Matrix. -/
@[simp] lemma matrix_distribute_two :
  2 * ((a₀, a₁, a₂, a₃), (a₄, a₅, a₆, a₇), (a₈, a₉, a₁₀, a₁₁), (a₁₂, a₁₃, a₁₄, a₁₅)) =
  (
    (2 * a₀, 2 * a₁, 2 * a₂, 2 * a₃),
    (2 * a₄, 2 * a₅, 2 * a₆, 2 * a₇),
    (2 * a₈, 2 * a₉, 2 * a₁₀, 2 * a₁₁),
    (2 * a₁₂, 2 * a₁₃, 2 * a₁₄, 2 * a₁₅)
  ) := rfl

/-
/-- The MOD sum of two equal matrices X is 2 times X. -/
@[simp] lemma mod_matrix_double : mod_matrix M M = 2 * M :=
begin
  unfold mod_matrix,
  simp only [mod_self],

  rw ← matrix_distribute_two
    M.fst.fst         M.fst.snd.fst         M.fst.snd.snd.fst         M.fst.snd.snd.snd
    M.snd.fst.fst     M.snd.fst.snd.fst     M.snd.fst.snd.snd.fst     M.snd.fst.snd.snd.snd
    M.snd.snd.fst.fst M.snd.snd.fst.snd.fst M.snd.snd.fst.snd.snd.fst M.snd.snd.fst.snd.snd.snd
    M.snd.snd.snd.fst M.snd.snd.snd.snd.fst M.snd.snd.snd.snd.snd.fst M.snd.snd.snd.snd.snd.snd,
  refl,
end
-/

/-- Convert a `matrixType` to a `list` as lists are easy to work sometimes than prods. -/
@[simp] def matrix_to_list : matrixType → list (bitvec word_len)
| m := [
  m.fst.fst, m.fst.snd.fst, m.fst.snd.snd.fst, m.fst.snd.snd.snd,
  m.snd.fst.fst, m.snd.fst.snd.fst, m.snd.fst.snd.snd.fst, m.snd.fst.snd.snd.snd,
  m.snd.snd.fst.fst, m.snd.snd.fst.snd.fst, m.snd.snd.fst.snd.snd.fst, m.snd.snd.fst.snd.snd.snd,
  m.snd.snd.snd.fst, m.snd.snd.snd.snd.fst, m.snd.snd.snd.snd.snd.fst, m.snd.snd.snd.snd.snd.snd
]

/-- Convert a list of bitvectors into a `matrixType`. Will panic if list size is < 16 -/
@[simp] def list_to_matrix : list (bitvec word_len) → matrixType
| l := (
  ((l.nth 0).iget, (l.nth 1).iget, (l.nth 2).iget, (l.nth 3).iget),
  ((l.nth 4).iget, (l.nth 5).iget, (l.nth 6).iget, (l.nth 7).iget),
  ((l.nth 8).iget, (l.nth 9).iget, (l.nth 10).iget, (l.nth 11).iget),
  ((l.nth 12).iget, (l.nth 13).iget, (l.nth 14).iget, (l.nth 15).iget)
)

end utils


/- Introduces evaluation contexts. -/

import sesh.term
import sesh.ren_sub

open term
open matrix

universes u v w

/- Oh no, additional axioms. -/
axiom heq.congr {α γ : Sort u} {β δ : Sort v} {f₁ : α → β} {f₂ : γ → δ} {a₁ : α} {a₂ : γ} (h₁ : f₁ == f₂) (h₂ : a₁ == a₂) : f₁ a₁ == f₂ a₂

axiom heq.dcongr {α β : Sort u} {γ δ : Sort v} {ε ζ : Sort w} {f₁ : Π α, γ → ε} {f₂ : Π β, δ → ζ} {a₁₁ : α} {a₁₂ : γ} {a₂₁ : β} {a₂₂ : δ} (h₁ : f₁ == f₂) (h₂ : a₁₁ == a₂₁) (h₃ : a₁₂ == a₂₂) : f₁ a₁₁ a₁₂ == f₂ a₂₁ a₂₂

/- The context except the hole consumes Γₑ resources, while the
   hole can consume arbitrary resources Γ. The term plugged in
   for the hole must have type A' (hence the hole is "typed") and
   the resulting term has type A. In every evaluation context,
   the hole has the same precontext as the overall expression,
   since GV never evaluates under binders. Therefore I don't have
   to allow a fully general (i.e. arbitrary precontext) typing
   context for the hole. -/
@[reducible]
def eval_ctx_fn {γ} (Γₑ: context γ) (A' A: tp): Type :=
  Π (Γ: context γ), term Γ A' → term (Γ + Γₑ) A

namespace eval_ctx_fn

@[reducible]
def apply {γ} {Γₑ Γ': context γ} {A' A: tp}
  (f: eval_ctx_fn Γₑ A' A)
  (Γ: context γ)
  (M: term Γ' A')
  (h: auto_param (Γ = Γ'+Γₑ) ``solve_context)
  : term Γ A :=
cast (by solve_context) $ f Γ' M

end eval_ctx_fn

/- A value of type (eval_ctx f) specifies that f is a valid function
   for an evaluation context. It's a Type rather than a Prop, because
   Prop can only eliminate into Prop but sometimes I need to extract
   the actual value of a constructor argument out of an eval_ctx. -/
inductive eval_ctx
  : Π {γ} {A' A: tp} (Γₑ: context γ), eval_ctx_fn Γₑ A' A → Type
| EHole:
  Π (γ: precontext) (A: tp),
  eval_ctx 0 (λ (Γ: context γ) (M: term Γ A), begin convert M, solve_context end)

| EAppLeft:
  Π {γ} {Γ₁ Γ₂: context γ} {A B: tp}
    {C'}
  (Γₑ: context γ)
  (hΓ: Γₑ = Γ₁ + Γ₂)
  (N: term Γ₂ A)
  (E: eval_ctx_fn Γ₁ C' $ A⊸B),
  eval_ctx Γ₁ E
  -------------------------------
→ eval_ctx Γₑ (λ Γ M, App _ (E Γ M) N)

| EAppRight:
  Π {γ} {Γ₁ Γ₂: context γ} {A B: tp}
    {A'} {V: term Γ₁ $ A⊸B}
  (Γₑ: context γ)
  (hΓ: Γₑ = Γ₁ + Γ₂)
  (hV: value V)
  (E: eval_ctx_fn Γ₂ A' A),
  eval_ctx Γ₂ E
  -------------------------------
→ eval_ctx Γₑ (λ Γ M, App _ V $ E Γ M)

| ELetUnit:
  Π {γ} {Γ₁ Γ₂: context γ} {A: tp}
    {A'}
  (Γₑ: context γ)
  (hΓ: Γₑ = Γ₁ + Γ₂)
  (N: term Γ₂ A)
  (E: eval_ctx_fn Γ₁ A' tp.unit),
  eval_ctx Γ₁ E
  ---------------------------------
→ eval_ctx Γₑ (λ Γ M, LetUnit _ (E Γ M) N)

| EPairLeft:
  Π {γ} {Γ₁ Γ₂: context γ} {A B: tp}
    {A'}
  (Γₑ: context γ)
  (hΓ: Γₑ = Γ₁ + Γ₂)
  (N: term Γ₂ B)
  (E: eval_ctx_fn Γ₁ A' A),
  eval_ctx Γ₁ E
  ------------------------------
→ eval_ctx Γₑ (λ Γ M, Pair _ (E Γ M) N)

| EPairRight:
  Π {γ} {Γ₁ Γ₂: context γ} {A B: tp}
    {B'} {V: term Γ₁ A}
  (Γₑ: context γ)
  (hΓ: Γₑ = Γ₁ + Γ₂)
  (hV: value V)
  (E: eval_ctx_fn Γ₂ B' B),
  eval_ctx Γ₂ E
  ------------------------------
→ eval_ctx Γₑ (λ Γ M, Pair _ V $ E Γ M)

| ELetPair:
  Π {γ} {Γ₁ Γ₂: context γ} {A B C: tp}
    {C'}
  (Γₑ: context γ)
  (hΓ: Γₑ = Γ₁ + Γ₂)
  (N: term (⟦1⬝A⟧::⟦1⬝B⟧::Γ₂) C)
  (E: eval_ctx_fn Γ₁ C' $ tp.prod A B),
  eval_ctx Γ₁ E
  ----------------------------------
→ eval_ctx Γₑ (λ Γ M, LetPair _ (E Γ M) N)

| EInl:
  Π {γ} {Γₑ: context γ} {A: tp}
    {A'}
  (B: tp)
  (E: eval_ctx_fn Γₑ A' A),
  eval_ctx Γₑ E
  -----------------------------
→ eval_ctx Γₑ (λ Γ M, Inl B $ E Γ M)

| EInr:
  Π {γ} {Γₑ: context γ} {B: tp}
    {B'}
  (A: tp)
  (E: eval_ctx_fn Γₑ B' B),
  eval_ctx Γₑ E
  ---------------------------
→ eval_ctx Γₑ (λ Γ M, Inr A $ E Γ M)

| ECase:
  Π {γ} {Γ₁ Γ₂: context γ} {A B C: tp}
    {D'}
  (Γₑ: context γ)
  (hΓ: Γₑ = Γ₁ + Γ₂)
  (M: term (⟦1⬝A⟧::Γ₂) C)
  (N: term (⟦1⬝B⟧::Γ₂) C)
  (E: eval_ctx_fn Γ₁ D' $ tp.sum A B),
  eval_ctx Γ₁ E
  --------------------------------
→ eval_ctx Γₑ (λ Γ L, Case _ (E Γ L) M N)

| EFork:
  Π {γ} {Γₑ: context γ} {S: sesh_tp}
    {T'}
  (E: eval_ctx_fn Γₑ T' $ S⊸End!),
  eval_ctx Γₑ E
  --------------------------
→ eval_ctx Γₑ (λ Γ x, Fork (E Γ x))

| ESendLeft:
  Π {γ} {Γ₁ Γ₂: context γ} {A: tp} {S: sesh_tp}
    {A'}
  (Γₑ: context γ)
  (hΓ: Γₑ = Γ₁ + Γ₂)
  (N: term Γ₂ $ !A⬝S)
  (E: eval_ctx_fn Γ₁ A' A),
  eval_ctx Γ₁ E
  ------------------------------
→ eval_ctx Γₑ (λ Γ M, Send _ (E Γ M) N)

| ESendRight:
  Π {γ} {Γ₁ Γ₂: context γ} {A: tp} {S: sesh_tp}
    {T'} {V: term Γ₁ A}
  (Γₑ: context γ)
  (hΓ: Γₑ = Γ₁ + Γ₂)
  (hV: value V)
  (E: eval_ctx_fn Γ₂ T' $ !A⬝S),
  eval_ctx Γ₂ E
  ------------------------------
→ eval_ctx Γₑ (λ Γ M, Send _ V $ E Γ M)

| ERecv:
  Π {γ} {Γₑ: context γ} {A: tp} {S: sesh_tp}
    {T'}
  (E: eval_ctx_fn Γₑ T' ?A⬝S),
  eval_ctx Γₑ E
  --------------------------
→ eval_ctx Γₑ (λ Γ M, Recv $ E Γ M)

| EWait:
  Π {γ} {Γₑ: context γ}
    {T'}
  (E: eval_ctx_fn Γₑ T' End?),
  eval_ctx Γₑ E
  --------------------------
→ eval_ctx Γₑ (λ Γ M, Wait $ E Γ M)

namespace eval_ctx
open matrix.vmul

/- The new function we're defining takes a hole term defined over
   an extended environment (rename ρ Γ) and returns the same expression,
   but well-typed under the extended environment. -/
def ext:
    Π {γ δ: precontext}  {A' A: tp}{Γ: context γ}
      {E: eval_ctx_fn Γ A' A}
    (ρ: ren_fn γ δ),
    eval_ctx Γ E
    -----------------------------------------------
  → Σ E': eval_ctx_fn (Γ ⊛ (λ B x, identity δ B $ ρ B x)) A' A,
      eval_ctx (Γ ⊛ (λ B x, identity δ B $ ρ B x)) E'
/- In each case we define what happens when the resulting
   renamed evaluation context is _applied_ to a hole-filling
   argument, which is the last matched variable (usually M).

   Most cases proceed by renaming parts of the evaluation context
   to make sense in the extended typing context and proving that
   the contexts still make sense.

   The return value of this function also carries the proof
   that the returned function is, in fact, an evaluation context. -/
| _ _ _ _ _ _ _ (EHole _ _) :=
begin
  rw [matrix.vmul.zero_vmul],
  exact ⟨_, EHole _ _⟩
end
| _ _ _ _ _ _ ρ (EAppLeft _ hΓ N _ E) :=
  let E' := ext ρ E in
  ⟨_, EAppLeft
    _
    (begin rw [hΓ, vmul_right_distrib] end)
    (rename ρ _ N)
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (EAppRight _ hΓ hV _ E) :=
  let E' := ext ρ E in
  ⟨_, EAppRight
    _
    (begin rw [hΓ, vmul_right_distrib] end)
    (hV.rename ρ)
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (ELetUnit _ hΓ N _ E) :=
  let E' := ext ρ E in
  ⟨_, ELetUnit
    _
    (begin rw [hΓ, vmul_right_distrib] end)
    (rename ρ _ N)
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (EPairLeft _ hΓ N _ E) :=
  let E' := ext ρ E in
  ⟨_, EPairLeft
    _
    (begin rw [hΓ, vmul_right_distrib] end)
    (rename ρ _ N)
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (EPairRight _ hΓ hV _ E) :=
  let E' := ext ρ E in
  ⟨_, EPairRight
    _
    (begin rw [hΓ, vmul_right_distrib] end)
    (hV.rename ρ)
    E'.fst
    E'.snd⟩
| γ _ _ _ _ _ ρ (ELetPair _ hΓ N _ E) :=
  let E' := ext ρ E in
  ⟨_, ELetPair
    _
    (begin rw [hΓ, vmul_right_distrib] end)
    (rename ((ρ.ext _).ext _) _ N)
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (EInl C _ E) :=
  let E' := ext ρ E in
  ⟨_, EInl
    C
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (EInr C _ E) :=
  let E' := ext ρ E in
  ⟨_, EInr
    C
    E'.fst
    E'.snd⟩
| γ _ _ _ _ _ ρ (ECase _ hΓ M N _ E) :=
  let E' := ext ρ E in
  ⟨_, ECase
    _
    (begin rw [hΓ, vmul_right_distrib] end)
    (rename (ρ.ext _) _ M)
    (rename (ρ.ext _) _ N)
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (EFork _ E) :=
  let E' := ext ρ E in
  ⟨_, EFork
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (ESendLeft _ hΓ N _ E) :=
  let E' := ext ρ E in
  ⟨_, ESendLeft
    _
    (begin rw [hΓ, vmul_right_distrib] end)
    (rename ρ _ N)
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (ESendRight _ hΓ hV _ E) :=
  let E' := ext ρ E in
  ⟨_, ESendRight
    _
    (begin rw [hΓ, vmul_right_distrib] end)
    (hV.rename ρ)
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (ERecv _ E) :=
  let E' := ext ρ E in
  ⟨_, ERecv
    E'.fst
    E'.snd⟩
| _ _ _ _ _ _ ρ (EWait _ E) :=
  let E' := ext ρ E in
  ⟨_, EWait
    E'.fst
    E'.snd⟩

def wrap: Π {γ} {A'' A' A: tp} {Γₑ Γₑ': context γ}
  (E: eval_ctx_fn Γₑ A'' A')
  (hE: eval_ctx Γₑ E)
  (E': eval_ctx_fn Γₑ' A' A)
  (hE': eval_ctx Γₑ' E')
  (Γ: context γ)
  (hΓ: Γ = Γₑ+Γₑ'),
  Σ E': eval_ctx_fn Γ A'' A,
      eval_ctx Γ E'
| _ _ _ _ _ _ E hE _ (EHole _ _) Γ hΓ :=
  cast (begin congr; simp [*, hΓ], congr' 1, simp [hΓ] end)
    (sigma.mk E hE)
| _ _ _ _ _ _ E hE _ (EAppLeft _ _ N E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, EAppLeft Γ (by solve_context) N EE'.fst EE'.snd⟩
| _ _ _ _ _ _ E hE _ (EAppRight _ _ hV E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, EAppRight Γ (by solve_context) hV EE'.fst EE'.snd⟩
| _ _ _ _ _ _ E hE _ (ELetUnit _ _ N E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, ELetUnit Γ (by solve_context) N EE'.fst EE'.snd⟩
| _ _ _ _ _ _ E hE _ (EPairLeft _ _ N E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, EPairLeft Γ (by solve_context) N EE'.fst EE'.snd⟩
| _ _ _ _ _ _ E hE _ (EPairRight _ _ hV E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, EPairRight Γ (by solve_context) hV EE'.fst EE'.snd⟩
| _ _ _ _ _ _ E hE _ (ELetPair _ _ N E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, ELetPair Γ (by solve_context) N EE'.fst EE'.snd⟩
| _ _ _ _ _ _ E hE _ (EInl B E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, EInl B (cast (by rw [hΓ]) EE'.fst) $ cast (begin
    congr' 1, exact hΓ.symm,
    h_generalize Hx: EE'.fst == x, exact Hx,
  end) EE'.snd⟩
| _ _ _ _ _ _ E hE _ (EInr A E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, EInr A (cast (by rw [hΓ]) EE'.fst) $ cast (begin
    congr' 1, exact hΓ.symm,
    h_generalize Hx: EE'.fst == x, exact Hx,
  end) EE'.snd⟩
| _ _ _ _ _ _ E hE _ (ECase _ _ M N E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, ECase Γ (by solve_context) M N EE'.fst EE'.snd⟩
| _ _ _ _ _ _ E hE _ (EFork E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, EFork (cast (by rw [hΓ]) EE'.fst) $ cast (begin
    congr' 1, exact hΓ.symm,
    h_generalize Hx: EE'.fst == x, exact Hx,
  end) EE'.snd⟩
| _ _ _ _ _ _ E hE _ (ESendLeft _ _ M E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, ESendLeft Γ (by solve_context) M EE'.fst EE'.snd⟩
| _ _ _ _ _ _ E hE _ (ESendRight _ _ hV E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, ESendRight Γ (by solve_context) hV EE'.fst EE'.snd⟩
| _ _ _ _ _ _ E hE _ (ERecv E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, ERecv (cast (by rw [hΓ]) EE'.fst) $ cast (begin
    congr' 1, exact hΓ.symm,
    h_generalize Hx: EE'.fst == x, exact Hx,
  end) EE'.snd⟩
| _ _ _ _ _ _ E hE _ (EWait E' hE') Γ hΓ :=
  let EE' := wrap E hE E' hE' _ rfl in
  ⟨_, EWait (cast (by rw [hΓ]) EE'.fst) $ cast (begin
    congr' 1, exact hΓ.symm,
    h_generalize Hx: EE'.fst == x, exact Hx,
  end) EE'.snd⟩

set_option pp.implicit true

lemma wrap_composes {γ} {Γ Γₑ Γₑ': context γ} {A'' A' A: tp}
    {M: term Γ A''}
    {E': eval_ctx_fn Γₑ A'' A'}
    {hE': eval_ctx Γₑ E'}
    {E: eval_ctx_fn Γₑ' A' A}
    {hE: eval_ctx Γₑ' E}
  (Γ': context γ)
  (hΓ': Γ' = Γ+Γₑ)
  (EM: term Γ' A')
  (Γ'': context γ)
  (hΓ'': Γ'' = Γ'+Γₑ')
  (EM': term Γ'' A)
  (hEM: EM = E'.apply Γ' M)
  (hEM': EM' = E.apply Γ'' EM)
  : EM' = (wrap E' hE' E hE _ rfl).fst.apply Γ'' M :=
begin
  induction hE; simp [*, wrap, eval_ctx_fn.apply],
  case EHole {
    h_generalize Hx: (E' Γ M) == x,
    h_generalize Hy: x == y,
    h_generalize Hx': (⟨E', hE'⟩: Σ E': eval_ctx_fn Γₑ A'' hE_A, eval_ctx Γₑ E') == x',
    congr' 1, simp [hΓ'],
    apply heq.trans Hy.symm, apply heq.trans Hx.symm,
    apply heq.congr, unfold sigma.fst,
    sorry
  },
  sorry
end

end eval_ctx

/- An evaluation context is a hole-replacing function
   together with a proof of its validity. -/
structure eval_ctx' {γ} (Γₑ: context γ) (A' A: tp) :=
(f: eval_ctx_fn Γₑ A' A)
(h: eval_ctx Γₑ f)

namespace eval_ctx'

def ext {γ δ: precontext} {Γ: context γ} {A' A: tp} (ρ: ren_fn γ δ) (E: eval_ctx' Γ A' A)
  : eval_ctx' (Γ ⊛ (λ B x, identity δ B $ ρ B x)) A' A :=
  let E' := eval_ctx.ext ρ E.h in
  ⟨E'.fst, E'.snd⟩

def wrap {γ} {Γₑ Γₑ': context γ} {A'' A' A: tp}
  (E: eval_ctx' Γₑ A'' A')
  (E': eval_ctx' Γₑ' A' A)
  (Γ: context γ)
  (hΓ: Γ = Γₑ+Γₑ')
  : eval_ctx' Γ A'' A :=
let EE' := eval_ctx.wrap E.f E.h E'.f E'.h Γ hΓ in
⟨EE'.fst, EE'.snd⟩

lemma wrap_composes {γ} {Γ Γₑ Γₑ': context γ} {A'' A' A: tp}
    {M: term Γ A''}
    {E': eval_ctx' Γₑ A'' A'}
    {E: eval_ctx' Γₑ' A' A}
  (Γ': context γ)
  (hΓ': Γ' = Γ+Γₑ)
  (EM: term Γ' A')
  (Γ'': context γ)
  (hΓ'': Γ'' = Γ'+Γₑ')
  (EM': term Γ'' A)
  (hE: EM = E'.f.apply Γ' M)
  (hE': EM' = E.f.apply Γ'' EM)
  : EM' = (E'.wrap E _ rfl).f.apply Γ'' M :=
sorry

end eval_ctx'

inductive term_reduces
  : ∀ {γ} {Γ: context γ} {A: tp}, term Γ A → term Γ A → Prop
infix ` ⟶M `:55 := term_reduces
| EvalLift:
  ∀ {γ} {Γ Γₑ: context γ} {A A': tp}
    {M M': term Γ A}
  (Γ': context γ)
  (E: eval_ctx' Γₑ A A')
  (EM EM': term Γ' A')
  (hΓ': Γ' = Γ+Γₑ)
  (hStep: M ⟶M M')
  (hEM: EM = E.f.apply Γ' M)
  (hEM': EM' = E.f.apply Γ' M'),
  -----------------------
  EM ⟶M EM'

| EvalLam:
  ∀ {γ} {Γ₁ Γ₂: context γ} {A B: tp}
    {V: term Γ₂ A}
  (Γ: context γ)
  (M: term (⟦1⬝A⟧::Γ₁) B)
  (hV: value V)
  (_: auto_param (Γ = Γ₁ + Γ₂) ``solve_context),
  ----------------------------------------------
  (App Γ (Abs M) V) ⟶M ssubst _ M V

| EvalUnit:
  ∀ {γ} {Γ: context γ} {A: tp}
  (M: term Γ A),
  ---------------------------------------
  (LetUnit Γ (Unit 0) M $ by simp) ⟶M M

| EvalPair:
  ∀ {γ} {Γ₁₁ Γ₁₂ Γ₂: context γ} {A B C: tp}
    {V: term Γ₁₁ A} {W: term Γ₁₂ B}
  (Γ₁ Γ: context γ)
  (hV: value V)
  (hW: value W)
  (M: term (⟦1⬝A⟧::⟦1⬝B⟧::Γ₂) C)
  (_: auto_param (Γ = Γ₁ + Γ₂) ``solve_context)
  (_: auto_param (Γ₁ = Γ₁₁ + Γ₁₂) ``solve_context),
  -----------------------------------------------------------
  (LetPair Γ (Pair Γ₁ V W $ by assumption) M $ by assumption)
  ⟶M
  dsubst _ M V W

| EvalInl:
  ∀ {γ} {Γ₁ Γ₂: context γ} {A B C: tp}
    {V: term Γ₁ A}
  (Γ: context γ)
  (hV: value V)
  (M: term (⟦1⬝A⟧::Γ₂) C)
  (N: term (⟦1⬝B⟧::Γ₂) C)
  (_: auto_param (Γ = Γ₁ + Γ₂) ``solve_context),
  ----------------------------------------------
  (Case Γ (Inl B V) M N) ⟶M ssubst _ M V

| EvalInr:
  ∀ {γ} {Γ₁ Γ₂: context γ} {A B C: tp}
    {V: term Γ₁ B}
  (Γ: context γ)
  (hV: value V)
  (M: term (⟦1⬝A⟧::Γ₂) C)
  (N: term (⟦1⬝B⟧::Γ₂) C)
  (_: auto_param (Γ = Γ₁ + Γ₂) ``solve_context),
  ----------------------------------------------
  (Case Γ (Inr A V) M N) ⟶M ssubst _ N V

infix ` ⟶M `:55 := term_reduces

/- This is what I thought initially (WRONG!):
   I _think_ there is barely any benefit to formalizing evaluation contexts
   because they would have to be defined for all kinds of terms anyway.

   Actually no, eval_ctx is _necessary_ to formalize configuration reduction
   without losing what's left of my sanity. -/


import tactic
import tactic.monotonicity
import tactic.norm_num
import category.basic
import category.serial
import data.serial.medium

universes u v w

def serial_inverse {α : Type u} (encode : α → put_m) (decode : get_m α) : Prop :=
∀ w, decode -<< encode w = pure w

class serial (α : Type u) :=
  (encode : α → put_m.{u})
  (decode : get_m α)
  (correctness : ∀ w, decode -<< encode w = pure w)

class serial1 (f : Type u → Type v) :=
  (encode : Π {α} [serial α], f α → put_m.{v})
  (decode : Π {α} [serial α], get_m (f α))
  (correctness : ∀ {α} [serial α] (w : f α), decode -<< encode w = pure w)

instance serial.serial1 {f α} [serial1 f] [serial α] : serial (f α) :=
{ encode := λ x, serial1.encode x,
  decode := serial1.decode f,
  correctness := serial1.correctness }

class serial2 (f : Type u → Type v → Type w) :=
  (encode : Π {α β} [serial α] [serial β], f α β → put_m.{w})
  (decode : Π {α β} [serial α] [serial β], get_m (f α β))
  (correctness : ∀ {α β} [serial α] [serial β] (w : f α β), decode -<< encode w = pure w)

instance serial.serial2 {f α β} [serial2 f] [serial α] [serial β] : serial (f α β) :=
{ encode := λ x, serial2.encode x,
  decode := serial2.decode f,
  correctness := serial2.correctness }

instance serial1.serial2 {f α} [serial2 f] [serial α] : serial1 (f α) :=
{ encode := λ β inst x, @serial2.encode _ _ α β _ inst x,
  decode := λ β inst, @serial2.decode f _ α β _ inst,
  correctness := λ β inst, @serial2.correctness _ _ α β _ inst }

export serial (encode decode)

namespace serial

open function

variables {α β σ γ : Type u} {ω : Type}

def serialize [serial α] (x : α) : list unsigned := (encode x).eval
def deserialize (α : Type u) [serial α] (bytes : list unsigned) : option α := (decode α).eval bytes

lemma encode_decode_bind [serial α]
  (f : α → get_m β) (f' : punit → put_m) (w : α) :
  (decode α >>= f) -<< (encode w >>= f') = f w -<< f' punit.star :=
by { rw [read_write_mono]; rw serial.correctness; refl }

lemma encode_decode_bind' [serial α]
  (f : α → get_m β) (w : α) :
  (decode α >>= f) -<< (encode w) = f w -<< pure punit.star :=
by { rw [read_write_mono_left]; rw serial.correctness; refl }

lemma encode_decode_pure
  (w w' : α) (u : punit) :
  (pure w) -<< (pure u) = pure w' ↔ w = w' :=
by split; intro h; cases h; refl

open ulift

protected def ulift.encode [serial α] (w : ulift.{v} α) : put_m :=
liftable1.up equiv.punit_equiv_punit (encode (down w))

protected def ulift.decode [serial α] : get_m (ulift α) :=
get_m.up ulift.up (decode α)

instance [serial α] : serial (ulift.{v u} α) :=
{ encode := ulift.encode
, decode := ulift.decode
, correctness :=
  by { introv, simp [ulift.encode,ulift.decode],
       rw up_read_write' _ equiv.ulift.symm,
       rw [serial.correctness], cases w, refl,
       intro, refl } }

instance unsigned.serial : serial unsigned :=
{ encode := λ w, put_m'.write w put_m'.pure
, decode := get_m.read get_m.pure
, correctness := by introv; refl }

def write_word (w : unsigned) : put_m.{u} :=
encode (up.{u} w)

@[simp] lemma loop_read_write_word {α β γ : Type u}
  (w : unsigned) (x : α) (f : α → unsigned → get_m (β ⊕ α)) (g : β → get_m γ)
  (rest : punit → put_m) :
  get_m.loop f g x -<< (write_word w >>= rest) =
  (f x w >>= @sum.rec _ _ (λ _, get_m γ) g (get_m.loop f g)) -<< rest punit.star := rfl

@[simp] lemma loop_read_write_word' {α β γ : Type u}
  (w : unsigned) (x : α) (f : α → unsigned → get_m (β ⊕ α)) (g : β → get_m γ)  :
  get_m.loop f g x -<< (write_word w) =
  (f x w >>= @sum.rec _ _ (λ _, get_m γ) g (get_m.loop f g)) -<< pure punit.star := rfl

def read_word : get_m.{u} (ulift unsigned) :=
decode _

def select_tag' (tag : unsigned) : list (unsigned × get_m α) → get_m α
| [] := get_m.fail
| ((w,x) :: xs) := if w = tag then x else select_tag' xs

def select_tag (xs : list (unsigned × get_m α)) : get_m α :=
do w ← read_word,
   select_tag' (down w) xs

@[simp]
lemma read_write_tag_hit {w w' : unsigned} {x : get_m α}
  {xs : list (unsigned × get_m α)} {y : put_m}
  (h : w = w') :
  select_tag ( (w,x) :: xs ) -<< (write_word w' >> y) = x -<< y :=
by subst w'; simp [select_tag,(>>),read_word,write_word,encode_decode_bind,select_tag']

lemma read_write_tag_hit' {w w' : unsigned} {x : get_m α}
  {xs : list (unsigned × get_m α)}
  (h : w = w') :
  select_tag ( (w,x) :: xs ) -<< (write_word w') = x -<< pure punit.star :=
by subst w'; simp [select_tag,(>>),read_word,write_word,encode_decode_bind',select_tag']

@[simp]
lemma read_write_tag_miss {w w' : unsigned} {x : get_m α}
  {xs : list (unsigned × get_m α)} {y : put_m}
  (h : w ≠ w') :
  select_tag ( (w,x) :: xs ) -<< (write_word w' >> y) = select_tag xs -<< (write_word w' >> y) :=
by simp [select_tag,(>>),read_word,write_word,encode_decode_bind,select_tag',*]

def recursive_parser {α} : ℕ → (get_m α → get_m α) → get_m α
| 0 _ := get_m.fail
| (nat.succ n) rec_fn := rec_fn $ recursive_parser n rec_fn

lemma recursive_parser_unfold {α} (n : ℕ) (f : get_m α → get_m α) (h : 1 ≤ n) :
  recursive_parser n f = f (recursive_parser (n-1) f) :=
by cases n; [ cases h, refl ]

attribute [simp] serial.correctness

end serial

structure serializer (α : Type u) (β : Type u) :=
(encoder : α → put_m.{u})
(decoder : get_m β)

namespace serializer

def valid_serializer {α} (x : serializer α α) :=
serial_inverse
      (serializer.encoder x)
      (serializer.decoder x)

lemma serializer.eq {α β} (x y : serializer α β)
  (h : x.encoder = y.encoder)
  (h' : x.decoder = y.decoder) :
  x = y :=
by cases x; cases y; congr; assumption

namespace serializer.seq

variables {α : Type u} {i j : Type u}
variables (x : serializer α (i → j))
variables (y : serializer α i)

def encoder := λ (k : α), x.encoder k >> y.encoder k
def decoder := x.decoder <*> y.decoder

end serializer.seq

instance {α : Type u} : applicative (serializer.{u} α) :=
{ pure := λ i x, { encoder := λ _, return punit.star, decoder := pure x }
, seq := λ i j x y,
  { encoder := serializer.seq.encoder x y
  , decoder := serializer.seq.decoder x y } }

section lawful_applicative

variables {α β : Type u} {σ : Type u}

@[simp]
lemma decoder_pure (x : β) :
  (pure x : serializer σ β).decoder = pure x := rfl

@[simp]
lemma decoder_map (f : α → β) (x : serializer σ α) :
  (f <$> x).decoder = f <$> x.decoder := rfl

@[simp]
lemma decoder_seq (f : serializer σ (α → β)) (x : serializer σ α) :
  (f <*> x).decoder = f.decoder <*> x.decoder := rfl

@[simp]
lemma encoder_pure (x : β) (w : σ) :
  (pure x : serializer σ β).encoder w = pure punit.star := rfl

@[simp]
lemma encoder_map (f : α → β) (w : σ) (x : serializer σ α) :
  (f <$> x : serializer σ β).encoder w = x.encoder w := rfl

@[simp]
lemma encoder_seq (f : serializer σ (α → β)) (x : serializer σ α) (w : σ) :
  (f <*> x : serializer σ β).encoder w = f.encoder w >> x.encoder w := rfl

end lawful_applicative

instance {α} : is_lawful_functor (serializer.{u} α) :=
by refine { .. }; intros; apply serializer.eq; try { ext }; simp [map_map]

instance {α} : is_lawful_applicative (serializer.{u} α) :=
by{  constructor; intros; apply serializer.eq; try { ext };
     simp [(>>),pure_seq_eq_map,seq_assoc,bind_assoc],  }

def ser_field {α β} [serial β] (f : α → β) : serializer α β :=
{ encoder := λ x, encode (f x)
, decoder := @decode _ _ }

variables {α β σ γ : Type u} {ω : Type}

def there_and_back_again
  (y : serializer γ α) (w : γ) : option α :=
y.decoder -<< y.encoder w

lemma there_and_back_again_seq [serial α]
  (x : serializer γ (α → β)) (f : α → β) (y : γ → α) (w : γ) (w' : β)
  (h' : there_and_back_again x w = pure f)
  (h  : w' = f (y w)) :
  there_and_back_again (x <*> ser_field y) w = pure w' :=
by { simp [there_and_back_again,(>>),seq_eq_bind_map] at *,
     rw [read_write_mono h',map_read_write],
     rw [ser_field,serial.correctness], subst w', refl }

@[simp]
lemma there_and_back_again_map [serial α]
  (f : α → β) (y : γ → α) (w : γ) :
  there_and_back_again (f <$> ser_field y) w = pure (f $ y w) :=
by rw [← pure_seq_eq_map,there_and_back_again_seq]; refl

@[simp]
lemma there_and_back_again_pure (x : β) (w : γ) :
  there_and_back_again (pure x) w =
  pure x := rfl

lemma valid_serializer_of_there_and_back_again
      {α : Type*} (y : serializer α α) :
  valid_serializer y ↔
  ∀ (w : α), there_and_back_again y w = pure w :=
by { simp [valid_serializer,serial_inverse],
     repeat { rw forall_congr, intro }, refl }

open ulift

def ser_field' {α β} [serial β] (f : α → β) : serializer.{max u v} α (ulift.{v} β) :=
ser_field (up ∘ f)

def of_serializer {α} (s : serializer α α) (h : ∀ w, there_and_back_again s w = pure w) : serial α :=
{ encode := s.encoder
, decode := s.decoder
, correctness := @h }

def of_serializer₁ {f : Type u → Type v}
  (s : Π α [serial α], serializer (f α) (f α))
  (h : ∀ α [serial α] w, there_and_back_again (s α) w = pure w) : serial1 f :=
{ encode := λ α inst, (@s α inst).encoder
, decode := λ α inst, (@s α inst).decoder
, correctness := @h }

def of_serializer₂ {f : Type u → Type v → Type w}
  (s : Π α β [serial α] [serial β], serializer (f α β) (f α β))
  (h : ∀ α β [serial α] [serial β] w, there_and_back_again (s α β) w = pure w) : serial2 f :=
{ encode := λ α β inst inst', (@s α β inst inst').encoder
, decode := λ α β inst inst', (@s α β inst inst').decoder
, correctness := @h }

end serializer


import tactic.choose

/- choice -/
example (h : ∀n m : ℕ, n < m → ∃i j, m = n + i ∨ m + j = n) : true :=
begin
  choose i j h using h,
  guard_hyp i : ∀n m : ℕ, n < m → ℕ,
  guard_hyp j : ∀n m : ℕ, n < m → ℕ,
  guard_hyp h : ∀ (n m : ℕ) (h : n < m), m = n + i n m h ∨ m + j n m h = n,
  trivial
end

example (h : ∀n m : ℕ, n < m → ∃i j, m = n + i ∨ m + j = n) : true :=
begin
  choose! i j h using h,
  guard_hyp i : ℕ → ℕ → ℕ,
  guard_hyp j : ℕ → ℕ → ℕ,
  guard_hyp h : ∀ (n m : ℕ), n < m → m = n + i n m ∨ m + j n m = n,
  trivial
end

example (h : ∀n m : ℕ, ∃i, ∀n:ℕ, ∃j, m = n + i ∨ m + j = n) : true :=
begin
  choose i j h using h,
  guard_hyp i : ℕ → ℕ → ℕ,
  guard_hyp j : ℕ → ℕ → ℕ → ℕ,
  guard_hyp h : ∀ (n m k : ℕ), m = k + i n m ∨ m + j n m k = k,
  trivial
end

-- Test `simp only [exists_prop]` gets applied after choosing.
-- Because of this simp, we need a non-rfl goal
example (h : ∀ n, ∃ k ≥ 0, n = k) : ∀ x : ℕ, 1 = 1 :=
begin
  choose u hu using h,
  guard_hyp hu : ∀ n, u n ≥ 0 ∧ n = u n,
  intro, refl
end

-- test choose with conjunction
example (h : ∀ i : ℕ, ∃ j, i < j ∧ j < i+i) : true :=
begin
  choose f h h' using h,
  guard_hyp f : ℕ → ℕ,
  guard_hyp h : ∀ (i : ℕ), i < f i,
  guard_hyp h' : ∀ (i : ℕ), f i < i + i,
  trivial,
end

-- test choose with nonempty instances
universe u
example {α : Type u} (p : α → Prop) (h : ∀ i : α, p i → ∃ j : α × α, p j.1) : true :=
begin
  choose! f h using h,
  guard_hyp f : α → α × α,
  guard_hyp h : ∀ (i : α), p i → p (f i).1,
  trivial,
end