/*
Copyright (c) 2013, Regents of the University of Alaska

All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
    * Neither the name of the Geographic Information Network of Alaska nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

This code was developed by Dan Stahlke for the Geographic Information Network of Alaska.
*/

#include <cassert>

#include <boost/lexical_cast.hpp>
#include <boost/numeric/conversion/cast.hpp>

#include "common.h"
#include "ndv.h"

using namespace dangdal;

struct Binning {
	Binning() :
		nbins(0),
		offset(0),
		scale(0)
	{ }

	int to_bin(double v) const {
		if(std::isinf(v) == -1) return 0;
		if(std::isinf(v) ==  1) return nbins-1;
		double bin_dbl = round((v-offset)/scale);
		if(std::isnan(bin_dbl)) fatal_error("nan in to_bin");
		if(bin_dbl < 0) return 0;
		if(bin_dbl > nbins-1) return nbins-1;
		int bin_int = int(bin_dbl);
		assert(bin_int >= 0 && bin_int < nbins);
		return bin_int;
	}

	double from_bin(int i) const {
		return double(i) * scale + offset;
	}

	int nbins;
	double offset;
	double scale;
};

struct Histogram {
	Histogram() :
		min(0), max(0), mean(0), stddev(0),
		data_count(0), ndv_count(0)
	{ }

	Binning binning;
	double min, max, mean, stddev;
	size_t data_count;
	size_t ndv_count;
	std::vector<size_t> counts;
};

std::vector<std::pair<double, double> > compute_minmax(
	const std::vector<GDALRasterBandH> &src_bands, const NdvDef &ndv_def,
	size_t w, size_t h
);
std::vector<Histogram> compute_histogram(
	const std::vector<GDALRasterBandH> &src_bands, const NdvDef &ndv_def,
	size_t w, size_t h, const std::vector<Binning> &binnings
);
void get_scale_from_percentile(
	const Histogram &histogram, int output_range,
	double from_percentile, double to_percentile,
	double *scale_out, double *offset_out
);
std::vector<uint8_t> invert_histogram_to_gaussian(const Histogram &histogram_in, double variance,
	int output_range);
void copyGeoCode(GDALDatasetH dst_ds, GDALDatasetH src_ds);

void usage(const std::string &cmdname) {
	printf("Usage: %s <options> src.tif dst.tif\n\n", cmdname.c_str());
	NdvDef::printUsage();
	printf(
"  -outndv <output_nodata_val>        Output no-data value\n"
"\n"
"Operation:\n"
"  -linear-stretch <target_avg> <target_stddev>      Linear stretch to a target range\n"
"  -percentile-range <from: 0.0-1.0> <to: 0.0-1.0>   Linear stretch using a percentile range of input\n"
"  -histeq <target_stddev>                           Histogram normalize to a target bell curve\n"
"  -dump-histogram                                   Just print the histogram to console\n"
"\n"
"Input can be any integer or floating type (but not complex).  Output is 8-bit.\n"
);
	exit(1);
}

int main(int argc, char *argv[]) {
	const std::string cmdname = argv[0];
	if(argc == 1) usage(cmdname);
	std::vector<std::string> arg_list = argv_to_list(argc, argv);

	std::string src_fn;
	std::string dst_fn;
	std::string output_format;

	int mode_histeq = 0;
	int mode_stddev = 0;
	double dst_avg = -1;
	double dst_stddev = -1;
	int mode_percentile = 0;
	int mode_dump_histogram = 0;
	double from_percentile = -1;
	double to_percentile = -1;
	int out_ndv = 0, set_out_ndv = 0;

	NdvDef ndv_def = NdvDef(arg_list);

	size_t argp = 1;
	while(argp < arg_list.size()) {
		const std::string &arg = arg_list[argp++];
		// FIXME - check for duplicate values
		if(arg[0] == '-') {
			try {
				if(arg == "-of") {
					if(argp == arg_list.size()) usage(cmdname);
					output_format = arg_list[argp++];
				} else if(arg == "-linear-stretch") {
					if(argp == arg_list.size()) usage(cmdname);
					dst_avg = boost::lexical_cast<double>(arg_list[argp++]);

					if(argp == arg_list.size()) usage(cmdname);
					dst_stddev = boost::lexical_cast<double>(arg_list[argp++]);

					mode_stddev = 1;
				} else if(arg == "-percentile-range") {
					if(argp == arg_list.size()) usage(cmdname);
					from_percentile = boost::lexical_cast<double>(arg_list[argp++]);

					if(argp == arg_list.size()) usage(cmdname);
					to_percentile = boost::lexical_cast<double>(arg_list[argp++]);

					mode_percentile = 1;
				} else if(arg == "-histeq") {
					if(argp == arg_list.size()) usage(cmdname);
					dst_stddev = boost::lexical_cast<double>(arg_list[argp++]);

					mode_histeq = 1;
				} else if(arg == "-dump-histogram") {
					mode_dump_histogram = 1;
				} else if(arg == "-outndv") {
					if(argp == arg_list.size()) usage(cmdname);
					int64_t ndv_long = boost::lexical_cast<int64_t>(arg_list[argp++]);
					if(ndv_long < 0 || ndv_long > 255) fatal_error("ndv must be in the range 0..255");
					out_ndv = boost::numeric_cast<uint8_t>(ndv_long);
					set_out_ndv++;
				} else {
					usage(cmdname);
				}
			} catch(boost::bad_lexical_cast &e) {
				fatal_error("cannot parse number given on command line");
			} catch(boost::bad_numeric_cast &e) {
				fatal_error("number given on command line out of range: %s", e.what());
			}
		} else {
			if(src_fn.empty()) {
				src_fn = arg;
			} else if(dst_fn.empty()) {
				dst_fn = arg;
			} else {
				usage(cmdname);
			}
		}
	}

	if(src_fn.empty()) usage(cmdname);
	if(dst_fn.empty() != (mode_dump_histogram > 0)) usage(cmdname);
	if(mode_percentile + mode_stddev + mode_histeq + mode_dump_histogram > 1) usage(cmdname);
	if(mode_stddev && (dst_avg < 0 || dst_stddev < 0)) usage(cmdname);
	if(mode_percentile && !(
		0 <= from_percentile &&
		from_percentile < to_percentile &&
		to_percentile <= 1)) usage(cmdname);

	if(output_format.empty()) output_format = "GTiff";

	GDALAllRegister();

	//////// open source ////////

	GDALDatasetH src_ds = GDALOpen(src_fn.c_str(), GA_ReadOnly);
	if(!src_ds) fatal_error("open failed");

	size_t w = GDALGetRasterXSize(src_ds);
	size_t h = GDALGetRasterYSize(src_ds);
	if(!w || !h) fatal_error("missing width/height");
	size_t src_band_count = GDALGetRasterCount(src_ds);
	printf("Input size is %zd, %zd, %zd\n", w, h, src_band_count);

	std::vector<size_t> bandlist;
	for(size_t i=0; i<src_band_count; i++) {
		bandlist.push_back(i+1);
	}
	size_t dst_band_count = bandlist.size();

	if(ndv_def.empty()) {
		ndv_def = NdvDef(src_ds, bandlist);
	}

	bool use_ndv = !ndv_def.empty();
	if(use_ndv && !set_out_ndv) {
		if(ndv_def.slabs.size() == 1) {
			const NdvSlab &slab = ndv_def.slabs[0];
			assert(slab.range_by_band.size());
			const NdvInterval &range = slab.range_by_band[0];
			if(range.first == range.second) {
				double v = range.first;
				out_ndv = (uint8_t)v;
				if((double)out_ndv == v) {
					set_out_ndv++;
				} else {
					printf("Cannot use %g as an NDV value because it cannot be cast to an 8-bit number.\n", v);
				}
			}
		}
		if(!set_out_ndv) {
			out_ndv = 0;
			set_out_ndv++;
		}
		printf("Output NDV is defaulting to %d.\n", out_ndv);
	}

	//////// open output ////////

	printf("Output size is %zd x %zd x %zd\n", w, h, dst_band_count);

	GDALDriverH dst_driver = GDALGetDriverByName(output_format.c_str());
	if(!dst_driver) fatal_error("unrecognized output format (%s)", output_format.c_str());

	//////// open bands ////////

	std::vector<GDALRasterBandH> src_bands;
	std::vector<GDALRasterBandH> dst_bands;

	for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
		src_bands.push_back(GDALGetRasterBand(src_ds, bandlist[band_idx]));
	}

	//////// find optimal binning ////////

	std::vector<Binning> binnings(dst_band_count);
	{
		// computed on demand
		std::vector<std::pair<double, double> > minmax;

		for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
			Binning &binning = binnings[band_idx];
			GDALDataType dt = GDALGetRasterDataType(src_bands[band_idx]);
			switch(dt) {
				case GDT_Byte:
					binning.nbins = 256;
					binning.offset = 0;
					binning.scale = 1;
					break;
				case GDT_UInt16:
					binning.nbins = 65536;
					binning.offset = 0;
					binning.scale = 1;
					break;
				case GDT_Int16:
					binning.nbins = 65536;
					binning.offset = -32768;
					binning.scale = 1;
					break;
				default:
					if(minmax.empty()) {
						printf("Computing min/max values...\n");
						minmax = compute_minmax(src_bands, ndv_def, w, h);
					}
					double min = minmax[band_idx].first;
					double max = minmax[band_idx].second;
					// a compromise between memory usage and datavalue resolution
					binning.nbins = 10000000;
					binning.offset = min;
					binning.scale = (max - min) / double(binning.nbins-1);
			}
		}
	}

	//////// compute lookup table ////////

	printf("\nComputing histogram...\n");
	std::vector<Histogram> histograms =
		compute_histogram(src_bands, ndv_def, w, h, binnings);

	for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
		Histogram &hg = histograms[band_idx];
		printf("band %zd: min=%g, max=%g, mean=%g, stddev=%g, valid_count=%zd, ndv_count=%zd\n",
			band_idx+1, hg.min, hg.max, hg.mean, hg.stddev, hg.data_count, hg.ndv_count);
		if(mode_dump_histogram) {
			for(int i=0; i<hg.binning.nbins; i++) {
				printf("bin %d: val=%g cnt=%zd\n",
					i, hg.binning.from_bin(i), hg.counts[i]);
			}
		}
	}
	if(mode_dump_histogram) {
		return 0;
	}

	//////// open output ////////

	GDALDatasetH dst_ds = GDALCreate(dst_driver, dst_fn.c_str(), w, h, dst_band_count, GDT_Byte, NULL);
	if(!dst_ds) fatal_error("couldn't create output");
	copyGeoCode(dst_ds, src_ds);

	for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
		dst_bands.push_back(GDALGetRasterBand(dst_ds, band_idx+1));
	}

	//////// compute tranformation parameters ////////

	const int output_range = 256;
	bool use_table; // otherwise, use linear
	std::vector<std::vector<uint8_t> > xform_table(dst_band_count);
	std::vector<double> lin_scales(dst_band_count);
	std::vector<double> lin_offsets(dst_band_count);

	if(mode_histeq) {
		use_table = true;
		for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
			xform_table[band_idx] = invert_histogram_to_gaussian(
				histograms[band_idx], dst_stddev, output_range);
		}
	} else{
		use_table = false;
		bool was_nop = false;
		if(mode_percentile) {
			for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
				get_scale_from_percentile(
					histograms[band_idx], output_range, from_percentile, to_percentile,
					&lin_scales[band_idx], &lin_offsets[band_idx]);
			}
		} else if(mode_stddev) {
			for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
				Histogram &hg = histograms[band_idx];
				lin_scales[band_idx] = hg.stddev ? dst_stddev / hg.stddev : 0;
				lin_offsets[band_idx] = hg.mean - dst_avg / lin_scales[band_idx];
			}
		} else { // no transformation
			printf("\nWarning: no transformation was specified!  I'll just cast the input to 8-bit.\n");
			was_nop = true;
			for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
				lin_scales[band_idx] = 1;
				lin_offsets[band_idx] = 0;
			}
		}
		if(!was_nop) {
			printf("\nLinear stretch:\n");
			for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
				double scale = lin_scales[band_idx];
				double offset = lin_offsets[band_idx];
				printf("band %zd: scale=%f, offset=%f, src_range=[%f, %f]\n",
					band_idx+1, scale, offset, offset, ((double)(output_range-1)/scale)+offset);
			}
		}
	}

	histograms.clear(); // free up memory

	if(use_table) {
		// avoid ndv in output for good pixels
		if(use_ndv) {
			for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
				for(size_t i=0; i<xform_table[band_idx].size(); i++) {
					uint8_t v = xform_table[band_idx][i];
					if(v == out_ndv) {
						if(out_ndv < output_range/2) v++;
						else v--;
					}
					xform_table[band_idx][i] = v;
				}
			}
		}
	}

	//////// do transformation ////////

	printf("\nComputing output...\n");

	int blocksize_x_int, blocksize_y_int;
	GDALGetBlockSize(src_bands[0], &blocksize_x_int, &blocksize_y_int);
	size_t blocksize_x = blocksize_x_int;
	size_t blocksize_y = blocksize_y_int;
	size_t block_len = blocksize_x*blocksize_y;

	std::vector<std::vector<double> > buf_in(dst_band_count);
	std::vector<std::vector<uint8_t> > buf_out(dst_band_count);
	for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
		buf_in[band_idx].resize(block_len);
		buf_out[band_idx].resize(block_len);
	}
	std::vector<uint8_t> ndv_mask(block_len);

	for(size_t boff_y=0; boff_y<h; boff_y+=blocksize_y) {
		size_t bsize_y = blocksize_y;
		if(bsize_y + boff_y > h) bsize_y = h - boff_y;
		for(size_t boff_x=0; boff_x<w; boff_x+=blocksize_x) {
			size_t bsize_x = blocksize_x;
			if(bsize_x + boff_x > w) bsize_x = w - boff_x;

			block_len = bsize_x*bsize_y;

			double progress =
				((double)boff_y * (double)w +
				(double)boff_x * (double)bsize_y) /
				((double)w * (double)h);
			GDALTermProgress(progress, NULL, NULL);

			for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
				GDALRasterIO(src_bands[band_idx], GF_Read, boff_x, boff_y, bsize_x, bsize_y,
					&buf_in[band_idx][0], bsize_x, bsize_y, GDT_Float64, 0, 0);
			}

			ndv_def.getNdvMask(buf_in, &ndv_mask[0], block_len);

			for(size_t band_idx=0; band_idx<dst_band_count; band_idx++) {
				double *p_in = &buf_in[band_idx][0];
				uint8_t *p_out = &buf_out[band_idx][0];
				uint8_t *p_ndv = &ndv_mask[0];
				if(use_table) {
					uint8_t *xform = &xform_table[band_idx][0];
					Binning binning = binnings[band_idx];

					for(size_t i=0; i<block_len; i++) {
						if(*p_ndv) {
							*p_out = out_ndv;
						} else {
							*p_out = xform[binning.to_bin(*p_in)];
							//printf("%g %d %d\n", *p_in, binning.to_bin(*p_in), *p_out);
						}
						p_in++; p_out++; p_ndv++;
					}
				} else {
					double scale = lin_scales[band_idx];
					double offset = lin_offsets[band_idx];
					for(size_t i=0; i<block_len; i++) {
						if(*p_ndv) {
							*p_out = out_ndv;
						} else {
							double out_dbl = (*p_in - offset) * scale;
							uint8_t v =
								(out_dbl < 0) ? 0 :
								(out_dbl > output_range-1) ? output_range-1 :
								uint8_t(out_dbl);
							if(v == out_ndv) {
								if(out_ndv < output_range/2) v++;
								else v--;
							}
							*p_out = v;
						}
						p_in++; p_out++; p_ndv++;
					}
				}

				GDALRasterIO(dst_bands[band_idx], GF_Write, boff_x, boff_y, bsize_x, bsize_y,
					&buf_out[band_idx][0], bsize_x, bsize_y, GDT_Byte, 0, 0);
			} // band
		} // block x
	} // block y

	GDALClose(src_ds);
	GDALClose(dst_ds);

	GDALTermProgress(1, NULL, NULL);

	return 0;
}

