package meerkat.mixer.network;

import org.junit.Before;
import org.junit.Test;

import java.util.Random;
import java.util.Set;
import java.util.TreeSet;

import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

/**
 * Generically Test a permutation network
 */
abstract public class PermutationNetworkTest {

    public static final int NUM_REPS = 10;
    Random rand;

    abstract protected PermutationNetwork getNewNetwork();

    PermutationNetwork network;

    @Before
    public void setup() {
        network = getNewNetwork();
        rand = new Random(1);
    }

    public static Set<Integer> getSequenceSet(int N) {
        Set<Integer> set = new TreeSet<>();

        for (int i = 0; i < N; ++i) {
            set.add(i);
        }
        return set;
    }

    public static Integer[] getSequenceArray(int N) {
        Integer[] arr = new Integer[N];

        for (int i = 0; i < N; ++i) {
            arr[i] = i;
        }
        return arr;
    }


    /**
     * Check if a given  network is actually a permutation network (i.e., always
     * implies a permutation regardless of 2x2 switch settings).
     */
    @Test
    public void isAlwaysAPermutation() {
        int numLayers = network.getNumLayers();
        int N = network.getNumInputs();

        for (int layer = 1; layer < numLayers; ++layer) {
            Set<Integer> unusedInputs = getSequenceSet(N);
            for (int i = 0; i < N; ++i) {
                unusedInputs.remove(network.getOutputIdxInPreviousLayer(layer, i));
            }
            assertTrue("Not a permutation! Didn't use: " + unusedInputs, unusedInputs.isEmpty());
        }
    }

    @Test
    public void forwardEqualsBackwards() {
        int numLayers = network.getNumLayers();
        int N = network.getNumInputs();

        for (int layer = 1; layer < numLayers; ++layer) {
            for (int i = 0; i < N; ++i) {
                int j = network.getOutputIdxInPreviousLayer(layer, i);
                assertEquals(String.format("Input %d in layer %d has problems", i, layer), i,
                        network.getInputIdxInNextLayer(layer - 1, j));
            }
        }
    }

    public static int[] convert(Integer[] in) {
        int[] out = new int[in.length];
        for (int i = 0; i < in.length; ++i)
            out[i] = in[i];

        return out;
    }

    @Test
    public void testRandomPermutations() throws Exception {
        for (int rep = 0; rep < NUM_REPS; ++rep) {
            Integer[] target = getSequenceArray(network.getNumInputs());
            RandomPermutation.permute(target, rand);

            network.setPermutation(convert(target));

            Integer[] id = getSequenceArray(network.getNumInputs());

            Integer[] out = new Integer[target.length];

            Util.permute(network, id, out);

            assertArrayEquals("Permutation mismatch: " + target + " != " + out, target, out);

        }
    }
}
