#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 void* mat_mul_thread(void *data) {
  // TODO: parallelize & optimize matrix multiplication
  int pid = *(int *)data;
  int slice = M/num_threads;
  int start = pid*slice;
  int end = pid == num_threads -1 ? M: (pid+1)*slice;
  float Aik;
  int bs=32;
  if(start>=M) return NULL;
  for(int kprim=0;kprim<K;kprim+=bs) {
	//for(int jj=0;jj<N;jj+=bs) {
	  for(int i=start;i<end;i++) {
		  for(int kk=kprim;kk<std::min(kprim+bs,K);kk++) {
			  Aik=A[i*K+kk];
			  for(int j=0;j<std::min(j+bs,N);j++) C[i*N+j]+=Aik*B[kk*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);
  	//pthread_create(&thread, NULL, mat_mul_thread, NULL);
  }

  for(int j=0;j<num_threads;j++) pthread_join(threads[j], NULL);
}