std::vector<std::pair<double, double> > compute_minmax(
	const std::vector<GDALRasterBandH> &src_bands, const NdvDef &ndv_def,
	size_t w, size_t h
) {
	size_t band_count = src_bands.size();
	std::vector<std::pair<double, double> > minmax(band_count);

	int blocksize_x_int, blocksize_y_int;
	GDALGetBlockSize(src_bands[0], &blocksize_x_int, &blocksize_y_int);
	size_t blocksize_x = blocksize_x_int;
	size_t blocksize_y = blocksize_y_int;
	size_t block_len = blocksize_x*blocksize_y;

	std::vector<std::vector<double> > buf_in(band_count);
	for(size_t band_idx=0; band_idx<band_count; band_idx++) {
		buf_in[band_idx].resize(block_len);
	}
	std::vector<uint8_t> ndv_mask(block_len);

	std::vector<bool> got_data(band_count);

	for(size_t boff_y=0; boff_y<h; boff_y+=blocksize_y) {
		size_t bsize_y = blocksize_y;
		if(bsize_y + boff_y > h) bsize_y = h - boff_y;
		for(size_t boff_x=0; boff_x<w; boff_x+=blocksize_x) {
			size_t bsize_x = blocksize_x;
			if(bsize_x + boff_x > w) bsize_x = w - boff_x;

			block_len = bsize_x*bsize_y;

			double progress =
				((double)boff_y * (double)w +
				(double)boff_x * (double)bsize_y) /
				((double)w * (double)h);
			GDALTermProgress(progress, NULL, NULL);

			for(size_t band_idx=0; band_idx<band_count; band_idx++) {
				GDALRasterIO(src_bands[band_idx], GF_Read, boff_x, boff_y, bsize_x, bsize_y,
					&buf_in[band_idx][0], bsize_x, bsize_y, GDT_Float64, 0, 0);
			}

			ndv_def.getNdvMask(buf_in, &ndv_mask[0], block_len);

			for(size_t band_idx=0; band_idx<band_count; band_idx++) {
				for(size_t i=0; i<block_len; i++) {
					if(ndv_mask[i]) continue;
					double v = buf_in[band_idx][i];
					if(std::isnan(v) || std::isinf(v)) continue;

					assert(!minmax.empty());
					double &min = minmax[band_idx].first;
					double &max = minmax[band_idx].second;
					if(!got_data[band_idx]) {
						min = v;
						max = v;
						got_data[band_idx] = true;
					}
					if(v < min) min = v;
					if(v > max) max = v;
				}
			}
		}
	}
	GDALTermProgress(1, NULL, NULL);

	return minmax;
}

std::vector<Histogram> compute_histogram(
	const std::vector<GDALRasterBandH> &src_bands, const NdvDef &ndv_def,
	size_t w, size_t h, const std::vector<Binning> &binnings
) {
	size_t band_count = src_bands.size();
	std::vector<Histogram> histograms(band_count);

	for(size_t band_idx=0; band_idx<band_count; band_idx++) {
		histograms[band_idx].binning = binnings[band_idx];
		histograms[band_idx].counts.assign(binnings[band_idx].nbins, 0);
	}

	int blocksize_x_int, blocksize_y_int;
	GDALGetBlockSize(src_bands[0], &blocksize_x_int, &blocksize_y_int);
	size_t blocksize_x = blocksize_x_int;
	size_t blocksize_y = blocksize_y_int;
	size_t block_len = blocksize_x*blocksize_y;

	std::vector<std::vector<double> > buf_in(band_count);
	for(size_t band_idx=0; band_idx<band_count; band_idx++) {
		buf_in[band_idx].resize(block_len);
	}
	std::vector<uint8_t> ndv_mask(block_len);

	std::vector<bool> first_valid_pixel(band_count, true);

	for(size_t boff_y=0; boff_y<h; boff_y+=blocksize_y) {
		size_t bsize_y = blocksize_y;
		if(bsize_y + boff_y > h) bsize_y = h - boff_y;
		for(size_t boff_x=0; boff_x<w; boff_x+=blocksize_x) {
			size_t bsize_x = blocksize_x;
			if(bsize_x + boff_x > w) bsize_x = w - boff_x;

			block_len = bsize_x*bsize_y;

			double progress =
				((double)boff_y * (double)w +
				(double)boff_x * (double)bsize_y) /
				((double)w * (double)h);
			GDALTermProgress(progress, NULL, NULL);

			for(size_t band_idx=0; band_idx<band_count; band_idx++) {
				GDALRasterIO(src_bands[band_idx], GF_Read, boff_x, boff_y, bsize_x, bsize_y,
					&buf_in[band_idx][0], bsize_x, bsize_y, GDT_Float64, 0, 0);
			}

			ndv_def.getNdvMask(buf_in, &ndv_mask[0], block_len);

			for(size_t band_idx=0; band_idx<band_count; band_idx++) {
				Histogram &hg = histograms[band_idx];
				double *p = &buf_in[band_idx][0];

				for(size_t i=0; i<block_len; i++) {
					if(ndv_mask[i]) {
						hg.ndv_count++;
					} else {
						double v = p[i];
						hg.counts[hg.binning.to_bin(v)]++;
						if(first_valid_pixel[band_idx]) {
							hg.min = hg.max = v;
							first_valid_pixel[band_idx] = false;
						}
						if(v < hg.min) hg.min = v;
						if(v > hg.max) hg.max = v;
					}
				}
			}
		}
	}
	GDALTermProgress(1, NULL, NULL);

	for(size_t band_idx=0; band_idx<band_count; band_idx++) {
		Histogram &hg = histograms[band_idx];
		double accum = 0;
		for(int i=0; i<hg.binning.nbins; i++) {
			size_t cnt = hg.counts[i];
			double v = hg.binning.from_bin(i);
			hg.data_count += cnt;
			accum += v * cnt;
		}
		hg.mean = accum / hg.data_count;

		double var_accum = 0;
		for(int i=0; i<hg.binning.nbins; i++) {
			size_t cnt = hg.counts[i];
			double v = hg.binning.from_bin(i);
			var_accum += (v-hg.mean) * (v-hg.mean) * cnt;
		}
		hg.stddev = sqrt(var_accum / hg.data_count);
	}

	return histograms;
}

void get_scale_from_percentile(
	const Histogram &histogram, int output_range,
	double from_percentile, double to_percentile,
	double *scale_out, double *offset_out
) {
	size_t start_count = (size_t)(histogram.data_count * from_percentile);
	size_t end_count = (size_t)(histogram.data_count * to_percentile);

	size_t cnt = 0;
	int from_idx = -1;
	int to_idx = -1;
	for(int i=0; i<histogram.binning.nbins; i++) {
		if(cnt <= start_count) from_idx = i;
		cnt += histogram.counts[i];
		if(cnt <= end_count) to_idx = i;
		else break;
	}
	if(from_idx<0 || to_idx<0) fatal_error("impossible: could not find window");
	if(from_idx == to_idx) { from_idx=0; to_idx=histogram.binning.nbins-1; }

	double from_val = histogram.binning.from_bin(from_idx);
	double to_val = histogram.binning.from_bin(to_idx);

	*scale_out = (double)(output_range-1) / (double)(to_val-from_val);
	*offset_out = from_val;
}

std::vector<double> gen_gaussian(double variance, int bin_count) {
	std::vector<double> arr(bin_count);
	double total = 0;
	for(int i=0; i<bin_count; i++) {
		if(variance == 0.0) {
			// if variance==0.0, just use a uniform distribution
			arr[i] = 1;
		} else {
			// Gaussian distribution.  Large variance gives
			// a more level curve, small variance gives
			// a curved peaked at center.
			double x = (double)(i-bin_count/2) / variance;
			arr[i] = exp(-x*x);
		}
		total += arr[i];
	}

	for(int i=0; i<bin_count; i++) arr[i] /= total;

	return arr;
}

