/*
    This file is part of Mitsuba, a physically based rendering system.

    Copyright (c) 2007-2012 by Wenzel Jakob and others.

    Mitsuba is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License Version 3
    as published by the Free Software Foundation.

    Mitsuba is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#include <mitsuba/render/util.h>
#include <mitsuba/core/timer.h>
#include <mitsuba/core/fresolver.h>
#include <mitsuba/core/plugin.h>
#include <boost/algorithm/string.hpp>
#if defined(WIN32)
#include <mitsuba/core/getopt.h>
#endif

MTS_NAMESPACE_BEGIN

class KDBench : public Utility {
public:
	void help() {
		cout << endl;
		cout << "Synopsis: kd-tree performance benchmark. Traces uniformly distributed rays" << endl;
		cout << "though the bounding sphere of a scene and reports the resulting number of" << endl;
		cout << "rays per second. The main intent of this utility is to optimize the kd-tree" << endl;
		cout << "construction parameters for particular scenes and machines." << endl;
		cout << endl;
		cout << "Usage: mtsutil kdbench [options] <Scene XML file or PLY file>" << endl;
		cout << "Options/Arguments:" << endl;
		cout << "   -h             Display this help text" << endl << endl;
		cout << "   -t value       Specify the SAH traversal cost" << endl << endl;
		cout << "   -i value       Specify the SAH intersection cost" << endl << endl;
		cout << "   -e value       Specify the SAH empty space bonus" << endl << endl;
		cout << "   -b value       Specify the number of min-max bins" << endl << endl;
		cout << "   -c true/false  Enable/disable primitive clipping (aka. \"perfect splits\")" << endl << endl;
		cout << "   -p true/false  Enable/disable parallel tree construction" << endl << endl;
		cout << "   -r true/false  Enable/disable retraction of bad splits" << endl << endl;
		cout << "   -l value       Specify the primitive count, below which a leaf node" << endl;
		cout << "                  will always be created" << endl << endl;
		cout << "   -d depth       Specify the maximum tree depth" << endl << endl;
	 	cout << "   -x value       Specify the number of primitives, at which the " << endl;
		cout << "                  builder will switch from (approximate) Min-Max " << endl;
		cout << "                  binning to the more accurate O(n log n) SAH-based " << endl;
		cout << "                  optimization method." << endl << endl;
		cout << "   -f             Try to empirically find the best SAH cost values by" << endl;
		cout << "                  fitting the cost model to collected performance data" << endl << endl;
		cout << "Examples:" << endl;
		cout << "  E.g. to build a tree for the Stanford bunny having a low SAH cost, type " << endl << endl;
		cout << "  $ mtsutil kdbench -e .9 -l1 -d48 -x100000 data/tests/bunny.ply" << endl << endl;
		cout << "  To get SAH costs comparable to [Wald and Havran 06], also specify -t15 -i20" << endl << endl;
		cout << "  The high -x paramer effectively disables Min-Max binning, which " << endl;
		cout << "  leads to a slower and more memory-intensive build, so don't try" << endl;
		cout << "  this on a huge model." << endl << endl;
	}

	int run(int argc, char **argv) {
		ref<FileResolver> fileResolver = Thread::getThread()->getFileResolver();
		int optchar;
		char *end_ptr = NULL;
		Float intersectionCost = -1, traversalCost = -1, emptySpaceBonus = -1;
		int stopPrims = -1, maxDepth = -1, exactPrims = -1, minMaxBins = -1;
		bool clip = true, parallel = true, retract = true, fitParameters = false;
		optind = 1;

		/* Parse command-line arguments */
		while ((optchar = getopt(argc, argv, "i:t:e:c:p:r:l:x:b:d:hf")) != -1) {
			switch (optchar) {
				case 'h': {
						help();
						return 0;
					}
					break;
				case 'f':
					fitParameters = true;
					break;
				case 'i':
					intersectionCost = (Float) strtod(optarg, &end_ptr);
					if (*end_ptr != '\0')
						SLog(EError, "Could not parse the intersection cost!");
					break;
				case 't':
					traversalCost = (Float) strtod(optarg, &end_ptr);
					if (*end_ptr != '\0')
						SLog(EError, "Could not parse the traversal cost!");
					break;
				case 'e':
					emptySpaceBonus = (Float) strtod(optarg, &end_ptr);
					if (*end_ptr != '\0')
						SLog(EError, "Could not parse the empty space bonus!");
					break;
				case 'l':
					stopPrims = strtol(optarg, &end_ptr, 10);
					if (*end_ptr != '\0')
						SLog(EError, "Could not parse the stopping primitive count!");
					break;
				case 'd':
					maxDepth = strtol(optarg, &end_ptr, 10);
					if (*end_ptr != '\0')
						SLog(EError, "Could not parse the maximum depth!");
					break;
				case 'x':
					exactPrims = strtol(optarg, &end_ptr, 10);
					if (*end_ptr != '\0')
						SLog(EError, "Could not parse the -e parameter!");
					break;
				case 'b':
					minMaxBins = strtol(optarg, &end_ptr, 10);
					if (*end_ptr != '\0')
						SLog(EError, "Could not parse the min-max bins parameter!");
					break;
				case 'c':
					if (strcmp(optarg, "true") == 0)
						clip = true;
					else if (strcmp(optarg, "false") == 0)
						clip = false;
					else
						SLog(EError, "Could not parse the clipping parameter!");
					break;
				case 'p':
					if (strcmp(optarg, "true") == 0)
						parallel = true;
					else if (strcmp(optarg, "false") == 0)
						parallel = false;
					else
						SLog(EError, "Could not parse the parallel build parameter!");
					break;
				case 'r':
					if (strcmp(optarg, "true") == 0)
						retract = true;
					else if (strcmp(optarg, "false") == 0)
						retract = false;
					else
						SLog(EError, "Could not parse the retraction parameter!");
					break;
			};
		}

		if (optind == argc || optind+1 < argc) {
			help();
			return 0;
		}

		ref<Scene> scene;
		ref<ShapeKDTree> kdtree;

		std::string lowercase = boost::to_lower_copy(std::string(argv[optind]));
		if (boost::ends_with(lowercase, ".xml")) {
			fs::path
				filename = fileResolver->resolve(argv[optind]),
				filePath = fs::absolute(filename).parent_path(),
				baseName = filename.stem();
			ref<FileResolver> frClone = fileResolver->clone();
			frClone->prependPath(filePath);
			Thread::getThread()->setFileResolver(frClone);
			scene = loadScene(argv[optind]);
			kdtree = scene->getKDTree();
		} else if (boost::ends_with(lowercase, ".ply")) {
			Properties props("ply");
			props.setString("filename", argv[optind]);
			ref<TriMesh> mesh;
			mesh = static_cast<TriMesh *> (PluginManager::getInstance()->
					createObject(MTS_CLASS(TriMesh), props));
			mesh->configure();
			kdtree = new ShapeKDTree();
			kdtree->addShape(mesh);
		} else {
			Log(EError, "The supplied scene filename must end in either PLY or XML!");
		}

		if (intersectionCost != -1)
			kdtree->setQueryCost(intersectionCost);
		if (traversalCost != -1)
			kdtree->setTraversalCost(traversalCost);
		if (emptySpaceBonus != -1)
			kdtree->setEmptySpaceBonus(emptySpaceBonus);
		if (stopPrims != -1)
			kdtree->setStopPrims(stopPrims);
		if (maxDepth != -1)
			kdtree->setMaxDepth(maxDepth);
		if (exactPrims != -1)
			kdtree->setExactPrimitiveThreshold(exactPrims);
		if (minMaxBins != -1)
			kdtree->setMinMaxBins(minMaxBins);
		kdtree->setClip(clip);
		kdtree->setRetract(retract);
		kdtree->setParallelBuild(parallel);

		/* Show some statistics, and make sure it roughly fits in 80cols */
		Logger *logger = Thread::getThread()->getLogger();
		DefaultFormatter *formatter = ((DefaultFormatter *) logger->getFormatter());
		logger->setLogLevel(EDebug);
		formatter->setHaveDate(false);

		if (scene)
			scene->initialize();
		else
			kdtree->build();

		BSphere bsphere(kdtree->getAABB().getBSphere());
		const size_t nRays = 5000000;

		if (!fitParameters) {
			Log(EInfo, "Bounding sphere: %s", bsphere.toString().c_str());
			Float best = 0;
			for (int j=0; j<3; ++j) {
				ref<Random> random = new Random();
				ref<Timer> timer = new Timer();
				size_t nIntersections = 0;

				Log(EInfo, "Shooting " SIZE_T_FMT " rays (1 thread, incoherent) ..", nRays);

				for (size_t i=0; i<nRays; ++i) {
					Point2 sample1(random->nextFloat(), random->nextFloat()),
						sample2(random->nextFloat(), random->nextFloat());
					Point p1 = bsphere.center + Warp::squareToUniformSphere(sample1) * bsphere.radius;
					Point p2 = bsphere.center + Warp::squareToUniformSphere(sample2) * bsphere.radius;
					Ray r(p1, normalize(p2-p1), 0.0f);

					Intersection its;
					if (kdtree->rayIntersect(r, its))
						nIntersections++;
				}

				Log(EInfo, "Found " SIZE_T_FMT " intersections in %i ms",
					nIntersections, timer->getMilliseconds());
				Float mrays = nRays / (timer->getMilliseconds() * (Float) 1000);
				Log(EInfo, "-> %.3f MRays/s", mrays);
				Log(EInfo, "");
				best = std::max(best, mrays);
			}
			Log(EInfo, "Best of three: %.3f MRays/s", best);
		} else {
			Float intersectionCost, traversalCost;
			kdtree->findCosts(intersectionCost, traversalCost);
		}

		Thread::getThread()->getLogger()->setLogLevel(EInfo);
		return 0;
	}

	MTS_DECLARE_UTILITY()
};

MTS_EXPORT_UTILITY(KDBench, "kd-tree performance benchmark")
MTS_NAMESPACE_END