import prop program

import data.finmap
import category_theory.category
import tactic.interactive
import tactic.tauto
import tactic.linarith
import tactic.monotonicity

namespace refinement

open category_theory
open separation separation.hProp memory

variables value : Type

-- local notation `heap` := heap value
-- local notation `hProp` := hProp value

def state := Σ t : Type*, t → hProp value

section value

variables {value}

def state.mk {s} (p : s → hProp value) : state value := ⟨s,p⟩

structure inst (x : state value) :=
 (σ : x.fst) (h : heap value)
 (shape : x.snd σ h)

end value

def hom (x y : state value) := inst x → set (inst y)

instance : category (state value) :=
{ hom := hom value,
  id := λ _, pure,
  comp := λ x y z A₀ A₁, A₀ >=> A₁,
  id_comp' := λ x y f, fish_pipe _,
  comp_id' := λ x y f, fish_pure _,
  assoc' := λ x y z w f g h, fish_assoc _ _ _ }

universes u

variables {value} {s s' α : Type u}

def sat (p : s → hProp value) (m : ST value s α) (q : α → s → hProp value) : Prop :=
∀ σ σ' h h' x, (x,σ',h') ∈ m.run (σ,h) → p σ h → q x σ' h'

variables {p : s → hProp value} {m : ST value s punit} {q : s → hProp value}
variables {p' : s' → hProp value} {m' : ST value s' punit} {q' : s' → hProp value}

open function (const)

notation `♯` := by repeat { ext : 1; try {refl} }

def attach (x : set α) : set { a // a ∈ x } :=
{ a | a.1 ∈ x }

def arrow (hm : sat p m (const _ q)) : state.mk p ⟶ state.mk q :=
show inst (state.mk p) → set (inst (state.mk q)), from
λ ⟨σ,h,hp⟩,
do ⟨⟨x₀,x₁,x₂⟩,hx⟩ ← attach (m.run (σ,h)),
   return (inst.mk x₁ x₂ (hm σ _ h _ x₀ hx hp))

def abs {s s' : Type*} (R : s → s' → Prop) {p : s → hProp value} {p' : s' → hProp value} : state.mk p ⟶ state.mk p' :=
show inst (state.mk p) → set (inst (state.mk p')), from
λ ⟨σ,h,hp⟩ ⟨σ',h',hp'⟩,
R σ σ'

def repr {s s' : Type*} (R : s → s' → Prop) {p : s → hProp value} {p' : s' → hProp value} : state.mk p' ⟶ state.mk p :=
show inst (state.mk p') → set (inst (state.mk p)), from
λ ⟨σ,h,hp⟩ ⟨σ',h',hp'⟩,
R σ' σ

def ref {X Y : state value} (a b : X ⟶ Y) : Prop :=
∀ i, a i ⊆ b i

notation `⦃` p `⦄` := inst (state.mk p)

def rel (m : ST value s punit) (x : ⦃p⦄) (y : ⦃q⦄) : Prop :=
(punit.star,y.1,y.2) ∈ m.run (x.1,x.2)

infix ` ⊑ `:50 := ref

-- #exit
-- lemma sound (R : s → s' → Prop) (hc : sat p m (const _ q)) (ha : sat p' m' (const _ q'))
--   (hh : arrow hc ≫ abs R ⊑ abs R ≫ arrow ha) :
--   ∀ (c : ⦃p⦄) (c' : ⦃q⦄) (a : ⦃p'⦄),
--     R c.1 a.1 → rel m c c' →
--     ∃ a' : ⦃q'⦄, rel m' a a' ∧ R c'.1 a'.1 :=
-- begin
--   introv hR hm,
--   let ha' := m'.run (a.1,a.2),
--   let a' : ⦃q'⦄ := ⟨ha'.2.1,ha'.2.2,ha _ _ _ _ _ ♯ a.shape⟩,
--   existsi [a'], split, repeat {ext : 1; try { refl <|> apply punit_eq }},
--   specialize @hh c a', cases c, simp [(≫),(>=>),arrow,return,abs] at hh,
--   specialize @hh a',
-- end

lemma sound' (R : s → s' → Prop) (hc : sat p m (const _ q)) (ha : sat p' m' (const _ q'))
  (hh : repr R ≫ arrow hc ⊑ arrow ha ≫ repr R) :
  ∀ (c : ⦃p⦄) (c' : ⦃q⦄) (a : ⦃p'⦄),
    R c.1 a.1 → rel m c c' →
    ∃ a' : ⦃q'⦄, rel m' a a' ∧ R c'.1 a'.1 :=
begin
  introv hR hm,
  let a' : set ⦃q'⦄ := arrow ha a,
  specialize @hh a c' ⟨_,_,hm⟩,
  casesm* ⦃ _ ⦄, -- squeeze_simp [arrow,(≫),(>=>),(>>=),attach,repr] at hh,
  simp only [category_struct.comp, fish, arrow, return, set.bUnion_singleton, and_imp, attach,
             exists_prop, set.mem_Union, set.bind_def, set.pure_def, exists_imp_distrib] at hh,
  revert hh, apply exists_imp_exists, rintros ⟨a,b,c⟩,
  simp only [repr, true_and, and_imp, subtype.val_prop, bex_imp_distrib, subtype.exists,
             set.mem_set_of_eq, exists_imp_distrib, prod.exists, rel, arrow, return, pure,
             set.mem_singleton_iff ],
  introv h₀ h₁ h₂ h₃, subst x_1, subst x_2, cases x,
  exact ⟨h₀,h₃⟩,
  simp [set.mem_singleton_iff],
  existsi c, cases c,
  apply set.ext, intro,
  simp only [arrow, set.Union, exists_prop, set.mem_Union, set.bind_def, subtype.exists, prod.exists],
  split,
  { rintro ⟨h₀,⟨⟨⟩,b,c,h₁,h₂,h₃⟩⟩, simp only [return, set.mem_singleton_iff, and_self, set.pure_def] at h₃, subst x, exact h₂, },
  { intro h, cases a, existsi [hR,punit.star,_,_,h,h],
    cases x, simp only [return, set.mem_singleton_iff, and_self, set.pure_def, eq_self_iff_true], },
end

open state_t

-- instance {α} : has_coe (option α) (ST value s α) := _

/-
 * framing
 * assignment
 * alloc
 * dealloc
 * sequencing
 * loop
-/

-- def spec_alloc (s) (free : set ptr) (p : hProp value) (m : ST value s α) (q : α → hProp value) : Prop :=
-- ∀ σ σ' h h' frame x, (x,σ',h') ∈ m.run (σ,h) → (∀ p ∈ free, ¬ p ∈ h) → holds h frame p → holds h' frame (q x)

-- lemma And_frame {α} {p : hProp value} {m : ST value s α} {q : α → hProp value}
--   (h : ) :
--   spec s p m q := _

-- lemma p_exists_one_point {α} {p : α → hProp value} (x : α)
--   (h : ∀ y, p y =*> [| x = y |] ⊛ True) :
--   p_exists p = p x :=
-- begin
--   ext; dsimp [p_exists]; split,
--   { rintros ⟨x',HH⟩, specialize h _ _ HH, admit },
--   apply @Exists.intro _ (λ x, p x h_1),
-- end

end refinement


/-
## MLIR types

This file implements support for builtin MLIR types, as well as conversion to
concrete Lean types and values.

In order to support the untyped SSA environment (which doesn't require name
definitions/uses to match as this is proven later), every concrete MLIR type
should be Inhabited so that `SSAEnvE.Get` can pretend to return default values.

Current properly-supported MLIR built-in types:

* Function type [τ₁ → τ₂]: Lean function
* Tuple type [(τ₁, ..., τₙ)]: n-ary product of underlying types
* Tensors [tensor<DxDx?xDx?xτ>]: RankedTensor type

Types that need improvements or refinements:

* Unsigned finite integer types [u32, etc]
  TODO: Model u32/etc with lean's Uint{8/16/32/64} or restart with Fin
* Integer types [i32, etc]: currently use Int
  TODO: Model i32/etc properly, probably restarting from Fin
* Float types [f16, f32, f64, f80, f128]. todo
  No idea how to model floats of different precisions in Lean?
* Vectors: could expand on tensors. todo
* Unranked tensors: todo
* User types: should use a typeclass
  TODO: Model user types
-/

import MLIRSemantics.Util.Arith
import MLIRSemantics.Util.List
import MLIRSemantics.Fitree

import MLIR.AST
open MLIR.AST

def shape_prod: List Nat → Nat :=
  List.foldr (·*·) 1

theorem shape_prod_nil: shape_prod (0::l) = 0 := by
  induction l <;> simp [shape_prod, List.foldr]

instance: OfNat Dimension (n: Nat) where
  ofNat := Dimension.Known n

/-
## Shape inference on literal tensors

This section defines shape verification and shape inference for TensorElem
(tensor literals), *excluding the case of uniform tensor literals*. The shape
inference is proven correct, and the `flatten` method is defined that exports
the tensor literal to a flat array suitable for use in a `RankedTensor`.

`RankedTensor` provides the functions that actually turn tensor literals into
ranked tensors and properly handle uniform tensor literals.

TODO: Integrate TensorElem invariants into the verifier
-/

namespace MLIR.AST.TensorElem

-- Check whether a tensor literal matches a concrete shape
def hasShape: TensorElem → List Nat → Bool
  | TensorElem.empty, _ =>
      false
  | TensorElem.int _, [] =>
      true
  | TensorElem.int _, _::_ =>
      false
  | TensorElem.bool _, [] =>
      true
  | TensorElem.bool _, _::_ =>
      false
  | TensorElem.float _, [] =>
      true
  | TensorElem.float _, _::_ =>
      false
  | TensorElem.nested l, rank::size =>
      l.length = rank ∧ l.all (hasShape . size)
  | TensorElem.nested _, [] =>
      false

-- Check whether a tensor literal has a particular data type
def hasType: TensorElem → MLIRTy → Bool
  | TensorElem.int _, MLIRTy.int _ =>
      -- TODO: Check bounds
      true
  | TensorElem.bool _, MLIRTy.int 1 =>
      true
  | TensorElem.float _, MLIRTy.float _ =>
      true
  | TensorElem.nested [], τ =>
      true
  | TensorElem.nested (e::l), τ =>
      e.hasType τ ∧ (TensorElem.nested l).hasType τ
  | _, _ =>
      false

def hasType_list_1 {l τ}: hasType (.nested l) τ → l.all (hasType . τ) := by
  induction l; simp
  case cons e l ih =>
    simp [hasType, List.all_cons]
    intro h
    simp [ih h.2, h.1]

def hasType_list_2 {l τ}: l.all (hasType . τ) → hasType (.nested l) τ := by
  induction l; simp [hasType]
  case cons e l ih =>
    simp [hasType, List.all_cons]
    intro h
    simp [ih h.2, h.1]

def mapWithType {α τ} l (f: (e: TensorElem) → (h: e.hasType τ) → α)
    (h: hasType (TensorElem.nested l) τ): List α :=
  match l, h with
  | [], h =>
      []
  | e::l, h =>
      let h₁ := (by simp [hasType] at h; apply h.1)
      let h₂ := (by simp [hasType] at h; apply h.2)
      f e h₁ :: mapWithType l f h₂

-- Shape inference function; this determines the unique shape that we allow a
-- non-uniform tensor can have (hasShape is more liberal with empty lists, but
-- the MLIR compiler is not)
def inferredShape: TensorElem → Option (List Nat)
  | TensorElem.empty =>
      none
  | TensorElem.int _ =>
      some []
  | TensorElem.bool _ =>
      some []
  | TensorElem.float _ =>
      some []
  | TensorElem.nested [] =>
      some [0]
  | TensorElem.nested (e::l) =>
      Option.bind (inferredShape e) fun s1 =>
      Option.bind (inferredShape (TensorElem.nested l)) fun s2 =>
      match s2 with
      | [] => none /- impossible -/
      | head :: tail =>
        if s1 = tail then some ((head+1) :: s1) else none

-- First let's prove the list case equivalent to a more readable form

theorem inferredShape_cons: ∀ head tail s_head s_tail,
    inferredShape (TensorElem.nested tail) = some (s_head :: s_tail) →
    inferredShape head = some s_tail →
    inferredShape (TensorElem.nested (head :: tail)) =
      some ((s_head+1) :: s_tail) := by
  intros head tail s_head s_tail H1 H2
  simp [inferredShape, H2, Option.bind, H1];

theorem inferredShape_cons_inv: ∀ {head tail s_head s_tail},
    inferredShape (TensorElem.nested (head::tail)) = some (s_head::s_tail) →
    s_head > 0 ∧
    inferredShape head = some s_tail ∧
    inferredShape (TensorElem.nested tail) = some ((s_head-1) :: s_tail) := by
  intros head tail s_head s_tail
  simp [inferredShape]
  cases inferredShape head <;> simp [Option.bind]
  case some head_shape =>
  cases inferredShape (TensorElem.nested tail) <;> simp [Option.bind]
  case some tail_shape =>
  cases tail_shape <;> simp
  case cons s_head' s_tail' =>
  apply dite (head_shape = s_tail')
  . intros Heq; rw [Heq]; simp
    intros H; rw [←H.1, Nat.add_sub_self_right, ←H.2]
    exact ⟨by simp_arith, rfl, rfl, rfl⟩
  . intros Hne; simp [Hne]

theorem inferredShape_list {l head tail}:
    inferredShape (TensorElem.nested l) = some (head::tail) →
    head = l.length ∧ l.all (inferredShape . = some tail) := by
  revert head tail; induction l <;> simp
  case nil =>
    intros head tile H; simp [inferredShape, List.all_nil] at *; simp [H.1]
  case cons head tail ih =>
    intros s_head s_tail H
    let H' := inferredShape_cons_inv H
    specialize (ih H'.2.2)
    constructor
    . simp [←ih.1, Nat.succ_eq_add_one, Nat.minus_plus_one H'.1]
    . simp [List.all_cons]
      constructor; exact H'.2.1; exact ih.2

theorem inferredShape_list_to_cons {l s}:
    inferredShape (TensorElem.nested l) = some s →
    ∃ tail, s = l.length :: tail := by
  cases s <;> simp [inferredShape]
  case nil =>
    cases l <;> simp [inferredShape]
    case cons head tail =>
      cases inferredShape head <;> simp [Option.bind]
      cases inferredShape (TensorElem.nested tail) <;> simp [Option.bind]
      case some.some s1 s2 =>
        cases s2 <;> simp
        case cons s2_head s2_tail =>
          apply dite (s1 = s2_tail) <;> intros H <;> simp [H]
  case cons s_head s_tail =>
    intro H
    let H' := inferredShape_list H
    apply Exists.intro s_tail
    exact ⟨H'.1, rfl⟩

-- We can now show that the shape inference function is correct

theorem hasShape_inferredShape_1:
  ∀ (e: TensorElem) (shape: List Nat),
    e.inferredShape = some shape → e.hasShape shape := by
  intro e
  -- Cannot use the [induction] tactic because TensorElem is a nested inductive
  -- and the tactic only supports recursors with a single motive
  apply @TensorElem.recOn
    (motive_1 := fun e =>
      ∀s, e.inferredShape = some s → e.hasShape s)
    (motive_2 := fun l =>
      ∀s, l.all (TensorElem.inferredShape . = some s) →
        l.all (TensorElem.hasShape . s))
  case int =>
    intros _ s H; cases s <;> simp [inferredShape, hasShape] at *
  case bool =>
    intros _ s H; cases s <;> simp [inferredShape, hasShape] at *
  case float =>
    intros _ s H; cases s <;> simp [inferredShape, hasShape] at *
  case nested =>
    intros l motive_2 s H
    let H' := inferredShape_list_to_cons H
    cases H'; case intro s_tail Hs =>
      rw [Hs]; rw [Hs] at H; clear Hs H'
      let H' := inferredShape_list H
      simp [hasShape, motive_2 _ H'.2]
  case empty =>
    simp [inferredShape]
  case nil =>
    intros s H; simp [List.all_nil]
  case cons =>
    intros head tail motive_1 ih s H; simp [List.all_cons] at *
    simp [motive_1 _ H.1, ih _ H.2]

end MLIR.AST.TensorElem

/-
### Vector types
TODO: Not modeled fully by lean-mlir right now
-/

/-
### Ranked tensors
TODO: Consider a different type KnownRankedTensor that we could project to if
TODO| we have known dimensions. Everything is conditioned by DimList.known...
-/

@[simp]
def shape_refines: List Nat → List Dimension → Bool
  | [], [] => true
  | size::shape, Dimension.Unknown::dim => shape_refines shape dim
  | size::shape, Dimension.Known d::dim => size = d && shape_refines shape dim
  | (_::_), [] => false
  | [], (_::_) => false

@[inline]
def DimList := List Dimension

@[simp]
def DimList.prod: DimList → Nat
  | [] => 1
  | Dimension.Known n :: D => n * prod D
  | Dimension.Unknown :: _ => 0

@[simp]
def DimList.project: DimList → List Nat
  | [] => []
  | Dimension.Known n :: D => n :: project D
  | Dimension.Unknown :: D => project D

@[simp]
def DimList.known: DimList → Bool
  | [] => true
  | Dimension.Known n :: D => known D
  | Dimension.Unknown :: _ => false

@[simp]
def DimList.default_refinement: DimList → List Nat
  | [] => []
  | Dimension.Known n :: D => n :: default_refinement D
  | Dimension.Unknown :: D => 0 :: default_refinement D

theorem dim_known_project_refines {D: DimList}:
    D.known → shape_refines D.project D := by
  intros h <;> induction D <;> simp
  case cons head tail ih =>
    cases head <;> simp at *; apply (ih h)

theorem dim_known_refines_inv {D: DimList} {S: List Nat}:
    D.known → shape_refines S D → D = S.map Dimension.Known := by
  intros Hknown; revert S; induction D <;> intros S Hrefines
  case nil =>
    cases S; simp [List.map]; simp at Hrefines
  case cons head tail ih =>
    cases S; simp at Hrefines
    simp [List.map]; cases head <;> simp at *
    rw [Hrefines.1, ←ih Hknown]; apply Hrefines.2

theorem dim_known_project_eq {D: DimList}:
    D.known → shape_refines S D → D.project = S := by
  intros Hknown Hrefines
  rw [dim_known_refines_inv Hknown Hrefines]
  clear D Hknown Hrefines
  induction S <;> simp; assumption

theorem dim_known_prod_refines {D: DimList}:
    D.known → shape_refines S D → shape_prod S = D.prod := by
  intros Hknown; revert S; induction D <;> intros S Hrefines <;> simp
  case nil =>
    cases S; simp; simp at Hrefines
  case cons head tail ih =>
    cases S; simp at Hrefines
    cases head <;> simp at *
    rw [←Hrefines.1, ←ih Hknown Hrefines.2]
    simp [shape_prod, List.foldr]

theorem dim_known_prod (D: DimList):
    D.known → shape_prod D.project = D.prod :=
  fun Hknown =>
    dim_known_prod_refines Hknown (dim_known_project_refines Hknown)

theorem default_refinement_refines (D: DimList):
    shape_refines D.default_refinement D := by
  induction D <;> simp
  case cons head _ ih =>
    cases head <;> simp <;> apply ih

-- Ranked tensors have known rank but possibly unknown dimensions. At runtime,
-- the size is fully determined. We store the visible type information in the
-- parameters of RankedTensor, and collect the fully-specified runtime data in
-- the structure fields.
--
-- We add an intermediate [size] parameter so that we can separate data
-- manipulation and proofs about the product of the dimensions, which would
-- otherwise be tricky due to [Fin] constantly carrying proofs.
structure RankedTensor (α: Type) (D: DimList) where
  -- Actual dimensions
  shape: List Nat
  -- Contents; we use a function for brevity
  -- TODO: RankedTensor: Consider a more computable data storage method
  data: List α
  -- Invariants: shape/dimension must be compatible, shape/size must match
  h_refines: shape_refines shape D
  h_data_size: data.length = shape_prod shape

theorem RankedTensor.eq_of_fields_eq (α D): ∀ (t₁ t₂: RankedTensor α D),
  t₁.shape = t₂.shape →
  t₁.data = t₂.data →
    t₁ = t₂ := by
  intros t₁ t₂ Hshape Hdata
  cases t₁; cases t₂; simp at *
  trivial

def RankedTensor.uniform {α} D (v: α): RankedTensor α D :=
  { shape       := D.default_refinement,
    data        := List.uniform v (shape_prod D.default_refinement),
    h_refines   := default_refinement_refines _,
    h_data_size := List.uniform_length _ _ }

def RankedTensor.default α D [Inhabited α]: RankedTensor α D :=
  RankedTensor.uniform D Inhabited.default

instance {α D} [Inhabited α]: Inhabited (RankedTensor α D) where
  default := RankedTensor.default α D

/-
### Unranked tensors
TODO: Unranked tensors?
-/

/-
### Evaluation of MLIR types
TODO: Not all MLIRTy types are correctly evaluated
-/

/- MLIRTy is a nested inductive type, thus defined with well-founded recursion.
   This prevents it from being reduced by the elaborator, so instead we define
   it manually with the recursor.
   See: https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/
   reduction.20of.20dependent.20return.20type/near/276044057 -/

@[reducible, simp_itree]
def MLIR.AST.MLIRTy.eval (τ: MLIRTy): Type :=
  @MLIRTy.rec
    (motive_1  := fun _ => Type) -- MemrefLayoutSpec
    (motive_2  := fun _ => Type) -- MLIRTy
    (motive_3  := fun _ => Type) -- AttrVal
    (motive_4  := fun _ => Type) -- AttrEntry
    (motive_5  := fun _ => Type) -- AttrDict
    (motive_6  := fun _ => Type) -- List MLIRTy
    (motive_7  := fun _ => Type) -- Option MemrefLayoutSpec
    (motive_8  := fun _ => Type) -- Option AttrVal
    (motive_9  := fun _ => Type) -- List AttrVal
    (motive_10 := fun _ => Type) -- List AttrEntry

    -- **MemrefLayoutSpec**

    -- MemrefLayoutSpec.stride
    (fun offset stride => Unit)
    -- MemrefLayoutSpec.attr (this is the one which loops everything)
    (fun attr eval_attr => Unit)

    -- **MLIRTy**

    -- MLIRTy.fn
    (fun τ₁ τ₂ eval_τ₁ eval_τ₂ => eval_τ₁ → eval_τ₂)
    -- MLIRTy.int
    (fun bitsize => Int)
    -- MLIRTy.float
    (fun bitsize => Float)
    -- MLIRTy.index
    Nat
    -- MLIRTy.tuple [Mapping motive_2 to motive_1]
    (fun _ ih => ih)
    -- MLIRTy.vector (todo)
    (fun Dfixed Dscaled τ eval_τ => Unit)
    -- MLIRTy.tensorRanked (todo)
    (fun D τ eval_τ => RankedTensor eval_τ D)
    -- MLIRTy.tensorUnranked (todo)
    (fun τ eval_τ => Unit)
    -- MLIRTy.memrefRanked (todo)
    (fun D τ eval_τ layout eval_layout memspace eval_memspace => Unit)
    -- MLIRTy.memrefUnranked (todo)
    (fun τ eval_τ memspace eval_memspace => Unit)
    -- MLIRTy.user (todo)
    (fun name => Unit)

    -- **AttrVal**

    -- AttrVal.symbol
    (fun symbol => Unit)
    -- AttrVal.str
    (fun str => Unit)
    -- AttrVal.int
    (fun int τ eval_τ => Unit)
    -- AttrVal.bool
    (fun bool => Unit)
    -- AttrVal.float
    (fun float τ eval_τ => Unit)
    -- AttrVal.type
    (fun τ eval_τ => Unit)
    -- AttrVal.dense
    (fun el τ eval_τ => Unit)
    -- AttrVal.affine
    (fun affine_map => Unit)
    -- AttrVal.list
    (fun list eval_list => Unit)
    -- AttrVal.nestedsymbol
    (fun a₁ eval_a₁ a₂ eval_a₂ => Unit)
    -- AttrVal.alias
    (fun str => Unit)
    -- AttrVal.dict
    (fun dict eval_dict => Unit)
    -- AttrVal.opaque
    (fun dialect value => Unit)
    -- AttrVal.opaqueElements
    (fun dialect value τ eval_τ => Unit)
    -- AttrVal.unit
    Unit

    -- **AttrEntry**

    -- AttrEntry.mk
    (fun key value eval_value => Unit)

    -- **AttrDict**

    -- AttrDict.mk
    (fun list eval_list => Unit)

    -- **List MLIRTy** (in MLIRTy.tuple)

    -- []
    Unit
    -- (τ::l)
    (fun τ l eval_τ eval_l =>
      match l with
      | [] => eval_τ
      | _  => eval_τ × eval_l)

    -- **Option MemrefLayoutSpec**

    Unit
    (fun layout eval_layout => Unit)

    -- **Option AttrVal**

    Unit
    (fun attr eval_attr => Unit)

    -- **List AttrVal**

    Unit
    (fun head tail eval_head eval_tail => Unit)

    -- **List AttrEntry**

    Unit
    (fun head tail eval_head eval_tail => Unit)

  -- The subject of the induction
  τ

def MLIR.AST.MLIRTy.default (τ: MLIRTy): τ.eval :=
  match τ with
  | MLIRTy.fn τ₁ τ₂ => (fun _ => τ₂.default)
  | MLIRTy.int _ => 0
  | MLIRTy.float _ => 0.0
  | MLIRTy.index => 0
  | MLIRTy.tuple [] => ()
  | MLIRTy.tuple [τ] => τ.default
  | MLIRTy.tuple (τ₁::τ₂::l) => (τ₁.default, (MLIRTy.tuple (τ₂::l)).default)
  | MLIRTy.vector _ _ _ => () /- todo -/
  | MLIRTy.tensorRanked D τ => @RankedTensor.default τ.eval D ⟨default τ⟩
  | MLIRTy.tensorUnranked τ => () /- todo -/
  | MLIRTy.memrefRanked D τ _ _ => () /- todo -/
  | MLIRTy.memrefUnranked τ _ => () /- todo -/
  | MLIRTy.user _ => () /- todo -/

instance (τ: MLIRTy): Inhabited τ.eval where
  default := τ.default

---

-- TODO: RankedTensor.ofTensorElem: account for typing?
-- Which allows us to flatten the term into a single array for a RankedTensor

namespace MLIR.AST.TensorElem

def flatten {τ} (e: TensorElem) (h: e.hasType τ): List τ.eval :=
  match e, τ with
  | TensorElem.int i, MLIRTy.int _ =>
      [i]
  | TensorElem.bool b, MLIRTy.int _ =>
      [if b then 1 else 0]
  | TensorElem.float f, MLIRTy.float _ =>
      [f]
  | TensorElem.nested [], _ =>
      []
  | TensorElem.nested (e::l), τ =>
      let h₁ := (by simp [hasType] at h; apply h.1)
      let h₂ := (by simp [hasType] at h; apply h.2)
      flatten e h₁ ++ flatten (TensorElem.nested l) h₂
  | _, _ =>
      [] -- TODO: Prove impossible

-- Once again, we prove a more friendly version of the list case first

theorem flatten_list {τ} (l: List TensorElem) (h: hasType (.nested l) τ):
    flatten (.nested l) h = (mapWithType l flatten h).join := by
  revert h
  induction l <;> intros h
  case nil =>
    simp [flatten, mapWithType, List.join]
  case cons _ _ ih =>
    simp [flatten, mapWithType, List.join, ih]

theorem flatten_size {τ} (e: TensorElem) (shape: List Nat):
    e.hasShape shape → (h: e.hasType τ) → (e.flatten h).length = shape_prod shape := by
  revert shape
  apply @TensorElem.recOn
    (motive_1 := fun e =>
      ∀s, e.hasShape s → (h: e.hasType τ) → (e.flatten h).length = shape_prod s)
    (motive_2 := fun l =>
      ∀s, l.all (TensorElem.hasShape . s) → (h: l.all (hasType . τ)) →
        (mapWithType l flatten (hasType_list_2 h)).join.length = l.length * shape_prod s)
    <;> simp <;> clear e
  case int =>
    intros i s Hshape Htype;
    cases τ <;> simp [hasType] at Htype
    cases s <;> simp [flatten, hasShape] at *
  case float =>
    intros i s Hshape Htype;
    cases τ <;> simp [hasType] at Htype
    cases s <;> simp [flatten, hasShape] at *
  case bool =>
    intros i s Hshape Htype;
    cases τ <;> simp [hasType] at Htype
    cases s <;> simp [flatten, hasShape] at *
  case nested =>
    intros l motive_2 s Hshape Htype
    cases s <;> simp [hasShape] at Hshape
    case cons s_head s_tail =>
    simp [TensorElem.flatten_list, shape_prod, List.foldr]
    simp [motive_2 s_tail Hshape.2 (hasType_list_1 Htype)]
    simp [shape_prod, Nat.mul_comm, Hshape.1]
  case empty =>
    intros s Hshape Htype
    simp [hasType] at Htype
  case nil =>
    intros _ Htype
    simp [mapWithType, List.join]
  case cons =>
    intros head tail motive_1 IH2 s Hshape Htype
    simp [List.map, List.join]
    rw [Nat.add_comm]
    simp [Nat.succ_eq_add_one, Nat.right_distrib]
    simp [List.all_cons] at Hshape
    simp [List.all_cons] at Htype
    simp [IH2 s Hshape.2 Htype.2]
    rw [motive_1 s Hshape.1 Htype.1]

inductive rankCompatibleWith (e: TensorElem) (D: DimList): MLIRTy → Type :=
  | UniformInt (i: Int) bitsize:
      -- TODO: Check range of uniform tensor value
      e = TensorElem.int i →
      e.rankCompatibleWith D (MLIRTy.int bitsize)
  | UniformBool (b: Bool):
      e = TensorElem.bool b →
      e.rankCompatibleWith D (MLIRTy.int 1)
  | UniformFloat (f: Float) bitsize:
      -- TODO: Check range of uniform tensor value
      e = TensorElem.float f →
      e.rankCompatibleWith D (MLIRTy.float bitsize)
  | HasShape s τ:
      e.hasShape s →
      shape_refines s D →
      e.rankCompatibleWith D τ

end MLIR.AST.TensorElem

def RankedTensor.ofTensorElem {τ} (D: DimList) (e: TensorElem)
    (Htype: e.hasType τ) (Hcompat: e.rankCompatibleWith D τ):
    RankedTensor τ.eval D:=
  match Hcompat with
  | .UniformInt i bitsize _ =>
      RankedTensor.uniform D i
  | .UniformBool b _ =>
      RankedTensor.uniform D (if b then 1 else 0)
  | .UniformFloat f bitsize _ =>
      RankedTensor.uniform D f
  | .HasShape s τ Hshape Hrefines =>
      { shape       := s,
        data        := e.flatten Htype,
        h_refines   := Hrefines,
        h_data_size := TensorElem.flatten_size e s Hshape Htype }


import ReactorModel.Objects.Reactor.Theorems.Equivalent

namespace ReactorType

theorem UniqueIDs.updated [ReactorType.WellFounded α] {rtr₁ rtr₂ : α}
    (u : LawfulUpdate cpt i f rtr₁ rtr₂) (h : UniqueIDs rtr₁) : UniqueIDs rtr₂ where
  allEq m₁ m₂ := open Member in
    h.allEq (.fromLawfulUpdate m₁ u) (.fromLawfulUpdate m₂ u) ▸ Equivalent.from_lawfulUpdate u m₁
      |>.trans (Equivalent.from_lawfulUpdate u m₂).symm
      |>.to_eq

end ReactorType


import ..phys.time.time
import ..phys.time_series.geom3d
import ..std.time_std
import ..std.geom3d_std
noncomputable theory

-- TODO: Should come from resp. std libraries and be distributed to them accordingly
namespace std
def time (p : K) : time time_std_space := mk_time time_std_space p
def duration (d : K) : duration time_std_space := mk_duration _ d
def position (x y z : K) : position3d geom3d_std_space := mk_position3d _ x y z
def displacement (x y z : K) : displacement3d geom3d_std_space := mk_displacement3d _ x y z
end std

/-
We need to assume a physical interpretation of the data
representing our coordinate system on time. We derive a new ACS from
"coordinated_universal_time_in_seconds" - see time_std.lean
for a more details on the coordinate system and physical interpretation
(note : this is a conventional UTC ACS expressed with units in seconds)

This example does not require any specialized timespaces beyond a simple
world time space.

(1) ORIGIN: We move the origin up to 1629311979

(2) BASIS VECTORS
    basis0
      - points to the future
      - unit length is 1 second (as in UTC)
(3) ACS is given by [Origin, b0]
-/

namespace utc
def origin := std.time 0   -- origin; first instant of January 1, 1970
def basis := std.duration 1    -- basis; "second;" the smallest non-variable unit in UTC
def frame := mk_time_frame origin basis -- recall why "time" is part of the constructor name? factor out?
def coords := mk_space frame  -- "cosys"?
def time (t : K) := mk_time coords t
def duration (d : K) := mk_duration coords d
end utc

/-
We need to assume a physical interpretation of the data
representing our coordinate system on geom3d. See geom3d_std.lean
for more details on the coordinate system and physical interpretation.
-/
-- Geometric world
namespace world  -- it's generic/parametric: for example, world -> Rice 440, as follows
def origin := std.position 0 0 0      -- looking in from doorway, the back lower left corner
def basis_0 := std.displacement 1 0 0 -- right/east along wall; unit is 1m; right
def basis_1 := std.displacement 0 1 0 -- to door along weset wall; 1m; right
def basis_2 := std.displacement 0 0 1 -- up along NW corner; 1m; right handed
def frame := mk_geom3d_frame origin basis_0 basis_1 basis_2
def coords := mk_geom3d_space frame
def position (x y z : K) := mk_position3d coords x y z
def displacement (x y z : K) := mk_displacement3d coords x y z
end world

/-
We define a base link ACS, which is intended to represent the ACS
centered at the robot, in terms of the world frame, and oriented at the robot,
in terms of the world frame.

We define the robot to be 3 meters to the right of the left wall, 4 meters
in north of the bottom wall, and 1 meter above the ground. It's rotated 1 radian
around it's yaw axis relative to the world's orientation.
-/

namespace base_link
def origin := std.position 3 4 1
def basis_0 := std.displacement 1 0 0
def basis_1 := std.displacement 0 1 0
def basis_2 := std.displacement 0 0 1
def frame := mk_geom3d_frame origin basis_0 basis_1 basis_2
def coords := mk_geom3d_space frame
def position (x y z : scalar) := mk_position3d coords x y z
def displacement (x y z : scalar) := mk_displacement3d coords x y z
end base_link

/-
We define a UTM coordinate space. UTM (Universal Transverse Mercator)
is a 2-D coordinate system on the globe that separates the globe into
slices separated into longitudinal strips.

Andrew - update these coordinates with more accurate representations.
-/

namespace UTM  -- it's generic/parametric: for example, world -> Rice 440, as follows
def origin := world.position 3 2 0      -- looking in from doorway, the back lower left corner
def basis_0 := world.displacement 1 0 0 -- right/east along wall; unit is 1m; right
def basis_1 := world.displacement 0 1 0 -- to door along weset wall; 1m; right
def basis_2 := world.displacement 0 0 0 -- up along NW corner; 1m; right handed
def frame := mk_geom3d_frame origin basis_0 basis_1 basis_2
def coords := mk_geom3d_space frame
def position (x y z : K) := mk_position3d coords x y z
def displacement (x y z : K) := mk_displacement3d coords x y z
end UTM

/-
We define a base_link_yaw_only frame. This is intended to represent
the base link frame, however, projected onto a 2-D coordinate sytem
(where there is only a yaw rotation).

Andrew - update these coordinates with more accurate representations.

This should, perhaps, derive from the UTC coordinates.
-/
namespace base_link_yaw_only
def origin := world.position 2 2 0
def basis_0 := world.displacement 0.54 (-0.84) 0
def basis_1 := world.displacement (0.84) (0.54) 0
def basis_2 := world.displacement 0 0 0
def frame := mk_geom3d_frame origin basis_0 basis_1 basis_2
def coords := mk_geom3d_space frame
def position (x y z : scalar) := mk_position3d coords x y z
def displacement (x y z : scalar) := mk_displacement3d coords x y z
end base_link_yaw_only

/-

    //! @brief Parameter that specifies the magnetic declination for the robot's environment.
    //!
    double magnetic_declination_;

    //! @brief UTM's meridian convergence
    //!
    //! Angle between projected meridian (True North) and UTM's grid Y-axis.
    //! For UTM projection (Ellipsoidal Transverse Mercator) it is zero on the equator and non-zero everywhere else.
    //! It increases as the poles are approached or as we're getting farther from central meridian.
    //!
    double utm_meridian_convergence_;

    //! @brief IMU's yaw offset
    //!
    //! Your IMU should read 0 when facing *magnetic* north. If it doesn't, this (parameterized) value gives the offset
    //! (NOTE: if you have a magenetic declination, use the parameter setting for that).
    //!
    double yaw_offset_;

    //! @brief Latest IMU orientation
    //!
    tf2::Quaternion transform_orientation_;

    //! @brief Latest IMU orientation
    //!
    tf2::Transform transform_world_pose_;

    //! @brief Holds the Cartesian->odom transform
    //!
    tf2::Transform cartesian_world_transform_;

    //! @brief Holds the cartesian (UTM or local ENU) pose that is used to compute the transform
    //!
    tf2::Transform transform_cartesian_pose_;

-/

/-
This is a transpiled version of the class containing the critical code, named "NavSatTransform". It contains (only)
the relevant class properties that are referenced in the methods relevant to where the isolated bug is.

-/

structure NavSatTransform :=
 (magnetic_declination_ : scalar)
 (utm_meridian_convergence_ : scalar)
 (yaw_offset_ : scalar)
 (transform_orientation_ : orientation3d world.coords)
 (transform_world_pose_  : geom3d_transform world.coords base_link.coords)
 (transform_cartesian_pose_ : geom3d_transform world.coords UTM.coords)

/-
    void NavSatTransform::getRobotOriginWorldPose(const tf2::Transform &gps_odom_pose,
                                                tf2::Transform &robot_odom_pose,
                                                const ros::Time &transform_time)
  {
    robot_odom_pose.setIdentity();
    tf2::Transform gps_offset_rotated;
    bool can_transform = RosFilterUtilities::lookupTransformSafe(tf_buffer_,
                                                                 base_link_frame_id_,
                                                                 gps_frame_id_,
                                                                 transform_time,
                                                                 transform_timeout_,
                                                                 gps_offset_rotated);

    if (can_transform)
    {
      tf2::Transform robot_orientation;
      can_transform = RosFilterUtilities::lookupTransformSafe(tf_buffer_,
                                                              world_frame_id_,
                                                              base_link_frame_id_,
                                                              transform_time,
                                                              transform_timeout_,
                                                              robot_orientation);
      if (can_transform)
      {
        gps_offset_rotated.setOrigin(tf2::quatRotate(robot_orientation.getRotation(), gps_offset_rotated.getOrigin()));
        gps_offset_rotated.setRotation(tf2::Quaternion::getIdentity());
        robot_odom_pose = gps_offset_rotated.inverse() * gps_odom_pose;
      }
      else
      {
      }
    }
    else
    {
    }
  }

-/

/-
Transpilation of "tf::Transform::setIdentity"
Notice that this method is "bound" to pose3d in Lean, so there is an interesting subtlety in how these would need to be
printed off (possibly with duplicates).

The other aspect to notice is that the purpose of this method, in C++, is an in-place assignment. We lose a particular aspect of the
original semantics here. We can see a re-assignment of the argument, p, below. However, the original semantics require an in-place assignment.
Thus, the caller expects that the member argument, p, has been re-assigned, but, as it is modeled,
the member argument is assigned *into* p0, and the enclosing scope does not have access to the updated value of p.
This is a broader issue and does not have a simple solution, such as forcibly convering this method and all call sites as assignments.
-/
def pose3d.setIdentity {f : geom3d_frame} {sp : geom3d_space f} (p : pose3d sp) : punit :=
  let p0 : pose3d sp := {
    position := inhabited.default _,
    orientation := inhabited.default _
  } in
  punit.star

/-

Just as above, transpilation of "getRPY" call from matrices or quaternions in ROS.
The intended semantics of this method should perform an in-place assignment of the arguments,
(roll pitch and yaw angles). Unfortunately, we face the same challenge as above.

(also, I haven't implemented the Lean orientation call to retrieve these values...will do so soon)
-/
def orientation3d.getRPY {f : geom3d_frame} {sp : geom3d_space f }
  (o : orientation3d sp)
  : scalar → scalar → scalar → punit :=
  λ s1 s2 s3,
  punit.star

/-

    void NavSatTransform::getRobotOriginWorldPose(const tf2::Transform &gps_odom_pose,
                                                tf2::Transform &robot_odom_pose,
                                                const ros::Time &transform_time)

This is the formalization of a helper function reference in the function "computeTransform",
formalized below, which describes the error of interest to the issue. There are no physical
type errors in this method, although formalizing raised some interesting issues.

It is defined as a "Lean-esque member function" of NavSatTransform.

For example, Poses and Transforms share the same type in C++, whereas they are modeled differently
-/
def NavSatTransform.getRobotOriginCartesianPose (nst : NavSatTransform)
  : geom3d_transform world.coords UTM.coords → pose3d world.coords → time utc.coords → punit
  :=
    λ gps_cartesian_pose, λ robot_cartesian_pose, λ transform_time,
    --robot_cartesian_pose.setIdentity();
    let setIdentityCall := robot_cartesian_pose.setIdentity in
    --tf2::Quaternion cartesian_orientation = transform_orientation_;
    let cartesian_orientation : orientation3d world.coords := nst.transform_orientation_ in
    --tf2::Matrix3x3 mat(cartesian_orientation);
    let mat : orientation3d world.coords := cartesian_orientation in

    /-
    double roll;
    double pitch;
    double yaw;
    -/
    let roll : scalar := inhabited.default _ in
    let pitch : scalar := inhabited.default _ in
    let yaw : scalar := inhabited.default _ in

    --mat.getRPY(roll, pitch, yaw);
    let getRPYCall := mat.getRPY roll pitch yaw in
    let yaw0 := yaw + nst.magnetic_declination_ + nst.yaw_offset_ + nst.utm_meridian_convergence_ in
    --yaw += (magnetic_declination_ + yaw_offset_ + utm_meridian_convergence_);
    /-
    cartesian_orientation.setRPY(roll, pitch, yaw);
    A note here: A more "literal" translation of the above line would involve a call to "setRPY",
    which, again, is an in-place assignment to cartesian_orientation. Fortunately, if we don't
    attempt that route, we can directly assign to the variable (SSA-style) using "mk_orientation3d_from_euler_angles",
    which carries the same semantics as the original C++.
    -/
    let cartesian_orientation0 := mk_orientation3d_from_euler_angles world.coords roll pitch yaw in

    /-
    offset.setOrigin(tf2::quatRotate(cartesian_orientation, offset.getOrigin()));
    offset.setRotation(tf2::Quaternion::getIdentity());

    robot_cartesian_pose = offset.inverse() * gps_cartesian_pose;

    We have no way of formalizing this code. "offset" is intended to be a transform - but we provide
    no way to set the "origin" or "rotation" of a Transform, and, it's not a simple fix, it's a bit beyond the scope of
    our formalization currently.
    -/

    punit.star

/-
void NavSatTransform::computeTransform()

This is the principal function in which the error of interest is modeled.
It is a member function of the C++ class "NavSatTransform", and here, in Lean, it is modeled
as a "member function" of the structure NavSatTransform.

The gist of the function and error, below, is that we will first "construct"
  the Transform from UTM → Base Link Yaw Only, T₁, and then we will "use" the member transform_world_pose_, T₂,
  of NavSatTransform, which is a transform from World → Base Link, in order to construct a transform from
  World → UTM, by taking T₂∘(T₁⁻¹). However, the domain of T₁⁻¹ does not match the codomain of T₂,
  yielding a type error in Lean that was not captured in the original code.
-/
def NavSatTransform.compute_transform
  (nst : NavSatTransform)
  : punit :=
  /-
      tf2::Transform transform_cartesian_pose_corrected;
  -/
  let transform_cartesian_pose_corrected : pose3d world.coords := inhabited.default _ in
  /-
  if (!use_manual_datum_)
      {
        getRobotOriginCartesianPose(transform_cartesian_pose_, transform_cartesian_pose_corrected, ros::Time(0));
      }
      else
      {
        transform_cartesian_pose_corrected = transform_cartesian_pose_;
      }

  Notice here, we simply assume an execution path. We assume that the former branch triggers, and thus, a call
  to the previously formalized member function of NavSatTransform is called
  -/
  let getRobotOriginCartesianPoseCall : punit
    := nst.getRobotOriginCartesianPose nst.transform_cartesian_pose_ transform_cartesian_pose_corrected (mk_time utc.coords 0) in
  /-
      tf2::Matrix3x3 mat(transform_orientation_);
  -/
  let mat : orientation3d world.coords := inhabited.default _ in

  /-
      double imu_roll;
      double imu_pitch;
      double imu_yaw;
      mat.getRPY(imu_roll, imu_pitch, imu_yaw);
  -/
  let imu_roll : scalar := inhabited.default _ in
  let imu_pitch : scalar := inhabited.default _ in
  let imu_yaw : scalar := inhabited.default _ in
  /-
  This does not correctly model the semantics of in-place assignment, as described earlier.
  -/
  let getRPYCall : punit := mat.getRPY imu_roll imu_pitch imu_yaw in

  --imu_yaw += (magnetic_declination_ + yaw_offset_ + utm_meridian_convergence_);
  /-

  -/
  let imu_yaw0 := imu_yaw + nst.magnetic_declination_ + nst.yaw_offset_ + nst.utm_meridian_convergence_ in

  /-
  tf2::Quaternion imu_quat;
  imu_quat.setRPY(0.0, 0.0, imu_yaw);

  -/
  let imu_quat : orientation3d base_link.coords := inhabited.default _ in
  let imu_quat0 : orientation3d base_link.coords := mk_orientation3d_from_euler_angles _ imu_roll imu_pitch imu_yaw in

/-

      cartesian_pose_with_orientation.setOrigin(transform_cartesian_pose_corrected.getOrigin());
      cartesian_pose_with_orientation.setRotation(imu_quat);

      We have no way to formalize these two lines. As described earlier, in our current architecture, a
      Transform is defined simply by providing two coordinate spaces, and we receive the *fixed* transform
      between those two spaces. In ROS, coordinate spaces *vary* over time, and thus, transforms are not necessarily
      fixed, and we can assign to their properties in an ad hoc manner, which often represents a transform at a particular time.
      Time Series, at least as we've currently modeled them, do not provide a solution to this limitation.

      To capture the semantics, I ignore the constructed transform_cartesian_pose_corrected origin and imu_quat, and
      construct the transform from UTM → Base Link Yaw Only as we conventionally construct it.
-/
  let cartesian_pose_with_orientation : geom3d_transform _ _ :=
    (UTM.coords.mk_geom3d_transform_to base_link_yaw_only.coords) in

  /-
  Here is where the error occurs: We attempt to compose a transform from World → Base Link, with the inverse of a transform from
  UTM → Base Link Yaw Only, by taking T₂∘(T₁⁻¹). However, the domain of T₁⁻¹ does not match the codomain of T₂,
  yielding a type error in Lean that was not captured in the original code.
  -/
  let cartestian_world_transform_ := (nst.transform_world_pose_.trans (cartesian_pose_with_orientation.symm)) in

  punit.star


import «smt-lean»

example {x y : ℤ} (h1 : ((x - y) = (x + (- y) + 1)))
 : false :=
begin
  veriT,
end


example (p q r : Prop) (hp : p) (hq : q) (hr : r) : p ∧ q ∧ r :=
begin
  split,
  any_goals { split },
  any_goals { assumption }
end


example (p q r : Prop) (hp : p) (hq : q) (hr : r) : p ∧ q ∧ r :=
begin
  repeat { any_goals { split } },
  all_goals { assumption }
end


/-
Copyright (c) 2018 Simon Hudon All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Simon Hudon

! This file was ported from Lean 3 source module tactic.scc
! leanprover-community/mathlib commit d6814c584384ddf2825ff038e868451a7c956f31
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
import Mathbin.Tactic.Tauto

/-!
# Strongly Connected Components

This file defines tactics to construct proofs of equivalences between a set of mutually equivalent
propositions. The tactics use implications transitively to find sets of equivalent propositions.

## Implementation notes

The tactics use a strongly connected components algorithm on a graph where propositions are
vertices and edges are proofs that the source implies the target. The strongly connected components
are therefore sets of propositions that are pairwise equivalent to each other.

The resulting strongly connected components are encoded in a disjoint set data structure to
facilitate the construction of equivalence proofs between two arbitrary members of an equivalence
class.

## Possible generalizations

Instead of reasoning about implications and equivalence, we could generalize the machinery to
reason about arbitrary partial orders.

## References

 * Tarjan, R. E. (1972), "Depth-first search and linear graph algorithms",
   SIAM Journal on Computing, 1 (2): 146–160, doi:10.1137/0201010
 * Dijkstra, Edsger (1976), A Discipline of Programming, NJ: Prentice Hall, Ch. 25.
 * <https://en.wikipedia.org/wiki/Disjoint-set_data_structure>

## Tags

graphs, tactic, strongly connected components, disjoint sets
-/

namespace Tactic

/-- `closure` implements a disjoint set data structure using path compression
optimization. For the sake of the scc algorithm, it also stores the preorder
numbering of the equivalence graph of the local assumptions.

The `expr_map` encodes a directed forest by storing for every non-root
node, a reference to its parent and a proof of equivalence between
that node's expression and its parent's expression. Given that data
structure, checking that two nodes belong to the same tree is easy and
fast by repeatedly following the parent references until a root is reached.
If both nodes have the same root, they belong to the same tree, i.e. their
expressions are equivalent. The proof of equivalence can be formed by
composing the proofs along the edges of the paths to the root.

More concretely, if we ignore preorder numbering, the set
`{ {e₀,e₁,e₂,e₃}, {e₄,e₅} }` is represented as:

```
e₀ → ⊥      -- no parent, i.e. e₀ is a root
e₁ → e₀, p₁ -- with p₁ : e₁ ↔ e₀
e₂ → e₁, p₂ -- with p₂ : e₂ ↔ e₁
e₃ → e₀, p₃ -- with p₃ : e₃ ↔ e₀
e₄ → ⊥      -- no parent, i.e. e₄ is a root
e₅ → e₄, p₅ -- with p₅ : e₅ ↔ e₄
```

We can check that `e₂` and `e₃` are equivalent by seeking the root of
the tree of each. The parent of `e₂` is `e₁`, the parent of `e₁` is
`e₀` and `e₀` does not have a parent, and thus, this is the root of its tree.
The parent of `e₃` is `e₀` and it's also the root, the same as for `e₂` and
they are therefore equivalent. We can build a proof of that equivalence by using
transitivity on `p₂`, `p₁` and `p₃.symm` in that order.

Similarly, we can discover that `e₂` and `e₅` aren't equivalent.

A description of the path compression optimization can be found at:
<https://en.wikipedia.org/wiki/Disjoint-set_data_structure#Path_compression>

-/
unsafe def closure :=
  ref (expr_map (Sum ℕ (expr × expr)))
#align tactic.closure tactic.closure

namespace closure

/-- `with_new_closure f` creates an empty `closure` `c`, executes `f` on `c`, and then deletes `c`,
returning the output of `f`. -/
unsafe def with_new_closure {α} : (closure → tactic α) → tactic α :=
  using_new_ref (expr_map.mk _)
#align tactic.closure.with_new_closure tactic.closure.with_new_closure

/-- `to_tactic_format cl` pretty-prints the `closure` `cl` as a list. Assuming `cl` was built by
`dfs_at`, each element corresponds to a node `pᵢ : expr` and is one of the folllowing:
- if `pᵢ` is a root: `"pᵢ ⇐ i"`, where `i` is the preorder number of `pᵢ`,
- otherwise: `"(pᵢ, pⱼ) : P"`, where `P` is `pᵢ ↔ pⱼ`.
Useful for debugging. -/
unsafe def to_tactic_format (cl : closure) : tactic format := do
  let m ← read_ref cl
  let l := m.toList
  let fmt ←
    l.mapM fun ⟨x, y⟩ =>
        match y with
        | Sum.inl y => f!"{(← x)} ⇐ {← y}"
        | Sum.inr ⟨y, p⟩ => f!"({(← x)}, {(← y)}) : {← infer_type p}"
  pure <| to_fmt fmt
#align tactic.closure.to_tactic_format tactic.closure.to_tactic_format

unsafe instance : has_to_tactic_format closure :=
  ⟨to_tactic_format⟩

/-- `(n,r,p) ← root cl e` returns `r` the root of the tree that `e` is a part of (which might be
itself) along with `p` a proof of `e ↔ r` and `n`, the preorder numbering of the root. -/
unsafe def root (cl : closure) : expr → tactic (ℕ × expr × expr)
  | e => do
    let m ← read_ref cl
    match m e with
      | none => do
        let p ← mk_app `` Iff.refl [e]
        pure (0, e, p)
      | some (Sum.inl n) => do
        let p ← mk_app `` Iff.refl [e]
        pure (n, e, p)
      | some (Sum.inr (e₀, p₀)) => do
        let (n, e₁, p₁) ← root e₀
        let p ← mk_app `` Iff.trans [p₀, p₁]
        modify_ref cl fun m => m e (Sum.inr (e₁, p))
        pure (n, e₁, p)
#align tactic.closure.root tactic.closure.root

/-- (Implementation of `merge`.) -/
unsafe def merge_intl (cl : closure) (p e₀ p₀ e₁ p₁ : expr) : tactic Unit := do
  let p₂ ← mk_app `` Iff.symm [p₀]
  let p ← mk_app `` Iff.trans [p₂, p]
  let p ← mk_app `` Iff.trans [p, p₁]
  modify_ref cl fun m => m e₀ <| Sum.inr (e₁, p)
#align tactic.closure.merge_intl tactic.closure.merge_intl

-- failed to format: unknown constant 'term.pseudo.antiquot'
/--
      `merge cl p`, with `p` a proof of `e₀ ↔ e₁` for some `e₀` and `e₁`,
      merges the trees of `e₀` and `e₁` and keeps the root with the smallest preorder
      number as the root. This ensures that, in the depth-first traversal of the graph,
      when encountering an edge going into a vertex whose equivalence class includes
      a vertex that originated the current search, that vertex will be the root of
      the corresponding tree. -/
    unsafe
  def
    merge
    ( cl : closure ) ( p : expr ) : tactic Unit
    :=
      do
        let q( $ ( e₀ ) ↔ $ ( e₁ ) ) ← infer_type p >>= instantiate_mvars
          let ( n₂ , e₂ , p₂ ) ← root cl e₀
          let ( n₃ , e₃ , p₃ ) ← root cl e₁
          if
            e₂ ≠ e₃
            then
            do
              if
                n₂ < n₃
                then
                do let p ← mk_app ` ` Iff.symm [ p ] cl p e₃ p₃ e₂ p₂
                else
                cl p e₂ p₂ e₃ p₃
            else
            pure ( )
#align tactic.closure.merge tactic.closure.merge

/-- Sequentially assign numbers to the nodes of the graph as they are being visited. -/
unsafe def assign_preorder (cl : closure) (e : expr) : tactic Unit :=
  modify_ref cl fun m => m.insert e (Sum.inl m.size)
#align tactic.closure.assign_preorder tactic.closure.assign_preorder

/-- `prove_eqv cl e₀ e₁` constructs a proof of equivalence of `e₀` and `e₁` if
they are equivalent. -/
unsafe def prove_eqv (cl : closure) (e₀ e₁ : expr) : tactic expr := do
  let (_, r, p₀) ← root cl e₀
  let (_, r', p₁) ← root cl e₁
  guard (r = r') <|> throwError "{(← e₀)} and {← e₁} are not equivalent"
  let p₁ ← mk_app `` Iff.symm [p₁]
  mk_app `` Iff.trans [p₀, p₁]
#align tactic.closure.prove_eqv tactic.closure.prove_eqv

/-- `prove_impl cl e₀ e₁` constructs a proof of `e₀ -> e₁` if they are equivalent. -/
unsafe def prove_impl (cl : closure) (e₀ e₁ : expr) : tactic expr :=
  cl.prove_eqv e₀ e₁ >>= iff_mp
#align tactic.closure.prove_impl tactic.closure.prove_impl

/-- `is_eqv cl e₀ e₁` checks whether `e₀` and `e₁` are equivalent without building a proof. -/
unsafe def is_eqv (cl : closure) (e₀ e₁ : expr) : tactic Bool := do
  let (_, r, p₀) ← root cl e₀
  let (_, r', p₁) ← root cl e₁
  return <| r = r'
#align tactic.closure.is_eqv tactic.closure.is_eqv

end closure

/-- mutable graphs between local propositions that imply each other with the proof of implication -/
@[reducible]
unsafe def impl_graph :=
  ref (expr_map (List <| expr × expr))
#align tactic.impl_graph tactic.impl_graph

/-- `with_impl_graph f` creates an empty `impl_graph` `g`, executes `f` on `g`, and then deletes
`g`, returning the output of `f`. -/
unsafe def with_impl_graph {α} : (impl_graph → tactic α) → tactic α :=
  using_new_ref (expr_map.mk (List <| expr × expr))
#align tactic.with_impl_graph tactic.with_impl_graph

namespace ImplGraph

-- failed to format: unknown constant 'term.pseudo.antiquot'
/--
      `add_edge g p`, with `p` a proof of `v₀ → v₁` or `v₀ ↔ v₁`, adds an edge to the implication
      graph `g`. -/
    unsafe
  def
    add_edge
    ( g : impl_graph ) : expr → tactic Unit
    |
      p
      =>
      do
        let t ← infer_type p
          match
            t
            with
            |
                q( $ ( v₀ ) → $ ( v₁ ) )
                =>
                do
                  is_prop v₀ >>= guardb
                    is_prop v₁ >>= guardb
                    let m ← read_ref g
                    let xs := ( m v₀ ) . getD [ ]
                    let xs' := ( m v₁ ) . getD [ ]
                    modify_ref g fun m => ( m v₀ ( ( v₁ , p ) :: xs ) ) . insert v₁ xs'
              |
                q( $ ( v₀ ) ↔ $ ( v₁ ) )
                =>
                do
                  let p₀ ← mk_mapp ` ` Iff.mp [ none , none , p ]
                    let p₁ ← mk_mapp ` ` Iff.mpr [ none , none , p ]
                    add_edge p₀
                    add_edge p₁
              | _ => failed
#align tactic.impl_graph.add_edge tactic.impl_graph.add_edge

section Scc

open List

parameter (g : expr_map (List <| expr × expr))

parameter (visit : ref <| expr_map Bool)

parameter (cl : closure)

/-- `merge_path path e`, where `path` and `e` forms a cycle with proofs of implication between
consecutive vertices. The proofs are compiled into proofs of equivalences and added to the closure
structure. `e` and the first vertex of `path` do not have to be the same but they have to be
in the same equivalence class. -/
unsafe def merge_path (path : List (expr × expr)) (e : expr) : tactic Unit := do
  let p₁ ← cl.prove_impl e Path.headI.fst
  let p₂ ← mk_mapp `` id [e]
  let path := (e, p₁) :: Path
  let (_, ls) ←
    Path.mapAccumLM
        (fun p p' => Prod.mk <$> mk_mapp `` Implies.trans [none, p'.1, none, p, p'.2] <*> pure p) p₂
  let (_, rs) ←
    Path.mapAccumRM
        (fun p p' => Prod.mk <$> mk_mapp `` Implies.trans [none, none, none, p.2, p'] <*> pure p')
        p₂
  let ps ← zipWithM (fun p₀ p₁ => mk_app `` Iff.intro [p₀, p₁]) ls.tail rs.dropLast
  ps cl
#align tactic.impl_graph.merge_path tactic.impl_graph.merge_path

/-- (implementation of `collapse`) -/
unsafe def collapse' : List (expr × expr) → List (expr × expr) → expr → tactic Unit
  | Acc, [], v => merge_path Acc v
  | Acc, (x, pr) :: xs, v => do
    let b ← cl.is_eqv x v
    let acc' := (x, pr) :: Acc
    if b then merge_path acc' v else collapse' acc' xs v
#align tactic.impl_graph.collapse' tactic.impl_graph.collapse'

/-- `collapse path v`, where `v` is a vertex that originated the current search
(or a vertex in the same equivalence class as the one that originated the current search).
It or its equivalent should be found in `path`. Since the vertices following `v` in the path
form a cycle with `v`, they can all be added to an equivalence class. -/
unsafe def collapse : List (expr × expr) → expr → tactic Unit :=
  collapse' []
#align tactic.impl_graph.collapse tactic.impl_graph.collapse

/-- Strongly connected component algorithm inspired by Tarjan's and
Dijkstra's scc algorithm. Whereas they return strongly connected
components by enumerating them, this algorithm returns a disjoint set
data structure using path compression. This is a compact
representation that allows us, after the fact, to construct a proof of
equivalence between any two members of an equivalence class.

 * Tarjan, R. E. (1972), "Depth-first search and linear graph algorithms",
   SIAM Journal on Computing, 1 (2): 146–160, doi:10.1137/0201010
 * Dijkstra, Edsger (1976), A Discipline of Programming, NJ: Prentice Hall, Ch. 25.
-/
unsafe def dfs_at : List (expr × expr) → expr → tactic Unit
  | vs, v => do
    let m ← read_ref visit
    let (_, v', _) ← cl.root v
    match m v' with
      | some tt => pure ()
      | some ff => collapse vs v
      | none => do
        cl v
        modify_ref visit fun m => m v ff
        let ns ← g v
        ns fun ⟨w, e⟩ => dfs_at ((v, e) :: vs) w
        modify_ref visit fun m => m v tt
        pure ()
#align tactic.impl_graph.dfs_at tactic.impl_graph.dfs_at

end Scc

/-- Use the local assumptions to create a set of equivalence classes. -/
unsafe def mk_scc (cl : closure) : tactic (expr_map (List (expr × expr))) :=
  with_impl_graph fun g =>
    using_new_ref (expr_map.mk Bool) fun visit => do
      let ls ← local_context
      ls fun l => try (g l)
      let m ← read_ref g
      m fun ⟨v, _⟩ => impl_graph.dfs_at m visit cl [] v
      pure m
#align tactic.impl_graph.mk_scc tactic.impl_graph.mk_scc

end ImplGraph

-- failed to format: unknown constant 'term.pseudo.antiquot'
unsafe
  def
    prove_eqv_target
    ( cl : closure ) : tactic Unit
    := do let q( $ ( p ) ↔ $ ( q ) ) ← target >>= whnf cl p q >>= exact
#align tactic.prove_eqv_target tactic.prove_eqv_target

-- failed to format: unknown constant 'term.pseudo.antiquot'
/--
      `scc` uses the available equivalences and implications to prove
      a goal of the form `p ↔ q`.

      ```lean
      example (p q r : Prop) (hpq : p → q) (hqr : q ↔ r) (hrp : r → p) : p ↔ r :=
      by scc
      ```
      -/
    unsafe
  def
    interactive.scc
    : tactic Unit
    :=
      closure.with_new_closure
        fun cl => do impl_graph.mk_scc cl let q( $ ( p ) ↔ $ ( q ) ) ← target cl p q >>= exact
#align tactic.interactive.scc tactic.interactive.scc

/-- Collect all the available equivalences and implications and
add assumptions for every equivalence that can be proven using the
strongly connected components technique. Mostly useful for testing. -/
unsafe def interactive.scc' : tactic Unit :=
  closure.with_new_closure fun cl => do
    let m ← impl_graph.mk_scc cl
    let ls := m.toList.map Prod.fst
    let ls' := Prod.mk <$> ls <*> ls
    ls' fun x => do
        let h ← get_unused_name `h
        try <| closure.prove_eqv cl x.1 x.2 >>= note h none
#align tactic.interactive.scc' tactic.interactive.scc'

/-- `scc` uses the available equivalences and implications to prove
a goal of the form `p ↔ q`.

```lean
example (p q r : Prop) (hpq : p → q) (hqr : q ↔ r) (hrp : r → p) : p ↔ r :=
by scc
```

The variant `scc'` populates the local context with all equivalences that `scc` is able to prove.
This is mostly useful for testing purposes.
-/
add_tactic_doc
  { Name := "scc"
    category := DocCategory.tactic
    declNames := [`` interactive.scc, `` interactive.scc']
    tags := ["logic"] }

end Tactic


import data.matrix.notation

import for_mathlib.snake_lemma2
import for_mathlib.short_exact_sequence

noncomputable theory

open category_theory
open category_theory.limits

universes v u

lemma preadditive.exact_of_iso_of_exact' {D : Type*} [category D] [abelian D]
  {A₁ B₁ C₁ A₂ B₂ C₂ : D}
  (f₁ : A₁ ⟶ B₁) (g₁ : B₁ ⟶ C₁) (f₂ : A₂ ⟶ B₂) (g₂ : B₂ ⟶ C₂)
  (α : A₁ ≅ A₂) (β : B₁ ≅ B₂) (γ : C₁ ≅ C₂) (hsq₁ : α.hom ≫ f₂ = f₁ ≫ β.hom)
  (hsq₂ : β.hom ≫ g₂ = g₁ ≫ γ.hom)
  (h : exact f₁ g₁) :
  exact f₂ g₂ :=
preadditive.exact_of_iso_of_exact f₁ g₁ f₂ g₂ (arrow.iso_mk α β hsq₁) (arrow.iso_mk β γ hsq₂) rfl h

namespace homological_complex

variables {C : Type u} [category.{v} C] [abelian C]
variables {ι : Type*} {c : complex_shape ι}

def mod_boundaries (A : homological_complex C c) (j : ι) : C :=
cokernel ((A.boundaries j).arrow)

def mod_boundaries_map {A B : homological_complex C c} (f : A ⟶ B) (j : ι) :
  A.mod_boundaries j ⟶ B.mod_boundaries j :=
cokernel.map _ _ (boundaries_map f j) (f.f j) $ by { rw image_subobject_map_arrow, refl }

@[simps]
def mod_boundaries_functor (j : ι) : homological_complex C c ⥤ C :=
{ obj := λ A, A.mod_boundaries j,
  map := λ A B f, mod_boundaries_map f j,
  map_id' := λ A,
  begin
    delta mod_boundaries mod_boundaries_map cokernel.map, ext,
    show cokernel.π (A.boundaries j).arrow ≫ _ = cokernel.π (A.boundaries j).arrow ≫ _,
    simp only [cokernel.π_desc, category.id_comp, id_f, category.comp_id],
  end,
  map_comp' := λ X Y Z f g,
  begin
    delta mod_boundaries mod_boundaries_map cokernel.map, ext,
    show cokernel.π (X.boundaries j).arrow ≫ _ = cokernel.π (X.boundaries j).arrow ≫ _,
    simp only [cokernel.π_desc, cokernel.π_desc_assoc, comp_f, category.assoc],
  end }
.

-- generalize to chain complexes over other shapes
@[simps]
def homology_to_mod_boundaries (i : ι) :
  homology_functor C c i ⟶ mod_boundaries_functor i :=
{ app := λ A, cokernel.map _ _ (𝟙 _) ((A.cycles i).arrow)
    (by { simp only [category.id_comp, image_to_kernel_arrow], }),
  naturality' := λ A B f,
  begin
    ext,
    simp only [homology_functor_map, mod_boundaries_functor_map, homology.π_map_assoc],
    delta mod_boundaries_map homology.π cokernel.map cycles,
    simp only [cokernel.π_desc, cokernel.π_desc_assoc, comp_f, category.assoc,
      kernel_subobject_map_arrow_assoc, hom.sq_from_left],
  end }
.

variables (A : homological_complex C c) (i j : ι) (hij : c.rel i j)

def delta_to_boundaries : A.X i ⟶ (A.boundaries j) :=
(X_prev_iso A hij).inv ≫ factor_thru_image_subobject _

instance delta_to_boundaries_epi : epi (delta_to_boundaries A i j hij) :=
epi_comp _ _

@[ext] lemma boundaries.ext' {X : C} (f g : (boundaries A j : C) ⟶ X)
  (h : factor_thru_image_subobject _ ≫ f = factor_thru_image_subobject _ ≫ g) : f = g :=
by rwa cancel_epi (factor_thru_image_subobject (A.d_to j)) at h

@[simp, reassoc] lemma delta_to_boundaries_comp_arrow :
  (delta_to_boundaries A i j hij) ≫ (boundaries A j).arrow = A.d i j :=
by rw [delta_to_boundaries, category.assoc, image_subobject_arrow_comp, X_prev_iso_comp_d_to]

@[simp, reassoc] lemma boundaries_arrow_comp_delta_to_boundaries :
  (boundaries _ i).arrow ≫ delta_to_boundaries A i j hij = 0 :=
begin
  ext,
  simp only [image_subobject_arrow_comp_assoc, category.assoc,
    delta_to_boundaries_comp_arrow, comp_zero, zero_comp,
    ← d_from_comp_X_next_iso A hij, reassoc_of (d_to_comp_d_from A)],
end

def delta_to_cycles : A.X i ⟶ (A.cycles j) :=
delta_to_boundaries _ i j hij ≫ boundaries_to_cycles _ _

@[simp, reassoc] lemma delta_to_cycles_comp_arrow :
  (delta_to_cycles A i j hij) ≫ (cycles A j).arrow = A.d i j :=
by rw [delta_to_cycles, category.assoc, image_to_kernel_arrow, delta_to_boundaries_comp_arrow]

@[simp, reassoc] lemma boundaries_arrow_comp_delta_to_cycles :
  (boundaries _ _).arrow ≫ delta_to_cycles A i j hij = 0 :=
by rw [delta_to_cycles, ← category.assoc, boundaries_arrow_comp_delta_to_boundaries, zero_comp]

@[simps]
def mod_boundaries_to_cycles : mod_boundaries_functor i ⟶ cycles_functor C c j :=
{ app := λ A, cokernel.desc _ (delta_to_cycles _ i j hij)
   (boundaries_arrow_comp_delta_to_cycles _ i j hij),
  naturality' := λ A B f,
  begin
    ext, show cokernel.π _ ≫ _ = cokernel.π _ ≫ _,
    simp only [homology_functor_map, mod_boundaries_functor_map, homology.π_map_assoc],
    delta mod_boundaries_map homology.π cokernel.map,
    simp only [category.assoc, cycles_functor_map, cycles_map_arrow, hom.comm,
      cokernel.π_desc_assoc, delta_to_cycles_comp_arrow_assoc, delta_to_cycles_comp_arrow]
  end }
.

@[simps]
def cycles_to_homology : cycles_functor C c i ⟶ homology_functor C c i :=
{ app := λ A, cokernel.π _,
  naturality' := λ A B f,
  begin
    simp only [cycles_functor_map, homology_functor_map],
    delta homology.map,
    rw cokernel.π_desc, refl,
  end }

open_locale zero_object

lemma _root_.option.eq_none_or_eq_some {α : Type*} : ∀ (o : option α), o = none ∨ ∃ a, o = some a
| option.none     := or.inl rfl
| (option.some a) := or.inr ⟨a, rfl⟩

lemma exact_next {A₁ A₂ A₃ : homological_complex C c} (f : A₁ ⟶ A₂) (g : A₂ ⟶ A₃)
  (i j : ι) (hij : c.rel i j) (h : exact (f.f j) (g.f j)) :
  exact (f.next i) (g.next i) :=
begin
  refine preadditive.exact_of_iso_of_exact' (f.f j) (g.f j) _ _
    (X_next_iso A₁ hij).symm (X_next_iso A₂ hij).symm (X_next_iso A₃ hij).symm _ _ h;
  simp only [hom.next_eq _ hij, iso.symm_hom, iso.inv_hom_id_assoc],
end

lemma exact_next' {A₁ A₂ A₃ : homological_complex C c} (f : A₁ ⟶ A₂) (g : A₂ ⟶ A₃) (i : ι)
  (h : ∀ n, exact (f.f n) (g.f n)) : exact (f.next i) (g.next i) :=
h _

lemma exact_prev {A₁ A₂ A₃ : homological_complex C c} (f : A₁ ⟶ A₂) (g : A₂ ⟶ A₃)
  (i j : ι) (hij : c.rel i j) (h : exact (f.f i) (g.f i)) :
  exact (f.prev j) (g.prev j) :=
begin
  refine preadditive.exact_of_iso_of_exact' (f.f i) (g.f i) _ _
    (X_prev_iso A₁ hij).symm (X_prev_iso A₂ hij).symm (X_prev_iso A₃ hij).symm _ _ h;
  simp only [hom.prev_eq _ hij, iso.symm_hom, iso.inv_hom_id_assoc],
end

lemma exact_prev' {A₁ A₂ A₃ : homological_complex C c} (f : A₁ ⟶ A₂) (g : A₂ ⟶ A₃) (j : ι)
  (h : ∀ n, exact (f.f n) (g.f n)) : exact (f.prev j) (g.prev j) :=
h _

lemma mono_next {A₁ A₂ : homological_complex C c} (f : A₁ ⟶ A₂)
  (i j : ι) (hij : c.rel i j) [mono (f.f j)] :
  mono (f.next i) :=
begin
  rw hom.next_eq _ hij,
  apply_with mono_comp { instances := ff },
  { apply_instance },
  { apply mono_comp }
end

instance mono_next' {A₁ A₂ : homological_complex C c} (f : A₁ ⟶ A₂)
  (i : ι) [∀ n, mono (f.f n)] :
  mono (f.next i) :=
by apply_assumption

lemma epi_prev {A₁ A₂ : homological_complex C c} (f : A₁ ⟶ A₂)
  (i j : ι) (hij : c.rel i j) [epi (f.f i)] :
  epi (f.prev j) :=
begin
  rw hom.prev_eq _ hij,
  apply_with epi_comp { instances := ff },
  { apply_instance },
  { apply epi_comp }
end

instance epi_prev' {A₁ A₂ : homological_complex C c} (f : A₁ ⟶ A₂)
  (j : ι) [∀ n, epi (f.f n)] :
  epi (f.prev j) :=
by apply_assumption

instance {A B : homological_complex C c} (f : A ⟶ B) [∀ n, epi (f.f n)] (i : ι) :
  epi (boundaries_map f i) :=
begin
  let sq := hom.sq_to f i,
  haveI : epi sq.left := by { dsimp, apply_instance, },
  apply_with (epi_of_epi (factor_thru_image_subobject _)) { instances := ff },
  suffices : factor_thru_image_subobject (A.d_to i) ≫
      boundaries_map f i =
    sq.left ≫ factor_thru_image_subobject (B.d_to i),
  { rw this, apply epi_comp, },
  ext,
  simp only [category.assoc, image_subobject_map_arrow, hom.sq_to_right,
    image_subobject_arrow_comp_assoc, hom.sq_to_left, image_subobject_arrow_comp, hom.comm_to],
end

lemma exact_kernel_subobject_arrow (A B : C) (f : A ⟶ B) : exact (kernel_subobject f).arrow f :=
by { rw [← kernel_subobject_arrow, exact_iso_comp], exact exact_kernel_ι }

lemma exact_cycles_arrow (A : homological_complex C c) (i : ι) :
  exact (cycles A i).arrow (d_from A i) :=
exact_kernel_subobject_arrow _ _ _

lemma exact_cycles_map {A₁ A₂ A₃ : homological_complex C c} (f : A₁ ⟶ A₂) (g : A₂ ⟶ A₃)
  (hfg : ∀ n, short_exact (f.f n) (g.f n)) (j : ι) :
  exact (cycles_map f j) (cycles_map g j) :=
begin
  have sq₁ :  d_from A₁ j ≫ f.next j = f.f j ≫ d_from A₂ j := (hom.comm_from _ _).symm,
  have sq₂ :  d_from A₂ j ≫ g.next j = g.f j ≫ d_from A₃ j := (hom.comm_from _ _).symm,
  suffices S : snake
    ↑(cycles A₁ j) ↑(cycles A₂ j) ↑(cycles A₃ j)
    (A₁.X j) (A₂.X j) (A₃.X j)
    _ _ _
    _ _ _
    (cycles_map f j) (cycles_map g j)
    (cycles _ j).arrow (cycles _ j).arrow (cycles _ j).arrow
    (f.f j) (g.f j)
    (A₁.d_from j) (A₂.d_from j) (A₃.d_from j)
    (f.next j) (g.next j)
    (cokernel.π $ A₁.d_from j) (cokernel.π $ A₂.d_from j) (cokernel.π $ A₃.d_from j)
    (cokernel.map _ _ _ _ sq₁) (cokernel.map _ _ _ _ sq₂),
  { exact S.six_term_exact_seq.pair },
  have hfg_epi := λ j, (hfg j).epi,
  have hfg_mono := λ j, (hfg j).mono,
  resetI,
  fsplit,
  { exact (hfg j).exact },
  { exact exact_next' _ _ _ (λ i, (hfg i).exact), },
  { refine (exact_cycles_arrow _ _).cons (abelian.exact_cokernel _).exact_seq, },
  { refine (exact_cycles_arrow _ _).cons (abelian.exact_cokernel _).exact_seq, },
  { refine (exact_cycles_arrow _ _).cons (abelian.exact_cokernel _).exact_seq, },
  { rw cycles_map_arrow, },
  { rw cycles_map_arrow, },
  { exact sq₁ },
  { exact sq₂ },
  { apply cokernel.π_desc, },
  { apply cokernel.π_desc, },
end

variables {A₁ A₂ A₃ : homological_complex C c} (f : A₁ ⟶ A₂) (g : A₂ ⟶ A₃)
variables (hfg : ∀ n, short_exact (f.f n) (g.f n))

lemma mono_cycles_map (hfg : ∀ n, short_exact (f.f n) (g.f n)) (i : ι) :
  mono (cycles_map f i) :=
begin
  apply_with (mono_of_mono _ (subobject.arrow _)) { instances := ff },
  rw cycles_map_arrow,
  haveI : mono (f.f i) := (hfg i).mono,
  apply mono_comp,
end

@[simp] lemma image_subobject_arrow {X : C} (S : subobject X) :
  image_subobject (S.arrow) = S :=
begin
  delta image_subobject,
  ext,
  swap,
  { exact limits.image_mono_iso_source _ },
  { simp }
end

@[simp] lemma kernel_subobject_cokernel.π {X : C} (S : subobject X) :
  kernel_subobject (cokernel.π S.arrow) = S :=
begin
  delta kernel_subobject,
  ext,
  swap,
  { exact (abelian.image_iso_image _).trans (limits.image_mono_iso_source _) },
  { simp }
end

lemma exact.congr {X₁ X₂ Y Z₁ Z₂ : C} (f₁ : X₁ ⟶ Y) (g₁ : Y ⟶ Z₁) (f₂ : X₂ ⟶ Y) (g₂ : Y ⟶ Z₂)
  (h : exact f₁ g₁) (him : image_subobject f₁ = image_subobject f₂)
  (hker : kernel_subobject g₁ = kernel_subobject g₂) :
  exact f₂ g₂ :=
by rwa [abelian.exact_iff_image_eq_kernel, ← him, ← hker, ← abelian.exact_iff_image_eq_kernel]

lemma exact_column :
exact_seq C [(kernel.ι (A.d_to j)), (A.d_to j), (cokernel.π (A.boundaries j).arrow)] :=
exact_kernel_ι.cons $
(exact.congr (boundaries A j).arrow _ _ _ (abelian.exact_cokernel _) (image_subobject_arrow _) rfl).exact_seq

lemma exact_mod_boundaries_map (hfg : ∀ n, short_exact (f.f n) (g.f n)) (j : ι) :
  exact (mod_boundaries_map f j) (mod_boundaries_map g j) :=
begin
  have sq1 : A₁.d_to j ≫ f.f j = f.prev j ≫ A₂.d_to j := (f.comm_to _).symm,
  have sq2 : A₂.d_to j ≫ g.f j = g.prev j ≫ A₃.d_to j := (g.comm_to _).symm,
  suffices S : snake
    -- the objects
         (kernel _)           (kernel _)           (kernel _)
        (A₁.X_prev j)         (A₂.X_prev j)         (A₃.X_prev j)
          (A₁.X j)             (A₂.X j)             (A₃.X j)
    (mod_boundaries _ j) (mod_boundaries _ j) (mod_boundaries _ j)
    -- the morphisms
    (kernel.map _ _ _ _ sq1) (kernel.map _ _ _ _ sq2)
    (kernel.ι $ A₁.d_to j) (kernel.ι $ A₂.d_to j) (kernel.ι $ A₃.d_to j)
    (f.prev j) (g.prev j)
    (A₁.d_to j) (A₂.d_to j) (A₃.d_to j)
    (f.f j) (g.f j)
    (cokernel.π _) (cokernel.π _) (cokernel.π _)
    (mod_boundaries_map f j) (mod_boundaries_map g j),
  { exact (S.six_term_exact_seq.drop 3).pair },
  have hfg_epi := λ n, (hfg n).epi,
  have hfg_mono := λ n, (hfg n).mono,
  resetI,
  fsplit,
  { exact exact_prev' _ _ _ (λ n, (hfg n).exact) },
  { exact (hfg j).exact },
  { apply exact_column },
  { apply exact_column },
  { apply exact_column },
  { simp },
  { simp },
  { exact sq1 },
  { exact sq2 },
  { simp [mod_boundaries_map] },
  { simp [mod_boundaries_map] }
end

lemma epi_mod_boundaries_map (hfg : ∀ n, short_exact (f.f n) (g.f n)) (i : ι) :
  epi (mod_boundaries_map g i) :=
begin
  apply_with (epi_of_epi (cokernel.π _)) { instances := ff },
  haveI : epi (g.f i) := (hfg i).epi,
  have : cokernel.π _ ≫ mod_boundaries_map g i = g.f i ≫ cokernel.π _ := cokernel.π_desc _ _ _,
  rw this,
  apply epi_comp,
end

lemma mono_homology_to_mod_boundaries :
  mono ((homology_to_mod_boundaries i).app A) :=
cokernel.map_mono_of_epi_of_mono
  (boundaries A i) (cycles A i)
  (boundaries A i) (A.X i)
  _ _ _ _ _

variables {C}

@[simp] lemma image_subobject_comp_eq_of_epi {X Y Z : C} (f : X ⟶ Y) (g : Y ⟶ Z) [epi f] :
  image_subobject (f ≫ g) = image_subobject g :=
begin
  delta image_subobject,
  haveI : is_iso (image.pre_comp f g) := is_iso_of_mono_of_epi _,
  ext, swap,
  { exact as_iso (image.pre_comp f g) },
  { simp only [as_iso_hom, image.pre_comp_ι], },
end

@[simp] lemma kernel_subobject_comp_eq_of_mono {X Y Z : C} (f : X ⟶ Y) (g : Y ⟶ Z) [mono g] :
  kernel_subobject (f ≫ g) = kernel_subobject f :=
begin
  delta kernel_subobject,
  ext, swap,
  { exact kernel_comp_mono f g },
  { simp only [kernel_comp_mono_hom, kernel.lift_ι] },
end

lemma exact_cycles_arrow_delta_to_cycles :
  exact (A.cycles i).arrow (delta_to_cycles A i j hij) :=
begin
  rw [category_theory.abelian.exact_iff_image_eq_kernel],
  dsimp [delta_to_cycles, delta_to_boundaries],
  simp only [image_subobject_arrow, kernel_subobject_comp_eq_of_mono],
  delta cycles,
  let g : ↑(A.boundaries j) ⟶ X_next A i := (A.boundaries j).arrow ≫ (X_next_iso _ hij).inv,
  haveI : mono g := mono_comp _ _,
  suffices aux : delta_to_boundaries _ i j hij ≫ g = d_from A i,
  { simp_rw [← aux, kernel_subobject_comp_eq_of_mono], refl, },
  simp only [delta_to_boundaries_comp_arrow_assoc, iso.comp_inv_eq, d_from_comp_X_next_iso],
end

lemma exact_homology_to_mod_boundaries_to_cycles :
  exact ((homology_to_mod_boundaries i).app A) ((mod_boundaries_to_cycles i j hij).app A) :=
begin
  let φ : homology A i ⟶ mod_boundaries A i :=
    limits.cokernel.desc _ ((kernel_subobject _).arrow ≫ (cokernel.π _)) (by simp),
  suffices S : snake
    (0:C) 0 0
    (A.boundaries i) (boundaries A i) 0
    (A.cycles i) (A.X i) (A.cycles j)
    (homology A i) (mod_boundaries A i) (A.cycles j)
    0 0
    0 0 0
    (𝟙 _) 0
    (boundaries_to_cycles _ _) (A.boundaries i).arrow 0
    (A.cycles i).arrow (delta_to_cycles _ i j hij)
    (homology.π _ _ _) (cokernel.π _) (𝟙 _)
    φ ((mod_boundaries_to_cycles i j hij).app A),
  { exact (S.six_term_exact_seq.drop 3).pair },
  letI : exact (cycles A i).arrow (delta_to_cycles A i j hij) :=
    exact_cycles_arrow_delta_to_cycles _ i j hij,
  letI : epi (homology.π (d_to A i) (d_from A i) (A.d_to_comp_d_from i)) := coequalizer.π_epi,
  fsplit,
  { rw ← epi_iff_exact_zero_right, apply_instance },
  { apply exact_cycles_arrow_delta_to_cycles },
  { exact (category_theory.exact_zero_mono _).cons (abelian.exact_cokernel _).exact_seq, },
  { exact (category_theory.exact_zero_mono _).cons (abelian.exact_cokernel _).exact_seq, },
  { exact (category_theory.exact_zero_mono _).cons (exact_zero_left_of_mono _).exact_seq, },
  { simp only [zero_comp] },
  { simp only [zero_comp] },
  { simp only [image_to_kernel_arrow, category.id_comp] },
  { simp only [boundaries_arrow_comp_delta_to_cycles, zero_comp], },
  { dsimp [homology.π], simp only [cokernel.π_desc] },
  { simp only [mod_boundaries_to_cycles_app, cokernel.π_desc, category.comp_id] },
end

lemma exact_mod_boundaries_to_cycles_to_homology :
  exact ((mod_boundaries_to_cycles i j hij).app A) ((cycles_to_homology j).app A)  :=
begin
  refine exact.congr (boundaries_to_cycles _ _) _ _ _ _ _ rfl,
  { exact abelian.exact_cokernel _, },
  { simp only [mod_boundaries_to_cycles_app],
    delta delta_to_cycles,
    rw [← image_subobject_comp_eq_of_epi (cokernel.π _)],
    simp only [cokernel.π_desc, image_subobject_comp_eq_of_epi], }
end

lemma epi_cycles_to_homology : epi ((cycles_to_homology j).app A) :=
coequalizer.π_epi

lemma exact_seq_column :
  exact_seq C
    [((homology_to_mod_boundaries i).app A₁),
     ((mod_boundaries_to_cycles i j hij).app A₁),
     ((cycles_to_homology j).app A₁)] :=
(exact_homology_to_mod_boundaries_to_cycles _ _ _ _).cons
  (exact_mod_boundaries_to_cycles_to_homology _ _ _ _).exact_seq

lemma snake (hfg : ∀ n, short_exact (f.f n) (g.f n)) (i j : ι) (hij : c.rel i j) :
  snake
  -- the objects
     (A₁.homology i)       (A₂.homology i)       (A₃.homology i)
  (A₁.mod_boundaries i) (A₂.mod_boundaries i) (A₃.mod_boundaries i)
      (A₁.cycles j)         (A₂.cycles j)         (A₃.cycles j)
     (A₁.homology j)       (A₂.homology j)       (A₃.homology j)
  -- the morphisms
  ((homology_functor _ _ i).map f) ((homology_functor _ _ i).map g)
  ((homology_to_mod_boundaries i).app A₁)
  ((homology_to_mod_boundaries i).app A₂)
  ((homology_to_mod_boundaries i).app A₃)
  ((mod_boundaries_functor i).map f) ((mod_boundaries_functor i).map g)
  ((mod_boundaries_to_cycles i j hij).app A₁)
  ((mod_boundaries_to_cycles i j hij).app A₂)
  ((mod_boundaries_to_cycles i j hij).app A₃)
  ((cycles_functor _ _ j).map f) ((cycles_functor _ _ j).map g)
  ((cycles_to_homology j).app A₁)
  ((cycles_to_homology j).app A₂)
  ((cycles_to_homology j).app A₃)
  ((homology_functor _ _ j).map f) ((homology_functor _ _ j).map g) :=
{ row_exact₁ := exact_mod_boundaries_map f g hfg _,
  row_exact₂ := exact_cycles_map f g hfg _,
  row_epi := epi_mod_boundaries_map f g hfg _,
  row_mono := mono_cycles_map f g hfg _,
  col_exact_a := exact_seq_column _ _ _,
  col_exact_b := exact_seq_column _ _ _,
  col_exact_c := exact_seq_column _ _ _,
  col_mono_a := mono_homology_to_mod_boundaries _ _,
  col_mono_b := mono_homology_to_mod_boundaries _ _,
  col_mono_c := mono_homology_to_mod_boundaries _ _,
  col_epi_a := epi_cycles_to_homology _ _,
  col_epi_b := epi_cycles_to_homology _ _,
  col_epi_c := epi_cycles_to_homology _ _,
  sq_a₀ := ((homology_to_mod_boundaries _).naturality _).symm,
  sq_b₀ := ((homology_to_mod_boundaries _).naturality _).symm,
  sq_a₁ := ((mod_boundaries_to_cycles _ _ _).naturality _).symm,
  sq_b₁ := ((mod_boundaries_to_cycles _ _ _).naturality _).symm,
  sq_a₂ := ((cycles_to_homology _).naturality _).symm,
  sq_b₂ := ((cycles_to_homology _).naturality _).symm }

def δ (hfg : ∀ n, short_exact (f.f n) (g.f n)) (i j : ι) (hij : c.rel i j) :
  homology A₃ i ⟶ homology A₁ j :=
(snake f g hfg i j hij).δ

lemma six_term_exact_seq (hfg : ∀ n, short_exact (f.f n) (g.f n)) (i j : ι) (hij : c.rel i j) :
  exact_seq C [
    (homology_functor _ _ i).map f, -- Hⁱ(A₁) ⟶ Hⁱ(A₂)
    (homology_functor _ _ i).map g, -- Hⁱ(A₂) ⟶ Hⁱ(A₃)
    δ f g hfg i j hij,              -- Hⁱ(A₃) ⟶ Hʲ(A₁)
    (homology_functor _ _ j).map f, -- Hʲ(A₁) ⟶ Hʲ(A₂)
    (homology_functor _ _ j).map g  -- Hʲ(A₁) ⟶ Hʲ(A₃)
  ] :=
(snake f g hfg i j hij).six_term_exact_seq

end homological_complex


/-
Copyright (c) 2021 Mario Carneiro. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.

Authors: Mario Carneiro, Gabriel Ebner
-/
import Std.Data.Nat.Lemmas
import Std.Data.List.Lemmas
import Std.Tactic.HaveI
import Std.Tactic.Simpa

local macro_rules | `($x[$i]'$h) => `(getElem $x $i $h)

@[simp] theorem getElem_fin [GetElem Cont Nat Elem Dom] (a : Cont) (i : Fin n) (h : Dom a i) :
    a[i] = a[i.1] := rfl

@[simp] theorem getElem?_fin [GetElem Cont Nat Elem Dom] (a : Cont) (i : Fin n)
    [Decidable (Dom a i)] : a[i]? = a[i.1]? := rfl

@[simp] theorem getElem!_fin [GetElem Cont Nat Elem Dom] (a : Cont) (i : Fin n)
    [Decidable (Dom a i)] [Inhabited Elem] : a[i]! = a[i.1]! := rfl

theorem getElem?_pos [GetElem Cont Idx Elem Dom]
    (a : Cont) (i : Idx) (h : Dom a i) [Decidable (Dom a i)] : a[i]? = a[i] := dif_pos h

theorem getElem?_neg [GetElem Cont Idx Elem Dom]
    (a : Cont) (i : Idx) (h : ¬Dom a i) [Decidable (Dom a i)] : a[i]? = none := dif_neg h

@[simp] theorem mkArray_data (n : Nat) (v : α) : (mkArray n v).data = List.replicate n v := rfl

namespace Array

attribute [simp] isEmpty uget

@[simp] theorem singleton_def (v : α) : singleton v = #[v] := rfl

@[simp] theorem toArray_data : (a : Array α) → a.data.toArray = a
  | ⟨l⟩ => ext' (data_toArray l)

@[simp] theorem get_eq_getElem (a : Array α) (i : Fin _) : a.get i = a[i.1] := rfl
@[simp] theorem get?_eq_getElem? (a : Array α) (i : Nat) : a.get? i = a[i]? := rfl
theorem getElem_fin_eq_data_get (a : Array α) (i : Fin _) : a[i] = a.data.get i := rfl

@[simp] theorem ugetElem_eq_getElem (a : Array α) {i : USize} (h : i.toNat < a.size) :
  a[i] = a[i.toNat] := rfl

theorem getElem?_eq_getElem (a : Array α) (i : Nat) (h : i < a.size) : a[i]? = a[i] :=
  getElem?_pos ..

theorem get?_len_le (a : Array α) (i : Nat) (h : a.size ≤ i) : a[i]? = none := by
  simp [getElem?_neg, h]

theorem getElem_mem_data (a : Array α) (h : i < a.size) : a[i] ∈ a.data := by
  simp [getElem_eq_data_get, List.get_mem]

theorem getElem?_eq_data_get? (a : Array α) (i : Nat) : a[i]? = a.data.get? i := by
  by_cases i < a.size <;> simp_all [getElem?_pos, getElem?_neg, List.get?_eq_get, eq_comm]; rfl

theorem get?_eq_data_get? (a : Array α) (i : Nat) : a.get? i = a.data.get? i :=
  getElem?_eq_data_get? ..

@[simp] theorem getD_eq_get? (a : Array α) (n d) : a.getD n d = (a.get? n).getD d := by
  simp [get?, getD]; split <;> simp

theorem get!_eq_getD [Inhabited α] (a : Array α) : a.get! n = a.getD n default := rfl

@[simp] theorem get!_eq_get? [Inhabited α] (a : Array α) : a.get! n = (a.get? n).getD default := by
  simp [get!_eq_getD]

@[simp] theorem back_eq_back? [Inhabited α] (a : Array α) : a.back = a.back?.getD default := by
  simp [back, back?]

@[simp] theorem back?_push (a : Array α) : (a.push x).back? = some x := by
  simp [back?, getElem?_eq_data_get?]

theorem back_push [Inhabited α] (a : Array α) : (a.push x).back = x := by simp

theorem get?_push_lt (a : Array α) (x : α) (i : Nat) (h : i < a.size) :
    (a.push x)[i]? = some a[i] := by
  rw [getElem?_pos, get_push_lt]

theorem get?_push_eq (a : Array α) (x : α) : (a.push x)[a.size]? = some x := by
  rw [getElem?_pos, get_push_eq]

@[simp] theorem data_set (a : Array α) (i v) : (a.set i v).data = a.data.set i.1 v := rfl

@[simp] theorem get_set_eq (a : Array α) (i : Fin a.size) (v : α) :
    (a.set i v)[i.1]'(by simp [i.2]) = v := by
  simp only [set, getElem_eq_data_get, List.get_set_eq]

@[simp] theorem get_set_ne (a : Array α) (i : Fin a.size) {j : Nat} (v : α) (hj : j < a.size)
    (h : i.1 ≠ j) : (a.set i v)[j]'(by simp [*]) = a[j] := by
  simp only [set, getElem_eq_data_get, List.get_set_ne h]

@[simp] theorem get?_set_eq (a : Array α) (i : Fin a.size) (v : α) :
    (a.set i v)[i.1]? = v := by simp [getElem?_pos, i.2]

@[simp] theorem get?_set_ne (a : Array α) (i : Fin a.size) {j : Nat} (v : α)
    (h : i.1 ≠ j) : (a.set i v)[j]? = a[j]? := by
  by_cases j < a.size <;> simp [getElem?_pos, getElem?_neg, *]

theorem get?_set (a : Array α) (i : Fin a.size) (j : Nat) (v : α) :
    (a.set i v)[j]? = if i.1 = j then some v else a[j]? := by
  if h : i.1 = j then subst j; simp [*] else simp [*]

theorem get_set (a : Array α) (i : Fin a.size) (j : Nat) (hj : j < a.size) (v : α) :
    (a.set i v)[j]'(by simp [*]) = if i = j then v else a[j] := by
  if h : i.1 = j then subst j; simp [*] else simp [*]

private theorem fin_cast_val (e : n = n') (i : Fin n) : e ▸ i = ⟨i.1, e ▸ i.2⟩ := by cases e; rfl

theorem swap_def (a : Array α) (i j : Fin a.size) :
    a.swap i j = (a.set i (a.get j)).set ⟨j.1, by simp [j.2]⟩ (a.get i) := by
  simp [swap, fin_cast_val]

theorem data_swap (a : Array α) (i j : Fin a.size) :
    (a.swap i j).data = (a.data.set i (a.get j)).set j (a.get i) := by simp [swap_def]

theorem get?_swap (a : Array α) (i j : Fin a.size) (k : Nat) : (a.swap i j)[k]? =
    if j = k then some a[i.1] else if i = k then some a[j.1] else a[k]? := by
  simp [swap_def, get?_set, ← getElem_fin_eq_data_get]

@[simp] theorem swapAt_def (a : Array α) (i : Fin a.size) (v : α) :
    a.swapAt i v = (a[i.1], a.set i v) := rfl

-- @[simp] -- FIXME: gives a weird linter error

@[simp] theorem data_pop (a : Array α) : a.pop.data = a.data.dropLast := by simp [pop]

@[simp] theorem pop_empty : (#[] : Array α).pop = #[] := rfl

@[simp] theorem pop_push (a : Array α) : (a.push x).pop = a := by simp [pop]

theorem SatisfiesM_foldrM [Monad m] [LawfulMonad m]
    {as : Array α} (motive : Nat → β → Prop)
    {init : β} (h0 : motive as.size init) {f : α → β → m β}
    (hf : ∀ i : Fin as.size, ∀ b, motive (i.1 + 1) b → SatisfiesM (motive i.1) (f as[i] b)) :
    SatisfiesM (motive 0) (as.foldrM f init) := by
  let rec go {i b} (hi : i ≤ as.size) (H : motive i b) :
    SatisfiesM (motive 0) (foldrM.fold f as 0 i hi b) := by
    unfold foldrM.fold; simp; split
    · next hi => exact .pure (hi ▸ H)
    · next hi =>
      split; {simp at hi}
      · next i hi' =>
        exact (hf ⟨i, hi'⟩ b H).bind fun _ => go _
  simp [foldrM]; split; {exact go _ h0}
  · next h => exact .pure (Nat.eq_zero_of_nonpos _ h ▸ h0)

theorem foldr_induction
    {as : Array α} (motive : Nat → β → Prop) {init : β} (h0 : motive as.size init) {f : α → β → β}
    (hf : ∀ i : Fin as.size, ∀ b, motive (i.1 + 1) b → motive i.1 (f as[i] b)) :
    motive 0 (as.foldr f init) := by
  have := SatisfiesM_foldrM (m := Id) (as := as) (f := f) motive h0
  simp [SatisfiesM_Id_eq] at this
  exact this hf

theorem mapM_eq_mapM_data [Monad m] [LawfulMonad m] (f : α → m β) (arr : Array α) :
    arr.mapM f = return mk (← arr.data.mapM f) := by
  rw [mapM, foldlM_eq_foldlM_data, ← List.foldrM_reverse, mkEmpty_eq]
  conv => rhs; rw [← List.reverse_reverse arr.data]
  induction arr.data.reverse with
  | nil => simp; rfl
  | cons a l ih => simp [ih]; simp [push]

theorem SatisfiesM_mapIdxM [Monad m] [LawfulMonad m] (as : Array α) (f : Fin as.size → α → m β)
    (motive : Nat → Prop) (h0 : motive 0)
    (p : Fin as.size → β → Prop)
    (hs : ∀ i, motive i.1 → SatisfiesM (p i · ∧ motive (i + 1)) (f i as[i])) :
    SatisfiesM
      (fun arr => motive as.size ∧ ∃ eq : arr.size = as.size, ∀ i h, p ⟨i, h⟩ (arr[i]'(eq ▸ h)))
      (Array.mapIdxM as f) := by
  let rec go {bs i j h} (h₁ : j = bs.size) (h₂ : ∀ i h h', p ⟨i, h⟩ bs[i]) (hm : motive j) :
    SatisfiesM
      (fun arr => motive as.size ∧ ∃ eq : arr.size = as.size, ∀ i h, p ⟨i, h⟩ (arr[i]'(eq ▸ h)))
      (Array.mapIdxM.map as f i j h bs) := by
    induction i generalizing j bs with simp [mapIdxM.map]
    | zero =>
      have := (Nat.zero_add _).symm.trans h
      exact .pure ⟨this ▸ hm, h₁ ▸ this, fun _ _ => h₂ ..⟩
    | succ i ih =>
      refine (hs _ (by exact hm)).bind fun b hb => ih (by simp [h₁]) (fun i hi hi' => ?_) hb.2
      simp at hi'; simp [get_push]; split
      · next h => exact h₂ _ _ h
      · next h => cases h₁.symm ▸ (Nat.le_or_eq_of_le_succ hi').resolve_left h; exact hb.1
  simp [mapIdxM]; exact go rfl (fun.) h0

theorem mapIdx_induction (as : Array α) (f : Fin as.size → α → β)
    (motive : Nat → Prop) (h0 : motive 0)
    (p : Fin as.size → β → Prop)
    (hs : ∀ i, motive i.1 → p i (f i as[i]) ∧ motive (i + 1)) :
    motive as.size ∧ ∃ eq : (Array.mapIdx as f).size = as.size,
      ∀ i h, p ⟨i, h⟩ ((Array.mapIdx as f)[i]'(eq ▸ h)) := by
  have := SatisfiesM_mapIdxM (m := Id) (as := as) (f := f) motive h0
  simp [SatisfiesM_Id_eq] at this
  exact this _ hs

theorem mapIdx_induction' (as : Array α) (f : Fin as.size → α → β)
    (p : Fin as.size → β → Prop) (hs : ∀ i, p i (f i as[i])) :
    ∃ eq : (Array.mapIdx as f).size = as.size,
      ∀ i h, p ⟨i, h⟩ ((Array.mapIdx as f)[i]'(eq ▸ h)) :=
  (mapIdx_induction _ _ (fun _ => True) trivial p fun _ _ => ⟨hs .., trivial⟩).2

@[simp] theorem size_mapIdx (a : Array α) (f : Fin a.size → α → β) : (a.mapIdx f).size = a.size :=
  (mapIdx_induction' (p := fun _ _ => True) (hs := fun _ => trivial)).1

@[simp] theorem getElem_mapIdx (a : Array α) (f : Fin a.size → α → β) (i : Nat) (h) :
    haveI : i < a.size := by simp_all
    (a.mapIdx f)[i]'h = f ⟨i, this⟩ a[i] :=
  (mapIdx_induction' _ _ (fun i b => b = f i a[i]) fun _ => rfl).2 i _

@[simp] theorem size_swap! (a : Array α) (i j) (hi : i < a.size) (hj : j < a.size) :
    (a.swap! i j).size = a.size := by simp [swap!, hi, hj]

@[simp] theorem size_reverse (a : Array α) : a.reverse.size = a.size := by
  let rec go (as : Array α) (i j) : (reverse.loop as i j).size = as.size := by
    rw [reverse.loop]
    if h : i < j then
      have := reverse.termination h
      simp [(go · (i+1) ⟨j-1, ·⟩), h]
    else simp [h]
  simp only [reverse]; split <;> simp [go]
termination_by _ => j - i

@[simp] theorem reverse_data (a : Array α) : a.reverse.data = a.data.reverse := by
  let rec go (as : Array α) (i j hj)
      (h : i + j + 1 = a.size) (h₂ : as.size = a.size)
      (H : ∀ k, as.data.get? k = if i ≤ k ∧ k ≤ j then a.data.get? k else a.data.reverse.get? k)
      (k) : (reverse.loop as i ⟨j, hj⟩).data.get? k = a.data.reverse.get? k := by
    rw [reverse.loop]; dsimp; split <;> rename_i h₁
    · have := reverse.termination h₁
      match j with | j+1 => ?_
      simp at *
      simp; rw [(go · (i+1) j)]
      · rwa [Nat.add_right_comm i]
      · simp [size_swap, h₂]
      · intro k
        rw [← getElem?_eq_data_get?, get?_swap]
        simp [getElem?_eq_data_get?, getElem_eq_data_get, ← List.get?_eq_get, H, Nat.le_of_lt h₁]
        split <;> rename_i h₂
        · simp [← h₂, Nat.not_le.2 (Nat.lt_succ_self _)]
          exact (List.get?_reverse' _ _ (Eq.trans (by simp_arith) h)).symm
        split <;> rename_i h₃
        · simp [← h₃, Nat.not_le.2 (Nat.lt_succ_self _)]
          exact (List.get?_reverse' _ _ (Eq.trans (by simp_arith) h)).symm
        simp only [Nat.succ_le, Nat.lt_iff_le_and_ne.trans (and_iff_left h₃),
          Nat.lt_succ.symm.trans (Nat.lt_iff_le_and_ne.trans (and_iff_left (Ne.symm h₂)))]
    · rw [H]; split <;> rename_i h₂
      · cases Nat.le_antisymm (Nat.not_lt.1 h₁) (Nat.le_trans h₂.1 h₂.2)
        cases Nat.le_antisymm h₂.1 h₂.2
        exact (List.get?_reverse' _ _ h).symm
      · rfl
  simp only [reverse]; split
  · match a with | ⟨[]⟩ | ⟨[_]⟩ => rfl
  · have := Nat.sub_add_cancel (Nat.le_of_not_le ‹_›)
    refine List.ext <| go _ _ _ _ (by simp [this]) rfl fun k => ?_
    split; {rfl}; rename_i h
    simp [← show k < _ + 1 ↔ _ from Nat.lt_succ (n := a.size - 1), this] at h
    rw [List.get?_eq_none.2 ‹_›, List.get?_eq_none.2 (a.data.length_reverse ▸ ‹_›)]
termination_by _ => j - i

@[simp] theorem size_ofFn_go {n} (f : Fin n → α) (i acc) :
    (ofFn.go f i acc).size = acc.size + (n - i) := by
  if hin : i < n then
    unfold ofFn.go
    have : 1 + (n - (i + 1)) = n - i :=
      Nat.sub_sub .. ▸ Nat.add_sub_cancel' (Nat.le_sub_of_add_le (Nat.add_comm .. ▸ hin))
    rw [dif_pos hin, size_ofFn_go f (i+1), size_push, Nat.add_assoc, this]
  else
    have : n - i = 0 := Nat.sub_eq_zero_of_le (Nat.le_of_not_lt hin)
    unfold ofFn.go
    simp [hin, this]
termination_by _ => n - i

@[simp] theorem size_ofFn (f : Fin n → α) : (ofFn f).size = n := by simp [ofFn]

theorem getElem_ofFn_go (f : Fin n → α) (i) {acc k}
    (hki : k < n) (hin : i ≤ n) (hi : i = acc.size)
    (hacc : ∀ j, ∀ hj : j < acc.size, acc[j] = f ⟨j, Nat.lt_of_lt_of_le hj (hi ▸ hin)⟩) :
    haveI : acc.size + (n - acc.size) = n := Nat.add_sub_cancel' (hi ▸ hin)
    (ofFn.go f i acc)[k]'(by simp [*]) = f ⟨k, hki⟩ := by
  unfold ofFn.go
  if hin : i < n then
    have : 1 + (n - (i + 1)) = n - i :=
      Nat.sub_sub .. ▸ Nat.add_sub_cancel' (Nat.le_sub_of_add_le (Nat.add_comm .. ▸ hin))
    simp only [dif_pos hin]
    rw [getElem_ofFn_go f (i+1) _ hin (by simp [*]) (fun j hj => ?hacc)]
    cases (Nat.lt_or_eq_of_le <| Nat.le_of_lt_succ (by simpa using hj)) with
    | inl hj => simp [get_push, hj, hacc j hj]
    | inr hj => simp [get_push, *]
  else
    simp [hin, hacc k (Nat.lt_of_lt_of_le hki (Nat.le_of_not_lt (hi ▸ hin)))]
termination_by _ => n - i

@[simp] theorem getElem_ofFn (f : Fin n → α) (i : Nat) (h) :
    (ofFn f)[i] = f ⟨i, size_ofFn f ▸ h⟩ :=
  getElem_ofFn_go _ _ _ (by simp) (by simp) fun.


/-
Copyright (c) 2022 Jannis Limperg. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Jannis Limperg
-/

import Aesop

set_option aesop.check.all true

-- This is a test case for postponed safe rules. When a safe rule assigns mvars,
-- it is not applied but instead postponed. Then we later try it again as an
-- unsafe rule.
--
-- It's hard to test this feature completely because we can't really tell from
-- the outside when a rule has been applied. But we can at least look at the
-- traces of the following test cases.

axiom T : Nat → Prop

@[aesop safe]
axiom t : T 0

example : ∃ (i : Nat), T i := by
  aesop

axiom U : Nat → Prop

@[aesop safe 0]
axiom u₁ : U 0

@[aesop safe 1]
axiom u₂ : U 1

example : ∃ i, U i ∧ i = 1 := by
  aesop


example (h : α = β) : h ▸ (a : α) = (b : β) := _
example (h : α = β) : id h ▸ (a : α) = (b : β) := _
example (h : α = β) : id h ▸ (a : α) = (b : β) := by simp
set_option pp.proofs.withType false
example (h : α = β) : id h ▸ (a : α) = (b : β) := _
set_option pp.proofs true
example (h : α = β) : id h ▸ (a : α) = (b : β) := _