std::vector<uint8_t> invert_histogram(
	const Histogram &src_h_in,
	const std::vector<double> &dst_h,
	size_t output_range
) {
	std::vector<size_t> src_h(src_h_in.counts);
	size_t pixel_count = 0;
	for(size_t i=0; i<src_h.size(); i++) pixel_count += src_h[i];

	std::vector<uint8_t> out_h(src_h.size());
	double src_total = 0;
	double dst_total = 0;
	uint8_t j = 0;
	for(size_t i=0; i<src_h.size(); i++) {
		out_h[i] = j;
		src_total += src_h[i];
		while(j<output_range-1 && dst_total < src_total) {
			dst_total += dst_h[j++] * (double)pixel_count;
		}
	}

	return out_h;
}

std::vector<uint8_t> invert_histogram_to_gaussian(
	const Histogram &histogram_in, double variance,
	int output_range
) {
	std::vector<double> gaussian = gen_gaussian(variance, output_range);
	return invert_histogram(histogram_in, gaussian, output_range);
}

void copyGeoCode(GDALDatasetH dst_ds, GDALDatasetH src_ds) {
	double affine[6];
	if(GDALGetGeoTransform(src_ds, affine) == CE_None) {
		GDALSetGeoTransform(dst_ds, affine);
	}
	GDALSetProjection(dst_ds, GDALGetProjectionRef(src_ds));
}


/*
 * File:   UCKneeIdentifier.cpp
 * Author: Cara Higgins
 *
 * Created on 20 October 2021, 16:12
 */

#include "UCKneeIdentifier.h"

#include <climits>

#include <boost/bind.hpp>

#include <common/partitionforests/images/GeneralFeature.h>
#include <common/dicom/volumes/DICOMVolume.h>
#include <common/util/ITKImageUtil.h>

namespace mp {

    //#################### CONSTRUCTORS ####################

    UCKneeIdentifier::UCKneeIdentifier(const DICOMVolume_CPtr& dicomVolume, const VolumeIPF_Ptr& volumeIPF)
    : FeatureIdentifier(dicomVolume, volumeIPF) {
    }

    //#################### PUBLIC METHODS ####################

    int UCKneeIdentifier::length() const {
        return 5;
    }

    //#################### PRIVATE METHODS ####################

    std::tuple<std::list<PFNodeID>, int> ucknee;

    // /**
    // @brief	Select region of interest (ROI) by setting window minima and maxima

    // @param[in]	volumeSize	    [x width, y height, z count]
    // @param[in]  distFromEdge    User-selected parameter indicating the minimum distance a node must be from the edge;
    //                             prevents selection of border nodes, which often have very bright values
    // @param[in]  windMinX        Pointer to an integer variable representing the minimum x-value for the ROI
    // @param[in]  windMinY
    // @param[in]  windMaxX
    // @param[in]  windMaxY
    // */
    // void selectROI(itk::Size<3> volumeSize, int distFromEdge, int* windMinX, int* windMinY, int* windMaxX, int* windMaxY) {
    //     int width = volumeSize[0];
    //     int height = volumeSize[1];
    //     float htow = height * 1.0 / width;

    //     *windMinY = distFromEdge;
    //     *windMinX = distFromEdge;
    //     *windMaxX = width-distFromEdge;
    //     *windMaxY = height-distFromEdge;

    //     // TODO: check that distFromEdge <= width / 4, height / 4
    //     // Determine region of focus based on image size
    //     if (0.75 < htow && htow < 1.5) { // Roughly square
    //         *windMinX = width / 4;
    //         *windMaxX = width * 3 / 4;

    //         *windMinY = height / 4;
    //         *windMaxY = height * 3 / 4;
    //     } else if (htow > 1.5) {
    //         *windMinY = height / 4;
    //         *windMaxY = height * 3 / 4;
    //     } else {
    //         *windMinX = width / 4;
    //         *windMaxX = width * 3 / 4;
    //     }

    //     std::cout<<"Height to weight ratio: "<<htow<<std::endl;
    //     std::cout<<"image (height x width) = "<<height<<" x "<<width<<std::endl;
    //     std::cout<<"window = ("<<*windMinX<<", "<<*windMinY<<") -> ("<<*windMaxX<<", "<<*windMaxY<<")"<<std::endl;
    // }

    void UCKneeIdentifier::execute_impl() {
        set_status("Identifying unicompartmental knee...");
        std::cout<<"OnMenuSegmentationSegmentUCKnee reached."<<std::endl;

        int windMinX, windMinY, windMaxX, windMaxY;
        // TODO: only call if ROI is not already initialized
        selectROI(dicom_volume()->size(), 10, &windMinX, &windMinY, &windMaxX, &windMaxY);

        std::vector<int> node_indices = volume_ipf()->node_indices(layerIndex);
        std::cout<<"Nodes in layer "<<layerIndex<<" = "<<node_indices.size()<<std::endl;

        std::tuple<std::list<PFNodeID>, int> nodesWithMaxMeanGray = max_mean_grey_with_filter(
            node_indices, windMinX, windMinY, windMaxX, windMaxY
        );

        std::list<PFNodeID> nodes = std::get<0>(nodesWithMaxMeanGray);
        int maxMeanGrey = std::get<1>(nodesWithMaxMeanGray);

        int threshold = maxMeanGrey - 20;

        std::list<component> connectedComps = connected_components(nodes, threshold);
        connectedComps.sort(compare_components);    // Sort connected components by voxel count

        // Initialize region to store unicompartmental knee nodes
        PartitionForestSelection_Ptr filledRegion(new PartitionForestSelectionT(volume_ipf()));

        // Select the two largest components and set the ucknee variable accordingly
        for (int i = 0; i < 2 && !connectedComps.empty(); i++) {
            component comp = connectedComps.front();
            connectedComps.pop_front();

            std::set<PFNodeID> compNodes(std::get<0>(comp).begin(), std::get<0>(comp).end());

            std::cout<<"calling merge_nonsibling_nodes"<<std::endl;
            compNodes = volume_ipf()->merge_nonsibling_nodes(compNodes, volume_ipf()->CheckPreconditions::CHECK_PRECONDITIONS);
            std::cout<<"merge_nonsibling_nodes executed"<<std::endl;
            PFNodeID compNode = *compNodes.begin();
            // PFNodeID compNode = volume_ipf()->merge_sibling_nodes(compNodes, volume_ipf()->CheckPreconditions::CHECK_PRECONDITIONS);

            ucknee.push_back(compNode);
            filledRegion->select_node(compNode);
        }

        VolumeIPFMultiFeatureSelection_Ptr multiFeatureSelection = get_multi_feature_selection();
        multiFeatureSelection->identify_selection(filledRegion, GeneralFeature::UC_KNEE);
        std::cout<<"Selected top 2 connected components"<<std::endl;
        get_bounding_boxes();
    }

    /**
    @brief	Compares two components based on their associated voxel count.

    @param[in]	first	The first component
    @param[in]  second  The second component

    @return	true if first has a voxel count greater than that of the second component.
    */
    bool UCKneeIdentifier::compare_components(const std::tuple<std::list<PFNodeID>, int>& first, const std::tuple<std::list<PFNodeID>, int>& second)
    {
        return std::get<1>(first) > std::get<1>(second);
    }

    /**
    @brief	Prints the bounding boxes for the two components of the UC knee prosthesis
    */
    void UCKneeIdentifier::get_bounding_boxes()
    {
        for (std::list<PFNodeID>::iterator it = ucknee.begin(); it != ucknee.end(); it++) {
            propertyMap propMap = volume_ipf()->branch_properties(*it).branch_property_map();

            int MinX = std::stoi(propMap.at("X Min"));
            int MinY = std::stoi(propMap.at("Y Min"));
            int MaxX = std::stoi(propMap.at("X Max"));
            int MaxY = std::stoi(propMap.at("Y Max"));

            std::cout<<"Bounding box: ("<<MinX<<", "<<MinY<<") -> ("<<MaxX<<", "<<MaxY<<")"<<std::endl;
        }
    }

    // std::tuple<std::list<PFNodeID>, int> UCKneeIdentifier::max_mean_grey_with_filter(
    //     std::vector<int> node_indices,
    //     int windMinX, int windMinY, int windMaxX, int windMaxY
    // ) {
    //     int meanGrey; int VoxelCount; int MinX; int MinY; int MaxX; int MaxY;
    //     int maxMeanGrey = 0;
    //     std::list<PFNodeID> nodes;

    //     for(std::vector<int>::iterator indexIt = node_indices.begin(); indexIt != node_indices.end(); indexIt++)
    //     {
    //         PFNodeID node = PFNodeID(layerIndex, *indexIt);
    //         propertyMap propMap = volume_ipf()->branch_properties(node).branch_property_map();

    //         meanGrey = std::stof(propMap.at("Mean Grey Value"));
    //         VoxelCount = std::stoi(propMap.at("Voxel Count"));
    //         MinX = std::stoi(propMap.at("X Min"));
    //         MinY = std::stoi(propMap.at("Y Min"));
    //         MaxX = std::stoi(propMap.at("X Max"));
    //         MaxY = std::stoi(propMap.at("Y Max"));

    //         if (
    //             VoxelCount <= 15000 && VoxelCount >= 10 && // Node must contain between 100 and 15,000 voxels
    //             windMinX < MinX && MinX < windMaxX && windMinY < MinY && MinY < windMaxY // Region lies within region of interest (not touching edges)
    //         ) {
    //             nodes.push_back(node);  // Add node to list of filtered nodes

    //             if (meanGrey > maxMeanGrey)
    //                 maxMeanGrey = meanGrey;
    //         }
    //     }

    //     std::cout<<"Max Mean Grey Value = "<<maxMeanGrey<<std::endl;

    //     return std::tuple<std::list<PFNodeID>, int>(nodes, maxMeanGrey);
    // }

    std::list<UCKneeIdentifier::component> UCKneeIdentifier::connected_components(std::list<PFNodeID> nodes, int threshold)
    {
        std::list<component> connectedComps;

        for (std::list<PFNodeID>::iterator it = nodes.begin(); it != nodes.end(); it++) {
            propertyMap propMap = volume_ipf()->branch_properties(*it).branch_property_map();
            int meanGrey = std::stof(propMap.at("Mean Grey Value"));

            if (meanGrey > threshold) { // Check that mean grey value is above the specified threshold
                PFNodeID node = *it;
                component* nodeAddedTo = nullptr; // Pointer to the component the node has been added to

                // Obtain set of adjacent nodes
                std::vector<int> adjNodes = volume_ipf()->adjacent_nodes(node);
                std::set<int> adjNodeSet(adjNodes.begin(), adjNodes.end());

                // Iterate through connected components
                for (std::list<component>::iterator compIt = connectedComps.begin(); compIt != connectedComps.end(); compIt++) {
                    std::list<PFNodeID> compNodes = std::get<0>(*compIt);

                    std::cout<<connectedComps.size()<<" connected components; "<<compNodes.size()<<" nodes in this"<<std::endl;

                    // TODO: use are_connected function
                    // Convert compNodes to a set of node indices
                    std::set<int> nodes;
                    for (std::list<PFNodeID>::iterator nodeIt = compNodes.begin(); nodeIt != compNodes.end(); nodeIt++) {
                        nodes.insert(nodeIt->index());
                    }

                    // Check if the node is adjacent to the component
                    if (volume_ipf()->are_connected(nodes, layerIndex)) {
                        std::cout<<"Connected"<<std::endl;
                        if (!nodeAddedTo) {
                            // If new node is adjacent to the component and hasn't been added to another component,
                            // add it to this component & update voxel count
                            compNodes.push_back(node);
                            std::get<0>(*compIt) = compNodes;
                            std::get<1>(*compIt) += std::stoi(propMap.at("Voxel Count"));
                            nodeAddedTo = &(*compIt);
                            std::cout<<"Node added to existing component."<<std::endl;
                        } else {
                            // Node has already been added to another component, but is also adjacent to this
                            // component, so the components should be merged
                            std::list<PFNodeID> nodeAddedToList = std::get<0>(*nodeAddedTo);
                            nodeAddedToList.splice(nodeAddedToList.begin(), compNodes);
                            std::get<0>(*nodeAddedTo) = nodeAddedToList;
                            std::get<1>(*nodeAddedTo) += std::get<1>(*compIt);

                            // Preserve compIt after erasing component
                            connectedComps.erase(compIt++); // Remove old component, after merging
                            compIt--;
                            std::cout<<"Merging components."<<std::endl;
                        }
                    }
                } // End connected component iteration

                if (!nodeAddedTo) { // If node not adjacent to any connected component, add singleton to connected components
                    component newComp = component(std::list<PFNodeID>(1, node), std::stoi(propMap.at("Voxel Count")));
                    connectedComps.push_back(newComp);
                    std::cout<<"Node added as a singleton component."<<std::endl;
                }
            }
        } // End node iteration

        return connectedComps;
    }

