-- Blase/Blase/AutoStructs/Basic.lean

import Mathlib.Algebra.Group.Nat.Range

import Blase.AutoStructs.BundledNfa

import Blase.FinEnum

open SetRel

abbrev State := Nat

structure RawCNFA (A : Type 0) [BEq A] [Hashable A] [DecidableEq A] [FinEnum A] where
  stateMax : State
  initials : Std.HashSet State
  finals : Std.HashSet State
  trans : Std.HashMap (State × A) (Std.HashSet State)
deriving Repr

section sim

variable {A : Type} [BEq A] [Hashable A] [DecidableEq A] [FinEnum A]

@[inline]
def RawCNFA.tr (m : RawCNFA A) s a := m.trans.getD (s, a) ∅

def RawCNFA.states (m : RawCNFA A) : Finset State := Finset.range m.stateMax

instance RawCNFA.statesFinset (m : RawCNFA A) : Fintype m.states := (Finset.range m.stateMax).fintypeCoeSort

end sim

section basics

variable {A : Type} [BEq A] [Hashable A] [DecidableEq A] [FinEnum A]

def RawCNFA.newState (m : RawCNFA A) : State × RawCNFA A :=
  let old := m.stateMax
  let m := { m with stateMax := old + 1 }
  (old, m)

def RawCNFA.addTrans (m : RawCNFA A) (a : A) (s s' : State) : RawCNFA A :=
  let ns := m.trans.getD (s, a) ∅
  let ns := ns.insert s'
  { m with trans :=  m.trans.insert (s, a) ns }

def RawCNFA.addInitial (m : RawCNFA A) (s : State) : RawCNFA A :=
  { m with initials := m.initials.insert s }

def RawCNFA.createSink (m : RawCNFA A) : State × RawCNFA A :=
  let (s, m) := m.newState
  let m := m.addInitial s
  let m := FinEnum.toList (α := A).foldl (init := m) fun m a =>
    m.addTrans a s s
  (s, m)

structure RawCNFA.WF (m : RawCNFA A) where
  initials_lt : ∀ {s}, s ∈ m.initials → s ∈ m.states
  finals_lt : ∀ {s}, s ∈ m.finals → s ∈ m.states
  trans_src_lt : ∀ s_a ∈ m.trans, s_a.1 ∈ m.states
  trans_tgt_lt : s' ∈ m.tr s a → s' ∈ m.states

-- Repository definitions used:
content: ...

-- Repository lemmas used:
content: @[aesop 50% unsafe]
theorem Std.HashMap.mem_of_getElem? [BEq K] [LawfulBEq K] [Hashable K] [LawfulHashable K] {m : Std.HashMap K V} {k : K} :
    m[k]? = some v → k ∈ m

content: theorem Std.HashMap.mem_iff_getElem? [BEq K] [LawfulBEq K] [Hashable K] [LawfulHashable K] [Inhabited V] {m : Std.HashMap K V} {k : K} :
     k ∈ m ↔ ∃ v, m[k]? = some v

-- Local definitions used:
content: abbrev State := Nat

content: structure RawCNFA (A : Type 0) [BEq A] [Hashable A] [DecidableEq A] [FinEnum A] where
  stateMax : State
  initials : Std.HashSet State
  finals : Std.HashSet State
  trans : Std.HashMap (State × A) (Std.HashSet State)
deriving Repr

content: @[inline]
def RawCNFA.tr (m : RawCNFA A) s a := m.trans.getD (s, a) ∅

content: def RawCNFA.states (m : RawCNFA A) : Finset State := Finset.range m.stateMax

content: instance RawCNFA.statesFinset (m : RawCNFA A) : Fintype m.states := (Finset.range m.stateMax).fintypeCoeSort

content: def RawCNFA.newState (m : RawCNFA A) : State × RawCNFA A :=
  let old := m.stateMax
  let m := { m with stateMax := old + 1 }
  (old, m)

content: def RawCNFA.addTrans (m : RawCNFA A) (a : A) (s s' : State) : RawCNFA A :=
  let ns := m.trans.getD (s, a) ∅
  let ns := ns.insert s'
  { m with trans :=  m.trans.insert (s, a) ns }

content: def RawCNFA.addInitial (m : RawCNFA A) (s : State) : RawCNFA A :=
  { m with initials := m.initials.insert s }

content: def RawCNFA.createSink (m : RawCNFA A) : State × RawCNFA A :=
  let (s, m) := m.newState
  let m := m.addInitial s
  let m := FinEnum.toList (α := A).foldl (init := m) fun m a =>
    m.addTrans a s s
  (s, m)

content: structure RawCNFA.WF (m : RawCNFA A) where
  initials_lt : ∀ {s}, s ∈ m.initials → s ∈ m.states
  finals_lt : ∀ {s}, s ∈ m.finals → s ∈ m.states
  trans_src_lt : ∀ s_a ∈ m.trans, s_a.1 ∈ m.states
  trans_tgt_lt : s' ∈ m.tr s a → s' ∈ m.states

-- Local lemmas used:
content: @[grind =, simp]
lemma RawCNFA.addInitial_tr {m : RawCNFA A} : (m.addInitial s'').tr s a = m.tr s a

content: @[grind =, simp, aesop 50% unsafe]
lemma states_addInitial (m : RawCNFA A) (s' : State) :
    (m.addInitial s').states = m.states

content: @[grind =] -- TODO: should I?
lemma addTrans_tr (m : RawCNFA A) [LawfulBEq A] {a b : A} {s₁ s₁' s₂ : State} :
    (m.addTrans a s₁ s₂).tr s₁' b = if s₁ = s₁' ∧ a = b then (m.tr s₁ a).insert s₂ else m.tr s₁' b

content: @[grind =, simp]
lemma mem_addTrans_tr (m : RawCNFA A) [LawfulBEq A] (a : A) (s1 s2 : State) :
    s' ∈ (m.addTrans a s1 s2).tr s b ↔
      (s = s1 ∧ s' = s2 ∧ b = a) ∨ s' ∈ m.tr s b

content: @[grind =, simp]
lemma newState_tr {m : RawCNFA A} : m.newState.2.tr s a = m.tr s a

content: @[grind =, simp, aesop 50% unsafe]
lemma states_newState (m : RawCNFA A) :
    m.newState.2.states = m.states ∪ { m.stateMax }

content: @[grind ., simp, aesop 50% unsafe]
lemma RawCNFA.WF.trans_src_lt' {m : RawCNFA A} (hwf : m.WF) :
    ∀ {s a}, (s, a) ∈ m.trans → s ∈ m.states

content: @[grind ., simp, aesop 50% unsafe]
lemma RawCNFA.WF.trans_src_lt'' [LawfulBEq A] {m : RawCNFA A} (hwf : m.WF) :
    ∀ {s a s'}, s' ∈ m.tr s a → s ∈ m.states

-- Target theorem and proof:
@[grind =, simp]
lemma createSink_trans [LawfulBEq A] {m : RawCNFA A} (hwf : m.WF) :
    s₂ ∈ m.createSink.2.tr s₁ a ↔
      (s₁ = m.stateMax ∧ s₂ = m.stateMax) ∨
      (s₁ ∈ m.states ∧ s₂ ∈ m.states ∧ s₂ ∈ m.tr s₁ a)
:= by
  unfold RawCNFA.createSink
  simp
  let motive (mᵢ : RawCNFA A) (hwf : mᵢ.WF) (as : List A) :=
    mᵢ.states = m.states ∪ {m.stateMax} →
    ∀ a,
      s₂ ∈ (as.foldl (λ (m' : RawCNFA A) a => m'.addTrans a m.stateMax m.stateMax)
              (init := mᵢ)).tr s₁ a ↔
        (s₁ = m.stateMax ∧ s₂ = m.stateMax ∧ a ∈ as) ∨
        (s₂ ∈ mᵢ.tr s₁ a)
  suffices h : motive (m.newState.2.addInitial m.stateMax) (by simp_all) (FinEnum.toList A) by
    simp only [FinEnum.mem_toList, and_true, motive] at h
    rw [h] <;> simp only [RawCNFA.addInitial_tr, newState_tr,
                          states_addInitial, states_newState]
    constructor; on_goal 2 => tauto
    rintro (h | h); (exact .inl h); right
    use (RawCNFA.WF.trans_src_lt'' hwf h),
        (RawCNFA.WF.trans_tgt_lt hwf h),
        h
  generalize_proofs hwf'; revert hwf'
  generalize (m.newState.2.addInitial m.stateMax) = mi
  induction FinEnum.toList A generalizing mi
  case nil =>
    rintro hwf'; simp [motive]
  case cons a as ih =>
    simp [motive]
    rintro hwf' hstates b
    simp [motive] at ih
    rw [ih]
    · simp only [mem_addTrans_tr]; tauto
    · simp [hwf', hstates]
    · simp [hstates]