#include "fits.h"

#ifdef WIN
#pragma warning (disable: 4018) // disable warning for comparing signed and unsigned
#endif /* WIN */

// Computes weighted centroid of one bin
void bin2d_weighted_centroid(Fits &x, Fits &y, Fits &density, double *xBar, double *yBar) {
	double mass = density.get_flux();
	Fits _tmp;
	_tmp.copy(density);
	_tmp.mul(x);
	*xBar = _tmp.get_flux() / mass;
	_tmp.copy(density);
	_tmp.mul(y);	
	*yBar = _tmp.get_flux() / mass;
}

// Implements equation (5) of Cappellari & Copin (2003)
double bin2d_roundness(Fits &x, Fits &y, double pixelSize) {
	double n = (double)(x.Nelements());
	double equivalentRadius = sqrt(n/M_PI)*pixelSize;
	double xBar = x.get_flux() / n;          // unweighted centroid here!
	double yBar = y.get_flux() / n;
	Fits _tmp1, _tmp2;
	_tmp1.copy(x);
	_tmp1.sub(xBar);
	_tmp1.mul(_tmp1);
	_tmp2.copy(y);
	_tmp2.sub(yBar);
	_tmp2.mul(_tmp2);
	_tmp1.add(_tmp2);
	double maxDistance = sqrt(_tmp1.get_max());
	return maxDistance / equivalentRadius - 1.0;
}

