-- Clean/Circuit/Loops.lean

import Clean.Circuit.Subcircuit

import Clean.Utils.Misc

variable {n m : ℕ} {F : Type} [Field F] {α β : Type}

namespace Circuit

variable {prop : Condition F}

namespace ForM

variable {circuit : α → Circuit F Unit} (xs : Vector α m) (constant : ConstantLength circuit) (n : ℕ)

end ForM

namespace MapM

variable {circuit : α → Circuit F β} {xs : Vector α m} [constant: ConstantLength circuit]
  {prop : Condition F}

end MapM

namespace FoldlM

@[reducible]
def prod (circuit : β → α → Circuit F β) : β × α → Circuit F β := fun t => circuit t.1 t.2

variable {env : Environment F} {prop : Condition F} {xs : Vector α m}
  {circuit : β → α → Circuit F β} {init : β} {constant : ConstantLength (prod circuit)}

def foldlAcc (n : ℕ) (xs : Vector α m) (circuit : β → α → Circuit F β) (init : β) (j : Fin m) : β :=
  Fin.foldl j (fun acc i => (circuit acc xs[i.val]).output (n + i*(circuit acc xs[i.val]).localLength)) init

variable {prop : Condition F}

section

variable {env : Environment F} {prop : Condition F} {m : ℕ}
  {Acc : ℕ → Type}
  {circuit : β → Fin m → Circuit F β} {init : β} {constant : ConstantLength (prod circuit)}

end

variable [Inhabited β] [Inhabited α] {h_const_out : ConstantOutput fun (t : β × α) => circuit t.1 t.2}

-- Repository definitions used:
content: def Circuit (F : Type) [Field F] (α : Type) := ℕ → α × List (Operation F)

content: inductive Operation (F : Type) [Field F] where
  | witness : (m : ℕ) → (compute : Environment F → Vector F m) → Operation F
  | assert : Expression F → Operation F
  | lookup : Lookup F → Operation F
  | subcircuit : {n : ℕ} → Subcircuit F n → Operation F

content: structure Condition (F : Type) [Field F] where
  witness (offset : ℕ) : (m : ℕ) → (Environment F → Vector F m) → Prop := fun _ _ => True
  assert (offset : ℕ) (_ : Expression F) : Prop := True
  lookup (offset : ℕ) (_ : Lookup F) : Prop := True
  subcircuit (offset : ℕ) {m : ℕ} (_ : Subcircuit F m) : Prop := True

content: inductive FlatOperation (F : Type) where
  | witness : (m : ℕ) → (Environment F → Vector F m) → FlatOperation F
  | assert : Expression F → FlatOperation F
  | lookup : Lookup F → FlatOperation F

content: structure Subcircuit (F : Type) [Field F] (offset : ℕ) where
  ops : List (FlatOperation F)

  
  
  
  Soundness : Environment F → Prop
  Completeness : Environment F → Prop
  UsesLocalWitnesses : Environment F → Prop

  
  
  localLength : ℕ

  
  imply_soundness : ∀ env,
    ConstraintsHoldFlat env ops → Soundness env

  
  implied_by_completeness : ∀ env, env.ExtendsVector (localWitnesses env ops) offset →
    Completeness env → ConstraintsHoldFlat env ops

  
  imply_usesLocalWitnesses : ∀ env, env.ExtendsVector (localWitnesses env ops) offset →
    UsesLocalWitnesses env

  
  localLength_eq : localLength = FlatOperation.localLength ops

content: @[circuit_norm]
def localWitnesses (env : Environment F) : (l : List (FlatOperation F)) → Vector F (localLength l)
  | [] => #v[]
  | witness _ compute :: ops => compute env ++ localWitnesses env ops
  | assert _ :: ops | lookup _ :: ops => localWitnesses env ops

content: @[circuit_norm]
def lookup {Row : TypeMap} [ProvableType Row] (table : Table F Row)  (entry : Row (Expression F)) : Circuit F Unit := fun _ =>
  ((), [.lookup { table := table.toRaw, entry := toElements entry }])

content: @[circuit_norm]
def Table.toRaw (table : Table F Row) : RawTable F where
  name := table.name
  arity := size Row
  Contains row := table.Contains (fromElements row)
  Soundness row := table.Soundness (fromElements row)
  Completeness row := table.Completeness (fromElements row)
  imply_soundness row := table.imply_soundness (fromElements row)
  implied_by_completeness row := table.implied_by_completeness (fromElements row)

content: structure RawTable (F : Type) where
  name : String
  arity : ℕ
  Contains : Vector F arity → Prop
  Soundness : Vector F arity → Prop
  Completeness : Vector F arity → Prop
  imply_soundness : ∀ row, Contains row → Soundness row
  implied_by_completeness : ∀ row, Completeness row → Contains row

