import ReactorModel.Objects.Reactor.Theorems.LawfulCoe

noncomputable section
open Classical

namespace ReactorType

inductive Equivalent [ReactorType α] : α → α → Prop
  | intro
    (mem_cpt?_iff : ∀ cpt j, (j ∈ cpt? cpt rtr₁) ↔ (j ∈ cpt? cpt rtr₂))
    (rcns_some_eq : ∀ {i r₁ r₂}, (rcns rtr₁ i = some r₁) → (rcns rtr₂ i = some r₂) → r₁ = r₂)
    (nest_equiv : ∀ {i n₁ n₂}, (nest rtr₁ i = some n₁) → (nest rtr₂ i = some n₂) → Equivalent n₁ n₂)
    : Equivalent rtr₁ rtr₂

namespace Equivalent

instance [ReactorType α] : HasEquiv α where
  Equiv := Equivalent

@[refl]
protected theorem refl [ReactorType.WellFounded α] {rtr : α} : rtr ≈ rtr := by
  induction rtr using ReactorType.WellFounded.induction
  case nest hi =>
    constructor <;> (intros; simp_all)
    exact hi _ _ ‹_›

variable [ReactorType α] {rtr rtr₁ : α}

@[symm]
protected theorem symm (e : rtr₁ ≈ rtr₂) : rtr₂ ≈ rtr₁ := by
  induction e
  case intro h₁ h₂ _ hi =>
    constructor <;> intros
    · exact h₁ ‹_› ‹_› |>.symm
    · exact h₂ ‹_› ‹_› |>.symm
    · exact hi ‹_› ‹_›

@[trans]
protected theorem trans (e₁ : rtr₁ ≈ rtr₂) (e₂ : rtr₂ ≈ rtr₃) : rtr₁ ≈ rtr₃ := by
  induction e₁ generalizing rtr₃; cases e₂
  case intro.intro h₁ h₂ _ hi h₁' h₂' h₃' =>
    constructor
    · intros; exact h₁ ‹_› ‹_› |>.trans (h₁' ‹_› ‹_›)
    · intro _ _ _ h _
      have ⟨_, h⟩ := Partial.mem_iff.mp <| h₁ .rcn ‹_› |>.mp $ Partial.mem_iff.mpr ⟨_, h⟩
      exact h₂ ‹_› h |>.trans (h₂' h ‹_›)
    · intro _ _ _ h _
      have ⟨_, h⟩ := Partial.mem_iff.mp <| h₁ .rtr ‹_› |>.mp $ Partial.mem_iff.mpr ⟨_, h⟩
      exact hi ‹_› h (h₃' h ‹_›)

theorem mem_cpt?_iff : (rtr₁ ≈ rtr₂) → (i ∈ cpt? cpt rtr₁ ↔ i ∈ cpt? cpt rtr₂)
  | intro h .. => h _ _

theorem rcns_some_eq : (rtr₁ ≈ rtr₂) → (rcns rtr₁ i = some r₁) → (rcns rtr₂ i = some r₂) → r₁ = r₂
  | intro _ h .. => h

theorem nest_equiv : (rtr₁ ≈ rtr₂) → (nest rtr₁ i = some n₁) → (nest rtr₂ i = some n₂) → n₁ ≈ n₂
  | intro _ _ h => h

theorem rcns_eq (e : rtr₁ ≈ rtr₂) : rcns rtr₂ = rcns rtr₁ := by
  funext i
  by_cases h₁ : i ∈ rcns rtr₁
  case pos =>
    have ⟨_, h₂⟩ := Partial.mem_iff.mp $ mem_cpt?_iff e (cpt := .rcn) |>.mp h₁
    have ⟨_, h₁⟩ := Partial.mem_iff.mp h₁
    exact rcns_some_eq e h₁ h₂ ▸ h₁ |>.symm ▸ h₂
  case neg =>
    have h₂ := Partial.mem_iff.not.mp $ mem_cpt?_iff e (cpt := .rcn) |>.not.mp h₁
    have h₁ := Partial.mem_iff.not.mp h₁
    simp [cpt?] at h₁ h₂
    simp [Option.eq_none_iff_forall_not_mem.mpr h₁, Option.eq_none_iff_forall_not_mem.mpr h₂]

theorem cpt?_some_iff (e : rtr₁ ≈ rtr₂) :
    (∃ o₁, cpt? cpt rtr₁ i = some o₁) ↔ (∃ o₂, cpt? cpt rtr₂ i = some o₂) := by
  simp [←Partial.mem_iff, mem_cpt?_iff e]

variable [Indexable α] {rtr₁ : α}

theorem obj?_rcn_eq (e : rtr₁ ≈ rtr₂) : rtr₁[.rcn] = rtr₂[.rcn] :=
  sorry

theorem mem_iff {i} (e : rtr₁ ≈ rtr₂) : (i ∈ rtr₁[cpt]) ↔ (i ∈ rtr₂[cpt]) := by
  sorry

theorem obj?_rtr_equiv (e : rtr₁ ≈ rtr₂) (h₁ : rtr₁[.rtr][i] = some n₁) (h₂ : rtr₂[.rtr][i] = some n₂) :
    n₁ ≈ n₂ := by
  sorry

theorem obj?_some_iff (e : rtr₁ ≈ rtr₂) :
    (∃ o₁, rtr₁[cpt][i] = some o₁) ↔ (∃ o₂, rtr₂[cpt][i] = some o₂) :=
  sorry

end Equivalent

theorem LawfulMemUpdate.equiv [ReactorType.WellFounded α] {rtr₁ : α}
    (u : LawfulMemUpdate cpt i f rtr₁ rtr₂) : rtr₁ ≈ rtr₂ := by
  induction u <;> constructor
  case final.mem_cpt?_iff e h₁ h₂ =>
    intro c j
    by_cases hc : c = cpt <;> try subst hc
    case neg => exact e.mem_iff (.inl hc)
    case pos =>
      by_cases hj : j = i <;> try subst hj
      case neg => exact e.mem_iff (.inr hj)
      case pos => simp [Partial.mem_iff, h₁, h₂]
  case final.rcns_some_eq e _ _ =>
    intro j _ _ h₁ h₂
    have h := e (c := .rcn) (j := j) (.inl $ by simp)
    simp_all [cpt?]
  case final.nest_equiv e _ _ =>
    intro j _ _ h₁ h₂
    have h := e (c := .rtr) (j := j) (.inl $ by simp)
    simp_all [cpt?]
    exact .refl
  case nest.mem_cpt?_iff j _ _ _ _ e h₁ h₂ _ _ =>
    intro c j'
    by_cases hc : c = .rtr <;> try subst hc
    case neg => exact e.mem_iff (.inl hc)
    case pos =>
      by_cases hj : j' = j <;> try subst hj
      case neg => exact e.mem_iff (.inr hj)
      case pos => simp [Partial.mem_iff, h₁, h₂]
  case nest.rcns_some_eq e h₁ h₂ _ _ =>
    intro j _ _ h₁ h₂
    have h := e (c := .rcn) (j := j) (.inl $ by simp)
    simp_all [cpt?]
  case nest.nest_equiv j _ _ _ _ e _ _ _ hi =>
    intro j' n₁' n₂' h₁' h₂'
    by_cases hj : j' = j <;> try subst hj
    case pos => simp_all [cpt?]; assumption
    case neg =>
      have := e (c := .rtr) (j := j') (.inr hj)
      simp_all [cpt?]
      exact .refl

theorem LawfulUpdate.equiv [ReactorType.WellFounded α] {rtr₁ : α} :
    (LawfulUpdate cpt i f rtr₁ rtr₂) → rtr₁ ≈ rtr₂
  | notMem _ h => h ▸ .refl
  | update u   => u.equiv

theorem LawfulUpdatable.equiv [LawfulUpdatable α] {rtr : α} :
    (Updatable.update rtr cpt i f) ≈ rtr :=
  Equivalent.symm (lawful rtr cpt i f).equiv

namespace Member

inductive Equivalent [ReactorType α] [ReactorType β] :
    {rtr₁ : α} → {rtr₂ : β} → (Member cpt i rtr₁) → (Member cpt i rtr₂) → Prop
  | final : Equivalent (.final h₁) (.final h₂)
  | nest {n₁ : α} {n₂ : β} {m₁ : Member cpt i n₁} {m₂ : Member cpt i n₂} :
    (h₁ : ReactorType.nest rtr₁ j = some n₁) → (h₂ : ReactorType.nest rtr₂ j = some n₂) →
    (Equivalent m₁ m₂) → Equivalent (.nest h₁ m₁) (.nest h₂ m₂)

namespace Equivalent

@[refl]
theorem refl [ReactorType.WellFounded α] {rtr : α} {m : Member cpt i rtr} :
    Equivalent m m := by
  induction rtr using ReactorType.WellFounded.induction
  case nest hi =>
    cases m
    case final  => exact .final
    case nest h => exact .nest _ _ (hi _ ⟨_, h⟩)

variable [ReactorType α] [ReactorType β]

theorem symm {rtr₁ : α} {rtr₂ : β} {m₁ : Member cpt i rtr₁} {m₂ : Member cpt i rtr₂}
    (e : Equivalent m₁ m₂) : (Equivalent m₂ m₁) := by
  induction e <;> constructor; assumption

theorem trans
    [ReactorType γ] {rtr₁ : α} {rtr₂ : β} {rtr₃ : γ}
    {m₁ : Member cpt i rtr₁} {m₂ : Member cpt i rtr₂} {m₃ : Member cpt i rtr₃}
    (e₁ : Equivalent m₁ m₂) (e₂ : Equivalent m₂ m₃) : (Equivalent m₁ m₃) := by
  induction e₁ generalizing m₃ rtr₃ <;> cases e₂ <;> constructor
  case nest.nest hi₁ _ _ _ _ hi₂ => exact hi₁ hi₂

-- Lemma for `to_eq`.
private theorem to_eq' {rtr₁ rtr₂ : α} {m₁ : Member cpt i rtr₁} {m₂ : Member cpt i rtr₂}
    (h : rtr₁ = rtr₂) (e : Equivalent m₁ m₂) : m₁ = cast (by simp [h]) m₂ := by
  induction e <;> subst h
  case final => rfl
  case nest m₁ _ h₁ _ hi h₂ =>
    injection h₁ ▸ h₂ with h
    simp [hi h, h]

theorem to_eq {rtr : α} {m₁ m₂ : Member cpt i rtr} (e : Equivalent m₁ m₂) : m₁ = m₂ :=
  e.to_eq' rfl

theorem from_lawfulCoe [ReactorType α] [ReactorType β] [LawfulCoe α β] {rtr : α}
    (m : Member cpt i rtr) : Equivalent m (m : Member cpt i (rtr : β)) := by
  induction m
  case final => constructor
  case nest e => simp [fromLawfulCoe, Equivalent.nest _ _ e]

end Equivalent

variable [ReactorType.WellFounded α] {rtr₁ : α}

def fromLawfulMemUpdate {rtr₁ : α} :
    (Member c j rtr₂) → (LawfulMemUpdate cpt i f rtr₁ rtr₂) → Member c j rtr₁
  | final h, u => final (Equivalent.mem_cpt?_iff u.equiv |>.mpr h)
  | nest h m (j := j), .final e _ _ =>
    nest (m := m) $ by
      have h' := e (c := .rtr) (j := j) (.inl $ by simp)
      simp [cpt?] at h'
      exact h'.symm ▸ h
  | nest h m (j := j₂), .nest e h₁ h₂ u (j := j₁) =>
      if hj : j₂ = j₁ then
        let m' := (hj ▸ h |>.symm.trans h₂ |> Option.some_inj.mp) ▸ m
        nest h₁ $ fromLawfulMemUpdate m' u
      else
        nest (m := m) $ by
          have h' := e (c := .rtr) (.inr hj)
          simp [cpt?] at h'
          exact h'.symm ▸ h

def fromLawfulUpdate (m : Member c j rtr₂) : (LawfulUpdate cpt i f rtr₁ rtr₂) → Member c j rtr₁
  | .notMem _ h => h ▸ m
  | .update u   => m.fromLawfulMemUpdate u

theorem Equivalent.from_lawfulMemUpdate (u : LawfulMemUpdate cpt i f rtr₁ rtr₂)
    (m : Member c j rtr₂) : Equivalent m (m.fromLawfulMemUpdate u) := by
  induction u <;> cases m <;> (simp [fromLawfulMemUpdate]; try exact .final)
  case final.nest e _ _ j _ _ hn =>
    have h := e (c := .rtr) (j := j) (.inl $ by simp)
    simp [cpt?] at h
    exact .nest hn (h ▸ hn) .refl
  case nest.nest e h₁ h₂ _ hi _ _ m hn =>
    split
    case inl hj =>
      subst hj
      cases Option.some_inj.mp $ hn.symm.trans h₂
      exact .nest hn h₁ (hi m)
    case inr hj =>
      have h := e (c := .rtr) (.inr hj)
      simp [cpt?] at h
      exact .nest hn (h.symm ▸ hn) .refl

theorem Equivalent.from_lawfulUpdate (u : LawfulUpdate cpt i f rtr₁ rtr₂)
    (m : Member c j rtr₂) : Equivalent m (m.fromLawfulUpdate u) := by
  cases u
  case notMem _ h => cases h; rfl
  case update u   => exact Equivalent.from_lawfulMemUpdate u m

end Member

theorem UniqueIDs.lift [ReactorType α] [ReactorType β] [LawfulCoe α β] {rtr : α}
    (h : UniqueIDs (rtr : β)) : UniqueIDs rtr where
  allEq m₁ m₂ :=
    h.allEq (.fromLawfulCoe m₁) (.fromLawfulCoe m₂) ▸ Member.Equivalent.from_lawfulCoe m₁
      |>.trans (Member.Equivalent.from_lawfulCoe m₂).symm
      |>.to_eq

instance [LawfulUpdatable α] [ind : Indexable β] [LawfulCoe α β] : Indexable α where
  unique_ids := UniqueIDs.lift ind.unique_ids

open Equivalent
variable [Indexable α] [Indexable β] {rtr rtr₁ : α}

namespace Dependency

theorem equiv (e : rtr₁ ≈ rtr₂) (d : j₁ <[rtr₂] j₂) : j₁ <[rtr₁] j₂ := by
  induction d with
  | prio h₁ h₂ h₃ =>
    -- TODO: The next 2 lines are a common pattern in the `updated` proofs. Perhaps create a
    --       (unidirectional) derivative of `Equivalent.obj?_some_iff` that includes equivalence.
    have ⟨_, h₁'⟩ := obj?_some_iff e |>.mpr ⟨_, h₁⟩
    have e := Equivalent.obj?_rtr_equiv e h₁' h₁
    exact prio h₁' (rcns_eq e ▸ h₂) (rcns_eq e ▸ h₃) ‹_› ‹_›
  | mutNorm h₁ h₂ h₃ =>
    have ⟨_, h₁'⟩ := obj?_some_iff e |>.mpr ⟨_, h₁⟩
    have e := Equivalent.obj?_rtr_equiv e h₁' h₁
    exact mutNorm h₁' (rcns_eq e ▸ h₂) (rcns_eq e ▸ h₃) ‹_› ‹_›
  | depOverlap h₁ h₂ =>
    exact depOverlap (e.obj?_rcn_eq.symm ▸ h₁) (e.obj?_rcn_eq.symm ▸ h₂) ‹_› ‹_› ‹_›
  | mutNest h₁ h₂ h₃ _ h₄ =>
    have ⟨_, h₁'⟩ := e.obj?_some_iff.mpr ⟨_, h₁⟩
    have e := Equivalent.obj?_rtr_equiv e h₁' h₁
    have ⟨_, h₂'⟩ := cpt?_some_iff e (cpt := .rtr) |>.mpr ⟨_, h₂⟩
    have h₄' := mem_cpt?_iff (Equivalent.nest_equiv e h₂' h₂) (cpt := .rcn) |>.mpr h₄
    exact mutNest h₁' h₂' (rcns_eq e ▸ h₃) ‹_› h₄'
  | trans _ _ d₁ d₂ =>
    exact trans d₁ d₂

theorem Acyclic.equiv (e : rtr₁ ≈ rtr₂) (a : Acyclic rtr₁) : Acyclic rtr₂ :=
  fun i d => absurd (d.equiv e) (a i)

end Dependency

namespace Wellformed

set_option hygiene false in
scoped macro "equiv_nested_proof " name:ident : term => `(
  fun hc hp =>
    have e := Equivalent.obj?_rtr_equiv ‹_› h₁ h₂
    have ⟨_, hc'⟩ := Equivalent.cpt?_some_iff e (cpt := .rtr) |>.mp ⟨_, hc⟩
    have e := Equivalent.nest_equiv e hc hc'
    $(Lean.mkIdentFrom name $ `ValidDependency ++ name.getId) hc'
    (Equivalent.mem_cpt?_iff e (cpt := .prt _) |>.mp hp)
)

theorem ValidDependency.equiv
    (e : rtr₁ ≈ rtr₂) (h₁ : rtr₁[.rtr][j] = some con₁) (h₂ : rtr₂[.rtr][j] = some con₂) :
    (ValidDependency con₁ rk dk d) → ValidDependency con₂ rk dk d
  | stv h           => stv $ mem_cpt?_iff (obj?_rtr_equiv e h₁ h₂) (cpt := .stv) |>.mp h
  | act h           => act $ mem_cpt?_iff (obj?_rtr_equiv e h₁ h₂) (cpt := .act) |>.mp h
  | prt h           => prt $ mem_cpt?_iff (obj?_rtr_equiv e h₁ h₂) (cpt := .prt _) |>.mp h
  | nestedIn hc hp  => (equiv_nested_proof nestedIn) hc hp
  | nestedOut hc hp => (equiv_nested_proof nestedOut) hc hp

set_option hygiene false in
scoped macro "equiv_prio_proof " name:ident rtr₁:ident rtr₂:ident : term => `(
  fun h₁ h₂ h₃ =>
    have ⟨_, h₁'⟩ := Equivalent.obj?_some_iff ‹$rtr₁ ≈ $rtr₂› |>.mpr ⟨_, h₁⟩
    have e := Equivalent.obj?_rtr_equiv ‹_› h₁' h₁
    $(Lean.mkIdentFrom name $ `Wellformed ++ name.getId)
      ‹_› h₁' (Equivalent.rcns_eq e ▸ h₂) (Equivalent.rcns_eq e ▸ h₃)
)

