package com.example.foo;

import java.util.Arrays;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ForkJoinTask;

/**
 * https://humanoidreadable.wordpress.com/2014/12/31/forkjoin-nonrecursive-task/
 * 
 * @author Claude Martin
 *
 */
public final class SomeClass {
	/** This is used to fork/join once per entry of the resulting matrix. */
	private static class MatrixProduct extends ForkJoinTask<int[][]> {
		private static final long serialVersionUID = 1527639909439152978L;
		private final int[][] A;
		private final int[][] B;
		private int[][] result = null;

		public MatrixProduct(final int[][] A, final int[][] B) {
			this.A = A;
			this.B = B;
		}

		@Override
		public int[][] getRawResult() {
			return result;
		}

		@Override
		protected void setRawResult(final int[][] value) {
			this.result = value;
		}

		// see https://en.wikipedia.org/wiki/Matrix_multiplication#Definition
		@Override
		protected boolean exec() {
			final int n = A.length;
			final int p = B[0].length;
			final int m = B.length; // == A[0].length
			this.result = new int[n][p];
			@SuppressWarnings("unchecked")
			final ForkJoinTask<Integer>[][] tasks = new ForkJoinTask[n][p];
			if (m != A[0].length) {
				throw new IllegalArgumentException("Can't calculate matrix product for given input");
			}
			if (n == 0 || p == 0)
				return true;
			for (int i = 0; i < n; i++)
				for (int j = 0; j < p; j++)
					tasks[i][j] = new CalculateEntry(A, B, i, j).fork();
			for (int i = 0; i < n; i++)
				for (int j = 0; j < p; j++)
					result[i][j] = tasks[i][j].join();
			return true;
		}
	}

	/** This is used to calculate each entry. It won't create more subtasks. */
	private static class CalculateEntry extends ForkJoinTask<Integer> {
		private static final long serialVersionUID = 2781862671220715328L;
		private final int[][] A;
		private final int[][] B;
		private final int i;
		private final int j;
		private Integer result = null;

		public CalculateEntry(final int[][] A, final int[][] B, final int i, final int j) {
			this.A = A;
			this.B = B;
			this.i = i;
			this.j = j;
		}

		@Override
		public Integer getRawResult() {
			return this.result;
		}

		@Override
		protected void setRawResult(final Integer value) {
			this.result = value;
		}

		@Override
		protected boolean exec() {
			int sum = 0;
			for (int k = 0; k < B.length; k++) {
				sum += A[i][k] * B[k][j];
			}
			this.setRawResult(sum);
			return true;
		}
	}

	public static void main(String[] args) {
		// Note: The matrices would have to be much larger to get better performance than 
		// doing the same without parallel execution. But then the output is hard to read. 
		int[][] A = { { 1, 2, 3 }, { 4, 5, 6 } };
		int[][] B = { { 7, 8 }, { 9, 10 }, { 11, 12 } };
		int[][] result = ForkJoinPool.commonPool().invoke(new MatrixProduct(A, B));
		for (int[] row : result)
			System.out.println(Arrays.toString(row));
	}
}