-- Blase/Blase/AutoStructs/Worklist.lean

import Blase.AutoStructs.Basic

open SetRel

section nfa

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

variable {S : Type} [Fintype S] [BEq S] [LawfulBEq S] [Hashable S] [DecidableEq S]

variable (inits : Array S) (final : S → Bool) (f : S → Array (A × S))

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

end nfa

section nfa'

variable {S : Type} [Fintype S] [BEq S] [LawfulBEq S] [Hashable S] [DecidableEq S]

variable (inits : Array S) (final : S → Bool) (f : S → Array (BitVec n × S))

end nfa'

section worklist

variable (A : Type) [BEq A] [LawfulBEq A] [Hashable A] [DecidableEq A] [FinEnum A]

variable (S : Type) [Fintype S] [BEq S] [LawfulBEq S] [Hashable S] [DecidableEq S]

structure worklist.St where
  m : RawCNFA A
  map : Std.HashMap S State := ∅
  worklist : Array S := ∅
  worklist_nodup : worklist.toList.Nodup
  worklist_incl : ∀ sa ∈ worklist, sa ∈ map

open List in

def worklist.St.addOrCreateState (st : worklist.St A S) (final? : Bool) (sa : S) : State × worklist.St A S :=
  match heq : st.map[sa]? with
  | some s => (s, st)
  | none =>
    let (s, m) := st.m.newState
    let m := if final? then m.addFinal s else m
    let map := st.map.insert sa s
    let worklist := st.worklist.push sa
    have worklist_nodup : worklist.toList.Nodup := by admit /- proof elided -/