content: structure StaticTable (F : Type) (Row : TypeMap) [ProvableType Row] where
  name : String
  length : ℕ
  row : Fin length → Row F
  
  
  
  index : Row F → ℕ
  Spec : Row F → Prop
  contains_iff : ∀ t, (∃ i, t = row i) ↔ Spec t

content: def Contains (table : StaticTable F Row) (row : Row F) :=
  ∃ i : Fin table.length, row = table.row i

content: instance ProvableType.fromStruct {α : TypeMap} [ProvableStruct α] : ProvableType α where
  size := combinedSize α
  toElements x :=
    toComponents x |> componentsToElements (components α) |>.cast combinedSize_eq.symm
  fromElements v :=
    v.cast combinedSize_eq |> componentsFromElements (components α) |> fromComponents
  fromElements_toElements x := by admit /- proof elided -/

content: @[reducible]
def TypeMap := Type → Type

content: structure Table (F : Type) (Row : TypeMap) [ProvableType Row] where
  name : String
   
  Contains : Row F → Prop

   
  Soundness : Row F → Prop
  Completeness : Row F → Prop

  imply_soundness : ∀ row, Contains row → Soundness row
  implied_by_completeness : ∀ row, Completeness row → Contains row

content: inductive Expression (F : Type) where
  | var : Variable F -> Expression F
  | const : F -> Expression F
  | add : Expression F -> Expression F -> Expression F
  | mul : Expression F -> Expression F -> Expression F

content: structure Variable (F : Type) where
  index : ℕ

content: def const (x : α F) : Var α F :=
  let values : Vector F _ := toElements x
  fromVars (values.map .const)

content: @[reducible] def Var (M : TypeMap) (F : Type) := M (Expression F)

content: @[circuit_norm]
def fromVars (vars : Vector (Expression F) (size M)) := fromElements vars

content: structure Lookup (F : Type) where
  table : RawTable F
  entry : Vector (Expression F) table.arity

content: class Witnessable (F : Type) [Field F] (value : outParam TypeMap) (var : TypeMap) [ProvableType value] where
  witness : ((Environment F) → value F) → Circuit F (var F)
  var_eq : var F = value (Expression F) := by admit /- proof elided -/

content: structure Environment (F : Type) where
  get : ℕ → F

content: @[circuit_norm]
def localLength : List (FlatOperation F) → ℕ
  | [] => 0
  | witness m _ :: ops => m + localLength ops
  | assert _ :: ops | lookup _ :: ops => localLength ops

content: def ConstraintsHoldFlat (eval : Environment F) : List (FlatOperation F) → Prop
  | [] => True
  | op :: ops => match op with
    | assert e => (eval e = 0) ∧ ConstraintsHoldFlat eval ops
    | lookup { table, entry } =>
      table.Contains (entry.map eval) ∧ ConstraintsHoldFlat eval ops
    | _ => ConstraintsHoldFlat eval ops

content: class ConstantLength (circuit : α → Circuit F β) where
  localLength : ℕ
  localLength_eq : ∀ (a : α) (n : ℕ), (circuit a).localLength n = localLength

content: @[circuit_norm]
def ConstantOutput (circuit : α → Circuit F β) [Inhabited α] :=
  ∀ (x : α) (n : ℕ), (circuit x).output n = (circuit default).output n

