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

def NFA.sa (_ : NFA α σ) := σ → Language α

structure NFA.correct (M : NFA α σ) (ζ : M.sa) (L : Language α) where
  cond1 : ∀ w, (w ∈ L ↔ ∃ q ∈ M.accept, w ∈ ζ q)
  cond2 : ∀ w q, q ∈ M.eval w ↔ w ∈ ζ q

abbrev BVNRel n := ∀ ⦃w⦄, List.Vector (BitVec w) n → Prop

def NFA'.sa (M : NFA' n) := M.σ → BVNRel n

structure NFA'.correct (M : NFA' n) (ζ : M.sa) (L : BVNRel n) where
  cond1 : ∀ ⦃w⦄ (bvn : List.Vector (BitVec w) n), (L bvn ↔ ∃ q ∈ M.M.accept, ζ q bvn)
  cond2 q : q ∈ M.M.start ↔ ζ q (List.Vector.replicate n .nil)
  cond3 q a {w} (bvn : List.Vector (BitVec w) n) : q ∈ M.M.stepSet { q | ζ q bvn } a ↔
              ζ q (List.Vector.ofFn fun k => BitVec.cons (a.getLsbD k) (bvn.get k))

structure NFA'.correct2 (M : NFA' 2) (ζ : M.sa2) (L : BVRel) where
  cond1 : ∀ (bv1 bv2 : BitVec w), (L bv1 bv2 ↔ ∃ q ∈ M.M.accept, ζ q bv1 bv2)
  cond2 q : q ∈ M.M.start ↔ ζ q .nil .nil
  cond3 q a w (bv1 bv2 : BitVec w) : q ∈ M.M.stepSet { q | ζ q bv1 bv2 } a ↔
              ζ q (BitVec.cons (a.getLsbD 0) bv1) (BitVec.cons (a.getLsbD 1) bv2)

section fsm

variable {arity : Type} [FinEnum arity]

open BitStream in

end fsm

section nfas_relations

inductive NFA.msbState : Type where
| i | f
deriving DecidableEq, Fintype

def NFA.msbStep (q : NFA.msbState) (a : BitVec 1) : List NFA.msbState :=
  match q, a with
  | .i, 0 => [.i]
  | .i, 1 => [.i, .f]
  | _, _ => []

def NFA.autMsbSet : NFA (BitVec 1) msbState where
  step s a := { s' | s' ∈ msbStep s a }
  start := {.i}
  accept := {.f}

def NFA'.autMsbSet : NFA' 1 := ⟨_, NFA.autMsbSet⟩

def NFA.msbLang : Language (BitVec 1) := { bvs  | bvs.getLast? = some 1 }

def NFA.msbSA (q : msbState) : Language (BitVec 1) :=
  match q with
  | .i => ⊤
  | .f => msbLang

def NFA.msbCorrect : NFA.autMsbSet.correct msbSA msbLang :=

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

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: @[simp]
def langMsb : Set (BitVecs 1) := { bvs | bvs.bvs.get 0 |>.msb }

content: structure BitVecs (n : Nat) where
  w : Nat
  bvs : List.Vector (BitVec w) n

content: def head (x : BitStream) : Bool      := x 0

content: def BitStream : Type := Nat → Bool

content: def enc (bvs : BitVecs n) : BitVecs' n :=
  (List.finRange bvs.w).map (fun i =>
    BitVec.ofFn (fun (k : Fin n) => (bvs.bvs.get k)[i]))

content: abbrev BitVecs' (n : Nat) := List (BitVec n)

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

content: @[simps]
def dec (bvs' : BitVecs' n) : BitVecs n where
  w := bvs'.length
  bvs := List.Vector.ofFn fun k => BitVec.ofFn fun i => bvs'[i].getLsbD k

