-- FormalSnarksProject/SNARKs/ToySnark.lean

import FormalSnarksProject.Models.AGMProofSystemInstantiation

import Mathlib.Algebra.Polynomial.Div

import FormalSnarksProject.ToMathlib.OptionEquivRight

import Mathlib.Algebra.MvPolynomial.Equiv

import FormalSnarksProject.SoundnessTactic.SoundnessProver

open scoped BigOperators Classical

section ToySnark

open MvPolynomial Option List

namespace ToySnark

inductive Vars : Type where
  | α : Vars
  | β : Vars
deriving Repr, BEq

inductive StmtEntries : Type where
  | x : StmtEntries
  | y : StmtEntries
  | z : StmtEntries
deriving Repr, BEq

inductive WitEntries : Type where
  | A : WitEntries
  | B : WitEntries
deriving Repr, BEq

local notation "Vars_α" => some Vars.α
local notation "Vars_β" => some Vars.β
local notation "Vars_x" => none

inductive Proof_G1_Idx : Type where
  | Pf : Proof_G1_Idx

def Proof_G2_Idx : Type := Empty

inductive PairingsIdx : Type where
  | lhs : PairingsIdx
  | rhs : PairingsIdx

inductive SRS_Elements_G2_Idx : Type where
  | α : SRS_Elements_G2_Idx
  | β : SRS_Elements_G2_Idx

noncomputable def ToySnark
     
    {F : Type} [Field F] :
    AGMProofSystemInstantiation F :=
  {
    Stmt := StmtEntries -> F
    Sample := Option Vars
    SRSElements_G1 := SRS_Elements_G1_Idx
    ListSRSElements_G1 :=
      [.α, .β]
    SRSElements_G2 := SRS_Elements_G2_Idx
    ListSRSElements_G2 :=
      [.α, .β]
    SRSElementValue_G1 := fun SRS_idx => match SRS_idx with
      | SRS_Elements_G1_Idx.α => MvPolynomial.X Vars_α
      | SRS_Elements_G1_Idx.β => MvPolynomial.X Vars_β
    SRSElementValue_G2 := fun SRS_idx => match SRS_idx with
      | SRS_Elements_G2_Idx.α => MvPolynomial.X Vars_α
      | SRS_Elements_G2_Idx.β => MvPolynomial.X Vars_β
    Proof_G1 := Proof_G1_Idx
    ListProof_G1 := [Proof_G1_Idx.Pf]
    Proof_G2 := Proof_G2_Idx
    ListProof_G2 := []
    EqualityChecks := Unit
    Pairings := fun _ => PairingsIdx
    ListPairings := fun _ => [PairingsIdx.lhs, PairingsIdx.rhs]
    verificationPairingSRS_G1 := fun stmt _ i SRS_idx => match i with
      | PairingsIdx.lhs => 0
      | PairingsIdx.rhs => match SRS_idx with
        | SRS_Elements_G1_Idx.α => stmt StmtEntries.z
        | SRS_Elements_G1_Idx.β => 0
    verificationPairingSRS_G2 := fun stmt _ i SRS_idx => match i with
      | PairingsIdx.lhs => match SRS_idx with
        | SRS_Elements_G2_Idx.α => stmt StmtEntries.x
        | SRS_Elements_G2_Idx.β => stmt StmtEntries.y
      | PairingsIdx.rhs => match SRS_idx with
        | SRS_Elements_G2_Idx.α => 0
        | SRS_Elements_G2_Idx.β => -1
    verificationPairingProof_G1 := fun stmt _ i pf => match i with
      | PairingsIdx.lhs => match pf with
        | Proof_G1_Idx.Pf => 1
      | PairingsIdx.rhs => match pf with
        | Proof_G1_Idx.Pf => 0
    verificationPairingProof_G2 := fun stmt _ i pf => 0
  }

section soundness

-- Repository definitions used:
content: syntax "integral_domain_tactic" : tactic

