/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.products.ov.mixnet.proofs.bg.zero;

import java.math.BigInteger;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
import com.scytl.products.ov.mixnet.commons.beans.proofs.ZeroProofAnswer;
import com.scytl.products.ov.mixnet.commons.beans.proofs.ZeroProofInitialMessage;
import com.scytl.products.ov.mixnet.commons.proofs.bg.commitments.CommitmentParams;
import com.scytl.products.ov.mixnet.commons.proofs.bg.commitments.PrivateCommitment;
import com.scytl.products.ov.mixnet.commons.proofs.bg.commitments.PublicCommitment;
import com.scytl.products.ov.mixnet.commons.tools.CommitmentTools;
import com.scytl.products.ov.mixnet.commons.tools.ExponentTools;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiation;
import com.scytl.products.ov.mixnet.commons.tools.RandomOracleHash;
import com.scytl.products.ov.mixnet.proofs.bg.hadamard.HadamardProductProofGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.product.ProductProofGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.state.StateHarvester;
import com.scytl.products.ov.mixnet.proofs.bg.state.StateLabel;

/**
 * Implements the parallel mode.
 */
public class ParallelZeroProofGenerator implements ZeroProofGenerator {

    private static final Logger LOG = LoggerFactory.getLogger("std");

    private final int n;

    private int m;

    private final BigInteger groupOrder;

    private final CommitmentParams comParams;

    private final MultiExponentiation multiExponentiation;

    public ParallelZeroProofGenerator(final int m, final int n, final BigInteger order,
            final CommitmentParams comParams, final MultiExponentiation multiExponentiation) {

        this.groupOrder = order;
        this.n = n;
        this.m = m;

        // Note: m must be set to 2 in this proof when the m of
        // the Hadamard proof is 1, but only in this case
        if (this.m == 1) {
        	this.m = 2;
        }
        this.comParams = comParams;
        this.multiExponentiation = multiExponentiation;
    }

    @Override
    public void generate(final StateHarvester harvester) {

        harvester.collect(Produced.RANDOM_ORACLE, new RandomOracleHash(groupOrder));

        LOG.info("\t PP ZP Preparing commitments...");
        prepareCommitments(harvester);
        LOG.info("\t PP ZP Preparing exponents...");

        prepareExponents(harvester);

        LOG.info("\t PP ZP Generating initial message...");
        generateInitialMessage(harvester);

        LOG.info("\t PP ZP Generating answer...");
        computeAnswer(harvester);
    }

    private void prepareCommitments(final StateHarvester harvester) {

        try {
            final PrivateCommitment[] ca =
                harvester.offer(ProductProofGenerator.Produced.PRIVATE_COMMITS_Ca, PrivateCommitment[].class);

            // computes the vector a_{1}...a_{m-1} , -1
            // simply copies a_1...a_m-1 (not copying a[0]) and adds a commitment of
            // -1 with randomness = 0

            final PrivateCommitment[] cD = harvester.offer(HadamardProductProofGenerator.Produced.PRIVATE_COMMITMENTS_D,
                PrivateCommitment[].class);

            PublicCommitment[] statementCB = CommitmentTools.makePublic(cD);

            final PrivateCommitment[] innerCB = new PrivateCommitment[m + 1];
            System.arraycopy(cD, 0, innerCB, 0, m);

            harvester.collect(Produced.PRIVATE_COMMITMENTS_cB, innerCB);

            final PrivateCommitment[] cZeroArgA = new PrivateCommitment[m];
            System.arraycopy(ca, 1, cZeroArgA, 0, m - 1);
            cZeroArgA[m - 1] = new PrivateCommitment(ExponentTools.getMinus1Vector(n, groupOrder),
                new Exponent(groupOrder, BigInteger.ZERO), comParams, multiExponentiation);

            PublicCommitment[] hadamardStatementCa =
                harvester.offer(HadamardProductProofGenerator.Produced.PUBLIC_COMMITMENTS_CA, PublicCommitment[].class);

            final PublicCommitment[] statementCA = new PublicCommitment[m];
            System.arraycopy(hadamardStatementCa, 1, statementCA, 0, m - 1);
            statementCA[m - 1] = cZeroArgA[m - 1].makePublicCommitment();

            final PrivateCommitment[] innerCA = new PrivateCommitment[m + 1];
            System.arraycopy(cZeroArgA, 0, innerCA, 1, m);

            harvester.collect(Produced.PUBLIC_COMMITMENTS_CA, statementCA);
            harvester.collect(Produced.PUBLIC_COMMITMENTS_CB, statementCB);
            harvester.collect(Produced.PRIVATE_COMMITMENTS_A, cZeroArgA);
            harvester.collect(Produced.PRIVATE_COMMITMENTS_cA, innerCA);

        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }
    }

