// code forked from:
// https://github.com/EspressoSystems/hyperplonk/tree/main/subroutines/src/poly_iop/sum_check
//
// Copyright (c) 2023 Espresso Systems (espressosys.com)
// This file is part of the HyperPlonk library.

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

//! This module implements the sum check protocol.

use crate::{
    transcript::Transcript,
    utils::virtual_polynomial::{VPAuxInfo, VirtualPolynomial},
};
use ark_ec::CurveGroup;
use ark_ff::PrimeField;
use ark_poly::univariate::DensePolynomial;
use ark_poly::{DenseMultilinearExtension, DenseUVPolynomial, Polynomial};
use ark_std::{end_timer, start_timer};
use std::{fmt::Debug, marker::PhantomData, sync::Arc};

use crate::utils::sum_check::structs::IOPProverMessage;
use crate::utils::sum_check::structs::IOPVerifierState;
use ark_ff::Field;
use espresso_subroutines::poly_iop::prelude::PolyIOPErrors;
use structs::{IOPProof, IOPProverState};

mod prover;
pub mod structs;
pub mod verifier;

/// A generic sum-check trait over a curve group
pub trait SumCheck<C: CurveGroup> {
    type VirtualPolynomial;
    type VPAuxInfo;
    type MultilinearExtension;

    type SumCheckProof: Clone + Debug + Default + PartialEq;
    type SumCheckSubClaim: Clone + Debug + Default + PartialEq;

    /// Extract sum from the proof
    fn extract_sum(proof: &Self::SumCheckProof) -> C::ScalarField;

    /// Generate proof of the sum of polynomial over {0,1}^`num_vars`
    ///
    /// The polynomial is represented in the form of a VirtualPolynomial.
    fn prove(
        poly: &Self::VirtualPolynomial,
        transcript: &mut impl Transcript<C>,
    ) -> Result<Self::SumCheckProof, PolyIOPErrors>;

    /// Verify the claimed sum using the proof
    fn verify(
        sum: C::ScalarField,
        proof: &Self::SumCheckProof,
        aux_info: &Self::VPAuxInfo,
        transcript: &mut impl Transcript<C>,
    ) -> Result<Self::SumCheckSubClaim, PolyIOPErrors>;
}

/// Trait for sum check protocol prover side APIs.
pub trait SumCheckProver<C: CurveGroup>
where
    Self: Sized,
{
    type VirtualPolynomial;
    type ProverMessage;

    /// Initialize the prover state to argue for the sum of the input polynomial
    /// over {0,1}^`num_vars`.
    fn prover_init(polynomial: &Self::VirtualPolynomial) -> Result<Self, PolyIOPErrors>;

    /// Receive message from verifier, generate prover message, and proceed to
    /// next round.
    ///
    /// Main algorithm used is from section 3.2 of [XZZPS19](https://eprint.iacr.org/2019/317.pdf#subsection.3.2).
    fn prove_round_and_update_state(
        &mut self,
        challenge: &Option<C::ScalarField>,
    ) -> Result<Self::ProverMessage, PolyIOPErrors>;
}

/// Trait for sum check protocol verifier side APIs.
pub trait SumCheckVerifier<C: CurveGroup> {
    type VPAuxInfo;
    type ProverMessage;
    type Challenge;
    type SumCheckSubClaim;

    /// Initialize the verifier's state.
    fn verifier_init(index_info: &Self::VPAuxInfo) -> Self;

    /// Run verifier for the current round, given a prover message.
    ///
    /// Note that `verify_round_and_update_state` only samples and stores
    /// challenges; and update the verifier's state accordingly. The actual
    /// verifications are deferred (in batch) to `check_and_generate_subclaim`
    /// at the last step.
    fn verify_round_and_update_state(
        &mut self,
        prover_msg: &Self::ProverMessage,
        transcript: &mut impl Transcript<C>,
    ) -> Result<Self::Challenge, PolyIOPErrors>;

    /// This function verifies the deferred checks in the interactive version of
    /// the protocol; and generate the subclaim. Returns an error if the
    /// proof failed to verify.
    ///
    /// If the asserted sum is correct, then the multilinear polynomial
    /// evaluated at `subclaim.point` will be `subclaim.expected_evaluation`.
    /// Otherwise, it is highly unlikely that those two will be equal.
    /// Larger field size guarantees smaller soundness error.
    fn check_and_generate_subclaim(
        &self,
        asserted_sum: &C::ScalarField,
    ) -> Result<Self::SumCheckSubClaim, PolyIOPErrors>;
}

/// A SumCheckSubClaim is a claim generated by the verifier at the end of
/// verification when it is convinced.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct SumCheckSubClaim<F: PrimeField> {
    /// the multi-dimensional point that this multilinear extension is evaluated
    /// to
    pub point: Vec<F>,
    /// the expected evaluation
    pub expected_evaluation: F,
}

#[derive(Clone, Debug, Default, Copy, PartialEq, Eq)]
pub struct IOPSumCheck<C: CurveGroup, T: Transcript<C>> {
    #[doc(hidden)]
    phantom: PhantomData<C>,
    #[doc(hidden)]
    phantom2: PhantomData<T>,
}