content: macro_rules
| `(tactic| simplify_mvpoly_option_eqn) =>
  `(tactic|
    simp only [monomial_zero', List.singleton_append, List.cons_append, List.append_assoc,
      List.map_cons, Sum.elim_inl, Sum.elim_inr, List.map_append, List.map_map, List.sum_cons,
      List.sum_append, List.map_nil, List.sum_nil, add_zero, Sum.elim_lam_const_lam_const, map_one,
      one_mul, map_zero, zero_mul, map_neg, neg_mul, neg_add_rev, zero_add, mul_zero,
      
      Function.comp, List.sum_map_zero] at eqn;
    simp only [mul_add, add_mul, List.sum_map_add] at eqn;
    
    simp only [
      
      mul_assoc,
      
      mul_left_comm (C _) (X (some _)) _, mul_left_comm (List.sum _) (X (some _)) _,
      mul_comm (C _) (X (some _)), mul_comm (List.sum _) (X (some _)),
      
      neg_mul, mul_neg,
      
      List.sum_map_mul_right, List.sum_map_mul_left] at eqn;

    
    trace "Converting to MvPolynomial over Polynomials";
    replace eqn := congr_arg (MvPolynomial.optionEquivRight F Vars) eqn;
    simp only [AlgEquiv.map_add, AlgEquiv.map_zero, AlgEquiv.map_mul, AlgEquiv.map_one,
      AlgEquiv.map_neg, AlgEquiv.list_map_sum, AlgEquiv.map_pow] at eqn;
    simp only [MvPolynomial.optionEquivRight_C, MvPolynomial.optionEquivRight_X_none, MvPolynomial.optionEquivRight_X_some, optionEquivRight_to_MvPolynomial_Option] at eqn;

    
    simp only [←MvPolynomial.C_mul, ←MvPolynomial.C_pow, ←MvPolynomial.C_add,
      MvPolynomial.sum_map_C] at eqn;

    simp only [MvPolynomial.X, C_apply, MvPolynomial.monomial_mul, one_mul, mul_one, add_zero, zero_add, mul_add, add_mul] at eqn
  )

content: macro_rules
| `(tactic| polynomial_ext) =>
  `(tactic| sorry) 

content: macro_rules
| `(tactic| integral_domain_tactic) =>
  `(tactic|
      trace "Call to integral_domain_tactic";
      
      
      
      simp_all (config := {decide := false, failIfUnchanged := false}) only [
            
            false_or, or_false, true_or, or_true, not_true, not_false_iff,
            
            add_zero, zero_add, mul_zero, zero_mul, mul_one, one_mul, neg_zero,
            
            neg_eq_zero, add_eq_zero_iff_eq_neg,
            eq_self_iff_true, eq_zero_of_zero_eq, one_ne_zero, mul_ne_zero_iff, zero_sub_eq_iff,
            
            mul_eq_zero];
      first
        
        | done
        
        | cases_or _ ∨ _
          all_goals integral_domain_tactic
        
        | skip
   )

content: inductive Vars : Type where
  | α : Vars
  | β_v : Vars
  | β_w : Vars
  | β_y : Vars
  | γ : Vars
deriving Repr, BEq

local notation "poly_α" => X (some Vars.α)
local notation "poly_β_v" => X (some Vars.β_v)
local notation "poly_β_w" => X (some Vars.β_w)
local notation "poly_β_y" => X (some Vars.β_y)
local notation "poly_γ" => X (some Vars.γ)
local notation "poly_s" => X (none)

content: inductive Vars : Type where
  | r_v : Vars
  | r_w : Vars
  | α_v : Vars
  | α_w : Vars
  | α_y : Vars
  | β : Vars
  | γ : Vars
deriving Repr, BEq

local notation "poly_r_v" => X (some Vars.r_v)
local notation "poly_r_w" => X (some Vars.r_w)
local notation "poly_α_v" => X (some Vars.α_v)
local notation "poly_α_w" => X (some Vars.α_w)
local notation "poly_α_y" => X (some Vars.α_y)
local notation "poly_β" => X (some Vars.β)
local notation "poly_γ" => X (some Vars.γ)
local notation "poly_s" => X (none)

content: inductive SRS_Elements_G1_Idx {n_stmt n_wit n_var : ℕ} : Type where
  | α : SRS_Elements_G1_Idx
  | β : SRS_Elements_G1_Idx
  | δ : SRS_Elements_G1_Idx
  | x_pow : Fin n_var → SRS_Elements_G1_Idx
  | x_pow_times_t : Fin (n_var - 1) → SRS_Elements_G1_Idx
  | y : Fin n_stmt → SRS_Elements_G1_Idx
  | q : Fin n_wit → SRS_Elements_G1_Idx

