//! Equality Proof

use bacteria::Transcript;
use mohan::{
    ser,
    dalek::{
        ristretto::RistrettoPoint,
        scalar::Scalar,
        constants::RISTRETTO_BASEPOINT_COMPRESSED
    }
};

use crate::ZeiError;

/// A `EqualityProof` is a proof of the equivalence of the discrete logarithm between two pairs of points.
#[allow(non_snake_case)]
pub struct EqualityProof {
    /// `c` is a `Scalar`
    /// \\(c=HASH(A, H, rH)\\)
    pub(crate) c: Scalar,
    /// `z` is a `Scalar`
    /// \\(z = (ac + r) \mod p\\) 
    pub(crate) z: Scalar,
}



impl EqualityProof {
    
    /// Proove two pederson commitments commit to same value 
    pub fn proove(
        c1: &RistrettoPoint,
        c1_blind: &Scalar, 
        c2: &RistrettoPoint, 
        c2_blind: &Scalar
    ) -> Result<EqualityProof, ZeiError> {

        let mut transcript = Transcript::new(b"zei_equality_proof");

        let base_h = crate::bulletproofs::hash_to_ristretto(RISTRETTO_BASEPOINT_COMPRESSED.as_bytes());

        //randomize transcript and commit private key
        let mut rng = transcript
            .build_rng()
            .finalize(&mut rand::thread_rng());

        let r = Scalar::random(&mut rng);
        // rH
        let rh = &r * &base_h;
        //A = aH
        let big_a = c1 - c2;

        // HASH(A, H, rH)
        let c = {
            //commit Public Point
            transcript.commit_point(b"A", &big_a.compress());
            //commit to our nonce
            transcript.commit_point(b"H", &base_h.compress());
            transcript.commit_point(b"factor", &rh.compress());
            //sample challenge
            transcript.challenge_scalar(b"c")
        };
        
        let a = c1_blind - c2_blind;
        // z = ac + r
        let z = (&a*&c) + r;

        Ok(
            EqualityProof {
                c: c,
                z: z
            }
        )
    }

    /// Verify a Equality Proof of two commitments
    pub fn verify(&self, c1: &RistrettoPoint, c2: &RistrettoPoint) -> Result<(), ZeiError> {
        let mut transcript = Transcript::new(b"zei_equality_proof");

        let base_h = crate::bulletproofs::hash_to_ristretto(RISTRETTO_BASEPOINT_COMPRESSED.as_bytes());
        
        //A = aH
        let big_a = c1 - c2;

        //zH - cA
        let factor = (&self.z * base_h) - (self.c * big_a);

        // HASH(A, H, rH)
        let c = {
            //commit Public Point
            transcript.commit_point(b"A", &big_a.compress());
            //commit to our nonce
            transcript.commit_point(b"H", &base_h.compress());
            transcript.commit_point(b"factor", &factor.compress());
            //sample challenge
            transcript.challenge_scalar(b"c")
        };

        if c == self.c {
            Ok(())
        } else {
            Err(ZeiError::VerificationError)
        }
    }
}

impl ser::Writeable for EqualityProof {
	fn write<W: ser::Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
        self.c.write(writer)?;
        self.z.write(writer)?;
        Ok(())
	}
}

impl ser::Readable for EqualityProof {
	fn read(reader: &mut dyn ser::Reader) -> Result<EqualityProof, ser::Error> {
		Ok(EqualityProof{
            c: Scalar::read(reader)?,
            z: Scalar::read(reader)?
        })
	}
}


#[cfg(test)]
mod test {
    use super::*;
    use crate::bulletproofs::PedersenGens;

    #[test]
    fn test_simple() { 
        let mut transcript = Transcript::new(b"zei_equality_proof_test");

        //randomize transcript and commit private key
        let mut rng = transcript
            .build_rng()
            .finalize(&mut rand::thread_rng());

        let value = Scalar::from(11111u64);

        let input_blind = Scalar::random(&mut rng);
        let output_blind = Scalar::random(&mut rng);

        //setup generators
        let pc_gens = PedersenGens::default();

        let inputs_comm = pc_gens.commit(value.clone(), input_blind.clone());
        let outputs_comm = pc_gens.commit(value.clone(), output_blind.clone());

        let proof = EqualityProof::proove(&inputs_comm, &input_blind, &outputs_comm, &output_blind).unwrap();

        assert!(proof.verify(&inputs_comm, &outputs_comm).is_ok());
    }

    #[test]
    fn test_simple_2() { 
        let mut transcript = Transcript::new(b"zei_equality_proof_test");

        //randomize transcript and commit private key
        let mut rng = transcript
            .build_rng()
            .finalize(&mut rand::thread_rng());

        let input_blind = Scalar::random(&mut rng);
        let output_blind = Scalar::random(&mut rng);
        
        let input_blind2 = Scalar::random(&mut rng);
        let output_blind2 = Scalar::random(&mut rng);

        //setup generators
        let pc_gens = PedersenGens::default();

        let inputs_comm = pc_gens.commit(Scalar::from(11u64), input_blind.clone());
        let inputs_comm2 = pc_gens.commit(Scalar::from(12u64), input_blind2.clone());
        let outputs_comm = pc_gens.commit(Scalar::from(11u64), output_blind.clone());
        let outputs_comm2 = pc_gens.commit(Scalar::from(1u64), output_blind2.clone());

        let total_in = inputs_comm + inputs_comm2;
        let total_out = outputs_comm + outputs_comm2;
        let total_in_r = input_blind + input_blind2;
        let total_out_r = output_blind + output_blind2;

        let proof = EqualityProof::proove(&total_in, &total_in_r, &total_out, &total_out_r).unwrap();

        assert!(proof.verify(&total_in, &total_out).is_err());
    }