theorem equiv (e : rtr₁ ≈ rtr₂) (wf : Wellformed rtr₁) : Wellformed rtr₂ where
  overlap_prio  := equiv_prio_proof overlap_prio rtr₁ rtr₂
  hazards_prio  := equiv_prio_proof hazards_prio rtr₁ rtr₂
  mutation_prio := equiv_prio_proof mutation_prio rtr₁ rtr₂
  acyclic_deps  := wf.acyclic_deps.equiv e
  valid_deps h₁ h₂ h₃ :=
    have ⟨_, h₁'⟩ := Equivalent.obj?_some_iff e |>.mpr ⟨_, h₁⟩
    have e := Equivalent.obj?_rtr_equiv e h₁' h₁
    have h₂' := Equivalent.rcns_eq e ▸ h₂
    wf.valid_deps h₁' h₂' h₃ |>.equiv ‹_› h₁' h₁
  unique_inputs h₁ h₂ _ h₃ :=
    have h₃' := Equivalent.mem_iff e |>.mpr h₃
    wf.unique_inputs (e.obj?_rcn_eq.symm ▸ h₁) (e.obj?_rcn_eq.symm ▸ h₂) ‹_› h₃'

end Wellformed
end ReactorType


import category_theory.preadditive.basic
import category_theory.abelian.projective
import category_theory.abelian.diagram_lemmas.four

import data.matrix.notation

import .abelian_category
import .fin_functor
import .split_exact

noncomputable theory

open category_theory
open category_theory.limits
open category_theory.preadditive

universes v u

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