content: inductive SRS_Elements_G2_Idx {n_stmt n_wit n_var : ℕ} : Type where
  | β : SRS_Elements_G2_Idx
  | γ : SRS_Elements_G2_Idx
  | δ : SRS_Elements_G2_Idx
  | x_pow : Fin n_var → SRS_Elements_G2_Idx

content: structure AGMProofSystemInstantiation (F : Type) [Field F] where
  Stmt Sample SRSElements Proof EqualityChecks : Type

  ListSRSElements : List SRSElements
  SRSElementValue : SRSElements → MvPolynomial Sample F
  ListProof : List Proof
  Pairings : EqualityChecks → Type
  ListPairings : (k : EqualityChecks) → List (Pairings k)
  verificationPairingSRS_G1 : Stmt -> (k : EqualityChecks) → Pairings k → SRSElements → F
  verificationPairingSRS_G2 : Stmt -> (k : EqualityChecks) → Pairings k → SRSElements → F
  verificationPairingProof_G1 : Stmt -> (k : EqualityChecks) → Pairings k → Proof → F
  verificationPairingProof_G2 : Stmt -> (k : EqualityChecks) → Pairings k → Proof → F

content: noncomputable def proof_element_G1_as_poly {F : Type} [Field F]
    (𝓟 : AGMProofSystemInstantiation F) (prover : 𝓟.Prover) (pf_elem : 𝓟.Proof_G1) :
    MvPolynomial (𝓟.Sample) F :=
  (𝓟.ListSRSElements_G1.map fun SRS_elem =>
          MvPolynomial.C (prover.fst pf_elem SRS_elem) * (𝓟.SRSElementValue_G1 SRS_elem)).sum

content: def Prover (F : Type) [Field F]
    (𝓟 : AGMProofSystemInstantiation F) : Type :=
  (𝓟.Proof -> 𝓟.SRSElements -> F) × (𝓟.Proof -> 𝓟.SRSElements -> F)

content: noncomputable def proof_element_G2_as_poly {F : Type} [Field F]
    (𝓟 : AGMProofSystemInstantiation F) (prover : 𝓟.Prover) (pf_elem : 𝓟.Proof_G2) :
    MvPolynomial (𝓟.Sample) F :=
  (𝓟.ListSRSElements_G2.map fun SRS_elem =>
          MvPolynomial.C (prover.snd pf_elem SRS_elem) * (𝓟.SRSElementValue_G2 SRS_elem)).sum

content: def soundness (F : Type) [Field F]
    (𝓟 : AGMProofSystemInstantiation F)
    (Wit : Type) (relation : 𝓟.Stmt -> Wit -> Prop)
    (extractor : 𝓟.Prover -> Wit) : Prop :=
   ∀ stmt : 𝓟.Stmt,
    ∀ prover : 𝓟.Prover,
      𝓟.verify prover stmt -> relation stmt (extractor prover)

content: def verify {F : Type} [Field F]
    (𝓟 : AGMProofSystemInstantiation F) (prover : 𝓟.Prover) (stmt : 𝓟.Stmt) : Prop :=
  (
    ∀ check_idx : 𝓟.EqualityChecks, 𝓟.check_poly prover stmt check_idx = 0
  )
  ∧
  ∀ pfs ∈ 𝓟.Identified_Proof_Elems,
    𝓟.proof_element_as_poly prover pfs.fst = 𝓟.proof_element_as_poly prover pfs.snd

content: noncomputable def check_poly {F : Type} [Field F]
    (𝓟 : AGMProofSystemInstantiation F) (prover : 𝓟.Prover) (stmt : 𝓟.Stmt) (check_idx : 𝓟.EqualityChecks) :
    MvPolynomial 𝓟.Sample F :=
  (
  (𝓟.ListPairings check_idx).map fun pairing =>
    𝓟.pairing_poly prover stmt check_idx pairing
  ).sum

