-- Blase/Blase/AutoStructs/FormulaToAuto.lean

import Batteries.Data.Fin.Basic

import Batteries.Data.Fin.Lemmas

import Blase.SingleWidth.Defs

import Blase.AutoStructs.Constructions

import Blase.AutoStructs.Defs

import Blase.AutoStructs.FiniteStateMachine

import Mathlib.Tactic.Ring

import Mathlib.Data.Nat.Size -- TODO: remove and get rid of shiftLeft_eq_mul_pow use

import Mathlib.Data.BitVec

import Mathlib.Tactic.FinCases

open Fin.NatCast

open Mathlib

section fsm

variable {arity : Type} [FinEnum arity]

open BitStream in

end fsm

section nfas_relations

def WidthPredicate.final? (wp : WidthPredicate) (n : Nat) (s : State) : Bool :=
  decide (wp.sat s n)

def RawCNFA.autWidth (wp : WidthPredicate) (n : Nat) : RawCNFA (BitVec 0) :=
  let m := (n+2).iterate f empty
  let m := m.addInitial 0
  m.addTrans (BitVec.zero 0) (n + 1) (n + 1)
where
  f m :=
    let (s, m) := m.newState
    let m := if wp.final? n s then m.addFinal s else m
    if s > 0 then m.addTrans (BitVec.zero 0) (s-1) s else m

def CNFA.autWidth (wp : WidthPredicate) (n : Nat) : CNFA 0 :=
  ⟨RawCNFA.autWidth wp n, RawCNFA.autWidth_wf⟩

def NFA.autWidth (wp : WidthPredicate) (n : Nat) : NFA (BitVec 0) (Fin (n+2)) where
  start := { 0 }
  accept := { s | wp.final? n s }
  step s₁ _ := { s₂ | if s₁ = Fin.last (n+1) then s₁ = s₂ else s₂ = s₁ + 1 }

def NFA'.autWidth (wp : WidthPredicate) (n : Nat) : NFA' 0 := ⟨_, NFA.autWidth wp n⟩