    UCKneeIdentifier::BranchProperties UCKneeIdentifier::calculate_component_properties(int layer, const std::set<int>& indices) const
    {
        std::vector<BranchProperties> nodeProperties;
        nodeProperties.reserve(indices.size());
        for(std::set<int>::const_iterator it=indices.begin(), iend=indices.end(); it!=iend; ++it)
        {
            PFNodeID node(layer, *it);
            nodeProperties.push_back(volume_ipf()->branch_properties(node));
        }
        return BranchProperties::combine_branch_properties(nodeProperties);
    }

    // std::set<PFNodeID> decompose_ucknee()
    // {

    //     // return split_node(const PFNodeID& node, const std::vector<std::set<int> >& groups, CheckPreconditions checkPreconditions = CHECK_PRECONDITIONS)
    // }
}


#include <cctbx/boost_python/flex_fwd.h>

#include <cctbx/miller/phase_integrator.h>
#include <boost/python/class.hpp>
#include <boost/python/args.hpp>

namespace cctbx { namespace miller { namespace boost_python {

namespace {

  struct phase_integrator_wrappers
  {
    typedef phase_integrator<> w_t;

    static void
    wrap()
    {
      using namespace boost::python;
      class_<w_t>("phase_integrator", no_init)
        .def(init<optional<unsigned> >(arg("n_steps")=360/5))
        .def("n_steps", &w_t::n_steps)
        .def("__call__",
          (std::complex<double>(w_t::*)(
            sgtbx::phase_info const&,
            hendrickson_lattman<> const&) const) &w_t::operator(), (
          arg("phase_info"), arg("hendrickson_lattman")))
        .def("__call__",
          (af::shared<std::complex<double> >(w_t::*)(
            sgtbx::space_group const&,
            af::const_ref<miller::index<> > const&,
            af::const_ref<hendrickson_lattman<> > const&) const)
              &w_t::operator(), (
          arg("space_group"),
          arg("miller_indices"),
          arg("hendrickson_lattman_coefficients")))
      ;
    }
  };

  struct phase_entropy_wrappers
  {
    typedef phase_entropy<> w_t;

    static void
    wrap()
    {
      using namespace boost::python;
      class_<w_t>("phase_entropy", no_init)
        .def(init<optional<unsigned> >(arg("n_steps")=360/5))
        .def("n_steps", &w_t::n_steps)
        .def("relative_entropy", &w_t::relative_entropy)
      ;
    }
  };

} // namespace <anoymous>

  void wrap_phase_integrator()
  {
    phase_integrator_wrappers::wrap();
    phase_entropy_wrappers::wrap();
  }

}}} // namespace cctbx::miller::boost_python


#pragma once

#include <Eigen/Core>
#include <steam.hpp>
#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/pose_with_covariance.hpp>

#include <cpo_interfaces/msg/tdcp.hpp>
#include <cpo_interfaces/srv/query_trajectory.hpp>

#include <deque>

using QueryTrajectory = cpo_interfaces::srv::QueryTrajectory;

struct CpoEdge {
  cpo_interfaces::msg::TDCP msg;    /// the TDCP msg spanning this edge
  lgmath::se3::Transformation T_ba; /// transformation estimate
  Eigen::Matrix<double, 6, 1> v_a;  /// velocity estimate at vertex a
  Eigen::Matrix<double, 6, 1> v_b;  /// velocity estimate at vertex a
};

/** \brief Class that subscribes to TDCP pseudo-measurements and outputs odometry estimates */
class CpoBackEnd : public rclcpp::Node {
 public:

  /** \brief Constructor */
  CpoBackEnd();

  /** \brief Get parameters from ROS2 and set appropriate fields */
  void getParams();

  /** \brief Helper/debugger function that prints total cost and number of terms for each group of cost terms
   * \param final  True if these are the final costs, if false assumed to be initial costs
   * */
  void printCosts(bool final);

  /** \brief Helper function */
  static Eigen::Vector3d toEigenVec3d(const geometry_msgs::msg::Vector3 &ros_vec);

 private:

  /** \brief Callback for TDCP msgs */
  void _tdcpCallback(cpo_interfaces::msg::TDCP::SharedPtr msg_in);

  /** \brief Timed callback to query trajectory if available and publish pose estimate */
  void _timedCallback();

  /** \brief Queries current trajectory and returns transform if valid */
  void _queryCallback(const std::shared_ptr<QueryTrajectory::Request> request,
                      std::shared_ptr<QueryTrajectory::Response> response);

  /** \brief Resets loss functions, cost terms, etc. at top of TDCP callback */
  void initializeProblem();

  /** \brief Clears window and starts over after major error occurs */
  void resetEstimator();

  /** \brief Store TDCP msg in our queue. Also checks times for dropped msgs */
  void addMsgToWindow(const cpo_interfaces::msg::TDCP::SharedPtr &msg);

  /** \brief Saves latest pose information to a CSV file for later analysis */
  void saveToFile(const lgmath::se3::Transformation &T_kg, double t_k) const;

  /** \brief Takes in Transforms and calls ROS publishers */
  void publishPose(const lgmath::se3::TransformationWithCovariance &T_ng);

  /** \brief Helper to convert lgmath Transform to ROS2 msg */
  static geometry_msgs::msg::PoseWithCovariance toPoseMsg(lgmath::se3::TransformationWithCovariance T);

  /** \brief Subscriber for TDCP msgs */
  rclcpp::Subscription<cpo_interfaces::msg::TDCP>::SharedPtr subscription_;

  /** \brief Publisher of integrated odometry in ENU frame */
  rclcpp::Publisher<geometry_msgs::msg::PoseWithCovariance>::SharedPtr
      enu_publisher_;

  /** \brief Whether to publish transform estimates at fixed rate (true) or after each msg received (false) */
  bool fixed_rate_publish_;

  /** \brief Optionally delay output of pose in timed callback so not always extrapolating */
  double publish_delay_;

  /** \brief Timer to publish odometry estimates at a fixed rate */
  rclcpp::TimerBase::SharedPtr publish_timer_;

  /** \brief Service to return the estimated relative transform between two times */
  rclcpp::Service<QueryTrajectory>::SharedPtr query_traj_srv_;

  /** \brief The fixed sensor-vehicle transform. Allows us to do estimation in the vehicle frame */
  steam::se3::FixedTransformEvaluator::ConstPtr tf_gps_vehicle_;

  /** \brief The STEAM optimizer */
  std::shared_ptr<steam::SolverBase> solver_;

  /** \brief The steam problem. */
  std::shared_ptr<steam::OptimizationProblem> problem_;

  /** \brief Cost terms associated with TDCP observations. */
  steam::ParallelizedCostTermCollection::Ptr tdcp_cost_terms_;

  /** \brief Cost terms associated with nonholonomic prior. */
  steam::ParallelizedCostTermCollection::Ptr nonholonomic_cost_terms_;

  /** \brief Cost terms associated with white-noise-on-acceleration motion prior */
  steam::ParallelizedCostTermCollection::Ptr smoothing_cost_terms_;

  /** \brief Cost term for prior on T_0g to help resolve roll, lock position */
  steam::ParallelizedCostTermCollection::Ptr pose_prior_cost_;

  /** \brief Loss function associated with TDCP costs */
  steam::LossFunctionBase::Ptr tdcp_loss_function_;

  /** \brief Loss function associated with nonholonomic costs */
  steam::LossFunctionBase::Ptr nonholonomic_loss_function_;

  /** \brief Loss function associated with roll prior */
  steam::LossFunctionBase::Ptr pp_loss_function_;

  /** \brief The steam trajectory, allows smoothing factors, velocity priors and pose extrapolation */
  std::shared_ptr<steam::se3::SteamTrajInterface> trajectory_;

  Eigen::Matrix<double, 6, 6> smoothing_factor_information_;

  Eigen::Matrix<double, 1, 1> tdcp_cov_;

  Eigen::Matrix<double, 4, 4> nonholonomic_cov_;

  Eigen::Matrix<double, 6, 6> pose_prior_cov_;

  /** \brief File path for saving a CSV file of our estimates */
  std::string results_path_;

  /** \brief Whether to print extra information about the optimization */
  bool steam_verbose_;

  /** \brief Whether to print the costs from each set of factors before/after optimizing */
  bool costs_verbose_;

  /** \brief Maximum number of solver iterations per optimization */
  uint steam_max_iterations_;

  /** \brief Our estimate of T_0g, stored to initialize the next optimization problem */
  lgmath::se3::Transformation init_pose_;

  /** \brief Keep track of whether we have T_0g estimate or need to get one from code solution */
  bool init_pose_estimated_ = false;

  /** \brief Store a window of TDCP messages and our state estimates. We use deque over queue to get access */
  std::deque<CpoEdge> edges_;

  /** \brief Size of the optimization window in msgs */
  uint window_size_;

  /** \brief If true, T_0g will be locked within a window once first estimated.
   *   Otherwise we'll apply a prior and re-estimate each loop.
   */
  bool lock_first_pose_;

  /** \brief Time period (in seconds) since last msg we will trust our odometry estimates */
  double traj_timeout_limit_;

  bool first_window_ = true;

};


/* Copyright Institute of Sound and Vibration Research - All rights reserved */

#include <librcl/hoa_allrad_gain_calculator.hpp>

#include <libvisr/signal_flow_context.hpp>

#include <libobjectmodel/object_vector.hpp>

#include <boost/test/unit_test.hpp>

#include <stdexcept>
#include <string>

namespace visr
{

// forward declaration
namespace panning
{
class LoudspeakerArray;
}
namespace pml
{
template<typename T>
class MatrixParameter;
}

namespace rcl
{
namespace test
{

using namespace objectmodel;

#if 0
namespace // unnamed
{
  class HoaCalc: public ril::AudioSignalFlow
  {
  public:
    HoaCalc( )
      : ril::AudioSignalFlow( 1024, 48000 )
      , mHoaCalc( *this, "HoaCalc" )
    {
    }

    void process( )
    {
      mHoaCalc.process( mObjVec, mPanningGains );
    }

    void setObjectVector( objectmodel::ObjectId id, objectmodel::Object const & obj );

  private:
    rcl::HoaAllRadGainCalculator mHoaCalc;

    objectmodel::ObjectVector mObjVec;

    efl::BasicMatrix<float> mPanningGains;

  };
}
#endif

BOOST_AUTO_TEST_CASE( InstantiateAllRad )
{
//  HoaCalc mFlow;
}

BOOST_AUTO_TEST_CASE( HoaAllradGainCalculatorSetup )
{
}

} // namespace test
} // namespace rcl
} // namespce visr


