-- Blase/Blase/KInduction/KInduction.lean

import Mathlib.Data.Bool.Basic

import Mathlib.Data.Fin.Basic

import Mathlib.Data.Finset.Basic

import Mathlib.Data.Finset.Defs

import Mathlib.Data.Multiset.FinsetOps

import Blase.Fast.BitStream

import Blase.Fast.Defs

import Blase.Fast.FiniteStateMachine

import Blase.Fast.Decide

import Blase.SingleWidth.Syntax

import Lean.Meta.ForEachExpr

import Lean.Meta.Tactic.Simp.BuiltinSimprocs.BitVec

import Blase.Fast.ForLean

import Blase.Vars

import Blase.EnvBitstream

import Lean

open Fin.NatCast

namespace ReflectVerif

open Lean Meta Elab Tactic

namespace BvDecide

open Std Sat AIG

def mkCarryAssignCircuitNAux {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (s : p.α) (n : Nat) : Circuit (Vars p.α arity (n + 1)) :=
    (p.nextStateCirc s).map fun v =>
      match v with
        | .inl t => Vars.stateN t n
        | .inr i => Vars.inputN i n

def mkCarryAssignCircuitN {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (n : Nat) :
  Circuit (Vars p.α arity (n + 1)) :=
  let carrys := FinEnum.toList p.α |>.map fun s =>
    
    Circuit.xor
      (mkCarryAssignCircuitNAux p s n)
      (Circuit.var true <| Vars.stateN s (n + 1))
  Circuit.bigOr carrys

def mkOutputAssignCircuitNAux {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (n : Nat) : Circuit (Vars p.α arity (n + 1)) :=
    (p.outputCirc).map fun v =>
      match v with
        | .inl s' => Vars.stateN s' n
        | .inr i => Vars.inputN i n

def mkOutputAssignCircuitN {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (n : Nat) :
  Circuit (Vars p.α arity (n + 1)) :=
    Circuit.xor
      (mkOutputAssignCircuitNAux p n)
      (Circuit.var true <| Vars.outputs ⟨n, by admit /- proof elided -/
      ⟩)

def mkStateNeqCircuit
  {arity : Type _} {i : Nat}
  [DecidableEq arity] [Fintype arity] [Hashable arity]
  (p : FSM arity) (s t : p.α → Circuit (Vars p.α arity i)) : Circuit (Vars p.α arity i) :=
  Circuit.bigAnd <| FinEnum.toList p.α |>.map fun a => ~~~ (s a) ^^^ (t a)

def mkStateUniqueCircuitN {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity]
  (p : FSM arity) (n : Nat) : Circuit (Vars p.α arity n) :=
  let sn : p.α → Circuit (Vars p.α arity n) := fun s =>
    Circuit.var true (Vars.stateN s n)
  let circs := (List.range n).attach |>.map fun ⟨i, hi⟩ =>
    let si : p.α → Circuit (Vars p.α arity n) := fun s =>
      Circuit.var true (Vars.stateN s i (by admit /- proof elided -/
      ))
    (mkStateNeqCircuit p si sn)
  Circuit.bigOr circs

structure KInductionCircuits {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity] (fsm : FSM arity) (n : Nat) where
  
  cInitCarryAssignCirc : Circuit (Vars fsm.α arity 0)
  
  cSuccCarryAssignCirc : Circuit (Vars fsm.α arity (n+2))
  
  cOutAssignCirc : Circuit (Vars fsm.α arity (n + 2))
  
  cStatesUniqueCirc : Circuit (Vars fsm.α arity n)

structure KInductionCircuits.IsLawful {arity : Type _}
    [DecidableEq arity] [Fintype arity] [Hashable arity] {fsm : FSM arity} {n : Nat}
  (circs : KInductionCircuits fsm n) where
  hCInitCarryAssignCirc :
    ∀ {env : Vars fsm.α arity 0 → Bool},
      (circs.cInitCarryAssignCirc.eval env = false)
      ↔ (∀ (s : fsm.α), fsm.initCarry s = env (Vars.state0 s))

  hCSuccCarryAssignCirc :
    ∀ {env : Vars fsm.α arity (n + 2) → Bool},
      (circs.cSuccCarryAssignCirc.eval env = false)
      ↔ (∀ (s : fsm.α) (i : Nat) (hi : i < n + 2),
        env (Vars.stateN s (i + 1)) =
          ((mkCarryAssignCircuitNAux fsm s i).map
            (fun v => v.castLe (by admit /- proof elided -/
            ))).eval env)
  hCOutAssignCirc :
    ∀ {env : Vars fsm.α arity (n + 2) → Bool},
      (circs.cOutAssignCirc.eval env = false)
      ↔ (∀ (i : Nat) (hi : i < n + 2),
        (fsm.outputCirc).eval
          (fun x => match x with
            | .inl s => env (Vars.stateN s i)
            | .inr j => env (Vars.inputN j i)) =
        env (Vars.outputs ⟨i, by admit /- proof elided -/
        ⟩))
  hCStatesUniqueCirc :
    ∀ {env : Vars fsm.α arity (n) → Bool},
      (circs.cStatesUniqueCirc.eval env = false)
      ↔ (∀ (i j : Nat) (hij : i < j ∧ j ≤ n),
        ∃ (s : fsm.α), env (Vars.stateN s i) ≠ env (Vars.stateN s j))

namespace KInductionCircuits

variable {arity : Type _}
  {fsm : FSM arity}

def castCircLe {n m : Nat} (c : Circuit (Vars fsm.α arity n)) (hnm : n ≤ m := by admit /- proof elided -/
) :
    Circuit (Vars fsm.α arity m) :=
  c.map (fun v => v.castLe hnm)

variable [DecidableEq arity] [Fintype arity] [Hashable arity]

def mkSucc
    (prev : KInductionCircuits fsm n) :
    KInductionCircuits fsm (n + 1) :=
  let cInitCarryAssignCirc := prev.cInitCarryAssignCirc
  { cInitCarryAssignCirc := cInitCarryAssignCirc
  , cSuccCarryAssignCirc :=
      (mkCarryAssignCircuitN fsm (n + 2)) |||
      (castCircLe prev.cSuccCarryAssignCirc)
  , cOutAssignCirc :=
      (mkOutputAssignCircuitN fsm (n + 2)) |||
      (castCircLe prev.cOutAssignCirc)
  , cStatesUniqueCirc :=
      mkStateUniqueCircuitN fsm (n + 1) |||
      (castCircLe prev.cStatesUniqueCirc)
  }

-- Repository definitions used:
content: def Vars.inputN (inp : ι) (k : Nat) {n : Nat} (hkn : k < n := by admit /- proof elided -/
) : Vars σ ι n :=
  .inputs (Inputs.mk ⟨k, by admit /- proof elided -/
  ⟩ inp)

content: inductive Vars (σ : Type) (ι : Type) (n : Nat)
| state (s : Inputs σ (n + 1))
| inputs (is : Inputs ι n)
| outputs (os : Fin n) 
deriving DecidableEq, Hashable

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: structure FSM (arity : Type) : Type 1 where
   
  ( α  : Type )
  [ i : FinEnum α ]
  [ h : Hashable α ]
  [ dec_eq : DecidableEq α ]
   
  ( initCarry : α → Bool )
   
  outputCirc : Circuit (α ⊕ arity)
  nextStateCirc : α → Circuit (α ⊕ arity)

content: def Vars.stateN (s : σ) (i : Nat) {n : Nat} (hin : i ≤ n := by admit /- proof elided -/
) : Vars σ ι n :=
  .state (Inputs.mk ⟨i, by admit /- proof elided -/
  ⟩ s)

content: structure Inputs (ι : Type) (n : Nat) : Type  where
  ix : Fin n
  input : ι
deriving DecidableEq, Hashable

content: def Vars.state0 (s : σ) {n : Nat} : Vars σ ι n :=
  .state (Inputs.mk ⟨0, by admit /- proof elided -/
  ⟩ s)

content: def bigOr {α : Type _}
    (cs : List (Circuit α)) : Circuit α :=
  match cs with
  | [] => Circuit.fals
  | c :: cs =>
    c ||| (Circuit.bigOr cs)

content: def bigAnd {α : Type _}
    (cs : List (Circuit α)) : Circuit α :=
  match cs with
  | [] => Circuit.tru
  | c :: cs =>
    c &&& (Circuit.bigAnd cs)

content: @[simp]
def eval : Circuit α → (α → Bool) → Bool
  | tru, _ => true
  | fals, _ => false
  | var b x, f => if b then f x else !(f x)
  | and c₁ c₂, f => (eval c₁ f) && (eval c₂ f)
  | or c₁ c₂, f => (eval c₁ f) || (eval c₂ f)
  | xor c₁ c₂, f => Bool.xor (eval c₁ f) (eval c₂ f)

content: def map : ∀ (_c : Circuit α) (_f : α → β), Circuit β
  | tru, _ => tru
  | fals, _ => fals
  | var b x, f => var b (f x)
  | and c₁ c₂, f => (map c₁ f) &&& (map c₂ f)
  | or c₁ c₂, f => (map c₁ f) ||| (map c₂ f)
  | xor c₁ c₂, f => (map c₁ f) ^^^ (map c₂ f)

content: def Vars.castLe {n m : Nat} (v : Vars σ ι n) (hnm : n ≤ m) : Vars σ ι m :=
  match v with
  | .state ss => .state (ss.castLe (by admit /- proof elided -/
  ))
  | .inputs is => .inputs (is.castLe hnm)
  | .outputs os =>
    .outputs (os.castLE (by admit /- proof elided -/
    ))

content: def castLe (i : Inputs ι n) (hn : n ≤ m) : Inputs ι m where
  ix := ⟨i.ix, by admit /- proof elided -/
  ⟩
  input := i.input

-- Repository lemmas used:
content: lemma eval_map {c : Circuit α} {f : α → β} {g : β → Bool} :
    eval (map c f) g = eval c (λ x => g (f x))

content: @[simp]
theorem eval_bigOr_eq_false_iff
    (cs : List (Circuit α)) (env : α → Bool):
    (Circuit.bigOr cs).eval env = false ↔
    (∀ (c : Circuit α), c ∈ cs → c.eval env = false)

content: @[simp]
theorem Vars.castLe_inputs_eq_inputs {n i m : Nat} (hi : i < n) (hnm : n ≤ m) :
  (Vars.inputN inp i hi : Vars σ ι n).castLe hnm =
  Vars.inputN inp i (by omega)

content: @[simp]
theorem Vars.castLe_stateN_eq_stateN  {n i m : Nat} (hi : i ≤ n) (hnm : n ≤ m) :
  (Vars.stateN s i : Vars σ ι n).castLe hnm =
  Vars.stateN s i (hin := by omega)

content: @[simp]
theorem Vars.castLe_castLe_eq_castLe_self {α : Type _} {p q r : Nat}
  (v : Vars α σ p) (h : p ≤ q) (h' : q ≤ r) :
    (v.castLe h).castLe h' = v.castLe (by omega)

-- Local definitions used:
content: def mkCarryAssignCircuitNAux {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (s : p.α) (n : Nat) : Circuit (Vars p.α arity (n + 1)) :=
    (p.nextStateCirc s).map fun v =>
      match v with
        | .inl t => Vars.stateN t n
        | .inr i => Vars.inputN i n

content: def mkCarryAssignCircuitN {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (n : Nat) :
  Circuit (Vars p.α arity (n + 1)) :=
  let carrys := FinEnum.toList p.α |>.map fun s =>
    
    Circuit.xor
      (mkCarryAssignCircuitNAux p s n)
      (Circuit.var true <| Vars.stateN s (n + 1))
  Circuit.bigOr carrys

content: def mkOutputAssignCircuitNAux {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (n : Nat) : Circuit (Vars p.α arity (n + 1)) :=
    (p.outputCirc).map fun v =>
      match v with
        | .inl s' => Vars.stateN s' n
        | .inr i => Vars.inputN i n

content: def mkOutputAssignCircuitN {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (n : Nat) :
  Circuit (Vars p.α arity (n + 1)) :=
    Circuit.xor
      (mkOutputAssignCircuitNAux p n)
      (Circuit.var true <| Vars.outputs ⟨n, by admit /- proof elided -/
      ⟩)

content: def mkStateNeqCircuit
  {arity : Type _} {i : Nat}
  [DecidableEq arity] [Fintype arity] [Hashable arity]
  (p : FSM arity) (s t : p.α → Circuit (Vars p.α arity i)) : Circuit (Vars p.α arity i) :=
  Circuit.bigAnd <| FinEnum.toList p.α |>.map fun a => ~~~ (s a) ^^^ (t a)

content: def mkStateUniqueCircuitN {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity]
  (p : FSM arity) (n : Nat) : Circuit (Vars p.α arity n) :=
  let sn : p.α → Circuit (Vars p.α arity n) := fun s =>
    Circuit.var true (Vars.stateN s n)
  let circs := (List.range n).attach |>.map fun ⟨i, hi⟩ =>
    let si : p.α → Circuit (Vars p.α arity n) := fun s =>
      Circuit.var true (Vars.stateN s i (by admit /- proof elided -/
      ))
    (mkStateNeqCircuit p si sn)
  Circuit.bigOr circs

content: structure KInductionCircuits {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity] (fsm : FSM arity) (n : Nat) where
  
  cInitCarryAssignCirc : Circuit (Vars fsm.α arity 0)
  
  cSuccCarryAssignCirc : Circuit (Vars fsm.α arity (n+2))
  
  cOutAssignCirc : Circuit (Vars fsm.α arity (n + 2))
  
  cStatesUniqueCirc : Circuit (Vars fsm.α arity n)

content: structure KInductionCircuits.IsLawful {arity : Type _}
    [DecidableEq arity] [Fintype arity] [Hashable arity] {fsm : FSM arity} {n : Nat}
  (circs : KInductionCircuits fsm n) where
  hCInitCarryAssignCirc :
    ∀ {env : Vars fsm.α arity 0 → Bool},
      (circs.cInitCarryAssignCirc.eval env = false)
      ↔ (∀ (s : fsm.α), fsm.initCarry s = env (Vars.state0 s))

  hCSuccCarryAssignCirc :
    ∀ {env : Vars fsm.α arity (n + 2) → Bool},
      (circs.cSuccCarryAssignCirc.eval env = false)
      ↔ (∀ (s : fsm.α) (i : Nat) (hi : i < n + 2),
        env (Vars.stateN s (i + 1)) =
          ((mkCarryAssignCircuitNAux fsm s i).map
            (fun v => v.castLe (by admit /- proof elided -/
            ))).eval env)
  hCOutAssignCirc :
    ∀ {env : Vars fsm.α arity (n + 2) → Bool},
      (circs.cOutAssignCirc.eval env = false)
      ↔ (∀ (i : Nat) (hi : i < n + 2),
        (fsm.outputCirc).eval
          (fun x => match x with
            | .inl s => env (Vars.stateN s i)
            | .inr j => env (Vars.inputN j i)) =
        env (Vars.outputs ⟨i, by admit /- proof elided -/
        ⟩))
  hCStatesUniqueCirc :
    ∀ {env : Vars fsm.α arity (n) → Bool},
      (circs.cStatesUniqueCirc.eval env = false)
      ↔ (∀ (i j : Nat) (hij : i < j ∧ j ≤ n),
        ∃ (s : fsm.α), env (Vars.stateN s i) ≠ env (Vars.stateN s j))

content: def castCircLe {n m : Nat} (c : Circuit (Vars fsm.α arity n)) (hnm : n ≤ m := by admit /- proof elided -/
) :
    Circuit (Vars fsm.α arity m) :=
  c.map (fun v => v.castLe hnm)

content: def mkSucc
    (prev : KInductionCircuits fsm n) :
    KInductionCircuits fsm (n + 1) :=
  let cInitCarryAssignCirc := prev.cInitCarryAssignCirc
  { cInitCarryAssignCirc := cInitCarryAssignCirc
  , cSuccCarryAssignCirc :=
      (mkCarryAssignCircuitN fsm (n + 2)) |||
      (castCircLe prev.cSuccCarryAssignCirc)
  , cOutAssignCirc :=
      (mkOutputAssignCircuitN fsm (n + 2)) |||
      (castCircLe prev.cOutAssignCirc)
  , cStatesUniqueCirc :=
      mkStateUniqueCircuitN fsm (n + 1) |||
      (castCircLe prev.cStatesUniqueCirc)
  }

-- Local lemmas used:
content: @[simp]
theorem mkCarryAssignCircuitNAux_eval_eq {arity : Type _}
    [DecidableEq arity]
    [Fintype arity]
    [Hashable arity]
    (p : FSM arity) (s : p.α) (n : Nat)
    {env : Vars p.α arity (n + 1) → Bool} :
    ((mkCarryAssignCircuitNAux p s n).eval env) = ((p.nextStateCirc s).eval
      (fun x => match x with | .inl x => env (Vars.stateN x n) | .inr x => env (Vars.inputN x n)))

content: theorem mkStateUniqueCircuitN_eq_false_iff {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity]
  (p : FSM arity) (n : Nat)
  {env : Vars p.α arity n → Bool} :
  ((mkStateUniqueCircuitN p n).eval env = false) ↔
  (∀ (i : Nat) (hi : i < n), ∃ (s : p.α), env (Vars.stateN s i) ≠ env (Vars.stateN s n))

-- Target theorem and proof:
theorem IsLawful_mkSucc_of_IsLawful {arity : Type _}
    [DecidableEq arity] [Fintype arity] [Hashable arity]
    {fsm : FSM arity} {n : Nat}
    (prev : KInductionCircuits fsm n)
    (hPrev : prev.IsLawful) :
    (mkSucc prev).IsLawful
where
  hCInitCarryAssignCirc := by
    simp only [mkSucc, castCircLe]
    exact hPrev.hCInitCarryAssignCirc
  hCSuccCarryAssignCirc := by
    simp only [mkSucc, castCircLe, Circuit.eval_map]
    simp [Circuit.eval_map]
    intros env
    constructor
    · intros h s i hi
      obtain ⟨h₁, h₂⟩ := h
      rw [hPrev.hCSuccCarryAssignCirc] at h₂
      by_cases hi : i < n + 2
      · simp only [Vars.castLe_stateN_eq_stateN, Circuit.eval_map,
        Vars.castLe_castLe_eq_castLe_self, mkCarryAssignCircuitNAux_eval_eq,
        Vars.castLe_inputs_eq_inputs] at h₂
        rw [h₂ s i hi]
      · have hi : i = n + 2 := by omega
        subst hi
        apply h₁
    · intros h
      constructor
      · intros s
        simp only at h ⊢
        rw [h s _ (by omega)]
      · rw [hPrev.hCSuccCarryAssignCirc]
        intros s i hi
        simp only at h
        simp only [Vars.castLe_stateN_eq_stateN, Circuit.eval_map,
          Vars.castLe_castLe_eq_castLe_self, mkCarryAssignCircuitNAux_eval_eq,
          Vars.castLe_inputs_eq_inputs]
        rw [h s i (by omega)]

  hCOutAssignCirc := by
    simp only [mkSucc, castCircLe]
    simp [Circuit.eval_map]
    intros env
    constructor
    · intros h i hi
      obtain ⟨h₁, h₂⟩ := h
      rw [hPrev.hCOutAssignCirc] at h₂
      by_cases hi : i < n + 2
      · simp at h₁ h₂ ⊢
        rw [h₂ i hi]
      · have : i = n + 2 := by omega
        subst this
        apply h₁
    · intros h
      constructor
      · simp only at h ⊢
        rw [h]
      · rw [hPrev.hCOutAssignCirc]
        intros i hi
        simp
        apply h
  hCStatesUniqueCirc := by
    simp only [mkSucc, castCircLe]
    simp [Circuit.eval_map]
    intros env
    constructor
    · intros h i j hij hjn
      simp [mkStateUniqueCircuitN_eq_false_iff] at h
      simp [hPrev.hCStatesUniqueCirc] at h
      obtain ⟨h₁, h₂⟩ := h
      by_cases hj : j ≤ n
      · grind
      · have : j = n + 1 := by
          omega
        subst this
        grind
    · intros h
      constructor
      · rw [mkStateUniqueCircuitN_eq_false_iff]
        intros i hi
        apply h (j := n + 1)
        omega
        omega
      · simp [hPrev.hCStatesUniqueCirc]
        intros i j hij jn
        simp at h ⊢
        apply h i j
        omega
        omega