package meerkat.mixer.network;

import java.util.*;

/**
 * Generate the Benes Network
 */
public class BenesNetwork implements PermutationNetwork
{
    /**
     * log (base 2) of number of inputs to Benes network (N = 2^{logN})
     */
    final int logN;

    /**
     * Values of switches set for a specific permutation.
     * switchValues[layer][switchNum] is true if the corresponding switch applies the identity permutation on its inputs.
     * layer can be in the range [0, 2*logN-1), switchNum in the range (0, 2^{logN-1}-1)
     */
    final boolean[][] switchValues;

    public BenesNetwork(int logN) {
        this.logN = logN;

        switchValues = new boolean[2*logN - 1][];
        for (int i = 0; i < switchValues.length; ++i)
            switchValues[i] = new boolean[1 << (logN - 1)];
    }

    /**
     * Return the index of the switch output in the previous layer that connects to a specified input in the current layer.
     * Each consecutive pair of outputs is from the same switch, and each consecutive pair of inputs goes to the same switch.
     * (that is, switch j has inputs/outputs  2j and 2j+1)
     *
     * @param logN log (base 2) of number of inputs to Benes network (N = 2^{logN})
     * @param layer current layer index. Must be between 0 and 2*logN-1 (layer -1 is the inputs and layer 2*logN-1 is the outputs)
     * @param inputIdx the input Idx for the current layer (must be between 0 and (1 << logN) - 1
     *
     * @return the requested index
     */
    public static int getOutputIdxInPreviousLayer(int logN, int layer, int inputIdx) {
        assert (layer >= 0) && (layer <= 2*logN - 1);
        assert (inputIdx >= 0) && (inputIdx < 1 << logN);

        if ((layer == 0) || (layer == 2*logN - 1) || (inputIdx & 1) == 0) {
            // layer 0 inputs and all even inputs
            // are connected straight "across" everywhere
            return inputIdx;
        }

        // --- Odd inputs are connected depending on layer ---

        int crossBit;
        if (layer < logN)
            crossBit = logN - layer;
        else
            crossBit = layer - logN + 1;

        return inputIdx ^ (1 << crossBit);
    }

    /**
     * Inverse of {@link #getOutputIdxInPreviousLayer(int, int, int)}
     * @param logN
     * @param layer
     * @param outputIdx
     * @return
     */
    public static int getInputIdxInNextLayer(int logN, int layer, int outputIdx) {
        return getOutputIdxInPreviousLayer(logN, layer + 1, outputIdx);
    }

    @Override
    public int getInputIdxInNextLayer(int layer, int outputIdx) {
        return getInputIdxInNextLayer(logN, layer, outputIdx);
    }

    @Override
    public int getNumInputs() {
        return 1 << logN;
    }

    @Override
    public int getNumLayers() {
        return 2*logN - 1;
    }

    @Override
    public int getOutputIdxInPreviousLayer(int layer, int inputIdx) {
        return getOutputIdxInPreviousLayer(logN, layer, inputIdx);
    }


    /**
     * Find in array
     * TODO: replace with more efficient data structure if this becomes a performance bottleneck
     * @param val
     * @param arr
     * @return
     */
    static private int find(int val, int[] arr) {
        for (int i = 0; i < arr.length; ++i) {
            if (arr[i] == val)
                return i;
        }
        return -1;
    }