content: noncomputable def pairing_poly {F : Type} [Field F]
    (𝓟 : AGMProofSystemInstantiation F) (prover : 𝓟.Prover) (stmt : 𝓟.Stmt) (check_idx : 𝓟.EqualityChecks) (pairing : 𝓟.Pairings check_idx) :
    MvPolynomial 𝓟.Sample F :=
  (
    ( 
      
      (
        (𝓟.ListProof.map fun pf_elem => 
          C (𝓟.verificationPairingProof stmt check_idx pairing pf_elem) 
            *
            
            𝓟.proof_element_as_poly prover pf_elem).sum
      )
      +
      ( 
        (𝓟.ListSRSElements.map fun SRS_elem =>
          C (𝓟.verificationPairingSRS stmt check_idx pairing SRS_elem) * (𝓟.SRSElementValue SRS_elem)).sum
      )
    )
    *
    ( 
      
      (
        (𝓟.ListProof.map fun pf_elem => 
          C (𝓟.verificationPairingProof stmt check_idx pairing pf_elem) 
            *
            
            𝓟.proof_element_as_poly prover pf_elem).sum
      )
      +
      ( 
        (𝓟.ListSRSElements.map fun SRS_elem =>
          C (𝓟.verificationPairingSRS stmt check_idx pairing SRS_elem) * (𝓟.SRSElementValue SRS_elem)).sum
      )
    )
  )

content: noncomputable def proof_element_as_poly {F : Type} [Field F]
    (𝓟 : AGMProofSystemInstantiation F) (prover : 𝓟.Prover) (pf_elem : 𝓟.Proof) :
    MvPolynomial (𝓟.Sample) F :=
  (𝓟.ListSRSElements.map fun SRS_elem =>
          MvPolynomial.C (prover.fst pf_elem SRS_elem) * (𝓟.SRSElementValue SRS_elem)).sum

content: noncomputable def to_MvPolynomial_Option {F : Type} [Field F] (V : Type) :
    Polynomial F →+* MvPolynomial (Option V) F
  where
    toFun p := Polynomial.eval₂ (MvPolynomial.C) (MvPolynomial.X none) p
    map_one' := by admit /- proof elided -/

-- Repository lemmas used:
content: lemma MvPolynomial.sum_map_C {σ A R : Type} [CommSemiring R] (l : List A) (f : A → R) :
    (l.map (fun (x : A) => C (σ := σ) (f x))).sum = C ((l.map f).sum)

content: theorem AlgEquiv.list_map_sum {R : Type uR} {A₁ : Type uA₁} {A₂ : Type uA₂}
    [CommSemiring R] [Semiring A₁] [Semiring A₂] [Algebra R A₁] [Algebra R A₂]
    (e : A₁ ≃ₐ[R] A₂) {ι : Type u_1} (f : ι → A₁) (l : List ι) :
    e (l.map (fun (x : ι) => f x)).sum = (l.map fun (x : ι) => e (f x)).sum

content: lemma optionEquivRight_to_MvPolynomial_Option {F V : Type} [Field F] (p : Polynomial F) :
    (MvPolynomial.optionEquivRight F V) (to_MvPolynomial_Option V p) = C p

content: lemma optionEquivRight_comp_to_MvPolynomial_Option {F V : Type} [Field F] :
    RingHom.comp (MvPolynomial.optionEquivRight F V).toRingEquiv.toRingHom (to_MvPolynomial_Option (F := F) V) = C

content: lemma to_MvPolynomial_Option_C {F V : Type} [Field F] (r : F) :
    to_MvPolynomial_Option V (Polynomial.C r) = MvPolynomial.C r

content: theorem Polynomial.hom_congr_vars {R : Type u} {S : Type v}
    [CommSemiring R] [CommSemiring S]
    {f₁ : Polynomial R →+* S} {f₂ : Polynomial R →+* S}
    (hC : RingHom.comp f₁ Polynomial.C = RingHom.comp f₂ Polynomial.C)
    (hv : f₁ (Polynomial.X) = f₂ (Polynomial.X)) :
    f₁ = f₂

content: lemma to_MvPolynomial_Option_X {F V : Type} [Field F] :
    to_MvPolynomial_Option V (Polynomial.X) = MvPolynomial.X (R := F) none

-- Local definitions used:
content: inductive Vars : Type where
  | α : Vars
  | β : Vars