//
//  TerrainWorld.hpp
//  Kessler Syndrome
//
//  Created by Shamyl Zakariya on 3/1/17.
//
//

#ifndef TerrainWorld_hpp
#define TerrainWorld_hpp

#include <cinder/app/App.h>
#include <boost/functional/hash.hpp>
#include <unordered_set>

#include "core/Core.hpp"
#include "core/Signals.hpp"

namespace elements {
    namespace terrain {

        typedef boost::geometry::model::d2::point_xy<double> dpoint2;
        typedef boost::geometry::model::polygon<dpoint2> dpolygon2;

        SMART_PTR(World);

        SMART_PTR(GroupBase);

        SMART_PTR(StaticGroup);

        SMART_PTR(DynamicGroup);

        SMART_PTR(Drawable);

        SMART_PTR(Shape);

        SMART_PTR(Anchor);

        SMART_PTR(Element);

        SMART_PTR(Attachment);

        /**
         Edges whos vertices are within 1/POLY_EDGE_PRECISION distance of eachother are considered congruent, and thus snapped.
         */

#define POLY_EDGE_PRECISION 10.0

        /**
         poly_edge represents an edge in a PolyLine2d. It is not meant to represent the specific vertices in
         a coordinate system but rather it is meant to make it easy to determine if two PolyShapes share an edge.
         As such the coordinates a and b are integers, representing the world space coords of the PolyShape outer
         contour at a given stage of precision.
         */
        struct poly_edge {
            ivec2 a, b;

            poly_edge(dvec2 m, dvec2 n) {
                a.x = static_cast<int>(lround(m.x * POLY_EDGE_PRECISION));
                a.y = static_cast<int>(lround(m.y * POLY_EDGE_PRECISION));
                b.x = static_cast<int>(lround(n.x * POLY_EDGE_PRECISION));
                b.y = static_cast<int>(lround(n.y * POLY_EDGE_PRECISION));

                // we want consistent ordering
                if (a.x == b.x) {
                    if (b.y < a.y) {
                        swap(a, b);
                    }
                } else {
                    if (b.x < a.x) {
                        swap(a, b);
                    }
                }
            }

            friend bool operator==(const poly_edge &e0, const poly_edge &e1) {
                return e0.a == e1.a && e0.b == e1.b;
            }

            friend bool operator<(const poly_edge &e0, const poly_edge &e1) {
                if (e0.a.x != e1.a.x) {
                    return e0.a.x < e1.a.x;
                } else if (e0.a.y != e1.a.y) {
                    return e0.a.y < e1.a.y;
                } else if (e0.b.x != e1.b.x) {
                    return e0.b.x < e1.b.x;
                }
                return e0.b.y < e1.b.y;
            }

            friend std::ostream &operator<<(std::ostream &lhs, const poly_edge &rhs) {
                return lhs << "(poly_edge " << dvec2(rhs.a.x / POLY_EDGE_PRECISION, rhs.a.y / POLY_EDGE_PRECISION)
                << " : " << dvec2(rhs.b.x / POLY_EDGE_PRECISION, rhs.b.y / POLY_EDGE_PRECISION) << ")";
            }
        };

    } // namespace terrain
}

#pragma mark -

namespace std {

    template<>
    struct hash<elements::terrain::poly_edge> {

        // make island::poly_edge hashable
        size_t operator()(const elements::terrain::poly_edge &e) const {
            size_t seed = 0;
            boost::hash_combine(seed, e.a.x);
            boost::hash_combine(seed, e.a.y);
            boost::hash_combine(seed, e.b.x);
            boost::hash_combine(seed, e.b.y);
            return seed;
        }
    };

}

namespace elements {
    namespace terrain {

#pragma mark - Material

        /**
         Describes basic physics material properties for a collision shape
         */
        struct material {
            cpFloat density;
            cpFloat friction;
            cpFloat collisionShapeRadius;
            cpShapeFilter filter;
            cpCollisionType collisionType;
            double minSurfaceArea;
            Color color;

            material() :
            density(1),
            friction(1),
            collisionShapeRadius(0),
            filter({0, 0, 0}),
            collisionType(0),
            minSurfaceArea(0.1),
            color(0.2, 0.2, 0.2) {
            }

            material(cpFloat d, cpFloat f, cpFloat csr, cpShapeFilter flt, cpCollisionType ct, double msa, Color c) :
            density(d),
            friction(f),
            collisionShapeRadius(csr),
            filter(flt),
            collisionType(ct),
            minSurfaceArea(msa),
            color(c) {
            }

            material(const material &c) :
            density(c.density),
            friction(c.friction),
            collisionShapeRadius(c.collisionShapeRadius),
            filter(c.filter),
            collisionType(c.collisionType),
            minSurfaceArea(c.minSurfaceArea),
            color(c.color) {
            }
        };

#pragma mark - DrawDispatcher

        class DrawDispatcher {
        public:

            struct collector {
                set <DrawableRef> visible;
                vector <DrawableRef> sorted;

                void remove(const DrawableRef &s) {
                    visible.erase(s);
                    sorted.erase(std::remove(sorted.begin(), sorted.end(), s), sorted.end());
                }
            };

            DrawDispatcher();

            virtual ~DrawDispatcher();

            void add(const DrawableRef &);

            void remove(const DrawableRef &);

            void moved(const DrawableRef &);

            void moved(Drawable *);

            void cull(const core::render_state &);

            void draw(const core::render_state &, const gl::GlslProgRef &shader);

            size_t visibleCount() const {
                return _collector.sorted.size();
            }

            const set <DrawableRef> &getVisible() const {
                return _collector.visible;
            }

            const vector <DrawableRef> &getVisibleSorted() const {
                return _collector.sorted;
            }

        private:

            /**
             render a run of shapes belonging to a common group
             returns iterator to last shape drawn
             */
            vector<DrawableRef>::iterator _drawGroupRun(vector<DrawableRef>::iterator first, vector<DrawableRef>::iterator storageEnd, const core::render_state &state, const gl::GlslProgRef &shader);

        private:

            cpSpatialIndex *_index;
            set <DrawableRef> _all;
            collector _collector;

        };

#pragma mark - World

        /**
         @class World
         World "owns" and manages Group instances, which in turn own and manage Shape instances.
         */
        class World : public enable_shared_from_this<World> {
        public:

            static void loadSvg(DataSourceRef svgData, dmat4 transform,
                                vector <ShapeRef> &shapes, vector <AnchorRef> &anchors, vector <ElementRef> &elements,
                                bool flip = true);

            static void march(const Channel8u isoSurface, double isoLevel, dmat4 transform, vector <ShapeRef> &shapes);

            static void march(const Channel8u isoSurface, const Channel8u anchorIsoSurface, double isoLevel, dmat4 transform, vector <ShapeRef> &shapes, vector <AnchorRef> &anchors);

            /**
             Partitions shapes in `shapes to a grid, with chunks of size paritionSize.
             The purpose of this is simple: You might have a large stage. You want to divide the geometry of the stage into manageable
             chunks for visibility culling and collision/physics performance.
             */
            static vector <ShapeRef> partition(const vector <ShapeRef> &shapes, double partitionSize);

            // return a unique ID; used to generate IDs for drawables, groups, etc
            static size_t nextId() {
                return _idCounter++;
            }

        public:

            World(core::SpaceAccessRef space, material worldMaterial, material anchorMaterial);

            virtual ~World();

            /**
             Build a world of dynamic and static shapes. Any shape overlapping an Anchor will be static.
             Pieces cut off a static shape will become dynamic provided they don't overlap another anchor.
             */

            void build(const vector <ShapeRef> &shapes,
                       const vector <AnchorRef> &anchors = vector<AnchorRef>(),
                       const vector <ElementRef> &elements = vector<ElementRef>());

            /**
             Perform a cut in world space from a to b, with half-thickness of radius.
             Discard any resultant geometry with less than minSurfaceArea
             */
            void cut(dvec2 a, dvec2 b, double radius, double minSurfaceArea = 1);

            /**
             Perform a cut in world space removing any geometry which overlaps the enclosed space of `polygonShape`
             Discard any resultant geometry with less than minSurfaceArea
             */
            void cut(const dpolygon2 &polygonShape, cpBB polygonShapeWorldBounds = cpBBInvalid, double minSurfaceArea = 0);

            void draw(const core::render_state &renderState);

            void step(const core::time_state &timeState);

            void update(const core::time_state &timeState);

            const set <DynamicGroupRef> &getDynamicGroups() const {
                return _dynamicGroups;
            }

            StaticGroupRef getStaticGroup() const {
                return _staticGroup;
            }

            const vector <AnchorRef> &getAnchors() const {
                return _anchors;
            }

            const vector <ElementRef> &getElements() const {
                return _elements;
            }

            ElementRef getElementById(string id) const;

            const set <AttachmentRef> &getOrphanedAttachments() const {
                return _orphanedAttachments;
            }

            core::SpaceAccessRef getSpace() const {
                return _space;
            }

            void setWorldMaterial(material wm) {
                _worldMaterial = wm;
            }

            material getWorldMaterial() const {
                return _worldMaterial;
            }

            void setAnchorMaterial(material am) {
                _anchorMaterial = am;
            }

            material getAnchorMaterial() const {
                return _anchorMaterial;
            }

            DrawDispatcher &getDrawDispatcher() {
                return _drawDispatcher;
            }

            const DrawDispatcher &getDrawDispatcher() const {
                return _drawDispatcher;
            }

            typedef function<bool(const DynamicGroupRef &)> DynamicGroupVisitor;

            /**
             Perform a culling of dynamic groups with less than a given surface area.
             All geometry of less than `minSurfaceArea will be collected and sorted,
             and `portion amount of that geometry (0 (none) to 1 (all)) will be discarded.
             Optional test argument is a functor called for each group smaller than minSurfaceArea to decide if it should be in the cull set.
             Returns number of groups culled.
             */
            size_t cullDynamicGroups(double minSurfaceArea, double portion = 1, const DynamicGroupVisitor &test = DynamicGroupVisitor());

            /**
             Gather DynamicGroups that have been sleeping for at least `minSleepTime (and if `test is provided, winnow to those who
             pass the test), then take the longest sleeping portion of that set and make them part of the static group.
             Think of this as a way of saying, "I want the rubble which has been still for a long time to become static environment"
             */
            size_t makeSleepingDynamicGroupsStatic(core::seconds_t minSleepTime, double portion = 1, const DynamicGroupVisitor &test = DynamicGroupVisitor());

            /**
             Get the Object which wraps this World for use in a Stage
             */
            core::ObjectRef getObject() const;

            /**
             Attempt to place an attachment in the World. Returns true if the position corresponded to solid geometry
             and the attachment was able to be parented to it.
             */
            bool addAttachment(const AttachmentRef &attachment, dvec2 worldPosition, double angle = 0);

            bool addAttachment(const AttachmentRef &attachment, dvec2 worldPosition, dvec2 rotation);

        protected:

            friend class TerrainObject;

            friend class Shape;

            friend class DynamicGroup;

            bool tryAddAttachment(const AttachmentRef &attachment, const GroupBaseRef &group, dvec2 worldPosition, dvec2 worldRotation);

            void addAttachment(const AttachmentRef &attachment, const GroupBaseRef &group, const ShapeRef shapeHint, dvec2 worldPosition, dvec2 worldRotation);

            void handleOrphanedAttachment(const AttachmentRef &attachment);

            void notifyCollisionShapesWillBeDestoyed(vector<cpShape *> shapes);

            void notifyBodyWillBeDestoyed(cpBody *body);

            void setObject(core::ObjectRef object);

            /**
             shapes: vector of shapes which were either newly created by a cut or are "neighbors" to cut shapes.
             parentage: mapping of shapes in `shapes to groups
             attachmentsToReparent: vector of attachments which were associated with cut shapes, which need new parentage
             */
            void build(const vector <ShapeRef> &shapes, const map <ShapeRef, GroupBaseRef> &parentage, vector <AttachmentRef> attachmentsToReparent);

