#![allow(non_snake_case)]
use curve25519_dalek::{ristretto::CompressedRistretto, scalar::Scalar};
use merlin::Transcript;
use rand::{seq::SliceRandom, thread_rng};
use std::time::{Duration, Instant};
use webb_bulletproofs::{r1cs::*, BulletproofGens, PedersenGens};

// Shuffle gadget (documented in markdown file)

/// A proof-of-shuffle.
struct ShuffleProof(R1CSProof);

impl ShuffleProof {
	fn gadget<CS: RandomizableConstraintSystem>(
		cs: &mut CS,
		x: Vec<Variable>,
		y: Vec<Variable>,
	) -> Result<(), R1CSError> {
		assert_eq!(x.len(), y.len());
		let k = x.len();

		if k == 1 {
			cs.constrain(y[0] - x[0]);
			return Ok(());
		}

		cs.specify_randomized_constraints(move |cs| {
			let z = cs.challenge_scalar(b"shuffle challenge");

			// Make last x multiplier for i = k-1 and k-2
			let (_, _, last_mulx_out) = cs.multiply(x[k - 1] - z, x[k - 2] - z);

			// Make multipliers for x from i == [0, k-3]
			let first_mulx_out = (0..k - 2).rev().fold(last_mulx_out, |prev_out, i| {
				let (_, _, o) = cs.multiply(prev_out.into(), x[i] - z);
				o
			});

			// Make last y multiplier for i = k-1 and k-2
			let (_, _, last_muly_out) = cs.multiply(y[k - 1] - z, y[k - 2] - z);

			// Make multipliers for y from i == [0, k-3]
			let first_muly_out = (0..k - 2).rev().fold(last_muly_out, |prev_out, i| {
				let (_, _, o) = cs.multiply(prev_out.into(), y[i] - z);
				o
			});

			// Constrain last x mul output and last y mul output to be equal
			cs.constrain(first_mulx_out - first_muly_out);

			Ok(())
		})
	}
}

impl ShuffleProof {
	/// Attempt to construct a proof that `output` is a permutation of `input`.
	///
	/// Returns a tuple `(proof, input_commitments || output_commitments)`.
	pub fn prove<'a, 'b>(
		pc_gens: &'b PedersenGens,
		bp_gens: &'b BulletproofGens,
		transcript: &'a mut Transcript,
		input: &[Scalar],
		output: &[Scalar],
	) -> Result<(ShuffleProof, Vec<CompressedRistretto>, Vec<CompressedRistretto>), R1CSError> {
		// Apply a domain separator with the shuffle parameters to the transcript
		// XXX should this be part of the gadget?
		let k = input.len();
		transcript.append_message(b"dom-sep", b"ShuffleProof");
		transcript.append_u64(b"k", k as u64);

		let mut prover = Prover::new(&pc_gens, transcript);

		// Construct blinding factors using an RNG.
		// Note: a non-example implementation would want to operate on existing
		// commitments.
		let mut blinding_rng = rand::thread_rng();

		let (input_commitments, input_vars): (Vec<_>, Vec<_>) = input
			.into_iter()
			.map(|v| prover.commit(*v, Scalar::random(&mut blinding_rng)))
			.unzip();

		let (output_commitments, output_vars): (Vec<_>, Vec<_>) = output
			.into_iter()
			.map(|v| prover.commit(*v, Scalar::random(&mut blinding_rng)))
			.unzip();

		ShuffleProof::gadget(&mut prover, input_vars, output_vars)?;
		println!(
			"For constraints {} and multipliers {} ",
			&prover.num_constraints(),
			&prover.num_multipliers()
		);
		let proof = prover.prove(&bp_gens)?;
		Ok((ShuffleProof(proof), input_commitments, output_commitments))
	}
}

impl ShuffleProof {
	/// Attempt to verify a `ShuffleProof`.
	pub fn verify<'a, 'b>(
		&self,
		pc_gens: &'b PedersenGens,
		bp_gens: &'b BulletproofGens,
		transcript: &'a mut Transcript,
		input_commitments: &Vec<CompressedRistretto>,
		output_commitments: &Vec<CompressedRistretto>,
	) -> Result<(), R1CSError> {
		// Apply a domain separator with the shuffle parameters to the transcript
		// XXX should this be part of the gadget?
		let k = input_commitments.len();
		transcript.append_message(b"dom-sep", b"ShuffleProof");
		transcript.append_u64(b"k", k as u64);

		let mut verifier = Verifier::new(transcript);

		let input_vars: Vec<_> = input_commitments.iter().map(|V| verifier.commit(*V)).collect();

		let output_vars: Vec<_> = output_commitments.iter().map(|V| verifier.commit(*V)).collect();

		ShuffleProof::gadget(&mut verifier, input_vars, output_vars)?;

		verifier.verify(&self.0, &pc_gens, &bp_gens)
	}
}

