#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <mpi.h>

#define MASTER 0

#define DIVIDE 111
#define UNIFY 222

int total_procs, my_id;

typedef struct {
	int cols;
	int rows;
	float *data;
} MATRIX;

MATRIX create_matrix( int rows, int cols ) {
	MATRIX m;

	m.rows = rows;
	m.cols = cols;

	m.data = ( float * ) malloc( sizeof( float ) * rows * cols );

	return m;
}

void show_matrix( MATRIX m ) {
	int i, j;
	
	for( i = 0; i < m.rows; i++ ) {
		for( j = 0; j < m.cols; j++ )
			printf( "%f ", m.data[i * m.cols + j] );
		printf( "\n" );
	}
}

void write_matrix( MATRIX m ) {
	int i, j;
	FILE *f = fopen( "OUTPUT", "w" );
	
	for( i = 0; i < m.rows; i++ ) {
		for( j = 0; j < m.cols; j++ )
			fprintf( f, "%f ", m.data[i * m.cols + j] );
		fprintf( f, "\n" );
	}
}

void read_matrix( char *path, MATRIX *mat ) {
	int i, j;

	char aux_string[50];
	char rows_string[10], cols_string[10];

	FILE *input_file = fopen( path, "r" );

	fgets( aux_string, 50, input_file );
	printf( "Leyendo la matriz: %s", aux_string );

	fgets( aux_string, 50, input_file );
	sscanf( aux_string, "%s %s", rows_string, cols_string );

	printf( "Dimension: %sx%s\n", rows_string, cols_string );

	*mat = create_matrix( atoi( rows_string ), atoi( cols_string ) );

	for( i = 0; i < mat->rows; i++ )
		for( j = 0; j < mat->cols; j++ )
			fscanf( input_file, "%f", &mat->data[i * mat->cols + j] );
	
}

void multiply_matrix( MATRIX *dest, MATRIX a, MATRIX b ) {
	int i, j, k;

	*dest = create_matrix( a.rows, b.cols );

	for( i = 0; i< a.rows; i++ )	{
		for( j = 0; j < b.cols; j++ ) {
			dest->data[i * b.cols + j] = 0;
			for( k = 0; k < a.cols; k++ )
				dest->data[i * b.cols + j] += a.data[i * a.cols + k] * b.data[k * b.cols + j];
		}
	}
}

void broadcast_matrix( MATRIX *m ) {
	int rows, cols;

	rows = m->rows;
	cols = m->cols;

	MPI_Bcast( &rows, 1, MPI_INT, MASTER, MPI_COMM_WORLD );
	MPI_Bcast( &cols, 1, MPI_INT, MASTER, MPI_COMM_WORLD );
	
	if( my_id == MASTER )
		MPI_Bcast( m->data, rows * cols, MPI_FLOAT, MASTER, MPI_COMM_WORLD );
	else {
		*m = create_matrix( rows, cols );
		MPI_Bcast( m->data, rows * cols, MPI_FLOAT, MASTER, MPI_COMM_WORLD );
	}
}

void SOR_matrix_divide_data( MATRIX mat1 , MATRIX mat2, MATRIX *partial_matrix) {
	MATRIX aux_matrix;
	MPI_Status status;
	int mod, div;
	int i, size, current_pos;
	int cols, rows;

	cols = mat1.cols;
	rows = mat1.rows;

	MPI_Bcast( &cols, 1, MPI_INT, MASTER, MPI_COMM_WORLD );
	MPI_Bcast( &rows, 1, MPI_INT, MASTER, MPI_COMM_WORLD );
	
	div = rows / total_procs;
	mod = rows % total_procs;

	if( my_id == MASTER ) {
		current_pos = div;
		
		aux_matrix = create_matrix( div, mat1.cols );
		memcpy( aux_matrix.data, mat1.data, sizeof( float ) * div * aux_matrix.cols );
		
		for( i = 1; i < total_procs; i++ ) {
			size = div;
			if( i <= mod )
				size++;

			if( size > 0 )
				MPI_Send( mat1.data + ( current_pos * mat1.cols ), size * mat1.cols, MPI_FLOAT, i, DIVIDE, MPI_COMM_WORLD );
			current_pos += size;
		}
	}
	else {
		if( my_id <= mod )
			div++;
		aux_matrix = create_matrix( div, cols );
		if( size > 0 )
			MPI_Recv( aux_matrix.data, aux_matrix.rows * aux_matrix.cols, MPI_FLOAT, MASTER, DIVIDE, MPI_COMM_WORLD, &status );
	}
	
	multiply_matrix( partial_matrix, aux_matrix, mat2 );
}

void SOR_unify_partials( MATRIX *result, MATRIX partial, int rows, int cols ) {
	MATRIX aux_matrix;
	MPI_Status status;
	int mod, div;
	int i, size, current_pos;

	if( my_id == MASTER ) {
		*result = create_matrix( rows, cols );

		div = result->rows / total_procs;
		mod = result->rows % total_procs;

		memcpy( result->data, partial.data, sizeof( float ) * partial.rows * partial.cols );

		current_pos = div;
		for( i = 1; i < total_procs; i++ ) {
			size = div;
			if( i<=mod )
				size++;

			if( size > 0 )
				MPI_Recv( result->data + ( current_pos * result->cols ), size * result->cols, MPI_FLOAT, i, UNIFY, MPI_COMM_WORLD, &status );

			current_pos += size;
		}
	}
	else {
		if( size > 0 )
			MPI_Send( partial.data, partial.rows * partial.cols, MPI_FLOAT, MASTER, UNIFY, MPI_COMM_WORLD );
	}
}

void init( int argc, char *argv[] ) {
	MPI_Init( &argc, &argv );
	MPI_Comm_size( MPI_COMM_WORLD, &total_procs );
	MPI_Comm_rank( MPI_COMM_WORLD, &my_id );
}

void read_from_files( MATRIX *mat1, MATRIX *mat2, char *argv[]) {
	if( my_id == MASTER ) {
		read_matrix( argv[1], mat1 );
		show_matrix( *mat1 );
		printf("\n");
		read_matrix( argv[2], mat2 );
		show_matrix( *mat2 );
	}
}

void show_result( MATRIX m ) {
	if( my_id == MASTER ) {
		printf( "\n" );
		printf( "Resultado :\n" );
		show_matrix( m );
		write_matrix( m );
	}
}

int main( int argc, char *argv[] ) {
	MATRIX mat1, mat2, partial, res;

	init( argc, argv );
	if( argc != 3 ) {
		if( my_id == MASTER )
			printf( "Uso del programa: nombre_del_ejecutable 1ra_matriz 2da_matriz\n" );

		exit(0);
	}
	read_from_files( &mat1, &mat2, argv );
	broadcast_matrix( &mat2 );
	SOR_matrix_divide_data( mat1, mat2, &partial );
	SOR_unify_partials( &res, partial, mat1.rows, mat2.cols );
	show_result( res );

	MPI_Finalize();

	return 0;
}