            vector <set<ShapeRef>> findShapeGroups(const vector <ShapeRef> &shapes, const map <ShapeRef, GroupBaseRef> &parentage);

            bool isShapeGroupStatic(const set <ShapeRef> shapeGroup, const GroupBaseRef &parentGroup);

        private:

            static size_t _idCounter;

            material _worldMaterial, _anchorMaterial;
            core::SpaceAccessRef _space;
            StaticGroupRef _staticGroup;
            set <DynamicGroupRef> _dynamicGroups;
            vector <AnchorRef> _anchors;
            vector <ElementRef> _elements;
            set <AttachmentRef> _orphanedAttachments;
            map <string, ElementRef> _elementsById;

            DrawDispatcher _drawDispatcher;
            gl::GlslProgRef _shader;

            core::ObjectWeakRef _object;

            // state to speed up adding attachments, to reduce lookup for neighboring attachment addition
            GroupBaseWeakRef _lastAttachmentGroup;
            ShapeWeakRef _lastAttachmentShape;
        };

#pragma mark - Attachment

        class Attachment {
        public:

            // signal fired when this attachment is orphaned from a group and moved from Group ownership to the World's orphan set
            // passes attachment's id and tag
            core::signals::signal<void(size_t /*id*/,size_t /*tag*/)> onOrphaned;

        public:

            // create an unparented Attachment. Position is assigned when calling World::addAttachment
            Attachment();
            virtual ~Attachment();

            // unique Id of this attachment
            size_t getId() const { return _id; }

            // get the Group this Attachment is attached to
            GroupBaseRef getGroup() const { return _group.lock(); }

            // return true if this Attachment has become detached from a Group and has become orphaned
            bool isOrphaned() const { return _orphaned; }

            // get transform (position/rotation) of this Attachment relative to its parent group
            dmat4 getLocalTransform() const { return _localTransform; }

            // get transform (position/rotation) of this attachment relative to world
            dmat4 getWorldTransform() const { return _worldTransform; }

            // get position of this attachment in coordinate space of its group
            dvec2 getLocalPosition() const;

            // get rotation of this attachment in coordinate space of its group
            dvec2 getLocalRotation() const;

            // get position of this attachment in world
            dvec2 getWorldPosition() const;

            // get rotation of this attachment in world
            dvec2 getWorldRotation() const;

            // To "kill" an attachment, call setFinished(true); it will be detached from its group
            // or if it's an orphan, will be detached from the world's orphan set. It will no longer
            // be updated and if no strong references are held, it will be destroyed.
            void setFinished(bool finished=true);

            // return true if this attachment is scheduled to be let go next timestep
            bool isFinished() const { return _finished; }

            // get the assigned tag value for this attachment
            size_t getTag() const { return _tag; }

            // set a tag value for this attachment. You can use this to assign custom data, like index or bitmasks, whatever.
            void setTag(size_t tag) { _tag = tag; }

            // return the step index of the last time this was moved.
            // this is updated at step() time, so any code monitoring Attachments in update() can
            // safely say: if (attachment->getStepIndexForLastPositionUpdate() == timeState.step) { /*whatever*/ }
            // always check for equality, not >=, <= etc.
            size_t getStepIndexForLastPositionUpdate() const { return _lastMovedAtStep; }

            ShapeRef getShapeHint() const { return _shapeHint.lock(); }

        private:

            friend class World;
            friend class GroupBase;
            friend class StaticGroup;
            friend class DynamicGroup;

            // called by World::addAttachment; returns true iff group != previous _group
            bool configure(const GroupBaseRef &group, dvec2 position, dvec2 rotation);

        private:

            size_t _id, _tag, _lastMovedAtStep;
            dmat4 _localTransform;
            dmat4 _worldTransform;
            GroupBaseWeakRef _group;
            GroupBase *_groupUnsafePtr;
            bool _finished, _orphaned;

            // shape hint is the shape that "contains" this attachment; used at
            // runtime to speed lookup when terrain destruction requires reassignment
            ShapeWeakRef _shapeHint;

        };

#pragma mark - GroupBase

        class GroupBase : public core::IChipmunkUserData {
        public:

            GroupBase(WorldRef world, material m, DrawDispatcher &dispatcher);

            virtual ~GroupBase();

            DrawDispatcher &getDrawDispatcher() const {
                return _drawDispatcher;
            }

            const core::SpaceAccessRef &getSpace() const {
                return _space;
            }

            const material &getMaterial() const {
                return _material;
            }

            virtual string getName() const {
                return _name;
            }

            virtual Color getColor(const core::render_state &state) const {
                return _material.color;
            }

            virtual cpHashValue getHash() const {
                return _hash;
            }

            virtual bool isStatic() const = 0;

            virtual bool isDynamic() const { return !isStatic(); }

            virtual cpBody *getBody() const = 0;

            virtual cpBB getBB() const = 0;

            virtual dmat4 getModelMatrix() const = 0;

            virtual dmat4 getInverseModelMatrix() const = 0;

            virtual dvec2 getPosition() const = 0;

            virtual double getAngle() const = 0;

            virtual double getSurfaceArea() const = 0;

            virtual set <ShapeRef> getShapes() const = 0;

            virtual void releaseShapes() = 0;

            virtual size_t getDrawingBatchId() const {
                return _drawingBatchId;
            }

            virtual void step(const core::time_state &timeState);

            virtual void update(const core::time_state &timeState);

            cpTransform getModelTransform() const {
                dmat4 mm = getModelMatrix();
                return cpTransform {mm[0].x, mm[0].y, mm[1].x, mm[1].y, mm[3].x, mm[3].y};
            }

            cpTransform getInverseModelTransform() const {
                dmat4 mvi = getInverseModelMatrix();
                return cpTransform {mvi[0].x, mvi[0].y, mvi[1].x, mvi[1].y, mvi[3].x, mvi[3].y};
            }

            WorldRef getWorld() const;

            // return shape which contains a given point in this group's local coordinate space
            virtual ShapeRef findShapeContainingLocalPoint(const dvec2 lp) const = 0;

            // return shape which contains a given point in the world coordinate space
            virtual ShapeRef findShapeContainingWorldPoint(const dvec2 wp) const { return findShapeContainingLocalPoint(getInverseModelMatrix() * wp); };

            const set <AttachmentRef> &getAttachments() const { return _attachments; }

            void clearAttachments() { _attachments.clear(); }

            // add an attachment to this Group's attachment set
            void addAttachment(const AttachmentRef &a) {
                _attachments.insert(a);
            }

            // remove an attachment from this Group's attachment set, returning true if it was removed
            virtual bool removeAttachment(const AttachmentRef &a) {
                if (_attachments.erase(a) > 0) {
                    a->_group.reset();
                    return true;
                }
                return false;
            }

            // IChipmunkUserData
            core::ObjectRef getObject() const override;

        protected:

            DrawDispatcher &_drawDispatcher;
            size_t _drawingBatchId;
            WorldWeakRef _world;
            core::SpaceAccessRef _space;
            set<AttachmentRef> _attachments;
            material _material;
            string _name;
            cpHashValue _hash;

        };

#pragma mark - StaticGroup

        class StaticGroup : public GroupBase, public enable_shared_from_this<StaticGroup> {
        public:
            StaticGroup(WorldRef world, material m, DrawDispatcher &dispatcher);

            ~StaticGroup();

            virtual bool isStatic() const override { return true; }

            cpBody *getBody() const override {
                return _body;
            }

            cpBB getBB() const override;

            dmat4 getModelMatrix() const override {
                return dmat4(1);
            }

            dmat4 getInverseModelMatrix() const override {
                return dmat4(1);
            }

            dvec2 getPosition() const override {
                return dvec2(0);
            }

            double getAngle() const override {
                return 0;
            }

            double getSurfaceArea() const override {
                return _surfaceArea;
            }

            set <ShapeRef> getShapes() const override {
                return _shapes;
            }

            void releaseShapes() override;

            ShapeRef findShapeContainingLocalPoint(const dvec2 lp) const override;

            void addShape(ShapeRef shape, double minShapeArea);

            void removeShape(ShapeRef shape);

        protected:
            cpBody *_body;
            set <ShapeRef> _shapes;
            mutable cpBB _worldBB;
            double _surfaceArea;
        };

#pragma mark - DynamicGroup

        class DynamicGroup : public GroupBase, public enable_shared_from_this<DynamicGroup> {
        public:
            DynamicGroup(WorldRef world, material m, DrawDispatcher &dispatcher);

            ~DynamicGroup();

            // GroupBase
            string getName() const override;

            Color getColor(const core::render_state &state) const override;

            virtual bool isStatic() const override { return false; }

            cpBody *getBody() const override {
                return _body;
            }

            cpBB getBB() const override {
                return _worldBB;
            }

            dmat4 getModelMatrix() const override {
                return _modelMatrix;
            }

            dmat4 getInverseModelMatrix() const override {
                return _inverseModelMatrix;
            }

            dvec2 getPosition() const override {
                return v2(_position);
            }

            double getAngle() const override {
                return static_cast<double>(_angle);
            }

            double getSurfaceArea() const override {
                return _surfaceArea;
            }

            set <ShapeRef> getShapes() const override {
                return _shapes;
            }

            void releaseShapes() override;

            void step(const core::time_state &timeState) override;
            void update(const core::time_state &timeState) override;

            ShapeRef findShapeContainingLocalPoint(const dvec2 lp) const override;

            // DynamicGroup
            // return true iff the physics body is sleeping
            bool isSleeping() const;

            // return the number of seconds that the body has been sleeping, or < 0 if it is awake
            core::seconds_t getSleepDuration() const {
                return _sleepDuration;
            }

        protected:
            friend class World;

            bool build(set <ShapeRef> shapes, const GroupBaseRef &parentGroup, double minShapeArea);

            // update position and angle to match underlying cpBody, returning true if we moved
            bool syncToCpBody();

        private:

            cpBody *_body;
            cpVect _position;
            cpFloat _angle;
            double _surfaceArea;
            cpBB _worldBB, _modelBB;
            dmat4 _modelMatrix, _inverseModelMatrix;
            core::seconds_t _sleepDuration;

            set <ShapeRef> _shapes;
        };

#pragma mark - Drawable

        class Drawable : public core::IChipmunkUserData, public enable_shared_from_this<Drawable> {
        public:
            Drawable();

            virtual ~Drawable();

            size_t getId() const {
                return _id;
            }

            virtual cpBB getBB() const = 0;

            virtual size_t getDrawingBatchId() const = 0;

            virtual size_t getLayer() const = 0;

            virtual dmat4 getModelMatrix() const = 0;

            virtual double getAngle() const = 0;

            virtual dvec2 getModelCentroid() const = 0;

            virtual const TriMeshRef &getTriMesh() const = 0;

            virtual const gl::VboMeshRef &getVboMesh() const = 0;

            virtual Color getColor(const core::render_state &state) const = 0;

            virtual bool shouldDraw(const core::render_state &state) const = 0;

            // get typed shared_from_this, e.g., shared_ptr<Shape> = shared_from_this_as<Shape>();
            template<typename T>
            shared_ptr<T const> shared_from_this_as() const {
                return static_pointer_cast<T const>(shared_from_this());
            }

            // get typed shared_from_this, e.g., shared_ptr<Shape> = shared_from_this_as<Shape>();
            template<typename T>
            shared_ptr <T> shared_from_this_as() {
                return static_pointer_cast<T>(shared_from_this());
            }

            void setWorld(WorldRef w);

            WorldRef getWorld() const;

            // IChipmunkUserData
            core::ObjectRef getObject() const override;

        private:

            size_t _id;
            WorldWeakRef _world;

        };

#pragma mark - Element