    private void prepareExponents(final StateHarvester harvester) {

        try {
            generateExponentsArrayAndMatrixFromCommitments(m, 0, Produced.PRIVATE_COMMITMENTS_cA, harvester,
                Produced.EXPONENTS_MATRIX_A, Produced.EXPONENTS_R);

            generateExponentsArrayAndMatrixFromCommitments(m, m, Produced.PRIVATE_COMMITMENTS_cB, harvester,
                Produced.EXPONENTS_MATRIX_B, Produced.EXPONENTS_S);

            final Exponent[] t = ExponentTools.getVectorRandomExponent(2 * m + 1, groupOrder);
            t[m + 1] = new Exponent(groupOrder, BigInteger.ZERO);

            harvester.collect(Produced.EXPONENTS_T, t);

        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }
    }

    private void generateExponentsArrayAndMatrixFromCommitments(final int m, final int substitutionIndex,
            final StateLabel requiredPrivateCommitmentsLabel, final StateHarvester harvester,
            final StateLabel generatedExponentsMatrixLabel, final StateLabel generatedExponentArrayLabel) {

        final PrivateCommitment[] requiredPrivateCommitments =
            harvester.offer(requiredPrivateCommitmentsLabel, PrivateCommitment[].class);

        final Exponent[][] generatedExponentsMatrix = new Exponent[m + 1][n];
        final Exponent[] generatedExponentArray = new Exponent[m + 1];

        // Resolve loop bounds depending on the substitutionIndex
        int loopStart = 0;
        int loopEnd = generatedExponentsMatrix.length - 1;
        if (substitutionIndex == 0) {
            loopStart = 1;
            loopEnd = generatedExponentsMatrix.length;
        }

        for (int i = loopStart; i < loopEnd; i++) {
            generatedExponentsMatrix[i] = requiredPrivateCommitments[i].getM();

            generatedExponentArray[i] = requiredPrivateCommitments[i].getR();
        }

        generatedExponentsMatrix[substitutionIndex] = ExponentTools.getVectorRandomExponent(n, groupOrder);

        generatedExponentArray[substitutionIndex] = ExponentTools.getRandomExponent(groupOrder);

        requiredPrivateCommitments[substitutionIndex] =
            new PrivateCommitment(generatedExponentsMatrix[substitutionIndex],
                generatedExponentArray[substitutionIndex], comParams, multiExponentiation);

        harvester.collect(generatedExponentsMatrixLabel, generatedExponentsMatrix);
        harvester.collect(generatedExponentArrayLabel, generatedExponentArray);
    }

