-- Blase/Blase/MultiWidth/GoodFSM.lean

import Blase.Fast.FiniteStateMachine

import Blase.Vars

import Blase.MultiWidth.Defs

import Blase.KInduction.KInduction

import Lean

namespace MultiWidth

def NatFSM.fsmUnaryIndexUle (a : NatFSM wcard tcard bcard ncard icard pcard v) :
    FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
  composeUnaryAux (FSM.ls true) a.toFsm

def fsmUnaryUle (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w)) : FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
 composeUnaryAux FSM.scanAnd (b.fsmUnaryIndexUle ||| ~~~ a.fsmUnaryIndexUle)

def fsmUnaryNeqUpto (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w)) : FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
  composeUnaryAux FSM.scanOr (a.fsmUnaryIndexUle ^^^ b.fsmUnaryIndexUle)

def fsmUltUnary
    (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w)) : FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
  composeBinaryAux' FSM.and (fsmUnaryUle a b) (fsmUnaryNeqUpto a b)

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

syntax "neg" : 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: def ls (b : Bool) : FSM Unit :=
  { α := Unit,
    initCarry := fun _ => b,
    nextStateCirc := fun () => Circuit.var true (inr ()),
    outputCirc := Circuit.var true (inl ())
  }

content: structure FSM (arity : Type) : Type 1 where
   
  ( α  : Type )
  [ i : FinEnum α ]
  [ h : Hashable α ]
  [ dec_eq : DecidableEq α ]
   
  ( initCarry : α → Bool )
   
  outputCirc : Circuit (α ⊕ arity)
  nextStateCirc : α → Circuit (α ⊕ arity)

content: inductive Circuit (α : Type u) : Type u
  | tru : Circuit α
  | fals : Circuit α
   
  | var : (positive: Bool) → α → Circuit α
  | and : Circuit α → Circuit α → Circuit α
  | or : Circuit α → Circuit α → Circuit α
  | xor : Circuit α → Circuit α → Circuit α
deriving Repr, DecidableEq

