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

variable {arity : Type} [FinEnum arity]

open BitStream in

end fsm

section nfas_relations

end nfas_relations

def liftOp n : Fin (n + 1) → Fin (n + 3) :=
  fun k =>
    if k = n then Fin.last (n+2) else k.castLE (by admit /- proof elided -/
    )

@[simp]
def liftOp_unchanged (k : Fin n) : liftOp n k.castSucc = k.castLE (by simp) :=

inductive TermBinop where
| and | or | xor | add | sub

def TermBinop.subst (op : TermBinop) (t₁ t₂ : Term) : Term :=
  match op with
  | .and => .and t₁ t₂
  | .or => .or t₁ t₂
  | .xor => .xor t₁ t₂
  | .add => .add t₁ t₂
  | .sub => .sub t₁ t₂

def TermBinop.openTerm (op : TermBinop) : Term := op.subst (.var 0) (.var 1)

@[simp]
def TermBinop.openTerm_arity (op : TermBinop) : op.openTerm.arity + 1 = 3 :=

def swapLastTwoBlock (x : Fin (n + 3)) : Fin (n + 3) :=
  if x = Fin.last (n+2) then n
  else if x = n+1 then Fin.last (n + 2)
  else if x = n then n + 1
  else x

-- Repository definitions used:
content: syntax "xor" : MLIR.Pretty.uniform_op

syntax "max" : MLIR.Pretty.uniform_op

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = const ($x)
      $[: $outer_type]? ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _)
      `(mlir_op| $res:mlir_op_operand = "const"()
          {val = $x:num : $outer_type} : ($outer_type) -> ($outer_type) )

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = li ($x)
     $[: $outer_type]?  ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _ )
      `(mlir_op| $res:mlir_op_operand = "li"()
          {imm = $x:num : $outer_type } : ($outer_type) -> ($outer_type))

content: macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithImmediate $reg1 , $x : $t)  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {imm = $x:num : $t} : ($t) -> ($t) )

content: macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithShamt $reg1 , $x  : $t )  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {shamt = $x:num : $t}  : ($t) -> ($t) )

content: inductive Term : Type
| var : Nat → Term
 
| zero : Term
 
| negOne : Term
 
| one : Term
 
| ofNat (n : Nat) : Term
 
| and : Term → Term → Term
 
| or : Term → Term → Term
 
| xor : Term → Term → Term
 
| not : Term → Term
 
| add : Term → Term → Term
 
| sub : Term → Term → Term
 
| neg : Term → Term


 
| shiftL : Term → Nat → Term


deriving Repr, Inhabited, Lean.ToExpr

content: def neg (x : BitStream) : BitStream :=
  fun n => (negAux x n).1

content: def negAux (x : BitStream) : Nat → Bool × Bool
  | 0 => (x 0, !(x 0))
  | n+1 =>
    let borrow := (negAux x n).2
    let a := x (n + 1)
    (xor (!a) borrow, !a && borrow)

content: def BitStream : Type := Nat → Bool

content: def sub (x y : BitStream) : BitStream :=
  fun n => (subAux x y n).1

content: def subAux (x y : BitStream) : Nat → Bool × Bool
  | 0 => (xor (x 0) (y 0), !(x 0) && y 0)
  | n+1 =>
    let borrow := (subAux x y n).2
    let a := x (n + 1)
    let b := y (n + 1)
    (xor a (xor b borrow), !a && b || ((!(xor a b)) && borrow))

content: def add (x y : BitStream) : BitStream :=
  fun n => (addAux x y n).1

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

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

content: def liftMaxSuccSucc2 (n m : Nat) : Fin (m + 1) → Fin (max n m + 3) :=
  fun k => if _ : k = Fin.last m then max n m + 1 else k.castLE (by admit /- proof elided -/
  )

content: def liftMaxSuccSucc1 (n m : Nat) : Fin (n + 1) → Fin (max n m + 3) :=
  fun k => if _ : k = Fin.last n then (max n m).cast else k.castLE (by admit /- proof elided -/
  )

content: def liftLast3 n : Fin 3 → Fin (n + 3)
| 0 => n
| 1 => n + 1
| 2 => Fin.last (n + 2)