        /**
         @class Element
         An element is a shape specified in the <g id="elements">...</g> section of a stage SVG file.
         Each element represents simply a shape in world space with an associated ID. Elements can be used
         as location markers for enemy spawn points, or possibly as trigger geometry.
         */
        class Element : public Drawable {
        public:

            static const size_t DRAWING_BATCH_ID = 0;
            static const size_t LAYER = 2;

            static ElementRef fromContour(string id, const PolyLine2d &contour) {
                return make_shared<Element>(id, contour);
            }

        public:

            Element(string id, const PolyLine2d &contour);

            virtual ~Element();

            // Element
            string getId() const {
                return _id;
            }

            dvec2 getPosition() const;

            // Drawable
            cpBB getBB() const override {
                return _bb;
            }

            size_t getDrawingBatchId() const override {
                return DRAWING_BATCH_ID;
            }

            size_t getLayer() const override {
                return LAYER;
            }

            dmat4 getModelMatrix() const override {
                return dmat4(1);
            } // identity

            double getAngle() const override {
                return 0;
            };

            dvec2 getModelCentroid() const override;

            const TriMeshRef &getTriMesh() const override {
                return _trimesh;
            }

            const gl::VboMeshRef &getVboMesh() const override {
                return _vboMesh;
            }

            Color getColor(const core::render_state &state) const override {
                return Color(1, 0, 1);
            }

            bool shouldDraw(const core::render_state &state) const override;

        private:

            cpBB _bb;
            string _id;
            TriMeshRef _trimesh;
            gl::VboMeshRef _vboMesh;

        };

#pragma mark - Anchor

        /**
         @class Anchor
         An Anchor is a single contour in world space which "anchors" shapes to be static and not dynamic bodies.
         A body is static if one of its shapes overlaps an anchor, and if the body's parentage has never been dynamic.
         This is to say, once a shape is severed from a static body and becomes dynamic, it can never be static again.
         */
        class Anchor : public Drawable {
        public:

            static const size_t DRAWING_BATCH_ID = 0;
            static const size_t LAYER = 1;

            static AnchorRef fromContour(const PolyLine2d &contour) {
                return make_shared<Anchor>(contour);
            }

            static vector <AnchorRef> fromContours(const vector <PolyLine2d> &contours) {
                vector<AnchorRef> anchors;
                anchors.reserve(contours.size());
                for (auto c : contours) {
                    anchors.push_back(make_shared<Anchor>(c));
                }
                return anchors;
            }

            static AnchorRef fromShape(const Shape2d &shape);

        public:

            Anchor(const PolyLine2d &contour);

            virtual ~Anchor();

            const PolyLine2d &getContour() const {
                return _contour;
            }

            // Drawable
            cpBB getBB() const override {
                return _bb;
            }

            size_t getDrawingBatchId() const override {
                return DRAWING_BATCH_ID;
            }

            size_t getLayer() const override {
                return LAYER;
            }

            // anchors are static, so always returns identity
            dmat4 getModelMatrix() const override {
                return dmat4(1);
            }

            // anchors are static, no rotation
            double getAngle() const override {
                return 0;
            };

            dvec2 getModelCentroid() const override;

            const TriMeshRef &getTriMesh() const override {
                return _trimesh;
            }

            const gl::VboMeshRef &getVboMesh() const override {
                return _vboMesh;
            }

            Color getColor(const core::render_state &state) const override {
                return _material.color;
            }

            bool shouldDraw(const core::render_state &state) const override {
                return true;
            }

        protected:

            friend class World;

            bool build(core::SpaceAccessRef space, material m);

        private:

            cpBody *_staticBody;
            vector<cpShape *> _shapes;

            cpBB _bb;
            material _material;
            PolyLine2d _contour;

            TriMeshRef _trimesh;
            gl::VboMeshRef _vboMesh;

        };

#pragma mark - Shape

        class Shape : public Drawable {
        public:

            static const size_t LAYER = 0;

            struct contour_pair {
                PolyLine2d world, model;

                // when constructed model and world are both equal
                contour_pair(const PolyLine2d &w) :
                world(w),
                model(w) {
                }
            };

            static ShapeRef fromContour(const PolyLine2d &outerContour);

            static vector <ShapeRef> fromContours(const vector <PolyLine2d> &contourSoup);

            static vector <ShapeRef> fromShapes(const vector <Shape2d> &shapeSoup);

        public:

            Shape(const PolyLine2d &shapeContour);

            Shape(const PolyLine2d &shapeContour, const std::vector<PolyLine2d> &holeContours);

            virtual ~Shape();

            const contour_pair &getOuterContour() const {
                return _outerContour;
            }

            const vector <contour_pair> &getHoleContours() const {
                return _holeContours;
            }

            GroupBaseRef getGroup() const {
                return _group.lock();
            }

            cpHashValue getGroupHash() const {
                return _groupHash;
            }

            dmat4 getModelMatrix() const override {
                if (GroupBaseRef group = _group.lock()) {
                    return group->getModelMatrix();
                } else {
                    return dmat4();
                }
            }

            dmat4 getInverseModelMatrix() const {
                if (GroupBaseRef group = _group.lock()) {
                    return group->getInverseModelMatrix();
                } else {
                    return dmat4();
                }
            }

            const vector<cpShape *> &getShapes(cpBB &shapesModelBB) const {
                shapesModelBB = _shapesModelBB;
                return _shapes;
            }

            bool hasValidTriMesh() const {
                return _trimesh && _trimesh->getNumTriangles() > 0;
            }

            cpBB getBB() const override;

            size_t getDrawingBatchId() const override {
                return _groupDrawingBatchId;
            }

            size_t getLayer() const override {
                return LAYER;
            }

            double getAngle() const override {
                return getGroup()->getAngle();
            }

            dvec2 getModelCentroid() const override {
                return _modelCentroid;
            }

            const TriMeshRef &getTriMesh() const override {
                return _trimesh;
            }

            const gl::VboMeshRef &getVboMesh() const override {
                return _vboMesh;
            }

            Color getColor(const core::render_state &state) const override {
                return getGroup()->getColor(state);
            }

            bool shouldDraw(const core::render_state &state) const override {
                return true;
            }

            double getSurfaceArea();

            const unordered_set<poly_edge> &getWorldSpaceContourEdges();

            cpBB getWorldSpaceContourEdgesBB();

            vector <ShapeRef> subtract(const dpolygon2 &polygonToSubtract) const;

            // return true iff the point localPoint (in this shape's group's coordinate space) is inside this shape
            bool isLocalPointInside(dvec2 localPoint);

        protected:

            friend class StaticGroup;

            friend class DynamicGroup;

            friend class World;

            void updateWorldSpaceContourAndBB();

            void setGroup(GroupBaseRef group);

            // build the trimesh, returning true iff we got > 0 triangles
            bool triangulate();

            void computeMassAndMoment(double density, double &mass, double &moment, double &area);

            void destroyCollisionShapes();

            const vector<cpShape *> &createCollisionShapes(cpBody *body, cpBB &shapesModelBB, double collisionShapeRadius);

        private:

            bool _worldSpaceShapeContourEdgesDirty;
            contour_pair _outerContour;
            vector <contour_pair> _holeContours;
            dvec2 _modelCentroid;

            cpBB _shapesModelBB;
            vector<cpShape *> _shapes;
            GroupBaseWeakRef _group;
            size_t _groupDrawingBatchId;
            cpHashValue _groupHash;

            unordered_set<poly_edge> _worldSpaceContourEdges;
            cpBB _worldSpaceContourEdgesBB;

            TriMeshRef _trimesh;
            gl::VboMeshRef _vboMesh;

            // the attachments which are anchored by being in this shape's geometry
            set <AttachmentRef> _attachments;
        };

    }
} // end namespace elements::terrain

#endif /* TerrainWorld_hpp */


#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <vector>
#include <map>
#include <Eigen/Geometry>
#include <arc_utilities/arc_helpers.hpp>
#include <arc_utilities/eigen_helpers.hpp>

#ifndef ABB_IRB1600_145_FK_FAST_HPP
#define ABB_IRB1600_145_FK_FAST_HPP

namespace ABB_IRB1600_145_FK_FAST
{
    const size_t ABB_IRB1600_145_NUM_ACTIVE_JOINTS = 6;
    const size_t ABB_IRB1600_145_NUM_LINKS = 8;

    const std::string ABB_IRB1600_145_ACTIVE_JOINT_1_NAME = "joint_1";
    const std::string ABB_IRB1600_145_ACTIVE_JOINT_2_NAME = "joint_2";
    const std::string ABB_IRB1600_145_ACTIVE_JOINT_3_NAME = "joint_3";
    const std::string ABB_IRB1600_145_ACTIVE_JOINT_4_NAME = "joint_4";
    const std::string ABB_IRB1600_145_ACTIVE_JOINT_5_NAME = "joint_5";
    const std::string ABB_IRB1600_145_ACTIVE_JOINT_6_NAME = "joint_6";

    const std::string ABB_IRB1600_145_LINK_1_NAME = "link_0";
    const std::string ABB_IRB1600_145_LINK_2_NAME = "link_1";
    const std::string ABB_IRB1600_145_LINK_3_NAME = "link_2";
    const std::string ABB_IRB1600_145_LINK_4_NAME = "link_3";
    const std::string ABB_IRB1600_145_LINK_5_NAME = "link_4";
    const std::string ABB_IRB1600_145_LINK_6_NAME = "link_5";
    const std::string ABB_IRB1600_145_LINK_7_NAME = "link_6";
    const std::string ABB_IRB1600_145_LINK_8_NAME = "link_7";

    typedef std::vector<Eigen::Isometry3d, Eigen::aligned_allocator<Eigen::Isometry3d>> VectorIsometry3d;