fn kshuffle_helper(k: usize) {
	use rand::Rng;

	// Common code
	let pc_gens = PedersenGens::default();
	let bp_gens = BulletproofGens::new((2 * k).next_power_of_two(), 1);

	let (proof, input_commitments, output_commitments) = {
		// Randomly generate inputs and outputs to kshuffle
		let mut rng = rand::thread_rng();
		let (min, max) = (0u64, std::u64::MAX);
		let input: Vec<Scalar> = (0..k).map(|_| Scalar::from(rng.gen_range(min, max))).collect();
		let mut output = input.clone();
		output.shuffle(&mut rand::thread_rng());

		let mut prover_transcript = Transcript::new(b"ShuffleProofTest");
		let start = Instant::now();
		let r = ShuffleProof::prove(&pc_gens, &bp_gens, &mut prover_transcript, &input, &output).unwrap();
		let end = start.elapsed();
		println!("Proving time for {} is {:?}", &input.len(), end);
		r
	};

	{
		let mut verifier_transcript = Transcript::new(b"ShuffleProofTest");
		let start = Instant::now();
		assert!(proof
			.verify(
				&pc_gens,
				&bp_gens,
				&mut verifier_transcript,
				&input_commitments,
				&output_commitments
			)
			.is_ok());
		let end = start.elapsed();
		println!("Verification time for {} is {:?}", input_commitments.len(), end);
	}
}

#[test]
fn shuffle_gadget_test_1() {
	kshuffle_helper(1);
}

#[test]
fn shuffle_gadget_test_2() {
	kshuffle_helper(2);
}

#[test]
fn shuffle_gadget_test_3() {
	kshuffle_helper(3);
}

#[test]
fn shuffle_gadget_test_4() {
	kshuffle_helper(4);
}

#[test]
fn shuffle_gadget_test_5() {
	kshuffle_helper(5);
}

#[test]
fn shuffle_gadget_test_6() {
	kshuffle_helper(6);
}

#[test]
fn shuffle_gadget_test_7() {
	kshuffle_helper(7);
}

#[test]
fn shuffle_gadget_test_24() {
	kshuffle_helper(24);
}

#[test]
fn shuffle_gadget_test_42() {
	kshuffle_helper(100);
}

/// Constrains (a1 + a2) * (b1 + b2) = (c1 + c2)
fn example_gadget<CS: ConstraintSystem>(
	cs: &mut CS,
	a1: LinearCombination,
	a2: LinearCombination,
	b1: LinearCombination,
	b2: LinearCombination,
	c1: LinearCombination,
	c2: LinearCombination,
) {
	let (_, _, c_var) = cs.multiply(a1 + a2, b1 + b2);
	cs.constrain(c1 + c2 - c_var);
}

// Prover's scope
fn example_gadget_proof(
	pc_gens: &PedersenGens,
	bp_gens: &BulletproofGens,
	a1: u64,
	a2: u64,
	b1: u64,
	b2: u64,
	c1: u64,
	c2: u64,
) -> Result<(R1CSProof, Vec<CompressedRistretto>), R1CSError> {
	let mut transcript = Transcript::new(b"R1CSExampleGadget");

	// 1. Create a prover
	let mut prover = Prover::new(pc_gens, &mut transcript);

	// 2. Commit high-level variables
	let (commitments, vars): (Vec<_>, Vec<_>) = [a1, a2, b1, b2, c1]
		.into_iter()
		.map(|x| prover.commit(Scalar::from(*x), Scalar::random(&mut thread_rng())))
		.unzip();

	// 3. Build a CS
	example_gadget(
		&mut prover,
		vars[0].into(),
		vars[1].into(),
		vars[2].into(),
		vars[3].into(),
		vars[4].into(),
		Scalar::from(c2).into(),
	);

	// 4. Make a proof
	let proof = prover.prove(bp_gens)?;

	Ok((proof, commitments))
}

// Verifier logic
fn example_gadget_verify(
	pc_gens: &PedersenGens,
	bp_gens: &BulletproofGens,
	c2: u64,
	proof: R1CSProof,
	commitments: Vec<CompressedRistretto>,
) -> Result<(), R1CSError> {
	let mut transcript = Transcript::new(b"R1CSExampleGadget");

	// 1. Create a verifier
	let mut verifier = Verifier::new(&mut transcript);

	// 2. Commit high-level variables
	let vars: Vec<_> = commitments.iter().map(|V| verifier.commit(*V)).collect();

	// 3. Build a CS
	example_gadget(
		&mut verifier,
		vars[0].into(),
		vars[1].into(),
		vars[2].into(),
		vars[3].into(),
		vars[4].into(),
		Scalar::from(c2).into(),
	);

	// 4. Verify the proof
	verifier
		.verify(&proof, &pc_gens, &bp_gens)
		.map_err(|_| R1CSError::VerificationError)
}

