//! Batch Equality Proof

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

/// A `BatchEqualityProof` is a proof of the equivalence of 
#[allow(non_snake_case)]
pub struct BatchEqualityProof {
    /// `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 BatchEqualityProof {
    
    /// Proove
    pub fn proove(
        cin_asset: &Vec<RistrettoPoint>,
        cin_blind: &Vec<Scalar>,
        cout_asset: &Vec<RistrettoPoint>,
        cout_blind: &Vec<Scalar>
    ) -> Result<BatchEqualityProof, ZeiError> {
        // setup transcript
        let mut transcript = Transcript::new(b"zei_equality_proof");

        // Feed All Input Commitments
        cin_asset.into_iter().map(|i|  transcript.commit_point(b"commitment", &i.compress()));
        // Feed All Output Commitments
        cout_asset.into_iter().map(|i|  transcript.commit_point(b"commitment", &i.compress()));

        let mut all_comms = vec![];
        all_comms.extend(cin_asset.iter().cloned());
        all_comms.extend(cout_asset.iter().cloned());

        let mut all_blinds = vec![];
        all_blinds.extend(cin_blind.iter().cloned());
        all_blinds.extend(cout_blind.iter().cloned());

        // holds the sumation 
        let mut big_d = RistrettoPoint::identity();
        // holds the sumation
        let mut a_blind : Scalar = Scalar::zero();

        //TODO: HANDLE ERRORR
        assert!(all_comms.len() == all_blinds.len());

        for i in 0..all_comms.len() {
            // Commit Scalar Index
            transcript.append_u64(b"beta_index", i as u64);
            // Sample inital challenge beta_i
            let beta_i = transcript.challenge_scalar(b"c");
            //D = beta_i(c_i - c_1)
            big_d = beta_i * (all_comms[i] - all_comms[0]);
            //a = beta_i(r_i - r_1)
            a_blind = beta_i * (all_blinds[i] - all_blinds[0]);
        }
        

        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;
       
        // HASH(A, H, rH)
        let c = {
            //commit Public Point
            transcript.commit_point(b"A", &big_d.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")
        };
        
        // z = ac + r
        let z = (&a_blind * &c) + r;

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

    }

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

        // Feed All Input Commitments
        cin_asset.into_iter().map(|i|  transcript.commit_point(b"commitment", &i.compress()));
        // Feed All Output Commitments
        cout_asset.into_iter().map(|i|  transcript.commit_point(b"commitment", &i.compress()));

        let mut all_comms = vec![];
        all_comms.extend(cin_asset.iter().cloned());
        all_comms.extend(cout_asset.iter().cloned());

        // holds the sumation 
        let mut big_d = RistrettoPoint::identity();

        for i in 0..all_comms.len() {
            // Commit Scalar Index
            transcript.append_u64(b"beta_index", i as u64);
            // Sample inital challenge beta_i
            let beta_i = transcript.challenge_scalar(b"c");
            //D = beta_i(c_i - c_1)
            big_d = beta_i * (all_comms[i] - all_comms[0]);
        }

        //
        let base_h = crate::bulletproofs::hash_to_ristretto(RISTRETTO_BASEPOINT_COMPRESSED.as_bytes());
        
        //zH - cD
        let factor = (&self.z * base_h) - (self.c * big_d);

        // HASH(A, H, rH)
        let c = {
            //commit Public Point
            transcript.commit_point(b"A", &big_d.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 BatchEqualityProof {
	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 BatchEqualityProof {
	fn read(reader: &mut dyn ser::Reader) -> Result<BatchEqualityProof, ser::Error> {
		Ok(BatchEqualityProof {
            c: Scalar::read(reader)?,
            z: Scalar::read(reader)?
        })
	}
}


#[cfg(test)]
mod test {
    use super::*;
    use crate::bulletproofs::PedersenGens;
    use crate::core::AssetId;
    use mohan::hash::H256;

    #[test]
    fn test_batch_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(11119989891u64);

        let mut input_blinds = vec![];
        let mut output_blinds = vec![];

        for i in 0..10 {
            input_blinds.push(Scalar::random(&mut rng));
            output_blinds.push(Scalar::random(&mut rng));
        }

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

        let mut inputs_comm = vec![];
        let mut outputs_comm = vec![];

        for i in 0..10 {
            inputs_comm.push(
                pc_gens.commit(value.clone(), input_blinds[i].clone())
            );
            outputs_comm.push(
                pc_gens.commit(value.clone(), output_blinds[i].clone())
            );
        }

        let proof = BatchEqualityProof::proove(&inputs_comm, &input_blinds, &outputs_comm, &output_blinds).unwrap();

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

    #[test]
    fn test_batch_simple_asset_mismatch() { 
        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 flav_id = AssetId::from_inner(H256::from_vec(b"satoshi"));
        let flav_id2 = AssetId::from_inner(H256::from_vec(b"turing"));

        let mut input_blinds = vec![];
        let mut output_blinds = vec![];

        for i in 0..10 {
            input_blinds.push(Scalar::random(&mut rng));
            output_blinds.push(Scalar::random(&mut rng));
        }

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

        let mut inputs_comm = vec![];
        let mut outputs_comm = vec![];

        for i in 0..10 {
            inputs_comm.push(
                pc_gens.commit(flav_id.into_scalar(), input_blinds[i].clone())
            );
            outputs_comm.push(
                pc_gens.commit(flav_id2.into_scalar(), output_blinds[i].clone())
            );
        }

        let proof = BatchEqualityProof::proove(&inputs_comm, &input_blinds, &outputs_comm, &output_blinds).unwrap();

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


}