deriving Repr, BEq

content: inductive StmtEntries : Type where
  | x : StmtEntries
  | y : StmtEntries
  | z : StmtEntries
deriving Repr, BEq

content: inductive WitEntries : Type where
  | A : WitEntries
  | B : WitEntries
deriving Repr, BEq

local notation "Vars_α" => some Vars.α
local notation "Vars_β" => some Vars.β
local notation "Vars_x" => none

content: inductive Proof_G1_Idx : Type where
  | Pf : Proof_G1_Idx

content: def Proof_G2_Idx : Type := Empty

content: inductive PairingsIdx : Type where
  | lhs : PairingsIdx
  | rhs : PairingsIdx

content: inductive SRS_Elements_G2_Idx : Type where
  | α : SRS_Elements_G2_Idx
  | β : SRS_Elements_G2_Idx

content: noncomputable def ToySnark
     
    {F : Type} [Field F] :
    AGMProofSystemInstantiation F :=
  {
    Stmt := StmtEntries -> F
    Sample := Option Vars
    SRSElements_G1 := SRS_Elements_G1_Idx
    ListSRSElements_G1 :=
      [.α, .β]
    SRSElements_G2 := SRS_Elements_G2_Idx
    ListSRSElements_G2 :=
      [.α, .β]
    SRSElementValue_G1 := fun SRS_idx => match SRS_idx with
      | SRS_Elements_G1_Idx.α => MvPolynomial.X Vars_α
      | SRS_Elements_G1_Idx.β => MvPolynomial.X Vars_β
    SRSElementValue_G2 := fun SRS_idx => match SRS_idx with
      | SRS_Elements_G2_Idx.α => MvPolynomial.X Vars_α
      | SRS_Elements_G2_Idx.β => MvPolynomial.X Vars_β
    Proof_G1 := Proof_G1_Idx
    ListProof_G1 := [Proof_G1_Idx.Pf]
    Proof_G2 := Proof_G2_Idx
    ListProof_G2 := []
    EqualityChecks := Unit
    Pairings := fun _ => PairingsIdx
    ListPairings := fun _ => [PairingsIdx.lhs, PairingsIdx.rhs]
    verificationPairingSRS_G1 := fun stmt _ i SRS_idx => match i with
      | PairingsIdx.lhs => 0
      | PairingsIdx.rhs => match SRS_idx with
        | SRS_Elements_G1_Idx.α => stmt StmtEntries.z
        | SRS_Elements_G1_Idx.β => 0
    verificationPairingSRS_G2 := fun stmt _ i SRS_idx => match i with
      | PairingsIdx.lhs => match SRS_idx with
        | SRS_Elements_G2_Idx.α => stmt StmtEntries.x
        | SRS_Elements_G2_Idx.β => stmt StmtEntries.y
      | PairingsIdx.rhs => match SRS_idx with
        | SRS_Elements_G2_Idx.α => 0
        | SRS_Elements_G2_Idx.β => -1
    verificationPairingProof_G1 := fun stmt _ i pf => match i with
      | PairingsIdx.lhs => match pf with
        | Proof_G1_Idx.Pf => 1
      | PairingsIdx.rhs => match pf with
        | Proof_G1_Idx.Pf => 0
    verificationPairingProof_G2 := fun stmt _ i pf => 0
  }

-- Local lemmas used:
content: lemma Vars.finsupp_eq_ext (f g : Vars →₀ ℕ) : f = g ↔
    f Vars.α = g Vars.α
    ∧ f Vars.β = g Vars.β

-- Target theorem and proof:
lemma soundness
    {F : Type} [Field F] :
    (AGMProofSystemInstantiation.soundness
      F
      (ToySnark
        (F := F))
      (WitEntries -> F)
      (fun (stmt : StmtEntries → F) (wit : WitEntries -> F) =>
        wit WitEntries.A * stmt StmtEntries.y = stmt StmtEntries.z -- - wit WitEntries.I
        ∨
        wit WitEntries.B * stmt StmtEntries.x = stmt StmtEntries.z -- - wit WitEntries.I
      )
      (fun prover i => prover.fst Proof_G1_Idx.Pf (if i = WitEntries.A then .α else .β))

    )
