package meerkat.mixer;

import com.google.protobuf.InvalidProtocolBufferException;
import meerkat.crypto.Encryption;
import meerkat.mixer.network.BenesNetwork;
import meerkat.mixer.network.PermutationNetwork;
import meerkat.mixer.network.RandomPermutation;
import meerkat.mixer.proofs.Mix2nizk;
import meerkat.protobuf.Crypto.EncryptionRandomness;
import meerkat.protobuf.Crypto.RerandomizableEncryptedMessage;
import meerkat.protobuf.Mixing.Mix2Proof;

import java.util.List;
import java.util.Random;

/**
 * Run Abe's permutation network mixing algorithm on set of n encrypted votes:
 *  0. n is power of two
 *  1. set switches according to Benes network on random permutation
 *  2. re encrypt and mix with respect to switches values (encryptor.rerandomize)
 *  3. generate zero knowledge proof on each re encrypted couple
 *  4. return proofs table + encryption table
 *
 */
public class MixGenerator {

    private final Mix2nizk.Prover prover;
    private final Encryption encryptor;

    /**
     *  constructor
     * @param prover
     * @param encryptor
     */
    public MixGenerator(Mix2nizk.Prover prover, Encryption encryptor) {
        this.prover = prover;
        this.encryptor = encryptor;
    }

    /**
     * return True iff n == 2 ^ k
     * @param n
     * @return
     */
    public static boolean isPowerOfTwo(int n){
        return (n & (n - 1)) == 0;
    }

    /**
     * initialize encryption table of size (layers + 1)* n
     * @param n number of votes
     * @param layers
     * @param ciphertexts encrypted votes
     * @return an initialized encryption table s.t first layer == given encrypted votes
     */
    private RerandomizableEncryptedMessage[][] initializeEncryptionTable(int n,int layers,List<RerandomizableEncryptedMessage> ciphertexts){
        // set first level of encryption
        RerandomizableEncryptedMessage[][] encryptionTable = new RerandomizableEncryptedMessage[layers + 1][n];
        for (int j = 0; j < n; j++) {
            encryptionTable[0][j] = ciphertexts.get(j);
        }
        return  encryptionTable;
    }

    /**
     * generate randomness for all rerandomize operations
     * @param n number of votes
     * @param layers
     * @return an initialized randomness table of size layers * n
     */
    private  EncryptionRandomness[][] generateRandomnessesForRerandomize(int n,int layers,Random random){
        EncryptionRandomness[][] randomnesses = new EncryptionRandomness[layers][n];
        for (int layer = 0; layer < layers; layer++)
        {
            for (int i = 0; i < n; i++) {
                randomnesses[layer][i] = encryptor.generateRandomness(random);;
            }
        }
        return randomnesses;
    }

    /**
     * generate new random mix network
     * @param logN log number of votes
     * @param random
     * @return new random mix network
     */
    private PermutationNetwork generateMixNetwork(int logN, Random random){
        BenesNetwork net = new BenesNetwork(logN);
        RandomPermutation perm = new RandomPermutation(1 << logN, random);
        net.setPermutation(perm.permutation);
        return net;
    }


