-- Blase/Blase/AutoStructs/Constructions.lean

import Mathlib.Data.Fintype.Prod

import Blase.AutoStructs.Worklist

import Mathlib.Tactic.ApplyFun

open SetRel

section sink

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

end sink

section generic_prod

variable {α} [BEq α] [Hashable α] [LawfulBEq α]

variable {β} [BEq β] [Hashable β] [LawfulBEq β]

variable {S₁ : Finset α} {S₂ : Finset β}

variable {γ} (f : S₁ → S₂ → γ) (hinj : Function.Injective2 f)

variable {m₁ : Std.HashSet α} (hm₁ : ∀ s₁ ∈ m₁, s₁ ∈ S₁)

variable {m₂ : Std.HashSet β} (hm₂ : ∀ s₂ ∈ m₂, s₂ ∈ S₂)

end generic_prod

section product

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

end product

section determinization

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

def RawCNFA.reverse (m : RawCNFA A) : RawCNFA A :=
  let m' := { stateMax := m.stateMax, trans := Std.HashMap.emptyWithCapacity m.trans.size, initials := m.finals, finals := m.initials}
  m.trans.fold (init := m') processState
where
  processState := fun m' (s, a) ss' =>
    ss'.fold (init := m') fun m' s' => m'.addTrans a s' s

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

content: abbrev State := Nat

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

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

content: def Std.HashSet.toSet [BEq α] [Hashable α] (m : HashSet α) : Set α := { x | x ∈ m }

content: def Std.HashMap.toPFun [BEq α] [Hashable α] (m : HashMap α β) (x : α) : Option β := m[x]?

-- Repository lemmas used:
content: theorem Std.HashSet.fold_induction [BEq α] [LawfulBEq α] [Hashable α]
  {f : β → α → β} {m : HashSet α} {motive : β → Set α → Prop} :
    motive b ∅ →
    (∀ b x s, x ∉ s → motive b s → motive (f b x) (s ∪ {x})) →
    motive (m.fold f b) m.toSet

content: theorem Std.HashSet.toSet_toList[BEq α] [LawfulBEq α] [Hashable α] (m : HashSet α) : m.toSet = { x | x ∈ m.toList }

content: theorem Std.HashMap.fold_induction [BEq α] [LawfulBEq α] [DecidableEq α] [Hashable α]
  {f : γ → α → β → γ} {m : HashMap α β} {motive : γ → (α → Option β) → Prop} :
    motive b (λ _ ↦ none) →
    (∀ b x y m, m x = none → motive b m → motive (f b x y) (Function.update m x y)) →
    motive (m.fold f b) m.toPFun

content: theorem Std.HashMap.toPFun_toList[BEq α] [LawfulBEq α] [Hashable α] (m : HashMap α β) :
    m.toPFun = λ k ↦ m.toList.find? (λ x ↦ x.1 == k) |>.map Prod.snd

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: def RawCNFA.reverse (m : RawCNFA A) : RawCNFA A :=
  let m' := { stateMax := m.stateMax, trans := Std.HashMap.emptyWithCapacity m.trans.size, initials := m.finals, finals := m.initials}
  m.trans.fold (init := m') processState
where
  processState := fun m' (s, a) ss' =>
    ss'.fold (init := m') fun m' s' => m'.addTrans a s' s

-- Local lemmas used:
content: lemma RawCNFA.reverse_spec_procesState {m : RawCNFA A} (hwf : m.WF) s₀ a₀ ss' (hs₀ : s₀ ∈ m.states) :
    let motive m' ss' :=
      (∀ s ∈ ss', s ∈ m.states) →
        m'.WF ∧ m'.stateMax = m.stateMax ∧ m'.initials = m.initials ∧ m'.finals = m.finals ∧
          ∀ s a s', s' ∈ m'.tr s a ↔ (s' ∈ m.tr s a ∨ s' = s₀ ∧ a = a₀ ∧ s ∈ ss')
    motive (RawCNFA.reverse.processState m (s₀, a₀) ss') ss'.toSet

-- Target theorem and proof:
lemma RawCNFA.reverse_spec {m : RawCNFA A} (hwf : m.WF) :
    let m' := m.reverse
    m'.WF ∧ m'.stateMax = m.stateMax ∧ m'.initials = m.finals ∧ m'.finals = m.initials ∧
      ∀ s a s', s' ∈ m'.tr s a ↔ s ∈ m.tr s' a
:= m.reverse
    m'.WF ∧ m'.stateMax = m.stateMax ∧ m'.initials = m.finals ∧ m'.finals = m.initials ∧
      ∀ s a s', s' ∈ m'.tr s a ↔ s ∈ m.tr s' a := by
  let motive (m' : RawCNFA A) (trs : (State × A) → Option (Std.HashSet State)) :=
    (∀ s a ss', trs (s, a) = some ss' → s ∈ m.states ∧ ∀ s' ∈ ss', s' ∈ m.states) →
    m'.WF ∧ m'.stateMax = m.stateMax ∧ m'.initials = m.finals ∧ m'.finals = m.initials ∧
      ∀ s a s', s' ∈ m'.tr s a ↔ ∃ ss', trs (s', a) = some ss' ∧ s ∈ ss' -- s' -a-> s is a transition
  suffices h : motive (m.reverse) (m.trans.toPFun) by
    specialize h (by rintro s a ss' heq; exact ⟨hwf.trans_src_lt _ (Std.HashMap.mem_of_getElem? heq), fun s' hin => hwf.trans_tgt_lt_internal heq hin⟩)
    simp_all

  apply Std.HashMap.fold_induction
  · simp [motive, tr]
    constructor <;> simp_all [tr]
  · rintro m' ⟨s, a⟩ ss' trs hnew ih
    rintro hsts
    specialize ih (by rintro s₀ a₀ ss₀ heq; apply hsts _ a₀; simp [Function.update_apply]; split <;> simp_all)
    rcases ih with ⟨hwf', hm', his', hfs', htrs'⟩
    specialize hsts s a ss' (by simp)
    have hs : s ∈ m'.states := by simp_all [states]
    have hss' : ∀ s ∈ ss'.toSet, s ∈ m'.states := by simp_all [states]
    obtain ⟨hwf'', hm'', his'', hfs'', htrs''⟩ := reverse_spec_procesState hwf' s a ss' hs hss'
    simp_all only [true_and]
    rintro s₁ b s₂
    by_cases hcond : s₂ = s ∧ b = a
    · rcases hcond with ⟨rfl, rfl, hin⟩
      simp_all
    · simp_all