package edu.rice.comp322;

import edu.rice.hj.api.HjFuture;
import edu.rice.hj.api.HjProcedure;
import edu.rice.hj.api.SuspendableException;

import java.util.Random;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.RecursiveTask;

import static edu.rice.hj.Module0.finish;
import static edu.rice.hj.Module0.launchHabaneroApp;
import static edu.rice.hj.Module1.future;

/**
 * Two-way Array sum example.
 */
public class ArraySum {
    /**
     * Constant <code>ERROR_MSG="Incorrect argument for array size"</code>.
     */
    public static final String ERROR_MSG = "Incorrect argument for array size (should be > 0), assuming default n";
    /**
     * Constant <code>DEFAULT_N=20_000_000</code>.
     */
    public static final int DEFAULT_N = 20_000_000;
    /**
     * Constant <code>THRESHOLD=1_000_000</code>.
     */
    public static final int THRESHOLD = 1_000_000;

    /**
     * <p>main.</p>
     *
     * @param args an array of {@link String} objects.
     */
    public static void main(final String[] args) throws Exception {
        // Initialization
        int n = parseLengthArgument(args);
        final double[] X = initializeArray(n);

        System.out.println("ReciprocalArraySum.main: starting...");
        for (int numRun = 0; numRun < 5; numRun++) {
            System.out.printf("  Run %d\n", numRun);

            timeExecution(() -> seqArraySum(X, 0, X.length), (p) -> printResults("seqArraySum", p));
            timeExecution(() -> parArraySumHabaneroJava(X, 0, X.length), (p) -> printResults("parArraySumHabaneroJava", p));
            timeExecution(() -> parArraySumForkJoin(X, 0, X.length), (p) -> printResults("parArraySumForkJoin", p));
            timeExecution(() -> parArraySumForkJoin2(X, 0, X.length), (p) -> printResults("parArraySumForkJoin2", p));
        }
        System.out.println("ReciprocalArraySum.main: ended.");

    }

    protected static int parseLengthArgument(final String[] argv) {
        int n;
        if (argv.length != 0) {
            try {
                n = Integer.parseInt(argv[0]);
                if (n <= 0) {
                    // Bad value of n
                    System.out.println(ERROR_MSG);
                    n = DEFAULT_N;
                }
            } catch (Throwable e) {
                System.out.println(ERROR_MSG);
                n = DEFAULT_N;
            }
        } else { // argv.length == 0
            n = DEFAULT_N;
        }
        return n;
    }

    protected static double[] initializeArray(final int n) {
        final double[] X = new double[n];
        final Random myRand = new Random(n);

        for (int i = 0; i < n; i++) {
            X[i] = myRand.nextInt(n);
            if (X[i] == 0.0) {
                i--;
            }
        }
        return X;
    }

    protected static void timeExecution(
        final Runnable body,
        final HjProcedure<Long> callback) throws Exception {

        final long startTime = System.nanoTime();

        body.run();

        final long endTime = System.nanoTime();
        final long execTimeInNanos = endTime - startTime;

        if (callback != null) {
            callback.apply(execTimeInNanos);
        }
    }

    protected static void printResults(final String name, final long timeInNanos) {
        System.out.printf("    %s completed in %8.3f milliseconds \n", name, timeInNanos / 1e6);
    }

    /**
     * <p>seqArraySum.</p>
     *
     * @param xArray an array of double.
     * @return a double.
     */
    protected static double seqArraySum(final double[] xArray, final int start, final int end) {
        double sum = 0;

        // Compute sum of array elements doing some work (square the root)
        for (int i = start; i < end; i++) {
            sum += Math.pow(2, Math.sqrt(xArray[i])) / (i + 1);
        }

        return sum;
    }