content: def Var (Γ : Context Ty) (ty : Ty) : Type := { v : VarName // Γ.hasType v ty }

content: inductive StateSpace (wcard tcard bcard ncard icard pcard : Nat)
| widthVar (v : Fin wcard)
| termVar (v : Fin tcard)
| predVar (v : Fin pcard)
| boolVar (v : Fin bcard)
deriving DecidableEq, Repr, Hashable

content: inductive Term
| ofNat (w : WidthExpr) (n : Nat) : Term
| var (v : Nat) (w : WidthExpr) : Term
| add (w : WidthExpr) (a b : Term) : Term
| zext (a : Term) (wnew : WidthExpr) : Term
| setWidth (a : Term) (wnew : WidthExpr) : Term
| sext (a : Term) (wnew : WidthExpr) : Term
| bor (w : WidthExpr) (a b : Term) : Term
| band (w : WidthExpr) (a b : Term) : Term
| bxor (w : WidthExpr) (a b : Term) : Term
| bnot (w : WidthExpr)  (a : Term) : Term
| boolVar (v : Nat) : Term
| boolConst (b : Bool) : Term
| shiftl (w : WidthExpr) (a : Term) (k : Nat) : Term
| bvOfBool (b : Term) : Term
| binWidthRel (k : WidthBinaryRelationKind) (wa wb : WidthExpr) : Term
| binRel (k : BinaryRelationKind) (w : WidthExpr)
    (a : Term) (b : Term) : Term
| or (p1 p2 : Term) : Term
| and (p1 p2 : Term) : Term
| pvar (v : Nat) : Term
| boolBinRel (k : BoolBinaryRelationKind)
    (a b : Term) : Term
deriving DecidableEq, Inhabited, Repr, Lean.ToExpr

content: def composeUnaryAux
    (p : FSM Unit)
    (q : FSM arity) :
    FSM arity :=
  p.compose
    arity
    _
    (λ _ => id)
    (λ _ => q)

content: def compose [FinEnum arity] [DecidableEq arity] [Hashable arity]
    (new_arity : Type)        
    (q_arity : arity → Type)  
    (vars : ∀ (a : arity), q_arity a → new_arity)
    
    
    (q : ∀ (a : arity), FSM (q_arity a)) : 
    FSM new_arity :=
  { α := p.α ⊕ (Σ a, (q a).α),
    i := by admit /- proof elided -/

content: structure NatFSM (wcard tcard bcard ncard icard pcard : Nat) (v : Nondep.WidthExpr) where
  toFsm : FSM (StateSpace wcard tcard bcard ncard icard pcard)

content: inductive WidthExpr where
| const : Nat → WidthExpr
| var : Nat → WidthExpr
| max : WidthExpr → WidthExpr → WidthExpr
| min : WidthExpr → WidthExpr → WidthExpr
| addK : WidthExpr → Nat → WidthExpr
| kadd : Nat → WidthExpr → WidthExpr
deriving Inhabited, Repr, Hashable, DecidableEq, Lean.ToExpr

content: def scanAnd  : FSM Unit :=
  {
   α := Unit,
   initCarry := fun () => true,
   outputCirc := Circuit.var true (inl ()) &&& Circuit.var true (inr ()),
   nextStateCirc := fun () => (Circuit.var true (inl ())) &&& (Circuit.var true (inr ()))
  }

content: def scanOr  : FSM Unit :=
  {
   α := Unit,
   initCarry := fun () => false,
   outputCirc := Circuit.var true (inl ()) ||| Circuit.var true (inr ()),
   nextStateCirc := fun () => Circuit.var true (inl ()) ||| Circuit.var true (inr ())
  }

content: def composeBinaryAux'
    (p : FSM Bool)
    (qtrue : FSM α)
    (qfalse : FSM α) :
    FSM α :=
  p.compose (α)
    (λ _ => α)
    (λ _ i => i)
    (λ b => match b with
      | true => qtrue
      | false => qfalse)

content: def and : FSM Bool :=
  { α := Empty,
    initCarry := Empty.elim,
    nextStateCirc := fun a => a.elim,
    outputCirc := Circuit.var true (inr true) &&& Circuit.var true (inr false),
  }

content: abbrev WidthExpr.Env (wcard : Nat) : Type :=
  Fin wcard → Nat

content: structure HWidthEnv {wcard tcard : Nat}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (wenv : Fin wcard → Nat) : Prop where
    heq_width : ∀ (v : Fin wcard),
      fsmEnv (StateSpace.widthVar v) = BitStream.ofNatUnary (wenv v)

content: structure HPredicateEnv {wcard tcard bcard ncard icard pcard : Nat}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (penv : Fin pcard → Prop) : Prop where
    heq_width : ∀ (v : Fin pcard),
      fsmEnv (StateSpace.predVar v) = BitStream.ofProp (penv v)

content: def BitStream : Type := Nat → Bool

content: abbrev ofNatUnary (n : Nat) : BitStream :=
  fun i => decide (i < n)

content: structure HNatFSMToBitstream {wcard : Nat} {v : WidthExpr wcard} {tcard : Nat} {bcard : Nat} {pcard : Nat}
   (fsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v)) : Prop where
  heq :
    ∀ (wenv : Fin wcard → Nat)
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
    (henv : HWidthEnv fsmEnv wenv) →
      fsm.toFsm.eval fsmEnv =
      BitStream.ofNatUnary (v.toNat wenv)

content: def WidthExpr.toNat (e : WidthExpr wcard) (env : WidthExpr.Env wcard) : Nat :=
  match e with
  | .const n => n
  | .var v => env v
  | .min v w => Nat.min (v.toNat env) (w.toNat env)
  | .max v w => Nat.max (v.toNat env) (w.toNat env)
  | .addK v k => v.toNat env + k
  | .kadd k v => k + v.toNat env

content: inductive WidthExpr (wcard : Nat) : Type
| const (n : Nat) :  WidthExpr wcard
| var : (v : Fin wcard) → WidthExpr wcard
| min : (v w : WidthExpr wcard) → WidthExpr wcard
| max : (v w : WidthExpr wcard) → WidthExpr wcard
| addK : (v : WidthExpr wcard) → (k : Nat) → WidthExpr wcard
| kadd : (k : Nat) → (v : WidthExpr wcard) → WidthExpr wcard

content: structure HPredFSMToBitStream {pcard : Nat}
  {tctx : Term.Ctx wcard tcard}
  {p : Term bcard ncard icard pcard tctx .prop}
  (fsm : TermFSM wcard tcard bcard ncard icard pcard
    (.ofDepTerm p)) : Prop where
  heq :
    ∀ {wenv : WidthExpr.Env wcard}
      (benv : Term.BoolEnv bcard)
      (nenv : Term.NatEnv ncard)
      (ienv : Term.IntEnv icard)
      (penv : Predicate.Env pcard) (tenv : tctx.Env wenv)
      (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
      (htenv : HTermEnv fsmEnv tenv benv) →
      (hpenv : HPredicateEnv fsmEnv penv) →
        p.toBV benv nenv ienv penv tenv  ↔ (fsm.toFsmZext.eval fsmEnv = .negOne)

content: def scanOr (s : BitStream) : BitStream := scanl false Bool.or s

content: abbrev scanl (init : Bool) (f : Bool → Bool → Bool) (s : BitStream) : BitStream :=
  fun n => match n with
    | 0 => f init (s 0)
    | n+1 => f (scanl init f s n) (s (n + 1))

content: def scanAnd (s : BitStream) : BitStream := scanl true Bool.and s

-- Repository lemmas used:
content: theorem scanOr_eq_decide (s : BitStream) (n : Nat) :
    s.scanOr n = decide (∃ (i : Nat), i ≤ n ∧ s i = true)

content: theorem scanOr_true_iff (s : BitStream) (n : Nat)
    : s.scanOr n = true ↔ ∃ (i : Nat), (i ≤ n) ∧ s i = true

content: theorem scanOr_false_iff (s : BitStream) (n : Nat) : s.scanOr n = false ↔ ∀ (i : Nat), (hi : i ≤ n) → s i = false

content: @[simp]
theorem scanOr_succ (s : BitStream) : scanOr s (n+1) = ((s.scanOr n) || s (n + 1))

content: theorem scanAnd_eq_decide (s : BitStream) (n : Nat) :
    s.scanAnd n = decide (∀ (i : Nat), i ≤ n → s i = true)

content: theorem scanAnd_true_iff (s : BitStream) (n : Nat) :
    s.scanAnd n = true ↔ ∀ (i : Nat), (hi : i ≤ n) → s i = true

content: @[simp] theorem scanAnd_succ (s : BitStream) : scanAnd s (n+1) = ((s.scanAnd n) && s (n + 1))

-- Local definitions used:
content: def NatFSM.fsmUnaryIndexUle (a : NatFSM wcard tcard bcard ncard icard pcard v) :
    FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
  composeUnaryAux (FSM.ls true) a.toFsm

content: def fsmUnaryUle (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w)) : FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
 composeUnaryAux FSM.scanAnd (b.fsmUnaryIndexUle ||| ~~~ a.fsmUnaryIndexUle)

content: def fsmUnaryNeqUpto (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w)) : FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
  composeUnaryAux FSM.scanOr (a.fsmUnaryIndexUle ^^^ b.fsmUnaryIndexUle)