    inline Eigen::Isometry3d Get_base_joint1_LinkJointTransform(const double joint_val)
    {
        Eigen::Translation3d pre_joint_translation(0.0, 0.0, 0.1245);
        Eigen::Quaterniond pre_joint_rotation(1.0, 0.0, 0.0, 0.0);
        Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
        Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
        Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitZ()));
        Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
        return (pre_joint_transform * joint_transform);
    }

    inline Eigen::Isometry3d Get_link_1_joint_2_LinkJointTransform(const double joint_val)
    {
        Eigen::Translation3d pre_joint_translation(0.15, -0.1395, 0.362);
        Eigen::Quaterniond pre_joint_rotation(1.0, 0.0, 0.0, 0.0);
        Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
        Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
        Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitY()));
        Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
        return (pre_joint_transform * joint_transform);
    }

    inline Eigen::Isometry3d Get_link_2_joint_3_LinkJointTransform(const double joint_val)
    {
        Eigen::Translation3d pre_joint_translation(0.0, 0.028, 0.7);
        Eigen::Quaterniond pre_joint_rotation(1.0, 0.0, 0.0, 0.0);
        Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
        Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
        Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitY()));
        Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
        return (pre_joint_transform * joint_transform);
    }

    inline Eigen::Isometry3d Get_link_3_joint_4_LinkJointTransform(const double joint_val)
    {
        Eigen::Translation3d pre_joint_translation(0.314, 0.107, 0.0);
        Eigen::Quaterniond pre_joint_rotation(1.0, 0.0, 0.0, 0.0);
        Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
        Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
        Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitX()));
        Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
        return (pre_joint_transform * joint_transform);
    }

    inline Eigen::Isometry3d Get_link_4_joint_5_LinkJointTransform(const double joint_val)
    {
        Eigen::Translation3d pre_joint_translation(0.286, 0.0, 0.0);
        Eigen::Quaterniond pre_joint_rotation(1.0, 0.0, 0.0, 0.0);
        Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
        Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
        Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitY()));
        Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
        return (pre_joint_transform * joint_transform);
    }

    inline Eigen::Isometry3d Get_link_5_joint_6_LinkJointTransform(const double joint_val)
    {
        Eigen::Translation3d pre_joint_translation(0.0, 0.0, 0.0);
        Eigen::Quaterniond pre_joint_rotation(1.0, 0.0, 0.0, 0.0);
        Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
        Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
        Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitX()));
        Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
        return (pre_joint_transform * joint_transform);
    }

    inline Eigen::Isometry3d Get_Fixed_link_6_joint_tool_LinkJointTransform(void)
    {
        Eigen::Translation3d pre_joint_translation(0.065, 0.0, 0.0);
        Eigen::Quaterniond pre_joint_rotation(0.7071067811865476, 0.0, 0.7071067811865476, 0.0);
        Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
        return pre_joint_transform;
    }

    inline VectorIsometry3d GetLinkTransforms(const std::vector<double>& configuration, const Eigen::Isometry3d& base_transform=Eigen::Isometry3d::Identity())
    {
        assert(configuration.size() == ABB_IRB1600_145_NUM_ACTIVE_JOINTS);
        VectorIsometry3d link_transforms(ABB_IRB1600_145_NUM_LINKS);
        link_transforms[0] = base_transform;
        link_transforms[1] = link_transforms[0] * Get_base_joint1_LinkJointTransform(configuration[0]);
        link_transforms[2] = link_transforms[1] * Get_link_1_joint_2_LinkJointTransform(configuration[1]);
        link_transforms[3] = link_transforms[2] * Get_link_2_joint_3_LinkJointTransform(configuration[2]);
        link_transforms[4] = link_transforms[3] * Get_link_3_joint_4_LinkJointTransform(configuration[3]);
        link_transforms[5] = link_transforms[4] * Get_link_4_joint_5_LinkJointTransform(configuration[4]);
        link_transforms[6] = link_transforms[5] * Get_link_5_joint_6_LinkJointTransform(configuration[5]);
        link_transforms[7] = link_transforms[6] * Get_Fixed_link_6_joint_tool_LinkJointTransform();
        return link_transforms;
    }

    inline VectorIsometry3d GetLinkTransforms(std::map<std::string, double> configuration, const Eigen::Isometry3d& base_transform=Eigen::Isometry3d::Identity())
    {
        std::vector<double> configuration_vector(ABB_IRB1600_145_NUM_ACTIVE_JOINTS);
        configuration_vector[0] = configuration[ABB_IRB1600_145_ACTIVE_JOINT_1_NAME];
        configuration_vector[1] = configuration[ABB_IRB1600_145_ACTIVE_JOINT_2_NAME];
        configuration_vector[2] = configuration[ABB_IRB1600_145_ACTIVE_JOINT_3_NAME];
        configuration_vector[3] = configuration[ABB_IRB1600_145_ACTIVE_JOINT_4_NAME];
        configuration_vector[4] = configuration[ABB_IRB1600_145_ACTIVE_JOINT_5_NAME];
        configuration_vector[5] = configuration[ABB_IRB1600_145_ACTIVE_JOINT_6_NAME];
        return GetLinkTransforms(configuration_vector, base_transform);
    }

    inline EigenHelpers::MapStringIsometry3d GetLinkTransformsMap(const std::vector<double>& configuration, const Eigen::Isometry3d& base_transform=Eigen::Isometry3d::Identity())
    {
        const EigenHelpers::VectorIsometry3d link_transforms = GetLinkTransforms(configuration, base_transform);
        EigenHelpers::MapStringIsometry3d link_transforms_map;
        link_transforms_map[ABB_IRB1600_145_LINK_1_NAME] = link_transforms[0];
        link_transforms_map[ABB_IRB1600_145_LINK_2_NAME] = link_transforms[1];
        link_transforms_map[ABB_IRB1600_145_LINK_3_NAME] = link_transforms[2];
        link_transforms_map[ABB_IRB1600_145_LINK_4_NAME] = link_transforms[3];
        link_transforms_map[ABB_IRB1600_145_LINK_5_NAME] = link_transforms[4];
        link_transforms_map[ABB_IRB1600_145_LINK_6_NAME] = link_transforms[5];
        link_transforms_map[ABB_IRB1600_145_LINK_7_NAME] = link_transforms[6];
        link_transforms_map[ABB_IRB1600_145_LINK_8_NAME] = link_transforms[7];
        return link_transforms_map;
    }

    inline EigenHelpers::MapStringIsometry3d GetLinkTransformsMap(const std::map<std::string, double>& configuration, const Eigen::Isometry3d& base_transform=Eigen::Isometry3d::Identity())
    {
        const EigenHelpers::VectorIsometry3d link_transforms = GetLinkTransforms(configuration, base_transform);
        EigenHelpers::MapStringIsometry3d link_transforms_map;
        link_transforms_map[ABB_IRB1600_145_LINK_1_NAME] = link_transforms[0];
        link_transforms_map[ABB_IRB1600_145_LINK_2_NAME] = link_transforms[1];
        link_transforms_map[ABB_IRB1600_145_LINK_3_NAME] = link_transforms[2];
        link_transforms_map[ABB_IRB1600_145_LINK_4_NAME] = link_transforms[3];
        link_transforms_map[ABB_IRB1600_145_LINK_5_NAME] = link_transforms[4];
        link_transforms_map[ABB_IRB1600_145_LINK_6_NAME] = link_transforms[5];
        link_transforms_map[ABB_IRB1600_145_LINK_7_NAME] = link_transforms[6];
        link_transforms_map[ABB_IRB1600_145_LINK_8_NAME] = link_transforms[7];
        return link_transforms_map;
    }
}

#endif // ABB_IRB1600_145_FK_FAST_HPP


/**
 * Copyright (c) 2009 Carnegie Mellon University.
 *     All rights reserved.
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing,
 *  software distributed under the License is distributed on an "AS
 *  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 *  express or implied.  See the License for the specific language
 *  governing permissions and limitations under the License.
 *
 *
 */

#ifndef PGIBBS_JT_SPLASH_SAMPLER_HPP
#define PGIBBS_JT_SPLASH_SAMPLER_HPP

#include <iostream>
#include <fstream>
#include <vector>
#include <map>
#include <set>
#include <deque>
#include <string>
#include <cassert>
#include <algorithm>

#include <boost/unordered_set.hpp>

// Including Standard Libraries

#include <graphlab/parallel/pthread_tools.hpp>
#include <graphlab/util/timer.hpp>

#include <graphlab.hpp>

#include "factorized_model.hpp"
#include "mrf.hpp"
#include "junction_tree.hpp"

/**
 * The settings for the jt_splash_sampler.  Originally these formed a
 * long list of arguments but since the order can easily introduce
 * bugs we switched to a struct.
 */
struct splash_settings {
  size_t ntrees;
  size_t max_tree_size;
  size_t max_tree_height;
  size_t max_tree_width;
  size_t max_factor_size;
  bool   priorities;
  size_t vanish_updates;
  size_t subthreads;

  splash_settings() :
    ntrees(2),
    max_tree_size(std::numeric_limits<size_t>::max()),
    max_tree_height(std::numeric_limits<size_t>::max()),
    max_tree_width(2),
    max_factor_size(std::numeric_limits<size_t>::max()),
    priorities(false),
    vanish_updates(10),
    subthreads(1) { }
};

/**
 * Run the jtsplash sampler
 *
 */
void run_jtsplash_sampler(mrf_graph_type& mrf_graph,
                          const std::string& jtsplash_results_fn,
                          const std::vector<double>& runtimes,
                          const bool draw_images,
                          const splash_settings& settings);

/**
 * This fairly complex update function assembles the clique factors by
 * conditioning on variables not in the tree.  Then it computes
 * messages at each clique to calibrate the junction tree.  Finally,
 * using the messages and the conditioned parents, it samples each
 * clique constructing new assignments to each variable.
 */
class jtree_update :
  public graphlab::iupdate_functor<jtree_graph_type, jtree_update> {
public:
  typedef graphlab::iupdate_functor<jtree_graph_type, jtree_update> base;
  jtree_update(mrf_graph_type* mrf_ptr = NULL) : mrf_ptr(mrf_ptr) { }
  mrf_graph_type* mrf_ptr;
  void operator()(base::icontext_type& context);
}; // end of class jtree_update

// Termination management
struct termination_condition {
  bool   error;
  float  finish_time_seconds;
  size_t target_nsamples;
  size_t target_ntrees;
  graphlab::atomic<size_t> atomic_nsamples;
  graphlab::atomic<size_t> atomic_ntrees;
  termination_condition();
  bool finished() const;
  void reset();
};

//! Predecleration
//! The jt worker executes splashes sequential within each thread.
class jt_builder :
  public graphlab::iupdate_functor<mrf_graph_type, jt_builder>{
public:
  typedef graphlab::iupdate_functor<mrf_graph_type, jt_builder> base;

  struct splash_state {
    size_t worker_id;
    splash_settings settings;
    // Tree building data structures
    size_t root_index;
    const std::vector<vertex_id_t>* root_perm_ptr;
    vertex_id_t current_root;
    //! track termination
    termination_condition* terminator_ptr;
    mrf_graph_type* graph_ptr;
    //! Track the collisions with the roots
    size_t ncollisions;
    //! Local junction tree graphlab core
    jtree_gl::core jt_core;
    /**
     * Local jt list used to build on the structure of the
     * jt_core.graph()
     */
    jtree_list jt_list;
    /**
     * Local data structures to reduce thread contention
     */
    std::deque<vertex_id_t> bfs_queue;
    graphlab::mutable_queue<size_t, double> priority_queue;
    boost::unordered_set<vertex_id_t> visited;
    factor_t clique_factor;
    factor_t product_of_marginals_factor;
    factor_t conditional_factor;
    factor_t marginal_factor;
  };

  std::set<splash_state*> state_set;
  jt_worker(splash_state* state_ptr = NULL);
  void operator+=(const jt_builder& other);

  //! The main loop
  void run();

private:
  //! Construct a single splash
  size_t splash_once();
  //! advance the root
  void advance_root();
  /**
   * Grab this vertex into the tree owned by worker id
   */
  bool is_vertex_available(vertex_id_t vid);
  /**
   * Grab this vertex into the tree owned by worker id
   */
  bool try_grab_vertex(iscope_type& scope);
  /**
   * Release the vertex
   */
  void release_vertex(iscope_type& scope);

  double score_vertex(vertex_id_t vid);
  double score_vertex_l1_diff(vertex_id_t vid);
  double score_vertex_log_odds(vertex_id_t vid);
  double score_vertex_lik(vertex_id_t vid);

  void grow_bfs_jtree();
  void grow_prioritized_jtree();
};  // End of JT worker

/**
 * The jt_splash_sampler implements the junction tree based Gibbs
 * sampler defined in:
 *
 *  Parallel Gibbs Sampling: From Colored Fields to Think Junction Trees
 *   by Joseph Gonzalez, Yucheng Low, Arthur Gretton, and Carlos Guestrin
 *
 */
class jt_splash_sampler {
public:
  typedef graphlab::general_scope_factory<mrf_graph_type>
  scope_factory_type;

private:
  std::vector<jt_worker*>     workers;
  scope_factory_type          scope_factory;
  std::vector< vertex_id_t >  root_perm;
  termination_condition       terminator;
public:
  jt_splash_sampler(mrf_graph_type& mrf_core,
                    const splash_settings& settings);
  ~jt_splash_sampler();

  /**
   * Get the number of times the splash sampler collided on a root.
   * This minor race event can lead to wasted cpu cycles but does not
   * affect the quality of the samples.
   */
  size_t total_collisions() const;

  /**
   * Get the total number of trees constructed on the last run
   */
  size_t total_trees() const;

  /**
   * Get the total number of single variable updates on the last run.
   */
  size_t total_samples() const;

  /** Run the splash sampler for a fixed number of seconds */
  void sample_seconds(float runtime_secs);
  /** Run the splash sampler for a fixed number of trees */
  void sample_trees(size_t total_trees);
  /**
   * Run the splash sampler for a fixed number of single variable *
   * updates
   */
  void sample_updates(size_t total_updates);
private:
  void run();
};

#endif