content: def induct {motive : {n : ℕ} → Vector α n → Sort u}
    (nil : motive #v[])
    (cons: ∀ {n : ℕ} (a : α) (as : Vector α n), motive as → motive (cons a as))
    {n : ℕ} (v : Vector α n) : motive v :=
  match v with
  | ⟨ .mk [], h ⟩ => by admit /- proof elided -/
  | ⟨ .mk (a :: as), h ⟩ => by admit /- proof elided -/

content: def cons (a : α) (v : Vector α n) : Vector α (n + 1) :=
  ⟨ .mk (a :: v.toList), by admit /- proof elided -/
  ⟩

content: def toList : Operations F → List (Operation F) := id

content: @[reducible, circuit_norm]
def Operations (F : Type) [Field F] := List (Operation F)

content: @[reducible, circuit_norm]
def operations (circuit : Circuit F α) (offset : ℕ) : Operations F :=
  (circuit offset).2

content: @[circuit_norm]
def localLength : Operation F → ℕ
  | .witness m _ => m
  | .assert _ => 0
  | .lookup _ => 0
  | .subcircuit s => s.localLength

content: @[circuit_norm]
def localLength : Operations F → ℕ
  | [] => 0
  | .witness m _ :: ops => m + localLength ops
  | .assert _ :: ops => localLength ops
  | .lookup _ :: ops => localLength ops
  | .subcircuit s :: ops => s.localLength + localLength ops

content: @[reducible, circuit_norm]
def output (circuit : Circuit F α) (offset : ℕ) : α :=
  (circuit offset).1

content: class ElaboratedCircuit (F : Type) (Input Output : TypeMap) [Field F] [ProvableType Input] [ProvableType Output] where
  main : Var Input F → Circuit F (Var Output F)

   
  localLength : Var Input F → ℕ

   
  localLength_eq : ∀ input offset, (main input).localLength offset = localLength input
    := by admit /- proof elided -/

content: @[reducible, circuit_norm]
def localLength (circuit : Circuit F α) (offset := 0) : ℕ :=
  Operations.localLength (circuit offset).2

-- Repository lemmas used:
content: theorem Fin.foldl_const (n : ℕ) (f : Fin n → α) (init : α) :
  Fin.foldl n (fun _ i => f i) init = match n with

content: theorem Fin.foldl_const_succ (n : ℕ) (f : Fin (n + 1) → α) (init : α) :
    Fin.foldl (n + 1) (fun _ i => f i) init = f (.last n)

content: theorem bind_operations_eq (f : Circuit F α) (g : α → Circuit F β) (n : ℕ) :
  (f >>= g).operations n = f.operations n ++ (g (f.output n)).operations (n + f.localLength n)

-- Local definitions used:
content: @[reducible]
def prod (circuit : β → α → Circuit F β) : β × α → Circuit F β := fun t => circuit t.1 t.2

content: def foldlAcc (n : ℕ) (xs : Vector α m) (circuit : β → α → Circuit F β) (init : β) (j : Fin m) : β :=
  Fin.foldl j (fun acc i => (circuit acc xs[i.val]).output (n + i*(circuit acc xs[i.val]).localLength)) init

-- Local lemmas used:
content: lemma Vector.foldlM_toList (xs : Vector α n) {m : Type → Type} [Monad m] (body : β → α → m β) (init : β) :
    xs.foldlM body init = xs.toList.foldlM body init

content: lemma ConstantLength.length_eq_default {circuit : α → Circuit F β} (_ : ConstantLength circuit) [Inhabited α] (a : α) (n : ℕ) :
   (circuit a).localLength n = (circuit default).localLength 0

content: lemma foldlM_cons (x : α) :
  (Vector.cons x xs).foldlM circuit init = (do
    let init' ← circuit init x
    xs.foldlM circuit init')

content: lemma foldlAcc_zero [NeZero m] : foldlAcc n xs circuit init 0 = init

content: lemma foldlAcc_cons_succ (i : Fin m) (x : α) [constant : ConstantLength (prod circuit)] :
  foldlAcc n (Vector.cons x xs) circuit init i.succ =
    foldlAcc (n + (circuit init x).localLength n) xs circuit ((circuit init x).output n) i

content: theorem operations_eq :
  (Vector.foldlM circuit init xs).operations n =
    (List.ofFn fun i => (circuit (foldlAcc n xs circuit init i) xs[i.val]).operations (n + i * constant.localLength)).flatten

content: theorem foldlAcc_const_succ (constant : ConstantLength (prod circuit))
  (h_const_out : ConstantOutput fun (t : β × α) => circuit t.1 t.2)
  (i : ℕ) (hi : i + 1 < m) :
  foldlAcc n xs circuit init ⟨ i + 1, hi ⟩ =
    (circuit default xs[i]).output (n + i * (circuit default default).localLength)

-- Target theorem and proof:
theorem operations_eq_const [NeZero m] (constant : ConstantLength (prod circuit))
    (h_const_out : ConstantOutput (prod circuit)) :
  (Vector.foldlM circuit init xs).operations n =
  (circuit init (xs[0]'(NeZero.pos m))).operations n ++
  (List.ofFn fun (⟨i, _⟩ : Fin (m - 1)) =>
    let k := (circuit default default).localLength
    let acc := (circuit default xs[i]).output (n + i*k)
    (circuit acc xs[i + 1]).operations (n + (i + 1)*k)).flatten
:= by
  rw [operations_eq]
  simp only
  set k := (circuit default default).localLength
  set k' := constant.localLength
  have : k' = k := by simp only [k]; rw [constant.localLength_eq (_, _)]
  rw [this]
  rcases m with rfl | m
  · nomatch ‹NeZero 0›
  rw [List.ofFn_succ, List.flatten_cons]
  simp only [foldlAcc_zero, Fin.val_zero, zero_mul, add_zero, Fin.val_succ, add_mul, one_mul,
    add_tsub_cancel_right, Nat.add_one_sub_one, Fin.cast_eq_self, List.append_cancel_left_eq]
  congr
  funext i
  rcases i with ⟨ i, hi ⟩
  simp only [Fin.succ_mk]
  rw [foldlAcc_const_succ constant h_const_out]