    /**
     * The recursive algorithm attributed to H. Stone by A. Waksman
     * (see pg. 161 in <a href="http://grid.cs.gsu.edu/~wkim/index_files/permutation_network.pdf">paper</a>)
     * @param permOut
     * @param level
     */
    private void setInternalPermutation(int[] permIn, int[] permOut, int idx, int level) {
        // Make sure idx is in the the range [0, 2^{level})
        assert (idx < (1 << level));
        assert (permOut.length == (1 << (logN - level)));
        assert (permIn.length == permOut.length);

        final int Nover2 = 1 << (logN - level - 1);

        // Base case
        if (level == logN - 1) {
            if (permOut[0] == permIn[0])
                switchValues[level][idx] = false;
            else
                switchValues[level][idx] = true;
            return;
        }

        int nextLen = permOut.length >>> 1;
        int[] upperPermIn = new int[nextLen];
        int[] lowerPermIn = new int[nextLen];

        int[] upperPermOut = new int[nextLen];
        int[] lowerPermOut = new int[nextLen];


        SortedSet<Integer> unmatchedIndices = new TreeSet<>();
        for (int i = 0; i < permIn.length; ++i) {
            unmatchedIndices.add(i);
        }

        /**
         * Where the set of switches for this level and index actually starts in switchValues
         */
        int blockStart = idx << (logN - level - 1);
        int numSwitches = 1 << (logN - level - 1);

        while (!unmatchedIndices.isEmpty()) {
            int i = unmatchedIndices.first();

            do {
                //assert ((i & 1) == 0); // Algorithm guarantees only even indices are unmatched at start of loop
                unmatchedIndices.remove(i);
                unmatchedIndices.remove(i ^ 1);

                int switchNum = i >>> 1;

                // index of outPerm[i] after the last switch layer.
                int iSwitched;

                if (((i & 1) == 0) && (switchNum < numSwitches / 2) || ((i & 1) == 1) && (switchNum >= numSwitches / 2)) {
                    // i is in the upper half and even, or in the lower half and odd, so
                    // switch must be "straight" to get i to upper half.
                    switchValues[2 * logN - level - 2][blockStart + switchNum] = false;
                    iSwitched = i;

                } else {
                    switchValues[2 * logN - level - 2][blockStart + switchNum] = true;
                    iSwitched = i ^ 1;
                }

                // Index of outPerm[i] in the output of the previous layer.
                int iConnectedTo = getOutputIdxInPreviousLayer(logN - level, 2 * (logN - level) - 2, iSwitched);
                int iPairConnectedTo = getOutputIdxInPreviousLayer(logN - level, 2 * (logN - level) - 2, iSwitched ^ 1);

                upperPermOut[iConnectedTo] = permOut[i];
                lowerPermOut[iPairConnectedTo - Nover2] = permOut[i ^ 1];

                // Index of permOut[i] before the first switch layer.
                int j = find(permOut[i], permIn);
                int jSwitchNum = j >>> 1;

                // Index of permOut[i] after the first switch layer
                int jSwitched;

                if ((((j & 1) == 0) && (jSwitchNum < numSwitches / 2)) || ((j & 1) == 1) && (jSwitchNum >= numSwitches / 2)) {
                    // Even output in the upper half, or odd output in the lower half
                    // so switch needs to be "straight" to get j to the upper half
                    switchValues[level][blockStart + jSwitchNum] = false;
                    jSwitched = j;
                } else {
                    // Otherwise switch needs to be "crossed" to get j to upper half
                    switchValues[level][blockStart + jSwitchNum] = true;
                    jSwitched = j ^ 1;
                }

                int jConnectedTo = getInputIdxInNextLayer(logN - level, 0, jSwitched);
                // must be in upper half
                assert(jConnectedTo < (1 << (logN - 1)));
                upperPermIn[jConnectedTo] = permOut[i];

                // Connect from j's pair to through the lower half
                int jPairConnectedTo = getInputIdxInNextLayer(logN - level, 0, jSwitched ^ 1);
                lowerPermIn[jPairConnectedTo - Nover2] = permIn[j ^ 1];

                int iPairNext = find(permIn[j ^ 1], permOut);
                i = iPairNext ^ 1;
            } while (unmatchedIndices.contains(i));
        }

        setInternalPermutation(upperPermIn, upperPermOut, 2*idx, level + 1);
        setInternalPermutation(lowerPermIn, lowerPermOut, 2*idx + 1, level + 1);
    }

    @Override
    public void setPermutation(int[] permutation) {
        int level = 0;
        int[] identity = new int[getNumInputs()];
        for (int i = 0; i < identity.length; ++i)
            identity[i] = i;

        setInternalPermutation(identity, permutation, 0, 0);
    }

    @Override
    public boolean isCrossed(int layer, int switchIdx) {
        return switchValues[layer][switchIdx];
    }

}