impl<C: CurveGroup, T: Transcript<C>> SumCheck<C> for IOPSumCheck<C, T> {
    type SumCheckProof = IOPProof<C::ScalarField>;
    type VirtualPolynomial = VirtualPolynomial<C::ScalarField>;
    type VPAuxInfo = VPAuxInfo<C::ScalarField>;
    type MultilinearExtension = Arc<DenseMultilinearExtension<C::ScalarField>>;
    type SumCheckSubClaim = SumCheckSubClaim<C::ScalarField>;

    fn extract_sum(proof: &Self::SumCheckProof) -> C::ScalarField {
        let start = start_timer!(|| "extract sum");
        let poly = DensePolynomial::from_coefficients_vec(proof.proofs[0].coeffs.clone());
        let res = poly.evaluate(&C::ScalarField::ONE) + poly.evaluate(&C::ScalarField::ZERO);
        end_timer!(start);
        res
    }

    fn prove(
        poly: &VirtualPolynomial<C::ScalarField>,
        transcript: &mut impl Transcript<C>,
    ) -> Result<IOPProof<C::ScalarField>, PolyIOPErrors> {
        transcript.absorb(&C::ScalarField::from(poly.aux_info.num_variables as u64));
        transcript.absorb(&C::ScalarField::from(poly.aux_info.max_degree as u64));
        let mut prover_state: IOPProverState<C> = IOPProverState::prover_init(poly)?;
        let mut challenge: Option<C::ScalarField> = None;
        let mut prover_msgs: Vec<IOPProverMessage<C::ScalarField>> =
            Vec::with_capacity(poly.aux_info.num_variables);
        for _ in 0..poly.aux_info.num_variables {
            let prover_msg: IOPProverMessage<C::ScalarField> =
                IOPProverState::prove_round_and_update_state(&mut prover_state, &challenge)?;
            transcript.absorb_vec(&prover_msg.coeffs);
            prover_msgs.push(prover_msg);
            challenge = Some(transcript.get_challenge());
        }
        if let Some(p) = challenge {
            prover_state.challenges.push(p)
        };
        Ok(IOPProof {
            point: prover_state.challenges,
            proofs: prover_msgs,
        })
    }

    fn verify(
        claimed_sum: C::ScalarField,
        proof: &IOPProof<C::ScalarField>,
        aux_info: &VPAuxInfo<C::ScalarField>,
        transcript: &mut impl Transcript<C>,
    ) -> Result<SumCheckSubClaim<C::ScalarField>, PolyIOPErrors> {
        transcript.absorb(&C::ScalarField::from(aux_info.num_variables as u64));
        transcript.absorb(&C::ScalarField::from(aux_info.max_degree as u64));
        let mut verifier_state = IOPVerifierState::verifier_init(aux_info);
        for i in 0..aux_info.num_variables {
            let prover_msg = proof.proofs.get(i).expect("proof is incomplete");
            transcript.absorb_vec(&prover_msg.coeffs);
            IOPVerifierState::verify_round_and_update_state(
                &mut verifier_state,
                prover_msg,
                transcript,
            )?;
        }

        IOPVerifierState::check_and_generate_subclaim(&verifier_state, &claimed_sum)
    }
}

#[cfg(test)]
pub mod tests {
    use std::sync::Arc;

    use ark_ff::Field;
    use ark_pallas::Fr;
    use ark_pallas::Projective;
    use ark_poly::DenseMultilinearExtension;
    use ark_poly::MultilinearExtension;
    use ark_std::test_rng;

    use crate::transcript::poseidon::poseidon_test_config;
    use crate::transcript::poseidon::PoseidonTranscript;
    use crate::transcript::Transcript;
    use crate::utils::sum_check::SumCheck;
    use crate::utils::virtual_polynomial::VirtualPolynomial;

    use super::IOPSumCheck;

    #[test]
    pub fn sumcheck_poseidon() {
        let mut rng = test_rng();
        let poly_mle = DenseMultilinearExtension::rand(5, &mut rng);
        let virtual_poly = VirtualPolynomial::new_from_mle(&Arc::new(poly_mle), Fr::ONE);
        let poseidon_config = poseidon_test_config::<Fr>();

        // sum-check prove
        let mut poseidon_transcript_prove: PoseidonTranscript<Projective> =
            PoseidonTranscript::<Projective>::new(&poseidon_config);
        let sum_check = IOPSumCheck::<Projective, PoseidonTranscript<Projective>>::prove(
            &virtual_poly,
            &mut poseidon_transcript_prove,
        )
        .unwrap();

        // sum-check verify
        let claimed_sum =
            IOPSumCheck::<Projective, PoseidonTranscript<Projective>>::extract_sum(&sum_check);
        let mut poseidon_transcript_verify: PoseidonTranscript<Projective> =
            PoseidonTranscript::<Projective>::new(&poseidon_config);
        let res_verify = IOPSumCheck::<Projective, PoseidonTranscript<Projective>>::verify(
            claimed_sum,
            &sum_check,
            &virtual_poly.aux_info,
            &mut poseidon_transcript_verify,
        );

        assert!(res_verify.is_ok());
    }
}