#include "stdio.h"
#include "stdlib.h"
#include "omp.h"
#include "math.h"
#include "openacc.h"

int main (int argc, char **argv);

void fillMatrix(int size, float **restrict A) {
//#pragma acc kernels loop collapse(2) pcopyin(A[0:size][0:size]) pcopyout(A[0:size][0:size]) gang(1000), vector(32)
   for (int i = 0; i < size; ++i) {
      for (int j = 0; j < size; ++j) {
        A[i][j] = ((float)i);
      }
   }
}

float** MatrixMult(int size, int nr, int nc, float **restrict A, float **restrict B,
        float **restrict C) {
//#pragma acc kernels pcopyin(A[0:size-1][0:size],B[0:size][0:size]) \
//  pcopyout(C[0:size][0:size])
#pragma omp parallel for default(none) shared(A,B,C,size)
   for (int i = 0; i < size; ++i) {
     for (int j = 0; j < size; ++j) {
       float tmp = 0.;
       for (int k = 0; k < size; ++k) {
          tmp += A[i][k] * B[k][j];
       }
       C[i][j] = tmp;
     }
   }
   return C;
}

float** MakeMatrix(int size, int nr, int nc, float **restrict arr) {
    int i;
    arr = (float **)malloc( sizeof(float *) * nr);
    arr[0] = (float *)malloc( sizeof(float) * nr * nc);
    for (i=1; i<nc; i++){
       arr[i] = (float *)(arr[i-1] + nc);
    }
    return arr;
}

void showMatrix(int size, int nr, int nc, float **restrict arr) {
   int i, j;
   for (i=0; i<size; i++){
      for (j=0; j<size; j++){
         printf("arr[%d][%d]=%f \n",i,j,arr[i][j]);
      }
   }
}

void copyMatrix(float **restrict A, float **restrict B, int nr, int nc){
//#pragma acc kernels loop copyin(A[0:nr][0:nc],B[0:nr][0:nc]) copyout(A[0:nr][0:nc]) gang(100) vector(32)
   for (int i=0; i<nc; ++i){
      for (int j=0; j<nr; ++j){
         A[i][j] = B[i][j];
      }
   }
}

int main (int argc, char **argv) {
   int i, j, k;
   float **A, **B, **C;
    
   if (argc != 4) {
      fprintf(stderr,"Use: %s size nIter verbose\n", argv[0]);
      return -1;
   }
   int size = atoi(argv[1]);
   int nIter = atoi(argv[2]);
   int verbose = atoi(argv[3]);
   int nr = size; //square matrix
   int nc = size; //square matrix
    
   if (nIter <= 0) {
      fprintf(stderr,"%s: Invalid nIter (%d)\n", argv[0],nIter);
      return -1;
   }
   A = (float**)MakeMatrix(size, nr, nc, A);
   fillMatrix(size, A);
   B = (float**)MakeMatrix(size, nr, nc, B);
   fillMatrix(size, B);
   C = (float**)MakeMatrix(size, nr, nc, C);
   if (verbose==1){
      printf("matrix A after filling: \n");
      showMatrix(size, nr, nc, A);
   }
   
// without this, A and B are copied to device at every iteration   
//#pragma acc data pcopyin(A[0:nr][0:nc],B[0:nr][0:nc],C[0:nr][0:nc]) pcopyout(C[0:nr][0:nc])
//{
   float startTime_tot = omp_get_wtime();
   for (int i=0; i<nIter; i++) {
      float startTime_iter = omp_get_wtime();
      C = MatrixMult(size, nr, nc, A, B, C);
      if (verbose==1) {
         printf("matrix A after MatrixMult(): \n");
         showMatrix(size, nr, nc, A);  // these will show garbage values unless the data region is removed
         printf("matrix C after MatrixMult(): \n");
         showMatrix(size, nr, nc, C);  // these will show garbage values unless the data region is removed
      }
      if (i%2==1) {
         copyMatrix(A, C, nr, nc); //multiply A by B and assign back to A on even iterations
         if (verbose==1){
            printf("matrix A after C gets copied to it: \n");
            showMatrix(size, nr, nc, A);  // these will show garbage values unless the data region is removed
         }
      }
      else {
         copyMatrix(B, C, nr, nc); //multiply A by B and assign back to B on odd iterations
         if (verbose==1){
            printf("matrix B after C gets copied to it: \n");
            showMatrix(size, nr, nc, B); // these will show garbage values unless the data region is removed
         }
      }
      float endTime_iter = omp_get_wtime();
      if (verbose==1) printf("%s iteration runtime %8.5g\n", argv[0], (endTime_iter-startTime_iter));
   }
   float endTime_tot = omp_get_wtime();
   printf("%s total runtime %8.5g\n", argv[0], (endTime_tot-startTime_tot));
//}
   if ((verbose==1) || (verbose==2)){
      printf("matrix C after all iterations\n");
      showMatrix(size, nr, nc, C); // this will show correct values outside of the data region
   }
   free(A); free(B); free(C); 
   return 0;
}