content: def BitVecs.cast (bvs : BitVecs n) (h : n = n') : BitVecs n' :=
  { w := bvs.w, bvs := h ▸ bvs.bvs }

content: def Valuation.cons {Γ : Ctxt Ty} {t : Ty} (x : toType t) (V : Γ.Valuation) :
    (Γ.cons t).Valuation :=

content: def Term.language (t : Term) : Set (BitVecs (t.arity + 1)) :=
  { bvs : BitVecs (t.arity + 1) | t.evalFinBV (fun n => bvs.bvs.get n) = bvs.bvs.get t.arity }

content: @[simp]
def Formula.arity : Formula → Nat
| width _ _ => 0
| atom _ t1 t2 => max t1.arity t2.arity
| msbSet t => t.arity
| unop _ φ => φ.arity
| binop _ φ1 φ2 => max φ1.arity φ2.arity

content: @[simp] def Term.evalFinBV (t : Term) (vars : Fin (arity t) → BitVec w) : BitVec w :=
  match t with
  | .var n => vars (Fin.last n)
  | .zero    => BitVec.zero w
  | .one     => 1
  | .negOne  => -1
  | .ofNat n => BitVec.ofNat _ n
  | .and t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ &&& x₂
  | .or t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ ||| x₂
  | .xor t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ ^^^ x₂
  | .not t     => ~~~(t.evalFinBV vars)
  
  | .add t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ + x₂
  | .sub t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ - x₂
  | .neg t       => -(t.evalFinBV vars)
  | .shiftL a n => (a.evalFinBV vars) <<< n

content: @[simp] def Term.arity : Term → Nat
| (var n) => n+1
| zero => 0
| one => 0
| negOne => 0
| ofNat _ => 0
| Term.and t₁ t₂ => max (arity t₁) (arity t₂)
| Term.or t₁ t₂ => max (arity t₁) (arity t₂)
| Term.xor t₁ t₂ => max (arity t₁) (arity t₂)
| Term.not t => arity t
| add t₁ t₂ => max (arity t₁) (arity t₂)
| sub t₁ t₂ => max (arity t₁) (arity t₂)
| neg t => arity t


| shiftL t .. => arity t

content: infixr:67 "::ᵥ" => Valuation.cons

-- Repository lemmas used:
content: @[simp]
lemma List.Vector.append_get_ge {x : List.Vector α n} {y : List.Vector α m} {i : Fin (n+m)} (hlt: n ≤ i) :
    (x ++ y).get i = y.get ((i.cast (Nat.add_comm n m) |>.subNat n hlt))

content: @[simp]
lemma List.Vector.append_get_lt {x : List.Vector α n} {y : List.Vector α m} {i : Fin (n+m)} (hlt: i < n) :
    (x ++ y).get i = x.get (i.castLT hlt)

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

-- Local definitions used:
content: def liftOp n : Fin (n + 1) → Fin (n + 3) :=
  fun k =>
    if k = n then Fin.last (n+2) else k.castLE (by admit /- proof elided -/
    )

content: @[simp]
def liftOp_unchanged (k : Fin n) : liftOp n k.castSucc = k.castLE (by simp) :=

content: inductive TermBinop where
| and | or | xor | add | sub

content: def TermBinop.subst (op : TermBinop) (t₁ t₂ : Term) : Term :=
  match op with
  | .and => .and t₁ t₂
  | .or => .or t₁ t₂
  | .xor => .xor t₁ t₂
  | .add => .add t₁ t₂
  | .sub => .sub t₁ t₂

content: def TermBinop.openTerm (op : TermBinop) : Term := op.subst (.var 0) (.var 1)

content: @[simp]
def TermBinop.openTerm_arity (op : TermBinop) : op.openTerm.arity + 1 = 3 :=

content: def swapLastTwoBlock (x : Fin (n + 3)) : Fin (n + 3) :=
  if x = Fin.last (n+2) then n
  else if x = n+1 then Fin.last (n + 2)
  else if x = n then n + 1
  else x

-- Local lemmas used:
content: lemma TermBinop.subst_arity' {op : TermBinop} : (op.subst t₁ t₂).arity + 1= t₁.arity ⊔ t₂.arity + 1

content: @[simp]
lemma BitVecs.cast_eq  (x : BitVecs n) (h : n = n') : h ▸ x = x.cast h

-- Target theorem and proof:
lemma TermBinop.alt_lang {t₁ t₂ : Term} (op : TermBinop) :
  (op.subst_arity' ▸ (op.subst t₁ t₂).language) =
    let lop : Set (BitVecs 3) := op.openTerm_arity ▸ op.openTerm.language
    let lop' : Set (BitVecs ((t₁.arity ⊔ t₂.arity) + 3)) := lop.lift (liftLast3 (max t₁.arity t₂.arity))
    let l₁ := t₁.language.lift (liftMaxSuccSucc1 t₁.arity t₂.arity)
    let l₂ := t₂.language.lift (liftMaxSuccSucc2 t₁.arity t₂.arity)
    let l := l₁ ∩ l₂ ∩ lop'
    l.proj (liftOp _)
:= op.openTerm_arity ▸ op.openTerm.language
    let lop' : Set (BitVecs ((t₁.arity ⊔ t₂.arity) + 3)) := lop.lift (liftLast3 (max t₁.arity t₂.arity))
    let l₁ := t₁.language.lift (liftMaxSuccSucc1 t₁.arity t₂.arity)
    let l₂ := t₂.language.lift (liftMaxSuccSucc2 t₁.arity t₂.arity)
    let l := l₁ ∩ l₂ ∩ lop'
    l.proj (liftOp _)
    := by
  simp [Term.language]
  ext bvs
  simp
  constructor
  · rintro heq
    let bvs' := bvs.bvs ++
                  (t₁.evalFinBV (λ n ↦ bvs.bvs.get n) ::ᵥ t₂.evalFinBV (λ n ↦ bvs.bvs.get n) ::ᵥ List.Vector.nil)
                |>.transport swapLastTwoBlock
    use ⟨_, bvs'⟩
    simp [bvs']
    split_ands
    · rw [liftMaxSuccSucc1]; simp
      conv =>
        enter [1, 2, n]
        rw [List.Vector.append_get_lt (by rcases n with ⟨n, hn⟩; simp_all; rw [Nat.mod_eq_of_lt (by omega)]; omega)]
        rfl

      rw [List.Vector.append_get_ge]
      on_goal 2 => apply Nat.le_of_eq; simp [Fin.add_def]; rw [Nat.mod_eq_of_lt (by omega)]

      generalize_proofs h₁ h₂ h₃ h₄
      have heq h : Fin.subNat (t₁.arity ⊔ t₂.arity + 1) (Fin.cast h₄ (↑(t₁.arity ⊔ t₂.arity) + 1)) h = 0 := by
         ext; simp [Fin.add_def]; rw [Nat.mod_eq_of_lt (by omega)]; omega
      simp [heq]
      congr; ext1 i; congr 1
      rcases i with ⟨i, hi⟩; simp [Fin.castLT]; ext; simp; repeat rw [Nat.mod_eq_of_lt (by omega)]
    · rw [liftMaxSuccSucc2]; simp
      conv =>
        enter [1, 2, n]
        rw [List.Vector.append_get_lt (by rcases n with ⟨n, hn⟩; simp_all; rw [Nat.mod_eq_of_lt (by omega)]; omega)]
        rfl
      generalize_proofs h₁ h₂ h₃ h₄
      have heq h : Fin.subNat (t₁.arity ⊔ t₂.arity + 1) (Fin.last (2 + t₁.arity ⊔ t₂.arity)) h = 1 := by
         ext; simp; omega
      simp [heq]
      congr; ext1 i; congr 1
      rcases i with ⟨i, hi⟩; simp [Fin.castLT]; ext; simp; repeat rw [Nat.mod_eq_of_lt (by omega)]
    · rw [BitVecs.cast_eq] at *
      simp [BitVecs.cast] at *
      rw [liftLast3]
      simp
      rw [List.Vector.append_get_lt (by simp; rw [Nat.mod_eq_of_lt (by omega)]; omega)]
      convert heq using 1
      have h n : n + 1 < n + 3 := by omega
      have h' n m k : n < m → n < (k ⊔ m) + 1 := by omega
      · cases op <;>
        · simp [openTerm, subst, liftLast3]; congr
          · rw [List.Vector.append_get_ge]
            · simp [Fin.subNat, Fin.add_def]; simp [Nat.mod_eq_of_lt, h]; congr! with ⟨n, hn⟩; ext; simp; omega
            · rw [Fin.add_def]; simp; rw [Nat.mod_eq_of_lt (by omega)]
          · have heq h : (Fin.subNat (t₁.arity ⊔ t₂.arity + 1) (Fin.last (2 + t₁.arity ⊔ t₂.arity)) h) = 1 := by
              simp [Fin.subNat]; omega
            simp [heq, List.Vector.get]; congr; ext1 ⟨i, hi⟩; simp; simp [Nat.mod_eq_of_lt (h' _ _ t₁.arity hi)]
      · congr!; simp [Fin.castLT, Fin.last]; omega
    · ext1
      · simp
      next i =>
        simp [liftOp]
        split_ifs with h
        · subst h
          simp
          rw [List.Vector.append_get_lt]
          on_goal 2 =>
            simp +arith
            repeat (rw [Nat.mod_eq_of_lt (by omega)])
            omega
          congr; ext; simp +arith
          omega
        · simp [h]
          rw [List.Vector.append_get_lt (by rcases i; simp_all [Fin.last]; rw [Nat.mod_eq_of_lt (by omega)]; omega)]
          rcases i with ⟨i, hi⟩; congr!; simp_all; omega
  · rintro ⟨bvs', ⟨⟨⟨heq₁, heq₂⟩, heq₃⟩, heq₄⟩⟩
    rw [BitVecs.cast_eq] at *
    simp [BitVecs.cast] at *
    rw [←heq₄]
    conv_rhs =>
      simp
      simp [liftOp]
      rfl
    simp [liftMaxSuccSucc1, liftMaxSuccSucc2] at heq₁ heq₂
    rw [liftLast3] at heq₃
    convert heq₃ using 1
    have h₁ : (t₁.evalFinBV fun i => bvs'.bvs.get (liftOp (t₁.arity ⊔ t₂.arity) (Fin.castLE (by omega) i).castSucc)) =
                  bvs'.bvs.get (liftLast3 (t₁.arity ⊔ t₂.arity) 0) := by
      simp only [liftOp_unchanged, Fin.castLE_castLE, liftLast3]; convert heq₁ using 1
    have h₂ : (t₂.evalFinBV fun i => bvs'.bvs.get (liftOp (t₁.arity ⊔ t₂.arity) (Fin.castLE (by omega) i).castSucc)) =
                  bvs'.bvs.get (liftLast3 (t₁.arity ⊔ t₂.arity) 1) := by
      simp only [liftOp_unchanged, Fin.castLE_castLE, liftLast3]; convert heq₂
    rcases op with _ | _ | _ <;>
    . simp [subst, openTerm] at *; congr