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

@[inline]
def product.prodArray' (a : Array γ) :=
  m₁.attachWith _ hm₁ |>.fold (init := a) fun is s1 =>
    m₂.attachWith _ hm₂ |>.fold (init := is) fun is s2 =>
      is.push (f s1 s2)

@[inline]
def product.prodArray := prodArray' f hm₁ hm₂ (Array.emptyWithCapacity <| m₁.size * m₂.size)

end generic_prod

section product

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

def product.inits (m₁ m₂ : CNFA n) :=
  product.prodArray Prod.mk @m₁.wf.initials_lt @m₂.wf.initials_lt

def product (final? : Bool → Bool → Bool) (m₁ m₂ : CNFA n) : CNFA n :=
  worklistRun (m₁.m.states × m₂.m.states) final (product.inits m₁ m₂)
    (by admit /- proof elided -/
    ) f
where final (ss : m₁.m.states × m₂.m.states) := final? (ss.1 ∈ m₁.m.finals) (ss.2 ∈ m₂.m.finals)
      f (ss : m₁.m.states × m₂.m.states) :=
        let (s1, s2) := ss
        (FinEnum.toList (α := BitVec n)).foldl (init := Array.empty) fun as a =>
          product.prodArray' (λ s₁ s₂ ↦ (a, (s₁, s₂)))
            (fun s' => m₁.wf.trans_tgt_lt (s := s1) (a := a)) (fun s' => m₂.wf.trans_tgt_lt (s := s2) (a := a)) as

noncomputable def to_prop (f : Bool → Bool → Bool) (p1 p2 : Prop) : Prop :=
  f (@Decidable.decide p1 (Classical.propDecidable _)) (@Decidable.decide p2 (Classical.propDecidable _))

-- Repository definitions used:
content: structure CNFA (n : Nat) where
  m : RawCNFA (BitVec n)
  wf : m.WF

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

content: abbrev State := Nat

content: def product (final? : Prop → Prop → Prop) (M N : NFA' n) : NFA' n where
  σ := _
  M := M.M.product final? N.M

content: structure NFA' (n : Nat) where
  σ : Type
  M : NFA (BitVec n) σ

content: def nfa : NFA A S where
  start := { sa | sa ∈ inits }
  accept := { sa | final sa }
  step sa a := { sa' | (a, sa') ∈ f sa }

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

axiom hashMap_missing : ∀ {P : Prop}, P

-- 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: @[simp]
lemma Std.HashSet.mem_toSet [BEq α] [Hashable α] (m : HashSet α) : x ∈ m.toSet ↔ x ∈ m

content: @[simp]
theorem Std.HashSet.mem_attachWith_mem [BEq α] [Hashable α] [LawfulBEq α] (m : HashSet α) {P H} (x : α) h :
    ⟨x, h⟩ ∈ m.attachWith P H ↔ x ∈ m

-- Local definitions used:
content: @[inline]
def product.prodArray' (a : Array γ) :=
  m₁.attachWith _ hm₁ |>.fold (init := a) fun is s1 =>
    m₂.attachWith _ hm₂ |>.fold (init := is) fun is s2 =>
      is.push (f s1 s2)

content: @[inline]
def product.prodArray := prodArray' f hm₁ hm₂ (Array.emptyWithCapacity <| m₁.size * m₂.size)

content: def product.inits (m₁ m₂ : CNFA n) :=
  product.prodArray Prod.mk @m₁.wf.initials_lt @m₂.wf.initials_lt

content: def product (final? : Bool → Bool → Bool) (m₁ m₂ : CNFA n) : CNFA n :=
  worklistRun (m₁.m.states × m₂.m.states) final (product.inits m₁ m₂)
    (by admit /- proof elided -/
    ) f
where final (ss : m₁.m.states × m₂.m.states) := final? (ss.1 ∈ m₁.m.finals) (ss.2 ∈ m₂.m.finals)
      f (ss : m₁.m.states × m₂.m.states) :=
        let (s1, s2) := ss
        (FinEnum.toList (α := BitVec n)).foldl (init := Array.empty) fun as a =>
          product.prodArray' (λ s₁ s₂ ↦ (a, (s₁, s₂)))
            (fun s' => m₁.wf.trans_tgt_lt (s := s1) (a := a)) (fun s' => m₂.wf.trans_tgt_lt (s := s2) (a := a)) as

content: noncomputable def to_prop (f : Bool → Bool → Bool) (p1 p2 : Prop) : Prop :=
  f (@Decidable.decide p1 (Classical.propDecidable _)) (@Decidable.decide p2 (Classical.propDecidable _))

-- Local lemmas used:
content: include hinj in
omit [BEq α] [Hashable α] [LawfulBEq α] in
lemma product.prodArray_spec_helper
    (is : Array γ) (hnd : is.toList.Nodup)
    (s : S₁) (hnew : ∀ s₂, f s s₂ ∉ is):
  let motive (a : Array γ) (S : Set S₂)  :=
    a.toList.Nodup ∧
    (∃ r, a.toList = is.toList ++ r ∧ (∀ z ∈ r, ∃ s₁ s₂, z = f s₁ s₂)) ∧
    ∀ s1 s2, f s1 s2 ∈ a ↔ s1 ≠ s ∧ f s1 s2 ∈ is ∨ s1 = s ∧ s2 ∈ S
  let body := m₂.attachWith _ hm₂ |>.fold (init := is) fun is s2 =>
      is.push (f s s2)
  motive body (m₂.attachWith _ hm₂).toSet