    /**
     * fills the encryption table with rerandomized encrypted votes.
     * @param mixNetwork switches table (boolean values)
     * @param encryptionTable an initialized encryption table s.t first layer == given encrypted votes
     * @param randomnesses an initialized randomness table of size layers * n, use for rerandomize operations
     * @throws InvalidProtocolBufferException
     */
    private void rerandomize(PermutationNetwork mixNetwork,
                             RerandomizableEncryptedMessage[][] encryptionTable,
                             EncryptionRandomness[][] randomnesses) throws InvalidProtocolBufferException {
        RerandomizableEncryptedMessage a,b;
        EncryptionRandomness r1,r2;
        int layers = mixNetwork.getNumLayers();

        int numSwitches =  mixNetwork.getNumInputs() >>> 1;

        for (int layer = 0; layer < layers; layer++) {
            for (int switchIdx = 0; switchIdx < numSwitches; ++switchIdx) {
                boolean isCrossed = mixNetwork.isCrossed(layer, switchIdx);

                int out0 = mixNetwork.getInputIdxInNextLayer(layer, 2 * switchIdx);
                int out1 = mixNetwork.getInputIdxInNextLayer(layer, 2 * switchIdx + 1);

                r1 = randomnesses[layer][2 * switchIdx];
                r2 = randomnesses[layer][2 * switchIdx + 1];

                a = encryptionTable[layer][2 * switchIdx];
                b = encryptionTable[layer][2 * switchIdx + 1];

                if (isCrossed) {
                    encryptionTable[layer + 1][out0] = encryptor.rerandomize(b, r2);
                    encryptionTable[layer + 1][out1] = encryptor.rerandomize(a, r1);
                } else {
                    encryptionTable[layer + 1][out0] = encryptor.rerandomize(a, r1);
                    encryptionTable[layer + 1][out1] = encryptor.rerandomize(b, r2);
                }

            }
        }
    }

    /**
     *  generate zero knowledge proof for each rerandomize encrypted votes couple in encryptionTable
     * @param mixNetwork switches table (boolean values) used for set encryption table
     * @param encryptionTable full encryption table
     * @param randomnesses randomness table of size layers * n, used for set encryption table
     * @return zero knowledge proofs table
     * @throws InvalidProtocolBufferException
     */
    private Mix2Proof[][] generateMix2ProofTable(PermutationNetwork mixNetwork
            , RerandomizableEncryptedMessage[][] encryptionTable
            , EncryptionRandomness[][] randomnesses) throws InvalidProtocolBufferException {
        int layers = mixNetwork.getNumLayers();
        int n = mixNetwork.getNumInputs();
        int numSwitches = n >> 1;
        Mix2Proof[][] proofsTable = new Mix2Proof[layers][numSwitches];

        RerandomizableEncryptedMessage a,b,c,d;
        EncryptionRandomness r1,r2;
        for (int layer = 0; layer < layers; layer++) {
            for (int switchIdx = 0; switchIdx < numSwitches; ++switchIdx) {
                boolean isCrossed = mixNetwork.isCrossed(layer, switchIdx);
                int out0 = mixNetwork.getInputIdxInNextLayer(layer, 2 * switchIdx);
                int out1 = mixNetwork.getInputIdxInNextLayer(layer, 2 * switchIdx + 1);

                a = encryptionTable[layer][2 * switchIdx];
                b = encryptionTable[layer][2 * switchIdx + 1];
                c = encryptionTable[layer + 1][out0];
                d = encryptionTable[layer + 1][out1];
                r1 = randomnesses[layer][2 * switchIdx];
                r2 = randomnesses[layer][2 * switchIdx + 1];

                proofsTable[layer][switchIdx] =
                        prover.prove(a, b, c, d, isCrossed, layer, switchIdx, out0, out1, r1, r2);
            }
        }
        return proofsTable;
    }

    /**
     * mix given encrypted votes using random
     * @param ciphertexts encrypted votes
     * @param random
     * @return meerkat.mixProverVerifier.network result
     * @throws InvalidProtocolBufferException
     */
    public MixerOutput mix(List<RerandomizableEncryptedMessage> ciphertexts,Random random) throws InvalidProtocolBufferException {

        int n = ciphertexts.size();
        assert (n > 1 && isPowerOfTwo(n));

        int logN = Integer.numberOfTrailingZeros(Integer.highestOneBit(n));

        PermutationNetwork net = generateMixNetwork(logN,random);

        int layers = net.getNumLayers();
        RerandomizableEncryptedMessage[][] encryptionTable = initializeEncryptionTable(n,layers,ciphertexts);
        EncryptionRandomness[][] randomnesses = generateRandomnessesForRerandomize(n,layers,random);

        rerandomize(net,encryptionTable,randomnesses);
        Mix2Proof[][] proofsTable = generateMix2ProofTable(net,encryptionTable,randomnesses);

        return new meerkat.mixer.MixerOutput(logN, proofsTable, encryptionTable);
    }
}