#include <iostream>
#include <cstdlib>
#include <cmath>
#include "Ttiming.h"
#include <mpi.h>
#include <stdio.h>

using namespace std;

void array_alloc(long **&arr, long arrsize)
{
	try
	{	
		arr = new long* [arrsize];	
		arr[0] = new long [arrsize*arrsize];
	
		for(long i = 1; i < arrsize; ++i)
			arr[i] = arr[i-1] + arrsize;
	}
	catch (bad_alloc& ex)
	{
		cerr << "Could not allocate memory for array" << endl;
		exit(1);
	}
}

void vect_alloc(long *&arr, long arrsize)
{
	try
	{	
		arr = new long [arrsize];	
	}
	catch (bad_alloc& ex)
	{
		cerr << "Could not allocate memory for vector" << endl;
		exit(1);
	}
}

int main(int argc, char *argv[])
{
	long **A=NULL, **B=NULL, *B_rot=NULL, *vect=NULL, *vect_c=NULL, *C_rot=NULL;
	long rozmiar=0, porcja_new=0, buff_max=0;
	char *endptr;
	TTiming tt;
	long i;
	MPI::Status  status;

	if (argc < 2)
	{
		cerr << "Usage: " << argv[0] << " <size>" << endl;
		exit(1);
	}
	
	rozmiar = strtol(argv[1], &endptr, 10);
	
	if (*endptr)
	{
		cerr << "Invalid array size format" << endl;
		exit(1);
	}	  

	if (rozmiar <= 0)
	{
		cerr << "The number of matrix dimension must be positive" << endl;
		exit(1);
	}
	  
	MPI::Init(argc, argv);

	int taskid = MPI::COMM_WORLD.Get_rank();
	int ntasks = MPI::COMM_WORLD.Get_size();

	//alokacja macierzy
	array_alloc(A, rozmiar);
	if (taskid == 0) 
	{
		array_alloc(B, rozmiar);
		vect_alloc(B_rot, rozmiar*rozmiar);

	}

	// porcja - ile wektorów dostaje jeden proces do liczenia
	// jeżeli liczba procesów większa od wektorów - każdy proces dostaje po jednym, niektóre wcale
	long porcja = (rozmiar>=ntasks) ? (long)round(rozmiar*1.0/ntasks) : 1;
	porcja *= rozmiar; //porcję mnożymy przez ilość elementów w jednym wektorze
	porcja_new = porcja;

	// dla kompatybilności liczę jeszcze rozmiar bufora odbioru,
	// np. jeżeli chunki będą jednolitego rozmiaru, a ostatni będzie większy
	// to bufor odbioru musi być zwiększony (żeby go nie ucinało)
	buff_max = (rozmiar - (porcja/rozmiar)*(ntasks-1))*rozmiar;
	if (porcja > buff_max)
		buff_max = porcja; // na wypadek, jeżeli jednak ostatni kawałek będzie mniejszy 
	
	if (taskid == 0)
		vect_alloc(C_rot, ntasks*buff_max);

	vect_alloc(vect, porcja);
	vect_alloc(vect_c, buff_max);
	if (taskid == 0)
	{
		// wypełnienie macierzy A liczbami "losowymi"
		for (long i=0; i<rozmiar; ++i)
			for (long j=0; j<rozmiar; ++j)
				//A[i][j] = 3*i+j+1; //test
				A[i][j] = (long)(sin(i) * i * j) % 10;
			
		// wypełnienie macierzy B liczbami "losowymi"
		for (long i=0; i<rozmiar; ++i)
			for (long j=0; j<rozmiar; ++j)
				//B[i][j] = j; //test
				B[i][j] = (long)(cos(j) *(i+j)) % 10;
				
		// dokonaj obracania macierzy
		for (long i=0; i<rozmiar; ++i)
			for (long j=0; j<rozmiar; ++j)
				B_rot[i*rozmiar+j] = B[j][i];
	}
	
	//wysyłanie macierzy A do wszystkich
	MPI::COMM_WORLD.Bcast(&A[0][0], rozmiar*rozmiar, MPI::LONG, 0);

	MPI::COMM_WORLD.Barrier();

	//dla wszystkich lub "rozmiar" procesów należy zadać porcję do obliczenia
	MPI::COMM_WORLD.Scatter(&B_rot[0], porcja, MPI::LONG, &vect[0], porcja, MPI::LONG, 0);
	
	//
	// część licząca ---------------------------------------------------------------------------------
	//
	
	MPI::COMM_WORLD.Barrier();
	if (taskid == 0)
		tt.Begin();

	if (taskid == ntasks -1) //przelicz porcję dla ostatniego procesu (w przypadku dzielenia z resztą)
	{
		porcja_new = (rozmiar - (porcja/rozmiar)*(ntasks-1))*rozmiar;
		
		// myk jest taki, że muszę podać jakąś niezerową porcję do gathera,
		// mimo, że proces nie ma żadnych (sensownych) danych do przetworzenia
		if (porcja_new > 0)
			porcja = porcja_new;
	}

	for (i=0; i<porcja_new; ++i)
	{
		vect_c[i] = 0;
		long a_row = i%rozmiar;
		long vect_col = (i/rozmiar)*rozmiar;
		for (long k=0; k<rozmiar; ++k)
		{
			vect_c[i] += A[a_row][k]*vect[vect_col];
			++vect_col;
		}
	}
	
	MPI::COMM_WORLD.Barrier();
	if (taskid == 0)
	{
		long elapsed = tt.End();	
		cout << "Time: " << elapsed << " ms" << endl;
	}
	
	// zwróć wszystko do programu głównego
	MPI::COMM_WORLD.Gather(&vect_c[0], buff_max, MPI::LONG, &C_rot[0], buff_max, MPI::LONG, 0);
	
	MPI::Finalize();
	
	exit(0);
}