content: include hinj in
lemma product.prodArray'_spec_full {aᵢ : Array γ} (hnd: aᵢ.toList.Nodup) (hnin : ∀ s₁ s₂, f s₁ s₂ ∉ aᵢ) :
    (product.prodArray' f hm₁ hm₂ aᵢ).toList.Nodup ∧
    (∃ r, (product.prodArray' f hm₁ hm₂ aᵢ).toList = aᵢ.toList ++ r ∧ (∀ z ∈ r, ∃ s₁ s₂, z = f s₁ s₂)) ∧
    ∀ s₁ s₂, f s₁ s₂ ∈ product.prodArray' f hm₁ hm₂ aᵢ ↔ (s₁.val ∈ m₁ ∧ s₂.val ∈ m₂)

content: include hinj in
lemma product.prodArray_spec_full :
    (product.prodArray f hm₁ hm₂).toList.Nodup ∧
    ∀ s₁ s₂, f s₁ s₂ ∈ product.prodArray f hm₁ hm₂ ↔ (s₁.val ∈ m₁ ∧ s₂.val ∈ m₂)

content: include hinj in
@[simp]
lemma product.prodArray_spec :
    ∀ s₁ s₂, f s₁ s₂ ∈ product.prodArray f hm₁ hm₂ ↔ (s₁.val ∈ m₁ ∧ s₂.val ∈ m₂)

content: @[simp]
lemma product.inits_spec :
    ∀ s₁ s₂, (s₁, s₂) ∈ inits m₁ m₂ ↔ (s₁.val ∈ m₁.m.initials ∧ s₂.val ∈ m₂.m.initials)

content: lemma product.f_spec {m₁ m₂ : CNFA n} {s₁ : m₁.m.states} {s₂ : m₂.m.states} :
    ∀ a s₁' s₂',
      (a, (s₁', s₂')) ∈ f m₁ m₂ (s₁, s₂) ↔ s₁'.val ∈ m₁.m.tr s₁ a ∧ s₂'.val ∈ m₂.m.tr s₂ a

-- Target theorem and proof:
lemma product.sim {m1 m2 : CNFA n}:
    m1.Sim M1 → m2.Sim M2 →
    (nfa (product.inits m1 m2) (final final? m1 m2) (f m1 m2)).Bisim (M1.M.product (to_prop final?) M2.M)
:= by
  rintro ⟨R₁, hsim₁⟩ ⟨R₂, hsim₂⟩
  let R : SetRel (m1.m.states × m2.m.states) (M1.σ × M2.σ) :=
    {((s₁, s₂), (q₁, q₂)) | s₁.val ~[R₁] q₁ ∧ s₂.val ~[R₂] q₂ }
  use R; constructor
  · rintro ⟨s₁, s₂⟩ ⟨q₁, q₂⟩ ⟨hR₁, hR₂⟩
    simp [nfa, to_prop, final]
    rw [←hsim₁.accept hR₁, ←hsim₂.accept hR₂]; congr
  · constructor
    · rintro ⟨s₁, s₂⟩ hstart
      simp [nfa, inits_spec] at hstart; rcases hstart with ⟨h₁, h₂⟩
      obtain ⟨q₁, hq₁, hR₁⟩ := hsim₁.initial₁ h₁
      obtain ⟨q₂, hq₂, hR₂⟩ := hsim₂.initial₁ h₂
      use (q₁, q₂); simp [NFA.product, R, *]
    · rintro ⟨q₁, q₂⟩ ⟨hst₁, hst₂⟩
      apply hsim₁.initial₂ at hst₁; obtain ⟨s₁, hi₁, hR₁⟩ := hst₁
      apply hsim₂.initial₂ at hst₂; obtain ⟨s₂, hi₂, hR₂⟩ := hst₂
      simp only [nfa, Set.mem_setOf_eq, Prod.exists, inits_spec]
      have hin₁ : s₁ ∈ m1.m.states := by apply m1.wf.initials_lt hi₁
      have hin₂ : s₂ ∈ m2.m.states := by apply m2.wf.initials_lt hi₂
      use ⟨s₁, hin₁⟩, ⟨s₂, hin₂⟩
      simp_all [R]
  · rintro ⟨s₁, s₂⟩ ⟨q₁, q₂⟩ a ⟨s₁', s₂'⟩ ⟨hR₁, hR₂⟩ hst
    simp [nfa, f_spec] at hst; rcases hst with ⟨hst₁, hst₂⟩
    obtain ⟨q₁', hst₁, hR₁'⟩ := hsim₁.trans_match₁ hR₁ hst₁
    obtain ⟨q₂', hst₂, hR₂'⟩ := hsim₂.trans_match₁ hR₂ hst₂
    simp [NFA.product]
    use q₁', q₂'
    simp_all [R]
  · rintro ⟨s₁, s₂⟩ ⟨q₁, q₂⟩ a ⟨q₁', q₂'⟩ ⟨hR₁, hR₂⟩ hst
    simp [NFA.product] at hst; rcases hst with ⟨hst₁, hst₂⟩
    obtain ⟨s₁', hst₁, hR₁'⟩ := hsim₁.trans_match₂ hR₁ hst₁ (by simp) (by simp)
    obtain ⟨s₂', hst₂, hR₂'⟩ := hsim₂.trans_match₂ hR₂ hst₂ (by simp) (by simp)
    have hin₁ : s₁' ∈ m1.m.states := by apply m1.wf.trans_tgt_lt hst₁;
    have hin₂ : s₂' ∈ m2.m.states := by apply m2.wf.trans_tgt_lt hst₂
    simp only [nfa, Set.mem_setOf_eq, Prod.exists, f_spec]
    use ⟨s₁', hin₁⟩, ⟨s₂', hin₂⟩
    simp_all [R]