def processOneElem (final : S → Bool) (s : State) (st : worklist.St A S) : A × S → worklist.St A S :=
  fun (a', sa') =>
    let (s', st') := st.addOrCreateState _ _ (final sa') sa'
    let m := st'.m.addTrans a' s s'
    { st' with m }

      open List in
      have hgrow : ∃ sas, st2.map.keys ~ (sas ++ st1.map.keys) ∧ st2.worklist.toList = st1.worklist.toList ++ sas := by
        rcases a with ⟨al⟩
        unfold st2
        generalize hst1 : st1 = x; clear hst1; revert x
        induction al with
        | nil => simp
        | cons asa al ih =>
          simp; simp at ih; intros st
          let wl' := processOneElem A S final s st asa
          rcases ih wl' with ⟨sas', h1', h2'⟩; clear ih
          rcases processOneElem_grow _ _ st final asa.1 asa.2 s with ⟨sas, h1, h2⟩
          use (sas ++ sas')
          constructor
          { simp [wl'] at h1'; apply list_perm_trick; exact h1'; exact h1 }
          { simp [wl', h2] at h2'; aesop }
      have hincl : ∀ k, k ∈ st1.map → k ∈ st2.map := by
        intros k; rcases hgrow with ⟨sas, hkeys, -⟩;
        have := @(List.perm_subset_iff_right hkeys st1.map.keys).mpr (by aesop) k; aesop
      have : st1.meas < st0.meas := by
        rcases heq' : sa? with ⟨⟩ | ⟨sa⟩
        { simp_all }
        apply Finset.card_lt_card
        simp [Finset.ssubset_iff, Finset.subset_iff]
        use sa
        simp [sa?] at heq'
        constructor
        { constructor
          { apply Array.mem_of_back? at heq'; apply st0.worklist_incl; assumption }
          { apply Array.not_elem_back_pop at heq' <;> simp_all +zetaDelta [Array.pop] } }
        constructor
        { right; apply Array.mem_of_back? at heq'; assumption }
        rintro sa hh; rcases hh with hnin | hin
        { simp +zetaDelta [hnin] }
        right
        exact Array.mem_of_mem_pop st0.worklist sa hin
      have : st2.meas ≤ st1.meas := by
        apply Finset.card_le_card
        simp +zetaDelta [Finset.subset_iff]
        intros sa' h
        rcases h with hnin | hin
        { left; simp [st1] at hincl; intros hc; apply hnin; apply hincl; assumption }
        by_cases hnew : sa' ∈ st0.map
        all_goals try (left; trivial)
        right
        simp [st1] at hgrow
        rcases hgrow with ⟨sas, hkeys2, hwl2⟩
        have hnin : sa'∉ sas := by
          intros hc
          have hdisj : st0.map.keys.Disjoint sas := by
            have : (sas ++ st0.map.keys).Nodup := by
              apply List.Perm.nodup
              assumption
              apply st2.map.keys_nodup
            simp [List.nodup_append_comm, List.disjoint_of_nodup_append, this]
          apply hdisj
          { simp_all [Std.HashMap.mem_keys]; apply hnew }
          { apply hc }
        rcases hin with ⟨hin⟩; simp_all +zetaDelta
      have : st2.meas < st0.meas := by omega
      go st2
    else
      st0.m -- never happens
  | none => st0.m -- never happens
  termination_by st0.meas

def worklist.St.visited (st : worklist.St A S) : Set S := { s : S | s ∈ st.map ∧ s ∉ st.worklist }

structure StInv (m : RawCNFA A) (map : Std.HashMap S State) where
  wf : m.WF
  map_states : ∀ (sa : S) s, map[sa]? = some s → s ∈ m.states
  map_surj : ∀ s : m.states, ∃ (sa : S), map[sa]? = some s.val
  map_inj : ∀ {s} {sa sa' : S}, map[sa]? = some s → map[sa']? = some s → sa = sa'

end worklist

section worklist_correct

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

variable {S : Type} [Fintype S] [BEq S] [LawfulBEq S] [Hashable S] [DecidableEq S]

variable (inits : Array S) (final : S → Bool) (f : S → Array (A × S))

def worklist.St.rel (st : worklist.St A S) : SetRel State S := {(s, sa) | st.map[sa]? = some s }

def worklist.St.D (st : worklist.St A S) : Set S := st.visited

abbrev worklist.St.sim {st : worklist.St A S} (T : Set (S × A × S)) :=
  st.m.Simul (nfa inits final f) st.rel st.D T

def processOneElem_mot (s : State) (sa : S) (n : ℕ) (st : worklist.St A S) : Prop :=
  st.map[sa]? = some s ∧
  sa ∈ st.visited ∧
  StInv A S st.m st.map ∧
  st.sim inits final f  {(sa1, a, sa') | sa1 = sa ∧ ∃ k ≥ n, (f sa)[k]? = some (a, sa') }

def processOneElem_inv {st : worklist.St A S} (s : State) (sa : S) (k : ℕ) :
    ∀ a sa', (f sa)[k]? = some (a, sa') →
    processOneElem_mot inits final f s sa k st →
    let st' := processOneElem A S final s st (a, sa')
    StInv A S st'.m st'.map :=

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

content: abbrev State := Nat

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

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: def RawCNFA.addFinal (m : RawCNFA A) (s : State) : RawCNFA A :=
  { m with finals := m.finals.insert s }

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 }

-- Repository lemmas used:
content: @[aesop 50% unsafe]
theorem Std.HashMap.get?_none_not_mem [BEq K] [LawfulBEq K] [Hashable K] [LawfulHashable K] {m : Std.HashMap K V} {k : K} : m.get? k = none → k ∉ m

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

content: @[grind =, simp]
lemma addTrans_tr_eq (m : RawCNFA A) [LawfulBEq A] (a : A) (s₁ s₂ : State) :
    (m.addTrans a s₁ s₂).tr s₁ a = (m.tr s₁ a).insert s₂

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

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: @[aesop 50% unsafe]
theorem Std.HashMap.getElem?_none_not_mem [BEq K] [LawfulBEq K] [Hashable K] [LawfulHashable K] {m : Std.HashMap K V} {k : K} :
    m[k]? = none → k ∉ m

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

content: structure worklist.St where
  m : RawCNFA A
  map : Std.HashMap S State := ∅
  worklist : Array S := ∅
  worklist_nodup : worklist.toList.Nodup
  worklist_incl : ∀ sa ∈ worklist, sa ∈ map

content: def worklist.St.addOrCreateState (st : worklist.St A S) (final? : Bool) (sa : S) : State × worklist.St A S :=
  match heq : st.map[sa]? with
  | some s => (s, st)
  | none =>
    let (s, m) := st.m.newState
    let m := if final? then m.addFinal s else m
    let map := st.map.insert sa s
    let worklist := st.worklist.push sa
    have worklist_nodup : worklist.toList.Nodup := by admit /- proof elided -/

content: def processOneElem (final : S → Bool) (s : State) (st : worklist.St A S) : A × S → worklist.St A S :=
  fun (a', sa') =>
    let (s', st') := st.addOrCreateState _ _ (final sa') sa'
    let m := st'.m.addTrans a' s s'
    { st' with m }

content: def worklist.St.visited (st : worklist.St A S) : Set S := { s : S | s ∈ st.map ∧ s ∉ st.worklist }

content: structure StInv (m : RawCNFA A) (map : Std.HashMap S State) where
  wf : m.WF
  map_states : ∀ (sa : S) s, map[sa]? = some s → s ∈ m.states
  map_surj : ∀ s : m.states, ∃ (sa : S), map[sa]? = some s.val
  map_inj : ∀ {s} {sa sa' : S}, map[sa]? = some s → map[sa']? = some s → sa = sa'

content: def worklist.St.rel (st : worklist.St A S) : SetRel State S := {(s, sa) | st.map[sa]? = some s }

content: def worklist.St.D (st : worklist.St A S) : Set S := st.visited

content: abbrev worklist.St.sim {st : worklist.St A S} (T : Set (S × A × S)) :=
  st.m.Simul (nfa inits final f) st.rel st.D T

content: def processOneElem_mot (s : State) (sa : S) (n : ℕ) (st : worklist.St A S) : Prop :=
  st.map[sa]? = some s ∧
  sa ∈ st.visited ∧
  StInv A S st.m st.map ∧
  st.sim inits final f  {(sa1, a, sa') | sa1 = sa ∧ ∃ k ≥ n, (f sa)[k]? = some (a, sa') }

content: def processOneElem_inv {st : worklist.St A S} (s : State) (sa : S) (k : ℕ) :
    ∀ a sa', (f sa)[k]? = some (a, sa') →
    processOneElem_mot inits final f s sa k st →
    let st' := processOneElem A S final s st (a, sa')
    StInv A S st'.m st'.map :=

-- Local lemmas used:
content: omit [Fintype S] [LawfulBEq A] in
lemma addOrCreate_preserves_map (st : worklist.St A S) (final? : Bool) (sa sa' : S) :
    let (_, st') := st.addOrCreateState _ _  final? sa'
    st.map[sa]? = some s →
    st'.map[sa]? = some s

content: omit [Fintype S] [LawfulBEq A] in
lemma processOneElem_preserves_map (st : worklist.St A S) (final : S → Bool) (a : A) (sa sa' : S) (s s' : State) :
    let st' := processOneElem _ _  final s st (a, sa')
    st.map[sa]? = some s' →
    st'.map[sa]? = some s'

content: omit [LawfulBEq A] [Fintype S] [DecidableEq S] in
lemma addOrCreateElem_visited final? (st : worklist.St A S) sa :
    st.addOrCreateState _ _ final? sa |>.2.visited = st.visited

content: omit [LawfulBEq A] [Fintype S] [DecidableEq S] in
lemma processOneElem_visited (st : worklist.St A S) :
    let st' := processOneElem _ _  final s st (a, sa')
    st'.visited = st.visited

content: omit [LawfulBEq A] [Fintype S] in
lemma processOneElem_map (st : worklist.St A S) (final : S → Bool) (a : A) (sa sa' : S) (s : State) :
  (processOneElem A S final s st (a, sa)).map[sa']? =
    match st.map[sa']? with
    | some s => some s
    | none => if sa = sa' then some st.m.stateMax else none

content: omit [LawfulBEq A] [Fintype S] [DecidableEq S] in
@[simp]
lemma processOneElem_initials (st : worklist.St A S) (final : S → Bool) (a : A) (sa : S) (s : State) :
    (processOneElem A S final s st (a, sa)).m.initials = st.m.initials

content: omit [LawfulBEq A] [Fintype S] [DecidableEq S] in
lemma processOneElem_finals (st : worklist.St A S) (final : S → Bool) (a : A) (sa : S) (s : State) :
    (processOneElem A S final s st (a, sa)).m.finals =
      if sa ∉ st.map ∧ final sa then st.m.finals.insert st.m.stateMax else st.m.finals

content: omit [Fintype S] [DecidableEq S] in
lemma processOneElem_tr (st : worklist.St A S) (final : S → Bool) (a b : A) (sa : S) (s s' : State) :
    if a = b ∧ s = s' then
      ∃ ssa, (processOneElem A S final s st (a, sa)).map[sa]? = some ssa ∧
        (processOneElem A S final s st (a, sa)).m.tr s' b =
        (st.m.tr s a |>.insert ssa)
    else
      (processOneElem A S final s st (a, sa)).m.tr s' b = st.m.tr s' b

content: omit [LawfulBEq A] [Fintype S] in
lemma processOneElem_rel {s₁ s₂ : State} :
    s₂ ~[(processOneElem A S final s₁ st (a, sa)).rel] sa' ↔
      (s₂ ~[st.rel] sa' ∨ (s₂ = st.m.stateMax ∧ sa' = sa ∧ st.map[sa']? = none))

content: omit [LawfulBEq A] [Fintype S] [LawfulBEq S] [DecidableEq S] in
lemma rel_in_states {st : worklist.St A S} (hinv : StInv A S st.m st.map) :
    s ~[st.rel] sa → s ∈ st.m.states

content: omit [LawfulBEq A] [Fintype S] in
lemma processOneElem_rel_preserve :
    s₂ ~[st.rel] sa' →
    s₂ ~[(processOneElem A S final s₁ st (a, sa)).rel] sa'

content: omit [LawfulBEq A] [Fintype S] in
lemma processOneElem_rel_preserve_olds :
    s₂ ~[(processOneElem A S final s₁ st (a, sa)).rel] sa' →
    s₂ ∈ st.m.states → s₂ ~[st.rel] sa'

-- Target theorem and proof:
omit [Fintype S] in
lemma processOneElem_spec {st : worklist.St A S} (s : State) (sa : S) (k : ℕ) :
    ∀ a sa', (f sa)[k]? = some (a, sa') →
    processOneElem_mot inits final f s sa k st →
    processOneElem_mot inits final f s sa (k+1) (processOneElem A S final s st (a, sa'))
:= by
  intro a sa' hf ⟨hmap, hvisited, inv, hsim⟩
  have hmem : ∀ s (sa : S), st.map[sa]? = some s → s ∈ st.m.states := by intros; apply inv.map_states; assumption
  have hwf : st.m.WF := by apply inv.wf
  have inv' := processOneElem_inv inits final f s sa k a sa' hf ⟨hmap, hvisited, inv, hsim⟩
  unfold processOneElem_mot
  constructor
  (rw [processOneElem_preserves_map]; assumption)
  constructor
  (rw [processOneElem_visited]; exact hvisited)
  use inv'; constructor
  { rw [processOneElem_finals]
    rintro s' q hR
    have hs' : s ∈ st.m.states := by apply hmem <;> assumption
    rw [processOneElem_rel] at hR
    rcases hR with hR | ⟨rfl, rfl, heq⟩
    · have heq := rel_in_states inv hR
      split_ifs with hcond
      · have hneq : st.m.stateMax ≠ s' := by rintro rfl; simp [RawCNFA.states] at heq
        simp [hneq]
        apply hsim.accept; assumption
      · apply hsim.accept; assumption
    · split_ifs with h
      · rcases h with ⟨_, hfin⟩
        simp [nfa, hfin]
      · simp [nfa]
        suffices hnin : q ∉ st.map by
          push_neg at h; specialize h hnin; simp_all
          rintro hc; apply hwf.finals_lt at hc; simp at hc
        exact Std.HashMap.getElem?_none_not_mem heq }
  { rintro s₁ hs₁; rw [processOneElem_initials] at hs₁
    obtain ⟨q, hq, hR⟩ := hsim.initial₁ hs₁
    use q, hq, processOneElem_rel_preserve final hR }
  { intros q hq; obtain ⟨s, hs, hR⟩ := hsim.initial₂ hq
    simp only [processOneElem_initials]
    use s, hs, (by exact processOneElem_rel_preserve final hR) }
  { rintro s₁ s₂ b q₁ hR htr
    have h := processOneElem_tr st final a b sa' s s₁
    split_ifs at h with hcond
    on_goal 2 => {
      rw [h] at htr
      apply processOneElem_rel_preserve_olds at hR
      specialize hR (RawCNFA.WF.trans_src_lt'' hwf htr)
      obtain ⟨q₂, hst, hrel⟩ := hsim.trans_match₁ hR htr
      use q₂; simp only [hst, true_and]
      exact processOneElem_rel_preserve final hrel }
    rcases hcond with ⟨rfl, rfl⟩
    rcases h with ⟨sₙ, hmap', htr'⟩
    rw [htr'] at htr; clear htr'
    simp only [Std.HashSet.mem_insert, beq_iff_eq] at htr
    rcases htr with rfl | htr
    on_goal 2 =>
      have hold := RawCNFA.WF.trans_src_lt'' hwf htr
      obtain ⟨q₂, hst, hrel⟩ := hsim.trans_match₁
        (by apply processOneElem_rel_preserve_olds final hR hold) htr
      use q₂, hst, processOneElem_rel_preserve final hrel
    use sa'; constructor
    · suffices heq : q₁ = sa by
        subst heq; apply Array.mem_of_getElem? hf
      apply processOneElem_preserves_map at hmap
      unfold worklist.St.rel at hR
      apply inv'.map_inj hR hmap
    · exact hmap' }
  { rintro s₁ b q₁ q₂ hR hs hD hnT
    simp only [ge_iff_le, Prod.mk.eta, Set.mem_setOf_eq, not_and, not_exists] at hnT
    have h := processOneElem_tr st final a b sa' s s₁
    split_ifs at h with hcond
    on_goal 2 =>
      have hR' : s₁ ~[st.rel] q₁ := by
        rw [processOneElem_rel] at hR; rcases hR with hR' | ⟨_, _, hnone⟩; exact hR'
        unfold worklist.St.D at hD; rw [processOneElem_visited] at hD
        rcases hD with ⟨hnin, -⟩
        apply Std.HashMap.getElem?_none_not_mem at hnone
        contradiction
      obtain ⟨s₂, hs', hR⟩ := hsim.trans_match₂ hR' hs
        (by simp_all [worklist.St.D, processOneElem_visited])
        (by simp; rintro rfl i hi hc
            by_cases heq: k = i
            · subst heq; apply hcond; constructor
              · simp [hc] at hf; simp [hf]
              · rw [hR'] at hmap; simp at hmap; exact hmap.symm
            · apply hnT rfl _ (by omega) hc)
      use s₂; simp only [h, hs', true_and]
      apply processOneElem_rel_preserve; assumption
    rcases hcond with ⟨rfl, rfl⟩
    rcases h with ⟨sₙ, hmap', htr'⟩
    obtain rfl : sa = q₁ := by
      apply processOneElem_preserves_map at hmap
      apply inv'.map_inj hmap hR
    simp only [htr', Std.HashSet.mem_insert, beq_iff_eq]
    by_cases heq : sa' = q₂
    · subst heq; use sₙ; simp only [true_or, true_and]; exact hmap'
    · have hold := hmem _ _ hmap
      apply processOneElem_rel_preserve_olds at hR
      obtain ⟨s₂, hs', hR⟩ := hsim.trans_match₂ (hR hold) hs
        (by simp_all [worklist.St.D, processOneElem_visited])
        (by simp; rintro i hi hc
            have _ : i ≠ k := by rintro rfl; simp_all
            apply hnT rfl i (by omega) hc)
      use s₂; simp only [hs', or_true, true_and]
      apply processOneElem_rel_preserve; assumption }