-- 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 BVRel := ∀ ⦃w⦄, BitVec w → BitVec w → Prop

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

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

def NFA'.sa2 (M : NFA' 2) := M.σ → BVRel

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

def langRel2 (R : BVRel) : Set (BitVecs 2) :=
  { bvs | R (bvs.bvs.get 0) (bvs.bvs.get 1) }

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

def NFA.autEq : NFA (BitVec 2) Unit :=
  { start := ⊤, accept := ⊤, step _ a := { _s' | if a = 0 ∨ a = 3 then true else false }}

def NFA'.autEq : NFA' 2 :=
  ⟨Unit, NFA.autEq⟩

def NFA'.eqRel : BVRel := fun _ x y => x = y

inductive NFA.unsignedCmpState : Type where
| eq | gt | lt
deriving Fintype, DecidableEq

def NFA.unsignedCmpStep (q : NFA.unsignedCmpState) (a : BitVec 2) : List NFA.unsignedCmpState :=
  match q, a with
  | .eq, 0 => [ .eq ] | .eq, 3 => [ .eq ] | .eq, 1 => [ .gt ] | .eq, 2 => [ .lt ]
  | .gt, 0 => [ .gt ] | .gt, 1 => [ .gt ] | .gt, 3 => [ .gt ] | .gt, 2 => [ .lt ]
  | .lt, 0 => [ .lt ] | .lt, 1 => [ .gt ] | .lt, 2 => [ .lt ] | .lt, 3 => [ .lt ]

def NFA.autUnsignedCmp (cmp: RelationOrdering) : NFA (BitVec 2) unsignedCmpState where
  step s a := { s' | s' ∈ unsignedCmpStep s a }
  start := {s | s = .eq }
  accept := { s | s ∈ match cmp with | .lt => [unsignedCmpState.lt] | .le => [.lt, .eq] | .gt => [.gt] | .ge => [.gt, .eq] }

def NFA'.autUnsignedCmp (cmp: RelationOrdering) : NFA' 2 :=
  ⟨_, NFA.autUnsignedCmp cmp⟩

def RelationOrdering.urel (cmp : RelationOrdering) : BVRel :=
  match cmp with
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .le => fun _ bv1 bv2 => bv1.ule bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1
  | .ge => fun _ bv1 bv2 => bv2.ule bv1

def NFA'.autUnsignedCmpSA (q : NFA.unsignedCmpState) : BVRel :=
  match q with
  | .eq => fun _ bv1 bv2 => bv1 = bv2
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1

inductive NFA.signedCmpState : Type where
| eq | gt | lt | ltfin | gtfin
deriving DecidableEq, Fintype

def NFA.signedCmpStep (q : NFA.signedCmpState) (a : BitVec 2) : List NFA.signedCmpState :=
  match q, a with
  | .eq, 0 => [ .eq ] | .eq, 3 => [.eq] | .eq, 1 => [.gt, .ltfin] | .eq, 2 => [ .lt, .gtfin ]
  | .gt, 0 => [ .gt, .gtfin ] | .gt, 1 => [ .gt, .ltfin ] | .gt, 3 => [ .gt, .gtfin ] | .gt, 2 => [ .lt, .gtfin ]
  | .lt, 0 => [ .lt, .ltfin ] | .lt, 1 => [ .gt, .ltfin ] | .lt, 2 => [ .lt, .gtfin ] | .lt, 3 => [ .lt, .ltfin ]
  | .gtfin, _ => ∅
  | .ltfin, _ => ∅

def NFA.autSignedCmp (cmp: RelationOrdering) : NFA (BitVec 2) signedCmpState where
  step s a := { s' | s' ∈ signedCmpStep s a }
  start := { s | s = signedCmpState.eq }
  accept := { s | s ∈ match cmp with | .lt => [NFA.signedCmpState.ltfin] | .le => [.ltfin, .eq] | .gt => [.gtfin] | .ge => [.gtfin, .eq] }

def NFA'.autSignedCmp (cmp: RelationOrdering) : NFA' 2 :=
  ⟨_, NFA.autSignedCmp cmp⟩

def RelationOrdering.srel (cmp : RelationOrdering) : BVRel :=
  match cmp with
  | .lt => fun _ bv1 bv2 => bv1.slt bv2
  | .le => fun _ bv1 bv2 => bv1.sle bv2
  | .gt => fun _ bv1 bv2 => bv2.slt bv1
  | .ge => fun _ bv1 bv2 => bv2.sle bv1

def NFA'.autSignedCmpSA (q : NFA.signedCmpState) : BVRel :=
  match q with
  | .eq => fun _ bv1 bv2 => bv1 = bv2
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1
  | .ltfin => fun _ bv1 bv2 => bv1.slt bv2
  | .gtfin => fun _ bv1 bv2 => bv2.slt bv1

end nfas_relations

def Relation.absAutOfRelation (rel : Relation) : NFA' 2 :=
  match rel with
  | .eq => NFA'.autEq
  | .unsigned cmp => NFA'.autUnsignedCmp cmp
  | .signed cmp => NFA'.autSignedCmp cmp

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

content: inductive RelationOrdering
| lt | le | gt | ge
deriving Repr, Fintype

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

content: inductive Relation
| eq
| signed (ord : RelationOrdering)
| unsigned (ord : RelationOrdering)
deriving Repr

content: def evalRelation (rel : Relation) {w} (bv1 bv2 : BitVec w) : Prop :=
  match rel with
  | .eq => bv1 = bv2
  | .signed .lt => bv1.slt bv2
  | .signed .le => bv1.sle  bv2
  | .signed .gt => bv2.slt bv1
  | .signed .ge => bv2.sle bv1
  | .unsigned .lt => bv1.ult bv2
  | .unsigned .le => bv1.ule bv2
  | .unsigned .gt => bv2.ult bv1
  | .unsigned .ge => bv2.ule bv1

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

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

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

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

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

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 BitVec.ofFn {w : Nat} (f : Fin w → Bool) : BitVec w :=
  BitVec.iunfoldr (fun i _ => ((), f i)) () |>.2

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

-- 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 BVRel := ∀ ⦃w⦄, BitVec w → BitVec w → Prop

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

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

content: def NFA'.sa2 (M : NFA' 2) := M.σ → BVRel

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

content: def langRel2 (R : BVRel) : Set (BitVecs 2) :=
  { bvs | R (bvs.bvs.get 0) (bvs.bvs.get 1) }

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: def NFA.autEq : NFA (BitVec 2) Unit :=
  { start := ⊤, accept := ⊤, step _ a := { _s' | if a = 0 ∨ a = 3 then true else false }}

content: def NFA'.autEq : NFA' 2 :=
  ⟨Unit, NFA.autEq⟩

content: def NFA'.eqRel : BVRel := fun _ x y => x = y

content: inductive NFA.unsignedCmpState : Type where
| eq | gt | lt
deriving Fintype, DecidableEq

content: def NFA.unsignedCmpStep (q : NFA.unsignedCmpState) (a : BitVec 2) : List NFA.unsignedCmpState :=
  match q, a with
  | .eq, 0 => [ .eq ] | .eq, 3 => [ .eq ] | .eq, 1 => [ .gt ] | .eq, 2 => [ .lt ]
  | .gt, 0 => [ .gt ] | .gt, 1 => [ .gt ] | .gt, 3 => [ .gt ] | .gt, 2 => [ .lt ]
  | .lt, 0 => [ .lt ] | .lt, 1 => [ .gt ] | .lt, 2 => [ .lt ] | .lt, 3 => [ .lt ]

content: def NFA.autUnsignedCmp (cmp: RelationOrdering) : NFA (BitVec 2) unsignedCmpState where
  step s a := { s' | s' ∈ unsignedCmpStep s a }
  start := {s | s = .eq }
  accept := { s | s ∈ match cmp with | .lt => [unsignedCmpState.lt] | .le => [.lt, .eq] | .gt => [.gt] | .ge => [.gt, .eq] }

content: def NFA'.autUnsignedCmp (cmp: RelationOrdering) : NFA' 2 :=
  ⟨_, NFA.autUnsignedCmp cmp⟩

content: def RelationOrdering.urel (cmp : RelationOrdering) : BVRel :=
  match cmp with
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .le => fun _ bv1 bv2 => bv1.ule bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1
  | .ge => fun _ bv1 bv2 => bv2.ule bv1

content: def NFA'.autUnsignedCmpSA (q : NFA.unsignedCmpState) : BVRel :=
  match q with
  | .eq => fun _ bv1 bv2 => bv1 = bv2
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1

content: inductive NFA.signedCmpState : Type where
| eq | gt | lt | ltfin | gtfin
deriving DecidableEq, Fintype

content: def NFA.signedCmpStep (q : NFA.signedCmpState) (a : BitVec 2) : List NFA.signedCmpState :=
  match q, a with
  | .eq, 0 => [ .eq ] | .eq, 3 => [.eq] | .eq, 1 => [.gt, .ltfin] | .eq, 2 => [ .lt, .gtfin ]
  | .gt, 0 => [ .gt, .gtfin ] | .gt, 1 => [ .gt, .ltfin ] | .gt, 3 => [ .gt, .gtfin ] | .gt, 2 => [ .lt, .gtfin ]
  | .lt, 0 => [ .lt, .ltfin ] | .lt, 1 => [ .gt, .ltfin ] | .lt, 2 => [ .lt, .gtfin ] | .lt, 3 => [ .lt, .ltfin ]
  | .gtfin, _ => ∅
  | .ltfin, _ => ∅

content: def NFA.autSignedCmp (cmp: RelationOrdering) : NFA (BitVec 2) signedCmpState where
  step s a := { s' | s' ∈ signedCmpStep s a }
  start := { s | s = signedCmpState.eq }
  accept := { s | s ∈ match cmp with | .lt => [NFA.signedCmpState.ltfin] | .le => [.ltfin, .eq] | .gt => [.gtfin] | .ge => [.gtfin, .eq] }

content: def NFA'.autSignedCmp (cmp: RelationOrdering) : NFA' 2 :=
  ⟨_, NFA.autSignedCmp cmp⟩

content: def RelationOrdering.srel (cmp : RelationOrdering) : BVRel :=
  match cmp with
  | .lt => fun _ bv1 bv2 => bv1.slt bv2
  | .le => fun _ bv1 bv2 => bv1.sle bv2
  | .gt => fun _ bv1 bv2 => bv2.slt bv1
  | .ge => fun _ bv1 bv2 => bv2.sle bv1

content: def NFA'.autSignedCmpSA (q : NFA.signedCmpState) : BVRel :=
  match q with
  | .eq => fun _ bv1 bv2 => bv1 = bv2
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1
  | .ltfin => fun _ bv1 bv2 => bv1.slt bv2
  | .gtfin => fun _ bv1 bv2 => bv2.slt bv1

content: def Relation.absAutOfRelation (rel : Relation) : NFA' 2 :=
  match rel with
  | .eq => NFA'.autEq
  | .unsigned cmp => NFA'.autUnsignedCmp cmp
  | .signed cmp => NFA'.autSignedCmp cmp

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

content: lemma NFA'.correct_spec {M : NFA' n} {ζ : M.sa} {L : BVNRel n} :
    M.correct ζ L → M.accepts = langRel L

content: lemma NFA'.correct2_spec {M : NFA' 2} {ζ : M.sa2} {L : BVRel} :
    M.correct2 ζ L → M.accepts = langRel2 L

content: lemma NFA'.autEq_correct : autEq.correct2 (fun _ => eqRel) eqRel

content: lemma BitVec.ule_iff_ult_or_eq {w : ℕ} (bv1 bv2 : BitVec w):
    (bv1.ule bv2) = true ↔ (bv1.ult bv2) = true ∨ bv1 = bv2

content: @[simp]
lemma ucmp_tricho {bv1 bv2 : BitVec w} : (bv2.ult bv1) = false → (bv1.ult bv2) = false → bv1 = bv2

content: lemma NFA'.autUnsignedCmp_correct cmp : autUnsignedCmp cmp |>.correct2 autUnsignedCmpSA cmp.urel

content: private lemma BitVec.sle_iff_slt_or_eq {w : ℕ} (bv1 bv2 : BitVec w):
    (bv1.sle bv2) = true ↔ (bv1.slt bv2) = true ∨ bv1 = bv2

content: lemma NFA'.autSignedCmp_correct cmp : autSignedCmp cmp |>.correct2 autSignedCmpSA cmp.srel

-- Target theorem and proof:
@[simp]
lemma autOfRelation_accepts (r : Relation) :
    r.absAutOfRelation.accepts = r.language
:= by
  simp [Relation.absAutOfRelation]
  rcases r with ⟨⟩ | ⟨cmp⟩ | ⟨cmp⟩ <;> simp
  · rw [NFA'.correct2_spec NFA'.autEq_correct]
    simp [langRel2, NFA'.eqRel, evalRelation]
  · rw [NFA'.correct2_spec (NFA'.autSignedCmp_correct cmp)]
    simp [langRel2, evalRelation, RelationOrdering.srel]
    cases cmp <;> simp
  · rw [NFA'.correct2_spec (NFA'.autUnsignedCmp_correct cmp)]
    simp [langRel2, evalRelation, RelationOrdering.urel]
    cases cmp <;> simp