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

abbrev Alphabet (arity: Type) [FinEnum arity] := BitVec (FinEnum.card arity + 1)

variable {arity : Type} [FinEnum arity]

def finFunToBitVec [fe : FinEnum carry] (c : carry → Bool) : BitVec (FinEnum.card carry) :=
  BitVec.ofFn fun i => c (fe.equiv.invFun i)

def bitVecToFinFun [FinEnum ar] (bv : BitVec $ FinEnum.card ar) : ar → Bool :=
  fun c => bv[FinEnum.equiv.toFun c]

def NFA.ofFSM (p : FSM arity) : NFA (Alphabet arity) (p.α → Bool) where
  start := { q | q = p.initCarry }
  accept := ⊤
  step s a := {s' |
    let (s'', b) := p.nextBit s (bitVecToFinFun (a.truncate $ FinEnum.card arity))
    s' = s'' ∧ a.msb = b }

open BitStream in

def CNFA.ofFSM (p : FSM arity) : CNFA (FinEnum.card arity + 1) :=
  worklistRun (BitVec (FinEnum.card p.α))
    (fun _ => true)
    #[finFunToBitVec p.initCarry]
    (by admit /- proof elided -/
    )
    f
  where

-- Repository definitions used:
content: def carry (initCarry : Bool) (x y : BitStream) : BitStream :=
  fun n => (addAux' initCarry x y n).2

content: def BitStream : Type := Nat → Bool

content: def addAux' (carryIn : Bool) (x y : BitStream) (i : Nat) :  Bool × Bool :=
  let carryIn : Bool := match i with
    | 0 => carryIn
    | i + 1 => (addAux' carryIn x y i).2
  Prod.swap (BitVec.adcb (x i) (y i) carryIn)

content: def BitVec.ofFn {w : Nat} (f : Fin w → Bool) : BitVec w :=
  BitVec.iunfoldr (fun i _ => ((), f i)) () |>.2

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

content: def worklistRun (final : S → Bool) (inits : Array S)
    (hinits : inits.toList.Nodup) (f : S → Array (BitVec n × S)) : CNFA n :=
  ⟨worklistRun' _ S final inits hinits f, worklistRun'_wf (BitVec n) S⟩

content: def worklistRun' (final : S → Bool) (inits : Array S) (hinits : inits.toList.Nodup) (f : S → Array (A × S)) : RawCNFA A :=
  let st0 := worklist.initState _ _ inits hinits final
  go st0
where go (st0 : worklist.St A S) : RawCNFA A :=
  if hemp : st0.worklist.isEmpty then st0.m else
  let sa? := st0.worklist.back?
  match heq : sa? with
  | some sa =>
    let wl := st0.worklist.pop
    let st1 := { st0 with worklist := wl,
                          worklist_nodup := by admit /- proof elided -/

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.initState (inits : Array S) (hinits : inits.toList.Nodup) (final? : S → Bool) : worklist.St A S :=
  let m := RawCNFA.empty (A := A)
  let mapm := inits.foldl (init := (Std.HashMap.emptyWithCapacity, m)) fun (map, m) sa =>
    let (s, m) := m.newState
    let m := m.addInitial s
    let m := if final? sa then m.addFinal s else m
    (map.insert sa s, m)
  let map := mapm.1
  let m := mapm.2
  let worklist_incl : ∀ sa ∈ inits, sa ∈ map :=

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

content: abbrev State := Nat

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

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.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: structure CNFA (n : Nat) where
  m : RawCNFA (BitVec n)
  wf : m.WF

content: structure FSM (arity : Type) : Type 1 where
   
  ( α  : Type )
  [ i : FinEnum α ]
  [ h : Hashable α ]
  [ dec_eq : DecidableEq α ]
   
  ( initCarry : α → Bool )
   
  outputCirc : Circuit (α ⊕ arity)
  nextStateCirc : α → Circuit (α ⊕ arity)

content: inductive Circuit (α : Type u) : Type u
  | tru : Circuit α
  | fals : Circuit α
   
  | var : (positive: Bool) → α → Circuit α
  | and : Circuit α → Circuit α → Circuit α
  | or : Circuit α → Circuit α → Circuit α
  | xor : Circuit α → Circuit α → Circuit α
deriving Repr, DecidableEq

content: def Var (Γ : Context Ty) (ty : Ty) : Type := { v : VarName // Γ.hasType v ty }

content: def bv2 : BitVec 4 := BitVec.ofNat 4 1 

content: def bv1 : BitVec 4 := BitVec.ofNat 4 5

-- Repository lemmas used:
content: @[ext]
theorem ext {x y : BitStream} (h : ∀ i, x i = y i) : x = y

-- Local definitions used:
content: abbrev Alphabet (arity: Type) [FinEnum arity] := BitVec (FinEnum.card arity + 1)

content: def finFunToBitVec [fe : FinEnum carry] (c : carry → Bool) : BitVec (FinEnum.card carry) :=
  BitVec.ofFn fun i => c (fe.equiv.invFun i)

content: def bitVecToFinFun [FinEnum ar] (bv : BitVec $ FinEnum.card ar) : ar → Bool :=
  fun c => bv[FinEnum.equiv.toFun c]

content: def NFA.ofFSM (p : FSM arity) : NFA (Alphabet arity) (p.α → Bool) where
  start := { q | q = p.initCarry }
  accept := ⊤
  step s a := {s' |
    let (s'', b) := p.nextBit s (bitVecToFinFun (a.truncate $ FinEnum.card arity))
    s' = s'' ∧ a.msb = b }

content: def CNFA.ofFSM (p : FSM arity) : CNFA (FinEnum.card arity + 1) :=
  worklistRun (BitVec (FinEnum.card p.α))
    (fun _ => true)
    #[finFunToBitVec p.initCarry]
    (by admit /- proof elided -/
    )
    f
  where

-- Local lemmas used:
content: @[simp]
theorem BitVec.cast_inj (h : w = w') {x y : BitVec w} : BitVec.cast h x = BitVec.cast h y ↔ x = y

content: @[simp]
theorem BitVec.append_inj {x1 x2 : BitVec w} {y1 y2 : BitVec w'} :
    x1 ++ y1 = x2 ++ y2 ↔ x1 = x2 ∧ y1 = y2

content: @[simp]
lemma BitVec.cons_inj : cons b1 bv1 = cons b2 bv2 ↔ (b1 = b2) ∧ bv1 = bv2

content: @[simp]
lemma bitVecToFinFun_rinv (c : carry → Bool) [FinEnum carry]:
    bitVecToFinFun (finFunToBitVec c) = c

content: @[simp]
lemma bitVecToFinFun_linv [FinEnum ar] (bv : BitVec $ FinEnum.card ar) :
    finFunToBitVec (bitVecToFinFun bv) = bv

content: @[simp]
lemma bitVecToFinFun_inj [FinEnum ar] : Function.Injective (bitVecToFinFun (ar := ar))

-- Target theorem and proof:
@[simp]
lemma CNFA.ofFSM.f_spec {p : FSM arity} {s s' : BitVec (FinEnum.card p.α)} :
    (a, s') ∈ f p s ↔ bitVecToFinFun s' ∈ (NFA.ofFSM p).step (bitVecToFinFun s) a
:= by
  let motive (as : List (BitVec (FinEnum.card arity))) := ∀ (acc : Array _) a s',
    ((a, s') ∈ as.foldl (init := acc) (process p s))
      ↔ (a, s') ∈ acc ∨ (a.setWidth (FinEnum.card arity) ∈ as) ∧
          bitVecToFinFun s' ∈ (NFA.ofFSM p).step (bitVecToFinFun s) a
  suffices h : motive (FinEnum.toList (BitVec (FinEnum.card arity))) by
    specialize h #[]
    simp at h
    rw [←h]
    rfl
  generalize FinEnum.toList (BitVec (FinEnum.card arity)) = qs
  induction qs
  case nil => simp [motive]
  case cons a as ih =>
    rintro acc b s'
    simp only [List.foldl_cons]; rw [ih]
    simp [process]
    constructor
    · rintro ((hacc | ⟨rfl, rfl⟩) | ⟨hin₁, hin₂⟩)
      · exact .inl hacc
      · right; simp [NFA.ofFSM]; constructor <;> rfl
      · right; simp_all only [or_true, and_self, motive]
    · rintro (hacc | ⟨(rfl | hold), hst⟩)
      · tauto
      · simp [NFA.ofFSM] at hst; left; right
        rcases hst with ⟨hs', hb⟩
        constructor
        · rw [←BitVec.cons_msb_setWidth b]
          simp_all only [BitVec.setWidth_cons, BitVec.cons_inj, and_true]; rfl
        · apply_fun bitVecToFinFun <;> simp only [hs', bitVecToFinFun_rinv, bitVecToFinFun_inj]
          rfl
      · tauto