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

end generic_prod

section product

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

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

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

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

-- Target theorem and proof:
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
:= by
  suffices heq :
      ∀ as a (hnd : as.toList.Nodup) (hnd' : (FinEnum.toList (α := BitVec n)).Nodup)
       (hnew : ∀ b s₁ s₂, (b, s₁, s₂) ∈ as → b ∉ (FinEnum.toList (α := BitVec n))) s₁' s₂',
        (a, (s₁', s₂')) ∈
          ((FinEnum.toList (α := BitVec n)).foldl (init := as) fun as a =>
          product.prodArray' (λ s₁' s₂' ↦ (a, (s₁', s₂')))
            (fun _ => m₁.wf.trans_tgt_lt (s := s₁) (a := a)) (fun _ => m₂.wf.trans_tgt_lt (s := s₂) (a := a)) as) ↔
        (a, (s₁', s₂')) ∈ as ∨ a ∈ (FinEnum.toList (α := BitVec n)) ∧ s₁'.val ∈ m₁.m.tr s₁ a ∧ s₂'.val ∈ m₂.m.tr s₂ a by
    rintro a s₁' s₂'; rw [f, heq]
    · simp; rintro h; apply Array.not_mem_empty at h; trivial
    · exact List.dedup_eq_self.mp rfl
    · exact FinEnum.nodup_toList
    · rintro _ _ _ h; apply Array.not_mem_empty at h; trivial
  induction (FinEnum.toList (α := BitVec n))
  case nil =>
    simp only [List.foldl_nil, List.not_mem_nil, false_and, Subtype.forall]
    intros _ _ _ _ _ _; tauto
  case cons a as ih =>
    rintro bs b hnd hnd' hnew s₁' s₂'
    obtain ⟨hnd'', ⟨r, hr, hf⟩, hin⟩ := prodArray'_spec_full (fun s₁' s₂' => (a, s₁', s₂'))
      (by rintro _ _ _ _ ⟨rfl, rfl, rfl⟩; tauto)
      (fun _ => m₁.wf.trans_tgt_lt (s := s₁)) (fun _ => m₂.wf.trans_tgt_lt (s := s₂) (a := a))
      hnd
      (by dsimp; rintro s₁ s₂ hin; apply hnew at hin; simp at hin)
    dsimp
    have hmem : ∀ b s₁' s₂', b ≠ a → (b, s₁', s₂') ∈ prodArray' (fun s₁' s₂' => (a, s₁', s₂'))
        (fun _ => m₁.wf.trans_tgt_lt (s := s₁) (a := a))
        (fun _ => m₂.wf.trans_tgt_lt (s := s₂) (a := a)) bs →
          (b, s₁', s₂') ∈ bs := by
      rintro b s₁' s₂' hneq hin
      rw [Array.mem_def] at hin ⊢; rw [hr] at hin; simp only [List.mem_append] at hin
      rcases hin with hin | hin; assumption
      obtain ⟨_, _, ⟨rfl, -, -⟩⟩ := hf _ hin; simp at hneq
    rw [ih] <;> clear ih; rotate_left
    · apply hnd''
    · exact (List.nodup_cons.mp hnd').2
    · rintro b s₁' s₂' hin; by_cases heq : b = a
      · subst heq; apply List.Nodup.notMem hnd'
      · specialize hnew b s₁' s₂' (hmem b s₁' s₂' heq hin); simp at hnew; simp [hnew]
    by_cases heq :  b = a
    · subst heq
      simp [hin]; clear hin; constructor
      · rintro (⟨h1, h2⟩ | ⟨_, h1, h2⟩) <;> right <;> simp [h1, h2]
      · rintro (hc | ⟨h1, h2⟩)
        · apply hnew at hc; simp at hc
        · left; simp [h1, h2]
    · simp only [List.mem_cons, heq, false_or]; constructor
      · rintro (hin | hin)
        · left; apply hmem _ _ _ heq hin
        · right; exact hin
      · rintro (hin | hin)
        · left; rw [Array.mem_def] at hin ⊢; rw [hr]; simp [hin]
        · right; exact hin