    /**
     * <p>parArraySumHabaneroJava.</p>
     * <p>
     * NOTE: you will need to add a "throws SuspendableException" clause to this method signature when it contains a
     * finish or any other blocking API.
     * </p>
     *
     * @param xArray an array of double.
     * @return a double.
     */
    protected static double parArraySumHabaneroJavaHelper(
        final double[] xArray, final int start, final int end)
        throws SuspendableException {

        final int fragmentSize = end - start;

        if (fragmentSize < THRESHOLD) {

            // sequential threshold cutoff
            return seqArraySum(xArray, start, end);

        } else {
            final int mid = (end + start) / 2;

            final HjFuture<Double> leftFuture = future(() -> {
                return parArraySumHabaneroJavaHelper(xArray, start, mid);
            });
            final HjFuture<Double> rightFuture = future(() -> {
                return parArraySumHabaneroJavaHelper(xArray, mid, end);
            });
            return leftFuture.get() + rightFuture.get();

        }
    }

    /**
     * <p>parArraySumHabaneroJava.</p>
     * <p>
     * NOTE: you will need to add a "throws SuspendableException" clause to this method signature when it contains a
     * finish or any other blocking API.
     * </p>
     *
     * @param xArray an array of double.
     * @return a double.
     */
    protected static double parArraySumHabaneroJava(
        final double[] xArray, final int start, final int end) {

        final double[] sum = {0.0};

        launchHabaneroApp(() -> {
            finish(() -> {
                sum[0] = parArraySumHabaneroJavaHelper(xArray, start, end);
            });
        });

        return sum[0];
    }

    protected static class ArraySumForkJoinTask extends RecursiveTask<Double> {

        private final double[] xArray;
        private final int start;
        private final int end;

        public ArraySumForkJoinTask(double[] xArray, int start, int end) {
            this.xArray = xArray;
            this.start = start;
            this.end = end;
        }

        @Override
        protected Double compute() {
            final int fragmentSize = end - start;

            if (fragmentSize < THRESHOLD) {

                // sequential threshold cutoff
                return seqArraySum(xArray, start, end);

            } else {

                final int mid = (end + start) / 2;
                final ArraySumForkJoinTask taskLeft = new ArraySumForkJoinTask(xArray, start, mid);
                final ArraySumForkJoinTask taskRight = new ArraySumForkJoinTask(xArray, mid, end);

                // taskLeft.fork(); taskRight.fork();
                invokeAll(taskLeft, taskRight);

                return taskLeft.join() + taskRight.join();

            }
        }
    }

    /**
     * <p>parArraySumHabaneroJava.</p>
     * <p>
     * NOTE: you will need to add a "throws SuspendableException" clause to this method signature when it contains a
     * finish or any other blocking API.
     * </p>
     *
     * @param xArray an array of double.
     * @return a double.
     */
    protected static double parArraySumForkJoin(final double[] xArray, final int start, final int end) {
        final ForkJoinPool commonPool = ForkJoinPool.commonPool();
        final ArraySumForkJoinTask task = new ArraySumForkJoinTask(xArray, start, end);
        return commonPool.invoke(task);
    }

    protected static class ArraySumForkJoinTask2 extends RecursiveTask<Double> {

        private final double[] xArray;
        private final int start;
        private final int end;

        public ArraySumForkJoinTask2(double[] xArray, int start, int end) {
            this.xArray = xArray;
            this.start = start;
            this.end = end;
        }

        @Override
        protected Double compute() {
            final int fragmentSize = end - start;

            if (fragmentSize < THRESHOLD) {

                // sequential threshold cutoff
                return seqArraySum(xArray, start, end);

            } else {

                final int mid = (end + start) / 2;
                final ArraySumForkJoinTask2 taskLeft = new ArraySumForkJoinTask2(xArray, start, mid);
                final ArraySumForkJoinTask2 taskRight = new ArraySumForkJoinTask2(xArray, mid, end);

                taskRight.fork();

                return taskLeft.compute() + taskRight.join();

            }
        }
    }

    /**
     * <p>parArraySumHabaneroJava.</p>
     * <p>
     * NOTE: you will need to add a "throws SuspendableException" clause to this method signature when it contains a
     * finish or any other blocking API.
     * </p>
     *
     * @param xArray an array of double.
     * @return a double.
     */
    protected static double parArraySumForkJoin2(final double[] xArray, final int start, final int end) {
        final ForkJoinPool commonPool = ForkJoinPool.commonPool();
        final ArraySumForkJoinTask2 task = new ArraySumForkJoinTask2(xArray, start, end);
        return commonPool.invoke(task);
    }
}