package lib;

import java.util.*;
import java.util.concurrent.*;

public class Matrix2D {
 //生成随机矩阵
 public static int[][] randomMatrix(int row, int col) {
 if (row < 0 || col < 0)
 throw new IllegalArgumentException("Number of rows or number of columns less than zero.");
 if (row == 0 || col == 0) return null;
 int[][] m = new int[row][col];

 ForkJoinPool pool = new ForkJoinPool();
 RandomMatrix2DTask masterTask = new RandomMatrix2DTask(m, 0, row, 0, col);
 pool.invoke(masterTask);

 return m;
 }

 //并行矩阵加法
 public static long[][] parallelMatAdd(int[][] matA, int[][] matB) {
 if (matA == null || matB == null) throw new NullPointerException();
 if (matA.length != matB.length || matA[0].length != matB[0].length)
 throw new IllegalArgumentException("Sizes of two matrices not equal.");
 long[][] resultMat = new long[matA.length][matA[0].length];

 ForkJoinPool pool = new ForkJoinPool();
 MatrixAddTask masterTask = new MatrixAddTask(resultMat, matA, matB, 0, matA.length, 0, matA[0].length);
 pool.invoke(masterTask);

 return resultMat;
 }

 //并行矩阵乘法
 public static double[][] parallelMatMul(int[][] matA, int[][] matB) {
 if (matA == null || matB == null) throw new NullPointerException();
 if (matA[0].length != matB.length)
 throw new IllegalArgumentException("nCol of matrix A not equals to nRow of matrix B.");
 double[][] resultMat = new double[matA.length][matB[0].length];

 ForkJoinPool pool = new ForkJoinPool();
 MatrixMulTask masterTask = new MatrixMulTask(resultMat, matA, matB, 0, resultMat.length, 0, resultMat[0].length);
 pool.invoke(masterTask);

 return resultMat;
 }
}

//用于并行创建随机矩阵的任务类
class RandomMatrix2DTask extends RecursiveAction {
 private final int[][] m;
 private final int rowStart;
 private final int colStart; //此任务对象处理矩阵的[rowStart,rowEnd-1]行,下同
 private final int rowEnd;
 private final int colEnd;
 private static final int THRESHOLD = 20;

 public RandomMatrix2DTask(int[][] m, int rowStart, int rowEnd, int colStart, int colEnd) {
 this.m = m;
 this.rowStart = rowStart;
 this.rowEnd = rowEnd;
 this.colStart = colStart;
 this.colEnd = colEnd;
 }

 public void compute() {
 //任务规模小,串行解决
 int nRow = rowEnd - rowStart;
 int nCol = colEnd - colStart;
 if (nRow <= THRESHOLD || nCol <= THRESHOLD) {
 Random r = new Random();
 for (int i = rowStart; i <= rowEnd - 1; i++) {
 for (int j = colStart; j <= colEnd - 1; j++) {
 m[i][j] = r.nextInt();
 }
 }
 return;
 }
 //拆分
 int rowMid = (rowStart + rowEnd) / 2;
 int colMid = (colStart + colEnd) / 2;
 RandomMatrix2DTask branchTaskA = new RandomMatrix2DTask(m, rowStart, rowMid, colStart, colMid);
 RandomMatrix2DTask branchTaskB = new RandomMatrix2DTask(m, rowStart, rowMid, colMid, colEnd);
 RandomMatrix2DTask branchTaskC = new RandomMatrix2DTask(m, rowMid, rowEnd, colStart, colMid);
 RandomMatrix2DTask branchTaskD = new RandomMatrix2DTask(m, rowMid, rowEnd, colMid, colEnd);
 invokeAll(branchTaskA, branchTaskB, branchTaskC, branchTaskD);
 }
}

//用于实现并行矩阵加法的任务类
class MatrixAddTask extends RecursiveAction {
 private final int[][] matA;
 private final int[][] matB;
 private final long[][] resultMat;
 private final int rowStart;
 private final int colStart; //此任务对象处理矩阵的[rowStart,rowEnd-1]行,下同
 private final int rowEnd;
 private final int colEnd;
 private static final int THRESHOLD = 20;

 public MatrixAddTask(long[][] resultMat, int[][] matA, int[][] matB, int rowStart, int rowEnd, int colStart, int colEnd) {
 this.resultMat = resultMat;
 this.matA = matA;
 this.matB = matB;
 this.rowStart = rowStart;
 this.rowEnd = rowEnd;
 this.colStart = colStart;
 this.colEnd = colEnd;
 }

 public void compute() {
 if (rowEnd - rowStart <= THRESHOLD || colEnd - colStart <= THRESHOLD) {
 for (int i = rowStart; i <= rowEnd - 1; i++) {
 for (int j = colStart; j <= colEnd - 1; j++) {
 resultMat[i][j] = matA[i][j] + matB[i][j];
 }
 }
 return;
 }

 int rowMid = (rowStart + rowEnd) / 2;
 int colMid = (colStart + colEnd) / 2;
 MatrixAddTask branchA = new MatrixAddTask(resultMat, matA, matB, rowStart, rowMid, colStart, colMid);
 MatrixAddTask branchB = new MatrixAddTask(resultMat, matA, matB, rowStart, rowMid, colMid, colEnd);
 MatrixAddTask branchC = new MatrixAddTask(resultMat, matA, matB, rowMid, rowEnd, colStart, colMid);
 MatrixAddTask branchD = new MatrixAddTask(resultMat, matA, matB, rowMid, rowEnd, colMid, colEnd);
 invokeAll(branchA, branchB, branchC, branchD);
 }
}

//用于实现并行矩阵乘法的任务类
class MatrixMulTask extends RecursiveAction {
 private final int[][] matA;
 private final int[][] matB;
 private final double[][] resultMat;
 private final int rowStart;
 private final int colStart; //此任务对象处理矩阵的[rowStart,rowEnd-1]行,下同
 private final int rowEnd;
 private final int colEnd;
 private static final int THRESHOLD = 5;

 public MatrixMulTask(double[][] resultMat, int[][] matA, int[][] matB, int rowStart, int rowEnd, int colStart, int colEnd) {
 this.resultMat = resultMat;
 this.matA = matA;
 this.matB = matB;
 this.rowStart = rowStart;
 this.rowEnd = rowEnd;
 this.colStart = colStart;
 this.colEnd = colEnd;
 }

 public void compute() {
 if (rowEnd - rowStart <= THRESHOLD || colEnd - colStart <= THRESHOLD) {
 for (int i = rowStart; i <= rowEnd - 1; i++) {
 for (int j = colStart; j <= colEnd - 1; j++) {
 //计算resultMat[i][j]
 double sum = 0;
 for (int k = 0; k <= matA[i].length - 1; k++) {
 sum += matA[i][k] * matB[k][j];
 }
 resultMat[i][j] = sum;
 }
 }
 return;
 }

 int rowMid = (rowStart + rowEnd) / 2;
 int colMid = (colStart + colEnd) / 2;
 MatrixMulTask branchA = new MatrixMulTask(resultMat, matA, matB, rowStart, rowMid, colStart, colMid);
 MatrixMulTask branchB = new MatrixMulTask(resultMat, matA, matB, rowStart, rowMid, colMid, colEnd);
 MatrixMulTask branchC = new MatrixMulTask(resultMat, matA, matB, rowMid, rowEnd, colStart, colMid);
 MatrixMulTask branchD = new MatrixMulTask(resultMat, matA, matB, rowMid, rowEnd, colMid, colEnd);
 invokeAll(branchA, branchB, branchC, branchD);
 }
}