def autWidth_equiv : (CNFA.autWidth wp n).m.states ≃ (NFA'.autWidth wp n).σ where
  toFun := fun ⟨s, hs⟩ =>
    Fin.mk s (by admit /- proof elided -/
    )
  invFun q := ⟨q.val, by admit /- proof elided -/
  ⟩
  left_inv := by admit /- proof elided -/

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

content: noncomputable def f : (ZMod q)[X] := X^(2^n) + 1

content: def RawCNFA.empty : RawCNFA A := {
  stateMax := 0
  initials := ∅
  finals := ∅
  trans := ∅
}

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

content: inductive WidthPredicate
| eq
| neq
| lt
| le
| gt
| ge
deriving Repr, Inhabited, Lean.ToExpr

content: instance instFinEnumBV : FinEnum (BitVec w) where
  card := 2^w
  equiv := {
    toFun := fun x => x.toFin
    invFun := fun x => BitVec.ofFin x
    left_inv := by admit /- proof elided -/

content: def CNFA.Sim (m : CNFA n) (M : NFA' n) :=
  m.m.Sim M.M

content: def RawCNFA.Sim (m : RawCNFA A) (A : NFA A S) := ∃ R, RawCNFA.Simul m A R ⊤ ∅

content: structure RawCNFA.Simul (m : RawCNFA A) (M : NFA A Q) (R : SetRel State Q) (D : Set Q) (T : Set (Q × A × Q)) where
  accept {s q} : s ~[R] q → (s ∈ m.finals ↔ q ∈ M.accept)
  initial₁ {s} : s ∈ m.initials → ∃ q ∈ M.start, s ~[R] q
  initial₂ {q} : q ∈ M.start → ∃ s ∈ m.initials, s ~[R] q
  trans_match₁ {s s' a q} : s ~[R] q → s' ∈ m.tr s a → ∃ q', q' ∈ M.step q a ∧ s' ~[R] q'
  trans_match₂ {s a q q'} : s ~[R] q → q' ∈ M.step q a → q ∈ D → (q, a, q') ∉ T → ∃ s', s' ∈ m.tr s a ∧ s' ~[R] q'

content: structure RawCNFA.SimulFun (m : RawCNFA A) (M : NFA A Q) (f : m.states ≃ Q)  where
  accept {q} : ((f.invFun q).val ∈ m.finals ↔ q ∈ M.accept)
  initial {q} : q ∈ M.start ↔ (f.invFun q).val ∈ m.initials
  trans_match {a q q'} : q' ∈ M.step q a ↔ (f.invFun q').val ∈ m.tr (f.invFun q) a

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

content: abbrev R := (ZMod q)[X] ⧸ (Ideal.span {f q n})

content: structure CNFA (n : Nat) where
  m : RawCNFA (BitVec n)
  wf : m.WF

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

content: abbrev zero   : BitStream := fun _ => false

content: def BitStream : Type := Nat → Bool

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

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: 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.newState (m : RawCNFA A) : State × RawCNFA A :=
  let old := m.stateMax
  let m := { m with stateMax := old + 1 }
  (old, m)

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

-- Repository lemmas used:
content: @[grind ., simp, aesop 50% unsafe]
lemma wf_addTrans [LawfulBEq A] (m : RawCNFA A) (hwf : m.WF) s a s' (hin : s ∈ m.states) (hin' : s' ∈ m.states) :
    (m.addTrans a s s').WF

content: @[grind =, simp]
lemma RawCNFA.same_stateMax (m : RawCNFA A) x y (z : Std.HashMap (State × A) (Std.HashSet State)) :
    (RawCNFA.mk m.stateMax x y z).states = m.states

content: @[grind =, simp, aesop 50% unsafe]
lemma newState_eq (m : RawCNFA A) :
    m.newState.1 = m.stateMax

content: @[grind =, simp]
lemma addInitial_stateMax {m : RawCNFA A} : (m.addInitial s).stateMax = m.stateMax

content: @[grind =, simp]
lemma addTrans_stateMax {m : RawCNFA A} : (m.addTrans a s s').stateMax = m.stateMax

content: lemma simulFun_sim {m : CNFA n} f :
    m.m.SimulFun M.M f → m.Sim M

content: lemma simulFun_sim_raw [LawfulBEq A] {m : RawCNFA A} (hwf : m.WF) f :
    m.SimulFun M f → m.Sim M

content: @[simp]
lemma RawCNFA.Simul.initial {m : RawCNFA A} {M : NFA A Q} (hsim : m.Simul M R ⊤ ∅) :
    R.set_eq m.initials.toSet M.start

-- Local definitions used:
content: def WidthPredicate.final? (wp : WidthPredicate) (n : Nat) (s : State) : Bool :=
  decide (wp.sat s n)

content: def RawCNFA.autWidth (wp : WidthPredicate) (n : Nat) : RawCNFA (BitVec 0) :=
  let m := (n+2).iterate f empty
  let m := m.addInitial 0
  m.addTrans (BitVec.zero 0) (n + 1) (n + 1)
where
  f m :=
    let (s, m) := m.newState
    let m := if wp.final? n s then m.addFinal s else m
    if s > 0 then m.addTrans (BitVec.zero 0) (s-1) s else m

content: def CNFA.autWidth (wp : WidthPredicate) (n : Nat) : CNFA 0 :=
  ⟨RawCNFA.autWidth wp n, RawCNFA.autWidth_wf⟩

content: def NFA.autWidth (wp : WidthPredicate) (n : Nat) : NFA (BitVec 0) (Fin (n+2)) where
  start := { 0 }
  accept := { s | wp.final? n s }
  step s₁ _ := { s₂ | if s₁ = Fin.last (n+1) then s₁ = s₂ else s₂ = s₁ + 1 }

content: def NFA'.autWidth (wp : WidthPredicate) (n : Nat) : NFA' 0 := ⟨_, NFA.autWidth wp n⟩

content: def autWidth_equiv : (CNFA.autWidth wp n).m.states ≃ (NFA'.autWidth wp n).σ where
  toFun := fun ⟨s, hs⟩ =>
    Fin.mk s (by admit /- proof elided -/
    )
  invFun q := ⟨q.val, by admit /- proof elided -/
  ⟩
  left_inv := by admit /- proof elided -/

-- Local lemmas used:
content: lemma RawCNFA.autWidth_spec {wp : WidthPredicate} :
  let m := RawCNFA.autWidth wp n
  m.WF ∧ m.stateMax = n+2 ∧
  (∀ s, s ∈ m.states → (s ∈ m.initials ↔ s = 0) ∧ (s ∈ m.finals ↔ wp.final? n s)) ∧
  (∀ s s', s ∈ m.states → s' ∈ m.states → (s' ∈ m.tr s 0 ↔ if s = n+1 then s = s' else s' = s + 1))

content: @[simp]
lemma CNFA.autWidth_states: s ∈ (autWidth wp n).m.states ↔ s < n+2

content: lemma CNFA.autWidth_initials : s ∈ (autWidth wp n).m.initials ↔ s = 0

content: lemma CNFA.autWidth_finals (hn : s < n + 2) : s ∈ (autWidth wp n).m.finals ↔ wp.final? n s

content: lemma CNFA.autWidth_tr (hs : s < n + 2) (hs' : s' < n + 2) : s' ∈ (autWidth wp n).m.tr s 0 ↔ if s = n+1 then s = s' else s' = s + 1

-- Target theorem and proof:
lemma CNFA.autWidth_spec : autWidth wp n |>.Sim (NFA'.autWidth wp n)
:= by
  apply simulFun_sim autWidth_equiv; simp [autWidth_equiv]; constructor
  · rintro q; simp_all [NFA'.autWidth, NFA.autWidth]; apply autWidth_finals (q.isLt)
  · rintro q; simp_all [NFA'.autWidth, NFA.autWidth]; rw [autWidth_initials]; exact
    eq_iff_eq_of_cmp_eq_cmp rfl
  · rintro a q q'; fin_cases a; simp_all [NFA'.autWidth, NFA.autWidth, instFinEnumBV];
    have h := @autWidth_tr q.val n q'.val wp q.isLt q'.isLt
    unfold State at *
    simp_all
    rcases q with ⟨q, hq⟩
    rcases q' with ⟨q', hq'⟩
    simp [Fin.last]
    split
    · rfl
    · rw [Fin.add_def]
      simp only [Fin.val_one, Fin.mk.injEq]
      rw [Nat.mod_eq_of_lt (by omega)]