@[ext]
structure short_exact_sequence [has_images 𝒞] [has_zero_morphisms 𝒞] [has_kernels 𝒞] :=
(fst snd trd : 𝒞)
(f : fst ⟶ snd)
(g : snd ⟶ trd)
[mono'  : mono f]
[epi'   : epi g]
(exact' : exact f g)

namespace short_exact_sequence

attribute [instance] mono' epi'

variables {𝒞} [has_images 𝒞] [has_zero_morphisms 𝒞] [has_kernels 𝒞]

@[simp, reassoc] lemma f_comp_g (A : short_exact_sequence 𝒞) : A.f ≫ A.g = 0 := A.exact'.w

@[ext]
structure hom (A B : short_exact_sequence 𝒞) :=
(fst : A.1 ⟶ B.1)
(snd : A.2 ⟶ B.2)
(trd : A.3 ⟶ B.3)
(sq1' : fst ≫ B.f = A.f ≫ snd . obviously)
(sq2' : snd ≫ B.g = A.g ≫ trd . obviously)

namespace hom

restate_axiom sq1' sq1
restate_axiom sq2' sq2

attribute [reassoc] sq1 sq2

end hom

instance : quiver (short_exact_sequence 𝒞) := ⟨hom⟩

def id (A : short_exact_sequence 𝒞) : A ⟶ A :=
{ fst := 𝟙 _,
  snd := 𝟙 _,
  trd := 𝟙 _,
  sq1' := by simp only [category.id_comp, category.comp_id],
  sq2' := by simp only [category.id_comp, category.comp_id], }

def comp {A B C : short_exact_sequence 𝒞} (f : A ⟶ B) (g : B ⟶ C) : A ⟶ C :=
{ fst := f.1 ≫ g.1,
  snd := f.2 ≫ g.2,
  trd := f.3 ≫ g.3,
  sq1' := by rw [category.assoc, hom.sq1, hom.sq1_assoc],
  sq2' := by rw [category.assoc, hom.sq2, hom.sq2_assoc], }

instance : category (short_exact_sequence 𝒞) :=
{ id := id,
  comp := λ A B C f g, comp f g,
  id_comp' := by { intros, ext; dsimp; apply category.id_comp, },
  comp_id' := by { intros, ext; dsimp; apply category.comp_id, },
  assoc' := by { intros, ext; dsimp; apply category.assoc, },
  .. (infer_instance : quiver (short_exact_sequence 𝒞)) }

@[simp] lemma id_fst (A : short_exact_sequence 𝒞) : hom.fst (𝟙 A) = 𝟙 A.1 := rfl
@[simp] lemma id_snd (A : short_exact_sequence 𝒞) : hom.snd (𝟙 A) = 𝟙 A.2 := rfl
@[simp] lemma id_trd (A : short_exact_sequence 𝒞) : hom.trd (𝟙 A) = 𝟙 A.3 := rfl

variables {A B C : short_exact_sequence 𝒞} (f : A ⟶ B) (g : B ⟶ C)

@[simp, reassoc] lemma comp_fst : (f ≫ g).1 = f.1 ≫ g.1 := rfl
@[simp, reassoc] lemma comp_snd : (f ≫ g).2 = f.2 ≫ g.2 := rfl
@[simp, reassoc] lemma comp_trd : (f ≫ g).3 = f.3 ≫ g.3 := rfl

variables (𝒞)

@[simps] def Fst : short_exact_sequence 𝒞 ⥤ 𝒞 :=
{ obj := fst, map := λ A B f, f.1 }

@[simps] def Snd : short_exact_sequence 𝒞 ⥤ 𝒞 :=
{ obj := snd, map := λ A B f, f.2 }

@[simps] def Trd : short_exact_sequence 𝒞 ⥤ 𝒞 :=
{ obj := trd, map := λ A B f, f.3 }

@[simps] def f_nat : Fst 𝒞 ⟶ Snd 𝒞 :=
{ app := λ A, A.f,
  naturality' := λ A B f, f.sq1 }

@[simps] def g_nat : Snd 𝒞 ⟶ Trd 𝒞 :=
{ app := λ A, A.g,
  naturality' := λ A B f, f.sq2 }

instance : has_zero_morphisms (short_exact_sequence 𝒞) :=
{ has_zero := λ A B, ⟨{ fst := 0, snd := 0, trd := 0 }⟩,
  comp_zero' := by { intros, ext; apply comp_zero },
  zero_comp' := by { intros, ext; apply zero_comp }, }
.

@[simp] lemma hom_zero_fst : (0 : A ⟶ B).1 = 0 := rfl

@[simp] lemma hom_zero_snd : (0 : A ⟶ B).2 = 0 := rfl

@[simp] lemma hom_zero_trd : (0 : A ⟶ B).3 = 0 := rfl

variables {𝒞}

protected def functor (A : short_exact_sequence 𝒞) : fin 3 ⥤ 𝒞 :=
fin3_functor_mk ![A.1, A.2, A.3] A.f A.g

def functor_map {A B : short_exact_sequence 𝒞} (f : A ⟶ B) :
  Π i, A.functor.obj i ⟶ B.functor.obj i
| ⟨0,h⟩ := f.1
| ⟨1,h⟩ := f.2
| ⟨2,h⟩ := f.3
| ⟨i+3,hi⟩ := by { exfalso, revert hi, dec_trivial }

meta def aux_tac : tactic unit :=
`[simp only [hom_of_le_refl, functor.map_id, category.id_comp, category.comp_id]]

lemma functor_map_naturality {A B : short_exact_sequence 𝒞} (f : A ⟶ B) :
  ∀ (i j : fin 3) (hij : i ≤ j),
    functor_map f i ≫ B.functor.map hij.hom = A.functor.map hij.hom ≫ functor_map f j
| ⟨0,hi⟩ ⟨0,hj⟩ hij := by aux_tac
| ⟨1,hi⟩ ⟨1,hj⟩ hij := by aux_tac
| ⟨2,hi⟩ ⟨2,hj⟩ hij := by aux_tac
| ⟨0,hi⟩ ⟨1,hj⟩ hij := f.sq1
| ⟨1,hi⟩ ⟨2,hj⟩ hij := f.sq2
| ⟨i+3,hi⟩ _ _ := by { exfalso, revert hi, dec_trivial }
| _ ⟨j+3,hj⟩ _ := by { exfalso, revert hj, dec_trivial }
| ⟨i+1,hi⟩ ⟨0,hj⟩ H := by { exfalso, revert H, dec_trivial }
| ⟨i+2,hi⟩ ⟨1,hj⟩ H := by { exfalso, revert H, dec_trivial }
| ⟨0,hi⟩ ⟨2,hj⟩ hij :=
begin
  have h01 : (0 : fin 3) ⟶ 1 := hom_of_le dec_trivial,
  have h12 : (1 : fin 3) ⟶ 2 := hom_of_le dec_trivial,
  calc functor_map f ⟨0, hi⟩ ≫ B.functor.map hij.hom
      = functor_map f ⟨0, hi⟩ ≫ B.functor.map h01 ≫ B.functor.map h12 : _
  ... = (functor_map f ⟨0, hi⟩ ≫ B.functor.map h01) ≫ B.functor.map h12 : by rw category.assoc
  ... = (A.functor.map h01 ≫ functor_map f _) ≫ B.functor.map h12 : _
  ... = A.functor.map h01 ≫ functor_map f _ ≫ B.functor.map h12 : category.assoc _ _ _
  ... = A.functor.map h01 ≫ A.functor.map h12 ≫ functor_map f _ : _
  ... = A.functor.map hij.hom ≫ functor_map f ⟨2, hj⟩ : _,
  { rw [← functor.map_comp], congr, },
  { congr' 1, exact f.sq1 },
  { congr' 1, exact f.sq2 },
  { rw [← functor.map_comp_assoc], congr, },
end

@[simps] def Functor : short_exact_sequence 𝒞 ⥤ fin 3 ⥤ 𝒞 :=
{ obj := short_exact_sequence.functor,
  map := λ A B f,
  { app := functor_map f,
    naturality' := λ i j hij, (functor_map_naturality f i j hij.le).symm },
  map_id' := λ A, by { ext i, fin_cases i; refl },
  map_comp' := λ A B C f g, by { ext i, fin_cases i; refl } }

end short_exact_sequence

namespace short_exact_sequence

variables {𝒞} [abelian 𝒞]
variables {A B C : short_exact_sequence 𝒞} (f : A ⟶ B) (g : B ⟶ C)

section iso

variables {A B C} (f g)

open_locale zero_object

/-- One form of the five lemma: if a morphism of short exact sequences has isomorphisms
as first and third component, then the second component is also an isomorphism. -/
lemma snd_is_iso (h1 : is_iso f.1) (h3 : is_iso f.3) : is_iso f.2 :=
@abelian.is_iso_of_is_iso_of_is_iso_of_is_iso_of_is_iso 𝒞 _ _
  0 A.1 A.2 A.3
  0 B.1 B.2 B.3
  0 A.f A.g
  0 B.f B.g
  0 f.1 f.2 f.3 (by rw [zero_comp, zero_comp]) f.sq1 f.sq2
  0 0
  0 0 0 (by rw [comp_zero, comp_zero])
  (exact_zero_left_of_mono _)
  A.exact'
  ((epi_iff_exact_zero_right _).mp infer_instance)
  (exact_zero_left_of_mono _)
  B.exact'
  ((epi_iff_exact_zero_right _).mp infer_instance) _ _ _ _

/-- One form of the five lemma: if a morphism `f` of short exact sequences has isomorphisms
as first and third component, then `f` itself is an isomorphism. -/
lemma is_iso_of_fst_of_trd (h1 : is_iso f.1) (h3 : is_iso f.3) : is_iso f :=
{ out :=
  begin
    haveI : is_iso f.2 := snd_is_iso f h1 h3,
    refine ⟨⟨inv f.1, inv f.2, inv f.3, _, _⟩, _, _⟩,
    { dsimp, simp only [is_iso.inv_comp_eq, f.sq1_assoc, category.comp_id, is_iso.hom_inv_id], },
    { dsimp, simp only [is_iso.inv_comp_eq, f.sq2_assoc, category.comp_id, is_iso.hom_inv_id], },
    { ext; dsimp; simp only [is_iso.hom_inv_id], },
    { ext; dsimp; simp only [is_iso.inv_hom_id], },
  end }

@[simps] def iso_of_components (f₁ : A.1 ≅ B.1) (f₂ : A.2 ≅ B.2) (f₃ : A.3 ≅ B.3)
  (sq1 : f₁.hom ≫ B.f = A.f ≫ f₂.hom) (sq2 : f₂.hom ≫ B.g = A.g ≫ f₃.hom) :
  A ≅ B :=
{ hom := ⟨f₁.hom, f₂.hom, f₃.hom, sq1, sq2⟩,
  inv :=
  begin
    refine ⟨f₁.inv, f₂.inv, f₃.inv, _, _⟩; dsimp,
    rw [iso.inv_comp_eq, ← category.assoc, iso.eq_comp_inv, sq1],
    rw [iso.inv_comp_eq, ← category.assoc, iso.eq_comp_inv, sq2],
  end,
  hom_inv_id' := by { ext; apply iso.hom_inv_id, },
  inv_hom_id' := by { ext; apply iso.inv_hom_id, } }

@[simps] def iso_of_components' (f₁ : A.1 ≅ B.1) (f₂ : A.2 ⟶ B.2) (f₃ : A.3 ≅ B.3)
  (sq1 : f₁.hom ≫ B.f = A.f ≫ f₂) (sq2 : f₂ ≫ B.g = A.g ≫ f₃.hom) :
  A ≅ B :=
let F : A ⟶ B := ⟨f₁.hom, f₂, f₃.hom, sq1, sq2⟩ in
{ hom := F,
  inv :=
  begin
    haveI : is_iso F.2 := snd_is_iso _ infer_instance infer_instance,
    refine ⟨f₁.inv, inv F.2, f₃.inv, _, _⟩; dsimp,
    rw [iso.inv_comp_eq, ← category.assoc, is_iso.eq_comp_inv, sq1],
    rw [is_iso.inv_comp_eq, ← category.assoc, iso.eq_comp_inv, sq2],
  end,
  hom_inv_id' := by { ext; try { apply iso.hom_inv_id, }, apply is_iso.hom_inv_id },
  inv_hom_id' := by { ext; try { apply iso.inv_hom_id, }, apply is_iso.inv_hom_id } }

end iso

section split

/-- A short exact sequence `0 ⟶ A₁ -f⟶ A₂ -g⟶ A₃ ⟶ 0` is *left split*
if there exists a morphism `φ : A₂ ⟶ A₁` such that `f ≫ φ = 𝟙 A₁`. -/
def left_split (A : short_exact_sequence 𝒞) : Prop :=
∃ φ : A.2 ⟶ A.1, A.f ≫ φ = 𝟙 A.1

/-- A short exact sequence `0 ⟶ A₁ -f⟶ A₂ -g⟶ A₃ ⟶ 0` is *right split*
if there exists a morphism `φ : A₂ ⟶ A₁` such that `f ≫ φ = 𝟙 A₁`. -/
def right_split (A : short_exact_sequence 𝒞) : Prop :=
∃ χ : A.3 ⟶ A.2, χ ≫ A.g = 𝟙 A.3

variables {𝒜 : Type*} [category 𝒜] [abelian 𝒜]

lemma exact_of_split {A B C : 𝒜} (f : A ⟶ B) (g : B ⟶ C) (χ : C ⟶ B) (φ : B ⟶ A)
  (hfg : f ≫ g = 0) (H : φ ≫ f + g ≫ χ = 𝟙 B) : exact f g :=
{ w := hfg,
  epi :=
  begin
    let ψ : (kernel_subobject g : 𝒜) ⟶ image_subobject f :=
      subobject.arrow _ ≫ φ ≫ factor_thru_image_subobject f,
    suffices : ψ ≫ image_to_kernel f g hfg = 𝟙 _,
    { convert epi_of_epi ψ _, rw this, apply_instance },
    rw ← cancel_mono (subobject.arrow _), swap, { apply_instance },
    simp only [image_to_kernel_arrow, image_subobject_arrow_comp, category.id_comp, category.assoc],
    calc (kernel_subobject g).arrow ≫ φ ≫ f
        = (kernel_subobject g).arrow ≫ 𝟙 B : _
    ... = (kernel_subobject g).arrow        : category.comp_id _,
    rw [← H, preadditive.comp_add],
    simp only [add_zero, zero_comp, kernel_subobject_arrow_comp_assoc],
  end }

-- move this
lemma exact_inl_snd (A B : 𝒜) : exact (biprod.inl : A ⟶ A ⊞ B) biprod.snd :=
exact_of_split _ _ biprod.inr biprod.fst biprod.inl_snd biprod.total

def mk_of_split {A B C : 𝒜} (f : A ⟶ B) (g : B ⟶ C) (φ : B ⟶ A) (χ : C ⟶ B)
  (hfg : f ≫ g = 0) (hφ : f ≫ φ = 𝟙 A) (hχ : χ ≫ g = 𝟙 C) (H : φ ≫ f + g ≫ χ = 𝟙 B) :
  short_exact_sequence 𝒜 :=
{ fst := A,
  snd := B,
  trd := C,
  f := f,
  g := g,
  mono' := by { haveI : mono (f ≫ φ), { rw hφ, apply_instance }, exact mono_of_mono f φ, },
  epi' := by { haveI : epi (χ ≫ g), { rw hχ, apply_instance }, exact epi_of_epi χ g, },
  exact' := exact_of_split f g χ φ hfg H }

def mk_of_split' {A B C : 𝒜} (f : A ⟶ B) (g : B ⟶ C)
  (H : ∃ (φ : B ⟶ A) (χ : C ⟶ B), f ≫ g = 0 ∧ f ≫ φ = 𝟙 A ∧ χ ≫ g = 𝟙 C ∧ φ ≫ f + g ≫ χ = 𝟙 B) :
  short_exact_sequence 𝒜 :=
mk_of_split f g H.some H.some_spec.some H.some_spec.some_spec.1 H.some_spec.some_spec.2.1
  H.some_spec.some_spec.2.2.1 H.some_spec.some_spec.2.2.2

@[simp] def mk_split (A B : 𝒜) : short_exact_sequence 𝒜 :=
{ fst := A,
  snd := A ⊞ B,
  trd := B,
  f := biprod.inl,
  g := biprod.snd,
  exact' := exact_inl_snd _ _ }

/-- A *splitting* of a short exact sequence `0 ⟶ A₁ -f⟶ A₂ -g⟶ A₃ ⟶ 0` is
an isomorphism to the short exact sequence `0 ⟶ A₁ ⟶ A₁ ⊕ A₃ ⟶ A₃ ⟶ 0`,
where the left and right components of the isomorphism are identity maps. -/
structure splitting (A : short_exact_sequence 𝒜) extends A ≅ (mk_split A.1 A.3) :=
(fst_eq_id : hom.1 = 𝟙 A.1)
(trd_eq_id : hom.3 = 𝟙 A.3)

/-- A short exact sequence `0 ⟶ A₁ -f⟶ A₂ -g⟶ A₃ ⟶ 0` is *split* if there exist
`φ : A₂ ⟶ A₁` and `χ : A₃ ⟶ A₂` such that:
* `f ≫ φ = 𝟙 A₁`
* `χ ≫ g = 𝟙 A₃`
* `χ ≫ φ = 0`
* `φ ≫ f + g ≫ χ = 𝟙 A₂`
-/
def split (A : short_exact_sequence 𝒜) : Prop :=
∃ (φ : A.2 ⟶ A.1) (χ : A.3 ⟶ A.2),
   A.f ≫ φ = 𝟙 A.1 ∧ χ ≫ A.g = 𝟙 A.3 ∧ χ ≫ φ = 0 ∧ φ ≫ A.f + A.g ≫ χ = 𝟙 A.2

lemma mk_split_split (A B : 𝒜) : (mk_split A B).split :=
⟨biprod.fst, biprod.inr, biprod.inl_fst, biprod.inr_snd, biprod.inr_fst, biprod.total⟩

lemma splitting.split {A : short_exact_sequence 𝒜} (i : splitting A) : A.split :=
begin
  refine ⟨i.hom.2 ≫ biprod.fst ≫ i.inv.1, i.hom.3 ≫ biprod.inr ≫ i.inv.2, _⟩,
  simp only [category.assoc, ← hom.sq1_assoc, hom.sq2], dsimp,
  simp only [biprod.inl_fst_assoc, biprod.inr_snd_assoc, category.comp_id, category.assoc,
    ← comp_fst, ← comp_snd_assoc, ← comp_trd, i.to_iso.hom_inv_id, i.to_iso.inv_hom_id],
  dsimp,
  simp only [true_and, biprod.inr_fst_assoc, zero_comp, eq_self_iff_true, comp_zero,
    category.id_comp],
  simp only [hom.sq1, ← hom.sq2_assoc, ← comp_add],
  simp only [← category.assoc, ← add_comp, biprod.total,
    category.comp_id, ← comp_snd, i.to_iso.hom_inv_id], refl,
end

def left_split.splitting {A : short_exact_sequence 𝒜} (h : A.left_split) : A.splitting :=
{ to_iso := iso_of_components' (iso.refl _) (biprod.lift h.some A.g) (iso.refl _)
    (by { dsimp, simp only [category.id_comp], ext,
      { simpa only [biprod.inl_fst, biprod.lift_fst, category.assoc] using h.some_spec.symm, },
      { simp only [exact.w, f_comp_g, biprod.lift_snd, category.assoc, exact_inl_snd] } })
    (by { dsimp, simp only [category.comp_id, biprod.lift_snd], }),
  fst_eq_id := rfl,
  trd_eq_id := rfl }

def right_split.splitting {A : short_exact_sequence 𝒜} (h : A.right_split) : A.splitting :=
{ to_iso := iso.symm $ iso_of_components' (iso.refl _) (biprod.desc A.f h.some) (iso.refl _)
    (by { dsimp, simp only [biprod.inl_desc, category.id_comp], })
    (by { dsimp, simp only [category.comp_id], ext,
      { simp only [exact.w, f_comp_g, biprod.inl_desc_assoc, exact_inl_snd] },
      { simpa only [biprod.inr_snd, biprod.inr_desc_assoc] using h.some_spec, } }),
  fst_eq_id := rfl,
  trd_eq_id := rfl }

lemma tfae_split (A : short_exact_sequence 𝒜) :
  tfae [A.left_split, A.right_split, A.split, nonempty A.splitting] :=
begin
  tfae_have : 3 → 1, { rintro ⟨φ, χ, hφ, hχ, hχφ, H⟩, exact ⟨φ, hφ⟩ },
  tfae_have : 3 → 2, { rintro ⟨φ, χ, hφ, hχ, hχφ, H⟩, exact ⟨χ, hχ⟩ },
  tfae_have : 4 → 3, { rintro ⟨i⟩, exact i.split, },
  tfae_have : 1 → 4, { intro h, exact ⟨h.splitting⟩ },
  tfae_have : 2 → 4, { intro h, exact ⟨h.splitting⟩ },
  tfae_finish
end

end split

end short_exact_sequence

namespace short_exact_sequence

open category_theory.preadditive

variables {𝒞} [preadditive 𝒞] [has_images 𝒞] [has_kernels 𝒞]
variables (A B : short_exact_sequence 𝒞)

local notation `π₁` := congr_arg _root_.prod.fst
local notation `π₂` := congr_arg _root_.prod.snd

protected def hom_inj (f : A ⟶ B) : (A.1 ⟶ B.1) × (A.2 ⟶ B.2) × (A.3 ⟶ B.3) := ⟨f.1, f.2, f.3⟩

protected lemma hom_inj_injective : function.injective (short_exact_sequence.hom_inj A B) :=
λ f g h, let aux := π₂ h in
by { ext; [have := π₁ h, have := π₁ aux, have := π₂ aux]; exact this, }

instance : has_add (A ⟶ B) :=
{ add := λ f g,
  { fst := f.1 + g.1,
    snd := f.2 + g.2,
    trd := f.3 + g.3,
    sq1' := by { rw [add_comp, comp_add, f.sq1, g.sq1], },
    sq2' := by { rw [add_comp, comp_add, f.sq2, g.sq2], } } }

instance : has_neg (A ⟶ B) :=
{ neg := λ f,
  { fst := -f.1,
    snd := -f.2,
    trd := -f.3,
    sq1' := by { rw [neg_comp, comp_neg, f.sq1], },
    sq2' := by { rw [neg_comp, comp_neg, f.sq2], } } }

instance : has_sub (A ⟶ B) :=
{ sub := λ f g,
  { fst := f.1 - g.1,
    snd := f.2 - g.2,
    trd := f.3 - g.3,
    sq1' := by { rw [sub_comp, comp_sub, f.sq1, g.sq1], },
    sq2' := by { rw [sub_comp, comp_sub, f.sq2, g.sq2], } } }

instance has_nsmul : has_smul ℕ (A ⟶ B) :=
{ smul := λ n f,
  { fst := n • f.1,
    snd := n • f.2,
    trd := n • f.3,
    sq1' := by rw [nsmul_comp, comp_nsmul, f.sq1],
    sq2' := by rw [nsmul_comp, comp_nsmul, f.sq2] } }

instance has_zsmul : has_smul ℤ (A ⟶ B) :=
{ smul := λ n f,
  { fst := n • f.1,
    snd := n • f.2,
    trd := n • f.3,
    sq1' := by rw [zsmul_comp, comp_zsmul, f.sq1],
    sq2' := by rw [zsmul_comp, comp_zsmul, f.sq2] } }

variables (𝒞)

instance : preadditive (short_exact_sequence 𝒞) :=
{ hom_group := λ A B, (short_exact_sequence.hom_inj_injective A B).add_comm_group _
  rfl (λ _ _, rfl) (λ _, rfl) (λ _ _, rfl) (λ _ _, rfl) (λ _ _, rfl),
  add_comp' := by { intros, ext; apply add_comp },
  comp_add' := by { intros, ext; apply comp_add }, }
.

instance Fst_additive : (Fst 𝒞).additive := {}
instance Snd_additive : (Snd 𝒞).additive := {}
instance Trd_additive : (Trd 𝒞).additive := {}

end short_exact_sequence

end category_theory


import Mathlib.Tactic.Basic
import Mathlib.Tactic.Linarith
import Mathlib.Tactic.Ring
import Mathlib.Tactic.LibrarySearch
import Mathlib.Tactic.Cases
import Mathlib.Data.Quot
import Std.Data.Int.Basic
import SSA.Experiment.SSA2TreeNoProof
import SSA.Experiment.SSARgn2TreeNoProof
import SSA.Experiment.SSARgnVar2TreeNoProof
import SSA.Experiment.MLIRFlat

open Std

#check RBMap
namespace AST

/-
Kinds of values. We must have 'pair' to take multiple arguments.
TODO: Does it make sense to make this a typeclass?
-/
inductive Kind where
| int : Kind
| nat : Kind
| float : Kind
| pair : Kind -> Kind -> Kind
| unit: Kind
deriving Inhabited, DecidableEq, BEq

instance : ToString Kind where
  toString k :=
    let rec go : Kind →String
    | .nat => "nat"
    | .int => "int"
    | .float => "float"
    | .unit => "unit"
    | .pair p q => s!"({go p}, {go q})"
    go k

-- A binding of 'name' with kind 'Kind'
structure Var where
  name : String
  kind : Kind
deriving Inhabited, DecidableEq, BEq

instance : ToString Var where
  toString x := "%" ++ x.name ++ ":" ++ toString x.kind

abbrev Var.unit : Var := { name := "_", kind := .unit }

-- compile time constant values.
inductive Const where
| int: Int → Const
| float: Float → Const
| unit: Const
| pair: Const → Const → Const
deriving BEq

instance : ToString Const where
  toString :=
    let rec go : Const → String
    | .int i => toString i
    | .float f => toString f
    | .unit => "()"
    | .pair p q => s!"({go p}, {go q})"
    go

-- Tag for variants of Op/Region hybrid
inductive OR: Type
| O -- Single Op
| Os -- Multiple Ops
| R -- Single Region
| Rs -- Multiple Regions

inductive OpName where
| arith.constant
| arith.add
| arith.sub
| arith.mul
| scf.if
| scf.for
| scf.twice
| scf.run
| unknown
deriving DecidableEq

instance : ToString OpName where
  toString
    | .arith.constant => "arith.constant"
    | .arith.add => "arith.add"
    | .arith.sub => "arith.sub"
    | .arith.mul => "arith.mul"
    | .scf.if => "scf.if"
    | .scf.for => "scf.for"
    | .scf.twice => "scf.twice"
    | .scf.run => "scf.run"
    | .unknown => "<unknown>"

@[simp]
def OpName.fromString : String -> OpName
  | "arith.constant" => OpName.arith.constant
  | "arith.add" => OpName.arith.add
  | "arith.sub" => OpName.arith.sub
  | "arith.mul" => OpName.arith.mul
  | "scf.if" => OpName.scf.if
  | "scf.for" => OpName.scf.for
  | "scf.twice" => OpName.scf.twice
  | "scf.run" => OpName.scf.run
  | _ => OpName.unknown

-- Tagged expressiontrees
inductive Expr: OR -> Type where
| opsone: Expr .O -> Expr .Os -- terminator in a sequence of Ops
| opscons: Expr .O -> Expr .Os -> Expr .Os -- cons cell 'op :: ops'
| regionsnil : Expr .Rs -- empty sequence of regions
| regionscons: Expr .R -> Expr .Rs -> Expr .Rs -- cons cell 'region :: regions'
| op (ret : Var)
   (name : OpName)
   (arg : Var)
   (regions : Expr .Rs)
   (const: Const): Expr .O -- '%ret:retty = 'name'(%var:varty) [regions*] {const}'
| tuple (ret: String) (a1 a2: Var): Expr .O -- %out = tuple (%v1, %v2)
| region (arg : Var) (ops : Expr .Os): Expr .R -- '{ ^entry(arg:argty) ops* }'

abbrev Op := Expr .O
abbrev Region := Expr .R
abbrev Regions := Expr .Rs
abbrev Ops := Expr .Os

abbrev Op.mk (ret: Var := Var.unit)
  (name: OpName)
  (arg: Var := Var.unit)
  (regions := Expr.regionsnil)
  (const := Const.unit): Op := Expr.op ret name arg regions const
-- Append an 'Op' to the end of the 'Ops' list.
def Ops.snoc: Ops → Op → Ops
| .opsone o, o' => .opscons o (.opsone o')
| .opscons o os, o' => .opscons o (Ops.snoc os o')

@[simp]
def Op.ret : Op → Var
| .op ret _ _ _ _ => ret
| .tuple retname arg1 arg2 =>
   { name := retname,
     kind := .pair arg1.kind arg2.kind : Var}

def Regions.isEmpty: Regions → Bool
| .regionsnil => True
| .regionscons _ _ => False

def Expr.format : Expr k → Format
| .opsone o => o.format
| .opscons o os => o.format ++ .line ++ os.format
| .regionsnil => Format.nil
| .regionscons r rs => r.format ++ .line ++ rs.format
| .op ret name arg rs const =>
    let constfmt : Format :=
      if const == .unit then ""
      else "{" ++ toString const ++ "}"
    let argfmt : Format :=
      if arg == Var.unit then ""
      else "(" ++ toString arg ++ ")"
    let rsfmt : Format :=
      if Regions.isEmpty rs then ""
      else (Format.nest 1 <| "[" ++ .line ++ Expr.format rs) ++ "]"
    .text (toString ret ++ " = " ++ toString name) ++
    argfmt ++ rsfmt ++ constfmt
| .tuple ret a1 a2 =>
    .text <|
      "%" ++ toString ret ++ " = " ++
      "(" ++ toString a1 ++ ", " ++ toString a2 ++ ")"
| .region arg ops =>
    let argfmt : Format :=
    if arg == Var.unit then "" else "^(" ++ toString arg ++ ")"
    "{" ++ argfmt ++
      Format.nest 2 (.line ++ ops.format) ++ .line ++ "}"
instance ⦃k: OR⦄: ToFormat (Expr k) where
  format := Expr.format

instance : ToString (Expr k) where
   toString := Format.pretty ∘ format

-- Lean type that corresponds to kind.
@[reducible, simp]
def Kind.eval: Kind -> Type
| .int => Int
| .nat => Nat
| .unit => Unit
| .float => Float
| .pair p q => p.eval × q.eval

instance Kind.bEqEval: (k : Kind) → BEq k.eval
| .int => inferInstanceAs (BEq Int)
| .nat => inferInstanceAs (BEq Nat)
| .unit => inferInstanceAs (BEq Unit)
| .float => inferInstanceAs (BEq Float)
| .pair p q => by
    letI : BEq p.eval := Kind.bEqEval _
    letI : BEq q.eval := Kind.bEqEval _
    infer_instance

def Kind.default (k: Kind): k.eval :=
  match k with
  | .int => 0
  | .nat => 0
  | .unit => ()
  | .float => 0.0
  | .pair p q => (p.default, q.default)

instance KindDefault (k: Kind) : Inhabited (k.eval) where
  default := Kind.default _

end AST

section Semantics
open AST

-- A kind and a value of that kind.
structure Val where
  kind: Kind
  val: kind.eval

instance : BEq Val where
  beq a b :=
    let rec go a b :=
    match a, b with
    | {kind := .int, val := a}, {kind := .int, val := b} => a == b
    | {kind := .float, val := a}, {kind := .float, val := b} => a == b
    | {kind := .unit, val := a}, {kind := .unit, val := b} => a == b
    | {kind := .pair p q, val := ⟨a, a'⟩}, {kind := .pair r s, val := ⟨b, b'⟩ } =>
      go ⟨p, a⟩ ⟨r, b⟩ && go ⟨q, a'⟩ ⟨s, b'⟩
    | _, _ => False
    go a b

def Val.unit : Val := { kind := Kind.unit, val := () }

def Val.toString (v: Val): String :=
  match v  with
  | {kind := .nat, val := val } =>
      let S : ToString Nat := inferInstance
      S.toString val
  | {kind := .int, val := val } =>
      let S : ToString Int := inferInstance
      S.toString val
  | {kind := .float, val := val } =>
      let S : ToString Float := inferInstance
      S.toString val
  | {kind := .unit, val := () } => "()"
  | {kind := .pair p q, val := (x, y) } =>
      let xstr := Val.toString ({ kind := p, val := x})
      let ystr := Val.toString { kind := q, val := y}
      s!"({xstr}, {ystr})"

instance : ToString Val where
  toString := Val.toString

-- The retun value of an SSA operation, with a name, kind, and value of that kind.
structure NamedVal extends Val where
  name : String
deriving BEq

def NamedVal.toString (nv: NamedVal): String :=
 s!"{nv.name} := {Val.toString nv.toVal}"

instance : ToString NamedVal where
  toString := NamedVal.toString

-- Given a 'Var' of kind 'kind', and a value of type 〚kind⟧, build a 'Val'
@[simp]
def AST.Var.toNamedVal (var: Var) (value: var.kind.eval): NamedVal :=
 { kind := var.kind, val := value, name := var.name }

@[simp]
def NamedVal.var (nv: NamedVal): Var :=
  { name := nv.name, kind := nv.kind }

-- Well typed environments; cons cells of
-- bindings of variables to values of type ⟦var.kind⟧
abbrev Env := (v: Var) → v.kind.eval

instance : Inhabited ((k: Kind) → Kind.eval k) where
  default :=
    let rec go := fun k =>
      match k with
      | .unit => ()
      | .pair p q => (go p, go q)
      | .float => 0
      | .int => 0
      | .nat => 0
    go

def Env.empty := fun (v : Var) =>
  let f : (k: Kind) → Kind.eval k := default
  f v.kind

-- truncation of a type that smashes everything into a single equivalence class.
inductive trunc (α: Type): α → α → Prop
| trunc: trunc α a a' -- smash everthing.

instance EquivalenceTrunc : Equivalence (trunc α) where
  refl _ := .trunc
  symm _  := .trunc
  trans _ _ := .trunc

instance SetoidTrunc (α : Type) : Setoid α where
   r := trunc α
   iseqv := EquivalenceTrunc

/-
The type of Error is computationally 'String', but logically just a point.
this allows us to ignore error states in proofs [they are all identified as equal],
while still allowing computationally relevant errors.
-/
abbrev ErrorKind : Type := Trunc String

/-
Cursed: We cast from 'ErrorKind' which is a quotient of 'String' into 'String'
since we know that these have the same RuntimeRep.
-/
unsafe def ErrorKind.toStringImpl (e: ErrorKind): String := unsafeCast e
@[implemented_by ErrorKind.toStringImpl]
partial def ErrorKind.toString (_: ErrorKind): String := "<<error>>"
instance : ToString ErrorKind where
  toString := ErrorKind.toString

abbrev ErrorKind.mk (s: String) : ErrorKind := Trunc.mk s

-- Coerce from regular strings into errors.
instance : Coe String ErrorKind where
  coe := ErrorKind.mk

@[simp]
def ErrorKind.subsingleton (e: ErrorKind) (e': ErrorKind): e = e' := by {
  apply Quotient.ind;
  intros a;
  apply Eq.symm;
  apply Quotient.ind;
  intros b;
  apply Quotient.sound;
  constructor;
}

-- Env → Except ErrorKind α
abbrev TopM (α : Type) : Type := ReaderT Env Id α

abbrev Env.set (var: Var) (val: var.kind.eval): Env → Env
| env => fun v => if H : v = var then H ▸ val else env v

-- We need to produce values of type '()' for eg. ops with zero agruments.
-- Here, we ensure that Env.get for a var of type 'Unit' will always succeed and return '()',
-- becuse there is not need not to. This allows us to use 'Unit' to signal zero arguments,
-- wthout have to make up a fake name for a variable of type 'Unit'
abbrev Env.get (var: Var) (e: Env): var.kind.eval := e var

abbrev ReaderT.get [Monad m]: ReaderT ρ m ρ := fun x => pure x
abbrev ReaderT.withEnv [Monad m] (f: ρ → ρ) (reader: ReaderT ρ m α): ReaderT ρ m α :=
  reader ∘ f

def TopM.get (var: Var): TopM var.kind.eval := ReaderT.get >>= (fun _ => (Env.get var))

-- the unit type will always successfully return '()'
theorem TopM.get_unit (name: String) (env: Env): TopM.get ⟨name, .unit⟩ env = () := by {
  simp[get, ReaderT.get, bind, ReaderT.bind, Except.bind, pure, Except.pure];
}

theorem TopM_get (name: String)  (kind: Kind) (k: kind.eval → TopM β):
  (TopM.get ⟨name, kind⟩) >>= k =  fun env => (k (env ⟨name, kind⟩)) env := rfl

def TopM.set (nv: NamedVal)  (k: TopM α): TopM α :=
  ReaderT.withEnv (Env.set nv.var nv.val) k

def Val.cast (val: Val) (t: Kind): t.eval :=
  if H : val.kind = t
  then .(H ▸ val.val)
  else default
-- Runtime values of arguments to an Op, This is the argument value,
-- the evaluated regions, and the constant.
structure Op' where
  argval : Val := ⟨.unit, ()⟩
  regions: List (Val → TopM Val) := []
  const: Const := .unit
  retkind: Kind

-- The single semantic unit, where the user provides the semantics
-- of a single op of a given 'name'.
structure Semantic where
  name: OpName
  run: (o : Op') → TopM o.retkind.eval

-- TODO: consider this design.
-- partial finitely supported function.
-- abbrev Semantics := (name: String) → Option (Semantic name)

abbrev Semantics := OpName → (o: Op') → TopM o.retkind.eval

instance : ToString Op' where
  toString x := "(" ++ toString x.argval ++ ")" ++
    " [" ++ "#" ++ toString x.regions.length ++ "]" ++
    " {" ++ toString x.const ++ "}" ++ " → " ++ toString x.retkind

@[reducible]
def AST.OR.denoteType: OR -> Type
| .O => TopM NamedVal
| .Os => TopM NamedVal
| .R => Val → TopM Val
| .Rs => List (Val → TopM Val) -- TODO: is 'List' here correct?

def AST.Expr.denote {kind: OR}
 (sem: Semantics): Expr kind → kind.denoteType
| .opsone o => AST.Expr.denote (kind := .O) sem o
| .opscons o os => do
    let retv ← AST.Expr.denote (kind := .O) sem o
    TopM.set retv (os.denote sem)
| .regionsnil => []
| .regionscons r rs => r.denote sem :: (rs.denote sem)
| .tuple ret arg1 arg2 => do
   let val1 ←TopM.get arg1
   let val2 ← TopM.get arg2
   return { name := ret,
            kind := .pair arg1.kind arg2.kind,
            val := (val1, val2)
          } -- build a pair
| .op ret name arg rs const => do
    let val ← TopM.get arg
    let op' : Op' :=
      { argval := ⟨arg.kind, val⟩
      , regions := rs.denote sem
      , const := const
      , retkind := ret.kind }
    let out ← sem name op'
    return ret.toNamedVal out
| .region arg ops => fun val => do
    -- TODO: improve dependent typing here
    let val' := val.cast arg.kind
    TopM.set (arg.toNamedVal val') (NamedVal.toVal <$> (ops.denote sem))

-- Write this separately for defEq purposes.
def Op'.fromOp (sem: Semantics)
  (ret: Var) (arg: Var) (argval: arg.kind.eval) (rs: Regions) (const: Const) : Op' :=
      { argval := ⟨arg.kind, argval⟩
      , regions := rs.denote sem
      , const := const
      , retkind := ret.kind
  }

def runRegion (sem : Semantics) (expr: AST.Expr .R)
(env :  Env := fun e => default)
(arg : Val := Val.unit) : Val :=
  (expr.denote sem arg).run env

theorem TopM_idempotent: ∀ (ma: TopM α) (k: α → α → TopM β),
 ma >>= (fun a => ma >>= (fun a' => k a a')) =
 ma >>= (fun a => k a a) := by {
  intros ma k;
  funext env;
  simp[ReaderT.bind];
  simp[bind];
  simp[ReaderT.bind];
}

theorem TopM_commutative: ∀ (ma: TopM α) (mb: TopM β) (k: α → β → TopM γ),
 ma >>= (fun a => mb >>= (fun b => k a b)) =
 mb >>= (fun b => ma >>= (fun a => k a b)) := by {
  intros ma mb k;
  funext env;
  simp[ReaderT.bind];
  simp[bind];
  simp[ReaderT.bind];
}

-- unfold Expr.denote for opscons
theorem Expr.denote_opscons:
  (Expr.opscons o os).denote sem =
    (o.denote sem >>= fun retv => TopM.set retv (os.denote sem)) := by { rfl; }

-- unfold Expr.denote for opsone
theorem Expr.denote_opsone:
  (Expr.opsone o).denote sem = o.denote sem := by { rfl; }

theorem Expr.denote_tuple:
  (Expr.tuple ret arg1 arg2).denote sem =
  do
   let val1 ←TopM.get arg1
   let val2 ← TopM.get arg2
   return { name := ret,
            kind := .pair arg1.kind arg2.kind,
            val := (val1, val2)
          } := by rfl

-- unfold Expr.denote for op
theorem Expr.denote_op:
  (Op.mk ret name arg rs const ).denote sem =
  TopM.get arg >>= fun val =>
    let op' : Op' := -- TODO: use Op'.fromOp
      { argval := ⟨arg.kind, val⟩
      , regions := rs.denote sem
      , const := const
      , retkind := ret.kind }
    sem name op' >>= fun out => fun env => ret.toNamedVal out
  := by { rfl; }

-- how to simply Expr.denote_op, assuming the environment lookup is correct.
theorem Expr.denote_op_success (ret: Var) (name: OpName) (arg: Var) (rs: Regions) (const: Const)
  (sem: Semantics)
  (env: Env)
  (argval: arg.kind.eval)
  (ARGVAL: env.get arg = argval)
  (outval : ret.kind.eval)
  (SEM: sem name (Op'.fromOp sem ret arg argval rs const) env = outval) :
  (Op.mk ret name arg rs const).denote sem env = (ret.toNamedVal outval)
  := by {
      rw[Expr.denote_op];
      simp[Op'.fromOp] at SEM;
      simp[Op'.fromOp, TopM.get, ReaderT.get, pure, bind, ReaderT.bind, Except.pure, Except.bind, ARGVAL, SEM];
      aesop
   }

-- 'opscons o os' @ env is the same as 'os @ (env[o.ret = o.val])'.
-- That is, if 'o' succeeds, the 'os' proceeds evaluation with 'env' which
-- has been updated for 'o.ret' with 'o.val'.
theorem Expr.denote_opscons_success_op (o: Op) (outval: (Op.ret o).kind.eval)
  (sem : Semantics) (env: Env)
  (OVAL: o.denote sem env =  (o.ret.toNamedVal outval)) :
  (Expr.opscons o os).denote sem env = os.denote sem (env.set o.ret outval) := by {
    rw[Expr.denote_opscons];
    simp[bind, ReaderT.bind, OVAL, Except.bind, TopM.set];
    sorry
}

-- If tuple succeeds, then:
-- 1. The arguments exists in 'env'.
-- 2. The value in 'env' is that of the tuple arguments, tupled up.
theorem Expr.denote_tuple_inv
  (SUCCESS: Expr.denote sem (Expr.tuple retname arg1 arg2) env = outval) :
    ∃ argval1, env.get arg1 = argval1 ∧
    ∃ argval2, env.get arg2 = argval2 ∧
    outval = { name := retname, kind := .pair arg1.kind arg2.kind, val := ⟨argval1, argval2 ⟩} := by {
    simp[Expr.denote, bind, ReaderT.bind, Except.bind] at SUCCESS;

      simp[TopM.get, bind, ReaderT.bind, Except.bind, ReaderT.get, pure, Except.pure] at SUCCESS ⊢;
      simp[pure, ReaderT.pure, Except.pure] at SUCCESS ⊢  ;
      aesop;
}

-- If op succeeds, then:
-- 1. the arguments exists in 'env'.
-- 2. 'sem' succeeded
-- 3. The value returned by 'Expr.denote' is the boxed value in 'sem'.
theorem Expr.denote_op_assign_inv
  (SUCCESS: Expr.denote sem (Expr.op ret name arg regions const) env = outval) :
    ∃ argval, env.get arg = argval ∧
    ∃ (outval_val : ret.kind.eval),
      (sem name (Op'.fromOp sem ret arg argval regions const) env = pure outval_val /\
       outval = ret.toNamedVal outval_val)   := by {
    -- rw[Expr.denote_op] at SUCCESS; -- TODO: find out why 'rw' does not work here.

    simp[Expr.denote, bind, ReaderT.bind, Except.bind] at SUCCESS;
      simp[TopM.get, bind, ReaderT.bind, Except.bind] at SUCCESS ⊢;
      simp[ReaderT.get, pure, Except.pure] at SUCCESS ⊢;
      simp[SUCCESS];
      simp[pure, Except.pure];
      simp[pure, Except.pure, ReaderT.pure] at SUCCESS;
      simp[SUCCESS];
      aesop
}

-- What we can say about a op in general.
theorem Expr.denote_regionsnil: Expr.regionsnil.denote sem = [] := rfl

theorem Expr.denote_regionscons:
  Expr.denote sem (.regionscons r rs) = r.denote sem :: rs.denote sem := rfl

end Semantics

namespace DSL
open AST

declare_syntax_cat dsl_op
declare_syntax_cat dsl_type
declare_syntax_cat dsl_ops
declare_syntax_cat dsl_region
declare_syntax_cat dsl_var
declare_syntax_cat dsl_var_name
declare_syntax_cat dsl_const
declare_syntax_cat dsl_kind

syntax sepBy1(ident, "×") :dsl_kind
syntax ident : dsl_var_name
syntax "%" dsl_var_name ":" dsl_kind : dsl_var
syntax dsl_var "="
      str
      ("("dsl_var ")")?
      ("[" dsl_region,* "]")?
      ("{" dsl_const "}")? ";": dsl_op
syntax "%" dsl_var_name "=" "tuple" "(" dsl_var "," dsl_var ")" ";": dsl_op
syntax num : dsl_const
syntax "{" ("^(" dsl_var ")" ":")?  dsl_op dsl_op*  "}" : dsl_region
syntax "[dsl_op|" dsl_op "]" : term
syntax "[dsl_ops|" dsl_op dsl_op* "]" : term
syntax "[dsl_region|" dsl_region "]" : term
syntax "[dsl_var|" dsl_var "]" : term
syntax "[dsl_kind|" dsl_kind "]" : term
syntax "[dsl_const|" dsl_const "]" : term
syntax "[dsl_var_name|" dsl_var_name "]" : term

@[simp]
def AST.Ops.fromList: Op → List Op → Ops
| op, [] => .opsone op
| op, op'::ops => .opscons op (AST.Ops.fromList op' ops)

@[simp]
def AST.Regions.fromList: List Region → Regions
| [] => .regionsnil
| r :: rs => .regionscons r (AST.Regions.fromList rs)

open Lean Macro in
def parseKind (k: TSyntax `ident) : MacroM (TSyntax `term) :=
  match k.getId.toString with
  | "int" => `(AST.Kind.int)
  | "nat" => `(AST.Kind.nat)
  | unk => (Macro.throwErrorAt k s!"unknown kind '{unk}'")

open Lean Macro in
macro_rules
| `([dsl_kind| $[ $ks ]×* ]) => do
    if ks.isEmpty
    then `(AST.Kind.unit)
    else
      let mut out ← parseKind ks[ks.size - 1]!
      for k in ks.pop.reverse do
        let cur ← parseKind k
        out ← `(AST.Kind.pair $cur $out)
      return out
| `([dsl_kind| $$($q:term)]) => `(($q : Kind))

open Lean Macro in
macro_rules
| `([dsl_var_name| $name:ident ]) =>
      return (Lean.quote name.getId.toString : TSyntax `term)
| `([dsl_var_name| $$($q:term)]) => `(($q : String))

macro_rules
| `([dsl_var| %$name:dsl_var_name : $kind:dsl_kind]) => do
    `({ name := [dsl_var_name| $name],
        kind := [dsl_kind| $kind] : Var})
| `([dsl_var| $$($q:term)]) => `(($q : Var))

macro_rules
| `([dsl_ops| $op $ops*]) => do
   let op_term ← `([dsl_op| $op])
   let ops_term ← ops.mapM (fun op => `([dsl_op| $op ]))
   `(AST.Ops.fromList $op_term [ $ops_term,* ])

macro_rules
| `([dsl_ops| $$($q)]) => `(($q : Ops))

macro_rules
| `([dsl_region| { $[ ^( $arg:dsl_var ): ]? $op $ops* } ]) => do
   let ops ← `([dsl_ops| $op $ops*])
   match arg with
   | .none => `(Expr.region Var.unit $ops)
   | .some arg => do
      let arg_term ← `([dsl_var| $arg ])
      `(Expr.region $arg_term $ops)
| `([dsl_region| $$($q)]) => `(($q : Region))

macro_rules
| `([dsl_const| $x:num ]) => `(Const.int $x)
| `([dsl_const| $$($q)]) => `(($q : Const))

open Lean Syntax in
macro_rules
| `([dsl_op| $res:dsl_var = $name:str
      $[ ( $arg:dsl_var ) ]?
      $[ [ $rgns,* ] ]?
      $[ { $const } ]?
      ;
      ]) => do
      let res_term ← `([dsl_var| $res])
      let arg_term ← match arg with
          | .some arg => `([dsl_var| $arg])
          | .none => `(AST.Var.unit)
      let name_term := name
      let rgns_term ← match rgns with
        | .none => `(.regionsnil)
        | .some rgns =>
           let rgns ← rgns.getElems.mapM (fun stx => `([dsl_region| $stx]))
           `(AST.Regions.fromList [ $rgns,* ])
      let const_term ←
        match const with
        | .none => `(Const.unit)
        | .some c => `([dsl_const| $c])
      `(Expr.op $res_term (OpName.fromString ($name_term : String)) $arg_term $rgns_term $const_term)
| `([dsl_op| $$($q)]) => `(($q : Op))

macro_rules
| `([dsl_op| %$res:dsl_var_name = tuple ( $arg1:dsl_var , $arg2:dsl_var) ; ]) => do
    let res_term ← `([dsl_var_name| $res])
    let arg1_term ← `([dsl_var| $arg1 ])
    let arg2_term ← `([dsl_var| $arg2 ])
    `(Expr.tuple $res_term $arg1_term $arg2_term)

def eg_kind_int := [dsl_kind| int]
#reduce eg_kind_int

def eg_kind_int_times_nat := [dsl_kind| int × nat]
#reduce eg_kind_int_times_nat

def eg_kind_int_times_nat_times_nat := [dsl_kind| int × nat × nat]
#reduce eg_kind_int_times_nat_times_nat

def eg_var : AST.Var := [dsl_var| %y : int]
#reduce eg_var

def eg_var_2 : AST.Var := [dsl_var| %$("y") : int ]
#reduce eg_var_2

def eg_var_3 : AST.Var :=
  let name := "name";  let ty := Kind.int;  [dsl_var| %$(name) : $(ty) ]
#reduce eg_var_3

def eg_op : AST.Op :=
  let name := "name";  let ty := Kind.int;
  [dsl_op| %res : int = "arith.sub" ( %$(name) : $(ty) ); ]
#reduce eg_var_3

section Unexpander

/-
Support for pretty printing our AST.Expr to make writing proofs easier.
-/
/-
@[app_unexpander Var.mk] def unexpandVar: Lean.PrettyPrinter.Unexpander
| `($_ $nm $kind) => `($nm $kind)
| _ => throw ()

@[app_unexpander Val.mk] def unexpandVal: Lean.PrettyPrinter.Unexpander
| `($_ $kind $val) => `($val $kind)
| _ => throw ()

@[app_unexpander Kind.pair] def unexpandKindPair: Lean.PrettyPrinter.Unexpander
| `($_ $l $r) => `($l × $r)
| _ => throw ()

@[app_unexpander Expr.tuple] def unexpandExprTuple: Lean.PrettyPrinter.Unexpander
| `($_ $ret $l $r) => `($ret = "tuple" ( $l , $r ))
| _ => throw ()

open Lean Macro PrettyPrinter in
@[app_unexpander Expr.op] def unexpandExprOp: Lean.PrettyPrinter.Unexpander
| `($_ $ret $name $arg regionsnil $const) => `($ret = "op" $name ( $arg ) $const)
-- | `($_ $ret $name $arg $regions Const.unit) => `($ret = "op" $name ( $arg ) $regions)
| `($_ $ret $name $arg $regions $const) => `($ret = "op" $name ( $arg ) $regions  $const)
| _ => sorry

@[app_unexpander Expr.opscons] def unexpandExprOpscons: Lean.PrettyPrinter.Unexpander
| `($_ $o $os) => `($o / $os)
| _ => throw ()

@[app_unexpander Expr.opsone] def unexpandExprOpsone: Lean.PrettyPrinter.Unexpander
| `($_ $o) => `($o)
| _ => throw ()

@[app_unexpander Expr.regionscons] def unexpandRegionsnil : Lean.PrettyPrinter.Unexpander
| `($_ ) => `(())

@[app_unexpander Expr.regionscons] def unexpandRegionsCons : Lean.PrettyPrinter.Unexpander
| `($_ $r Expr.regionsnil) => `($r)
| _ => throw ()

@[app_unexpander Const.unit] def unexpandConstUnit : Lean.PrettyPrinter.Unexpander
| `($_) => `(())

@[app_unexpander Const.int] def unexpandConstInt : Lean.PrettyPrinter.Unexpander
| `($_ $i:num)  => `($i)
| _ => throw ()
-/
end Unexpander

end DSL

namespace Arith

def const : Semantic := {
  name := .arith.constant
  run := fun o => match o with
  | { const := .int x, retkind := .int} => return x
  | _ => default
}

def add : Semantic := {
  name := .arith.add
  run := fun o => match o with
  | { retkind := .int,
      argval := ⟨.pair .int .int, (x, y)⟩ } => return x + y
  | _ => default
}

def sub : Semantic := {
  name := .arith.sub
  run := fun o => match o with
  | { retkind := .int,
      argval := ⟨.pair .int .int, (x, y)⟩ } => return x - y
  | _ => default
}

def sem: Semantics :=
  fun name =>
    match name with
    | .arith.constant => const.run
    | .arith.add => add.run
    | .arith.sub => sub.run
    | _ => default

def eg_region_sub :=
 [dsl_region| {
   %one : int = "arith.constant" {1};
   %two : int = "arith.constant" {2};
   %t = tuple(%one : int , %two : int);
   %x : int = "arith.sub"(%t : int × int);
 }]
#reduce eg_region_sub

end Arith

namespace Scf -- 'structured control flow'

def repeatM (n: Nat) (f: Val → TopM Val) : Val → TopM Val :=
  n.repeat (f >=> .) pure

theorem repeatM.peel_left: ∀ (n: Nat) (f: Val → TopM Val),
  repeatM (n+1) f = f >=> repeatM n f := by {
    intros n f;
    simp[repeatM, Nat.repeat];
}
-- This theorem is now useless
-- theorem repeatM.peel_left: repeatM (n+1) f = f >=> repeatM n f := rfl

-- kleisli arrow is associative.
theorem kleisliRight.assoc {m: Type → Type}
  [M: Monad m] [LM: LawfulMonad m]
  {α β γ δ: Type}
  (f: α → m β) (g: β → m γ) (h: γ → m δ) :
  (f >=> g) >=> h = f >=> (g >=> h) := by {
  funext x;
  simp[Bind.kleisliRight];
}

-- pure is left identity
@[simp]
theorem kleisliRight.pure_id_left {m: Type → Type}
  [M: Monad m] [LM: LawfulMonad m]
  {α β: Type}
  (f: α → m β) :
  pure >=> f = f := by {
  funext x;
  simp[Bind.kleisliRight];
}

-- pure is left identity
@[simp]
theorem kleisliRight.pure_id_right {m: Type → Type}
  [M: Monad m] [LM: LawfulMonad m]
  {α β: Type}
  (f: α → m β) :
  f >=> pure = f:= by {
  funext x;
  simp[Bind.kleisliRight];
}

-- The above shows that (>=>, pure) forms a category.

@[simp] theorem repeatM.zero: repeatM 0 f = pure := rfl
@[simp] theorem repeatM.one: repeatM 1 f = f := by simp [repeatM, Nat.repeat]
@[simp] theorem repeatM.succ: repeatM (Nat.succ n) f = f >=> repeatM n f := rfl

-- composing repeatM is the same as adding the repeat
-- counter.
theorem repeatM.compose: (repeatM n f) >=> (repeatM m f) = repeatM (n+m) f := by
  revert m
  induction n <;> intros m <;> simp
  case succ n' IH =>
    simp [Nat.succ_add, kleisliRight.assoc, IH]

theorem repeatM.peel_right: repeatM (n+1) f = repeatM n f >=> f := by
  induction n <;> simp at *
  case succ _ IH => rw [kleisliRight.assoc, IH]

theorem repeatM.commute_f: repeatM n f >=> f = f >=> repeatM n f := by
  simp [←peel_right]

-- pull a function that commutes with f to the left of repeatM
@[simp] -- pull simple stuff to the left.
theorem repeatM.commuting_pull_left
  {f g: Val → TopM Val} (COMMUTES: f >=> g = g >=> f) :
  (repeatM n f) >=> g = g >=> repeatM n f := by
  induction n <;> simp
  case succ n' IH =>
    simp [kleisliRight.assoc, IH]
    simp [←kleisliRight.assoc, COMMUTES]

theorem repeatM.commuting_pull_right
  {f g: Val → TopM Val} (COMMUTES: f >=> g = g >=> f) :
  g >=> repeatM n f = repeatM n f >=> g := (commuting_pull_left COMMUTES).symm

-- TODO: decision procedure for free theory of
-- commuting functions?
theorem repeatM.commuting_commute
  {f g: Val → TopM Val} (COMMUTES: f >=> g = g >=> f) :
  repeatM n f >=> repeatM m g = repeatM m g >=> repeatM n f := by
  induction n <;> simp
  case succ n' IH =>
    rw [kleisliRight.assoc, IH]
    simp [←kleisliRight.assoc]
    rw [commuting_pull_left COMMUTES.symm]

-- to show (f >=> k) = (f >=> h), it suffices to show that k = h
theorem kleisli.cancel_left [M: Monad m] (f: α → m β) (k h : β → m γ) (H: k = h) :
  f >=> k = f >=> h := by {
    rw[H];
}

theorem repeatM.commuting_compose
  {f g: Val → TopM Val} (COMMUTES: f >=> g = g >=> f) :
  (repeatM n f) >=> repeatM n g = repeatM n (f >=> g) := by
  induction n <;> simp
  case succ n' IH =>
    rw [←kleisliRight.assoc _ g _, kleisliRight.assoc _ _ g]
    rw [commuting_pull_left COMMUTES]
    rw [←kleisliRight.assoc f g _, kleisliRight.assoc (f >=> g) _ _, IH]

def if_ : Semantic := {
  name := .scf.if
  run := fun o => match o with
  | { argval := ⟨.int, cond⟩,
      regions := [rthen, relse],
      retkind := .int -- hack: we assume that we always return ints.
    } => do
      let rrun := if cond ≠ 0 then rthen else relse
      let v ← rrun Val.unit
      return v.cast .int
  | _ => default
}

def twice : Semantic := {
  name := .scf.twice
  run := fun o => match o with
  | { regions := [r],
      retkind := .int
    } => do
      let v ← r Val.unit
      let w ← r Val.unit -- run a region twice, check that it is same as running once.
      return w.cast .int
  | _ => default
}

def run : Semantic := {
  name := .scf.run
  run := fun o => match o with
  | { regions := [r],
      retkind := .int
    } => do
      let v ← r Val.unit
      return v.cast .int
  | _ => default
}

def for_ : Semantic := {
  name := .scf.for
  run := fun o => match o with
  | { regions := [r],
      retkind := .int,
      argval := ⟨.pair .nat .int, ⟨niters, init⟩⟩
    } =>
      repeatM niters r ⟨.int, init⟩ >=> Val.cast (t := .int)
  | _ => default
}

def sem : Semantics :=
  fun name =>
    match name with
    | .scf.if => if_.run
    | .scf.twice => twice.run
    | .scf.for => for_.run
    | _ => default

end Scf

namespace Examples
open AST
def eg_arith_shadow : Region := [dsl_region| {
  %x : int = "arith.constant"{0};
  %x : int = "arith.constant"{1};
  %x : int = "arith.constant"{2};
  %x : int = "arith.constant"{3};
  %x : int = "arith.constant"{4};
}]

def eg_scf_ite_true : Region := [dsl_region| {
  %cond : int = "arith.constant"{1};
  %out : int = "scf.if" (%cond : int) [{
    %out_then : int = "arith.constant"{42};
  }, {
    %out_else : int = "arith.constant"{0};
  }];
}]

-- f: Unit → a ∼ a
def eg_scf_ite_false : Region := [dsl_region| {
  %cond : int = "arith.constant"{0};
  %out : int = "scf.if" (%cond : int) [{
    %out_then : int = "arith.constant"{42};
  }, {
    %out_else : int = "arith.constant"{0};
  }];
}]

#print eg_scf_ite_false

def eg_scf_run_twice : Region := [dsl_region| {
  %x : int = "arith.constant"{41};
  %x : int = "scf.twice" [{
      %one : int = "arith.constant"{1};
      %xone = tuple(%x : int, %one : int);
      %x : int = "arith.add"(%xone : int × int);
  }];
}]

def eg_scf_well_scoped : Region := [dsl_region| {
  %x : int = "arith.constant"{41};
  %one : int = "arith.constant"{1}; -- one is outside.
  %x : int = "scf.twice" [{
      %xone = tuple(%x : int, %one : int); -- one is accessed here.
      %x : int = "arith.add"(%xone : int × int);
  }];
}]

def eg_scf_ill_scoped : Region := [dsl_region| {
  %x : int = "arith.constant"{41};
  %x : int = "scf.twice" [{
      %y : int = "arith.constant"{42};
  }];
  -- %y should NOT be accessible here
  %out = tuple(%y : int, %y : int);
}]

def eg_scf_for: Region := [dsl_region| {
  %x : int = "arith.constant"{10};  -- 10 iterations
  %init : int = "arith.constant"{32}; -- start value
  %xinit = tuple(%x : int, %init : int);
  %out : int = "scf.for"(%xinit : int × int)[{
    ^(%xcur : int):
      %one : int = "arith.constant"{1};
      %xcur_one = tuple(%xcur : int, %one : int);
      %xnext : int = "arith.add"(%xcur_one : int × int);
  }];
}]

end Examples

namespace Rewriting
open AST

/- Replace the final Op with the new seqence of Ops -/
def Ops.replaceOne (os: Ops) (new: Ops) : Ops :=
  match os with
  | .opsone o => new
  | .opscons o os => .opscons o (Ops.replaceOne os new)

structure Peephole where
  -- findbegin: TypingCtx -- free variables.
  find : Ops -- stuff in the pattern. The last op is to be replaced.
  replace : Ops -- replacement ops. can use 'findbegin'.
  sem: Semantics
  -- TODO: Once again, we need to reason 'upto error'.
  replaceCorrect: ∀ (env: Env),
      find.denote sem env = replace.denote sem env

      -- (FIND: find.denote sem env = origval),
      -- (replace.denote sem) env =  origval
end Rewriting

namespace RewritingHoare
open AST
open Rewriting
/-
A theory of rewriting, plus helper lemmas phrased along a hoare logic.
-/
-- https://softwarefoundations.cis.upenn.edu/plf-current/Hoare.html
def assertion (t: Type) : Type := t → Prop
def assertion.implies (P Q: assertion T) : Prop := ∀ (e: T), P e → Q e

notation P:80 "->>" Q:80 => assertion.implies P Q

abbrev assertion.iff (P Q: assertion T) : Prop := P ->> Q ∧ Q ->> P
notation P:80 "<<->>" Q:80 => assertion.iff P Q

def assertion.and (P Q: assertion T) : assertion T := fun (v: T) => P v ∧ Q v
notation P:80 "h∧" Q:80 => assertion.and P Q -- how does this work?

def assertion.prop (p: Prop): assertion T := fun (_v: T) => p

def assertion.mapsto (v: Var) (val: v.kind.eval): assertion Env :=
  fun (e: Env) => e.get v =  val

def assertion.maps (v: Var): assertion Env :=
  fun (e: Env) => ∃ (val: v.kind.eval), (e.get v) = val

notation "h[" v:80 "↦" val:80 "]" => assertion.mapsto v val
notation "hprop(" p:80 ")" => assertion.prop p
notation "h[" v:80 "↦" "?" "]" => assertion.maps v
-- def hoare_triple_region (P: assertion Env) (r: Expr .R) (Q: assertion (Except ErrorKind NamedVal)) : Prop :=
--  ∀ (e: Env) (sem : (o : Op') → TopM o.retkind.eval) (v: Val), P e -> Q (r.denote sem v e)

-- more conceptual proof, still does not use full hoare logic machinery, only a fragment.
section SubXXHoare
-- x - x ~= 0

theorem Int.sub_n_n (n: Int) : n - n = 0 := by {
  linarith
}

def sub_x_x_equals_zero (res: String) (arg: String) (pairname: String) : Peephole := {
  find := [dsl_ops|
    % $(pairname) = tuple( %$(arg) : int, %$(arg) : int);
    % $(res) : int = "arith.sub" ( %$(pairname) : int × int);
  ]
  replace := [dsl_ops|
    % $(res) : int = "arith.constant" {0};
  ]
  sem := Arith.sem,
  replaceCorrect := fun env => by {
    simp[Expr.denote, TopM.get, TopM.set, ReaderT.get, bind, Env.get, ReaderT.bind, pure, ReaderT.pure, Semantic.run,
      Arith.sem, Arith.sub, Arith.const, Env.set] at *;
    }
 }

end SubXXHoare

section ForFusionHoare
-- peel loop iterations out.

-- e with x as subexpr
-- (fun v => e' v) x = e
def pure_rewrite [M: Monad m] [LM: LawfulMonad m]
  (ma: m α)
  (maval: α)
  (MAVAL: ma = pure maval)
  (compute: α -> m β) :
  compute maval = ma >>= compute := by {
    simp[MAVAL];
}

def for_fusion (r: Region) (n m : String)
  (n_plus_m n_m: String)
  (out1 out2: String)
  (n_init m_out1: String)
  (init n_plus_m_init: String) : Peephole := {
  find := [dsl_ops|
    % $(n_init) = tuple( %$(n) : nat, %$(init) : int);
    % $(out1) : int = "scf.for" ( %$(n_init) : nat × int) [$(r)];
    % $(m_out1) = tuple( %$(m) : nat, %$(out1) : int);
    % $(out2) : int = "scf.for" ( %$(m_out1) : nat × int) [$(r)];
  ]
  replace := [dsl_ops|
    % $(n_m) = tuple( %$(n) : nat, %$(m) : nat);
    % $(n_plus_m) : nat = "arith.add" ( %$(n_m) : nat × nat);
    % $(n_plus_m_init) = tuple( %$(n_plus_m) : nat, %$(init) : int);
    % $(out2) : int = "scf.for" ( %$(n_plus_m_init) : nat × int) [$(r)];
  ]
  sem :=
    fun name =>
    match name with
    | .arith.add => Arith.add.run
    | .scf.for =>  Scf.for_.run
    | _ => fun op env => return default

  replaceCorrect := fun env1 => by {
    simp;
    generalize SEM:((fun name =>
      match (motive := OpName → (o : Op') → TopM (Kind.eval o.retkind)) name with
      | OpName.arith.add => Arith.add.run
      | OpName.scf.for => Scf.for_.run
      | x => fun op env => default)) = sem;

    simp only[Expr.denote_opscons, Expr.denote_tuple, TopM_get];
    simp[TopM_get];
    simp[TopM_get];
    -- simp does not simplify 'TopM_get' again.
    sorry
    -- rw[Expr.denote_opscons];
    -- rw[Expr.denote_tuple];
    -- simp[Expr.denote];
    -- sorry

    -- simp[TopM.get, ReaderT.get,  Env.get, bind, ReaderT.bind, TopM.set, pure, ReaderT.pure, Scf.for_, Env.set, Expr.denote];
    -- simp[Expr.denote, TopM.get, TopM.set, ReaderT.get, bind, Env.get, ReaderT.bind, pure, ReaderT.pure, Semantic.run];
  }
}
end ForFusionHoare

end RewritingHoare

namespace Theorems

end Theorems

namespace Test
def runRegionTest
  (sem: Semantics)
  (r: AST.Region)
  (expected: Val)
  (arg: Val := Val.unit )
  (env: Env := Env.empty) : IO Bool := do
    IO.println r
    let v : Val := (r.denote sem arg).run env
    if v == expected
    then
      IO.println s!"{v}. OK"; return True
    else
        IO.println s!"ERROR: computed '{v}', expected '{expected}'."
        return False

-- there is no notion of failure.
-- def runRegionTestXfail :
--   (sem: Semantics) →
--   (r: AST.Region) →
--   (arg: Val := Val.unit ) →
--   (env: Env := Env.empty) → IO Bool :=
--   fun sem r arg env => do
--     IO.println r
--     let v : Val := (r.denote sem arg).run env
--     match v? with
--     | .ok v =>
--         IO.println s!"ERROR: expected failure, but succeeded with '{v}'.";
--         return False
--     | .error e =>
--       IO.println s!"OK: Succesfully xfailed '{e}'."
--       return True

end Test

open Arith Scf in
def arith_plus_scf.sem : Semantics := fun name =>
    match name with
    | .arith.constant => const.run
    | .arith.add => add.run
    | .arith.sub => sub.run
    | .scf.if => if_.run
    | .scf.for => for_.run
    | .scf.twice => twice.run
    | .scf.run => run.run
    | _ => default

open Arith Scf DSL Examples Test in
def main : IO UInt32 := do
  let tests :=
  [runRegionTest
    (sem := Arith.sem)
    (r := eg_region_sub)
    (expected := { kind := .int, val := -1}),
  runRegionTest
    (sem := Arith.sem)
    (r := eg_arith_shadow)
    (expected := { kind := .int, val := 4}),
  runRegionTest
    (sem := arith_plus_scf.sem)
    (r := eg_scf_ite_true)
    (expected := { kind := .int, val := 42}),
  runRegionTest
    (sem := arith_plus_scf.sem)
    (r := eg_scf_ite_false)
    (expected := { kind := .int, val := 0}),
  runRegionTest
    (sem := arith_plus_scf.sem)
    (r := eg_scf_run_twice)
    (expected := { kind := .int, val := 42}),
  runRegionTest
    (sem := arith_plus_scf.sem)
    (r := eg_scf_well_scoped)
    (expected := { kind := .int, val := 42}),
  runRegionTest
    (sem := arith_plus_scf.sem)
    (r := eg_scf_for)
    (expected := { kind := .int, val := 42})]
  let mut total := 0
  let mut correct := 0
  for t in tests do
    total := total + 1
    IO.println s!"---Test {total}---"
    let pass? ← t
    if pass? then correct := correct + 1
  IO.println "---"
  IO.println s!"Tests: {correct} successful/{total}"
  return (if correct == total then 0 else 1)


import category_theory.preadditive
import category_theory.abelian.projective
import category_theory.abelian.diagram_lemmas.four

import data.matrix.notation

import for_mathlib.abelian_category
import for_mathlib.fin_functor
import for_mathlib.split_exact

noncomputable theory

open category_theory
open category_theory.limits
open category_theory.preadditive

universes v u

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

@[ext]
structure short_exact_sequence [has_images 𝒞] [has_zero_morphisms 𝒞] [has_kernels 𝒞] :=
(fst snd trd : 𝒞)
(f : fst ⟶ snd)
(g : snd ⟶ trd)
[mono'  : mono f]
[epi'   : epi g]
(exact' : exact f g)

namespace short_exact_sequence

attribute [instance] mono' epi'

variables {𝒞} [has_images 𝒞] [has_zero_morphisms 𝒞] [has_kernels 𝒞]

@[simp, reassoc] lemma f_comp_g (A : short_exact_sequence 𝒞) : A.f ≫ A.g = 0 := A.exact'.w

@[ext]
structure hom (A B : short_exact_sequence 𝒞) :=
(fst : A.1 ⟶ B.1)
(snd : A.2 ⟶ B.2)
(trd : A.3 ⟶ B.3)
(sq1' : fst ≫ B.f = A.f ≫ snd . obviously)
(sq2' : snd ≫ B.g = A.g ≫ trd . obviously)

namespace hom

restate_axiom sq1' sq1
restate_axiom sq2' sq2

attribute [reassoc] sq1 sq2

end hom

instance : quiver (short_exact_sequence 𝒞) := ⟨hom⟩

def id (A : short_exact_sequence 𝒞) : A ⟶ A :=
{ fst := 𝟙 _,
  snd := 𝟙 _,
  trd := 𝟙 _,
  sq1' := by simp only [category.id_comp, category.comp_id],
  sq2' := by simp only [category.id_comp, category.comp_id], }

def comp {A B C : short_exact_sequence 𝒞} (f : A ⟶ B) (g : B ⟶ C) : A ⟶ C :=
{ fst := f.1 ≫ g.1,
  snd := f.2 ≫ g.2,
  trd := f.3 ≫ g.3,
  sq1' := by rw [category.assoc, hom.sq1, hom.sq1_assoc],
  sq2' := by rw [category.assoc, hom.sq2, hom.sq2_assoc], }

instance : category (short_exact_sequence 𝒞) :=
{ id := id,
  comp := λ A B C f g, comp f g,
  id_comp' := by { intros, ext; dsimp; apply category.id_comp, },
  comp_id' := by { intros, ext; dsimp; apply category.comp_id, },
  assoc' := by { intros, ext; dsimp; apply category.assoc, },
  .. (infer_instance : quiver (short_exact_sequence 𝒞)) }

@[simp] lemma id_fst (A : short_exact_sequence 𝒞) : hom.fst (𝟙 A) = 𝟙 A.1 := rfl
@[simp] lemma id_snd (A : short_exact_sequence 𝒞) : hom.snd (𝟙 A) = 𝟙 A.2 := rfl
@[simp] lemma id_trd (A : short_exact_sequence 𝒞) : hom.trd (𝟙 A) = 𝟙 A.3 := rfl

variables {A B C : short_exact_sequence 𝒞} (f : A ⟶ B) (g : B ⟶ C)

@[simp, reassoc] lemma comp_fst : (f ≫ g).1 = f.1 ≫ g.1 := rfl
@[simp, reassoc] lemma comp_snd : (f ≫ g).2 = f.2 ≫ g.2 := rfl
@[simp, reassoc] lemma comp_trd : (f ≫ g).3 = f.3 ≫ g.3 := rfl

variables (𝒞)

@[simps] def Fst : short_exact_sequence 𝒞 ⥤ 𝒞 :=
{ obj := fst, map := λ A B f, f.1 }

@[simps] def Snd : short_exact_sequence 𝒞 ⥤ 𝒞 :=
{ obj := snd, map := λ A B f, f.2 }

@[simps] def Trd : short_exact_sequence 𝒞 ⥤ 𝒞 :=
{ obj := trd, map := λ A B f, f.3 }

@[simps] def f_nat : Fst 𝒞 ⟶ Snd 𝒞 :=
{ app := λ A, A.f,
  naturality' := λ A B f, f.sq1 }

@[simps] def g_nat : Snd 𝒞 ⟶ Trd 𝒞 :=
{ app := λ A, A.g,
  naturality' := λ A B f, f.sq2 }

instance : has_zero_morphisms (short_exact_sequence 𝒞) :=
{ has_zero := λ A B, ⟨{ fst := 0, snd := 0, trd := 0 }⟩,
  comp_zero' := by { intros, ext; apply comp_zero },
  zero_comp' := by { intros, ext; apply zero_comp }, }
.

@[simp] lemma hom_zero_fst : (0 : A ⟶ B).1 = 0 := rfl

@[simp] lemma hom_zero_snd : (0 : A ⟶ B).2 = 0 := rfl

@[simp] lemma hom_zero_trd : (0 : A ⟶ B).3 = 0 := rfl

variables {𝒞}

protected def functor (A : short_exact_sequence 𝒞) : fin 3 ⥤ 𝒞 :=
fin3_functor_mk ![A.1, A.2, A.3] A.f A.g

def functor_map {A B : short_exact_sequence 𝒞} (f : A ⟶ B) :
  Π i, A.functor.obj i ⟶ B.functor.obj i
| ⟨0,h⟩ := f.1
| ⟨1,h⟩ := f.2
| ⟨2,h⟩ := f.3
| ⟨i+3,hi⟩ := by { exfalso, revert hi, dec_trivial }

meta def aux_tac : tactic unit :=
`[simp only [hom_of_le_refl, functor.map_id, category.id_comp, category.comp_id]]

lemma functor_map_naturality {A B : short_exact_sequence 𝒞} (f : A ⟶ B) :
  ∀ (i j : fin 3) (hij : i ≤ j),
    functor_map f i ≫ B.functor.map hij.hom = A.functor.map hij.hom ≫ functor_map f j
| ⟨0,hi⟩ ⟨0,hj⟩ hij := by aux_tac
| ⟨1,hi⟩ ⟨1,hj⟩ hij := by aux_tac
| ⟨2,hi⟩ ⟨2,hj⟩ hij := by aux_tac
| ⟨0,hi⟩ ⟨1,hj⟩ hij := f.sq1
| ⟨1,hi⟩ ⟨2,hj⟩ hij := f.sq2
| ⟨i+3,hi⟩ _ _ := by { exfalso, revert hi, dec_trivial }
| _ ⟨j+3,hj⟩ _ := by { exfalso, revert hj, dec_trivial }
| ⟨i+1,hi⟩ ⟨0,hj⟩ H := by { exfalso, revert H, dec_trivial }
| ⟨i+2,hi⟩ ⟨1,hj⟩ H := by { exfalso, revert H, dec_trivial }
| ⟨0,hi⟩ ⟨2,hj⟩ hij :=
begin
  have h01 : (0 : fin 3) ⟶ 1 := hom_of_le dec_trivial,
  have h12 : (1 : fin 3) ⟶ 2 := hom_of_le dec_trivial,
  calc functor_map f ⟨0, hi⟩ ≫ B.functor.map hij.hom
      = functor_map f ⟨0, hi⟩ ≫ B.functor.map h01 ≫ B.functor.map h12 : _
  ... = (functor_map f ⟨0, hi⟩ ≫ B.functor.map h01) ≫ B.functor.map h12 : by rw category.assoc
  ... = (A.functor.map h01 ≫ functor_map f _) ≫ B.functor.map h12 : _
  ... = A.functor.map h01 ≫ functor_map f _ ≫ B.functor.map h12 : category.assoc _ _ _
  ... = A.functor.map h01 ≫ A.functor.map h12 ≫ functor_map f _ : _
  ... = A.functor.map hij.hom ≫ functor_map f ⟨2, hj⟩ : _,
  { rw [← functor.map_comp], congr, },
  { congr' 1, exact f.sq1 },
  { congr' 1, exact f.sq2 },
  { rw [← functor.map_comp_assoc], congr, },
end

@[simps] def Functor : short_exact_sequence 𝒞 ⥤ fin 3 ⥤ 𝒞 :=
{ obj := short_exact_sequence.functor,
  map := λ A B f,
  { app := functor_map f,
    naturality' := λ i j hij, (functor_map_naturality f i j hij.le).symm },
  map_id' := λ A, by { ext i, fin_cases i; refl },
  map_comp' := λ A B C f g, by { ext i, fin_cases i; refl } }

end short_exact_sequence

namespace short_exact_sequence

variables {𝒞} [abelian 𝒞]
variables {A B C : short_exact_sequence 𝒞} (f : A ⟶ B) (g : B ⟶ C)

section iso

variables {A B C} (f g)

open_locale zero_object

/-- One form of the five lemma: if a morphism of short exact sequences has isomorphisms
as first and third component, then the second component is also an isomorphism. -/
lemma snd_is_iso (h1 : is_iso f.1) (h3 : is_iso f.3) : is_iso f.2 :=
@abelian.is_iso_of_is_iso_of_is_iso_of_is_iso_of_is_iso 𝒞 _ _
  0 A.1 A.2 A.3
  0 B.1 B.2 B.3
  0 A.f A.g
  0 B.f B.g
  0 f.1 f.2 f.3 (by rw [zero_comp, zero_comp]) f.sq1 f.sq2
  0 0
  0 0 0 (by rw [comp_zero, comp_zero])
  (exact_zero_left_of_mono _)
  A.exact'
  ((epi_iff_exact_zero_right _).mp infer_instance)
  (exact_zero_left_of_mono _)
  B.exact'
  ((epi_iff_exact_zero_right _).mp infer_instance) _ _ _ _

/-- One form of the five lemma: if a morphism `f` of short exact sequences has isomorphisms
as first and third component, then `f` itself is an isomorphism. -/
lemma is_iso_of_fst_of_trd (h1 : is_iso f.1) (h3 : is_iso f.3) : is_iso f :=
{ out :=
  begin
    haveI : is_iso f.2 := snd_is_iso f h1 h3,
    refine ⟨⟨inv f.1, inv f.2, inv f.3, _, _⟩, _, _⟩,
    { dsimp, simp only [is_iso.inv_comp_eq, f.sq1_assoc, category.comp_id, is_iso.hom_inv_id], },
    { dsimp, simp only [is_iso.inv_comp_eq, f.sq2_assoc, category.comp_id, is_iso.hom_inv_id], },
    { ext; dsimp; simp only [is_iso.hom_inv_id], },
    { ext; dsimp; simp only [is_iso.inv_hom_id], },
  end }

@[simps] def iso_of_components (f₁ : A.1 ≅ B.1) (f₂ : A.2 ≅ B.2) (f₃ : A.3 ≅ B.3)
  (sq1 : f₁.hom ≫ B.f = A.f ≫ f₂.hom) (sq2 : f₂.hom ≫ B.g = A.g ≫ f₃.hom) :
  A ≅ B :=
{ hom := ⟨f₁.hom, f₂.hom, f₃.hom, sq1, sq2⟩,
  inv :=
  begin
    refine ⟨f₁.inv, f₂.inv, f₃.inv, _, _⟩; dsimp,
    rw [iso.inv_comp_eq, ← category.assoc, iso.eq_comp_inv, sq1],
    rw [iso.inv_comp_eq, ← category.assoc, iso.eq_comp_inv, sq2],
  end,
  hom_inv_id' := by { ext; apply iso.hom_inv_id, },
  inv_hom_id' := by { ext; apply iso.inv_hom_id, } }

@[simps] def iso_of_components' (f₁ : A.1 ≅ B.1) (f₂ : A.2 ⟶ B.2) (f₃ : A.3 ≅ B.3)
  (sq1 : f₁.hom ≫ B.f = A.f ≫ f₂) (sq2 : f₂ ≫ B.g = A.g ≫ f₃.hom) :
  A ≅ B :=
let F : A ⟶ B := ⟨f₁.hom, f₂, f₃.hom, sq1, sq2⟩ in
{ hom := F,
  inv :=
  begin
    haveI : is_iso F.2 := snd_is_iso _ infer_instance infer_instance,
    refine ⟨f₁.inv, inv F.2, f₃.inv, _, _⟩; dsimp,
    rw [iso.inv_comp_eq, ← category.assoc, is_iso.eq_comp_inv, sq1],
    rw [is_iso.inv_comp_eq, ← category.assoc, iso.eq_comp_inv, sq2],
  end,
  hom_inv_id' := by { ext; try { apply iso.hom_inv_id, }, apply is_iso.hom_inv_id },
  inv_hom_id' := by { ext; try { apply iso.inv_hom_id, }, apply is_iso.inv_hom_id } }

end iso

section split

/-- A short exact sequence `0 ⟶ A₁ -f⟶ A₂ -g⟶ A₃ ⟶ 0` is *left split*
if there exists a morphism `φ : A₂ ⟶ A₁` such that `f ≫ φ = 𝟙 A₁`. -/
def left_split (A : short_exact_sequence 𝒞) : Prop :=
∃ φ : A.2 ⟶ A.1, A.f ≫ φ = 𝟙 A.1

/-- A short exact sequence `0 ⟶ A₁ -f⟶ A₂ -g⟶ A₃ ⟶ 0` is *right split*
if there exists a morphism `φ : A₂ ⟶ A₁` such that `f ≫ φ = 𝟙 A₁`. -/
def right_split (A : short_exact_sequence 𝒞) : Prop :=
∃ χ : A.3 ⟶ A.2, χ ≫ A.g = 𝟙 A.3

variables {𝒜 : Type*} [category 𝒜] [abelian 𝒜]

lemma exact_of_split {A B C : 𝒜} (f : A ⟶ B) (g : B ⟶ C) (χ : C ⟶ B) (φ : B ⟶ A)
  (hfg : f ≫ g = 0) (H : φ ≫ f + g ≫ χ = 𝟙 B) : exact f g :=
{ w := hfg,
  epi :=
  begin
    let ψ : (kernel_subobject g : 𝒜) ⟶ image_subobject f :=
      subobject.arrow _ ≫ φ ≫ factor_thru_image_subobject f,
    suffices : ψ ≫ image_to_kernel f g hfg = 𝟙 _,
    { convert epi_of_epi ψ _, rw this, apply_instance },
    rw ← cancel_mono (subobject.arrow _), swap, { apply_instance },
    simp only [image_to_kernel_arrow, image_subobject_arrow_comp, category.id_comp, category.assoc],
    calc (kernel_subobject g).arrow ≫ φ ≫ f
        = (kernel_subobject g).arrow ≫ 𝟙 B : _
    ... = (kernel_subobject g).arrow        : category.comp_id _,
    rw [← H, preadditive.comp_add],
    simp only [add_zero, zero_comp, kernel_subobject_arrow_comp_assoc],
  end }

-- move this
lemma exact_inl_snd (A B : 𝒜) : exact (biprod.inl : A ⟶ A ⊞ B) biprod.snd :=
exact_of_split _ _ biprod.inr biprod.fst biprod.inl_snd biprod.total

def mk_of_split {A B C : 𝒜} (f : A ⟶ B) (g : B ⟶ C) (φ : B ⟶ A) (χ : C ⟶ B)
  (hfg : f ≫ g = 0) (hφ : f ≫ φ = 𝟙 A) (hχ : χ ≫ g = 𝟙 C) (H : φ ≫ f + g ≫ χ = 𝟙 B) :
  short_exact_sequence 𝒜 :=
{ fst := A,
  snd := B,
  trd := C,
  f := f,
  g := g,
  mono' := by { haveI : mono (f ≫ φ), { rw hφ, apply_instance }, exact mono_of_mono f φ, },
  epi' := by { haveI : epi (χ ≫ g), { rw hχ, apply_instance }, exact epi_of_epi χ g, },
  exact' := exact_of_split f g χ φ hfg H }

def mk_of_split' {A B C : 𝒜} (f : A ⟶ B) (g : B ⟶ C)
  (H : ∃ (φ : B ⟶ A) (χ : C ⟶ B), f ≫ g = 0 ∧ f ≫ φ = 𝟙 A ∧ χ ≫ g = 𝟙 C ∧ φ ≫ f + g ≫ χ = 𝟙 B) :
  short_exact_sequence 𝒜 :=
mk_of_split f g H.some H.some_spec.some H.some_spec.some_spec.1 H.some_spec.some_spec.2.1
  H.some_spec.some_spec.2.2.1 H.some_spec.some_spec.2.2.2

@[simp] def mk_split (A B : 𝒜) : short_exact_sequence 𝒜 :=
{ fst := A,
  snd := A ⊞ B,
  trd := B,
  f := biprod.inl,
  g := biprod.snd,
  exact' := exact_inl_snd _ _ }

/-- A *splitting* of a short exact sequence `0 ⟶ A₁ -f⟶ A₂ -g⟶ A₃ ⟶ 0` is
an isomorphism to the short exact sequence `0 ⟶ A₁ ⟶ A₁ ⊕ A₃ ⟶ A₃ ⟶ 0`,
where the left and right components of the isomorphism are identity maps. -/
structure splitting (A : short_exact_sequence 𝒜) extends A ≅ (mk_split A.1 A.3) :=
(fst_eq_id : hom.1 = 𝟙 A.1)
(trd_eq_id : hom.3 = 𝟙 A.3)

/-- A short exact sequence `0 ⟶ A₁ -f⟶ A₂ -g⟶ A₃ ⟶ 0` is *split* if there exist
`φ : A₂ ⟶ A₁` and `χ : A₃ ⟶ A₂` such that:
* `f ≫ φ = 𝟙 A₁`
* `χ ≫ g = 𝟙 A₃`
* `χ ≫ φ = 0`
* `φ ≫ f + g ≫ χ = 𝟙 A₂`
-/
def split (A : short_exact_sequence 𝒜) : Prop :=
∃ (φ : A.2 ⟶ A.1) (χ : A.3 ⟶ A.2),
   A.f ≫ φ = 𝟙 A.1 ∧ χ ≫ A.g = 𝟙 A.3 ∧ χ ≫ φ = 0 ∧ φ ≫ A.f + A.g ≫ χ = 𝟙 A.2

lemma mk_split_split (A B : 𝒜) : (mk_split A B).split :=
⟨biprod.fst, biprod.inr, biprod.inl_fst, biprod.inr_snd, biprod.inr_fst, biprod.total⟩

lemma splitting.split {A : short_exact_sequence 𝒜} (i : splitting A) : A.split :=
begin
  refine ⟨i.hom.2 ≫ biprod.fst ≫ i.inv.1, i.hom.3 ≫ biprod.inr ≫ i.inv.2, _⟩,
  simp only [category.assoc, ← hom.sq1_assoc, hom.sq2], dsimp,
  simp only [biprod.inl_fst_assoc, biprod.inr_snd_assoc, category.comp_id, category.assoc,
    ← comp_fst, ← comp_snd_assoc, ← comp_trd, i.to_iso.hom_inv_id, i.to_iso.inv_hom_id],
  dsimp,
  simp only [true_and, biprod.inr_fst_assoc, zero_comp, eq_self_iff_true, comp_zero,
    category.id_comp],
  simp only [hom.sq1, ← hom.sq2_assoc, ← comp_add],
  simp only [← category.assoc, ← add_comp, biprod.total,
    category.comp_id, ← comp_snd, i.to_iso.hom_inv_id], refl,
end

def left_split.splitting {A : short_exact_sequence 𝒜} (h : A.left_split) : A.splitting :=
{ to_iso := iso_of_components' (iso.refl _) (biprod.lift h.some A.g) (iso.refl _)
    (by { dsimp, simp only [category.id_comp], ext,
      { simpa only [biprod.inl_fst, biprod.lift_fst, category.assoc] using h.some_spec.symm, },
      { simp only [exact.w, f_comp_g, biprod.lift_snd, category.assoc, exact_inl_snd] } })
    (by { dsimp, simp only [category.comp_id, biprod.lift_snd], }),
  fst_eq_id := rfl,
  trd_eq_id := rfl }

def right_split.splitting {A : short_exact_sequence 𝒜} (h : A.right_split) : A.splitting :=
{ to_iso := iso.symm $ iso_of_components' (iso.refl _) (biprod.desc A.f h.some) (iso.refl _)
    (by { dsimp, simp only [biprod.inl_desc, category.id_comp], })
    (by { dsimp, simp only [category.comp_id], ext,
      { simp only [exact.w, f_comp_g, biprod.inl_desc_assoc, exact_inl_snd] },
      { simpa only [biprod.inr_snd, biprod.inr_desc_assoc] using h.some_spec, } }),
  fst_eq_id := rfl,
  trd_eq_id := rfl }

lemma tfae_split (A : short_exact_sequence 𝒜) :
  tfae [A.left_split, A.right_split, A.split, nonempty A.splitting] :=
begin
  tfae_have : 3 → 1, { rintro ⟨φ, χ, hφ, hχ, hχφ, H⟩, exact ⟨φ, hφ⟩ },
  tfae_have : 3 → 2, { rintro ⟨φ, χ, hφ, hχ, hχφ, H⟩, exact ⟨χ, hχ⟩ },
  tfae_have : 4 → 3, { rintro ⟨i⟩, exact i.split, },
  tfae_have : 1 → 4, { intro h, exact ⟨h.splitting⟩ },
  tfae_have : 2 → 4, { intro h, exact ⟨h.splitting⟩ },
  tfae_finish
end

end split

end short_exact_sequence

namespace short_exact_sequence

open category_theory.preadditive

variables {𝒞} [preadditive 𝒞] [has_images 𝒞] [has_kernels 𝒞]
variables (A B : short_exact_sequence 𝒞)

local notation `π₁` := congr_arg _root_.prod.fst
local notation `π₂` := congr_arg _root_.prod.snd

protected def hom_inj (f : A ⟶ B) : (A.1 ⟶ B.1) × (A.2 ⟶ B.2) × (A.3 ⟶ B.3) := ⟨f.1, f.2, f.3⟩

protected lemma hom_inj_injective : function.injective (short_exact_sequence.hom_inj A B) :=
λ f g h, let aux := π₂ h in
by { ext; [have := π₁ h, have := π₁ aux, have := π₂ aux]; exact this, }

instance : has_add (A ⟶ B) :=
{ add := λ f g,
  { fst := f.1 + g.1,
    snd := f.2 + g.2,
    trd := f.3 + g.3,
    sq1' := by { rw [add_comp, comp_add, f.sq1, g.sq1], },
    sq2' := by { rw [add_comp, comp_add, f.sq2, g.sq2], } } }

instance : has_neg (A ⟶ B) :=
{ neg := λ f,
  { fst := -f.1,
    snd := -f.2,
    trd := -f.3,
    sq1' := by { rw [neg_comp, comp_neg, f.sq1], },
    sq2' := by { rw [neg_comp, comp_neg, f.sq2], } } }

instance : has_sub (A ⟶ B) :=
{ sub := λ f g,
  { fst := f.1 - g.1,
    snd := f.2 - g.2,
    trd := f.3 - g.3,
    sq1' := by { rw [sub_comp, comp_sub, f.sq1, g.sq1], },
    sq2' := by { rw [sub_comp, comp_sub, f.sq2, g.sq2], } } }

instance has_nsmul : has_smul ℕ (A ⟶ B) :=
{ smul := λ n f,
  { fst := n • f.1,
    snd := n • f.2,
    trd := n • f.3,
    sq1' := by rw [nsmul_comp, comp_nsmul, f.sq1],
    sq2' := by rw [nsmul_comp, comp_nsmul, f.sq2] } }

instance has_zsmul : has_smul ℤ (A ⟶ B) :=
{ smul := λ n f,
  { fst := n • f.1,
    snd := n • f.2,
    trd := n • f.3,
    sq1' := by rw [zsmul_comp, comp_zsmul, f.sq1],
    sq2' := by rw [zsmul_comp, comp_zsmul, f.sq2] } }

variables (𝒞)

instance : preadditive (short_exact_sequence 𝒞) :=
{ hom_group := λ A B, (short_exact_sequence.hom_inj_injective A B).add_comm_group _
  rfl (λ _ _, rfl) (λ _, rfl) (λ _ _, rfl) (λ _ _, rfl) (λ _ _, rfl),
  add_comp' := by { intros, ext; apply add_comp },
  comp_add' := by { intros, ext; apply comp_add }, }
.

instance Fst_additive : (Fst 𝒞).additive := {}
instance Snd_additive : (Snd 𝒞).additive := {}
instance Trd_additive : (Trd 𝒞).additive := {}

end short_exact_sequence

end category_theory


import category_theory.isomorphism
import homotopy_theory.formal.cylinder.hep

import .category
import .colimits
import .cylinder
import .interval_endpoints
import .pair
import .pushout_lemmas

open set

open category_theory
open category_theory.category
local notation f ` ∘ `:80 g:80 := g ≫ f

open homotopy_theory.cylinder

namespace homotopy_theory.topological_spaces
open homotopy_theory.topological_spaces.Top
local notation `Top` := Top.{0}

-- The classical definition of a cofibration between topological
-- spaces as a map which satisfies the homotopy extension property.
def cofibration {A X : Top} (j : A ⟶ X) : Prop := hep 0 j

/-

* A cofibration in Top is an embedding. [Strøm, Note on cofibrations,
  Theorem 1.] Proof: Suppose j : A → X is a cofibration. Form the
  mapping cylinder Z as the pushout shown below, with induced map
  k : Z → IX.

      j
    A → X = X
  i₀↓   ↓   ↓i₀
   IA → Z → IX → Z
  i₁↑   ↑   ↑i₁
    A = A → X
          j

  By the homotopy extension property, we can find a map r from IX back
  to Z. So Z → X is the inclusion of a retract, hence in particular an
  embedding. The map i₀ : A → IA has closed image, so IA → Z is a
  homeomorphism onto its image away from the image of i₀. Thus the
  composition

      i₁
    A → IA → Z → IX

  is an embedding; but it equals i₁ ∘ j : A → IX, so j : A → X is an
  embedding as well.

-/

variables {A X : Top} {j : A ⟶ X}
local notation `i` := i.{0}

lemma embedding_i {ε} : embedding (i ε @> A) :=
embedding_of_embedding_comp (p @> A) embedding_id

lemma closed_i {ε} : is_closed (set.range (i ε @> A)) :=
have is_closed {p : Top.prod A I01 | p.snd ∈ ({I01_of_endpoint ε} : set I01)}, from
  continuous_iff_is_closed.mp continuous_snd _ is_closed_singleton,
begin
  convert this, ext p, cases p with a t,
  change (∃ a', (a', _) = (a, t)) ↔ _,
  simpa using eq_comm
end

lemma disjoint_i₀_i₁ : set.range (i 0 @> A) ∩ set.range (i 1 @> A) = ∅ :=
begin
  apply set.eq_empty_of_subset_empty, intros p hp,
  cases p with a t, rcases hp with ⟨⟨_, hp₀⟩, ⟨_, hp₁⟩⟩,
  have hp₀' : 0 = t := congr_arg prod.snd hp₀,
  have hp₁' : 1 = t := congr_arg prod.snd hp₁,
  have : (0 : I01) = 1 := hp₀'.trans hp₁'.symm,
  exact absurd (congr_arg subtype.val this)
    (show ¬(0 : ℝ) = 1, by norm_num)
end

lemma embedding_of_cofibration (h : cofibration j) : embedding j :=
let po := has_pushouts.pushout (i 0 @> A) j,
    Z := po.ob,
    k : Z ⟶ I.obj X :=
      po.is_pushout.induced (I &> j) (i 0 @> X) ((i 0).naturality j).symm,
    ⟨r, hr₀, hr₁⟩ := h Z po.map₁ po.map₀ po.is_pushout.commutes.symm in
have _ := hr₀.symm,
have hr : r ∘ k = 𝟙 _, by apply po.is_pushout.uniqueness; { rw ←assoc, simpa },
have e_z_ix : embedding k, from
  embedding_of_embedding_comp r (by rw hr; exact embedding_id),
have e_a_z : embedding (po.map₀ ∘ i 1 @> A), from
  comp_embedding_of_embedding_of_disjoint
    po.is_pushout (or.inl closed_i) embedding_i disjoint_i₀_i₁,
have embedding (k ∘ (po.map₀ ∘ i 1 @> A)), from e_z_ix.comp e_a_z,
have embedding (i 1 @> X ∘ j), begin
  convert this using 2,
  transitivity,
  exact (i 1).naturality j,
  simp
end,
embedding_of_embedding_comp _ this

lemma cofibration_iff_cofibered_of_embedding (e : embedding j) :
  cofibration j ↔ (pair.mk X (range j)).cofibered :=
let j' := Top.factor_through_incl j (range j) (subset.refl _) in
show hep 0 j ↔ hep 0 _, from
mem_iff_mem_of_isomorphic
  (homeomorphism_to_image_of_embedding e)
  (iso.refl X)
  (by ext p; refl)

lemma cofibration_iff_cofibered :
  cofibration j ↔ embedding j ∧ (pair.mk X (range j)).cofibered :=
iff.intro
  (assume h,
    have e : embedding j := embedding_of_cofibration h,
    ⟨e, (cofibration_iff_cofibered_of_embedding e).mp h⟩)
  (assume h, (cofibration_iff_cofibered_of_embedding h.1).mpr h.2)

section relative_cylinder
noncomputable theory

variables (j) (ha : is_closed (range j))
variable (hj : cofibration j)

lemma relative_cylinder' : ∃ Po : pushout (∂I &> j) (ii @> A),
  cofibration (Po.is_pushout.induced (ii @> X) (I &> j) (ii.naturality _)) ∧
  is_closed (range (Po.is_pushout.induced (ii @> X) (I &> j) (ii.naturality _))) :=
let P : pair := pair.mk X (range j) in
let j_ : homeomorphism A P.subspace :=
  homeomorphism_to_image_of_embedding (embedding_of_cofibration hj) in
let po := pair.po P I_01 ha I_01.is_closed in
let po' := Is_pushout_of_isomorphic po.transpose (∂I &> j) (ii @> A)
  ((∂I.map_iso j_).trans prod_doubleton) prod_doubleton (I.map_iso j_)
  (by apply coprod.uniqueness; refl)
  (by apply coprod.uniqueness; refl) in
let ind := po'.induced (ii @> X) (I &> j) (ii.naturality _) in
have ind = pair.incl _, begin
  dsimp [ind], apply po'.uniqueness,
  { apply coprod.uniqueness; simpa },
  { simpa },
end,
⟨⟨(P ⊗ I_01).subspace, _, _, po'⟩,
 begin
   change cofibration ind ∧ is_closed (range ind), rw this,
   refine ⟨prod_I_01_cofibered P ha (cofibration_iff_cofibered.mp hj).2, _⟩,
   convert @pair.prod.is_closed P I_01 ha I_01.is_closed using 1,
   apply subtype.range_val
 end⟩

end relative_cylinder

end homotopy_theory.topological_spaces


inductive Foo where
  | mk : Nat → Foo
  | boo : String → Foo

instance : ToString Foo where
  toString o := match o with
    | .mk n => aux1 n
    | .boo s => aux2 s
where
  aux1 (n : Nat) : String :=
    s!".mk {n}"
  aux2 (s : String) : String :=
    s!".boo {s}"

example : toString (Foo.mk 10) = ".mk 10" := rfl


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

/- The Yoneda embedding, as a functor `yoneda : C ⥤ (Cᵒᵖ ⥤ Type v₁)`,
   along with an instance that it is `fully_faithful`.

   Also the Yoneda lemma, `yoneda_lemma : (yoneda_pairing C) ≅ (yoneda_evaluation C)`. -/

import category_theory.natural_transformation
import category_theory.opposites
import category_theory.types
import category_theory.fully_faithful
import category_theory.natural_isomorphism

namespace category_theory

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

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

def yoneda : C ⥤ (Cᵒᵖ ⥤ Type v₁) :=
{ obj := λ X,
  { obj := λ Y, unop Y ⟶ X,
    map := λ Y Y' f g, f.unop ≫ g,
    map_comp' := λ _ _ _ f g, begin ext1, dsimp at *, erw [category.assoc] end,
    map_id' := λ Y, begin ext1, dsimp at *, erw [category.id_comp] end },
  map := λ X X' f, { app := λ Y g, g ≫ f } }

def coyoneda : Cᵒᵖ ⥤ (C ⥤ Type v₁) :=
{ obj := λ X,
  { obj := λ Y, unop X ⟶ Y,
    map := λ Y Y' f g, g ≫ f,
    map_comp' := λ _ _ _ f g, begin ext1, dsimp at *, erw [category.assoc] end,
    map_id' := λ Y, begin ext1, dsimp at *, erw [category.comp_id] end },
  map := λ X X' f, { app := λ Y g, f.unop ≫ g },
  map_comp' := λ _ _ _ f g, begin ext1, ext1, dsimp at *, erw [category.assoc] end,
  map_id' := λ X, begin ext1, ext1, dsimp at *, erw [category.id_comp] end }

namespace yoneda
@[simp] lemma obj_obj (X : C) (Y : Cᵒᵖ) : (yoneda.obj X).obj Y = (unop Y ⟶ X) := rfl
@[simp] lemma obj_map (X : C) {Y Y' : Cᵒᵖ} (f : Y ⟶ Y') :
  (yoneda.obj X).map f = λ g, f.unop ≫ g := rfl
@[simp] lemma map_app {X X' : C} (f : X ⟶ X') (Y : Cᵒᵖ) :
  (yoneda.map f).app Y = λ g, g ≫ f := rfl

lemma obj_map_id {X Y : C} (f : op X ⟶ op Y) :
  ((@yoneda C _).obj X).map f (𝟙 X) = ((@yoneda C _).map f.unop).app (op Y) (𝟙 Y) :=
by obviously

@[simp] lemma naturality {X Y : C} (α : yoneda.obj X ⟶ yoneda.obj Y)
  {Z Z' : C} (f : Z ⟶ Z') (h : Z' ⟶ X) : f ≫ α.app (op Z') h = α.app (op Z) (f ≫ h) :=
begin erw [functor_to_types.naturality], refl end

instance yoneda_fully_faithful : fully_faithful (@yoneda C _) :=
{ preimage := λ X Y f, (f.app (op X)) (𝟙 X),
  injectivity' := λ X Y f g p,
  begin
    injection p with h,
    convert (congr_fun (congr_fun h (op X)) (𝟙 X)); dsimp; simp,
  end }

/-- Extensionality via Yoneda. The typical usage would be
```
-- Goal is `X ≅ Y`
apply yoneda.ext,
-- Goals are now functions `(Z ⟶ X) → (Z ⟶ Y)`, `(Z ⟶ Y) → (Z ⟶ X)`, and the fact that these
functions are inverses and natural in `Z`.
```
-/
def ext (X Y : C)
  (p : Π {Z : C}, (Z ⟶ X) → (Z ⟶ Y)) (q : Π {Z : C}, (Z ⟶ Y) → (Z ⟶ X))
  (h₁ : Π {Z : C} (f : Z ⟶ X), q (p f) = f) (h₂ : Π {Z : C} (f : Z ⟶ Y), p (q f) = f)
  (n : Π {Z Z' : C} (f : Z' ⟶ Z) (g : Z ⟶ X), p (f ≫ g) = f ≫ p g) : X ≅ Y :=
@preimage_iso _ _ _ _ yoneda _ _ _ _
  (nat_iso.of_components (λ Z, { hom := p, inv := q, }) (by tidy))

-- We need to help typeclass inference with some awkward universe levels here.
instance prod_category_instance_1 : category ((Cᵒᵖ ⥤ Type v₁) × Cᵒᵖ) :=
category_theory.prod.{(max u₁ v₁)  v₁} (Cᵒᵖ ⥤ Type v₁) Cᵒᵖ

instance prod_category_instance_2 : category (Cᵒᵖ × (Cᵒᵖ ⥤ Type v₁)) :=
category_theory.prod.{v₁ (max u₁ v₁)} Cᵒᵖ (Cᵒᵖ ⥤ Type v₁)

end yoneda

namespace coyoneda
@[simp] lemma obj_obj (X : Cᵒᵖ) (Y : C) : (coyoneda.obj X).obj Y = (unop X ⟶ Y) := rfl
@[simp] lemma obj_map {X' X : C} (f : X' ⟶ X) (Y : Cᵒᵖ) :
  (coyoneda.obj Y).map f = λ g, g ≫ f := rfl
@[simp] lemma map_app (X : C) {Y Y' : Cᵒᵖ} (f : Y ⟶ Y') :
  (coyoneda.map f).app X = λ g, f.unop ≫ g := rfl
end coyoneda

class representable (F : Cᵒᵖ ⥤ Type v₁) :=
(X : C)
(w : yoneda.obj X ≅ F)

variables (C)
open yoneda

def yoneda_evaluation : Cᵒᵖ × (Cᵒᵖ ⥤ Type v₁) ⥤ Type (max u₁ v₁) :=
evaluation_uncurried Cᵒᵖ (Type v₁) ⋙ ulift_functor.{u₁}

@[simp] lemma yoneda_evaluation_map_down
  (P Q : Cᵒᵖ × (Cᵒᵖ ⥤ Type v₁)) (α : P ⟶ Q) (x : (yoneda_evaluation C).obj P) :
  ((yoneda_evaluation C).map α x).down = α.2.app Q.1 (P.2.map α.1 x.down) := rfl

def yoneda_pairing : Cᵒᵖ × (Cᵒᵖ ⥤ Type v₁) ⥤ Type (max u₁ v₁) :=
functor.prod yoneda.op (functor.id (Cᵒᵖ ⥤ Type v₁)) ⋙ functor.hom (Cᵒᵖ ⥤ Type v₁)

@[simp] lemma yoneda_pairing_map
  (P Q : Cᵒᵖ × (Cᵒᵖ ⥤ Type v₁)) (α : P ⟶ Q) (β : (yoneda_pairing C).obj P) :
  (yoneda_pairing C).map α β = yoneda.map α.1.unop ≫ β ≫ α.2 := rfl

def yoneda_lemma : yoneda_pairing C ≅ yoneda_evaluation C :=
{ hom :=
  { app := λ F x, ulift.up ((x.app F.1) (𝟙 (unop F.1))),
    naturality' :=
    begin
      intros X Y f, ext1, ext1,
      cases f, cases Y, cases X,
      dsimp at *, simp at *,
      erw [←functor_to_types.naturality,
           obj_map_id,
           functor_to_types.naturality,
           functor_to_types.map_id]
    end },
  inv :=
  { app := λ F x,
    { app := λ X a, (F.2.map a.op) x.down,
      naturality' :=
      begin
        intros X Y f, ext1,
        cases x, cases F,
        dsimp at *,
        erw [functor_to_types.map_comp]
      end },
    naturality' :=
    begin
      intros X Y f, ext1, ext1, ext1,
      cases x, cases f, cases Y, cases X,
      dsimp at *,
      erw [←functor_to_types.naturality, functor_to_types.map_comp]
    end },
  hom_inv_id' :=
  begin
    ext1, ext1, ext1, ext1, cases X, dsimp at *,
    erw [←functor_to_types.naturality,
         obj_map_id,
         functor_to_types.naturality,
         functor_to_types.map_id], refl,
  end,
  inv_hom_id' :=
  begin
    ext1, ext1, ext1,
    cases x, cases X,
    dsimp at *,
    erw [functor_to_types.map_id]
  end }.

variables {C}

@[simp] def yoneda_sections (X : C) (F : Cᵒᵖ ⥤ Type v₁) : (yoneda.obj X ⟹ F) ≅ ulift.{u₁} (F.obj (op X)) :=
nat_iso.app (yoneda_lemma C) (op X, F)

omit 𝒞
@[simp] def yoneda_sections_small {C : Type u₁} [small_category C] (X : C) (F : Cᵒᵖ ⥤ Type u₁) : (yoneda.obj X ⟹ F) ≅ F.obj (op X) :=
yoneda_sections X F ≪≫ ulift_trivial _

end category_theory