// Implements steps (i)-(v) in section 5.1 of Cappellari & Copin (2003)
void bin2d_accretion(Fits &x, Fits &y, Fits &signal, Fits &noise, double targetSN, Fits &cclass, double *pixelSize) {
	Fits unBinned, binned;
	Fits _tmp, _tmp1, _tmp2;
	Fits currentBin;
//	long ldummy;
	long n = x.Nelements();
	int m;
	double minDist;
	double SN;
	Fits w;
	int num, nc;
	long maxnum;
	double xBar, yBar;
	int k;
	Fits nextBin;
	double SNOld, roundness;
	int j;
	int ind;
	
	cclass.create(n, 1, I4); // will contain the bin number of each given pixel
	Fits good;
	good.create(n, 1, I1);   // will contain 1 if the bin has been accepted as good

// For each point, find the distance to all other points and select the minimum.
// This is a robus way of determining the pixel size of unbinned data.

	double dx = 1e30;
	for (j = 0; j <= n-2; j++) {
		double dd;
		x.extractRange(_tmp1, j+2, n, -1, -1, -1, -1); // fortran indices!!!
		y.extractRange(_tmp2, j+2, n, -1, -1, -1, -1);
		_tmp1 -= x[j];
		_tmp2 -= y[j];
		_tmp1 *= _tmp1;
		_tmp2 *= _tmp2;
		_tmp1 += _tmp2;
		dd = _tmp1.get_min();
		if (dd < dx) dx = dd;
	}
	*pixelSize = sqrt(dx);

	_tmp.copy(signal);
	_tmp.div(noise);
	currentBin.create(1, 1, I4);
	currentBin.i4data[0] = _tmp.maxLinearIndex(&SN); // Start from the pixel with highest S/N

// Rough estimate of the expected final bin number.
// This value is only used to have a feeling of the expected
// remaining computation time when binning very big dataset.

	num = w.where(_tmp, "<", targetSN);
	nc = _tmp.Nelements() - num;
	_tmp.mul(_tmp);
	maxnum = (long)(_tmp.get_flux() / (targetSN*targetSN) + .5) + nc;

// The first bin will be assigned CLASS = 1
// With N pixels there will be at most N bins

	for (ind = 1; ind <= n; ind++) {
		dp_output("Bin: %i / %i\n", ind, maxnum);
		cclass.i4data[currentBin.i4data[0]] = ind; // Here currentBin is still made of one pixel

		xBar = x[currentBin.i4data[0]];
		yBar = y[currentBin.i4data[0]]; // Centroid of one pixels

		while (TRUE) {
			m = unBinned.where(cclass, "==", 0);

			if (m == 0) break; // Stops if all pixels are binned
// Find the unbinned pixel closest to the centroid of the current bin

			_tmp1.extractLinearRange(x, unBinned);
			_tmp2.extractLinearRange(y, unBinned);
			_tmp1 -= xBar;
			_tmp2 -= yBar;
			_tmp1 *= _tmp1;
			_tmp2 *= _tmp2;
			_tmp1 += _tmp2;
			k = _tmp1.minLinearIndex(&minDist);
// Find the distance from the closest pixel to the current bin
			_tmp1.extractLinearRange(x, currentBin);
			_tmp2.extractLinearRange(y, currentBin);
			_tmp1 -= x[unBinned.i4data[k]];
			_tmp2 -= y[unBinned.i4data[k]];
			_tmp1 *= _tmp1;
			_tmp2 *= _tmp2;
			_tmp1 += _tmp2;
			minDist = _tmp1.get_min();

// Estimate the `roundness' of the POSSIBLE new bin
			nextBin.create(currentBin.Nelements() + 1, 1, I4);
			for (j = 0; j < currentBin.Nelements(); j++)
				nextBin.i4data[j] = currentBin.i4data[j];
			nextBin.i4data[nextBin.Nelements() - 1] = unBinned.i4data[k];

			_tmp1.extractLinearRange(x, nextBin);
			_tmp2.extractLinearRange(y, nextBin);
			roundness = bin2d_roundness(_tmp1, _tmp2, *pixelSize);

// Compute the S/N one would obtain by adding
// the CANDIDATE pixel to the current bin

			SNOld = SN;
			_tmp1.extractLinearRange(signal, nextBin);
			_tmp2.extractLinearRange(noise, nextBin);
			_tmp2 *= _tmp2;
			_tmp1 /= sqrt(_tmp2.get_flux());
			SN = _tmp1.get_flux();

// Test whether the CANDIDATE pixel is connected to the
// current bin, whether the POSSIBLE new bin is round enough
// and whether the resulting S/N would get closer to targetSN

			if ((sqrt(minDist) > 1.2* (*pixelSize)) || 
				(roundness > 0.3) ||
				(fabs(SN-targetSN) > fabs(SNOld-targetSN))) {
					if (SNOld > 0.8*targetSN) {
						for (int j = 0; j < currentBin.Nelements(); j++) {
							good.i1data[currentBin.i4data[j]] = 1;
						}
					}
				break;
			}

// If all the above tests are negative then accept the CANDIDATE pixel,
// add it to the current bin, and continue accreting pixels

			cclass.i4data[unBinned.i4data[k]] = ind;
			currentBin.copy(nextBin);

// Update the centroid of the current bin
			_tmp1.extractLinearRange(x, currentBin);
			_tmp2.extractLinearRange(y, currentBin);
			_tmp.extractLinearRange(signal, currentBin);
			bin2d_weighted_centroid(_tmp1, _tmp2, _tmp, &xBar, &yBar);
		}

// Get the centroid of all the binned pixels
    	m = unBinned.where(cclass, "==", 0);
		if (m == 0) break; // Stop if all pixels are binned
		binned.where(cclass, "!=", 0);
		_tmp1.extractLinearRange(x, binned);
		_tmp2.extractLinearRange(y, binned);
		_tmp.extractLinearRange(signal, binned);
		bin2d_weighted_centroid(_tmp1, _tmp2, _tmp, &xBar, &yBar);

// Find the closest unbinned pixel to the centroid of all
// the binned pixels, and start a new bin from that pixel.
		_tmp1.extractLinearRange(x, unBinned);
		_tmp2.extractLinearRange(y, unBinned);
		_tmp1 -= xBar;
		_tmp1 *= _tmp1;
		_tmp2 -= yBar;
		_tmp2 *= _tmp2;
		_tmp1 += _tmp2;
		k = _tmp1.minLinearIndex(&minDist);
		currentBin.create(1, 1, I4);
		currentBin.i4data[0] = unBinned.i4data[k]; // The bin is initially made of one pixel
		SN = signal[currentBin.i4data[0]]/noise[currentBin.i4data[0]];
	}

// Set to zero all bins that did not reach the target S/N
	for (j = 0; j < cclass.Nelements(); j++)
		cclass.i4data[j] *= good.i1data[j];
}

