/*
 * 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.multiexp.reduction;

import java.math.BigInteger;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.List;
import java.util.concurrent.ForkJoinTask;

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.elgamal.bean.CiphertextImpl;
import com.scytl.cryptolib.elgamal.cryptoapi.Ciphertext;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
import com.scytl.products.ov.mixnet.commons.beans.proofs.MultiExponentiationBasicProofAnswer;
import com.scytl.products.ov.mixnet.commons.beans.proofs.MultiExponentiationBasicProofInitialMessage;
import com.scytl.products.ov.mixnet.commons.beans.proofs.MultiExponentiationReductionAnswer;
import com.scytl.products.ov.mixnet.commons.beans.proofs.MultiExponentiationReductionInitialMessage;
import com.scytl.products.ov.mixnet.commons.concurrent.LoopParallelizerAction;
import com.scytl.products.ov.mixnet.commons.concurrent.processor.NextCiphertextsAndCAPrimeProcessor;
import com.scytl.products.ov.mixnet.commons.homomorphic.Randomness;
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.config.ProofsGeneratorConfigurationParams;
import com.scytl.products.ov.mixnet.commons.tools.ExponentTools;
import com.scytl.products.ov.mixnet.commons.tools.LoopTools;
import com.scytl.products.ov.mixnet.commons.tools.LoopTools.Range;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiation;
import com.scytl.products.ov.mixnet.commons.tools.RandomOracleHash;
import com.scytl.products.ov.mixnet.commons.tools.RandomnessTools;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.ComputeAllE;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.MultiExponentiationProofGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.basic.BasicInitialMessageGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.basic.BasicProofGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.basic.ParallelBasicAnswerGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.shuffle.SecondCommitmentGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.state.CommonStateLabels;
import com.scytl.products.ov.mixnet.proofs.bg.state.StateHarvester;

/**
 * Implements the parallel mode.
 */
public class ParallelReductionProofGenerator implements ReductionProofGenerator {

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

    private final ProofsGeneratorConfigurationParams config;

    private final int commitmentLength;

    private final BasicProofGenerator basicProofGenerator;

    private final CommitmentParams comParams;

    private final ComputeAllE computeAllE;

    private final MultiExponentiation multiExponentiation;

    private final int concurrencyLevel;

    public ParallelReductionProofGenerator(final ProofsGeneratorConfigurationParams config, final int commitmentLength,
            final BasicProofGenerator basicProofGenerator, final CommitmentParams comParams,
            final ComputeAllE computeAllE, final MultiExponentiation multiExponentiation, final int concurrencyLevel) {

        super();

        this.comParams = comParams;
        this.config = config;
        this.commitmentLength = commitmentLength;
        this.basicProofGenerator = basicProofGenerator;

        this.computeAllE = computeAllE;
        this.multiExponentiation = multiExponentiation;
        this.concurrencyLevel = concurrencyLevel;
    }

    @Override
    public void generate(final StateHarvester harvester)
            throws GeneralCryptoLibException {

        LOG.info("\t MEP R using concurrency level " + concurrencyLevel);

        int numiterationsleft = harvester.offer(CommonStateLabels.REDUCT_ITERATIONS_LEFT, Integer.class);

        final Deque<IterationResults> bootstrap = new ArrayDeque<>();

        // With futures all of them could be linked.
        boolean keepLooping = numiterationsleft > 0;
        while (keepLooping) {

            LOG.info("\t MEP R " + numiterationsleft + " Generating initial message...");
            generateInitialMessage(harvester);
            LOG.info("\t MEP R " + numiterationsleft + " Generating answer...");
            prepareAnswer(harvester);

            // grab iteration's generated
            final Exponent s = harvester.offer(Produced.S, Exponent.class);
            final Exponent[] b = harvester.offer(Produced.B, Exponent[].class);

            final MultiExponentiationReductionInitialMessage iniReduct =
                harvester.offer(Produced.REDUCT_MULTIEXP_INIT_MSG, MultiExponentiationReductionInitialMessage.class);

            bootstrap.add(new IterationResults(s, b, iniReduct));

            numiterationsleft--;
            keepLooping = numiterationsleft > 0;
        }

        LOG.info("\t MEP R Generating basic multi-exponentiation proof...");
        basicProofGenerator.generate(harvester);

        LOG.info("\t MEP R Constructing the recursive answer...");
        constructRecursiveAnswer(harvester, bootstrap);
    }

