#include "mat_mul.h"

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

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

static void* mat_mul_thread(void *data) {
  // TODO: parallelize & optimize matrix multiplication
int me = *(int *)data;
//int bs = 4096;
int bs = 2048;
//int inum = M_div+M_remain;

if(me==(num_threads-1))
  M_remain = M % num_threads;
else
  M_remain = 0;

int inum = M_div+M_remain;

  for (int kk = 0; kk < K; kk=std::min(kk+K_div,K)) {
  	for (int i = 0; i < inum; ++i) {
       	    for (int jj = 0;  jj < N; jj+=bs) {
       		for (int k = kk; k < std::min(kk+K_div,K); ++k) {
        //for (int k = kk; k < std::min(kk+K_div,K); k=k+2) {
      			for (int j = jj; j < std::min(jj+bs,N); ++j) {
        			  C[(i + me*M_div) * N + j ] += A[(i+me*M_div) * K + k ] * B[k * N + j];
         // C[(i + me*M_div) * N + j ] += A[(i+me*M_div) * K + k+1 ] * B[(k+1) * 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;
  
  M_div = M / num_threads;
  K_div = 32; //or 64;   
  M_remain = M % num_threads;

  // TODO: create '_num_threads' pthreads
  pthread_t *thread;
  thread = (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(&thread[i], NULL, mat_mul_thread, pid);
  }
  for(int i=0; i< num_threads; i++)
	pthread_join(thread[i], NULL);
}