// Implements steps (vi)-(vii) in section 5.1 of Cappellari & Copin (2003)
void bin2d_reassign_bad_bins(Fits &x, Fits &y, Fits &signal, Fits &noise, double targetSN, Fits &cclass, Fits &xnode, Fits &ynode) {

// Find the centroid of all succesful bins.
// CLASS = 0 are unbinned pixels which are excluded.
	Fits area, good, r, p, bad;
	Fits _tmp, _tmp1, _tmp2;
	int k, nnodes, m, index, j;
	double xBar, yBar, max, tmp;

	max = cclass.get_max();
	area.histogram(cclass, 1., max);
	r.histogram_indices(cclass, 1.0, max);
	for (j = 0; j < r.Nelements(); j++) r.i4data[j]--;
	nnodes = good.where(area, ">", 0.0); // Obtain the index of the good bins
	xnode.create(nnodes, 1, R8);
	ynode.create(nnodes, 1, R8);
	for (j = 0; j < nnodes; j++) {
		k = good.i4data[j];
		r.extractLinearIndex(p, r.i4data[k], r.i4data[k+1]-1); // Find subscripts of pixels in bin k.
		_tmp1.extractLinearRange(x, p);
		_tmp2.extractLinearRange(y, p);
		_tmp.extractLinearRange(signal, p);
		bin2d_weighted_centroid(_tmp1, _tmp2, _tmp, &xBar, &yBar);
		xnode.r8data[j] = xBar;
		ynode.r8data[j] = yBar;
	}

// Reassign pixels of bins with S/N < targetSN
// to the closest centroid of a good bin

	m = bad.where(cclass, "==", 0);
	for (j = 0; j < m; j++) {
		_tmp1.copy(xnode);
		_tmp1 *= -1.;
		_tmp1 += x[bad.i4data[j]];
		_tmp1 *= _tmp1;
		_tmp2.copy(ynode);
		_tmp2 *= -1.;
		_tmp2 += y[bad.i4data[j]];
		_tmp2 *= _tmp2;
		_tmp1 += _tmp2;
		index = _tmp.minLinearIndex(&tmp);
		cclass.i4data[bad.i4data[j]] = good.i4data[index] + 1;
	}

// Recompute all centroids of the reassigned bins.
// These will be used as starting points for the CVT.

	max = cclass.get_max();
	area.histogram(cclass, 0.0, max, 1.0);
	r.histogram_indices(cclass, 0.0, max, 1.0);
	for (j = 0; j < r.Nelements(); j++) r.i4data[j]--;
	nnodes = good.where(area, ">", 0.0); // Re-obtain the index of the good bins
	for (j = 0; j < nnodes-1; j++) {
		k = good.i4data[j];
		r.extractLinearIndex(p, r.i4data[k], r.i4data[k+1]-1); // Find subscripts of pixels in bin k.
		_tmp1.extractLinearRange(x, p);
		_tmp2.extractLinearRange(y, p);
		_tmp.extractLinearRange(signal, p);
		bin2d_weighted_centroid(_tmp1, _tmp2, _tmp, &xBar, &yBar);
		xnode.r8data[j] = xBar;
		ynode.r8data[j] = yBar;
	}
}

// Implements the modified Lloyd algorithm
// in section 4.1 of Cappellari & Copin (2003)
void bin2d_cvt_equal_mass(Fits &x, Fits &y, Fits &signal, Fits &noise, Fits &xnode, Fits &ynode, int *iter) {
	unsigned long npixels;
	int index, m, k, j;
	double tmp, xBar, yBar, diff, min, max;
	Fits cclass, dens, xnodeOld, ynodeOld, _tmp, _tmp1, _tmp2, area, r, w, p;

	npixels = signal.Nelements();
	cclass.create(npixels, 1, I4);
	dens.copy(signal);
	dens /= noise;
	dens *= dens;      // See beginning of section 4.1

	*iter = 1;
	do {
		xnodeOld.copy(xnode);
		ynodeOld.copy(ynode);

// Computes Voronoi Tessellation of the pixels grid

		for (j = 0; j < npixels; j++) {
			_tmp1.copy(xnode);
			_tmp1 *= -1.;
			_tmp1 += x[j];
			_tmp1 *= _tmp1;
			_tmp2.copy(ynode);
			_tmp2 *= -1.;
			_tmp2 += y[j];
			_tmp2 *= _tmp2;
			_tmp1 += _tmp2;
			index = _tmp1.minLinearIndex(&tmp);
			cclass.i4data[j] = index;
		}

// Computes centroids of the bins, weighted by dens^2.
// Exponent 2 on the density produces equal-mass Voronoi bins.
		cclass.get_minmax(&min, &max);
		area.histogram(cclass, min, max);
		r.histogram_indices(cclass, min, max);
		for (j = 0; j < r.Nelements(); j++) r.i4data[j]--;
		m = w.where(area, ">", 0.0); // Check for zero-size Voronoi bins
		for (j = 0; j < m; j++) {
			k = w.i4data[j];
			r.extractLinearIndex(p, r.i4data[k], r.i4data[k+1]-1); // Find subscripts of pixels in bin k.
			_tmp1.extractLinearRange(x, p);
			_tmp2.extractLinearRange(y, p);
			_tmp.extractLinearRange(dens, p);
			bin2d_weighted_centroid(_tmp1, _tmp2, _tmp, &xBar, &yBar);
			xnode.r8data[k] = xBar;
			ynode.r8data[k] = yBar;
		}
		
		_tmp1.copy(xnode);
		_tmp1 -= xnodeOld;
		_tmp1 *= _tmp1;
		_tmp2.copy(ynode);
		_tmp2 -= ynodeOld;
		_tmp2 *= _tmp2;
		_tmp1 += _tmp2;
		
		diff = _tmp1.get_flux();
		(*iter)++;

		dp_output("Iter:  %i,  Diff:  %f\n", *iter, diff);
	} while (diff != 0.0);

// Only return the generators of the nonzero Voronoi bins
	_tmp.copy(xnode);
	xnode.extractLinearRange(_tmp, w);
	_tmp.copy(ynode);
	ynode.extractLinearRange(_tmp, w);
}

