// In this example, we will commit to a single polynomial, open it first at one point, and then batched at two points, and finally verify the proofs.
// We will use the KZG10 polynomial commitment scheme, following the approach from Marlin.

use ark_poly_commit::{Polynomial, marlin_pc::MarlinKZG10, LabeledPolynomial, PolynomialCommitment, QuerySet, Evaluations};
use ark_bls12_377::Bls12_377;
use ark_crypto_primitives::sponge::poseidon::{PoseidonSponge, PoseidonConfig};
use ark_crypto_primitives::sponge::CryptographicSponge;
use ark_ec::pairing::Pairing;
use ark_ff::UniformRand;
use ark_std::test_rng;
use ark_poly::{univariate::DensePolynomial, DenseUVPolynomial};
use rand_chacha::ChaCha20Rng;
use ark_ff::PrimeField;

type UniPoly_377 = DensePolynomial<<Bls12_377 as Pairing>::ScalarField>;
type Sponge_Bls12_377 = PoseidonSponge<<Bls12_377 as Pairing>::ScalarField>;
type PCS = MarlinKZG10<Bls12_377, UniPoly_377, Sponge_Bls12_377>;

let rng = &mut test_rng();

let max_degree = 16; // max degree supported by the scheme with the given public parameters generated by the setup here.

// 1. PolynomialCommitment::setup
// The setup procedure in this example is for demonstration purposes only - typically a setup ceremony would be run to generate the public parameters.
let pp = PCS::setup(max_degree, None, rng).unwrap();

let degree = 10; //degree of our polynomial
let secret_poly = UniPoly_377::rand(degree, rng);

let point_1 = <Bls12_377 as Pairing>::ScalarField::rand(rng);
let point_2 = <Bls12_377 as Pairing>::ScalarField::rand(rng);

let label = String::from("secret_poly");
let labeled_poly = LabeledPolynomial::new(
   label.clone(),
   secret_poly.clone(),
   Some(degree),
   Some(2), // we will open a univariate poly at two points
);

// TODO: replace by https://github.com/arkworks-rs/crypto-primitives/issues/112.
fn test_sponge<F: PrimeField>() -> PoseidonSponge<F> {
   let full_rounds = 8;
   let partial_rounds = 31;
   let alpha = 17;

   let mds = vec![
      vec![F::one(), F::zero(), F::one()],
      vec![F::one(), F::one(), F::zero()],
      vec![F::zero(), F::one(), F::one()],
   ];

   let mut v = Vec::new();
   let mut ark_rng = test_rng();

   for _ in 0..(full_rounds + partial_rounds) {
      let mut res = Vec::new();

      for _ in 0..3 {
         res.push(F::rand(&mut ark_rng));
      }
      v.push(res);
   }
   let config = PoseidonConfig::new(full_rounds, partial_rounds, alpha, mds, v, 2, 1);
   PoseidonSponge::new(&config)
}
let mut test_sponge = test_sponge::<<Bls12_377 as Pairing>::ScalarField>();

// 2. PolynomialCommitment::trim
// Since the setup produced pp with a max degree of 16, and our poly is of degree 10, we can trim the SRS to tailor it to this example.
let (ck, vk) = PCS::trim(&pp, degree, 2, Some(&[degree])).unwrap(); 

// 3. PolynomialCommitment::commit
// The prover commits to the polynomial using their committer key `ck`.
let (comms, states) = PCS::commit(&ck, [&labeled_poly], Some(rng)).unwrap(); 

// 4a. PolynomialCommitment::open
// Opening proof at a single point.
let proof_single = PCS::open(&ck, [&labeled_poly], &comms, &point_1, &mut (test_sponge.clone()), &states, None).unwrap(); 

// 5a. PolynomialCommitment::check
// Verifying the proof at a single point, given the commitment, the point, the claimed evaluation, and the proof.
assert!(PCS::check(&vk, &comms, &point_1, [secret_poly.evaluate(&point_1)], &proof_single, &mut (test_sponge.clone()), Some(rng)).unwrap()); 

let mut query_set = QuerySet::new();
let mut values = Evaluations::new();
for (i, point) in [point_1, point_2].iter().enumerate() {
   query_set.insert((label.clone(), (format!("{}", i), point.clone())));
   let value = secret_poly.evaluate(&point);
   values.insert((label.clone(), point.clone()), value);
}

// 4b. PolynomialCommitment::batch_open
// Some schemes support batch opening proofs. Generate a single proof for opening the polynomial at multiple points.
let proof_batched = PCS::batch_open(
   &ck,
   [&labeled_poly],
   &comms,
   &query_set,
   &mut (test_sponge.clone()),
   &states,
   Some(rng),
).unwrap();

// 5b. PolynomialCommitment::batch_check
assert!(PCS::batch_check(
   &vk,
   &comms,
   &query_set,
   &values,
   &proof_batched,
   &mut (test_sponge.clone()),
   rng,
).unwrap());