content: def accepts (M : NFA' n) : Set (BitVecs n) := dec '' M.accepts'

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

content: def accepts' (M : NFA' n) : Set (BitVecs' n) := M.M.accepts

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

content: @[simp]
lemma dec_enc' : dec (enc bvs) = bvs

content: @[simp]
lemma dec_enc : Function.RightInverse (α := BitVecs' n) enc dec

content: lemma dec_enc_w (bvs : BitVecs n) : (dec (enc bvs)).w = bvs.w

content: @[simp]
theorem BitVec.ofFn_getElem {w : Nat} (f : Fin w → Bool) {i : Nat} (hi : i < w) :
    (BitVec.ofFn f)[i] = f ⟨i, hi⟩

content: @[simp]
theorem BitVec.ofFn_getLsbD {w : Nat} {f : Fin w → Bool} {i : Nat} (hi : i < w) :
    (BitVec.ofFn f).getLsbD i = f ⟨i, hi⟩

content: theorem BitVec.ofFn_getLsbD_fin {w : Nat} {f : Fin w → Bool} {i : Fin w} :
    (BitVec.ofFn f).getLsbD i = f i

-- Local definitions used:
content: def NFA.sa (_ : NFA α σ) := σ → Language α

content: structure NFA.correct (M : NFA α σ) (ζ : M.sa) (L : Language α) where
  cond1 : ∀ w, (w ∈ L ↔ ∃ q ∈ M.accept, w ∈ ζ q)
  cond2 : ∀ w q, q ∈ M.eval w ↔ w ∈ ζ q

content: abbrev BVNRel n := ∀ ⦃w⦄, List.Vector (BitVec w) n → Prop

content: def NFA'.sa (M : NFA' n) := M.σ → BVNRel n

content: structure NFA'.correct (M : NFA' n) (ζ : M.sa) (L : BVNRel n) where
  cond1 : ∀ ⦃w⦄ (bvn : List.Vector (BitVec w) n), (L bvn ↔ ∃ q ∈ M.M.accept, ζ q bvn)
  cond2 q : q ∈ M.M.start ↔ ζ q (List.Vector.replicate n .nil)
  cond3 q a {w} (bvn : List.Vector (BitVec w) n) : q ∈ M.M.stepSet { q | ζ q bvn } a ↔
              ζ q (List.Vector.ofFn fun k => BitVec.cons (a.getLsbD k) (bvn.get k))

content: structure NFA'.correct2 (M : NFA' 2) (ζ : M.sa2) (L : BVRel) where
  cond1 : ∀ (bv1 bv2 : BitVec w), (L bv1 bv2 ↔ ∃ q ∈ M.M.accept, ζ q bv1 bv2)
  cond2 q : q ∈ M.M.start ↔ ζ q .nil .nil
  cond3 q a w (bv1 bv2 : BitVec w) : q ∈ M.M.stepSet { q | ζ q bv1 bv2 } a ↔
              ζ q (BitVec.cons (a.getLsbD 0) bv1) (BitVec.cons (a.getLsbD 1) bv2)

content: inductive NFA.msbState : Type where
| i | f
deriving DecidableEq, Fintype

content: def NFA.msbStep (q : NFA.msbState) (a : BitVec 1) : List NFA.msbState :=
  match q, a with
  | .i, 0 => [.i]
  | .i, 1 => [.i, .f]
  | _, _ => []

content: def NFA.autMsbSet : NFA (BitVec 1) msbState where
  step s a := { s' | s' ∈ msbStep s a }
  start := {.i}
  accept := {.f}

content: def NFA'.autMsbSet : NFA' 1 := ⟨_, NFA.autMsbSet⟩

content: def NFA.msbLang : Language (BitVec 1) := { bvs  | bvs.getLast? = some 1 }

content: def NFA.msbSA (q : msbState) : Language (BitVec 1) :=
  match q with
  | .i => ⊤
  | .f => msbLang

content: def NFA.msbCorrect : NFA.autMsbSet.correct msbSA msbLang :=

-- Local lemmas used:
content: lemma NFA.correct_spec {M : NFA α σ} {ζ : M.sa} {L : Language α} :
    M.correct ζ L → M.accepts = L

-- Target theorem and proof:
@[simp]
lemma autMsbSet_accepts : NFA'.autMsbSet.accepts = langMsb
:= by
  simp [NFA'.accepts, NFA'.accepts', NFA'.autMsbSet]
  rw [NFA.correct_spec NFA.msbCorrect, NFA.msbLang]
  ext bvs; simp only [BitVec.ofNat_eq_ofNat, Set.mem_image, Set.mem_setOf_eq]
  constructor
  · rintro ⟨bvs', hl, heq⟩
    have _ : bvs'.length ≠ 0 := by cases bvs'; tauto; simp
    rw [←heq]
    simp [dec]
    rw [BitVec.msb_eq_getLsbD_last]
    rw [BitVec.ofFn_getLsbD (by omega)]
    simp
    rw [List.getLast?_eq_getElem?] at hl
    rw [List.getElem?_eq_getElem (by omega)] at hl
    injection hl
    simp_all only [BitVec.getElem_one, decide_true]
  · intros h; use enc bvs
    simp only [dec_enc', and_true]
    simp [enc]
    have hw : bvs.w ≠ 0 := by
      rcases bvs with ⟨w, bvs⟩; rintro rfl
      simp_all [BitVec.eq_nil (bvs.head)]
    use ⟨bvs.w - 1, by omega⟩
    simp; rw [List.getLast?_eq_getElem?]
    simp; constructor
    · rw [List.getElem?_eq_getElem (by simp; omega)]; simp
    · ext i hi; rw [BitVec.ofFn_getElem _ (by omega)]
      rw [BitVec.msb_eq_getLsbD_last] at h
      simp [←BitVec.getLsbD_eq_getElem]
      obtain rfl : i = 0 := by omega
      simp_all