// Recomputes Voronoi Tessellation of the pixels grid to make sure
// that the class number corresponds to the proper Voronoi generator.
// This is done to take into account possible zero-size Voronoi bins
// in output from the previous CVT.
void bin2d_compute_useful_bin_quantities(Fits &x, Fits &y, Fits &signal, Fits &noise, Fits &xnode, Fits &ynode, Fits &cclass, Fits &xBar, Fits &yBar, Fits &sn, Fits &area) {
	int npix, index, nbins, j;
	Fits _tmp1, _tmp2, _tmp, p, r;
	double tmp, min, max, xb, yb;
	
	npix = x.Nelements();
	cclass.create(npix, 1, I4); // will contain the bin number of each given pixel
	for (j = 0; j < npix; j++) {
		_tmp1.copy(xnode);
		_tmp1 *= -1.;
		_tmp1 += x[j];
		_tmp1 *= _tmp1;
		_tmp2.copy(ynode);
		_tmp2 *= -1.;
		_tmp2 += y[j];
		_tmp2 *= _tmp2;
		_tmp1 += _tmp2;
		index = _tmp1.minLinearIndex(&tmp);
		cclass.i4data[j] = index;
	}

// At the end of the computation evaluate the bin luminosity-weighted
// centroids (xbar,ybar) and the corresponding final S/N of each bin.

	cclass.get_minmax(&min, &max);
	area.histogram(cclass, min, max);
	r.histogram_indices(cclass, min, max);
	for (j = 0; j < r.Nelements(); j++) r.i4data[j]--;
	nbins = xnode.Nelements();
	xBar.create(nbins, 1, R8);
	yBar.create(nbins, 1, R8);
	sn.create(nbins, 1, R8);
	for (j = 0; j < nbins; j++) {
		r.extractLinearIndex(p, r.i4data[j], r.i4data[j+1]-1); // Find subscripts of pixels in bin j.
		_tmp1.extractLinearRange(x, p);
		_tmp2.extractLinearRange(y, p);
		_tmp.extractLinearRange(signal, p);
		
		bin2d_weighted_centroid(_tmp1, _tmp2, _tmp, &xb, &yb);
		xBar.r8data[j] = xb;
		yBar.r8data[j] = yb;
		_tmp1.extractLinearRange(noise, p);
		_tmp1 *= _tmp1;
		sn.r8data[j] = _tmp.get_flux() / sqrt(_tmp1.get_flux());
	}
}

