#include "convolution.h"
#include <mpi.h>
#include <stdio.h>
#include "util.h"

static float *input, *output, *filter;
static int N, C, H, W;
static int K, R, S;
static int OH, OW;
static int pad;
static int dilation;
static int stride;
static int mpi_rank, mpi_world_size;

void convolution_core(
		float *_input, float *_output, float *_filter, 
		int _N, int _C, int _H, int _W,
		int _K, int _R, int _S,
		int _pad, int _dilation, int _stride) {
	//input = _input;
	//output = _output;
	//filter = _filter;

	//OH = (H + 2 * pad - dilation * (R - 1) - 1) / stride + 1;
	//OW = (W + 2 * pad - dilation * (S - 1) - 1) / stride + 1;
	//
	//printf("convolution core N : %d\n", _N);

//#pragma omp parallel for  //num_threads(num_threads)
//#pragma omp parallel for collapse(2) // 72
//#pragma omp parallel for collapse(3) // 74
#pragma omp parallel for collapse(4) // 91
//#pragma omp parallel for collapse(5) // 
	for (int n = 0; n < _N; ++n) {
		for (int k = 0; k < K; ++k) {
			for (int oh = 0; oh < OH; ++oh) {
				for (int ow = 0; ow < OW; ++ow) {
					
					float o = 0.0f;
					for (int c = 0; c < C; ++c) {
						for (int r = 0; r < R; ++r) {
							for (int s = 0; s < S; ++s) {
								int h = oh * stride - pad + r * dilation;
								int w = ow * stride - pad + s * dilation;
								if (h < 0 || h >= H || w < 0 || w >= W) continue;
								float i = input[n * C * H * W + c * H * W + h * W + w];
								float f = filter[k * C * R * S + c * R * S + r * S + s];
								o += i * f;
							}
						}
					}
					output[n * K * OH * OW + k * OH * OW + oh * OW + ow] = o;

				}
			}
		}
	}

}

void convolution(
		float *_input, float *_output, float *_filter, 
		int _N, int _C, int _H, int _W,
		int _K, int _R, int _S,
		int _pad, int _dilation, int _stride) {


	N = _N;
	C = _C;
	H = _H;
	W = _W;
	K = _K;
	R = _R;
	S = _S;
	pad = _pad;
	dilation = _dilation;
	stride = _stride;

	OH = (H + 2 * pad - dilation * (R - 1) - 1) / stride + 1;
	OW = (W + 2 * pad - dilation * (S - 1) - 1) / stride + 1;

	MPI_Request mpi_request;
	MPI_Status mpi_status;

	int input_size = N * C * H * W;
	int input_middle = (N/2) * C * H * W;

	int filter_size = K * C * R * S;

	int output_size = N * K * OH * OW;
	int output_middle = (N/2) * K * OH * OW;

	int N_size = N;

	//printf("mpi world size : %d\n", mpi_world_size);

	if(mpi_rank == 0){
		input = _input;
		output = _output;
		filter = _filter;

		if(mpi_world_size == 2){
			MPI_Isend(input + input_middle, input_size - input_middle, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &mpi_request);
			MPI_Isend(filter, filter_size, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &mpi_request);
			N_size = N/2;
			//printf("0 N_size : %d\n", N_size);
		}
	}
	else{
		input = (float *) aligned_alloc(32, sizeof(float) * (input_size - input_middle));
		filter = (float *) aligned_alloc(32, sizeof(float) * filter_size);
		output = (float *) aligned_alloc(32, sizeof(float) * (output_size - output_middle));

		MPI_Recv(input, input_size - input_middle, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &mpi_status);
		MPI_Recv(filter, filter_size, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &mpi_status);

		N_size = N - N/2;
		//printf("1 N_size : %d\n", N_size);
	}

	//printf("call core : %d\n", N_size);
	convolution_core( input, output, filter, 
			N_size, C, H, W,
			K, R, S,
			pad, dilation, stride);
	//printf("core finished\n");

	if(mpi_world_size == 2){
		if(mpi_rank == 0){
			//printf("rank 0 recv : %d\n", output_size - output_middle);
			MPI_Recv(output + output_middle, output_size - output_middle, MPI_FLOAT, 1, 1, MPI_COMM_WORLD, &mpi_status);
		}
		else{
			//printf("rank 1 send : %d\n", output_size - output_middle);
			MPI_Isend(output, output_size - output_middle, MPI_FLOAT, 0, 1, MPI_COMM_WORLD, &mpi_request);

		}
	}

	// input_size = N * C * H * W;
	// filter_size = K * C * R * S;
	// output_size = N * K * OH * OW;

}


void convolution_init(
		int _N, int _C, int _H, int _W,
		int _K, int _R, int _S,
		int _pad, int _dilation, int _stride) {
	N = _N; C = _C; H = _H; W = _W;
	K = _K; R = _R; S = _S;
	pad = _pad;
	dilation = _dilation;
	stride = _stride;

	MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
	MPI_Comm_size(MPI_COMM_WORLD, &mpi_world_size);
}

void convolution_final(
		int _N, int _C, int _H, int _W,
		int _K, int _R, int _S,
		int _pad, int _dilation, int _stride) {


}

//printf("mpi rank 0\n");
//printf("0: N : %d\n", _N);
//printf("0: C : %d\n", _C);
//printf("0: H : %d\n", _H);
//printf("0: W : %d\n", _W);
//printf("0: pad : %d\n", _pad);
//printf("0: dilation : %d\n", _dilation);
//printf("0: stride : %d\n", _stride);
//printf("0: input : %f / %p\n", input[0], input);
//printf("0: input middle : %f / %p\n", input[input_middle], input);
//printf("0: output : %f / %p\n", output[0], output);
//printf("0: filter : %f / %p\n", filter[0], filter);

//printf("mpi rank 1\n");
//printf("1: N : %d\n", N);
//printf("1: C : %d\n", C);
//printf("1: H : %d\n", H);
//printf("1: W : %d\n", W);
//printf("1: pad : %d\n", pad);
//printf("1: dilation : %d\n", dilation);
//printf("1: stride : %d\n", stride);
//printf("1: input : %f %p\n", input[0], input);
//printf("1: output : %f %p\n", output[0], output);
//printf("1: filter : %f %p\n", filter[0], filter);