-- LeanMLIR/LeanMLIR/Transforms/DCE.lean

import LeanMLIR.Framework

import Mathlib.Tactic.Linarith

import Mathlib.Tactic.DepRewrite

open Ctxt (Var Valuation Hom)

namespace DCE

structure DeleteRange (Γ : Ctxt Ty) where
   
  start : Fin (Γ.length + 1)
   
  num : Fin (Γ.length + 1 - start.val)

open DCE (DeleteRange)

namespace DeleteRange

def appendInl {Γ : Ctxt Ty} {ts : List Ty}
    (r : DeleteRange Γ) : DeleteRange (⟨ts⟩ ++ Γ) where
  start := ⟨r.start + ts.length, by admit /- proof elided -/
  ⟩
  num := ⟨r.num, by admit /- proof elided -/
  ⟩

section Lemmas

end Lemmas

end DeleteRange

end DCE

open DCE (DeleteRange)

def Ctxt.delete (Γ : Ctxt Ty) (vs : DeleteRange Γ) : Ctxt Ty :=
  Ctxt.ofList <| List.ofFn (n := Γ.length - vs.num.val) fun i =>
    have := vs.start.prop
    if hi : i.val < vs.start then
      Γ[i.val]
    else
      Γ[i.val + vs.num]

section Lemmas

variable {Γ : Ctxt Ty}

end Lemmas

def Hom.delete {Γ : Ctxt Ty} (delv : DeleteRange Γ) : Hom (Γ.delete delv) Γ :=
  fun t' v =>
    let idx :=
      if v.val < delv.start then
        v.val
      else
        v.val + delv.num
    ⟨idx, by admit /- proof elided -/
    ⟩

