/*
 * 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;

import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.commons.io.IOUtils;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.cryptoapi.Ciphertext;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.cryptolib.primitives.service.PrimitivesService;
import com.scytl.products.ov.mixnet.commons.ballots.ElGamalEncryptedBallots;
import com.scytl.products.ov.mixnet.commons.beans.ShuffleOutput;
import com.scytl.products.ov.mixnet.commons.beans.proofs.ShuffleProof;
import com.scytl.products.ov.mixnet.commons.configuration.file.FileBasedApplicationConfig;
import com.scytl.products.ov.mixnet.commons.configuration.file.FileBasedMixerConfig;
import com.scytl.products.ov.mixnet.commons.configuration.file.FileBasedVerifierConfig;
import com.scytl.products.ov.mixnet.commons.configuration.locations.LocationsProvider;
import com.scytl.products.ov.mixnet.commons.homomorphic.impl.GjosteenElGamal;
import com.scytl.products.ov.mixnet.commons.io.ElGamalEncryptedBallotsLoader;
import com.scytl.products.ov.mixnet.commons.io.ElGamalEncryptedBallotsWriter;
import com.scytl.products.ov.mixnet.commons.io.ElgamalPublicKeyWriter;
import com.scytl.products.ov.mixnet.commons.io.JSONProofsReader;
import com.scytl.products.ov.mixnet.commons.io.JSONProofsWriter;
import com.scytl.products.ov.mixnet.commons.proofs.bg.commitments.CommitmentParams;
import com.scytl.products.ov.mixnet.commons.proofs.bg.config.ProofsGeneratorConfigurationParams;
import com.scytl.products.ov.mixnet.commons.tools.CiphertextTools;
import com.scytl.products.ov.mixnet.commons.tools.ExponentTools;
import com.scytl.products.ov.mixnet.commons.tools.MatrixArranger;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiation;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiationImpl;
import com.scytl.products.ov.mixnet.commons.validation.EncryptedBallotsDuplicationValidator;
import com.scytl.products.ov.mixnet.commons.validation.EncryptedBallotsValidator;
import com.scytl.products.ov.mixnet.commons.validation.EncryptedBallotsValidatorManager;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.ComputeAllE;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.ParallelComputeAllE;
import com.scytl.products.ov.mixnet.proofs.bg.shuffle.ParallelPrivateAndPublicCommitmentsGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.shuffle.PrivateAndPublicCommitmentsGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.shuffle.ShuffleProofGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.shuffle.ShuffleProofVerifier;
import com.scytl.products.ov.mixnet.shuffler.ElGamalShuffler;
import com.scytl.products.ov.mixnet.shuffler.permutation.PermutationGenerator;
import com.scytl.products.ov.mixnet.shuffler.permutation.Permutator;

public class BGMixnetMultiplePhisITest {

    private static ZpSubgroup zp;

    private static ElGamalShuffler shuffler;

    private static int concurrencyLevel;

    private static CommitmentParams commitmentParams;

    private static GjosteenElGamal elgamal;

    private static ProofsGeneratorConfigurationParams config;

    private static LocationsProvider locations;

    private static JSONProofsReader proofsReader;

    private static JSONProofsWriter proofsWriter;

    private static ElGamalEncryptedBallotsLoader elgamalEncryptedBallotsLoader;

    private static MultiExponentiation limMultiExpo;

    private static CiphertextTools ciphertextTools;

    private static ComputeAllE computeAllE;

    private static int mu;

    private static int numiterations;

    @BeforeClass
    public static void setUp() throws IOException {

        final FileBasedApplicationConfig mixnetConfig = new FileBasedApplicationConfig();
        mixnetConfig.setConcurrencyLevel(1);
        mixnetConfig.setMinimumBatchSize(1);
        mixnetConfig.setNumberOfParallelBatches(1);
        mixnetConfig.setNumOptions(1);
        final FileBasedMixerConfig mixerConfig = new FileBasedMixerConfig();
        final Map<String, String> inputMap = new HashMap<>();
        final FileBasedVerifierConfig verifierConfig = new FileBasedVerifierConfig();
        verifierConfig.setInputDirectory("target/output/");
        inputMap.put("encryptionParameters", "src/test/resources/input/encryption_parameters");
        mixerConfig.setInput(inputMap);
        mixerConfig.setOutputDirectory("target/output/");
        mixnetConfig.setMixer(mixerConfig);

        locations = LocationsProvider.build(FileBasedApplicationConfig.convertToStreamConfig(mixnetConfig));

        try {

            InputStream resourceAsStream =
                BGMixnetMultiplePhisITest.class.getClassLoader().getResourceAsStream("input/encryption_parameters.json");
            zp = ZpSubgroup.fromJson(IOUtils.toString(resourceAsStream, StandardCharsets.UTF_8));
        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }

        mu = 2;

        numiterations = 0;

        concurrencyLevel = Runtime.getRuntime().availableProcessors();

        proofsReader = new JSONProofsReader();

        proofsWriter = new JSONProofsWriter();

        final List<EncryptedBallotsValidator> listValidators = new ArrayList<>();
        listValidators.add(new EncryptedBallotsDuplicationValidator());
        EncryptedBallotsValidatorManager validator = new EncryptedBallotsValidatorManager(listValidators);
        elgamalEncryptedBallotsLoader = new ElGamalEncryptedBallotsLoader(validator);

        Files.createDirectories(Paths.get(mixerConfig.getOutputDirectory()));

        limMultiExpo = MultiExponentiationImpl.getInstance();

        ciphertextTools = new CiphertextTools(limMultiExpo);

        computeAllE = new ParallelComputeAllE(ciphertextTools);
    }

    private void performMixingAndAssert(int numVotes, final int numPhis, int m, int n) throws Exception {

        // ///////////////////////////////////////////////////////////////////////
        //
        // Perform initializations specific to the particular test that is
        // running
        //
        // ///////////////////////////////////////////////////////////////////////

        ElGamalPublicKey multiElementPublicKey = createMultiElementPublicKey(numPhis, zp);

        commitmentParams = new CommitmentParams(zp, n);

        elgamal = new GjosteenElGamal(zp, multiElementPublicKey);

        shuffler = new ElGamalShuffler(new Permutator(), new PermutationGenerator());

        config = new ProofsGeneratorConfigurationParams(m, n, elgamal, numiterations, mu, zp);

        final String inputFolder = "src/test/resources/input/n" + numVotes + "_p256_q32_phis" + numPhis;

        String filenameEncryptionParameters = "src/test/resources/input/encryption_parameters";
        String filenameEncryptedBallots = inputFolder + "/encrypted_ballots";
        String filenamePublicKey = inputFolder + "/public_key";
        String filenamePrivateKey = inputFolder + "/private_key";

        final FileBasedApplicationConfig mixnetConfig = new FileBasedApplicationConfig();
        final FileBasedMixerConfig mixerConfig = new FileBasedMixerConfig();
        final Map<String, String> inputMap = new HashMap<>();
        inputMap.put("encryptionParameters", filenameEncryptionParameters);
        inputMap.put("encryptedBallots", filenameEncryptedBallots);
        inputMap.put("publicKey", filenamePublicKey);
        inputMap.put("ballotDecryptionKey", filenamePrivateKey);
        mixerConfig.setInput(inputMap);
        mixerConfig.setOutputDirectory("target/");
        mixnetConfig.setMixer(mixerConfig);
        final FileBasedVerifierConfig verifierConfig = new FileBasedVerifierConfig();
        verifierConfig.setInputDirectory("target/");
        mixnetConfig.setVerifier(verifierConfig);
        mixnetConfig.setMinimumBatchSize(1000);

        locations = LocationsProvider.build(FileBasedApplicationConfig.convertToStreamConfig(mixnetConfig));

        // /////////////////////////////////////////////////
        //
        // Load encrypted ballots
        //
        // /////////////////////////////////////////////////

        final ElGamalEncryptedBallots encryptedBallots =
            elgamalEncryptedBallotsLoader.loadCSV(zp, locations.getMixerLocationsProvider().getEncryptedBallotsInput());

        // /////////////////////////////////////////////////
        //
        // Perform Shuffle
        //
        // /////////////////////////////////////////////////

        final ShuffleOutput<ElGamalEncryptedBallots> shuffleOutput =
            shuffler.sequentialShuffle(zp, multiElementPublicKey, concurrencyLevel, encryptedBallots, null, null);

        // /////////////////////////////////////////////////
        //
        // Generate proofs
        //
        // /////////////////////////////////////////////////

        final Ciphertext[][] encryptedCiphertexts = MatrixArranger.arrangeInCiphertextMatrix(encryptedBallots, m, n);

        final Ciphertext[][] reencryptedCiphertexts =
            MatrixArranger.arrangeInCiphertextMatrix(shuffleOutput.getReEncryptedBallots(), m, n);

        final PrivateAndPublicCommitmentsGenerator privateAndPublicCommitmentsGenerator =
            createPrivateAndPublicCommitmentsGenerator();

        final ShuffleProofGenerator prover = new ShuffleProofGenerator(config, commitmentParams, encryptedCiphertexts,
            shuffleOutput.getPermutation(), privateAndPublicCommitmentsGenerator, limMultiExpo, concurrencyLevel,
            reencryptedCiphertexts, shuffleOutput.getExponents(), computeAllE);

        final ShuffleProof shuffleProof = prover.generate();

        // /////////////////////////////////////////////////
        //
        // Store shuffled ballots
        //
        // /////////////////////////////////////////////////

        ElGamalEncryptedBallotsWriter.CSV.write(
            locations.getMixerLocationsProvider().getReEncryptedBallotsOutputOfBatch(1),
            shuffleOutput.getReEncryptedBallots(), new PrimitivesService().getRawMessageDigest());

        // /////////////////////////////////////////////////
        //
        // Store Proofs
        //
        // /////////////////////////////////////////////////

        proofsWriter.write(locations.getMixerLocationsProvider().getProofsOutputOfBatch(1), shuffleProof);

        // /////////////////////////////////////////////////
        //
        // Store the commitment parameters
        //
        // /////////////////////////////////////////////////

        commitmentParams
            .serializeToStream(locations.getMixerLocationsProvider().getCommitmentParametersOutputOfBatch(1));

        // /////////////////////////////////////////////////
        //
        // Store the public key
        //
        // /////////////////////////////////////////////////

        ElgamalPublicKeyWriter.serializeToStream(multiElementPublicKey,
            locations.getMixerLocationsProvider().getPublicKeyOutputOfBatch(1));

        // /////////////////////////////////////////////////
        //
        // Verify the proofs
        //
        // /////////////////////////////////////////////////
        mixnetConfig.getMixer().setOutputDirectory(null);
        locations = LocationsProvider.build(FileBasedApplicationConfig.convertToStreamConfig(mixnetConfig));

        final ElGamalEncryptedBallots encryptedBallots_VerifierCopy =
            elgamalEncryptedBallotsLoader.loadCSV(zp, locations.getMixerLocationsProvider().getEncryptedBallotsInput());

        final Ciphertext[][] originalCiphertexts_VerifierCopy =
            MatrixArranger.arrangeInCiphertextMatrix(encryptedBallots_VerifierCopy, m, n);

        final ElGamalEncryptedBallots reencryptedBallots_VerifierCopy = elgamalEncryptedBallotsLoader.loadCSV(zp,
            locations.getVerifierLocationsProvider().getReEncryptedBallotsInputOfBatch(1));

        final Ciphertext[][] reencryptedCiphertexts_VerifierCopy =
            MatrixArranger.arrangeInCiphertextMatrix(reencryptedBallots_VerifierCopy, m, n);

        final ShuffleProofVerifier verifier = new ShuffleProofVerifier(zp, elgamal, commitmentParams,
            originalCiphertexts_VerifierCopy, reencryptedCiphertexts_VerifierCopy, m, n, mu, numiterations,
            ciphertextTools, limMultiExpo);

        final ShuffleProof shuffleProof_VerifierCopy =
            proofsReader.read(locations.getVerifierLocationsProvider().getProofsInputOfBatch(1));

        final boolean testResult = verifier.verifyProof(shuffleProof_VerifierCopy.getInitialMessage(),
            shuffleProof_VerifierCopy.getFirstAnswer(), shuffleProof_VerifierCopy.getSecondAnswer());

        final String errorMsg = "proofs failed to validate";
        Assert.assertTrue(errorMsg, testResult);
    }

    private ElGamalPublicKey createMultiElementPublicKey(final int numPhis, final ZpSubgroup zp) {

        try {
            final List<ZpGroupElement> publicKeyAsListOfGroupElements = new ArrayList<>();

            for (int i = 0; i < numPhis; i++) {
                publicKeyAsListOfGroupElements
                    .add(zp.getGenerator().exponentiate(ExponentTools.getRandomExponent(zp.getQ())));
            }

            return new ElGamalPublicKey(publicKeyAsListOfGroupElements, zp);

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

    private PrivateAndPublicCommitmentsGenerator createPrivateAndPublicCommitmentsGenerator() {
        return new ParallelPrivateAndPublicCommitmentsGenerator();
    }

    @Test
    public void perform_mixing_when_encrypted_ballots_has_two_phis() throws Exception {

        int numVotes = 512;

        int numPhis = 2;

        int m = 2;

        int n = 256;

        performMixingAndAssert(numVotes, numPhis, m, n);
    }

    @Test
    public void perform_mixing_when_encrypted_ballots_has_ten_phis() throws Exception {

        int numVotes = 512;

        int numPhis = 10;

        int m = 2;

        int n = 256;

        performMixingAndAssert(numVotes, numPhis, m, n);
    }

    @Test
    public void more_phis_than_votes_10phis_2votes() throws Exception {

        int numVotes = 2;

        int numPhis = 10;

        int m = 1;

        int n = 2;

        performMixingAndAssert(numVotes, numPhis, m, n);
    }
}
