// Copyright (c) 2022 Espresso Systems (espressosys.com)
// This file is part of the Jellyfish library.

// You should have received a copy of the MIT License
// along with the Jellyfish library. If not, see <https://mit-license.org/>.

use ark_bls12_381::{Bls12_381, Fr as Fr381};
use ark_ec::{
    PairingEngine,
};
use ark_ff::PrimeField;
use ark_std::{UniformRand};
use jf_plonk::{
    circuit::{Arithmetization, Circuit, PlonkCircuit},
    errors::PlonkError,
    proof_system::{PlonkKzgSnark, Snark},
    transcript::StandardTranscript,
};
use rand_chacha::ChaCha20Rng;
use ark_std::rand::SeedableRng;

#[allow(non_snake_case)]
fn main() -> Result<(), PlonkError> {
    let mut rng = ChaCha20Rng::from_seed([0u8; 32]);
    let x = Fr381::rand(&mut rng);
    let circuit = range_circuit::<Fr381, Bls12_381>(x)?;
    let srs_size = circuit.srs_size()?;
    let srs = PlonkKzgSnark::<Bls12_381>::universal_setup(srs_size, &mut rng)?;
    
    let (pk, vk) = PlonkKzgSnark::<Bls12_381>::preprocess(&srs, &circuit)?;
    
    let proof = PlonkKzgSnark::<Bls12_381>::prove::<_, _, StandardTranscript>(
        &mut rng, &circuit, &pk, None,
    )?;

    let public_inputs = circuit.public_input().unwrap();
    let extra_transcript_init_msg = None;
    assert!(PlonkKzgSnark::<Bls12_381>::verify::<StandardTranscript>(
        &vk,
        &public_inputs,
        &proof,
        extra_transcript_init_msg,
    )
    .is_ok());

    Ok(())
}


#[allow(non_snake_case)]
fn range_circuit<F:PrimeField, PairingCurve:PairingEngine>(
    x: F,
) -> Result<PlonkCircuit<F>, PlonkError>
{
    let mut circuit: PlonkCircuit<F> = PlonkCircuit::new_turbo_plonk();
    let x_var = circuit.create_variable(x)?;
    circuit.range_gate(x_var, 10)?;
    circuit.finalize_for_arithmetization()?;
    Ok(circuit)
}