#include <stdio.h>
#include <string.h>
#include <mpi.h>
#include "bin_reduce.h"

int int_log2(int x) {
	int r=0;
	while (x >>= 1) {
		r++;
	}
	return r;
}

int Bin_Reduce(const void *sendbuf, void *recvbuf, int count,
		MPI_Datatype datatype, MPI_Op op, int root, MPI_Comm comm) {
	if (root != 0) {
		fprintf(stderr, "Sorry, root!=0 not allowed");
		return -1;
	}

	int r, p, size;
	MPI_Status status;

	MPI_Comm_rank(comm, &r);
	MPI_Comm_size(comm, &p);
	MPI_Type_size(datatype, &size);

	int tree_depth = int_log2(p) + 1;
	int i = 0;
	int depth;
	int parent = 0;

	// maximum possible number of nodes in a subtree with current depth
	int max_nodes = ((1 << tree_depth) - 1) / 2;

	void *recv_left;
	void *recv_right;

	if (r == root) {
		recv_right = recvbuf;
	} else {
		MPI_Alloc_mem(count * size, MPI_INFO_NULL, &recv_right);
	}

	for (depth = 1; i != r; depth++) {
		parent = i;
		if (r > i + max_nodes) {
			i += max_nodes + 1;
		} else {
			i++;
		}
		max_nodes /= 2;
	}

	if (depth != tree_depth && r + 1 < p) {
		MPI_Alloc_mem(count * size, MPI_INFO_NULL, &recv_left);
		MPI_Recv(recv_left, count, datatype, r + 1, 0, comm, &status);
		MPI_Reduce_local(sendbuf, recv_left, count, datatype, op);

		if (r + max_nodes + 1 < p) {
			MPI_Recv(recv_right, count, datatype, r + max_nodes + 1, 0, comm,
					&status);
			MPI_Reduce_local(recv_left, recv_right, count, datatype, op);
		} else {
			memcpy(recv_right, recv_left, count * size);
		}

		MPI_Free_mem(recv_left);
	} else {
		memcpy(recv_right, sendbuf, count * size);
	}

	if (r != root) {
		MPI_Send(recv_right, count, datatype, parent, 0, comm);
		MPI_Free_mem(recv_right);
	}

	return 0;
}