void voronoi(Fits &signal, Fits &noise, Fits &apply, double targetSN, Fits &result, int returnwhat) {
	double pixelSize;
	Fits cclass, xNode, yNode, xbar, ybar, sn, area, x, y, ssignal, snoise;
	int iter, i, j, count, z;

// Create x and y vectors
	x.copy(signal);
	x.setType(I4);
	y.copy(x);
	
	for (i = 1; i <= signal.Naxis(1); i++) {
		for (j = 1; j <= signal.Naxis(2); j++) {
			x.i4data[x.F_I(i, j)] = i-1;
			y.i4data[y.F_I(i, j)] = j-1;
		}
	}
	x.setNaxis(1, x.Nelements());
	x.setNaxis(2, 1);
	x.setNaxis(3, 1);
	y.setNaxis(1, y.Nelements());
	y.setNaxis(2, 1);
	y.setNaxis(3, 1);

// exclude pixels with noise <= 0
	count = 0;
	ssignal.copy(signal);
	ssignal.setType(R8);
	snoise.copy(noise);
	snoise.setType(R8);
	ssignal.setNaxis(1, x.Nelements());
	ssignal.setNaxis(2, 1);
	ssignal.setNaxis(3, 1);
	snoise.setNaxis(1, y.Nelements());
	snoise.setNaxis(2, 1);
	snoise.setNaxis(3, 1);
	for (i = 0; i < x.Nelements(); i++) {
		if (noise[i] > 0.0) {
			x.i4data[count] = x.i4data[i];
			y.i4data[count] = y.i4data[i];
			ssignal.r8data[count] = signal[i];
			snoise.r8data[count] = noise[i];
			count++;
		}
	}
	x.resize(count, 1);
	y.resize(count, 1);
	ssignal.resize(count, 1);
	snoise.resize(count, 1);

	dp_output("Bin-accretion...\n");
	bin2d_accretion(x, y, ssignal, snoise, targetSN, cclass, &pixelSize);
	dp_output("%i initial bins.\n", (int)cclass.get_max());
	dp_output("Reassign bad bins...\n");
	bin2d_reassign_bad_bins(x, y, ssignal, snoise, targetSN, cclass, xNode, yNode);
	dp_output("%i good bins.\n", xNode.Nelements());
	dp_output("Modified Lloyd algorithm...\n");
	bin2d_cvt_equal_mass(x, y, ssignal, snoise, xNode, yNode, &iter);
	dp_output("%i iterations.\n", iter-1);
	bin2d_compute_useful_bin_quantities(x, y, ssignal, snoise, xNode, yNode, cclass, xbar, ybar, sn, area);

	Fits image, binval, w, p;
	image.copy(signal);
	image.setType(R8);
	
	if (apply.Naxis(3) < 2) {
		count = 0;
		ssignal.copy(apply);
		ssignal.setType(R8);
		ssignal.setNaxis(1, ssignal.Nelements());
		ssignal.setNaxis(2, 1);
		ssignal.setNaxis(3, 1);
//		for (i = 0; i < x.Nelements(); i++) {
                for (i = 0; i < signal.Nelements(); i++) {
                        if (noise[i] > 0.0) {
				ssignal.r8data[count] = ssignal[i];
				count++;
			}
		}
		ssignal.resize(count, 1);

// find mean values & positions of bins
		binval.create(cclass.get_max() + 1, 1, R8);
		for (i = 0; i < binval.Nelements(); i++) {
			w.where(cclass, "==", i);
			p.extractLinearRange(ssignal, w);
			binval.r8data[i] = p.get_avg();
			if (returnwhat == 1) binval.r8data[i] = (double)i;
		}

// create output image - binned
		image = 0.0;
		for (i = 0; i < ssignal.Nelements(); i++) {
			image.r8data[image.C_I((int)x[i], (int)y[i])] = binval.r8data[cclass.i4data[i]];
		}
		result.copy(image);
	} else {
		result.copy(apply);
		for (z = 1; z <= apply.Naxis(3); z++) {
			count = 0;
			apply.extractRange(ssignal, -1, -1, -1, -1, z, z);
			ssignal.setType(R8);
			ssignal.setNaxis(1, ssignal.Nelements());
			ssignal.setNaxis(2, 1);
			ssignal.setNaxis(3, 1);
//			for (i = 0; i < x.Nelements(); i++) {
                        for (i = 0; i < signal.Nelements(); i++) {
                                if (noise[i] > 0.0) {
					ssignal.r8data[count] = ssignal[i];
					count++;
				}
			}
			ssignal.resize(count, 1);

// find mean values & positions of bins
			binval.create(cclass.get_max() + 1, 1, R8);
			for (i = 0; i < binval.Nelements(); i++) {
				w.where(cclass, "==", i);
				p.extractLinearRange(ssignal, w);
				binval.r8data[i] = p.get_avg();
				if (returnwhat == 1) binval.r8data[i] = (double)i;
			}

// create output image - binned
			image = 0.0;
			for (i = 0; i < ssignal.Nelements(); i++) {
				image.r8data[image.C_I((int)x[i], (int)y[i])] = binval.r8data[cclass.i4data[i]];
			}
			result.setRange(image, -1, -1, -1, -1, z, z);
		}
	}
}