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

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

end nfas_relations

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

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)

inductive TermUnop where
| neg | not | shiftL (k : Nat)

def TermUnop.openTerm (op : TermUnop) : Term :=
  match op with
  | .neg => .neg (.var 0)
  | .not => .not (.var 0)
  | .shiftL k => .shiftL (.var 0) k

@[simp]
def TermUnop.openTerm_arity' (op : TermUnop) : op.openTerm.arity + 1 = 2 :=

def TermUnop.subst (op : TermUnop) (t : Term) : Term :=
  match op with
  | .neg => .neg t
  | .not => .not t
  | .shiftL k => .shiftL t k

def swapLastTwo (x : Fin (n + 2)) : Fin (n + 2) :=
  if x = Fin.last (n + 1) then n else if x = n then Fin.last (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 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 BitStream : Type := Nat → Bool

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: 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: structure BitVecs (n : Nat) where
  w : Nat
  bvs : List.Vector (BitVec w) n

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: def BitVec.transport (f : Fin n2 → Fin n1) (bv : BitVec n1) : BitVec n2 :=
  BitVec.ofFn fun i => bv.getLsbD (f i)

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

content: def List.Vector.transport (v : Vector α m) (f : Fin n → Fin m) : Vector α n :=
  Vector.ofFn fun i => v.get (f i)

content: def BitVecs'.transport (f : Fin n → Fin m) (bvs' : BitVecs' m): BitVecs' n :=
  bvs'.map fun bv => bv.transport f

content: def BitVecs.transport (f : Fin n → Fin m) (bvs : BitVecs m) : BitVecs n :=
  { w := bvs.w, bvs := bvs.bvs.transport f }

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

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

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: @[simp]
lemma BitVecs.transport_getElem {bvs : BitVecs m} (f : Fin n → Fin m) (i : Fin n) :
    (bvs.transport f).bvs.get i = bvs.bvs.get (f i)

-- Local definitions used:
content: inductive NFA.msbState : Type where
| i | f
deriving DecidableEq, Fintype

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

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: inductive TermUnop where
| neg | not | shiftL (k : Nat)

content: def TermUnop.openTerm (op : TermUnop) : Term :=
  match op with
  | .neg => .neg (.var 0)
  | .not => .not (.var 0)
  | .shiftL k => .shiftL (.var 0) k

content: @[simp]
def TermUnop.openTerm_arity' (op : TermUnop) : op.openTerm.arity + 1 = 2 :=

content: def TermUnop.subst (op : TermUnop) (t : Term) : Term :=
  match op with
  | .neg => .neg t
  | .not => .not t
  | .shiftL k => .shiftL t k

content: def swapLastTwo (x : Fin (n + 2)) : Fin (n + 2) :=
  if x = Fin.last (n + 1) then n else if x = n then Fin.last (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

content: lemma Fin.natAdd_zero' [h : NeZero m] : Fin.natAdd (m := m) n 0 = n

content: @[simp]
lemma TermUnop.subst_arity' {op : TermUnop} : (op.subst t).arity + 1 = t.arity + 1

-- Target theorem and proof:
lemma TermUnop.alt_lang {t : Term} (op : TermUnop) :
  (op.subst_arity' ▸ (op.subst t).language) =
    let lop : Set (BitVecs 2) := op.openTerm_arity' ▸ op.openTerm.language
    let lop' : Set (BitVecs (t.arity + 2)) := lop.lift (λ i ↦ i.natAdd t.arity)
    let lt : Set (BitVecs (t.arity + 2)) := t.language.lift (λ i ↦ i.castLE (by omega))
    let l := lt ∩ lop'
    l.proj (liftUnop t.arity)
:= op.openTerm_arity' ▸ op.openTerm.language
    let lop' : Set (BitVecs (t.arity + 2)) := lop.lift (λ i ↦ i.natAdd t.arity)
    let lt : Set (BitVecs (t.arity + 2)) := t.language.lift (λ i ↦ i.castLE (by omega))
    let l := lt ∩ lop'
    l.proj (liftUnop t.arity)
    := by
  simp [Term.language]
  ext bvs
  simp

  generalize_proofs h₁ h₂ h₃ h₄ h₅
  constructor
  · rintro heq
    let bvs' := bvs.bvs ++ (t.evalFinBV (λ n ↦ bvs.bvs.get n) ::ᵥ List.Vector.nil) |>.transport swapLastTwo
    use ⟨_, bvs'⟩
    split_ands
    · simp [bvs']
      have heq : (swapLastTwo (Fin.castLE h₃ (Fin.last t.arity))) = Fin.last (t.arity + 1) := by
        simp [swapLastTwo]
        split_ifs with h₁ h₂
        · exfalso; rw [Fin.ext_iff] at h₁
          simp at h₁
        · rfl
        · exfalso; apply h₂; ext; simp; exact Eq.symm (Nat.mod_eq_of_lt h₃)
      rw [heq]
      rw [List.Vector.append_get_ge]
      on_goal 2 => simp
      simp
      congr
      ext1 x
      rw [List.Vector.append_get_lt]
      on_goal 2 => simp +arith
      congr!
    · rw [BitVecs.cast_eq] at heq ⊢
      unfold BitVecs.cast
      simp
      simp [BitVecs.cast] at heq
      have hget : bvs'.get (Fin.natAdd t.arity 1) = bvs.bvs.get t.arity := by
        simp [bvs']
        rw [List.Vector.append_get_lt]
        on_goal 2 => simp +arith; exact Nat.mod_le t.arity (t.arity + 2)
        · congr
          ext
          simp +arith
      rw [hget]
      convert heq using 1
      · have hbvs' : bvs'.get t.arity = t.evalFinBV fun n => bvs.bvs.get n := by
          simp [bvs']
        simp [openTerm]
        rcases op <;> simp [hbvs', Fin.natAdd_zero'] <;> rfl
      · congr!;  ext; simp
    · ext1
      · simp
      next i =>
        simp [bvs', liftUnop]
        split_ifs with h
        · subst h
          simp
          rw [List.Vector.append_get_lt]
          on_goal 2 => simp +arith; exact Nat.mod_le t.arity (t.arity + 2)
          congr
          ext; simp +arith
        · simp [h]
          congr!
  · rintro ⟨bvs', ⟨⟨heq₁, heq₂⟩, heq₃⟩⟩
    rw [BitVecs.cast_eq] at *
    simp [BitVecs.cast] at *
    rw [←heq₃]
    conv_rhs =>
      simp only [BitVecs.transport_getElem]
      simp [liftUnop]
      rw [ite_cond_eq_true]
      rfl
      tactic => simp only [eq_iff_iff, iff_true]; ext1; simp
    convert heq₂ using 1
    rcases op with _ | _ | _ <;>
    . simp [subst, openTerm] at *; congr