import tactic
import category_theory.functor
import data.W.basic
import category_theory.closed.types
import algebra.category.CommRing.basic
import algebra.category.Module.basic

universes w x u v

open category_theory

variables (𝒞 : Type u) [category.{v} 𝒞]

@[protect_proj] structure struc : Type (max u v (w+1) (x+1)) :=
( F : 𝒞 → Type w )
[ cat : category.{x} (sigma F) ]
( fst_map : Π {A B : sigma F} (f : A ⟶ B), A.1 ⟶ B.1 )
( fst_map_id : ∀ (A : sigma F), fst_map (𝟙 A) = 𝟙 A.1 )
( fst_map_comp : ∀ {A B C : sigma F} (f : A ⟶ B) (g : B ⟶ C),
    fst_map (f ≫ g) = fst_map f ≫ fst_map g )

namespace struc

instance : has_coe_to_fun (struc 𝒞) (λ _, 𝒞 → Type w) :=
{ coe := struc.F }

variables {𝒞} {F : struc 𝒞}

instance : category (sigma F) := F.cat

def fst : sigma F ⥤ 𝒞 :=
{ obj := sigma.fst,
  map := F.fst_map,
  map_id' := F.fst_map_id,
  map_comp' := F.fst_map_comp }

instance (X : 𝒞) : category_struct (F X) :=
{ hom := λ A B, { f : sigma.mk X A ⟶ ⟨X, B⟩ // fst.map f = 𝟙 X },
  id := λ A, ⟨𝟙 _, by simp; refl⟩,
  comp := λ A B C f g, ⟨f.1 ≫ g.1, by erw [functor.map_comp, f.2, g.2, category.comp_id]⟩ }

instance (X : 𝒞) : category (F X) :=
{ comp_id' := λ _ _ _, subtype.ext (category.comp_id _),
  id_comp' := λ _ _ _, subtype.ext (category.id_comp _),
  assoc' := λ _ _ _ _ _ _ _, subtype.ext (category.assoc _ _ _) }

open opposite

def of_functor (F : 𝒞 ⥤ Type w) : struc 𝒞 :=
{ F := F.obj,
  cat :=
  { hom := λ A B, {f : A.1 ⟶ B.1 // F.map f A.2 = B.2 },
    id := λ A, ⟨𝟙 A.1, by simp⟩,
    comp := λ A B C f g, ⟨f.1 ≫ g.1, by simp [f.prop, g.prop]⟩,
    comp_id' := λ _ _ _, subtype.ext (category.comp_id _),
    id_comp' := λ _ _ _, subtype.ext (category.id_comp _),
    assoc' := λ _ _ _ _ _ _ _, subtype.ext (category.assoc _ _ _) },
  fst_map := λ _ _, subtype.val,
  fst_map_id := by intros; refl,
  fst_map_comp := by intros; refl }

def Module₂ : struc Ring :=
{ F := λ R, Module R,
  cat :=
  { hom := λ A B, Σ f : A.1 ⟶ B.1, A.2 →ₛₗ[f] B.2,
    id := λ A, ⟨𝟙 A.1, linear_map.id⟩,
    comp := λ A B C f g, ⟨f.1 ≫ g.1,
      @linear_map.comp _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ ⟨rfl⟩ g.2 f.2⟩,
    comp_id' := by { intros, cases f, cases f_fst, cases f_snd, refl },
    id_comp' := by { intros, cases f, cases f_fst, cases f_snd, refl },
    assoc' := by { intros, refl } },
  fst_map := λ _ _ f, f.fst,
  fst_map_id := by intros; refl,
  fst_map_comp := by intros; refl }

def of_category (𝒟 : Type*) [category 𝒟] : struc 𝒞 :=
{ F := λ _, 𝒟,
  cat :=
  { hom := λ A B, (A.1 ⟶ B.1) × (A.2 ⟶ B.2),
    id := λ _, (𝟙 _, 𝟙 _),
    comp := λ A B C f g, (f.1 ≫ g.1, f.2 ≫ g.2),
    comp_id' := λ A B f, prod.ext (category.comp_id _) (category.comp_id _),
    id_comp' := λ A B f, prod.ext (category.id_comp _) (category.id_comp _),
    assoc' := λ A B C D f g h, prod.ext (category.assoc _ _ _) (category.assoc _ _ _) },
  fst_map := λ _ _, prod.fst,
  fst_map_id := λ _, rfl,
  fst_map_comp := by intros; refl }

variable (𝒞)

def type : struc 𝒞 := of_category (Type v)

def prop : struc 𝒞 := of_category Prop

lemma hcongr {α α' : Sort*}
  {β : α → Sort*} {β' : α' → Sort*} {f : Π a, β a}
  {g : Π a, β' a} (hβ : β == β')
  (a a') (h : f == g) (ha : a == a') :
  f a == g a' :=
begin
  have := type_eq_of_heq ha,
  subst this,
  simp at *,
  substs hβ ha,
  simp at *,
  subst h
end

def sigma_pi (F : 𝒞 ⥤ Type) (G : struc (sigma (of_functor F))) : struc 𝒞 :=
{ F := λ X, Π a : F.obj X, G.F ⟨X, a⟩,
  cat :=
  { hom := λ A B, Σ (f : A.1 ⟶ B.1),
      Π (a : of_functor F A.1) (b : of_functor F B.1) (hab : b = F.map f a),
      sigma.mk (sigma.mk A.1 a) (A.2 a) ⟶ sigma.mk (sigma.mk B.1 b) (B.2 b),
    id := λ X, ⟨𝟙 X.1, λ x y h, cast (by simp [F.map_id] at h; rw h)
        (𝟙 (sigma.mk (sigma.mk X.1 x) (X.2 x)))⟩,
    comp := λ X Y Z f g, ⟨f.1 ≫ g.1,
        λ a b h, cast (by simp) (f.2 a _ rfl ≫ g.2 (F.map f.1 a) b (by simp [h]))⟩,
    comp_id' := λ X Y f, begin
        cases f with f₁ f₂,
        ext,
        { simp },
        { refl },
        { intros a a' h,
          rw heq_iff_eq at h,
          subst a',
          dsimp,
          apply function.hfunext,
          { refl },
          { intros b b' h,
            rw [heq_iff_eq] at h,
            subst b',
            apply function.hfunext,
            simp,
            intros _ h _,
            subst h,
            simp } }
      end,
    id_comp' := λ X Y f, begin
        cases f with f₁ f₂,
        ext,
        { simp },
        { refl },
        { intros a a' h,
          dsimp,
          rw heq_iff_eq at h,
          subst a',
          apply function.hfunext,
          { refl },
          { intros b b' h,
            rw heq_iff_eq at h,
            subst b',
            apply function.hfunext,
            { simp * at * },
            { intros,
              simp * at *,
              convert category.id_comp (f₂ a b a'),
              { simp },
              { rw [F.map_id],
                refl },
              { simp },
              { simp } } } }
      end,
    assoc' := λ W X Y Z f g h, begin
        ext, simp [category.assoc],
        intros a a' h,
        rw [heq_iff_eq] at h,
        subst h,
        simp,
        apply function.hfunext,
        { refl },
        { intros b b' h,
          rw heq_iff_eq at h,
          subst b',
          apply function.hfunext,
          { simp [category.assoc] },
          { intros c c' h,
            simp,
            dsimp,
            congr,
            { simp },
            { rw F.map_comp, refl },
            { apply hcongr,
              apply function.hfunext,
              rw F.map_comp; refl,
              intros,
              rw [F.map_comp],
              refl,
              rw [F.map_comp],
              refl,
              exact proof_irrel_heq _ _ },
            { apply hcongr,
              apply function.hfunext,
              rw F.map_comp; refl,
              intros,
              rw [F.map_comp],
              refl,
              rw [F.map_comp],
              refl,
              exact proof_irrel_heq _ _ } } }
      end },
  fst_map := λ _ _ f, f.fst,
  fst_map_id := by intros; refl,
  fst_map_comp := by intros; refl }

example : 1 = 1 := rfl

def sigma_arrow (F : 𝒞 ⥤ Type) (G : struc 𝒞) : struc 𝒞 :=
{ F := λ X, F.obj X → G X,
  cat :=
  { hom := λ A B, Σ (f : A.1 ⟶ B.1),
      Π (a : of_functor F A.1) (b : of_functor F B.1) (h : b = F.map f a),
      { g : sigma.mk A.1 (A.2 a) ⟶ sigma.mk B.1 (B.2 b) // fst.map g = f } ,
    id := λ X, ⟨𝟙 X.1, λ x y h, ⟨cast (by simp [h]) (𝟙 (sigma.mk X.1 (X.2 x))),
      begin simp, end⟩⟩,
    comp := λ X Y Z f g, ⟨f.1 ≫ g.1,
        λ x z h, cast (by simp [h]) (f.2 x (F.map f.1 x) rfl ≫ g.2 (F.map f.1 x) z (by simp [h]))⟩,
    comp_id' := λ X Y f,
      begin
        cases f with f₁ f₂,
        ext,
        { simp },
        { refl },
        { intros a a' h,
          rw heq_iff_eq at h,
          subst a',
          apply function.hfunext,
          { refl },
          { intros b b' h,
            rw heq_iff_eq at h,
            subst b',
            dsimp,
            apply function.hfunext,
            { simp },
            { intros _ h _,
              subst h,
              simp } } }
      end,
    id_comp' := λ X Y f, begin
        cases f with f₁ f₂,
        ext,
        { simp },
        { refl },
        { intros a a' h,
          dsimp,
          rw heq_iff_eq at h,
          subst a',
          apply function.hfunext,
          { refl },
          { intros b b' h,
            rw heq_iff_eq at h,
            subst b',
            apply function.hfunext,
            { simp * at * },
            { intros,
              simp * at *,
              convert category.id_comp (f₂ a b a'),
              { simp },
              { simp },
              { simp } } } }
      end,
    assoc' := λ W X Y Z f g h, begin
        ext, simp [category.assoc],
        intros a a' h,
        rw [heq_iff_eq] at h,
        subst h,
        simp,
        apply function.hfunext,
        { refl },
        { intros b b' h,
          rw heq_iff_eq at h,
          subst b',
          apply function.hfunext,
          { simp [category.assoc] },
          { intros c c' h,
            simp,
            dsimp,
            congr,
            { simp },
            { apply hcongr,
              apply function.hfunext,
              rw F.map_comp; refl,
              intros,
              rw [F.map_comp],
              refl,
              rw [F.map_comp],
              refl,
              exact proof_irrel_heq _ _ },
            { apply hcongr,
              apply function.hfunext,
              rw F.map_comp; refl,
              intros,
              rw [F.map_comp],
              refl,
              rw [F.map_comp],
              refl,
              exact proof_irrel_heq _ _ } } }
      end },
  fst_map := λ _ _ f, f.fst,
  fst_map_id := by intros; refl,
  fst_map_comp := by intros; refl }

-- def sigma_arrow (F : struc 𝒞) (G : struc 𝒞) : struc 𝒞 :=
-- { F := λ X, Σ (i : F X → G X),
--     (Π (a b : F X), (sigma.mk X a ⟶ ⟨X, b⟩) →
--       { f : sigma.mk X (i a) ⟶ ⟨X, i b⟩ // fst.map f = 𝟙 X}),
--   cat :=
--   { hom := λ A B, Σ (f : A.1 ⟶ B.1), Π (a : F A.1) (b : F B.1),
--       (sigma.mk A.1 a ⟶ sigma.mk B.1 b) →
--       { g : (sigma.mk A.1 (A.2.1 a)) ⟶ (sigma.mk B.1 (B.2.1 b)) // fst.map g = f },
--     id := λ A, ⟨𝟙 _, λ a b f, A.2.2 a b f⟩,
--     comp := λ A B C f g, ⟨f.1 ≫ g.1, λ a c h,
--       begin
--         have := sigma.snd f a,

--       end⟩,
--     comp_id' := sorry,
--     id_comp' := sorry,
--     assoc' := sorry },

--   fst_map := λ _ _ f, f.fst,
--   fst_map_id := by intros; refl,
--   fst_map_comp := by intros; refl }

-- def sigma_pi₂ (F : struc 𝒞) (G : struc (sigma F)) : struc 𝒞 :=
-- { F := λ X, Σ (i : Π a : F X, G.F ⟨X, a⟩),
--     (Π (a b : F X) (f : sigma.mk X a ⟶ sigma.mk X b),
--       { g : sigma.mk (sigma.mk X a) (i a) ⟶ ⟨⟨X, b⟩, i b⟩ // fst.map g = f }),
--   cat :=
--   { hom := λ A B, Σ (f : A.1 ⟶ B.1), (Π (a : F A.1) (b : F B.1),
--       (sigma.mk A.1 a ⟶ sigma.mk B.1 b) →
--         (sigma.mk (sigma.mk A.1 a) (A.2.1 a) ⟶ ⟨⟨B.1, b⟩, B.2.1 b⟩)),
--     id := λ A, ⟨𝟙 _, λ a b f, (A.2.2 a b f).1⟩,
--     comp := λ A B C f g, ⟨f.1 ≫ g.1, λ a c h,
--       begin
--         have := sigma.snd f a,

--       end⟩,
--     comp_id' := sorry,
--     id_comp' := sorry,
--     assoc' := sorry },

--   fst_map := λ _ _ f, f.fst,
--   fst_map_id := by intros; refl,
--   fst_map_comp := by intros; refl }

end struc


import algebra.homology.quasi_iso
import algebra.homology.short_complex.pseudoelements
import for_mathlib.algebra.homology.hom_complex_shift
import category_theory.triangulated.triangulated
import for_mathlib.algebra.homology.homological_complex_limits

noncomputable theory
open category_theory category_theory.category category_theory.limits
  category_theory.pretriangulated

@[simp]
lemma category_theory.limits.biprod.is_zero_iff {C : Type*} [category C]
  [has_zero_morphisms C] (A B : C)
  [has_binary_biproduct A B] : is_zero (biprod A B) ↔ is_zero A ∧ is_zero B :=
begin
  split,
  { intro h,
    simp only [is_zero.iff_id_eq_zero],
    split,
    { rw ← cancel_mono (biprod.inl : _ ⟶ A ⊞ B),
      apply h.eq_of_tgt, },
    { rw ← cancel_mono (biprod.inr : _ ⟶ A ⊞ B),
      apply h.eq_of_tgt, }, },
  { rintro ⟨h₁, h₂⟩,
    rw is_zero.iff_id_eq_zero,
    ext1,
    { apply h₁.eq_of_src, },
    { apply h₂.eq_of_src, }, },
end

open category_theory category_theory.category category_theory.limits
  category_theory.pretriangulated

namespace cochain_complex

variables {C : Type*} [category C]

section preadditive

variables [preadditive C]
  {F G : cochain_complex C ℤ} [∀ p, has_binary_biproduct (F.X (p+1)) (G.X p)]
  (φ : F ⟶ G)

open hom_complex

include φ

def mapping_cone : cochain_complex C ℤ :=
{ X := λ i, F.X (i+1) ⊞ G.X i,
  d := λ i j, begin
    by_cases i+1 = j,
    { exact -biprod.fst ≫ F.d _ _ ≫ biprod.inl +
        biprod.fst ≫ (cochain.of_hom φ).v (i+1) j (by simpa only [add_zero, ← h]) ≫ biprod.inr +
        biprod.snd ≫ G.d _ _ ≫ biprod.inr,
      },
    { exact 0, },
  end,
  shape' := λ i j (hij : i+1 ≠ j), by rw dif_neg hij,
  d_comp_d' := λ i j k (hij : i+1=j) (hjk : j+1=k), begin
    simp only [dif_pos hij, hjk, dif_pos rfl],
    substs hij hjk,
    ext1,
    { dsimp,
      simp only [assoc, preadditive.comp_add, preadditive.add_comp, preadditive.neg_comp,
        preadditive.comp_neg, biprod.inl_fst_assoc, biprod.inl_snd_assoc, zero_comp,
        F.d_comp_d_assoc, neg_zero, zero_add, biprod.inr_fst_assoc, biprod.inr_snd_assoc,
        comp_zero, add_zero, cochain.of_hom_v, φ.comm_assoc, add_left_neg], },
    { simp only [assoc, preadditive.comp_add, preadditive.add_comp, preadditive.neg_comp,
        preadditive.comp_neg, biprod.inr_fst_assoc, biprod.inr_snd_assoc, zero_comp,
        comp_zero, neg_zero, zero_add, G.d_comp_d_assoc], },
  end, }

omit φ

namespace mapping_cone

include φ

lemma X_is_zero_iff (n : ℤ) :
  is_zero ((mapping_cone φ).X n) ↔ is_zero (F.X (n+1)) ∧ is_zero (G.X n) :=
biprod.is_zero_iff _ _

def inl : cochain F (mapping_cone φ) (-1) :=
cochain.mk (λ p q hpq, (cochain.of_hom (𝟙 F)).v p (q+1) (by linarith) ≫ biprod.inl)

def inr : G ⟶ mapping_cone φ :=
cocycle.hom_of
  (cocycle.mk (cochain.mk (λ p q hpq,
    (cochain.of_hom (𝟙 G)).v p q hpq ≫ biprod.inr)) 1 (zero_add 1)
    begin
      ext1 p _ rfl,
      dsimp [mapping_cone],
      simp only [cochain.zero_v,
        δ_v 0 1 (zero_add 1) _ p _ rfl p (p+1) (by linarith) rfl,
        zero_add, cochain.mk_v, cochain.of_hom_v, homological_complex.id_f, id_comp,
        ε_1, neg_smul, one_zsmul, dif_pos rfl, preadditive.comp_add,
        preadditive.comp_neg, biprod.inr_fst_assoc, zero_comp, neg_zero,
        biprod.inr_snd_assoc, add_right_neg],
    end)

def fst : cocycle (mapping_cone φ) F 1 :=
cocycle.mk (cochain.mk (λ p q hpq, biprod.fst ≫
  (cochain.of_hom (𝟙 F)).v (p+1) q (by simpa only [add_zero] using hpq))) 2 (by linarith)
  begin
    ext1 p q hpq,
    have hpq' : q = p + 1 + 1 := by linarith,
    subst hpq',
    dsimp [mapping_cone],
    simp only [δ_v 1 2 (by linarith) _ p (p+1+1) (by linarith) (p+1) (p+1) (by linarith) rfl,
      dif_pos rfl, ε_succ, ε_1, neg_neg, one_smul, cochain.mk_v, cochain.of_hom_v, id_comp,
      homological_complex.id_f, assoc, preadditive.add_comp, preadditive.comp_add, comp_id,
      preadditive.neg_comp, preadditive.comp_neg, biprod.inl_fst, biprod.inr_fst, comp_zero,
      add_zero, add_right_neg],
  end

def snd : cochain (mapping_cone φ) G 0 :=
cochain.mk (λ p q hpq, biprod.snd ≫ (cochain.of_hom (𝟙 G)).v p q hpq)

@[simp, reassoc]
lemma inl_fst (p q : ℤ) (hpq : p = q+1) :
  (inl φ).v p q (by rw [hpq, int.add_neg_one, add_tsub_cancel_right]) ≫
     (fst φ : cochain (mapping_cone φ) F 1).v q p hpq = 𝟙 _ :=
begin
  subst hpq,
  dsimp [inl, fst],
  simp only [cochain.of_hom_v, homological_complex.id_f, id_comp, biprod.inl_fst_assoc],
  erw [cochain.of_hom_v, homological_complex.id_f],
end

@[simp, reassoc]
lemma inl_snd (p q : ℤ) (hpq : q = p+(-1)) :
  (inl φ).v p q hpq ≫ (snd φ).v q q (add_zero q).symm = 0 :=
begin
  dsimp [inl, snd],
  simp only [assoc, biprod.inl_snd_assoc, zero_comp, comp_zero],
end

@[simp, reassoc]
lemma inr_fst (p q : ℤ) (hpq : q = p+1) :
  (inr φ).f p ≫ (fst φ : cochain (mapping_cone φ) F 1).v p q hpq = 0 :=
begin
  subst hpq,
  dsimp [inr, fst],
  simp only [cochain.of_hom_v, homological_complex.id_f, id_comp, biprod.inr_fst_assoc, zero_comp],
end

@[simp, reassoc]
lemma inr_snd (p : ℤ) :
  (inr φ).f p ≫ (snd φ).v p p (add_zero p).symm = 𝟙 _ :=
begin
  dsimp [inr, snd],
  simp only [assoc, biprod.inr_snd, cochain.of_hom_v, homological_complex.id_f, comp_id],
end

lemma id (p q : ℤ) (hpq : q = p+1) :
  (fst φ : cochain (mapping_cone φ) F 1).v p q hpq ≫
    (inl φ).v q p (by rw [hpq, int.add_neg_one, add_tsub_cancel_right]) +
  (snd φ).v p p (add_zero p).symm ≫ (inr φ).f p = 𝟙 _ :=
begin
  subst hpq,
  dsimp [inl, inr, fst, snd],
  simp only [cochain.of_hom_v, homological_complex.id_f, id_comp, assoc],
  erw [cochain.of_hom_v, homological_complex.id_f, id_comp],
  apply biprod.total,
end

@[reassoc]
lemma inl_d (n₁ n₂ n₃ : ℤ) (h₁₂ : n₁ = n₂ + (-1)) (h₂₃ : n₂ = n₃ + (-1)) :
  (inl φ).v n₂ n₁ h₁₂ ≫ (mapping_cone φ).d n₁ n₂ =
    φ.f n₂ ≫ (inr φ).f n₂ - F.d n₂ n₃ ≫ (inl φ).v _ _ h₂₃ :=
begin
  have hn₂ : n₂ = n₁ + 1 := by linarith,
  have hn₃ : n₃ = n₁ + 1 + 1 := by linarith,
  substs hn₂ hn₃,
  dsimp [mapping_cone, inl, inr],
  simp only [dif_pos rfl, add_zero, cochain.of_hom_v, homological_complex.id_f, id_comp,
    preadditive.comp_add, preadditive.comp_neg, biprod.inl_fst_assoc,
    biprod.inl_snd_assoc, zero_comp],
  erw [cochain.of_hom_v, neg_add_eq_sub],
end

@[simp, reassoc]
lemma inr_d (n n' : ℤ) :
  (inr φ).f n ≫ (mapping_cone φ).d n n' =
    G.d n n' ≫ (inr φ).f n' :=
begin
  by_cases h : n+1 = n',
  { subst h,
    dsimp [inr, mapping_cone],
    simp only [dif_pos rfl],
    simp only [cochain.of_hom_v, homological_complex.id_f, id_comp, preadditive.comp_add,
      preadditive.comp_neg, biprod.inr_fst_assoc, zero_comp, neg_zero,
      biprod.inr_snd_assoc, zero_add], },
  { change ¬ (complex_shape.up ℤ).rel n n' at h,
    simp only [homological_complex.shape _ _ _ h, zero_comp, comp_zero], },
end

attribute [irreducible] mapping_cone inl inr fst snd

@[simps]
def X_iso (n i : ℤ) (hi : i = n+1) [has_binary_biproduct (F.X i) (G.X n)] :
  (mapping_cone φ).X n ≅ F.X i ⊞ G.X n :=
{ hom := (fst φ : cochain (mapping_cone φ) F 1).v n i hi ≫ biprod.inl +
    (snd φ).v n n (add_zero n).symm   ≫ biprod.inr,
  inv := biprod.fst ≫ (inl φ).v i n (by linarith) + biprod.snd ≫ (inr φ).f n,
  hom_inv_id' := by simp only [add_zero, zero_add, preadditive.comp_add,
    preadditive.add_comp_assoc, assoc, biprod.inl_fst, comp_id, biprod.inr_fst,
    comp_zero, biprod.inl_snd, biprod.inr_snd, id],
  inv_hom_id' := begin
      ext1,
      { simp only [assoc, preadditive.comp_add, preadditive.add_comp,
          biprod.inl_fst_assoc, biprod.inl_snd_assoc, zero_comp, add_zero, comp_id,
          inl_fst_assoc, inl_snd_assoc], },
      { simp only [assoc, preadditive.comp_add, preadditive.add_comp,
          biprod.inr_fst_assoc, biprod.inr_snd_assoc, zero_comp, zero_add, comp_id,
          inr_fst_assoc, inr_snd_assoc], },
    end, }

@[simp]
lemma inl_comp_fst :
  (inl φ).comp (fst φ : cochain (mapping_cone φ) F 1) (neg_add_self 1).symm =
    cochain.of_hom (𝟙 _) :=
begin
  ext n,
  simp only [cochain.comp_v _ _ (neg_add_self 1).symm n (n-1) n (by linarith) (by linarith),
    inl_fst, cochain.of_hom_v, homological_complex.id_f],
end

@[simp]
lemma inl_comp_snd :
  (inl φ).comp (snd φ) (add_zero _).symm = 0 :=
begin
  ext n,
  simp only [cochain.comp_zero_cochain, inl_snd, cochain.zero_v],
end

@[simp]
lemma inr_comp_fst :
  (cochain.of_hom (inr φ)).comp (fst φ : cochain (mapping_cone φ) F 1) (zero_add 1).symm = 0 :=
by tidy

@[simp]
lemma inr_comp_snd :
  (cochain.of_hom (inr φ)).comp
    (snd φ : cochain (mapping_cone φ) G 0) (zero_add 0).symm = cochain.of_hom (𝟙 _) :=
by tidy

@[simps]
def δ_as_cocycle : cocycle (mapping_cone φ) F 1 :=
-fst φ

def δ : mapping_cone φ ⟶ F⟦(1 : ℤ)⟧ :=
cocycle.hom_of (cocycle.right_shift (δ_as_cocycle φ) 1 0 (zero_add 1).symm)

@[simp, priority 1100]
lemma inr_δ : inr φ ≫ δ φ = 0 :=
begin
  ext n,
  dsimp only [δ],
  simp only [homological_complex.comp_f, cocycle.hom_of_f, cochain.neg_v,
    cocycle.right_shift_coe, δ_as_cocycle_coe, homological_complex.zero_f_apply,
    hom_complex.cochain.right_shift_v _ 1 0 (zero_add 1).symm n n (by linarith) _ rfl,
    preadditive.neg_comp, preadditive.comp_neg, inr_fst_assoc, zero_comp, neg_zero],
end

@[simp]
lemma inl_δ :
  (inl φ).comp (cochain.of_hom (δ φ)) (add_zero _).symm =
  -(cochain.of_hom (𝟙 F)).right_shift _ _ (add_neg_self 1).symm :=
begin
  /- TODO deduplicate the proof of this and the lemma above -/
  ext p q hpq,
  simp only [cochain.comp_zero_cochain, cochain.of_hom_v, δ,
    cocycle.hom_of_f, cocycle.right_shift_coe, δ_as_cocycle_coe,
    hom_complex.cochain.right_shift_v _ 1 0 (zero_add 1).symm q q (by linarith) p (by linarith),
    hom_complex.cochain.right_shift_v _ 1 (-1) (add_neg_self 1).symm p q hpq p (by linarith),
    cochain.neg_v, preadditive.comp_neg, preadditive.neg_comp, cochain.neg_v,
    inl_fst_assoc, homological_complex.id_f, id_comp],
end

variable {φ}

lemma to_ext_iff {A : C} {n : ℤ} (f g : A ⟶ (mapping_cone φ).X n) (n' : ℤ) (hn' : n' = n+1) :
  f = g ↔ f ≫ (fst φ : cochain (mapping_cone φ) F 1).v n n' hn' =
    g ≫ (fst φ : cochain (mapping_cone φ) F 1).v n n' hn' ∧
    f ≫ (snd φ).v n n (add_zero n).symm = g ≫ (snd φ).v n n (add_zero n).symm :=
begin
  split,
  { rintro rfl,
    tauto, },
  { intro hfg,
    rw [← cancel_mono (𝟙 ((mapping_cone φ).X n))],
    simp only [← id _ _ _ hn', preadditive.comp_add, reassoc_of hfg.1, reassoc_of hfg.2], },
end

lemma from_ext_iff {A : C} {n : ℤ} (f g : (mapping_cone φ).X n ⟶ A)
  (n' : ℤ) (h : n' = n+1) :
  f = g ↔ (inl φ).v n' n (by rw [h, int.add_neg_one, add_tsub_cancel_right]) ≫ f =
    (inl φ).v n' n (by rw [h, int.add_neg_one, add_tsub_cancel_right]) ≫ g ∧
    (inr φ).f n ≫ f = (inr φ).f n ≫ g :=
begin
  haveI : has_binary_biproduct (F.X n') (G.X n) := by { subst h, apply_instance, },
  split,
  { rintro rfl,
    tauto, },
  { intro hfg,
    rw [← cancel_epi (𝟙 ((mapping_cone φ).X n))],
    simp only [← id _ _ _ h, preadditive.add_comp, assoc, hfg.1, hfg.2], },
end

variable (φ)

@[reassoc]
lemma d_fst (n₁ n₂ n₃ : ℤ) (h₁₂ : n₂ = n₁ + 1) (h₂₃ : n₃ = n₂ + 1) :
  (mapping_cone φ).d n₁ n₂ ≫ (fst φ : cochain (mapping_cone φ) F 1).v n₂ n₃ h₂₃ =
  -(fst φ : cochain (mapping_cone φ) F 1).v n₁ n₂ h₁₂ ≫ F.d n₂ n₃ :=
by simp only [from_ext_iff _ _ _ h₁₂, inl_d_assoc _ n₁ n₂ n₃ (by linarith) (by linarith),
  assoc, preadditive.sub_comp, inr_fst, comp_zero, inl_fst, comp_id, zero_sub,
  preadditive.comp_neg, inl_fst_assoc, inr_d_assoc, inr_fst_assoc, zero_comp, neg_zero,
  eq_self_iff_true, and_self]

@[reassoc]
lemma d_snd (n₁ n₂ : ℤ) (h₁₂ : n₂ = n₁ + 1) :
  (mapping_cone φ).d n₁ n₂ ≫ (snd φ).v n₂ n₂ (add_zero n₂).symm =
    (fst φ : cochain (mapping_cone φ) F 1).v n₁ n₂ h₁₂ ≫ φ.f n₂ +
    (snd φ).v n₁ n₁ (add_zero n₁).symm ≫ G.d n₁ n₂ :=
by simp only [from_ext_iff _ _ _ h₁₂, assoc,
  inl_d_assoc _ n₁ n₂ (n₂+1) (by linarith) (by linarith),
  preadditive.sub_comp, inl_snd, comp_zero, sub_zero, preadditive.comp_add,
  inl_snd_assoc, zero_comp, add_zero, inl_fst_assoc, inr_snd, comp_id,
  inr_d_assoc, inr_fst_assoc, zero_add, inr_snd_assoc,
  eq_self_iff_true, and_self]

@[simp]
lemma δ_inl :
  hom_complex.δ (-1) 0 (inl φ) = cochain.of_hom (φ ≫ inr φ) :=
begin
  ext p,
  simp only [δ_v (-1) 0 (neg_add_self 1) _ p p (add_zero p).symm _ _ rfl rfl,
    inl_d φ (p-1) p (p+1) (by linarith)( by linarith),
    add_left_neg, ε_0, one_zsmul, sub_add_cancel, cochain.of_hom_comp,
    cochain.comp_zero_cochain, cochain.of_hom_v],
end

@[simp]
lemma δ_snd :
  hom_complex.δ 0 1 (snd φ) = -(fst φ : cochain (mapping_cone φ) F 1).comp
    (cochain.of_hom φ) (add_zero 1).symm :=
begin
  ext p q hpq,
  simp only [δ_v 0 1 (zero_add 1) _ p q hpq p q (by linarith) hpq, d_snd _ _ _ hpq,
    zero_add, add_zero, neg_neg, neg_zero, neg_eq_zero, add_tsub_cancel_right, ε_1,
    smul_add, neg_smul, one_zsmul, add_neg_cancel_comm_assoc, cochain.neg_v,
    cochain.comp_zero_cochain, cochain.of_hom_v],
  abel,
end

omit φ
lemma _root_.int.two_eq_one_add_one : (2 : ℤ) = 1+1 := by linarith
lemma _root_.int.one_eq_two_add_neg_one : (1 : ℤ) = 2+(-1) := by linarith

lemma of_d_eq : cochain.of_d (mapping_cone φ) =
  -((fst φ : cochain (mapping_cone φ) F 1).comp (cochain.of_d F)
    int.two_eq_one_add_one).comp (inl φ) int.one_eq_two_add_neg_one +
  ((fst φ : cochain (mapping_cone φ) F 1).comp (cochain.of_hom φ) (add_zero 1).symm).comp
      (cochain.of_hom (inr φ)) (add_zero 1).symm +
  ((snd φ).comp (cochain.of_d G) (zero_add 1).symm).comp
    (cochain.of_hom (inr φ)) (add_zero 1).symm :=
begin
  ext p q hpq,
  simp only [from_ext_iff _ _ _ hpq,
    cochain.of_d_v, inl_d φ p q (q+1) (by linarith) (by linarith), cochain.add_v,
    preadditive.comp_add, cochain.comp_assoc_of_third_is_zero_cochain, cochain.comp_zero_cochain,
    cochain.of_hom_v, inl_fst_assoc, cochain.neg_v, inl_snd_assoc, zero_comp,
    cochain.comp_assoc_of_first_is_zero_cochain, cochain.zero_cochain_comp, preadditive.comp_neg,
    cochain.comp_v _ _ int.one_eq_two_add_neg_one p (q+1) q (by linarith) (by linarith),
    cochain.comp_v _ _ _root_.int.two_eq_one_add_one p q (q+1) hpq rfl, assoc, add_zero,
    inl_fst_assoc, inr_d, inr_fst_assoc, neg_zero, zero_add, inr_snd_assoc, sub_eq_neg_add,
    eq_self_iff_true, and_true],
end

variable {φ}

lemma to_decomposition {A : C} {n : ℤ} (f : A ⟶ (mapping_cone φ).X n)
  (n' : ℤ) (h : n' = n+1) :
  ∃ (x : A ⟶ F.X n') (y : A ⟶ G.X n), f = x ≫
    (inl φ : cochain F (mapping_cone φ) (-1)).v n' n (by rw [h, int.add_neg_one, add_tsub_cancel_right])
      + y ≫ (inr φ).f n :=
begin
  refine ⟨f ≫ (fst φ : cochain (mapping_cone φ) F 1).v _ _ (by linarith), f ≫ (snd φ).v n n (by linarith), _⟩,
  have h := f ≫= id φ n n' h,
  rw comp_id at h,
  nth_rewrite 0 ← h,
  simp only [preadditive.comp_add, assoc],
end

lemma cochain_ext {K : cochain_complex C ℤ} {m m' : ℤ}
  (y₁ y₂ : cochain (mapping_cone φ) K m) (hm' : m = m'+1) :
  y₁ = y₂ ↔ (inl φ).comp y₁ (show m' = -1+m, by rw [hm', neg_add_cancel_comm_assoc]) =
    (inl φ).comp y₂ (show m' = -1+m, by rw [hm', neg_add_cancel_comm_assoc]) ∧
    (cochain.of_hom (inr φ)).comp y₁ (zero_add m).symm =
      (cochain.of_hom (inr φ)).comp y₂ (zero_add m).symm :=
begin
  split,
  { rintro rfl,
    tauto, },
  { rintro ⟨h₁, h₂⟩,
    ext p q hpq,
    replace h₁ := cochain.congr_v h₁ (p+1) q (by linarith),
    replace h₂ := cochain.congr_v h₂ p q (by linarith),
    simp only [cochain.comp_v _ _ (show m' = -1+m, by linarith) (p+1) p q (by linarith) hpq] at h₁,
    simp only [cochain.zero_cochain_comp, cochain.of_hom_v] at h₂,
    rw [from_ext_iff _ _ (p+1) rfl, h₁, h₂],
    tauto, },
end

lemma cochain_ext' {K : cochain_complex C ℤ} {m m' : ℤ}
  (y₁ y₂ : cochain K (mapping_cone φ) m) (hm' : m' = m+1) :
  y₁ = y₂ ↔ y₁.comp (fst φ : cochain (mapping_cone φ) F 1) hm' =
    y₂.comp (fst φ : cochain (mapping_cone φ) F 1) hm' ∧
    y₁.comp (snd φ) (add_zero m).symm =
      y₂.comp (snd φ) (add_zero m).symm :=
begin
  split,
  { rintro rfl,
    tauto, },
  { rintro ⟨h₁, h₂⟩,
    ext p q hpq,
    replace h₁ := cochain.congr_v h₁ p (q+1) (by linarith),
    simp only [cochain.comp_v _ _ hm' p q (q+1) (by linarith) (by linarith)] at h₁,
    replace h₂ := cochain.congr_v h₂ p q (by linarith),
    simp only [cochain.comp_zero_cochain] at h₂,
    rw [to_ext_iff _ _ (q+1) rfl, h₂, h₁],
    tauto, },
end

variable (φ)

@[simp]
def ι' := (homotopy_category.quotient _ _).map (inr φ)

def δ' : (homotopy_category.quotient _ _).obj (mapping_cone φ) ⟶
  ((homotopy_category.quotient _ _).obj F)⟦(1 : ℤ)⟧ :=
(homotopy_category.quotient _ _).map (δ φ)

def desc_cochain {K : cochain_complex C ℤ} {n m : ℤ} (α : cochain F K m) (β : cochain G K n)
  (h : m+1=n) :
  cochain (mapping_cone φ) K n :=
(fst φ : cochain (mapping_cone φ) F 1).comp α (show n = 1+m, by rw [← h, add_comm])
  + (snd φ).comp β (zero_add n).symm

@[simp, reassoc]
lemma inl_desc_cochain_v {K : cochain_complex C ℤ} {n m : ℤ}
  (α : cochain F K m) (β : cochain G K n) (h : m+1=n) (p₁ p₂ p₃ : ℤ)
    (h₁₂ : p₂ = p₁ + (-1)) (h₂₃ : p₃ = p₂ + n) :
  (inl φ).v p₁ p₂ h₁₂ ≫ (desc_cochain φ α β h).v p₂ p₃ h₂₃ =
      α.v p₁ p₃ (by rw [h₂₃, h₁₂, ← h, int.add_neg_one, sub_add_add_cancel]) :=
begin
  dsimp [desc_cochain],
  simp only [add_zero, cochain.zero_cochain_comp, preadditive.comp_add, zero_comp,
    cochain.comp_v _ _ (show n = 1 + m, by linarith) p₂ p₁ p₃ (by linarith) (by linarith),
    inl_fst_assoc, inl_snd_assoc],
end

@[simp, reassoc]
lemma inr_desc_cochain_v {K : cochain_complex C ℤ} {n m : ℤ}
  (α : cochain F K m) (β : cochain G K n) (h : m+1=n) (p₁ p₂ : ℤ)
    (h₁₂ : p₂ = p₁ + n) :
  (inr φ).f p₁ ≫ (desc_cochain φ α β h).v p₁ p₂ h₁₂ =
      β.v p₁ p₂ h₁₂ :=
begin
  dsimp [desc_cochain],
  simp only [cochain.zero_cochain_comp, preadditive.comp_add, inr_snd_assoc, add_left_eq_self,
    cochain.comp_v _ _ (show n = 1 + m, by linarith) p₁ (p₁ + 1) p₂ rfl (by linarith),
    inr_fst_assoc, zero_comp],
end

@[simp]
lemma inl_desc_cochain {K : cochain_complex C ℤ} {n m : ℤ}
  (α : cochain F K m) (β : cochain G K n) (h : m+1=n) :
  (inl φ).comp (desc_cochain φ α β h)
    (show m = -1+n, by rw [← h, neg_add_cancel_comm_assoc]) = α :=
begin
  ext p q hpq,
  simp only [cochain.comp_v _ _ (show m = -1 + n, by linarith)
    p (p-1) q (by linarith) (by linarith), inl_desc_cochain_v],
end

@[simp]
lemma inr_desc_cochain {K : cochain_complex C ℤ} {n m : ℤ}
  (α : cochain F K m) (β : cochain G K n) (h : m+1=n) :
  (cochain.of_hom (inr φ)).comp
    (desc_cochain φ α β h) (zero_add n).symm = β  :=
begin
  ext p q hpq,
  simp only [cochain.comp_v _ _ (zero_add n).symm p p q (add_zero p).symm hpq,
    cochain.of_hom_v, inr_desc_cochain_v],
end

lemma δ_desc_cochain {K : cochain_complex C ℤ} {n m n' : ℤ} (α : cochain F K m) (β : cochain G K n)
  (h : m+1=n) (hn' : n+1 = n') : hom_complex.δ n n' (desc_cochain φ α β h) =
  (fst φ : cochain (mapping_cone φ) F 1).comp (hom_complex.δ m n α +
    ε (n+1) • (cochain.of_hom φ).comp β (zero_add n).symm) (by rw [← hn', add_comm]) +
    (snd φ).comp (hom_complex.δ n n' β) (zero_add n').symm :=
begin
  ext p q hpq,
  simp only [from_ext_iff _ _ (p+1) rfl,
    δ_v n n' hn' _ p q hpq (q-1) (p+1) rfl rfl, cochain.add_v,
    cochain.comp_v _ _ (show n' = 1+n, by linarith) p (p+1) q rfl (by linarith),
    zero_add, neg_zero, add_zero, ε_succ, neg_smul, preadditive.comp_add,
    inl_desc_cochain_v_assoc, preadditive.comp_neg, linear.comp_smul, cochain.neg_v,
    cochain.zsmul_v, cochain.zero_cochain_comp, cochain.of_hom_v, inl_fst_assoc,
    inl_snd_assoc, zero_comp, inr_desc_cochain_v_assoc, inr_d_assoc, inr_desc_cochain_v,
    inr_fst_assoc, smul_zero, inr_snd_assoc, smul_sub, show m = n-1, by linarith,
    inl_d_assoc φ p (p+1) (p+2) (by linarith) (by linarith),
    δ_v m n h _ (p+1) q (by linarith) (q-1) (p+2) rfl (by linarith),
    preadditive.sub_comp, assoc, inl_desc_cochain_v, ε_sub, ε_1, mul_neg, mul_one, neg_neg],
  exact ⟨by abel, rfl⟩,
end

def desc_cocycle {K : cochain_complex C ℤ} {n m : ℤ} (α : cochain F K m) (β : cocycle G K n)
  (h : m+1=n) (eq : hom_complex.δ m n α =
    ε n • (cochain.of_hom φ).comp (β : cochain G K n) (zero_add n).symm) :
  cocycle (mapping_cone φ) K n :=
cocycle.mk (desc_cochain φ α (β : cochain G K n) h) (n+1) rfl
  (by simp only [δ_desc_cochain φ α (β : cochain G K n) h rfl, ε_add, ε_1, mul_neg, mul_one, eq,
    neg_smul, ← sub_eq_add_neg, sub_self, cochain.comp_zero, zero_add,
    cocycle.δ_eq_zero, cochain.comp_zero])

@[simp]
lemma desc_cocycle_coe {K : cochain_complex C ℤ} {n m : ℤ} (α : cochain F K m) (β : cocycle G K n)
  (h : m+1=n) (eq : hom_complex.δ m n α = ε n • (cochain.of_hom φ).comp β.1 (zero_add n).symm) :
(desc_cocycle φ α β h eq : cochain (mapping_cone φ) K n) =
  desc_cochain φ α β h := rfl

def desc {K : cochain_complex C ℤ} (α : cochain F K (-1)) (β : G ⟶ K)
  (eq : hom_complex.δ (-1) 0 α = cochain.of_hom (φ ≫ β)) :
  mapping_cone φ ⟶ K :=
cocycle.hom_of (desc_cocycle φ α (cocycle.of_hom β) (neg_add_self 1)
  (by simp only [eq, ε_0, cochain.of_hom_comp, subtype.val_eq_coe, cocycle.of_hom_coe, one_zsmul]))

@[simp, reassoc]
lemma inl_desc_v {K : cochain_complex C ℤ} (α : cochain F K (-1)) (β : G ⟶ K)
  (eq : hom_complex.δ (-1) 0 α = cochain.of_hom (φ ≫ β)) (p q : ℤ) (hpq : q = p + (-1)) :
  (inl φ).v p q hpq ≫ (desc φ α β eq).f q = α.v p q hpq :=
begin
  dsimp only [desc],
  simp only [cocycle.hom_of_f, desc_cocycle_coe, inl_desc_cochain_v],
end

@[simp]
lemma inl_desc {K : cochain_complex C ℤ} (α : cochain F K (-1)) (β : G ⟶ K)
  (eq : hom_complex.δ (-1) 0 α = cochain.of_hom (φ ≫ β)) :
  (inl φ).comp (cochain.of_hom (desc φ α β eq)) (add_zero _).symm = α :=
by tidy

@[simp, reassoc]
lemma inr_desc_f {K : cochain_complex C ℤ} (α : cochain F K (-1)) (β : G ⟶ K)
  (eq : hom_complex.δ (-1) 0 α = cochain.of_hom (φ ≫ β)) (n : ℤ):
  (inr φ).f n ≫ (desc φ α β eq).f n = β.f n :=
begin
  dsimp only [desc],
  simp only [cocycle.hom_of_f, desc_cocycle_coe, cocycle.of_hom_coe,
    inr_desc_cochain_v, cochain.of_hom_v],
end

@[simp, reassoc]
lemma inr_desc {K : cochain_complex C ℤ} (α : cochain F K (-1)) (β : G ⟶ K)
  (eq : hom_complex.δ (-1) 0 α = cochain.of_hom (φ ≫ β)) :
  inr φ ≫ desc φ α β eq = β :=
begin
  dsimp only [desc],
  ext n,
  simp only [homological_complex.comp_f, cocycle.hom_of_f, desc_cocycle_coe,
    cocycle.of_hom_coe, inr_desc_cochain_v, cochain.of_hom_v],
end

lemma desc_f {K : cochain_complex C ℤ} (α : cochain F K (-1)) (β : G ⟶ K)
  (eq : hom_complex.δ (-1) 0 α = cochain.of_hom (φ ≫ β)) (n n' : ℤ) (hn' : n' = n+1) :
  (desc φ α β eq).f n =
    (fst φ : cochain (mapping_cone φ) F 1).v n n' hn' ≫
      α.v n' n (by { rw [hn', int.add_neg_one, add_tsub_cancel_right]}) +
      (snd φ).v n n (add_zero n).symm ≫ β.f n :=
by simp only [from_ext_iff _ _ _ hn', add_zero, inl_desc_v, preadditive.comp_add,
  inl_fst_assoc, inl_snd_assoc, zero_comp, eq_self_iff_true, inr_desc_f,
  inr_fst_assoc, inr_snd_assoc, zero_add, and_self]

def desc_homotopy {K : cochain_complex C ℤ} (f₁ f₂ : mapping_cone φ ⟶ K)
  (γ₁ : cochain F K (-2)) (γ₂ : cochain G K (-1))
  (h₁ : (inl φ).comp (cochain.of_hom f₁) (add_zero (-1)).symm =
    hom_complex.δ (-2) (-1) γ₁ + (cochain.of_hom φ).comp γ₂ (zero_add _).symm +
    (inl φ).comp (cochain.of_hom f₂) (add_zero (-1)).symm)
  (h₂ : cochain.of_hom (inr φ ≫ f₁) =
    hom_complex.δ (-1) 0 γ₂ + cochain.of_hom (inr φ ≫ f₂)) :
  homotopy f₁ f₂ :=
(equiv_homotopy _ _).symm
begin
  refine ⟨desc_cochain _ γ₁ γ₂ (by linarith), _⟩,
  rw [δ_desc_cochain φ γ₁ γ₂ (by linarith) (neg_add_self 1),
    cochain_ext _ _ (show (0 : ℤ) = -1 +1 , by linarith)],
  split,
  { rw [cochain.comp_add, h₁],
    nth_rewrite 0 cochain.comp_add,
    simp only [← cochain.comp_assoc _ _ _ (neg_add_self 1).symm (add_neg_self 1).symm
        (show (-1 : ℤ) = (-1) +1 + (-1), by linarith), inl_comp_fst, cochain.id_comp,
        neg_add_self, ε_0, one_smul, ← cochain.comp_assoc_of_second_is_zero_cochain,
        inl_comp_snd, cochain.zero_comp, add_zero], },
  { rw [cochain.comp_add, ← cochain.of_hom_comp, ← cochain.of_hom_comp, h₂],
    nth_rewrite 0 cochain.comp_add,
    simp only [← hom_complex.cochain.comp_assoc_of_first_is_zero_cochain,
      inr_comp_fst, cochain.zero_comp, zero_add, inr_comp_snd,
      cochain.id_comp], },
end

def lift_cochain {K : cochain_complex C ℤ}
  {n m : ℤ} (α : cochain K F m) (β : cochain K G n) (h : n+1=m) :
  cochain K (mapping_cone φ) n :=
α.comp (inl φ) (by linarith) + β.comp (cochain.of_hom (inr φ)) (by linarith)

@[simp, reassoc]
lemma lift_cochain_fst_v {K : cochain_complex C ℤ}
  {n m : ℤ} (α : cochain K F m) (β : cochain K G n) (h : n+1=m) (p₁ p₂ p₃ : ℤ)
  (h₁₂ : p₂ = p₁ + n) (h₂₃ : p₃ = p₂ + 1) :
  (lift_cochain φ α β h).v p₁ p₂ h₁₂ ≫ (fst φ : cochain (mapping_cone φ) F 1).v p₂ p₃ h₂₃ =
    α.v p₁ p₃ (by rw [h₂₃, h₁₂, ← h, add_assoc])  :=
begin
  dsimp only [lift_cochain],
  simp only [cochain.add_v, add_zero, cochain.comp_zero_cochain, cochain.of_hom_v,
    subtype.val_eq_coe, preadditive.add_comp, assoc, inr_fst, comp_zero,
    cochain.comp_v _ _ (show n = m+(-1), by linarith) p₁ p₃ p₂ (by linarith) (by linarith),
    inl_fst, comp_id],
end

@[simp, reassoc]
lemma lift_cochain_snd_v {K : cochain_complex C ℤ}
  {n m : ℤ} (α : cochain K F m) (β : cochain K G n) (h : n+1=m)
    (p₁ p₂ : ℤ) (h₁₂ : p₂ = p₁ + n) :
  (lift_cochain φ α β h).v p₁ p₂ h₁₂ ≫ (snd φ).v p₂ p₂ (add_zero p₂).symm =
    β.v p₁ p₂ h₁₂ :=
begin
  dsimp [lift_cochain],
  simp only [cochain.comp_zero_cochain, cochain.of_hom_v, preadditive.add_comp, assoc,
    cochain.comp_v _ _ (show n = m+(-1), by linarith) p₁ (p₁+m) p₂ rfl (by linarith),
    inr_snd, comp_id, add_left_eq_self, inl_snd, comp_zero],
end

@[simp]
lemma lift_cochain_fst {K : cochain_complex C ℤ}
  {n m : ℤ} (α : cochain K F m) (β : cochain K G n) (h : n+1=m)  :
  (lift_cochain φ α β h).comp (fst φ : cochain (mapping_cone φ) F 1) h.symm = α :=
begin
  ext p q hpq,
  simp only [cochain.comp_v _ _ h.symm p (p+n) q rfl (by linarith), lift_cochain_fst_v],
end

@[simp]
lemma lift_cochain_snd {K : cochain_complex C ℤ}
  {n m : ℤ} (α : cochain K F m) (β : cochain K G n) (h : n+1=m)  :
  (lift_cochain φ α β h).comp (snd φ) (add_zero n).symm = β :=
begin
  ext p q hpq,
  simp only [cochain.comp_zero_cochain, lift_cochain_snd_v],
end

lemma δ_lift_cochain {K : cochain_complex C ℤ}
  {n m : ℤ} (α : cochain K F m) (β : cochain K G n) (h : n+1=m) (m' : ℤ) (hm' : m = m'+(-1)) :
  hom_complex.δ n m (lift_cochain φ α β h) =
    -(hom_complex.δ m m' α).comp (inl φ) hm' +
    (hom_complex.δ n m β + α.comp (cochain.of_hom φ) (add_zero m).symm).comp
      (cochain.of_hom (inr φ)) (add_zero m).symm :=
begin
  ext p q hpq,
  simp only [to_ext_iff _ _ (q+1) rfl, δ_v n m h _ p q hpq _ _ rfl rfl, cochain.add_v,
    cochain.comp_v _ _ hm' p (q+1) q (by linarith) (by linarith),
    δ_v m m' (by linarith) _ p  (q+1) (by linarith) q (p+1) (by linarith) rfl,
    cochain.neg_v, cochain.comp_zero_cochain, cochain.of_hom_v,
    preadditive.add_comp, assoc, preadditive.zsmul_comp, lift_cochain_fst_v, inl_fst, inr_fst,
    preadditive.neg_comp, preadditive.comp_neg, comp_zero, smul_zero, add_zero,
    d_fst φ (q-1) q (q+1) (by linarith) rfl, lift_cochain_fst_v_assoc, comp_id, neg_add, h,
    ε_succ, neg_smul, neg_neg, inl_snd, neg_zero, zero_add, d_snd φ (q-1) q (by linarith),
    preadditive.comp_add, lift_cochain_snd_v_assoc, inr_snd, lift_cochain_snd_v],
  refine ⟨rfl, _⟩,
  have : ∀ (x y z : K.X p ⟶ G.X q), x +y +z = y+z +x := λ x y z, by abel,
  apply this,
end

def lift_cocycle {K : cochain_complex C ℤ}
  {n m : ℤ} (α : cocycle K F m) (β : cochain K G n) (h : n+1=m)
  (hαβ : hom_complex.δ n m β + (α : cochain K F m).comp (cochain.of_hom φ) (add_zero m).symm = 0) :
  cocycle K (mapping_cone φ) n :=
cocycle.mk (lift_cochain φ (α : cochain K F m) β h) _ h
  (by simp only [δ_lift_cochain φ _ _ h (m+1) (by linarith), hαβ, cochain.zero_comp, add_zero,
    cocycle.δ_eq_zero, neg_zero])

@[simp]
def lift_cocycle_coe {K : cochain_complex C ℤ}
  {n m : ℤ} (α : cocycle K F m) (β : cochain K G n) (h : n+1=m)
  (hαβ : hom_complex.δ n m β + (α : cochain K F m).comp (cochain.of_hom φ) (add_zero m).symm = 0) :
  (lift_cocycle φ α β h hαβ : cochain K (mapping_cone φ) n) =
    lift_cochain φ (α : cochain K F m) β h := rfl

def lift {K : cochain_complex C ℤ} (α : cocycle K F 1) (β : cochain K G 0)
  (hαβ : hom_complex.δ 0 1 β + (α : cochain K F 1).comp (cochain.of_hom φ) (add_zero 1).symm = 0) :
   K ⟶ mapping_cone φ :=
cocycle.hom_of (lift_cocycle φ α β (zero_add 1) hαβ)

@[simp, reassoc]
lemma lift_fst_f {K : cochain_complex C ℤ} (α : cocycle K F 1) (β : cochain K G 0)
  (hαβ : hom_complex.δ 0 1 β + (α : cochain K F 1).comp (cochain.of_hom φ) (add_zero 1).symm = 0)
  (n n' : ℤ) (hnn' : n' = n+1) :
    (lift φ α β hαβ).f n ≫
      (fst φ : cochain (mapping_cone φ) F 1).v n n' hnn' = (α : cochain K F 1).v n n' hnn' :=
begin
  dsimp only [lift],
  simp only [cocycle.hom_of_f, lift_cocycle_coe, lift_cochain_fst_v],
end

@[simp]
lemma lift_fst {K : cochain_complex C ℤ} (α : cocycle K F 1) (β : cochain K G 0)
  (hαβ : hom_complex.δ 0 1 β + (α : cochain K F 1).comp (cochain.of_hom φ) (add_zero 1).symm = 0) :
  (cochain.of_hom (lift φ α β hαβ)).comp
    (fst φ : cochain (mapping_cone φ) F 1) (zero_add 1).symm =
      (α : cochain K F 1) :=
begin
  ext p q hpq,
  simp only [cochain.zero_cochain_comp, cochain.of_hom_v, lift_fst_f],
end

@[simp, reassoc]
lemma lift_snd_f {K : cochain_complex C ℤ} (α : cocycle K F 1) (β : cochain K G 0)
  (hαβ : hom_complex.δ 0 1 β + (α : cochain K F 1).comp (cochain.of_hom φ) (add_zero 1).symm = 0) (n : ℤ) :
  (lift φ α β hαβ).f n ≫ (snd φ).v n n (add_zero n).symm =
    β.v n n (add_zero n).symm :=
begin
  dsimp only [lift],
  simp only [cocycle.hom_of_f, lift_cocycle_coe, lift_cochain_snd_v],
end

@[simp]
lemma lift_snd {K : cochain_complex C ℤ} (α : cocycle K F 1) (β : cochain K G 0)
  (hαβ : hom_complex.δ 0 1 β + (α : cochain K F 1).comp (cochain.of_hom φ) (add_zero 1).symm = 0) :
  (cochain.of_hom (lift φ α β hαβ)).comp
    (snd φ) (add_zero 0).symm = β :=
begin
  dsimp only [lift],
  simp only [cocycle.cochain_of_hom_hom_of_eq_coe, lift_cocycle_coe, lift_cochain_snd],
end

lemma lift_desc_f {K L : cochain_complex C ℤ} (α : cocycle K F 1) (β : cochain K G 0)
  (hαβ : hom_complex.δ 0 1 β + (α : cochain K F 1).comp (cochain.of_hom φ) (add_zero 1).symm = 0)
  (α' : cochain F L (-1)) (β' : G ⟶ L) (eq : hom_complex.δ (-1) 0 α' = cochain.of_hom (φ ≫ β'))
  (n n' : ℤ) (hnn' : n' = n+1) :
  (lift φ α β hαβ).f n ≫ (desc φ α' β' eq).f n =
    (α : cochain K F 1).v n n' hnn' ≫ α'.v n' n (by { rw [hnn', int.add_neg_one, add_tsub_cancel_right], }) +
      β.v n n (add_zero n).symm ≫ β'.f n :=
begin
  rw [← id_comp ((desc φ α' β' eq).f n), ← id φ _ _ hnn'],
  simp only [preadditive.add_comp, assoc, inl_desc_v, inr_desc_f, preadditive.comp_add,
    lift_fst_f_assoc, lift_snd_f_assoc],
end

lemma lift_f {K : cochain_complex C ℤ} (α : cocycle K F 1) (β : cochain K G 0)
  (hαβ : hom_complex.δ 0 1 β + (α : cochain K F 1).comp (cochain.of_hom φ) (add_zero 1).symm = 0) (n n' : ℤ)
    (hn' : n' = n+1) :
    (lift φ α β hαβ).f n = (α : cochain K F 1).v n n' hn' ≫
      (inl φ).v n' n (by rw [hn', int.add_neg_one, add_tsub_cancel_right]) +
    β.v n n (add_zero n).symm ≫ (inr φ).f n :=
by simp only [to_ext_iff _ _ _ hn', add_zero, lift_fst_f, preadditive.add_comp, assoc,
  inl_fst, comp_id, inr_fst, comp_zero, eq_self_iff_true, lift_snd_f, inl_snd,
  inr_snd, zero_add, and_self]

def lift_homotopy {K : cochain_complex C ℤ} (f₁ f₂ : K ⟶ mapping_cone φ)
  (γ₁ : cochain K F 0) (γ₂ : cochain K G (-1))
  (h₁ : (cochain.of_hom f₁).comp (fst φ :
    cochain (mapping_cone φ) F 1) (zero_add 1).symm = -hom_complex.δ 0 1 γ₁ +
      (cochain.of_hom f₂).comp (fst φ : cochain (mapping_cone φ) F 1) (zero_add 1).symm)
  (h₂ : (cochain.of_hom f₁).comp (snd φ) (add_zero 0).symm =
    hom_complex.δ (-1) 0 γ₂ + γ₁.comp (cochain.of_hom φ) (zero_add 0).symm +
    (cochain.of_hom f₂).comp (snd φ) (add_zero 0).symm) :
  homotopy f₁ f₂ :=
(equiv_homotopy _ _).symm
begin
  refine ⟨lift_cochain φ γ₁ γ₂ (neg_add_self 1), _⟩,
  simp only [δ_lift_cochain φ _ _ _ 1 (show (0 : ℤ) = 1 +(-1), by linarith),
    cochain_ext' _ _ (zero_add 1).symm],
  split,
  { simp only [add_zero, cochain.add_comp, cochain.neg_comp,
      cochain.comp_assoc_of_second_is_zero_cochain, inr_comp_fst,
      cochain.comp_zero,
      cochain.comp_assoc _ _ _ (add_neg_self 1).symm (neg_add_self 1).symm
      (show (1 : ℤ) = 1+(-1)+1, by linarith),
      inl_comp_fst, cochain.comp_id, h₁], },
  { simp only [zero_add, neg_zero, cochain.add_comp, cochain.comp_assoc_of_third_is_zero_cochain,
      cochain.neg_comp, inl_comp_snd, cochain.comp_zero, inr_comp_snd, cochain.comp_id, h₂], },
end

section

variables {K₁ K₂ L₁ L₂ : cochain_complex C ℤ}
  [∀ p, has_binary_biproduct (K₁.X (p+1)) (L₁.X p)]
  [∀ p, has_binary_biproduct (K₂.X (p+1)) (L₂.X p)]
  (f₁ : K₁ ⟶ L₁) (f₂ : K₂ ⟶ L₂) (τ₁ : K₁ ⟶ K₂) (τ₂ : L₁ ⟶ L₂) (comm : f₁ ≫ τ₂ = τ₁ ≫ f₂)

include comm

def map : mapping_cone f₁ ⟶ mapping_cone f₂ :=
desc f₁ ((cochain.of_hom τ₁).comp (inl f₂) (zero_add _).symm)
  (τ₂ ≫ inr f₂)
begin
  rw [δ_comp_of_first_is_zero_cochain _ _ _ (neg_add_self 1), δ_inl,
    cocycle.δ_cochain_of_hom, cochain.zero_comp, smul_zero, add_zero, cochain.of_hom_comp f₂,
    ← assoc f₁, ← cochain.of_hom_comp, ← cochain.of_hom_comp, ← assoc, comm],
end

lemma inr_comp_map :
  inr f₁ ≫ map _ _ _ _ comm =
    τ₂ ≫ inr f₂ :=
begin
  apply hom_complex.cochain.of_hom_injective,
  rw cochain_ext' _ _ (zero_add 1).symm,
  dsimp only [map],
  split,
  { simp only [inr_desc, cochain.of_hom_comp,
      cochain.comp_assoc_of_second_is_zero_cochain, inr_comp_fst,
      inr_fst], },
  { simp only [inr_desc, cochain.of_hom_comp, inr_snd,
      cochain.comp_assoc_of_third_is_zero_cochain, inr_comp_snd], },
end

lemma map_comp_δ :
  map _ _ _ _ comm ≫ δ f₂ =
  δ f₁ ≫ τ₁⟦1⟧' :=
begin
  apply hom_complex.cochain.of_hom_injective,
  rw cochain_ext _ _(neg_add_self 1).symm,
  dsimp only [map],
  split,
  { simp only [cochain.of_hom_comp, ← hom_complex.cochain.comp_assoc_of_second_is_zero_cochain,
      inl_desc, hom_complex.cochain.comp_assoc_of_first_is_zero_cochain,
      inl_δ, cochain.comp_neg, cochain.of_hom_comp],
    ext p q hpq,
    have hp : p = q+1 := by linarith,
    subst hp,
    simp only [cochain.neg_v, cochain.zero_cochain_comp, cochain.of_hom_v,
      cochain.neg_comp, cochain.comp_zero_cochain, shift_functor_map_f', neg_inj],
    erw cochain.right_shift_v (cochain.of_hom _) 1 (-1)
      (by linarith) (q+1) q (by linarith) (q+1) (by linarith),
    erw cochain.right_shift_v (cochain.of_hom _) 1 (-1)
      (by linarith) (q+1) q (by linarith) (q+1) (by linarith),
    simp only [shift_functor_obj_X_iso, cochain.of_hom_v, homological_complex.id_f,
      homological_complex.X_iso_of_eq_refl, id_comp],
    dsimp [iso.refl],
    rw [comp_id, id_comp], },
  { rw [cochain.of_hom_comp, ← hom_complex.cochain.comp_assoc_of_first_is_zero_cochain,
      ← cochain.of_hom_comp, inr_desc, ← cochain.of_hom_comp, assoc,
      inr_δ, comp_zero, cochain.of_hom_zero, ← cochain.of_hom_comp, ← assoc,
      inr_δ, zero_comp, cochain.of_hom_zero], },
end

end

example : ℕ := 42

section

variables {K L : cochain_complex C ℤ} (f : K ⟶ L) {D : Type*} [category D] [preadditive D]
  [∀ p, has_binary_biproduct (K.X (p+1)) (L.X p)] (Φ : C ⥤ D) [functor.additive Φ]
  [∀ p, has_binary_biproduct (((Φ.map_homological_complex (complex_shape.up ℤ)).obj K).X (p + 1))
    (((Φ.map_homological_complex (complex_shape.up ℤ)).obj  L).X p)]

@[simps]
def map_iso : (Φ.map_homological_complex _).obj (mapping_cone f) ≅
  mapping_cone ((Φ.map_homological_complex _).map f) :=
{ hom := mapping_cone.lift _ (cocycle.map (mapping_cone.fst f) Φ)
    ((mapping_cone.snd f).map Φ) (by simp),
  inv := mapping_cone.desc _ ((mapping_cone.inl f).map Φ)
      ((Φ.map_homological_complex _).map (mapping_cone.inr f)) (by simp),
  hom_inv_id' := begin
    ext n,
    simpa only [homological_complex.comp_f, homological_complex.id_f,
      lift_desc_f _ _ _ _ _ _ _ n (n+1) rfl, cocycle.map_coe, cochain.map_v,
      functor.map_homological_complex_map_f, ← Φ.map_comp, ← Φ.map_add,
      mapping_cone.id, Φ.map_id],
  end,
  inv_hom_id' := hom_complex.cochain.of_hom_injective begin
    ext n,
    simp only [cochain.of_hom_comp, cochain.comp_zero_cochain, cochain.of_hom_v,
      homological_complex.id_f, from_ext_iff _ _ (n+1) rfl, to_ext_iff _ _ (n+1) rfl,
      assoc, lift_fst_f, cocycle.map_coe, cochain.map_v, inl_desc_v_assoc, id_comp,
      inl_fst, inr_desc_f_assoc, functor.map_homological_complex_map_f, inr_fst,
      lift_snd_f, inl_snd, inr_snd, ← Φ.map_comp, Φ.map_zero, Φ.map_id],
    tauto,
  end, }

end

end mapping_cone

end preadditive

section abelian

open hom_complex

variables [abelian C] {S : short_complex (cochain_complex C ℤ)} (ex : S.short_exact)

include ex

lemma degreewise_exact (n : ℤ) :
  (S.map (homological_complex.eval C (complex_shape.up ℤ) n)).short_exact :=
ex.map_of_exact (homological_complex.eval C (complex_shape.up ℤ) n)

def from_mapping_cone_of_ses : mapping_cone S.f ⟶ S.X₃ :=
mapping_cone.desc S.f 0 S.g (by simp)

@[simp, reassoc]
lemma inr_from_mapping_cone_of_ses (n : ℤ) :
  (mapping_cone.inr S.f).f n ≫ (from_mapping_cone_of_ses ex).f n = S.g.f n :=
begin
  dsimp only [from_mapping_cone_of_ses],
  simp only [mapping_cone.inr_desc_f],
end

@[simp, reassoc]
lemma inl_from_mapping_cone_of_ses (p q : ℤ) (hpq : q = p + (-1)) :
  (mapping_cone.inl S.f).v p q hpq ≫ (from_mapping_cone_of_ses ex).f q = 0 :=
begin
  dsimp only [from_mapping_cone_of_ses],
  simp only [mapping_cone.inl_desc_v, cochain.zero_v],
end

@[simp, reassoc]
lemma inr_mapping_cone_comp_from_mapping_cone_of_ses :
  mapping_cone.inr S.f ≫ from_mapping_cone_of_ses ex = S.g :=
begin
  ext n : 2,
  simp only [homological_complex.comp_f, inr_from_mapping_cone_of_ses],
end

instance from_mapping_cone_of_ses_quasi_iso : quasi_iso (from_mapping_cone_of_ses ex) :=
⟨λ n, begin
  rw is_iso_homology_map_iff_short_complex_quasi_iso'
    (from_mapping_cone_of_ses ex) (show (n-1)+1=n, by linarith) rfl,
  change is_iso _,
  haveI : ∀ (n : ℤ), mono (S.f.f n) :=
    λ n, (ex.map_of_exact (homological_complex.eval _ _ n)).mono_f,
  rw is_iso_iff_mono_and_epi,
  split,
  { rw short_complex.mono_homology_map_iff,
    dsimp,
    intros A x₂ hxy z hz,
    obtain ⟨x, y, rfl⟩ := mapping_cone.to_decomposition x₂ _ rfl,
    simp only [preadditive.add_comp, assoc, mapping_cone.inr_d, preadditive.comp_sub,
      mapping_cone.inl_d S.f n (n+1) (n+1+1) (by linarith) (by linarith)] at hxy,
    obtain ⟨hx, hy⟩ := (mapping_cone.to_ext_iff _ _ _ rfl).mp hxy,
    simp only [preadditive.add_comp, preadditive.sub_comp, assoc, mapping_cone.inr_fst,
      comp_zero, mapping_cone.inl_fst, comp_id, zero_sub, add_zero, zero_comp, neg_eq_zero] at hx,
    simp only [preadditive.add_comp, preadditive.sub_comp, assoc, mapping_cone.inr_snd, comp_id,
      mapping_cone.inl_snd, comp_zero, sub_zero, zero_comp, ← eq_neg_iff_add_eq_zero] at hy,
    clear hxy,
    simp only [preadditive.add_comp, assoc, inr_from_mapping_cone_of_ses,
      inl_from_mapping_cone_of_ses, comp_zero, zero_add] at hz,
    haveI : epi (S.g.f (n-1)) := (ex.map_of_exact (homological_complex.eval _ _ _)).epi_g,
    obtain ⟨A', π, hπ, z', hz'⟩ := abelian.pseudo_surjective_of_epi' (S.g.f (n-1)) z,
    have ex' := (ex.map_of_exact (homological_complex.eval _ _ n)),
    haveI := ex'.mono_f,
    let w : A' ⟶ S.X₁.X n := ex'.exact.lift (π ≫ y - z' ≫ S.X₂.d _ _) begin
      dsimp,
      simp only [preadditive.sub_comp, assoc, hz, reassoc_of hz',
        homological_complex.hom.comm, sub_self],
    end,
    have hw : w ≫ S.f.f n = _ := ex'.exact.lift_f _ _,
    refine ⟨A', π, hπ, w ≫ (mapping_cone.inl S.f).v n (n-1) (show n-1 = n+(-1), by refl) + z' ≫ (mapping_cone.inr S.f).f (n-1),
      (mapping_cone.to_ext_iff _ _ _ rfl).mpr ⟨_, _⟩⟩,
    { simp only [assoc, preadditive.add_comp, mapping_cone.inr_fst, comp_zero, add_zero,
        mapping_cone.inl_fst, comp_id, mapping_cone.inr_d_assoc,
        mapping_cone.inl_d_assoc S.f (n-1) n (n+1) (by refl) (by linarith),
        preadditive.sub_comp, preadditive.comp_sub, ← cancel_mono (S.f.f (n+1)), zero_comp],
      simp only [← S.f.comm, reassoc_of hw, preadditive.sub_comp, assoc, homological_complex.d_comp_d,
        comp_zero, sub_zero, zero_sub, hy, preadditive.comp_neg], },
    { simp only [assoc, preadditive.comp_add, preadditive.add_comp, mapping_cone.inl_snd, comp_zero,
        zero_add, mapping_cone.inr_snd, comp_id, mapping_cone.inr_d_assoc, preadditive.comp_sub,
        preadditive.sub_comp, hw,
        mapping_cone.inl_d S.f (n-1) n (n+1) (show n-1 = n+(-1), by refl) (by linarith)],
        abel, }, },
  { rw short_complex.epi_homology_map_iff,
    dsimp,
    intros A z hz,
    haveI : epi (S.g.f n) := (ex.map_of_exact (homological_complex.eval _ _ _)).epi_g,
    obtain ⟨A', π, hπ, y, hy⟩ := abelian.pseudo_surjective_of_epi' (S.g.f n) z,
    have ex' := (ex.map_of_exact (homological_complex.eval _ _ (n+1))),
    haveI := ex'.mono_f,
    let x : A' ⟶ S.X₁.X (n+1) := ex'.exact.lift (y ≫ S.X₂.d _ _) begin
      dsimp,
      simp only [assoc, ← S.g.comm, ← reassoc_of hy, hz, comp_zero],
    end,
    have hx : x ≫ S.f.f (n+1) = _ := ex'.exact.lift_f _ _,
    have hdx : x ≫ S.X₁.d (n+1) (n+1+1) = 0,
    { simp only [← cancel_mono (S.f.f (n+1+1)), assoc, zero_comp, ← S.f.comm, reassoc_of hx,
        homological_complex.d_comp_d, comp_zero], },
    refine ⟨A', π, hπ, y ≫ (mapping_cone.inr S.f).f n -
      x ≫ (mapping_cone.inl S.f).v (n+1) n (show n = (n+1)+(-1), by linarith), _, _⟩,
    { simp only [preadditive.sub_comp, assoc, mapping_cone.inr_d, ← reassoc_of hx,
        mapping_cone.inl_d S.f n (n+1) (n+1+1) (by linarith) (by linarith), preadditive.comp_sub,
        reassoc_of hdx, zero_comp, sub_zero, sub_self], },
    { exact ⟨0, by simp only [hy, preadditive.sub_comp, assoc, inr_from_mapping_cone_of_ses,
        inl_from_mapping_cone_of_ses, comp_zero, sub_zero, zero_comp, add_zero]⟩, }, },
end⟩

end abelian

end cochain_complex


import Architectural.proofObligations
import Architectural.lang
import tactic

open LANG PORTS tactic

local infix ` OR `:50 := LANG.disj
local infix ` & `:50 := LANG.conj
variables {Var : Type} [fintype Var] [decidable_eq Var] [reflected Var]

theorem asdff (x : Trace PORTS) {A : LANG} : x ∈ (@AssertionLang.sem LANG PORTS _ _ _ A) ↔ x ∈ LANG.sem A := by {refl,}

theorem fsadfdsa (x : Trace PORTS) (A B : LANG) : x ∈ LANG.sem (A & B).always ↔ x ∈ LANG.sem (A).always ∧ x ∈ LANG.sem (B).always :=
begin
simp, rw forall_and_distrib,
end

theorem disj_comm_guarded (x : Trace PORTS) (A B : LANG) : x ∈ LANG.sem (A OR B).always ↔ x ∈ LANG.sem (B OR A).always :=
begin
simp,
end
theorem conj_comm_guarded (x : Trace PORTS) (A B : LANG) : x ∈ LANG.sem (A & B).always ↔ x ∈ LANG.sem (B & A).always :=
begin
split, all_goals {intros h i, replace h := h i, cases h, split, assumption, assumption,},
end

theorem fasdjbasdf {x : Trace PORTS} (A B : LANG) : x ∈ LANG.sem (A & B) ↔ x ∈ LANG.sem A ∧ x ∈ LANG.sem B :=
begin
rw conj_def, simp,
end

theorem foobars {x : Trace PORTS} (A B C D : LANG) :
                            x ∈ LANG.sem ((A & B & C).always & D.always)
              ↔
                            x ∈ LANG.sem (((A & B).always &C.always) & D.always)  :=
begin
repeat {rw fasdjbasdf,  rw fsadfdsa,},
end

theorem portation_left {x : Trace PORTS} (A B : LANG) : (x ∈ @AssertionLang.sem LANG PORTS _ _ _ (A & B).always) → x ∈ (A).always.sem :=
begin
intro h, intro i, rw forall_conj_distrib_mem at h, cases h with h1 h2, apply h1 i,
end

theorem portation_right {x : Trace PORTS} (A B : LANG) : (x ∈ @AssertionLang.sem LANG PORTS _ _ _ (A & B).always) → x ∈ (B).always.sem :=
begin
intro h, intro i, rw forall_conj_distrib_mem at h, cases h with h1 h2, apply h2 i,
end

------

meta structure model_info :=
(del : list (expr × expr))
(comps : list expr)

meta def preprocess_rpo_fst : tactic (name × name) := do
  x ← tactic.get_unused_name `x,
  H ← tactic.get_unused_name `H,
  tactic.intro x,
  `[simp],
  `[rw AssertionLang.impl_def],
  tactic.intro H,
  `[rw [list_conj_iff, get_nfs] at H, simp at H, rw toMap at H],
  return ⟨x, H⟩

meta def decompose_conj : expr →  expr → tactic (unit)
| h `(and %%left %%right) := do
    left_nm ← get_unused_name,
    right_nm ← get_unused_name,
    cases_core h [left_nm, right_nm] transparency.semireducible,
    right_e ← get_local right_nm,
    decompose_conj right_e right
| _ _ := do return ()

meta def decompose_conj_left : expr →  expr → tactic (unit)
| h `(and %%left %%right) := do
    left_nm ← get_unused_name,
    right_nm ← get_unused_name,
    cases_core h [left_nm, right_nm] transparency.semireducible,
    left_e ← get_local left_nm,
    decompose_conj left_e left
| _ _ := do return ()

meta def collect_assertions : list expr → tactic (list expr)
| (h::t ) := do
  τ ← infer_type h,
  match τ with
  | `(_ ∈ _) := do
                l ← collect_assertions t,
                return $ [h]++l
  | `(_ ∈ _ → _) := do
                l ← collect_assertions t,
                return $ [h]++l
  | _ := collect_assertions t
  end
| [] := return []

theorem  IFf_pos (c : Prop) (H : decidable c) : ∀ {α : Type} {t e  :α}, c → ite c t e = t := by {intros a b c h, apply if_pos h,}
theorem  IFf_neg  : Π (c : Prop) [H : decidable c], ∀ {α : Type} {t e  :α}, ¬c → (@ite _ c H t e) = e := by {intros a b c h a f, apply if_neg, exact f}

meta def get_conds_aux : expr → tactic (list expr)
| `(@ite %%A %%B %%C %%D %%E) := do
        let cond := [B],
        r ← get_conds_aux E,
        return $ cond ++ r
| _ := return []

meta def get_conds : expr → tactic (list expr)
| `(%%x ∈ AssertionLang.sem (%%E)) :=
  match E with
    | `(Contract.nf (option.iget %%Y)) := get_conds_aux Y
    | `(Contract.A (option.iget %%Y)) := get_conds_aux Y
    | _ := return []
  end
| `(_ → %%x ∈ AssertionLang.sem (%%E)) :=
  match E with
    | `(Contract.nf (option.iget %%Y)) := get_conds_aux Y
    | _ := return []
  end
| x :=  return []

meta def GFSAUIHasd : list expr → tactic (list (list expr))
| (h::t) := do
    -- tactic.trace "here",
    -- τ ← infer_type h,
    -- tactic.trace τ,
    l← get_conds h, t ← GFSAUIHasd t, return (l::t)
| _ := do return []

meta def aiuhsfd : list expr → list (expr × bool)
| [] :=  []
| (h::t) := match h with
           | `(eq %%A %%B) :=
              if B = A then  ([(h,tt)]++(aiuhsfd t)) else ([(h,ff)]++(aiuhsfd t))
           | _ :=  []
          end

meta def sjdioaf (hyp : expr) : list (expr × bool) → tactic unit
| [] := return ()
| ((e, b)::t) := do
                  if b then do
                  let ef := expr.mk_app `(IFf_pos) [e],
                  rewrite_hyp ef hyp, return  ()
                  else do
                  sjdioaf t

meta def suihfa : list (expr × list (expr × bool)) → tactic unit
| [] := return ()
| ((e, l)::xs) := do a ← sjdioaf e l,
                  suihfa xs

meta def extract_contracts : expr →  tactic unit := λ h, do
  infer_type h >>= decompose_conj h, dedup,
  tactic.repeat `[rw Map.find_val at *],
  xs ← local_context >>= collect_assertions,
  l ← xs.mmap infer_type >>= GFSAUIHasd,
  let ls := xs.zip (l.map aiuhsfd),
  suihfa ls,
  iterate `[rw if_neg at *],
  any_goals `[dec_trivial],
  `[rw nf_def at *, simp at *],
  return ()

meta def collect_consequents : list (expr × expr) → (list (expr × expr))
| []   := []
| ((h, τ)::t) := match τ with
                | `(%%P → %%Q) := (collect_consequents t).cons (h,Q)
                | _ := collect_consequents t
                end

meta def collect_vars_aux : expr → list expr
| `(LANG.atom %%a) := [a]
| `(LANG.lt %%a %%v) := [a,v]
| `(LANG.neg %%A) := collect_vars_aux A
| `(LANG.conj %%A %%B) := collect_vars_aux A ++ collect_vars_aux B
| `(LANG.disj %%A %%B) := collect_vars_aux A ++ collect_vars_aux B
| `(LANG.always %%A) := collect_vars_aux A
|_ := []

meta def collect_vars : expr → list expr
| `(_ ∈ AssertionLang.sem %%P) := collect_vars_aux P
| _ := []

meta def find_matches_aux : expr → list (expr × expr) → list (expr × expr)
| h ((a,b)::l) := if h=a then [(b,h)] else if h=b then [(a,h)]else (find_matches_aux h l)
| _ _ := []

meta def find_matches (del : list (expr × expr))  : list expr → list (expr × expr)
| [] := []
| (h::t) := find_matches_aux h del++(find_matches t )
-- find_matches_aux es del

open tactic.interactive («have»)

meta def mk_sync_apply (x : expr) : (expr × expr) →  tactic unit
| (a,b) := do
 -- ((synchronize' x fault_PWMFlow_LACU fault_armFlow_armController) (atom fault_PWMFlow_LACU).neg.always).mpr
 `(_ ∈ AssertionLang.sem %%FOO) ← target,
 let e := expr.mk_app `(synchronize') [x, b, a, FOO],
--  tactic.trace e,
 τ ← infer_type e,
 match τ with
 | `(%%GOAL → %%FOO) := do
   h ← get_unused_name `h,
   «have» h ``(%%GOAL) ``(by dec_trivial),
   h_e ← get_local h,
   let e' := expr.mk_app e [h_e],
    τ ← infer_type e',
   match τ  with
   | `(%%l ↔ %%r) := do
   let e'' := `(@iff.mpr %%l %%r %%e'),
   apply e'',
   `[simp],
   clear h_e,
   return ()
   | _ :=  return ()
   end
 | _ := return ()
 end

meta def try_application : list expr → tactic unit
| [] := return ()
| (h::t) :=  (do apply h, return ()) <|> try_application t

meta def synchronize_ports (x : name) (Γ : model_info): tactic unit := do
  tgt ← target,
  let cs := collect_vars tgt,
  let ls := find_matches Γ.del cs, -- vars to sub will be on LEFT
  x ← get_local x,
  ls.mmap (mk_sync_apply x),
  return ()

meta def decompose_env_assumptions (x_nm : name) : tactic unit := do
  A ← get_unused_name `A,
  intro A,
  A_e ← get_local A, x ← get_local x_nm,
  let e := expr.mk_app `(asdff) [x],
  tactic.rewrite_hyp e A_e,
  let e' := expr.mk_app `(fsadfdsa) [x],
  tactic.repeat (do A_e ← get_local A,tactic.rewrite_hyp e' A_e, return ()),
  A_e ← get_local A,
  tactic.repeat (do A_e ← get_local A, cases_core A_e [A] transparency.semireducible, return ()),
  return ()

meta def try_rewrites : expr → tactic unit
| `(_ ∈ AssertionLang.sem _) := do `[rw asdff],  tactic.repeat (`[rw fsadfdsa]), return ()
| _ := do tactic.repeat (`[rw fsadfdsa]), return ()

meta def check_for_portation_aux (e : expr) : list expr → tactic unit
| (h::t) := match h with
            | `(_ → (_ ∈ AssertionLang.sem (%%A & %%B).always)) :=
              if A = e then do `[apply portation_left], return ()
              else if B = e then do `[apply portation_right], return ()
              else check_for_portation_aux t
            | _ := check_for_portation_aux t
            end
| [] := return ()

meta def check_for_portation_aux' (e : expr) : list expr → tactic unit
| (h::t) := match h with
            | `(_ → _ → (_ ∈ AssertionLang.sem (%%A & %%B).always)) :=
              if A = e then do `[apply portation_left], return ()
              else if B = e then do `[apply portation_right], return ()
              else check_for_portation_aux' t
            | _ := check_for_portation_aux' t
            end
| [] := return ()

meta def check_for_portation : tactic unit := do
`(_ ∈ LANG.sem (LANG.always %%FOO)) ← target,
local_context >>= (list.mmap infer_type) >>= check_for_portation_aux FOO

meta def check_for_portation' : tactic unit := do
`(_ ∈ LANG.sem (LANG.always %%FOO)) ← target,
local_context >>= (list.mmap infer_type) >>= check_for_portation_aux' FOO

meta def find_assumption : list expr → tactic bool
| [] := return ff
| (h::t) := do apply h >> return tt <|> find_assumption t

meta def try_terminal_rws : tactic unit := do
 `[rwa disj_comm_guarded] <|> `[rwa conj_comm_guarded] <|> return ()

meta def core_rpo_loop (x : name) (Γ : model_info) : tactic unit := do
check_for_portation,
local_context >>= try_application,
synchronize_ports x Γ,
target >>= try_rewrites,
b ← local_context >>= find_assumption,
if !b then try_terminal_rws else return ()

meta def solve_rpo_fst_new (Γ : model_info) : tactic unit := do
  ⟨x, H⟩ ← preprocess_rpo_fst,
  get_local H >>= extract_contracts,
  decompose_env_assumptions x,
  synchronize_ports x Γ,
  local_context >>= try_application,
  synchronize_ports x Γ,
  target >>= try_rewrites,
  repeat `[split],
  all_goals $ core_rpo_loop x Γ,
  try (`[repeat {split}, repeat {assumption}]), -- can fix.
  return ()

meta def make_cases (H : name) : tactic unit := do
  H_e ← get_local H,
  τ ← infer_type H_e,
  match τ with
  |`(or _ _) := do cases_core H_e [] transparency.semireducible,
                   all_goals (make_cases ),
                   return ()
  | _ := return ()
  end

meta def find_this_component (Γ : model_info) : list expr → tactic expr
| [] := return default
| (h::t) := do
            τ ← infer_type h,
            match τ with
            | `(_ = %%FOO) :=
                  if FOO ∈ Γ.comps then return FOO else find_this_component t
            | _ := find_this_component t
            end

meta def get_comp_contract (H : expr)  (comp : expr): tactic unit := do
  h ← get_unused_name,
  let e : expr := (expr.mk_app H [comp]),
  τ ← infer_type e,
  tactic.local_proof h (τ) (tactic.exact e),
  return ()

meta def get_other_comp_contracts (Γ : model_info) (H2 : expr) : tactic expr := do
  e ← local_context >>= find_this_component Γ,
  let l := Γ.comps.erase e,
  l.mmap (get_comp_contract H2),
  return e

meta def rewrite_this_comp  (comp : expr ) : list expr → tactic unit
| (h::t) := do
          τ ← infer_type h,
          match τ with
          | `(_ = %%FOO) := if FOO = comp then subst h else rewrite_this_comp t
          | _ := rewrite_this_comp t
          end
| [] := return ()

meta def decompose_fst_hyp (Γ : model_info) : tactic unit := do
  Ha ← tactic.get_unused_name `Ha,
  `[rw AssertionLang.impl_def],
  tactic.intro Ha,
  `[rw AssertionLang.conj_def at *],
  H_e ← get_local Ha,
  H1 ← get_unused_name `H1,
  H2 ← get_unused_name `H2,
  cases_core H_e [H1,H2] transparency.semireducible,
  `[rw list_conj_iff at *, simp only [list.mem_map, forall_exists_index, and_imp, forall_apply_eq_imp_iff₂] at *],
   this_comp ← get_local H2 >>= get_other_comp_contracts Γ,
   local_context >>= rewrite_this_comp this_comp,
   H1_e ← get_local H1,
   tactic.rewrite_hyp `(asdff) H1_e,
   repeat `[rw fsadfdsa at *],
  tactic.iterate (do H1_e ← get_local H1, cases_core H1_e [H1] transparency.semireducible),
  return ()

meta def clear_goal_rewrites_aux : expr → tactic unit
| `(ite (%%X = %%Y) %%B %%C) :=  if X = Y then do `[rw if_pos] else do `[rw if_neg], clear_goal_rewrites_aux C
| _ := do  return ()

meta def clear_goal_of_rewrites : tactic unit := do
 tgt ← target,
 match tgt with
 | `(_ ∈ AssertionLang.sem (Contract.A (option.iget (%%FOO)))) := clear_goal_rewrites_aux FOO
 | _ := do tactic.trace "hmm", return ()
 end

meta def rpo_snd_core : tactic unit := do
  repeat (`[rw Map.find_val at *]),
  clear_goal_of_rewrites,
  xs ← local_context >>= collect_assertions,
  l ← xs.mmap infer_type >>= GFSAUIHasd,
  let ls := xs.zip (l.map aiuhsfd),
  suihfa ls,
  clear_goal_of_rewrites,
  iterate `[rw if_neg at *],
  any_goals `[dec_trivial],
  `[rw nf_def at *, simp only [Map.find_val] at *],
  try (`[rw asdff, repeat {rw fsadfdsa}, rw ← asdff, repeat {split}]),
  return ()

meta def check_for_semantic_unfolding : expr → tactic unit
| `(_ ∈ AssertionLang.sem _) := do return ()
| _ := do try `[rw ← asdff]

meta def finish_conjunctions : tactic unit := do
 any_goals $ `[rw fsadfdsa],
 any_goals $ split,
 any_goals $ assumption,
 any_goals $ `[rw ← asdff], assumption,
 return ()

meta def new_tac_finisher (x : name) (Γ : model_info) : tactic unit := do
  synchronize_ports x Γ,
  `[rw asdff],
  try_terminal_rws,
  -- finish_conjunctions,
  -- `[rw fsadfdsa],
    -- split, -- copied from above
    -- assumption,rw ← asdff,assumption],
  return ()

meta def solve_rpo_snd_new (Γ : model_info) : tactic unit := do
  S ← tactic.get_unused_name `S,
  H ← tactic.get_unused_name `H,
  x ← tactic.get_unused_name `x,
  tactic.intro S, tactic.intro H, tactic.intro x,
  `[simp at *],
  make_cases H,
  all_goals $ decompose_fst_hyp Γ,
  all_goals $ rpo_snd_core,
  all_goals $ try $ `[rw ← asdff],
  all_goals (do x_e ← get_local x,synchronize_ports x Γ),
  all_goals $ `[rw asdff],
  all_goals $ local_context >>= try_application,
  any_goals `[dec_trivial],
  any_goals $ try_terminal_rws,
  all_goals $ try (check_for_portation'),
  all_goals $ local_context >>= try_application,
  any_goals $ `[dec_trivial],
  all_goals $ new_tac_finisher x Γ,
  all_goals $ (do `[rw fsadfdsa, split], assumption <|> `[rw ← asdff], assumption),
  return ()


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

Notation for operators defined at Prelude.lean
-/
prelude
import Init.NotationExtra

namespace Lean.Parser.Tactic.Conv

/-- `conv` is the syntax category for a "conv tactic", where "conv" is short
for conversion. A conv tactic is a program which receives a target, printed as
`| a`, and is tasked with coming up with some term `b` and a proof of `a = b`.
It is mainly used for doing targeted term transformations, for example rewriting
only on the left side of an equality. -/
declare_syntax_cat conv (behavior := both)

syntax convSeq1Indented := sepBy1IndentSemicolon(conv)
syntax convSeqBracketed := "{" withoutPosition(sepByIndentSemicolon(conv)) "}"
-- Order is important: a missing `conv` proof should not be parsed as `{ <missing> }`,
-- automatically closing goals
syntax convSeq := convSeqBracketed <|> convSeq1Indented

/-- The `*` occurrence list means to apply to all occurrences of the pattern. -/
syntax occsWildcard := "*"

/--
A list `1 2 4` of occurrences means to apply to the first, second, and fourth
occurrence of the pattern.
-/
syntax occsIndexed := num+

/-- An occurrence specification, either `*` or a list of numbers. The default is `[1]`. -/
syntax occs := atomic("(" &"occs") " := " (occsWildcard <|> occsIndexed) ") "

/--
`with_annotate_state stx t` annotates the lexical range of `stx : Syntax` with
the initial and final state of running tactic `t`.
-/
scoped syntax (name := withAnnotateState)
  "with_annotate_state " rawStx ppSpace conv : conv

/-- `skip` does nothing. -/
syntax (name := skip) "skip" : conv

/-- Traverses into the left subterm of a binary operator.
(In general, for an `n`-ary operator, it traverses into the second to last argument.) -/
syntax (name := lhs) "lhs" : conv

/-- Traverses into the right subterm of a binary operator.
(In general, for an `n`-ary operator, it traverses into the last argument.) -/
syntax (name := rhs) "rhs" : conv

/-- Reduces the target to Weak Head Normal Form. This reduces definitions
in "head position" until a constructor is exposed. For example, `List.map f [a, b, c]`
weak head normalizes to `f a :: List.map f [b, c]`. -/
syntax (name := whnf) "whnf" : conv

/-- Expands let-declarations and let-variables. -/
syntax (name := zeta) "zeta" : conv

/-- Puts term in normal form, this tactic is meant for debugging purposes only. -/
syntax (name := reduce) "reduce" : conv

/-- Performs one step of "congruence", which takes a term and produces
subgoals for all the function arguments. For example, if the target is `f x y` then
`congr` produces two subgoals, one for `x` and one for `y`. -/
syntax (name := congr) "congr" : conv

/--
* `arg i` traverses into the `i`'th argument of the target. For example if the
  target is `f a b c d` then `arg 1` traverses to `a` and `arg 3` traverses to `c`.
* `arg @i` is the same as `arg i` but it counts all arguments instead of just the
  explicit arguments. -/
syntax (name := arg) "arg " "@"? num : conv

/-- `ext x` traverses into a binder (a `fun x => e` or `∀ x, e` expression)
to target `e`, introducing name `x` in the process. -/
syntax (name := ext) "ext" (colGt ident)* : conv

/-- `change t'` replaces the target `t` with `t'`,
assuming `t` and `t'` are definitionally equal. -/
syntax (name := change) "change " term : conv

/-- `delta id1 id2 ...` unfolds all occurrences of `id1`, `id2`, ... in the target.
Like the `delta` tactic, this ignores any definitional equations and uses
primitive delta-reduction instead, which may result in leaking implementation details.
Users should prefer `unfold` for unfolding definitions. -/
syntax (name := delta) "delta " (colGt ident)+ : conv

/--
* `unfold foo` unfolds all occurrences of `foo` in the target.
* `unfold id1 id2 ...` is equivalent to `unfold id1; unfold id2; ...`.
Like the `unfold` tactic, this uses equational lemmas for the chosen definition
to rewrite the target. For recursive definitions,
only one layer of unfolding is performed. -/
syntax (name := unfold) "unfold " (colGt ident)+ : conv

/--
* `pattern pat` traverses to the first subterm of the target that matches `pat`.
* `pattern (occs := *) pat` traverses to every subterm of the target that matches `pat`
  which is not contained in another match of `pat`. It generates one subgoal for each matching
  subterm.
* `pattern (occs := 1 2 4) pat` matches occurrences `1, 2, 4` of `pat` and produces three subgoals.
  Occurrences are numbered left to right from the outside in.

Note that skipping an occurrence of `pat` will traverse inside that subexpression, which means
it may find more matches and this can affect the numbering of subsequent pattern matches.
For example, if we are searching for `f _` in `f (f a) = f b`:
* `occs := 1 2` (and `occs := *`) returns `| f (f a)` and `| f b`
* `occs := 2` returns `| f a`
* `occs := 2 3` returns `| f a` and `| f b`
* `occs := 1 3` is an error, because after skipping `f b` there is no third match.
-/
syntax (name := pattern) "pattern " (occs)? term : conv

/-- `rw [thm]` rewrites the target using `thm`. See the `rw` tactic for more information. -/
syntax (name := rewrite) "rewrite" (config)? rwRuleSeq : conv

/-- `simp [thm]` performs simplification using `thm` and marked `@[simp]` lemmas.
See the `simp` tactic for more information. -/
syntax (name := simp) "simp" (config)? (discharger)? (&" only")?
  (" [" withoutPosition((simpStar <|> simpErase <|> simpLemma),*) "]")? : conv

/--
`dsimp` is the definitional simplifier in `conv`-mode. It differs from `simp` in that it only
applies theorems that hold by reflexivity.

Examples:

```lean
example (a : Nat): (0 + 0) = a - a := by
  conv =>
    lhs
    dsimp
    rw [← Nat.sub_self a]
```
-/
syntax (name := dsimp) "dsimp " (config)? (discharger)? (&"only ")?
  ("[" withoutPosition((simpErase <|> simpLemma),*) "]")? : conv

/-- `simp_match` simplifies match expressions. For example,
```
match [a, b] with
| [] => 0
| hd :: tl => hd
```
simplifies to `a`. -/
syntax (name := simpMatch) "simp_match" : conv

/-- Executes the given tactic block without converting `conv` goal into a regular goal. -/
syntax (name := nestedTacticCore) "tactic'" " => " tacticSeq : conv

/-- Focuses, converts the `conv` goal `⊢ lhs` into a regular goal `⊢ lhs = rhs`, and then executes the given tactic block. -/
syntax (name := nestedTactic) "tactic" " => " tacticSeq : conv

/-- Executes the given conv block without converting regular goal into a `conv` goal. -/
syntax (name := convTactic) "conv'" " => " convSeq : tactic

/-- `{ convs }` runs the list of `convs` on the current target, and any subgoals that
remain are trivially closed by `skip`. -/
syntax (name := nestedConv) convSeqBracketed : conv

/-- `(convs)` runs the `convs` in sequence on the current list of targets.
This is pure grouping with no added effects. -/
syntax (name := paren) "(" withoutPosition(convSeq) ")" : conv

/-- `rfl` closes one conv goal "trivially", by using reflexivity
(that is, no rewriting). -/
macro "rfl" : conv => `(conv| tactic => rfl)

/-- `done` succeeds iff there are no goals remaining. -/
macro "done" : conv => `(conv| tactic' => done)

/-- `trace_state` prints the current goal state. -/
macro "trace_state" : conv => `(conv| tactic' => trace_state)

/-- `all_goals tac` runs `tac` on each goal, concatenating the resulting goals, if any. -/
macro (name := allGoals) tk:"all_goals " s:convSeq : conv =>
  `(conv| tactic' => all_goals%$tk conv' => $s)

/--
`any_goals tac` applies the tactic `tac` to every goal, and succeeds if at
least one application succeeds.
-/
macro (name := anyGoals) tk:"any_goals " s:convSeq : conv =>
  `(conv| tactic' => any_goals%$tk conv' => $s)

/--
* `case tag => tac` focuses on the goal with case name `tag` and solves it using `tac`,
  or else fails.
* `case tag x₁ ... xₙ => tac` additionally renames the `n` most recent hypotheses
  with inaccessible names to the given names.
* `case tag₁ | tag₂ => tac` is equivalent to `(case tag₁ => tac); (case tag₂ => tac)`.
-/
macro (name := case) tk:"case " args:sepBy1(caseArg, " | ") arr:" => " s:convSeq : conv =>
  `(conv| tactic' => case%$tk $args|* =>%$arr conv' => ($s); all_goals rfl)

/--
`case'` is similar to the `case tag => tac` tactic, but does not ensure the goal
has been solved after applying `tac`, nor admits the goal if `tac` failed.
Recall that `case` closes the goal using `sorry` when `tac` fails, and
the tactic execution is not interrupted.
-/
macro (name := case') tk:"case' " args:sepBy1(caseArg, " | ") arr:" => " s:convSeq : conv =>
  `(conv| tactic' => case'%$tk $args|* =>%$arr conv' => $s)

/--
`next => tac` focuses on the next goal and solves it using `tac`, or else fails.
`next x₁ ... xₙ => tac` additionally renames the `n` most recent hypotheses with
inaccessible names to the given names.
-/
macro "next " args:binderIdent* " => " tac:convSeq : conv => `(conv| case _ $args* => $tac)

/--
`focus tac` focuses on the main goal, suppressing all other goals, and runs `tac` on it.
Usually `· tac`, which enforces that the goal is closed by `tac`, should be preferred.
-/
macro (name := focus) tk:"focus " s:convSeq : conv => `(conv| tactic' => focus%$tk conv' => $s)

/-- `conv => cs` runs `cs` in sequence on the target `t`,
resulting in `t'`, which becomes the new target subgoal. -/
syntax (name := convConvSeq) "conv" " => " convSeq : conv

/-- `· conv` focuses on the main conv goal and tries to solve it using `s`. -/
macro dot:patternIgnore("·" <|> ".") s:convSeq : conv => `(conv| {%$dot ($s) })

/-- `fail_if_success t` fails if the tactic `t` succeeds. -/
macro (name := failIfSuccess) tk:"fail_if_success " s:convSeq : conv =>
  `(conv| tactic' => fail_if_success%$tk conv' => $s)

/-- `rw [rules]` applies the given list of rewrite rules to the target.
See the `rw` tactic for more information. -/
macro "rw" c:(config)? s:rwRuleSeq : conv => `(conv| rewrite $[$c]? $s)

/-- `erw [rules]` is a shorthand for `rw (config := { transparency := .default }) [rules]`.
This does rewriting up to unfolding of regular definitions (by comparison to regular `rw`
which only unfolds `@[reducible]` definitions). -/
macro "erw" s:rwRuleSeq : conv => `(conv| rw (config := { transparency := .default }) $s)

/-- `args` traverses into all arguments. Synonym for `congr`. -/
macro "args" : conv => `(conv| congr)
/-- `left` traverses into the left argument. Synonym for `lhs`. -/
macro "left" : conv => `(conv| lhs)
/-- `right` traverses into the right argument. Synonym for `rhs`. -/
macro "right" : conv => `(conv| rhs)
/-- `intro` traverses into binders. Synonym for `ext`. -/
macro "intro" xs:(colGt ident)* : conv => `(conv| ext $xs*)

syntax enterArg := ident <|> ("@"? num)

/-- `enter [arg, ...]` is a compact way to describe a path to a subterm.
It is a shorthand for other conv tactics as follows:
* `enter [i]` is equivalent to `arg i`.
* `enter [@i]` is equivalent to `arg @i`.
* `enter [x]` (where `x` is an identifier) is equivalent to `ext x`.
For example, given the target `f (g a (fun x => x b))`, `enter [1, 2, x, 1]`
will traverse to the subterm `b`. -/
syntax "enter" " [" (colGt enterArg),+ "]": conv
macro_rules
  | `(conv| enter [$i:num]) => `(conv| arg $i)
  | `(conv| enter [@$i]) => `(conv| arg @$i)
  | `(conv| enter [$id:ident]) => `(conv| ext $id)
  | `(conv| enter [$arg, $args,*]) => `(conv| (enter [$arg]; enter [$args,*]))

/-- The `apply thm` conv tactic is the same as `apply thm` the tactic.
There are no restrictions on `thm`, but strange results may occur if `thm`
cannot be reasonably interpreted as proving one equality from a list of others. -/
-- TODO: error if non-conv subgoals?
macro "apply " e:term : conv => `(conv| tactic => apply $e)

/-- `first | conv | ...` runs each `conv` until one succeeds, or else fails. -/
syntax (name := first) "first " withPosition((colGe "|" convSeq)+) : conv

/-- `try tac` runs `tac` and succeeds even if `tac` failed. -/
macro "try " t:convSeq : conv => `(conv| first | $t | skip)

macro:1 x:conv tk:" <;> " y:conv:0 : conv =>
  `(conv| tactic' => (conv' => $x:conv) <;>%$tk (conv' => $y:conv))

/-- `repeat convs` runs the sequence `convs` repeatedly until it fails to apply. -/
syntax "repeat" convSeq : conv
macro_rules
  | `(conv| repeat $seq) => `(conv| first | ($seq); repeat $seq | rfl)

/--
`conv => ...` allows the user to perform targeted rewriting on a goal or hypothesis,
by focusing on particular subexpressions.

See <https://leanprover.github.io/theorem_proving_in_lean4/conv.html> for more details.

Basic forms:
* `conv => cs` will rewrite the goal with conv tactics `cs`.
* `conv at h => cs` will rewrite hypothesis `h`.
* `conv in pat => cs` will rewrite the first subexpression matching `pat` (see `pattern`).
-/
-- HACK: put this at the end so that references to `conv` above
-- refer to the syntax category instead of this syntax
syntax (name := conv) "conv " (" at " ident)? (" in " (occs)? term)? " => " convSeq : tactic

end Lean.Parser.Tactic.Conv


import Lean
import Mathlib.Tactic.Simps.Basic
import Mathlib.Tactic.PermuteGoals
import Mathlib.Tactic.Classical

open Lean Elab Parser Term Meta Tactic

syntax (name := seq) "seq" tacticSeq : tactic

section Source
  -- modified from `Lean.Elab.Tactic.Simp`
  def traceSimpCall' (stx : Syntax) (usedSimps : Simp.UsedSimps) : MetaM Syntax := do
    let mut stx := stx
    if stx[3].isNone then
      stx := stx.setArg 3 (mkNullNode #[mkAtom "only"])
    let mut args := #[]
    let mut localsOrStar := some #[]
    let lctx ← getLCtx
    let env ← getEnv
    for (thm, _) in usedSimps.toArray.qsort (·.2 < ·.2) do
      match thm with
      | .decl declName => -- global definitions in the environment
        if env.contains declName && !simpOnlyBuiltins.contains declName then
          args := args.push (← `(Parser.Tactic.simpLemma| $(mkIdent (← unresolveNameGlobal declName)):ident))
      | .fvar fvarId => -- local hypotheses in the context
        if let some ldecl := lctx.find? fvarId then
          localsOrStar := localsOrStar.bind fun locals =>
            if !ldecl.userName.isInaccessibleUserName &&
                (lctx.findFromUserName? ldecl.userName).get!.fvarId == ldecl.fvarId then
              some (locals.push ldecl.userName)
            else
              none
        -- Note: the `if let` can fail for `simp (config := {contextual := true})` when
        -- rewriting with a variable that was introduced in a scope. In that case we just ignore.
      | .stx _ thmStx => -- simp theorems provided in the local invocation
        args := args.push ⟨thmStx⟩
      | .other _ => -- Ignore "special" simp lemmas such as constructed by `simp_all`.
        pure ()     -- We can't display them anyway.
    if let some locals := localsOrStar then
      args := args ++ (← locals.mapM fun id => `(Parser.Tactic.simpLemma| $(mkIdent id):ident))
    else
      args := args.push ⟨(← `(Parser.Tactic.simpStar| *))⟩
    let argsStx := if args.isEmpty then #[] else #[mkAtom "[", (mkAtom ",").mkSep args, mkAtom "]"]
    stx := stx.setArg 4 (mkNullNode argsStx)
    return stx

  def dsimpLocation' (ctx : Simp.Context) (loc : Location) : TacticM Syntax := do
    match loc with
    | Location.targets hyps simplifyTarget =>
      withMainContext do
        let fvarIds ← getFVarIds hyps
        go fvarIds simplifyTarget
    | Location.wildcard =>
      withMainContext do
        go (← (← getMainGoal).getNondepPropHyps) (simplifyTarget := true)
  where
    go (fvarIdsToSimp : Array FVarId) (simplifyTarget : Bool) : TacticM Syntax := do
      let mvarId ← getMainGoal
      let (result?, usedSimps) ← dsimpGoal mvarId ctx (simplifyTarget := simplifyTarget) (fvarIdsToSimp := fvarIdsToSimp)
      match result? with
      | none => replaceMainGoal []
      | some mvarId => replaceMainGoal [mvarId]
      traceSimpCall' (← getRef) usedSimps

  def getMainGoal' : TacticM (MVarId × List MVarId) := do
  loop (← getGoals)
    where
  loop : List MVarId → TacticM (MVarId × List MVarId)
    | [] => throwNoGoalsToBeSolved
    | mvarId :: mvarIds => do
      if (← mvarId.isAssigned) then
        loop mvarIds
      else
        setGoals (mvarId :: mvarIds)
        return (mvarId, mvarIds)

/--
  Searches for a metavariable `g` s.t. `tag` is its exact name.
  If none then searches for a metavariable `g` s.t. `tag` is a suffix of its name.
  If none, then it searches for a metavariable `g` s.t. `tag` is a prefix of its name. -/
def findTag? (mvarIds : List MVarId) (tag : Name) : TacticM (Option MVarId) := do
  match (← mvarIds.findM? fun mvarId => return tag == (← mvarId.getDecl).userName) with
  | some mvarId => return mvarId
  | none =>
  match (← mvarIds.findM? fun mvarId => return tag.isSuffixOf (← mvarId.getDecl).userName) with
  | some mvarId => return mvarId
  | none => mvarIds.findM? fun mvarId => return tag.isPrefixOf (← mvarId.getDecl).userName

def getCaseGoals (tag : TSyntax `Lean.binderIdent) : TacticM (MVarId × List MVarId) := do
  let gs ← getUnsolvedGoals
  let g ← if let `(Lean.binderIdent| $tag:ident) := tag then
    let tag := tag.getId
    let some g ← findTag? gs tag | throwError "tag not found"
    pure g
  else
    getMainGoal
  return (g, gs.erase g)

def matchAltTac := Term.matchAlt (rhsParser := matchRhs)

end Source

def traceGoalsAt (stx : TSyntax `tactic) : TacticM Unit := do
  let gs ← getUnsolvedGoals
  withRef stx <| addRawTrace (goalsToMessageData gs)

def traceTacticCallAt (stx : TSyntax `tactic) (tac : TSyntax `tactic) : TacticM Unit := do
  withRef stx <| addRawTrace m!"[TACTIC] {tac}"

#check Split.applyMatchSplitter

partial def evalTacticWithTrace : TSyntax `tactic → TacticM Unit
  /- Dealing with bracketing -/
  | `(tactic| { $[$tacs]* }) => do
    for tac in tacs do
      evalTacticWithTrace tac
  | `(tactic| ( $[$tacs]* )) => do
    for tac in tacs do
      evalTacticWithTrace tac
  /- Dealing with focused goals -/
  | `(tactic| · $[$tacs]*) => do
    let (mainGoal, otherGoals) ← getMainGoal'
    setGoals [mainGoal]
    for tac in tacs do
      evalTacticWithTrace tac
    setGoals otherGoals
  | `(tactic| focus $[$tacs]*) => do
    let (mainGoal, otherGoals) ← getMainGoal'
    setGoals [mainGoal]
    for tac in tacs do
      evalTacticWithTrace tac
    setGoals otherGoals
  /- Handle trace for the `classical` tactic -/
  | `(tactic| classical $[$tacs]*) => do
      modifyEnv Meta.instanceExtension.pushScope
      Meta.addInstance ``Classical.propDecidable .local 10
      try for tac in tacs do evalTacticWithTrace tac
      finally modifyEnv Meta.instanceExtension.popScope
   /- Trace `simp` calls with the complete list of theorems used -/
  | stx@`(tactic| simp%$tk $(config)? $(discharger)? $[only%$o]? $[[$args,*]]? $(loc)?) => do
    traceGoalsAt stx
    let { ctx, dischargeWrapper } ← withMainContext <| mkSimpContext stx (eraseLocal := false)
    let usedSimps ← dischargeWrapper.with fun discharge? =>
      simpLocation ctx discharge? (expandOptLocation stx.raw[5])
    traceTacticCallAt stx ⟨← traceSimpCall' stx usedSimps⟩
    traceGoalsAt stx
  | stx@`(tactic| simp_all%$tk $(config)? $(discharger)? $[only%$o]? $[[$args,*]]?) => do
    traceGoalsAt stx
    let { ctx, .. } ← mkSimpContext stx (eraseLocal := true) (kind := .simpAll) (ignoreStarArg := true)
    let (result?, usedSimps) ← simpAll (← getMainGoal) ctx
    match result? with
    | none => replaceMainGoal []
    | some mvarId => replaceMainGoal [mvarId]
    traceTacticCallAt stx ⟨← traceSimpCall' stx usedSimps⟩
    traceGoalsAt stx
  | stx@`(tactic| dsimp%$tk $(config)? $[only%$o]? $[[$args,*]]? $(loc)?) => do
    traceGoalsAt stx
    let { ctx, .. } ← withMainContext <| mkSimpContext stx (eraseLocal := false) (kind := .dsimp)
    traceTacticCallAt stx ⟨← dsimpLocation' ctx (expandOptLocation stx.raw[5])⟩
    traceGoalsAt stx
  /- Treat a rewrite sequence as a sequence of individual rewrites, with a trace provided at each step -/
  | `(tactic| rw $[$cfg]? [$rs,*] $[$loc]?) => do
    for r in (rs : TSyntaxArray `Lean.Parser.Tactic.rwRule) do
      traceGoalsAt ⟨r.raw⟩
      let rtac ← `(tactic| rw $[$cfg]? [$r] $[$loc]?)
      evalTactic rtac
      traceTacticCallAt ⟨r.raw⟩ rtac
      traceGoalsAt ⟨r.raw⟩
  | `(tactic| erw [$rs,*] $[$loc]?) => do
    for r in (rs : TSyntaxArray `Lean.Parser.Tactic.rwRule) do
      traceGoalsAt ⟨r.raw⟩
      let rtac ← `(tactic| erw [$r] $[$loc]?)
      evalTactic rtac
      traceTacticCallAt ⟨r.raw⟩ rtac
      traceGoalsAt ⟨r.raw⟩
  | `(tactic| rwa [$rs,*] $[$loc]?) => do
    for r in (rs : TSyntaxArray `Lean.Parser.Tactic.rwRule) do
      traceGoalsAt ⟨r.raw⟩
      let rtac ← `(tactic| rw [$r] $[$loc]?)
      evalTactic rtac
      traceTacticCallAt ⟨r.raw⟩ rtac
      traceGoalsAt ⟨r.raw⟩
    `(tactic| assumption) >>= evalTactic ∘ TSyntax.raw
  /- Annotate `apply` and `exact` applications with the type information -/
  | stx@`(tactic| apply $v) => do
    traceGoalsAt stx
    evalTactic stx
    let trm ← Tactic.elabTerm v none
    let typ ← inferType trm
    let (mainGoal, otherGoals) ← getMainGoal'
    let newGoals ← mainGoal.apply trm
    setGoals <| newGoals ++ otherGoals
    let typ ← instantiateMVars typ
    let typStx ← PrettyPrinter.delab typ
    `(tactic| apply ($v : $typStx)) >>= traceTacticCallAt stx
    traceGoalsAt stx
  | stx@`(tactic| exact $v) => do
    traceGoalsAt stx
    evalTactic stx
    let trm ← Tactic.elabTerm v none
    let typ ← inferType trm
    let typStx ← PrettyPrinter.delab typ
    `(tactic| exact ($v : $typStx)) >>= traceTacticCallAt stx
    traceGoalsAt stx
  /- Handling `match`, `induction` and `cases` -/
  | stx@`(tactic| case $[$tag $hs*]|* =>%$arr $tac:tacticSeq) => do
    for tag in tag, h in hs do
      traceGoalsAt ⟨arr⟩
      traceTacticCallAt ⟨arr⟩ stx -- TODO (unimportant) remove the `tacticSeq` from `stx`
      let (g, _) ← getCaseGoals tag
      withRef arr <| addRawTrace (goalsToMessageData [g])
    let stx ← `(tactic| case $[$tag $hs*]|* =>%$arr seq $tac:tacticSeq)
    evalTactic stx.raw
  | stx@`(tactic| case' $[$tag $hs*]|* =>%$arr $tac:tacticSeq) => do
    for tag in tag, h in hs do
      traceGoalsAt ⟨arr⟩
      traceTacticCallAt ⟨arr⟩ stx -- TODO (unimportant) remove the `tacticSeq` from `stx`
      let (g, _) ← getCaseGoals tag
      withRef arr <| addRawTrace (goalsToMessageData [g])
    let stx ← `(tactic| case' $[$tag $hs*]|* =>%$arr seq $tac:tacticSeq)
  | `(tactic| induction $[$ts],* $[using $id:ident]?  $[generalizing $gs*]? with $[$tac]? $is*) => do
    let is' : TSyntaxArray ``inductionAlt ←
      is.mapM <|
        fun
          | `(inductionAlt| $il* => $ts:tacticSeq) => `(inductionAlt| $il* => seq $ts)
          | i => return ⟨i⟩
    let stx' ← `(tactic| induction $[$ts],* $[using $id:ident]?  $[generalizing $gs*]? with $[$tac]? $is'*)
    evalTactic stx'
  | `(tactic| cases $[$cs],* $[using $id:ident]? with $[$tac]? $is*) => do
    let is' : TSyntaxArray ``inductionAlt ←
      is.mapM <|
        fun
          | `(inductionAlt| $il* => $ts:tacticSeq) => `(inductionAlt| $il* => seq $ts)
          | i => return ⟨i⟩
    let stx' ← `(tactic| cases $[$cs],* $[using $id:ident]? with $[$tac]? $is'*)
    evalTactic stx'
  | `(tactic| match $[$gen]? $[$motive]? $discrs,* with $alts:matchAlt*) => do
    let alts' : TSyntaxArray ``matchAlt ←
      alts.mapM <|
        fun
          | `(matchAltTac| | $[$pats,*]|* => $rhs:tacticSeq) => do
              let alt ← `(matchAltTac| | $[$pats,*]|* => seq $rhs)
              return ⟨alt⟩
          | alt =>  return ⟨alt⟩
      let stx' ← `(tactic| match $[$gen]? $[$motive]? $discrs,* with $alts':matchAlt*)
      evalTactic stx'
  /- Display the expected type in `have` and `let` statements -/
  | stx@`(tactic| have $[$x:ident]? := $prf) => do
    traceGoalsAt stx
    evalTactic stx
    let trm ← Tactic.elabTerm prf none
    let typ ← inferType trm
    let typStx ← PrettyPrinter.delab typ
    `(tactic| have $[$x:ident]? : $typStx := $prf) >>= traceTacticCallAt stx
    traceGoalsAt stx
  | stx@`(tactic| let $x:ident := $val) => do
    traceGoalsAt stx
    evalTactic stx
    let trm ← Tactic.elabTerm val none
    let typ ← inferType trm
    let typStx ← PrettyPrinter.delab typ
    `(tactic| let $x:ident : $typStx := $val) >>= traceTacticCallAt stx
    traceGoalsAt stx
  /- Otherwise, evaluate the tactic normally -/
  | stx@`(tactic| $tac) => do
    traceGoalsAt stx
    evalTactic tac
    traceTacticCallAt stx tac
    traceGoalsAt stx

@[tactic seq]
def traceSequence : Tactic := fun s => do
  -- Leonardo de Moura's code for extracting the list of tactics
  match s with
  | `(tactic| seq $[$tacs]*) =>
    for tac in tacs do
      evalTacticWithTrace tac
  | _ => evalTactic <| ← `(tactic.sorry)

-- an example of the `seq` tactic
example (h : x = y) : 0 + x = y ∧ 1 = 1 := by
  seq
    rw [Nat.zero_add, (h : x = y)]
  refine' ⟨_, _⟩
  focus
    rw [←h, h]
  · rfl

-- a deep copy of Lean's `by` tactic, called `by'`
syntax (name := byTactic') "by' " tacticSeq : term

@[term_elab byTactic'] def elabByTactic' : TermElab := fun stx expectedType? => do
  match expectedType? with
  | some expectedType =>
    let mvar ← mkFreshExprMVar expectedType MetavarKind.syntheticOpaque
    let mvarId := mvar.mvarId!
    let ref ← getRef
    registerSyntheticMVar ref mvarId <| SyntheticMVarKind.tactic stx (← saveContext)
    return mvar
  | none =>
    tryPostpone
    throwError ("invalid 'by\'' tactic, expected type has not been provided")

example : 1 + 1 = 2 := by' -- the new `by'` syntax can be used to replace `by`
  focus
  rfl

-- intercepting the `by` tactic to output intermediate trace data
-- the `by'` clone is needed here to avoid infinite recursion
macro_rules
  | `(by $t) => `(by' seq $t)

section Test

set_option linter.unreachableTactic false

-- the `by` tactic now generates trace data by default
example (h : x = y) : x + 0 + x = x + y ∧ 1 = 1 := by
  have := (rfl : 1 = 1)
  let a := 5
  simp at this
  refine' ⟨_, _⟩
  · apply Eq.symm
    apply Eq.symm
    erw [h, ← h, h, ← h]
    simp_all
  · apply Eq.symm
    apply Eq.symm
    subst h
    rfl
  done

example : ∀ n : Nat, n = n := by
  intro n
  let x : ∀ m : ℕ, m = m := by
    intro a
    rfl
  match n with
  | .zero => rfl
  | .succ _ => rfl

example : P ∧ Q ↔ Q ∧ P := by
  constructor
  · intro ⟨x, y⟩
    constructor
    · assumption
    · assumption
  · intro ⟨_, _⟩
    constructor
    · assumption
    · assumption

example : ∀ n : Nat, n = n := by
  intro n
  induction n
  case zero => rfl
  case succ _ ih => simp

example : ∀ n : Nat, n + n = n + n := by
  intro n
  induction n with
    | zero => rfl
    | succ _ _ => rfl

example : ∀ n : Nat, n + n = n + n := by
  intro n
  cases n with
    | zero =>
      let h := (rfl : 1 = 1)
      rfl
    | succ _ => rfl

example : ∀ n : Nat, n + n = n + n := by
  intro n
  match n with
    | .zero => rfl
    | .succ _ => rfl

end Test


import Lean
import Mathlib.Control.Writer
import PremiseSelection.Utils
open Lean
/-!

# Theorem feature extraction

Input: the goal state
Ouput: the theorem statement as an expr

 -/
open Std

def Std.RBMap.modify' (k : κ) (fn : Option α → Option α) (r : RBMap κ α cmp) :=
  match fn <| r.find? k with
  | none => r.erase k
  | some v => r.insert k v

def Std.RBMap.mergeBy (fn : κ → α → α → α) (r1 r2 : RBMap κ α cmp) :  RBMap κ α cmp :=
  r2.foldl (fun r1 k v2 => r1.modify' k (fun | none => some v2 | some v1 => some (fn k v1 v2))) r1

namespace PremiseSelection

def Multiset (α : Type) [Ord α] := Std.RBMap α Nat compare

variable {α : Type} [Ord α]

def Multiset.empty : Multiset α := mkRBMap _ _ _

instance : EmptyCollection  (Multiset α) :=  ⟨Multiset.empty⟩

instance : Append  (Multiset α) where
  append x y := x.mergeBy (fun _ => (·+·)) y

def Multiset.add : Multiset α → α → Multiset α
  | m, a => m.modify' a (fun | none => some 1 | some v => some (v + 1))

def Multiset.singleton : α → Multiset α
  | a => Multiset.empty |>.add a

instance : Ord Name := ⟨Name.quickCmp⟩

structure Bigram where
  fst : Name
  snd : Name
  deriving Ord

structure Trigram where
  fst : Name
  snd : Name
  trd : Name
  deriving Ord

instance : ToJson Bigram where
  toJson b := s!"{b.fst}/{b.snd}"

instance : ToString Bigram where
  toString b := s!"{b.fst}/{b.snd}"

instance : ToString Trigram where
  toString t := s!"{t.fst}/{t.snd}/{t.trd}"

instance : ToJson Trigram where
  toJson t := s!"{t.fst}/{t.snd}/{t.trd}"

structure StatementFeatures where
  /-- Just the constant's names and how frequently they arise. -/
  nameCounts : Multiset Name := ∅
  bigramCounts : Multiset Bigram := ∅
  trigramCounts : Multiset Trigram := ∅

instance : ForIn M (Multiset α) (α × Nat) :=
  show ForIn _ (Std.RBMap _ _ _) _ by infer_instance

def Multiset.toList (m : Multiset α) : List α :=
  m.foldl (fun l x _ => x :: l) []

def Multiset.toHFeatures [ToString α] (m : Multiset α) : Array String :=
  Array.mk <| m.toList.map (s!"H:{·}")

def Multiset.toTFeatures [ToString α] (m : Multiset α) : Array String :=
  Array.mk <| m.toList.map (s!"T:{·}")

instance [ToJson α] : ToJson (Multiset α) where
  toJson m := Json.arr (Array.mk (m.toList.map toJson))

instance : EmptyCollection StatementFeatures := ⟨{}⟩
instance : Append StatementFeatures where
  append x y := {
    nameCounts := x.nameCounts ++ y.nameCounts
    bigramCounts := x.bigramCounts ++ y.bigramCounts
    trigramCounts := x.trigramCounts ++ y.trigramCounts
  }

instance : ToJson StatementFeatures where
  toJson f := Json.mkObj [
    ("nameCounts", toJson f.nameCounts),
    ("bigramCounts", toJson f.bigramCounts),
    ("trigramCounts", toJson f.trigramCounts)
  ]

def StatementFeatures.mkName : Name → StatementFeatures
  | n => {nameCounts := Multiset.singleton n}

def StatementFeatures.mkBigram : Name → Name → StatementFeatures
  | n1, n2 => {bigramCounts := Multiset.singleton ⟨n1, n2⟩}

def StatementFeatures.mkTrigram : Name → Name → Name → StatementFeatures
  | n1, n2, n3 => {trigramCounts := Multiset.singleton ⟨n1, n2, n3⟩}

def StatementFeatures.toHFeatures (f : StatementFeatures) : Array String :=
  f.nameCounts.toHFeatures ++
  f.bigramCounts.toHFeatures ++
  f.trigramCounts.toHFeatures

def StatementFeatures.toTFeatures (f : StatementFeatures) : Array String :=
  f.nameCounts.toTFeatures ++
  f.bigramCounts.toTFeatures ++
  f.trigramCounts.toTFeatures

def immediateName (e : Expr) : Option Name :=
  if let .const n _ := e then
    some n
  else if let some n := e.natLit? then
    some <| toString n
  else
    none

def getHeadName? (e : Expr) : Option Name := do
  immediateName <| e.getAppFn

def visitFeature (e : Expr) : WriterT StatementFeatures MetaM Unit  := do
  --let ppe ← Lean.PrettyPrinter.ppExpr e
  if let some n := immediateName e then
    tell <| StatementFeatures.mkName n
  if e.isApp then
    e.withApp (fun f args => do
      if let some n1 := immediateName f then
        for arg in args do
          if let some n2 := getHeadName? arg then
            tell <| StatementFeatures.mkBigram n1 n2
            if arg.isApp then
              arg.withApp (fun f args => do
                if let some n2 := immediateName f then
                  for arg in args do
                    if let some n3 := getHeadName? arg then
                      tell <| StatementFeatures.mkTrigram n1 n2 n3
              )
        for p in args.toList.allPairs do
          if let some n2 := getHeadName? p.1 then
            if let some n3 := getHeadName? p.2 then
              tell <| StatementFeatures.mkTrigram n1 n2 n3
    )
  return ()

def getStatementFeatures (e : Expr) : MetaM StatementFeatures := do
  let ((), features) ← WriterT.run <| forEachExpr visitFeature e
  return features

open Lean.Meta

def getArgsFeatures (args : List Expr) : MetaM (Array StatementFeatures) := do
  let mut argsFeats := #[]
  for arg in args do
    let argType ← inferType arg
    if (← inferType argType).isProp then
      let argFeats ← getStatementFeatures argType
      if ! argFeats.nameCounts.isEmpty then
        argsFeats := argsFeats ++ #[argFeats]
  return argsFeats

def getThmAndArgsFeatures (e : Expr)
  : MetaM (StatementFeatures × Array StatementFeatures) := do
  forallTelescope e <| fun args thm => do
      let thmFeats ← getStatementFeatures thm
      let argsFeats ← getArgsFeatures args.data
      return (thmFeats, argsFeats)

end PremiseSelection