    private void generateInitialMessage(final StateHarvester harvester) {

        final RandomOracleHash randomOracle = harvester.offer(Produced.RANDOM_ORACLE, RandomOracleHash.class);

        final PrivateCommitment[] cA = harvester.offer(Produced.PRIVATE_COMMITMENTS_cA, PrivateCommitment[].class);

        final PrivateCommitment[] cB = harvester.offer(Produced.PRIVATE_COMMITMENTS_cB, PrivateCommitment[].class);

        final Exponent[] t = harvester.offer(Produced.EXPONENTS_T, Exponent[].class);

        final Exponent[] d = computeD(harvester);
        final PrivateCommitment[] cD = new PrivateCommitment[2 * m + 1];
        for (int k = 0; k < 2 * m + 1; k++) {
            cD[k] = new PrivateCommitment(new Exponent[] {d[k] }, t[k], comParams, multiExponentiation);
        }

        final ZeroProofInitialMessage ini = new ZeroProofInitialMessage(cA[0], cB[m], cD);

        final PublicCommitment[] statementCa =
            harvester.offer(Produced.PUBLIC_COMMITMENTS_CA, PublicCommitment[].class);
        final PublicCommitment[] statementCb =
            harvester.offer(Produced.PUBLIC_COMMITMENTS_CB, PublicCommitment[].class);

        randomOracle.addDataToRO(statementCa);
        randomOracle.addDataToRO(statementCb);
        randomOracle.addDataToRO(ini);

        harvester.collect(Produced.INITIAL_MSG, ini);
    }

    private Exponent[] computeD(final StateHarvester harvester) {

        try {
            final Exponent[][] a = harvester.offer(Produced.EXPONENTS_MATRIX_A, Exponent[][].class);

            final Exponent[][] b = harvester.offer(Produced.EXPONENTS_MATRIX_B, Exponent[][].class);

            final Exponent challengeInnerProduct =
                harvester.offer(HadamardProductProofGenerator.Produced.CHALLENGE_INNER_PRODUCT, Exponent.class);

            // We compute (and commit to) d[k]
            // According to the paper by Bayer and Groth,
            // d[k] is computed as
            // \sum_{0\leq i,j\leq m; j=(m-k+i)} \vec{a}_i \innerProduct \vec{b}_j
            //
            // The computation makes a difference between k
            // smaller or equal than m and greater than m, in order to fit the
            // calculation in a single loop without if-else clauses

            final Exponent[] d = new Exponent[2 * m + 1];
            Exponent zero = new Exponent(groupOrder, BigInteger.ZERO);
            for (int k = 0; k < 2 * m + 1; k++) {
                Exponent aux = zero;
                if (k < m + 1) {
                    for (int i = 0, j = (m - k) + i; i < m + 1 && j < m + 1 && j > -1; i++, j++) {
                        aux = aux.add(ExponentTools.innerProduct(a[i], b[j], groupOrder, challengeInnerProduct));
                    }
                } else {
                    for (int j = 0, i = j - (m - k); i < m + 1 && j < m + 1 && i > -1; i++, j++) {
                        aux = aux.add(ExponentTools.innerProduct(a[i], b[j], groupOrder, challengeInnerProduct));
                    }
                }
                d[k] = aux;
            }
            return d;

        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }
    }

    private static void computeAnswer(final StateHarvester harvester) {

        final RandomOracleHash randomOracle = harvester.offer(Produced.RANDOM_ORACLE, RandomOracleHash.class);

        final Exponent challengeX = randomOracle.getHash();

        final Exponent[][] a = harvester.offer(Produced.EXPONENTS_MATRIX_A, Exponent[][].class);
        Exponent[] answerA = ExponentTools.evaluateVecPolyTransposeFrom1(a, challengeX);

        final Exponent[][] b = harvester.offer(Produced.EXPONENTS_MATRIX_B, Exponent[][].class);
        Exponent[] answerB = ExponentTools.evaluateVecBackwardsPolyTransposeFrom1(b, challengeX);

        final Exponent[] s = harvester.offer(Produced.EXPONENTS_S, Exponent[].class);
        Exponent answerS = ExponentTools.evaluateBackwardsPolynomialFrom1(s, challengeX);

        final Exponent[] r = harvester.offer(Produced.EXPONENTS_R, Exponent[].class);
        Exponent answerR = ExponentTools.evaluatePolynomialFrom1(r, challengeX);

        final Exponent[] t = harvester.offer(Produced.EXPONENTS_T, Exponent[].class);
        Exponent answerT = ExponentTools.evaluatePolynomialFrom1(t, challengeX);

        harvester.collect(Produced.ANSWER, new ZeroProofAnswer(answerA, answerB, answerR, answerS, answerT));
    }
}
