#include <stdio.h>
#include <string.h>
#include <mpi.h>
#include <float.h>
#include <getopt.h>
#include <paths.h>
#include <stdlib.h>

#include "global.h"
#include "pathname.h"
#include "horizon.h"
#include "datagrid.h"
#include "dac.h"
#include "rtimer.h"

const int CHUNK_SIZE = 50;
const int IDENTIFIER_SIZE = 20;

const int WORK_TAG = 1;
const int DIE_TAG = 2;

const int DO_PRINT = 0;

// Function Declarations
int minimum(int a, int b);
void printUsage();
struct GlobalInfo* setGlobal(char* input, char* output, dim_t numRows, dim_t numCols, float NODATA_value, int verbosity);

/* TODO:
	Make SMP work!!!!
	Make check path work
	Make the code way less ugly
	Make some logging for processor efficiency, make scripts to run, check idle time
*/

int main(int argc, char *argv[]) {

	MPI_Init(&argc, &argv); // Initialize MPI World

	int my_rank; // Rank of specific processor
	MPI_Comm_rank(MPI_COMM_WORLD, &my_rank); // Get current processor rank

	int comm_sz; // Number of processors
	MPI_Comm_size(MPI_COMM_WORLD, &comm_sz); // Get number of processors

	// Get the name of the processor
	char processor_name[MPI_MAX_PROCESSOR_NAME];
	int namelen;
	MPI_Get_processor_name(processor_name, &namelen);

	// Put the name of the processor into the identifier that will be printed
	char myIdentifier[IDENTIFIER_SIZE];
	sprintf(myIdentifier, "[%d on %s]", my_rank, processor_name);

	// Save the current time
	double startTime = 0, endTime = 0;
	if (my_rank == 0)
		startTime = MPI_Wtime();

	// Get command line options
	const struct option options[] = {
		{"elev", 1, NULL, 'e'},
		{"verbose", 0, NULL, 'v'},
		/* sentinel */
		{0, 0, 0, 0}
	};

	float elev = 0;
	int verbosity = DEFAULT;
	char c;
	char *error;
	
	while ((c = getopt_long(argc, argv, "e:v", options, NULL)) >= 0) {
		
		switch (c) {
			case 'e':
				elev = strtod(optarg, &error);
				break;
			case 'v':
				verbosity = HIGH;
				break;
			case '?': // Bad option
				printUsage();
				return -1;
			default:
				assert(NULL);
		}

		if (optarg && error == optarg) {
			if (my_rank == 0) {
				printf("Error processing a numerical input\n");
				printUsage();
			}
			return -1;
		}
	}

	assert(c == -1);
	
	// Check arguments
	if (argc - optind != 2) {
		if (my_rank == 0) {
			printf("Invalid number of arguments (expected 2, got %i)\n", argc - optind);
			printUsage();
		}
		return -1;
	}

	char *input  = argv[optind];
	char *output = argv[optind+1]; 
	
	// check input file 
	if (check_path(input)) {
		printf("Test");
		if (my_rank == 0)
			printf("Path: %s error\n", input);
		return -1;
	}     

	// If it's the master, print out our arguments
	if (my_rank == 0) {
		printf("\nRunning with arguments:\n");
		switch(verbosity) {
			case HIGH: 
				printf("Verbosity set to HIGH\n");
			break; 
			case DEFAULT: 
				printf("Verbosity set to DEFAULT\n"); 
			break; 
			case SUPPRESS: 
			 	printf("Verbosity set to SUPPRESS\n");
			break;
			default: 
				printf("Verbosity not set at all\n"); 
		}
		printf("Elevation above ground: %f\n", elev);

		printf("Input file: %s \nOutput file: %s\n", input, output);
	}

	// First, load the grid
	DataSet *elevDataSet = dg_load_grid(input);
	Grid elevGrid = elevDataSet->grid; // Check if i use this

	int numRows = elevGrid.hd.nrow;
	int numCols = elevGrid.hd.ncol;
	float NODATA_value = elevGrid.hd.NODATA_value;

	// Variables for MPI sending and receiving
	int result[CHUNK_SIZE];
	MPI_Status status;

	// Master Process
	if (my_rank == 0) {

		// Make a new visibility grid, intialized to zero
		DataSet *visDataSet = dg_init(elevGrid.hd);
		dg_set(elevDataSet, visDataSet, 0);

		// Create an array of indexes for master to hold start/end for each slave process
		int *index = (int *) malloc(sizeof(int) * 2 * comm_sz);
		assert(index);

		// Loop variables
		int startIndex = 0;
		int endIndex;

		int resultStart, resultEnd;
		int r, c;

		int task[2];

		// Count how many chunks are completed
		int numChunks = (numRows * numCols)/CHUNK_SIZE;
		int chunk = 0;

		// Seed the slaves
		for (int i = 1; i < comm_sz; i++) {

			endIndex = startIndex + CHUNK_SIZE; // End index not inclusive as far as slave knows
			
			// Check that it's not over the limit
			if (endIndex > numRows * numCols)
				endIndex = numRows * numCols;

			index[2*i] = task[0] = startIndex;
			index[2*i + 1] = task[1] = endIndex;

			if (verbosity > DEFAULT)
				printf("%s Telling slave process %d to count from %d up to and not including %d.\n", 
						myIdentifier, i, startIndex, endIndex);


			MPI_Send(&task, 2, MPI_INT, i, 0, MPI_COMM_WORLD);

			startIndex = endIndex;

			chunk++;
			printf("%d/%d chunks completed.\n", chunk, numChunks);
		}

		// Now, wait for tasks to finish their task, and send them a new one
		while (startIndex < numRows * numCols) {

			// Wait for processors to return that they've completed a task.
			MPI_Recv(result, CHUNK_SIZE, MPI_INT, MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);

			// Process results
			resultStart = index[2*status.MPI_SOURCE];
			resultEnd = index[2*status.MPI_SOURCE + 1];

			for (int i = 0; resultStart + i < resultEnd; i++) {
				r = (i + resultStart) / numCols;
				c = (i + resultStart) % numCols;

				if (verbosity > DEFAULT && result[i] != NODATA_value)
					printf("%s Just learned from %d that point (%d,%d) has viewshed %d.\n", 
							myIdentifier, status.MPI_SOURCE, r, c, result[i]);
				
				datagrid_set(visDataSet->grid, r, c, result[i]);
			}

			endIndex = startIndex + CHUNK_SIZE; // End index not inclusive as far as slave knows
			
			// Check that it's not over the limit
			if (endIndex > numRows * numCols)
				endIndex = numRows * numCols;

			index[2*status.MPI_SOURCE] = task[0] = startIndex;
			index[2*status.MPI_SOURCE + 1] = task[1] = endIndex;

			if (verbosity > DEFAULT)
				printf("%s Telling slave process %d to count from %d up to and not including %d.\n", 
						myIdentifier, status.MPI_SOURCE, startIndex, endIndex);


			MPI_Send(task, 2, MPI_INT, status.MPI_SOURCE, 0, MPI_COMM_WORLD);

			startIndex = endIndex; 

			chunk++;
			printf("%d/%d chunks completed.\n", chunk, numChunks);

		}

		for (int i = 1; i < comm_sz; i++) {

			// Wait for processors to return that they've completed their last task.
			MPI_Recv(result, CHUNK_SIZE, MPI_INT, MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);

			// Process results
			resultStart = index[2*status.MPI_SOURCE];
			resultEnd = index[2*status.MPI_SOURCE + 1];

			for (int i = 0; resultStart + i < resultEnd; i++) {
				r = (i + resultStart) / numCols;
				c = (i + resultStart) % numCols;

				if (verbosity > DEFAULT && result[i] != NODATA_value)
					printf("%s Just learned from %d that point (%d,%d) has viewshed %d.\n", 
							myIdentifier, status.MPI_SOURCE, r, c, result[i]);
				
				datagrid_set(visDataSet->grid, r, c, result[i]);
			}

			if (verbosity > DEFAULT)
				printf("%s Process %d was just killed.\n", myIdentifier, status.MPI_SOURCE);

			// Send final task to die
			MPI_Send(0, 0, MPI_INT, status.MPI_SOURCE, DIE_TAG, MPI_COMM_WORLD);

			chunk++;
			printf("%d/%d chunks completed.\n", chunk, numChunks);
		}

		// Write to file

		printf("%s\n", output);
		printf("\n");
		int error = dg_store_grid(visDataSet, output);
		if (error) {
			printf("Error saving grid to file %s\n", output);
			exit(-1);
		}
	}
	else {
		
		int index[2];

		int start, end;

		while (1) {

			// Receive a command from the master, and save those variables.
			MPI_Recv(&index, 2, MPI_INT, 0, MPI_ANY_TAG, MPI_COMM_WORLD, &status);

			start = index[0];
			end = index[1];

			// If the tag is DIE_TAG, we want to end
			if (status.MPI_TAG == DIE_TAG)
				break;

			if (verbosity > DEFAULT)
				printf("%s Just told to calculate from %d up to and not including %d.\n", myIdentifier, start, end);

			for (int i = 0; i < CHUNK_SIZE; i++) {
				result[i] = 0;
			}

			int i, index, r, c, count;
			DataSet *tempDataSet;

			#pragma omp parallel shared(result, elevDataSet, elevGrid, elev) private(index, r, c, tempDataSet, count)
			{	

				tempDataSet = dg_init(elevGrid.hd);

				#pragma omp for
				for (i = 0; i < (end - start); i++) {

				index = i + start; // Index of point to compute in viewshed
				r = index / numCols;
				c = index % numCols;

				dg_set(elevDataSet, tempDataSet, 1);

				ViewPoint vp;
				vp.r = r;
				vp.c = c;
				vp.vp_elev = datagrid_get(elevGrid, (index_t) r, (index_t) c);
				vp.vp_aboveground = elev;

				count = NODATA_value;
				
				if (vp.vp_elev != NODATA_value) {
					count = inmem_horizon_iter_viewshed(elevDataSet, tempDataSet, vp, SUPPRESS);
				}

				result[i] = (int) count;

				if (verbosity > DEFAULT && count != NODATA_value)
	                    printf("%s Telling master that point (%d,%d) has viewshed %d.\n",
	                                    myIdentifier, r, c, result[i]);
				}
			}
			
			MPI_Send(&result, CHUNK_SIZE, MPI_INT, 0, 0, MPI_COMM_WORLD);

			if (verbosity > DEFAULT)
				printf("%s Telling master I finished calculating, return calculations. \n", myIdentifier);

			// Wait for the next task
		}
	}

	if (my_rank == 0) {
		endTime = MPI_Wtime();
		printf("\nRuntime: %f seconds.\n\n", endTime-startTime);
	}

	MPI_Finalize();
	return 0;
}

int minimum(int a, int b) { 
	if (a > b) {
		return b;
	}
	else {
		return a;	
	}
}

void printUsage() {
	const char usage[] =
	"Usage: %s [options] ELEV.asc VMAP.asc\n"
	"\n"
	"Options:\n"
	"  -e, --elev       Optional float.  Viewpoint elevation layer above the\n"
	"                   terrain.  Defaults to 0.\n"
	"  -v, --verbose    Optional. Sets verbosity higher than the default.\n";

	printf("%s", usage);
}