    #[test]
    fn test_simple_3() { 
        let mut transcript = Transcript::new(b"zei_equality_proof_test");

        //randomize transcript and commit private key
        let mut rng = transcript
            .build_rng()
            .finalize(&mut rand::thread_rng());

        let input_blind = Scalar::random(&mut rng);
        let output_blind = Scalar::random(&mut rng);
        
        let input_blind2 = Scalar::random(&mut rng);
        let output_blind2 = Scalar::random(&mut rng);

        //setup generators
        let pc_gens = PedersenGens::default();

        let inputs_comm = pc_gens.commit(Scalar::from(11u64), input_blind.clone());
        let inputs_comm2 = pc_gens.commit(Scalar::from(11u64), input_blind2.clone());
        let outputs_comm = pc_gens.commit(Scalar::from(11u64), output_blind.clone());
        let outputs_comm2 = pc_gens.commit(Scalar::from(11u64), output_blind2.clone());

        let total_in = inputs_comm + inputs_comm2;
        let total_out = outputs_comm + outputs_comm2;
        let total_in_r = input_blind + input_blind2;
        let total_out_r = output_blind + output_blind2;

        let proof = EqualityProof::proove(&total_in, &total_in_r, &total_out, &total_out_r).unwrap();

        assert!(proof.verify(&total_in, &total_out).is_ok());
    }

    #[test]
    fn test_simple_4() { 
        let mut transcript = Transcript::new(b"zei_equality_proof_test");

        //randomize transcript and commit private key
        let mut rng = transcript
            .build_rng()
            .finalize(&mut rand::thread_rng());

        let input_blind = Scalar::random(&mut rng);
        let output_blind = Scalar::random(&mut rng);
        
        let input_blind2 = Scalar::random(&mut rng);
        

        //setup generators
        let pc_gens = PedersenGens::default();

        let inputs_comm = pc_gens.commit(Scalar::from(11u64), input_blind.clone());
        let inputs_comm2 = pc_gens.commit(Scalar::from(11u64), input_blind2.clone());
        let outputs_comm = pc_gens.commit(Scalar::from(11u64), output_blind.clone());
      

        let total_in = inputs_comm + inputs_comm2;
        let total_out = outputs_comm;
        let total_in_r = input_blind + input_blind2;
        let total_out_r = output_blind;

        let proof = EqualityProof::proove(&total_in, &total_in_r, &total_out, &total_out_r).unwrap();

        assert!(proof.verify(&total_in, &total_out).is_err());
    }

    #[test]
    fn test_simple_5() { 
        let mut transcript = Transcript::new(b"zei_equality_proof_test");

        //randomize transcript and commit private key
        let mut rng = transcript
            .build_rng()
            .finalize(&mut rand::thread_rng());

        let input_blind = Scalar::random(&mut rng);
        let output_blind = Scalar::random(&mut rng);
        
        let input_blind2 = Scalar::random(&mut rng);
        

        //setup generators
        let pc_gens = PedersenGens::default();

        let inputs_comm = pc_gens.commit(Scalar::from(11u64), input_blind.clone());
        let inputs_comm2 = pc_gens.commit(Scalar::from(11u64), input_blind2.clone());
        let outputs_comm = pc_gens.commit(Scalar::from(1000u64), output_blind.clone());
      

        let total_in = inputs_comm + inputs_comm2;
        let total_out = outputs_comm;
        let total_in_r = input_blind + input_blind2;
        let total_out_r = output_blind;

        let proof = EqualityProof::proove(&total_in, &total_in_r, &total_out, &total_out_r).unwrap();

        assert!(proof.verify(&total_in, &total_out).is_err());
    }

    #[test]
    fn test_simple_not_equal() { 
        let mut transcript = Transcript::new(b"zei_equality_proof_test");

        //randomize transcript and commit private key
        let mut rng = transcript
            .build_rng()
            .finalize(&mut rand::thread_rng());

        let value = Scalar::from(11111u64);
        let value2 = Scalar::from(111111u64);

        let input_blind = Scalar::random(&mut rng);
        let output_blind = Scalar::random(&mut rng);

        //setup generators
        let pc_gens = PedersenGens::default();

        let inputs_comm = pc_gens.commit(value.clone(), input_blind.clone());
        let outputs_comm = pc_gens.commit(value2.clone(), output_blind.clone());

        let proof = EqualityProof::proove(&inputs_comm, &input_blind, &outputs_comm, &output_blind).unwrap();

        assert!(proof.verify(&inputs_comm, &outputs_comm).is_err());
    }

    


}