-- SSA/Projects/CIRCT/Stream/WeakBisim.lean

import SSA.Projects.CIRCT.Stream.Stream

import Mathlib.Logic.Function.Iterate

import Mathlib.Data.Stream.Init

namespace CIRCTStream

namespace Stream

namespace Bisim

scoped infix:50 " ~ " => Bisim

end Bisim

open Bisim

def stuck (α : Type) : Stream α := Stream'.const none

noncomputable def nonesUntilSome (x : Stream α) (not_stuck : x ≠ stuck α) : Nat :=
  prop.choose
where
  prop : ∃ i, (x.drop i).head.isSome ∧ ∀ j < i, x.get j = none := by admit /- proof elided -/

noncomputable def dropLeadingNones (x : Stream α) (not_stuck : x ≠ stuck α) : Stream α:=
  x.drop (nonesUntilSome x not_stuck)

open Classical in

noncomputable def removeNone (x : Stream α) : Stream α :=
  Stream.corec x fun x =>
    if h : x ≠ stuck α then
      let x := x.dropLeadingNones h
      (x.head, x.tail)
    else
      (none, x)

open Classical in

open Classical in

open Classical in

-- 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) )

content: def Stream (β : Type) := Stream' (Option β)

content: def head : Stream α → Option α := Stream'.head

content: def tail : Stream α → Stream α := Stream'.tail

content: def corec {α} {β} (s0 : β) (f : β → (Option α × β)) : Stream α :=
  Stream'.corec (f · |>.fst) (f · |>.snd) s0

-- 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: def stuck (α : Type) : Stream α := Stream'.const none

content: noncomputable def nonesUntilSome (x : Stream α) (not_stuck : x ≠ stuck α) : Nat :=
  prop.choose
where
  prop : ∃ i, (x.drop i).head.isSome ∧ ∀ j < i, x.get j = none := by admit /- proof elided -/

content: noncomputable def dropLeadingNones (x : Stream α) (not_stuck : x ≠ stuck α) : Stream α:=
  x.drop (nonesUntilSome x not_stuck)

content: noncomputable def removeNone (x : Stream α) : Stream α :=
  Stream.corec x fun x =>
    if h : x ≠ stuck α then
      let x := x.dropLeadingNones h
      (x.head, x.tail)
    else
      (none, x)

-- Local lemmas used:
content: theorem nonesUntilSome_spec (x : Stream α) (not_stuck : x ≠ stuck α) :
    (dropLeadingNones x not_stuck).head.isSome
    ∧ ∀ j < nonesUntilSome x not_stuck, x.get j = none

content: @[simp] theorem head_removeNone (x : Stream α) :
    x.removeNone.head =
      if h : x ≠ stuck α then
        (x.dropLeadingNones h).head
      else
        none

-- Target theorem and proof:
theorem removeNone_equiv (x : Stream α) :
    x.removeNone ~ x
:= by
  apply Bisim.coinduct (· = ·.removeNone)
  · rintro _ x rfl
    · use 0
      simp only [Nat.zero_add, Nat.not_lt_zero, false_implies, implies_true, true_and]
      by_cases x_eq_stuck : x = stuck α
      case pos =>
        subst x_eq_stuck
        refine ⟨0, ?_, ?_, by intros; contradiction⟩
        · show tail _ = removeNone (stuck α).tail; simp
        · show head _ = none; simp
      case neg =>
        have ⟨_, h2⟩ := nonesUntilSome_spec x x_eq_stuck
        refine ⟨nonesUntilSome x x_eq_stuck, ?_, ?_, h2⟩
        · show x.removeNone.tail = removeNone (Stream'.drop _ x)
          have (w : Nat) : x.drop (w + 1) = tail (x.drop w) := by simp [tail]
          rw [this]
          simp [x_eq_stuck]
          rfl
        · show x.removeNone.head = _
          simp only [head_removeNone, ne_eq, x_eq_stuck, dropLeadingNones]
          simp [head]
  · rfl