def Deleted (Γ: Ctxt Ty) (vs : DeleteRange Γ) (Γ' : Ctxt Ty) : Prop :=
  Γ' = Γ.delete vs

def Deleted.toHom (h : Deleted Γ r Γ') : Γ'.Hom Γ :=
  fun _ v => Hom.delete r (v.castCtxt h)

-- Repository definitions used:
content: structure Ctxt (Ty : Type) : Type where
  ofList :: toList : List Ty
  
  deriving Repr, Lean.ToExpr, DecidableEq

content: @[grind=]
def length (Γ : Ctxt Ty) : Nat := Γ.toList.length

content: inductive Ty (q : Nat) (n : Nat)
  | index : Ty q n
  | integer : Ty q n
  | tensor : Ty q n
  | polynomialLike : Ty q n
  deriving DecidableEq, Repr

content: def Hom.append {ζ : Ctxt Ty} (f : Γ.Hom Δ) : Hom (ζ ++ Γ) (ζ ++ Δ) :=
  fun _ => Var.appendCases
    (fun v => v.appendInl)
    (fun v => (f v).appendInr)

content: @[elab_as_elim]
def appendCases
    {motive : (Γ ++ Δ).Var t → Sort u}
    (left : (v : Var Γ t) → motive (appendInl v))
    (right : (v : Var Δ t) → motive (appendInr v)) :
    (v : (Γ ++ Δ).Var t) → motive v
  | ⟨idx, h⟩ =>
    if hv : idx < Γ.length then
      left ⟨idx, by admit /- proof elided -/
      ⟩
    else
      let v' : Var _ _ := ⟨idx - Γ.length, by admit /- proof elided -/
      ⟩
      have eq : v'.appendInr = ⟨idx, h⟩ := by admit /- proof elided -/

content: def appendInl (v : Γ.Var t) : (Γ ++ Δ).Var t :=
  ⟨v.val, by admit /- proof elided -/
  ⟩

content: def Var (Γ : Ctxt Ty) (t : Ty) : Type :=
  { i : Nat // Γ[i]? = some t }

content: def toHom (d : Diff Γ₁ Γ₂) : Hom Γ₁ Γ₂ :=
  fun _ v => ⟨v.val + d.val, d.property v.property⟩

content: abbrev Hom (Γ Γ' : Ctxt Ty) := ⦃t : Ty⦄ → Γ.Var t → Γ'.Var t

content: def Diff (Γ₁ Γ₂ : Ctxt Ty) : Type :=
  {d : Nat // Diff.Valid Γ₁ Γ₂ d}

content: @[simp]
abbrev Diff.Valid (Γ₁ Γ₂ : Ctxt Ty) (d : Nat) : Prop :=
  ∀ {i t}, Γ₁[i]? = some t → Γ₂[i+d]? = some t

-- Repository lemmas used:
content: @[simp] theorem ofList_append {ts us : List Ty} :
    Ctxt.ofList ts ++ Ctxt.ofList us = Ctxt.ofList (ts ++ us)

content: @[simp, grind=] theorem length_ofList : (ofList ts).length = ts.length

content: @[simp, grind=] theorem getElem?_ofList (i : Nat) : (ofList ts)[i]? = ts[i]?

content: @[simp, grind=] theorem getElem_ofList (i : Nat) (h : _) : (ofList ts)[i]'h = ts[i]'h

content: theorem val_lt (v : Γ.Var t) : v.val < Γ.length

content: @[simp, grind=] theorem castCtxt_rfl (h : Γ = Γ) : v.castCtxt h = v

content: @[simp, grind=] theorem val_castCtxt : (castCtxt h v).val = v.val

-- Local definitions used:
content: structure DeleteRange (Γ : Ctxt Ty) where
   
  start : Fin (Γ.length + 1)
   
  num : Fin (Γ.length + 1 - start.val)

content: def appendInl {Γ : Ctxt Ty} {ts : List Ty}
    (r : DeleteRange Γ) : DeleteRange (⟨ts⟩ ++ Γ) where
  start := ⟨r.start + ts.length, by admit /- proof elided -/
  ⟩
  num := ⟨r.num, by admit /- proof elided -/
  ⟩

content: def Ctxt.delete (Γ : Ctxt Ty) (vs : DeleteRange Γ) : Ctxt Ty :=
  Ctxt.ofList <| List.ofFn (n := Γ.length - vs.num.val) fun i =>
    have := vs.start.prop
    if hi : i.val < vs.start then
      Γ[i.val]
    else
      Γ[i.val + vs.num]

content: def Hom.delete {Γ : Ctxt Ty} (delv : DeleteRange Γ) : Hom (Γ.delete delv) Γ :=
  fun t' v =>
    let idx :=
      if v.val < delv.start then
        v.val
      else
        v.val + delv.num
    ⟨idx, by admit /- proof elided -/
    ⟩

content: def Deleted (Γ: Ctxt Ty) (vs : DeleteRange Γ) (Γ' : Ctxt Ty) : Prop :=
  Γ' = Γ.delete vs

content: def Deleted.toHom (h : Deleted Γ r Γ') : Γ'.Hom Γ :=
  fun _ v => Hom.delete r (v.castCtxt h)

-- Local lemmas used:
content: @[simp] theorem Ctxt.delete_append_appendInl {Γ : Ctxt Ty} {us : List Ty}
    {r : DeleteRange Γ} :
    (⟨us⟩ ++ Γ).delete r.appendInl = ⟨us⟩ ++ (Γ.delete r)

-- Target theorem and proof:
@[simp] lemma Deleted.toHom_append {Γ Γ' : Ctxt Ty} {vs : DeleteRange Γ}
    (DEL : Deleted (⟨us⟩ ++ Γ) vs.appendInl (⟨us⟩ ++ Γ')) :
    DEL.toHom
    = have DEL' : Deleted Γ vs Γ'
:= by
        rcases Γ'
        simp only [Deleted, Ctxt.delete_append_appendInl] at DEL ⊢
        injection DEL
        simp_all [Ctxt.delete]
      DEL'.toHom.append := by
  have DEL' : Deleted Γ vs Γ' := by
    rcases Γ'
    simp only [Deleted, Ctxt.delete_append_appendInl] at DEL ⊢
    injection DEL
    simp_all [Ctxt.delete]
  subst DEL'
  funext t v
  apply Subtype.ext
  simp only [toHom, Hom.delete, Var.val_castCtxt, DeleteRange.val_start_appendInl,
    DeleteRange.val_num_appendInl, Hom.append, Var.castCtxt_rfl]
  cases v using Var.appendCases with
  | right _  => simp; grind
  | left v =>
    have := v.val_lt
    simp; grind