#include "mat_mul.h"

#include <cstdlib>
#include <cstdio>
#include <pthread.h>

static float *A, *B, *C;
static int M, N, K;
static int num_threads;

static void* mat_mul_thread(void *data) {
  // TODO: parallelize & optimize matrix multiplication

  int pid = * (int *) data;

  int M_div = M / num_threads;
  int start = pid * M_div;
  int end;

  if(pid==num_threads - 1)  	end = M;
  else 				end = (pid+1)*M_div;

  float temp;
  int bs = 50;
  int k_min =0;


  for (int kk = 0; kk < K; kk += bs) 
  {
 	 for (int i = start; i < end; i++) 
	 {
   		if(K>=(kk+bs)) k_min = kk+bs;
		else k_min = K;
		for (int k = kk; k < k_min; k++) 
		{
			temp = A[i * K + k];	  
			for (int j = 0; j < N; j++) 
			{
				C[i*N+(j+0)] += temp*B[k*N+(j+0)];
								
			}  
		          
		}
		    
	 }
   }


/*  
  for (int i = 0; i < M; i++) 
  {
	for (int k = 0; k < N; k++) 
    	{     
		float temp = A[i*K + k];
		for (int j = start; j < end; j=j+8) 
	      	{
	  		C[i*N+(j+0)] += temp*B[k*N+(j+0)];
			C[i*N+(j+1)] += temp*B[k*N+(j+1)];
			C[i*N+(j+2)] += temp*B[k*N+(j+2)];
			C[i*N+(j+3)] += temp*B[k*N+(j+3)];
			C[i*N+(j+4)] += temp*B[k*N+(j+4)];
			C[i*N+(j+5)] += temp*B[k*N+(j+5)];
			C[i*N+(j+6)] += temp*B[k*N+(j+6)];
			C[i*N+(j+7)] += temp*B[k*N+(j+7)];
	
	
		}
	}
  }
*/		  
	
/*
for (int i = 0; i < M; ++i) {
      for (int j = 0; j < N; ++j) {
        for (int k = 0; k < K; ++k) {
          C[i * N + j] += A[i * K + k] * B[k * N + j];
        }
    }
  }
*/
  return NULL;
}

void mat_mul(float *_A, float *_B, float *_C, int _M, int _N, int _K, int _num_threads) {
  A = _A, B = _B, C = _C;
  M = _M, N = _N, K = _K;
  num_threads = _num_threads;

  // TODO: create '_num_threads' pthreads
  
  pthread_t * threads;
  threads = (pthread_t *) malloc(sizeof(pthread_t) * _num_threads);

 
 for (int i = 0; i < num_threads; i++) 
 {
    int * pid = (int *) malloc(sizeof(int));
	    
    *pid = i;
     pthread_create(&threads[i], NULL, mat_mul_thread, pid);
 }

  for (int i = 0; i < _num_threads; i++) 
  {
  	pthread_join(threads[i], NULL);
  }
  
  /*
  pthread_t thread;
  pthread_create(&thread, NULL, mat_mul_thread, NULL);
  pthread_join(thread, NULL);
  */
}