-- SSA/Projects/Tensor1D/Tensor1D.lean

import LeanMLIR.Framework

import LeanMLIR.Util

import Mathlib.Tactic.Linarith

namespace Tensor1D

abbrev Index := ℕ

structure Tensor1d (α : Type) [Inhabited α] where
  size : Index
  val :  Index → α
  spec : ∀ (ix: Index), ix >= size -> val ix = default

def Tensor1d.extract [Inhabited α] (t: Tensor1d α)
  (left: Index) (len: Index) : Tensor1d α :=
  let right := if (left + len) < t.size then left + len else 0
  let size := right - left
  { size := size,
    val := fun ix =>
    if left + len < t.size
    then if (ix < len) then t.val (ix + left) else default
    else default,
    spec := by admit /- proof elided -/

def Tensor1d.insertslice  [Inhabited α] (t: Tensor1d α)
  (sliceix: Nat)
  (slice : Tensor1d α) : Tensor1d α where
  size := if sliceix > t.size then 0 else t.size + slice.size
  val := fun ix =>
    if sliceix > t.size then default 
    else if ix >= t.size + slice.size then default 
    else
      let go (ix: Nat) : α :=
        if ix < sliceix then t.val sliceix
        else if ix < sliceix + slice.size then slice.val (ix - sliceix)
        else t.val (ix - (sliceix + slice.size))
      go ix
  spec := by admit /- proof elided -/

-- Repository definitions used:
content: syntax "neg" : MLIR.Pretty.uniform_op

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = const ($x)
      $[: $outer_type]? ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _)
      `(mlir_op| $res:mlir_op_operand = "const"()
          {val = $x:num : $outer_type} : ($outer_type) -> ($outer_type) )

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = li ($x)
     $[: $outer_type]?  ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _ )
      `(mlir_op| $res:mlir_op_operand = "li"()
          {imm = $x:num : $outer_type } : ($outer_type) -> ($outer_type))

content: macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithImmediate $reg1 , $x : $t)  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {imm = $x:num : $t} : ($t) -> ($t) )

content: macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithShamt $reg1 , $x  : $t )  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {shamt = $x:num : $t}  : ($t) -> ($t) )

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: abbrev Index := ℕ

content: structure Tensor1d (α : Type) [Inhabited α] where
  size : Index
  val :  Index → α
  spec : ∀ (ix: Index), ix >= size -> val ix = default

content: def Tensor1d.extract [Inhabited α] (t: Tensor1d α)
  (left: Index) (len: Index) : Tensor1d α :=
  let right := if (left + len) < t.size then left + len else 0
  let size := right - left
  { size := size,
    val := fun ix =>
    if left + len < t.size
    then if (ix < len) then t.val (ix + left) else default
    else default,
    spec := by admit /- proof elided -/

content: def Tensor1d.insertslice  [Inhabited α] (t: Tensor1d α)
  (sliceix: Nat)
  (slice : Tensor1d α) : Tensor1d α where
  size := if sliceix > t.size then 0 else t.size + slice.size
  val := fun ix =>
    if sliceix > t.size then default 
    else if ix >= t.size + slice.size then default 
    else
      let go (ix: Nat) : α :=
        if ix < sliceix then t.val sliceix
        else if ix < sliceix + slice.size then slice.val (ix - sliceix)
        else t.val (ix - (sliceix + slice.size))
      go ix
  spec := by admit /- proof elided -/

-- Local lemmas used:
content: theorem not_lt_is_geq {a b: Nat} (NOT_LT: ¬ (a < b)): a >= b

-- Target theorem and proof:
theorem extractslice_insertslice [Inhabited α]
  (t: Tensor1d α)
  (sliceix: Nat)
  (slice: Tensor1d α)
  (CORRECT: ((t.insertslice sliceix slice).extract sliceix slice.size).size ≠ 0)
  : (t.insertslice sliceix slice).extract sliceix slice.size = slice
:= by {
    simp[Tensor1d.insertslice, Tensor1d.extract]
    cases slice
    simp;
    rename_i slicesize sliceval spec
    by_cases A : (t.size < sliceix) <;> simp[A]

    case pos => simp[Tensor1d.insertslice, Tensor1d.extract, A] at CORRECT ;
    case neg =>
      have B : t.size >= sliceix := not_lt_is_geq A

      by_cases C:(sliceix < t.size) <;> simp[C]
      case neg => simp [Tensor1d.insertslice, Tensor1d.extract, A, C] at CORRECT
      case pos =>
          funext ix
          by_cases D: (ix < slicesize) <;> simp[D]
          case neg =>
            -- here we fail, because we do not know that 'slice' behaves like a
            -- real tensor that returns 'default' outside of its range.
            -- This is something we need to add into the spec of a Tensor.
            have E : ix >= slicesize := by simp[Index] at *; linarith
            simp[spec _ E]
          case pos =>
            try simp
            by_cases E:(t.size + slicesize <= ix + sliceix) <;> simp[E]
            case pos =>
              have CONTRA : False := by simp[Index] at *; linarith;
              simp at CONTRA;
            case neg =>
              intros K
              have CONTRA : False := by simp[Index] at *; linarith
              simp at CONTRA
}