content: def fsmUltUnary
    (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w)) : FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
  composeBinaryAux' FSM.and (fsmUnaryUle a b) (fsmUnaryNeqUpto a b)

-- Local lemmas used:
content: private theorem decide_eq_eq_decide_iff_decide {P Q : Prop}
  [Decidable P] [Decidable Q] :
  (decide P = decide Q) = decide (P ↔ Q)

content: private theorem not_decide_eq_decide_lnot {P : Prop}
  [Decidable P] :
    (!(decide P)) = (decide (¬ P))

content: private theorem decide_and_decide_eq_decide {P Q : Prop}
  [Decidable P] [Decidable Q] :
  (decide P && decide Q) = decide (P ∧ Q)

content: private theorem decide_or_decide_eq_decide {P Q : Prop}
  [Decidable P] [Decidable Q] :
  (decide P || decide Q) = decide (P ∨ Q)

content: @[simp]
private theorem min_eq_of_not_le {a b : Nat} (hab : ¬ a ≤ b) : min a b = b

content: @[simp]
private theorem min_eq_of_not_le' {a b : Nat} (hab : ¬ a ≤ b) : min b a = b

content: @[simp]
theorem HNatFSMToBitstream.fsmIndexUle_eval_eq
    (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    (henv : HWidthEnv fsmEnv wenv)
    (ha : HNatFSMToBitstream a) :
    (NatFSM.fsmUnaryIndexUle a).eval fsmEnv = fun i =>
    decide (i ≤ v.toNat wenv)

content: theorem eval_fsmUnaryUle_eq_decide
    (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    (henv : HWidthEnv fsmEnv wenv)
    (ha : HNatFSMToBitstream a) (hb : HNatFSMToBitstream b) :
    ((fsmUnaryUle a b).eval fsmEnv) i =
    decide (min i (v.toNat wenv) ≤ min i (w.toNat wenv))

content: @[simp]
theorem eval_fsmUnaryUle_eq_lt_or_decide
    (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    (henv : HWidthEnv fsmEnv wenv)
    (ha : HNatFSMToBitstream a) (hb : HNatFSMToBitstream b) :
    ((fsmUnaryUle a b).eval fsmEnv) i =
    decide (i ≤ min (v.toNat wenv) (w.toNat wenv) ∨ (v.toNat wenv) ≤ (w.toNat wenv))

content: @[simp]
theorem eval_fsmUnaryNeqUpto_eq_decide
    (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    (henv : HWidthEnv fsmEnv wenv)
    (ha : HNatFSMToBitstream a) (hb : HNatFSMToBitstream b) :
    ((fsmUnaryNeqUpto a b).eval fsmEnv) i =
    (decide (min i (v.toNat wenv) ≠ min i (w.toNat wenv)))

-- Target theorem and proof:
theorem eval_fsmUltUnary_eq_decide
    (a : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v))
    (b : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    (henv : HWidthEnv fsmEnv wenv)
    (ha : HNatFSMToBitstream a) (hb : HNatFSMToBitstream b) :
    ((fsmUltUnary a b).eval fsmEnv) i =
   (decide (min i (v.toNat wenv) < min i (w.toNat wenv)))
:= by
  simp [fsmUltUnary]
  rw [eval_fsmUnaryUle_eq_lt_or_decide (wenv := wenv) (henv := henv) (ha := ha) (hb := hb)]
  rw [eval_fsmUnaryNeqUpto_eq_decide (wenv := wenv) (henv := henv) (ha := ha) (hb := hb)]
  simp
  generalize v.toNat wenv = v'
  generalize w.toNat wenv = w'
  simp only [not_decide_eq_decide_lnot,
    decide_and_decide_eq_decide,
    decide_or_decide_eq_decide, decide_eq_decide]
  omega