package edu.rice.hj.example.comp322;

import edu.rice.hj.api.HjRegion;

import static edu.rice.hj.Module1.*;

/**
 * <p>OneDimAveragingBarrier class.</p>
 *
 * @author Shams Imam (shams@rice.edu)
 */
public class OneDimAveragingBarrier {

    /**
     * Constant <code>initialOutput</code>
     */
    public static double[] initialOutput;

    /**
     * <p>main.</p>
     *
     * @param args an array of {@link java.lang.String} objects.
     */
    public static void main(final String[] args) {

        final int tasks = (args.length > 0) ? Integer.parseInt(args[0]) : 2;
        final int n = (args.length > 1) ? Integer.parseInt(args[1]) : 200000;
        final int iterations = (args.length > 2) ? Integer.parseInt(args[2]) : 20000;
        final int rounds = (args.length > 3) ? Integer.parseInt(args[3]) : 5;
        printParams(tasks, n, iterations, rounds);

        {
            // initial run to set the output to be equal to the sequential run
            final OneDimAveragingBarrier initialObj = new OneDimAveragingBarrier(n);
            initialObj.runSequential(iterations);
            setOutput(initialObj);
        }

        System.out.println("Timed executions:");
        finish(() -> {
            for (int r = 0; r < rounds; r++) {
                final OneDimAveragingBarrier serialBody = new OneDimAveragingBarrier(n);
                final OneDimAveragingBarrier forallBodyChunked = new OneDimAveragingBarrier(n);
                final OneDimAveragingBarrier forallBodyBarrier = new OneDimAveragingBarrier(n);

                System.out.println(" Round: " + r + (r == 0 ? " [ignore: warm up for JIT]" : ""));

                timeIt("Sequential", () -> {
                    serialBody.runSequential(iterations);
                }, () -> {
                    serialBody.validateOutput();
                });

                timeIt(String.format("Forall-Chunked [tasks=%d]", tasks), () -> {
                    forallBodyChunked.runForallChunked(iterations, tasks);
                }, () -> {
                    forallBodyChunked.validateOutput();
                });

                timeIt(String.format("Forall-barrier [tasks=%d]", tasks), () -> {
                    forallBodyBarrier.runForallBarrier(iterations, tasks);
                }, () -> {
                    forallBodyBarrier.validateOutput();
                });

                System.out.println();
            }
        });
    }

    /**
     * <p>setOutput.</p>
     *
     * @param oneDimAveraging a {@link edu.rice.hj.example.comp322.OneDimAveragingBarrier} object.
     */
    public static void setOutput(final OneDimAveragingBarrier oneDimAveraging) {
        if (OneDimAveragingBarrier.initialOutput != null) {
            System.out.println("Warning: initialOutput has already been set.");
        }
        OneDimAveragingBarrier.initialOutput = oneDimAveraging.myVal;
    }

    /**
     * <p>printParams.</p>
     *
     * @param tasks      a int.
     * @param n          a int.
     * @param iterations a int.
     * @param rounds     a int.
     */
    public static void printParams(int tasks, int n, int iterations, int rounds) {
        System.out.println("Configuration: ");
        System.out.println("  # tasks for parallel run: " + tasks);
        System.out.println("  Array size n: " + n);
        System.out.println("  # iterations: " + iterations);
        System.out.println("  Rounds: " + rounds + " (to reduce JIT overhead)");
        System.out.println();
    }

    private static void timeIt(
            final String label,
            final Runnable actualBody,
            final Runnable postExecBody) {

        final long s = System.currentTimeMillis();
        actualBody.run();
        final long e = System.currentTimeMillis();

        postExecBody.run();
        System.out.printf("%25s Time: %6d ms. \n", label, (e - s));
    }

    public double[] myNew, myVal;
    public int n;

    /**
     * <p>Constructor for OneDimAveragingBarrier.</p>
     *
     * @param n a int.
     */
    public OneDimAveragingBarrier(int n) {
        this.n = n;
        this.myNew = new double[n + 2];
        this.myVal = new double[n + 2];
        this.myVal[n + 1] = 1.0;
    }

    /**
     * <p>validateOutput.</p>
     */
    public void validateOutput() {
        if (OneDimAveragingBarrier.initialOutput == null) {
            System.out.println("initialOutput is null");
            return;
        } else if (myVal == null) {
            System.out.println("myVal is null");
            return;
        }

        for (int i = 0; i < n + 2; i++) {
            double init = OneDimAveragingBarrier.initialOutput[i];
            double curr = myVal[i];
            if (init != curr) {
                System.out.println("Diff: myVal[" + i + "]=" + curr + " != initialOutput[" + i + "]=" + init);
            }
        }
    }

    private int ceilDiv(int n, int d) {
        int m = n / d;
        if (m * d == n) {
            return m;
        } else {
            return (m + 1);
        }
    }

    /**
     * <p>runSequential.</p>
     *
     * @param iterations a int.
     */
    public void runSequential(int iterations) {
        for (int iter = 0; iter < iterations; iter++) {
            for (int j = 1; j <= n; j++) {
                myNew[j] = (myVal[j - 1] + myVal[j + 1]) / 2.0;
            }
            double[] temp = myNew;
            myNew = myVal;
            myVal = temp;
        }
    }

    /**
     * <p>runForall.</p>
     *
     * @param iterations a int.
     */
    public void runForall(final int iterations) {
        for (int iter = 0; iter < iterations; iter++) {
            forall(1, n, (j) -> {
                myNew[j] = (myVal[j - 1] + myVal[j + 1]) / 2.0;
            });
            double[] temp = myNew;
            myNew = myVal;
            myVal = temp;
        }
    }

    /**
     * <p>runForallChunked.</p>
     *
     * @param iterations a int.
     * @param tasks      a int.
     */
    public void runForallChunked(final int iterations, final int tasks) {
        final HjRegion.HjRegion1D iterSpace = newRectangularRegion1D(1, n);

        for (int iter = 0; iter < iterations; iter++) {
            forallChunked(1, n, n / tasks, (j) -> {
                myNew[j] = (myVal[j - 1] + myVal[j + 1]) / 2.0;
            });
            double[] temp = myNew;
            myNew = myVal;
            myVal = temp;
        }
    }

    /**
     * <p>runForallBarrier.</p>
     *
     * @param iterations a int.
     * @param tasks      a int.
     */
    public void runForallBarrier(final int iterations, final int tasks) {
        final HjRegion.HjRegion1D iterSpace = newRectangularRegion1D(1, n);

        forallPhased(0, tasks - 1, (Integer i) -> {
            final HjRegion.HjRegion1D myGroup = myGroup(i, iterSpace, tasks);

            for (int iter = 0; iter < iterations; iter++) {
                final double[] myVal = (iter % 2 == 0) ? this.myVal : this.myNew;
                final double[] myNew = (iter % 2 == 0) ? this.myNew : this.myVal;
                forseq(myGroup, (j) -> {
                    myNew[j] = (myVal[j - 1] + myVal[j + 1]) / 2.0;
                });

                next();  // BARRIER
            }
        });
    }
}