/*
 * 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.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

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

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.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.ApplicationConfig;
import com.scytl.products.ov.mixnet.commons.configuration.MatrixDimensions;
import com.scytl.products.ov.mixnet.commons.configuration.MixerConfig;
import com.scytl.products.ov.mixnet.commons.configuration.VerifierConfig;
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.proofs.bg.commitments.CommitmentParams;
import com.scytl.products.ov.mixnet.commons.proofs.bg.config.ProofsGeneratorConfigurationParams;
import com.scytl.products.ov.mixnet.commons.tests.CommonTestingUtilities;
import com.scytl.products.ov.mixnet.commons.tests.categories.UltraSlowTest;
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.MatrixCalculator;
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;

@Category(UltraSlowTest.class)
public class BGMixnetIEdgeCasesTest extends BaseBGMixnet {

    private static final String PARAMETER_FOLDER = "src/test/resources/input/n512_p256_q32";

    private static int concurrencyLevel;

    private static ZpSubgroup zp;

    private static ElGamalShuffler shuffler;

    private static GjosteenElGamal elgamal;

    private static ElGamalEncryptedBallots encryptedBallots;

    private static MultiExponentiation limMultiExpo;

    private static CiphertextTools ciphertextTools;

    private static ComputeAllE computeAllE;

    private static LocationsProvider locations;

    private static ElGamalPublicKey pubKey;

    @BeforeClass
    public static void setUp() throws IOException {

        limMultiExpo = MultiExponentiationImpl.getInstance();

        ciphertextTools = new CiphertextTools(limMultiExpo);

        computeAllE = new ParallelComputeAllE(ciphertextTools);
    }

    @Test
    public void perform_mixing_when_n_mod_mu_is_not_equal_to_zero_and_iterations_are_zero() throws Exception {

        int numBallots = 28;

        int numSubkeys = 1;

        int m = 4;

        int n = 7;

        int mu = 2;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 1);
    }

    @Test
    public void perform_mixing_when_n_mod_mu_is_not_equal_to_zero_and_iterations_are_two() throws Exception {

        int numBallots = 28;

        int numSubkeys = 1;

        int m = 4;

        int n = 7;

        int mu = 2;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 1);
    }

    @Test
    public void perform_mixing_when_m_mod_mu_is_not_equal_to_zero_and_iterations_are_zero() throws Exception {

        int numBallots = 52;

        int numSubkeys = 1;

        int m = 4;

        int n = 13;

        int mu = 8;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 1);
    }

    @Test
    public void perform_mixing_for_N40_k1_m1_n40() throws Exception {

        int numBallots = 40;

        int numSubkeys = 1;

        int m = 1;

        int n = 40;

        int mu = 2;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 1);
    }

    @Test
    public void perform_mixing_for_N512_k1_m1_n512() throws Exception {

        int numBallots = 512;

        int numSubkeys = 1;

        int m = 1;

        int n = 512;

        int mu = 2;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 1);
    }

    @Test
    public void performing_mixing_for_all_numbers_of_encrypted_ballots_up_to_double_the_block_size() throws Exception {

        int numSubkeys = 1;

        int mu = 1;

        int numiterations = 0;

        int minBlockSize = 16;

        MatrixCalculator matrixCalculator = new MatrixCalculator();

        for (int numBallots = 2; numBallots < (2 * minBlockSize); numBallots++) {

            int m;
            int n;

            if (BigInteger.valueOf(numBallots).isProbablePrime(100)) {
                m = 1;
                n = numBallots;
            } else {

                List<MatrixDimensions> combinations = matrixCalculator.generateCombinationsWithoutOnes(numBallots);

                MatrixDimensions combination = matrixCalculator.getMatrixWithSmallestM(combinations);

                m = combination.getNumberOfRows();
                n = combination.getNumberOfColumns();
            }

            System.out.println("Mixing with numBallots: " + numBallots + ", m: " + m + ", n: " + n);

            performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 1);
        }
    }

    @Test
    public void perform_mixing_for_N20_k10_m2_n10() throws Exception {

        int numBallots = 20;

        int numSubkeys = 10;

        int m = 2;

        int n = 10;

        int mu = 2;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 10);
    }

    @Test
    public void perform_mixing_for_N20_k10_m10_n2() throws Exception {

        int numBallots = 20;

        int numSubkeys = 10;

        int m = 10;

        int n = 2;

        int mu = 2;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 10);
    }

    @Test
    public void perform_mixing_for_N20_k10_m4_n5() throws Exception {

        int numBallots = 20;

        int numSubkeys = 10;

        int m = 4;

        int n = 5;

        int mu = 2;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 10);
    }

    @Test
    public void perform_mixing_for_N20_k10_m5_n4() throws Exception {

        int numBallots = 20;

        int numSubkeys = 10;

        int m = 5;

        int n = 4;

        int mu = 2;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 10);
    }

    @Test
    public void perform_mixing_for_N20_k10_m1_n20() throws Exception {

        int numBallots = 20;

        int numSubkeys = 10;

        int m = 1;

        int n = 20;

        int mu = 2;

        int numiterations = 0;

        performMixingAndAssert(numBallots, numSubkeys, m, n, mu, numiterations, 10);
    }

    private void performMixingAndAssert(final int numBallots, int numSubkeys, final int m, final int n, final int mu,
            final int numiterations, int numPhis)
            throws IOException, GeneralCryptoLibException {

        String parameterFolder = PARAMETER_FOLDER;
        if (numPhis > 1) {
            parameterFolder += "_phis" + String.valueOf(numPhis);
        }

        final ApplicationConfig mixnetConfig = new ApplicationConfig();
        final MixerConfig mixerConfig = new MixerConfig();
        mixerConfig.setEncryptionParametersInput(
            new FileInputStream(Paths.get("src/test/resources/input/encryption_parameters.json").toFile()));
        mixerConfig.setEncryptedBallotsInput(
            new FileInputStream(Paths.get(parameterFolder + "/encrypted_ballots" + ".csv").toFile()));
        List<OutputStream> mixedBallots = new ArrayList<>();
        mixedBallots.add(new ByteArrayOutputStream());
        List<InputStream> mixedBallotsInput = new ArrayList<>();
        mixedBallotsInput.add(new ByteArrayInputStream("".getBytes()));
        mixerConfig.setDecryptedVotesOutputOfBatch(mixedBallots);
        mixerConfig.setEncryptedBallotsOutputOfBatch(mixedBallots);
        mixnetConfig.setMixer(mixerConfig);
        final VerifierConfig verifierConfig = new VerifierConfig();
        verifierConfig.setDecryptedVotesInputOfBatch(mixedBallotsInput);
        mixnetConfig.setVerifier(verifierConfig);

        locations = LocationsProvider.build(mixnetConfig);

        try {
            zp = ZpSubgroup.fromJson(IOUtils.toString(
                locations.getMixerLocationsProvider().getEncryptionParametersInput(), StandardCharsets.UTF_8));
        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }

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

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

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

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

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

        // /////////////////////////////////////////////////
        //
        // Create a key with the required number of subkeys
        //
        // /////////////////////////////////////////////////

        // Create ElGamal public keys (we dont need the private key)
        final List<ZpGroupElement> publicKeyGroupElements = new ArrayList<>();

        try {
            for (int i = 0; i < numSubkeys; i++) {
                publicKeyGroupElements.add( //
                    zp.getGenerator().exponentiate( //
                        ExponentTools.getRandomExponent(zp.getQ())));
            }

            pubKey = new ElGamalPublicKey(publicKeyGroupElements, zp);

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

        elgamal = new GjosteenElGamal(zp, pubKey);

        // /////////////////////////////////////////////////
        //
        // Obtain the required number of encrypted ballots
        //
        // /////////////////////////////////////////////////

        ElGamalEncryptedBallots filteredEncryptedBallots =
            CommonTestingUtilities.filterBallots(encryptedBallots, numBallots);

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

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

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

        CommitmentParams commitmentParams = new CommitmentParams(zp, n);

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

        final Ciphertext[][] ciphertexts = MatrixArranger.arrangeInCiphertextMatrix(filteredEncryptedBallots, m, n);

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

        final PrivateAndPublicCommitmentsGenerator privateAndPublicCommitmentsGenerator =
            new ParallelPrivateAndPublicCommitmentsGenerator();

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

        final ShuffleProof shuffleProof = prover.generate();

        // /////////////////////////////////////////////////
        //
        // Verify the proofs
        //
        // /////////////////////////////////////////////////

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

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

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

        final ShuffleProof shuffleProof_VerifierCopy = shuffleProof;

        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);
    }
}