    private void generateInitialMessage(final StateHarvester harvester) {

        final int mu = config.getMu();

        final int range = 2 * mu - 1;

        final ForkJoinTask<PrivateCommitment[]> calculationCb = ForkJoinTask.adapt(() -> {
            
            final BigInteger order = config.getZpGroupParams().getQ();
            
            // b
            final Exponent[][] b = new Exponent[range][config.getCryptosystem().getNumberOfMessages()];
            
            Exponent zero = new Exponent(order, BigInteger.ZERO);
            
            for (int i = 0; i < b.length; i++) {
                for (int j = 0; j < b[i].length; j++) {
                    if (i == (mu - 1)) {
                        b[i][j] = zero;
                    } else {
                        b[i][j] = ExponentTools.getRandomExponent(order);
                    }
                }
            }
            
            harvester.collect(Produced.EXPONENTS_MULTIARR_B, b);
            
            // s
            final Exponent[] s = ExponentTools.getVectorRandomExponent(range, order);
            s[mu - 1] = zero;
            
            harvester.collect(Produced.EXPONENTS_ARR_S, s);
            
            // cb
            final PrivateCommitment[] cb = new PrivateCommitment[range];
            for (int k = 0; k < range; k++) {
                cb[k] = new PrivateCommitment(b[k], s[k], comParams, multiExponentiation);
            }
            
            return cb;
        });
        
        calculationCb.fork();

        // tau
        final Randomness rho =
            harvester.offer(MultiExponentiationProofGenerator.Produced.CALCULATED_RANDOM_EXPONENT, Randomness.class);

        final Randomness[] tau = config.getCryptosystem().getVectorFreshRandomness(range);

        tau[mu - 1] = rho;
        harvester.collect(Produced.RANDOMNESS_TAU_ARR, tau);

        // Computing E's without randomizing factors
        // will compute the m' "inner E's" using the interpolation technique
        // will then multiply them

        final CiphertextImpl[][] ciphertexts = harvester.offer(CommonStateLabels.CIPHERTEXTS, CiphertextImpl[][].class);

        final int m = ciphertexts.length;
        final int mprime = m / mu;

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

        final Exponent[][] a = new Exponent[m][commitmentLength];
        for (int i = 0; i < m; i++) {
            a[i] = cA[i].getM();
        }

        final PrivateCommitment[] cb;
        if (calculationCb.tryUnfork()) {
            cb = calculationCb.invoke();
        } else {
            cb = calculationCb.join();
        }
        
        // If cb has been retrieved, _b is already available:
        final Exponent[] b = harvester.offer(Produced.EXPONENTS_MULTIARR_B, Exponent[].class);

        final Ciphertext[] cE = generateE(mu, range, b, tau, ciphertexts, mprime, a);

        final MultiExponentiationReductionInitialMessage ini =
            new MultiExponentiationReductionInitialMessage(cb, cE);

        final RandomOracleHash rO = new RandomOracleHash(config.getZpGroupParams().getQ());
        rO.addDataToRO(ini);

        harvester.collect(Produced.MULTIEXP_REDUCT_RO, rO);
        harvester.collect(Produced.REDUCT_MULTIEXP_INIT_MSG, ini);
    }

    private Ciphertext[] generateE(final int mu, final int range, final Exponent[] b, final Randomness[] tau,
            final Ciphertext[][] ciphertexts, final int mprime, final Exponent[][] a) {

        Ciphertext[] cE = new CiphertextImpl[range];

        for (int l = 0; l < mprime; l++) {

            // first need to prepare vecC, a for computing the inner ones
            final Ciphertext[][] vecCl = new Ciphertext[mu][commitmentLength];
            final Exponent[][] al = new Exponent[mu][commitmentLength];

            for (int i = 0; i < mu; i++) {
                vecCl[i] = ciphertexts[l * mu + i];
                al[i] = a[l * mu + i];
            }
            // now compute E without randomizing factors
            final Ciphertext[] eAux =
                computeAllE.computeAllENoRandomizing(al, vecCl, config.getZpGroupParams().getQ());
            for (int i = 0; i < eAux.length; i++) {

                try {
                    cE[i] = l == 0 ? eAux[i] : cE[i].multiply(eAux[i]);
                } catch (GeneralCryptoLibException ex) {
                    throw new CryptoLibException(ex);
                }
            }
        }

        cE = computeAllE.addRandomizingEncryption(cE, b, tau, config.getCryptosystem(), concurrencyLevel);
        return cE;
    }

    private void prepareAnswer(final StateHarvester harvester) {

        final RandomOracleHash localRandomOracle = harvester.offer(Produced.MULTIEXP_REDUCT_RO, RandomOracleHash.class);

        final Exponent challengeX = localRandomOracle.getHash();

        ForkJoinTask<Void> task = ForkJoinTask.adapt(() -> {
            generateNextIterationInput(harvester, challengeX);
            return null;
        });
        task.fork();
        
        generateBAndS(harvester, challengeX);
        
        if (task.tryUnfork()) {
            task.invoke();
        } else {
            task.join();
        }
    }

