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

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

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

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

-- Repository definitions used:
content: inductive RelationOrdering
| lt | le | gt | ge
deriving Repr, Fintype

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

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

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: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
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: def NFA'.sa2 (M : NFA' 2) := M.σ → BVRel

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

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

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

-- Target theorem and proof:
lemma NFA'.autSignedCmp_correct cmp : autSignedCmp cmp |>.correct2 autSignedCmpSA cmp.srel
:= by
  let getState {w} (bv1 bv2 : BitVec w) : NFA.signedCmpState :=
    if bv2.ult bv1 then .gt else if bv1.ult bv2 then .lt else .eq
  constructor <;> simp [NFA.autSignedCmp, autSignedCmp, autSignedCmpSA, RelationOrdering.srel]
  · cases cmp <;> simp [BitVec.sle_iff_slt_or_eq]; tauto
  · rintro (_ | _ | _) <;> simp
  · rintro (_ | _ | _) a w bv1 bv2 <;> simp [NFA.stepSet]
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro ⟨_, _⟩; use .eq; simp; fin_cases a <;> simp [instFinEnumBV] at * <;> tauto
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro _; fin_cases a <;> simp [NFA.signedCmpStep, instFinEnumBV] at *
        · use .gt; simp_all
        · use (getState bv1 bv2); simp [getState]; split_ifs <;> simp_all; apply ucmp_tricho <;> assumption
        · use .gt; simp_all
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro _; fin_cases a <;> simp [NFA.signedCmpStep, instFinEnumBV] at *
        · use .lt; simp_all
        · use (getState bv1 bv2); simp [getState]; split_ifs <;> simp_all; apply ucmp_tricho <;> assumption
        · use .lt; simp_all
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro _; fin_cases a <;> simp [NFA.signedCmpStep, instFinEnumBV] at *
        · use .lt; simp_all
        · use (getState bv1 bv2); simp [getState]; split_ifs <;> simp_all; apply ucmp_tricho <;> assumption
        · use .lt; simp_all
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro _; fin_cases a <;> simp [NFA.signedCmpStep, instFinEnumBV] at *
        · use .gt; simp_all
        · use (getState bv1 bv2); simp [getState]; split_ifs <;> simp_all; apply ucmp_tricho <;> assumption
        · use .gt; simp_all