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

def langRel (R : BVNRel n) : Set (BitVecs n) :=
  { bvs | R bvs.bvs }

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)

-- Repository definitions used:
content: structure BitVecs (n : Nat) where
  w : Nat
  bvs : List.Vector (BitVec w) n

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

content: def BitStream : Type := Nat → Bool

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

-- 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: def langRel (R : BVNRel n) : Set (BitVecs n) :=
  { bvs | R bvs.bvs }

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)

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

content: @[simp]
lemma in_enc : x ∈ enc '' S ↔ dec x ∈ S

content: @[simp]
lemma dec_snoc_in_langRel {n} {R : BVNRel n} {w : BitVecs' n} {a : BitVec n} :
    dec (w ++ [a]) ∈ langRel R ↔
      R (List.Vector.ofFn fun k => .cons (a.getLsbD k) ((dec w).bvs.get k))

-- Target theorem and proof:
lemma NFA'.correct_spec {M : NFA' n} {ζ : M.sa} {L : BVNRel n} :
    M.correct ζ L → M.accepts = langRel L
:= by
  rintro ⟨h1, h2, h3⟩
  simp [accepts, accepts']
  have heq : dec '' (enc '' langRel L) = langRel L := by simp
  rw [←heq]
  congr!
  suffices h : M.M.correct (fun q => enc '' langRel (ζ q)) (enc '' langRel L) by
    apply NFA.correct_spec h
  constructor
  · intros w; rw [in_enc]; simp [langRel, h1]; simp_rw [@in_enc _ _ w]; rfl
  intros w; induction w using List.reverseRecOn
  case nil =>
    intros q; simp only [NFA.eval_nil]; rw [in_enc]; simp [h2, langRel]
  case append_singleton w a ih =>
    rintro q
    simp only [NFA.eval_append_singleton]
    rw [in_enc]
    have h : M.M.eval w = { q | w ∈ enc '' langRel (ζ q) } := by
      ext; rw [ih]; dsimp; rfl
    rw [dec_snoc_in_langRel]
    rw [h]; simp_rw [in_enc]
    simp [langRel, h3]