:= by
  unfold AGMProofSystemInstantiation.soundness AGMProofSystemInstantiation.verify AGMProofSystemInstantiation.proof_element_G1_as_poly AGMProofSystemInstantiation.proof_element_G2_as_poly
  intros stmt prover eqns'
  rcases eqns' with ⟨eqns, null⟩
  have eqn := eqns ()
  clear eqns null

  -- Step 1: Obtain the coefficient equations of the mv_polynomials
  simp_rw [ToySnark] at eqn
  simp only [monomial_zero', List.singleton_append, List.cons_append, List.append_assoc,
    List.map_cons, Sum.elim_inl, Sum.elim_inr, List.map_append, List.map_map, List.sum_cons,
    List.sum_append, List.map_nil, List.sum_nil, add_zero, Sum.elim_lam_const_lam_const, map_one,
    one_mul, map_zero, zero_mul, map_neg, neg_mul, neg_add_rev, zero_add, mul_zero,
    -- Note: everything above is @simp tagged
    Function.comp, List.sum_map_zero] at eqn

  simp only [mul_add, add_mul, List.sum_map_add] at eqn

  -- Move all the X (some _) terms to the left, and out of sums
  simp only [
    -- Associativity to obtain a right-leaning tree
    mul_assoc,
    -- Commutativity lemmas to move X (some _) to the left
    mul_left_comm (C _) (X (some _)) _, mul_left_comm (List.sum _) (X (some _)) _,
    mul_comm (C _) (X (some _)), mul_comm (List.sum _) (X (some _)),
    -- Move negations to the bottom
    neg_mul, mul_neg,
    -- Move constant multiplications (which the X (some _) terms should be) out of sums
    List.sum_map_mul_right, List.sum_map_mul_left] at eqn


  -- I apply MvPolynomial.optionEquivRight *here*,
  -- so that we can treat polynomials in Vars_X as constants
  trace "Converting to MvPolynomial over Polynomials"
  replace eqn := congr_arg (MvPolynomial.optionEquivRight F Vars) eqn
  simp only [AlgEquiv.map_add, AlgEquiv.map_zero, AlgEquiv.map_mul, AlgEquiv.map_one,
    AlgEquiv.map_neg, AlgEquiv.list_map_sum, AlgEquiv.map_pow] at eqn
  simp only [MvPolynomial.optionEquivRight_C, MvPolynomial.optionEquivRight_X_none, MvPolynomial.optionEquivRight_X_some, optionEquivRight_to_MvPolynomial_Option] at eqn
  -- Move Cs back out so we can recognize the monomials
  -- simp only [←MvPolynomial.C_mul, ←MvPolynomial.C_pow, ←MvPolynomial.C_add,
  --   MvPolynomial.sum_map_C] at eqn

  simp only [MvPolynomial.X, C_apply, MvPolynomial.monomial_mul, one_mul, mul_one, add_zero, zero_add, mul_add, add_mul] at eqn

  trace "Applying individual coefficients"


  have h20 := congr_arg (coeff (Finsupp.single Vars.α 2 + Finsupp.single Vars.β 0)) eqn
  have h11 := congr_arg (coeff (Finsupp.single Vars.α 1 + Finsupp.single Vars.β 1)) eqn
  have h02 := congr_arg (coeff (Finsupp.single Vars.α 0 + Finsupp.single Vars.β 2)) eqn

  clear eqn

  trace "Distribute coefficient-taking over terms"
  simp only [coeff_monomial, coeff_add, coeff_neg, coeff_zero] at h20 h11 h02

  trace "Simplifying coefficient expressions"
  simp only [Vars.finsupp_eq_ext, Finsupp.single_apply, Finsupp.add_apply] at h20 h11 h02

  simp [ite_true, ite_self, add_zero, ite_false, and_self, zero_add,
    one_ne_zero, and_false, false_and, add_eq_zero, mul_eq_zero,
    add_right_eq_self, zero_ne_one, and_true, true_and, neg_zero] at h20 h11 h02 ⊢


  -- Completely remove references to Polynomial
  simp only [add_neg_eq_zero, Polynomial.C_inj, ←Polynomial.C_add, ←Polynomial.C_mul] at h20 h11 h02

  integral_domain_tactic