#include <iostream>
#include <sstream>
#include <cstdlib> // needed for atoi and atof
#include "ppm.h"
using namespace std;
using namespace math144_ppm;

/*
 * Runs the partial differential equation corresponding to pde_num on "img".
 * For the steady state level set equation, alpha and tolerance are used.
 * For the rest of the equations, "iter" iterations is used, and "tot_iter"
 * is used for saving purposes.
 * Returns  0 on success.
 * Returns -1 on failure.
 */
int run_pde(image &img, string filename, unsigned int pde_num, double delta_t,
       unsigned int iter, unsigned int tot_iter, double alpha, double tolerance)
{
	/*
	 * Note: the steady state level set equation uses alpha and tolerance.
	 * The other equations use iter and tot_iter.
	 */

	/* Starting building the filename to save as */
	stringstream savename;
	savename << filename.substr(0,filename.find(".ppm")) << "_";

	/* Call the PDE */
	switch (pde_num) {
		case 1:
			savename << "heat";
			cout << "Running the heat equation"
			     << " on " << filename << " " << iter
			     << " time(s) with delta-t = " << delta_t << endl;
			heat(img, delta_t, iter);
			break;
		case 2:
			savename << "lvlset";
			cout << "Running the level set equation"
			     << " on " << filename << " " << iter
			     << " time(s) with delta-t = " << delta_t << endl;
			lvlset(img, delta_t, iter);
			break;
		case 3:
			savename << "stdylvlset";
			cout << "Running the steady state level set equation"
			     << " on " << filename
			     << " with delta-t = " << delta_t
			     << " and alpha = " << alpha
			     << " and tolerance = " << tolerance << endl;
			stdylvlset(img, delta_t, alpha, tolerance);
			break;
		case 4:
			savename << "shock";
			cout << "Running the shock filter equation"
			     << " on " << filename << " " << iter
			     << " time(s) with delta-t = " << delta_t << endl;
			shock(img, delta_t, iter);
			break;
		default:
			cerr << "Error selecting pde_name to apply." << endl;
			return -1;
	}

	/* Save the image */
	if (pde_num == 3) {
		/* The steady state level set equation doesn't have iterations */
		savename << "_" << delta_t << "_" << alpha << "_" << tolerance << ".ppm";
	} else {
		savename << "_" << delta_t << "_" << tot_iter << ".ppm";
	}
	img.save(savename.str());
	cout << "Image saved as \"" << savename.str() << "\"." << endl << endl;

	return 0;
}

int main(int argc, char *argv[])
{
	/* Print the header */
	cout << "Math 144 - Differential Equations" << endl
	     << "Image Processing Project" << endl << endl;

	/*
	 * Check whether we're using the menu selection interface, or command
	 * line delta_teters.
	 */
	string filename;
	double delta_t;
	unsigned int pde_num;
	if (argc == 1) {
		/* Use menu selection interface */

		/* Pick the source image */
		int img_num;
		cout << "1. beatles.ppm" << endl
		     << "2. tri.ppm" << endl
		     << "3. trilgt.ppm" << endl
		     << "4. trimed.ppm" << endl
		     << "5. trihvy.ppm" << endl
		     << "6. dobbs.ppm" << endl
		     << "7. love.ppm" << endl
		     << "8. lovens.ppm" << endl
		     << "9. Other" << endl;
		do {
		     	cout << "Please Select an Image to Modify (1-9): ";
			cin >> img_num;
		} while (img_num < 1 || img_num > 9);

		switch (img_num) {
			case 1:
				filename = "beatles.ppm";
				break;
			case 2:
				filename = "tri.ppm";
				break;
			case 3:
				filename = "trilgt.ppm";
				break;
			case 4:
				filename = "trimed.ppm";
				break;
			case 5:
				filename = "trihvy.ppm";
				break;
			case 6:
				filename = "dobbs.ppm";
				break;
			case 7:
				filename = "love.ppm";
				break;
			case 8:
				filename = "lovens.ppm";
				break;
			case 9:
				cout << "Please enter the filename: ";
				cin >> filename;
				break;
			default:
				cerr << "Error selecting image." << endl;
				return -1;
		}
		cout << endl;

		/* Pick the PDE */
		cout << "1. Heat Equation" << endl
		     << "2. Level Set Equation" << endl
		     << "3. Steady State Level Set Equation" << endl
		     << "4. Shock Filter Equation" << endl;

		do {
		     	cout << "Please Select an Equation to Apply (1-4): ";
			cin >> pde_num;
		} while (pde_num < 1 || pde_num > 4);

		/* Get the value of delta-t */
		cout << endl << "Please enter a value for delta-t: ";
		cin >> delta_t;
		cout << endl;
	} else if (argc == 4) {
		/* Use the command line arguments */
		filename = argv[1];
		pde_num = atoi(argv[2]);
		delta_t = atof(argv[3]);
	} else {
		/* Incorrect command line argument usage */
		cout << "Usage: ./imagepro [FILENAME PDE delta_t]" << endl
		     << "FILENAME is the name of the file to use" << endl
		     << "PDE is an integer 1-4 corresponding to" << endl
		     << "  1. Heat Equation" << endl
		     << "  2. Level Set Equation" << endl
		     << "  3. Steady State Level Set Equation" << endl
		     << "  4. Shock Filter Equation" << endl
		     << "delta_t is delta-t" << endl
		     << "Use no command-line arguments for menu selection"
		     << " (Recommended for new users)." << endl;
		return -1;
	}

	/* Open the image */
	image img(filename);

	/* Keep running the PDE for the specified number of iterations */
	unsigned int tot_iter = 0;
	if (pde_num == 3) {
		/*
		 * We're running the steady state level set equation.  It
		 * doesn't make sense to have a specified number of
		 * iterations.  Get alpha and tolerance instead.
		 */
		double alpha, tolerance;
		cout << "Please enter the value for alpha: ";
		cin >> alpha;
		cout << endl;
		cout << "Please enter the tolerance: ";
		cin >> tolerance;
		cout << endl;

		/* Run the PDE */
		if (run_pde(img, filename, pde_num, delta_t, 0, tot_iter, alpha, tolerance) == -1) {
			return -1;
		}
		
		/* Uncomment this to save a comma delimited printout */
		// img.save_csv();
	} else {
		while (1) {
			/* Get the number of additional iterations */
			cout << "Please enter the number of additonal iterations ("
			     << tot_iter << " so far, 0 to exit): ";
			unsigned int iter;
			cin >> iter;
			if (iter <= 0) {
				break;
			}

			/* Calculate the total number of iterations, then run the PDE */
			tot_iter += iter;
			if (run_pde(img, filename, pde_num, delta_t, iter, tot_iter, 0, 0) == -1) {
				return -1;
			}

			/* Uncomment this to save a comma delimited printout */
			// img.save_csv();
		}
	}

	return 0;
}