    private void generateNextIterationInput(final StateHarvester harvester, final Exponent challengeX) {

        final CiphertextImpl[][] ciphertexts = harvester.offer(CommonStateLabels.CIPHERTEXTS, CiphertextImpl[][].class);

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

        final int m = ciphertexts.length;
        final int mu = config.getMu();
        final int mprime = m / mu;

        final CiphertextImpl[][] nextCiphertexts = new CiphertextImpl[mprime][commitmentLength];

        final PrivateCommitment[] cAprime = new PrivateCommitment[mprime];

        // Ideally, the number of tasks should be independent from the number of
        // processors. In this case we use it this way to calculate the ranges.
        final List<Range> ranges = LoopTools.asRangesList(mprime, concurrencyLevel);

        final NextCiphertextsAndCAPrimeProcessor processor = new NextCiphertextsAndCAPrimeProcessor(nextCiphertexts,
            cAprime, ciphertexts, mu, cA, challengeX, config.getZpGroupParams().getQ(), commitmentLength);

        new LoopParallelizerAction(processor, ranges).invoke();

        // rho
        final Randomness[] tau = harvester.offer(Produced.RANDOMNESS_TAU_ARR, Randomness[].class);
        final Randomness rhoPrime = RandomnessTools.evaluatePolynomialFrom1(tau, challengeX);

        harvester.collect(MultiExponentiationProofGenerator.Produced.CALCULATED_RANDOM_EXPONENT, rhoPrime);
        harvester.collect(CommonStateLabels.CIPHERTEXTS, nextCiphertexts);
        harvester.collect(SecondCommitmentGenerator.Produced.PRIVATE_COMMITS, cAprime);
    }

    private static void generateBAndS(final StateHarvester harvester, final Exponent challengeX) {

        final ForkJoinTask<Exponent> calculationS = ForkJoinTask.adapt(() -> {
            final Exponent[] s = harvester.offer(Produced.EXPONENTS_ARR_S, Exponent[].class);

            return ExponentTools.evaluatePolynomialFrom1(s, challengeX);
        });
        calculationS.fork();
        harvester.collect(Produced.S, calculationS);

        final ForkJoinTask<Exponent[]> calculationB = ForkJoinTask.adapt(() -> {
            final Exponent[][] b = harvester.offer(Produced.EXPONENTS_MULTIARR_B, Exponent[][].class);

            return ExponentTools.evaluateVecPolyTransposeFrom1(b, challengeX);
        });
        calculationB.fork();
        harvester.collect(Produced.B, calculationB);
    }

    private static void constructRecursiveAnswer(final StateHarvester harvester,
            final Deque<IterationResults> iterationsResults) {

        final MultiExponentiationBasicProofInitialMessage iniBasic =
            harvester.offer(BasicInitialMessageGenerator.Produced.BASIC_MULTIEXP_INIT_MSG,
                MultiExponentiationBasicProofInitialMessage.class);

        final MultiExponentiationBasicProofAnswer ansBasic = harvester.offer(
            ParallelBasicAnswerGenerator.Produced.BASIC_MULTIEXP_ANSWER, MultiExponentiationBasicProofAnswer.class);

        // Getting the results from the last iteration
        final IterationResults lastIterationResults = iterationsResults.pollLast();

        // Constructs the answer of the last iteration, which requires the basic
        // initial message and answer.

        // So.. the last iteration initial message and answer are kept, to keep
        // creating the following answers.
        MultiExponentiationReductionAnswer previousReductAnswer = new MultiExponentiationReductionAnswer(
            lastIterationResults.getB(), lastIterationResults.getS(), iniBasic, ansBasic);
        MultiExponentiationReductionInitialMessage previousReductInitMsg = lastIterationResults.getInitMsg();

        while (!iterationsResults.isEmpty()) {

            final IterationResults results = iterationsResults.pollLast();

            // Creating the answer with the previous initial message and answer.
            previousReductAnswer = new MultiExponentiationReductionAnswer(results.getB(), results.getS(),
                previousReductInitMsg, previousReductAnswer);

            // Keeping the initial message for the next iteration.
            previousReductInitMsg = results.getInitMsg();
        }

        // publish
        harvester.collect(Produced.REDUCT_MULTIEXP_ANSWER, previousReductAnswer);
        harvester.collect(Produced.REDUCT_MULTIEXP_INIT_MSG, previousReductInitMsg);
    }

    private static class IterationResults {

        private final Exponent s;

        private final Exponent[] b;

        private final MultiExponentiationReductionInitialMessage initMsg;

        /**
         * @param s
         * @param b
         * @param initMsg
         */
        public IterationResults(final Exponent s, final Exponent[] b,
                final MultiExponentiationReductionInitialMessage initMsg) {
            super();
            this.s = s;
            this.b = b;
            this.initMsg = initMsg;
        }

        /**
         * @return Returns the s.
         */
        public Exponent getS() {
            return s;
        }

        /**
         * @return Returns the b.
         */
        public Exponent[] getB() {
            return b;
        }

        /**
         * @return Returns the initMsg.
         */
        public MultiExponentiationReductionInitialMessage getInitMsg() {
            return initMsg;
        }
    }

}