fn example_gadget_roundtrip_helper(a1: u64, a2: u64, b1: u64, b2: u64, c1: u64, c2: u64) -> Result<(), R1CSError> {
	// Common
	let pc_gens = PedersenGens::default();
	let bp_gens = BulletproofGens::new(128, 1);

	let (proof, commitments) = example_gadget_proof(&pc_gens, &bp_gens, a1, a2, b1, b2, c1, c2)?;

	example_gadget_verify(&pc_gens, &bp_gens, c2, proof, commitments)
}

fn example_gadget_roundtrip_serialization_helper(
	a1: u64,
	a2: u64,
	b1: u64,
	b2: u64,
	c1: u64,
	c2: u64,
) -> Result<(), R1CSError> {
	// Common
	let pc_gens = PedersenGens::default();
	let bp_gens = BulletproofGens::new(128, 1);

	let (proof, commitments) = example_gadget_proof(&pc_gens, &bp_gens, a1, a2, b1, b2, c1, c2)?;

	let proof = proof.to_bytes();

	let proof = R1CSProof::from_bytes(&proof)?;

	example_gadget_verify(&pc_gens, &bp_gens, c2, proof, commitments)
}

#[test]
fn example_gadget_test() {
	// (3 + 4) * (6 + 1) = (40 + 9)
	assert!(example_gadget_roundtrip_helper(3, 4, 6, 1, 40, 9).is_ok());
	// (3 + 4) * (6 + 1) != (40 + 10)
	assert!(example_gadget_roundtrip_helper(3, 4, 6, 1, 40, 10).is_err());
}

#[test]
fn example_gadget_serialization_test() {
	// (3 + 4) * (6 + 1) = (40 + 9)
	assert!(example_gadget_roundtrip_serialization_helper(3, 4, 6, 1, 40, 9).is_ok());
	// (3 + 4) * (6 + 1) != (40 + 10)
	assert!(example_gadget_roundtrip_serialization_helper(3, 4, 6, 1, 40, 10).is_err());
}

// Range Proof gadget

/// Enforces that the quantity of v is in the range [0, 2^n).
pub fn range_proof<CS: ConstraintSystem>(
	cs: &mut CS,
	mut v: LinearCombination,
	v_assignment: Option<u64>,
	n: usize,
) -> Result<(), R1CSError> {
	let mut exp_2 = Scalar::one();
	for i in 0..n {
		// Create low-level variables and add them to constraints
		let (a, b, o) = cs.allocate_multiplier(v_assignment.map(|q| {
			let bit: u64 = (q >> i) & 1;
			((1 - bit).into(), bit.into())
		}))?;

		// Enforce a * b = 0, so one of (a,b) is zero
		cs.constrain(o.into());

		// Enforce that a = 1 - b, so they both are 1 or 0.
		cs.constrain(a + (b - 1u64));

		// Add `-b_i*2^i` to the linear combination
		// in order to form the following constraint by the end of the loop:
		// v = Sum(b_i * 2^i, i = 0..n-1)
		v = v - b * exp_2;

		exp_2 = exp_2 + exp_2;
	}

	// Enforce that v = Sum(b_i * 2^i, i = 0..n-1)
	cs.constrain(v);

	Ok(())
}

#[test]
fn range_proof_gadget() {
	use rand::{thread_rng, Rng};

	let mut rng = thread_rng();
	let m = 3; // number of values to test per `n`

	for n in [2, 10, 32, 63].iter() {
		let (min, max) = (0u64, ((1u128 << n) - 1) as u64);
		let values: Vec<u64> = (0..m).map(|_| rng.gen_range(min, max)).collect();
		for v in values {
			assert!(range_proof_helper(v.into(), *n).is_ok());
		}
		assert!(range_proof_helper((max + 1).into(), *n).is_err());
	}
}

fn range_proof_helper(v_val: u64, n: usize) -> Result<(), R1CSError> {
	// Common
	let pc_gens = PedersenGens::default();
	let bp_gens = BulletproofGens::new(128, 1);

	// Prover's scope
	let (proof, commitment) = {
		// Prover makes a `ConstraintSystem` instance representing a range proof gadget
		let mut prover_transcript = Transcript::new(b"RangeProofTest");
		let mut rng = rand::thread_rng();

		let mut prover = Prover::new(&pc_gens, &mut prover_transcript);

		let (com, var) = prover.commit(v_val.into(), Scalar::random(&mut rng));
		assert!(range_proof(&mut prover, var.into(), Some(v_val), n).is_ok());

		let proof = prover.prove(&bp_gens)?;

		(proof, com)
	};

	// Verifier makes a `ConstraintSystem` instance representing a merge gadget
	let mut verifier_transcript = Transcript::new(b"RangeProofTest");
	let mut verifier = Verifier::new(&mut verifier_transcript);

	let var = verifier.commit(commitment);

	// Verifier adds constraints to the constraint system
	assert!(range_proof(&mut verifier, var.into(), None, n).is_ok());

	// Verifier verifies proof
	Ok(verifier.verify(&proof, &pc_gens, &bp_gens)?)
}
