main.cpp

#define _CRT_SECURE_NO_WARNINGS

#include <iostream>
#include <vector>
#include <cstdio>
#include <cstdint>
#include <cmath>
#include <cfloat>
#include <queue>
#include <algorithm>
#include <numeric>
#include <unordered_set>
#include <atomic>

#include "octree/Octree.hpp"
#include "estimator/graph.cpp"
#include "estimator/point.cpp"

std::vector<estimator::Point> points;
std::vector<estimator::Point> normals;
// Pauly's surface variation lambda0 / (lambda0 + lambda1 + lambda2) per point
std::vector<float> variations;

// Neighborhood size for tangent plane estimation and the Riemannian graph.
// Scale-free, unlike a fixed search radius (Hoppe et al. use k-NN as well).
const uint32_t K_NEIGHBORS = 16;
// Strength of the Euclidean distance term in the propagation cost. Penalizes
// long edges so orientation doesn't jump across gaps between surface sheets.
const float DISTANCE_WEIGHT = 0.1f;
// Strength of the surface variation term in the propagation cost. Prefers
// propagating orientation through flat regions before creases.
const float CURVATURE_WEIGHT = 0.5f;

// Find the k nearest neighbors of points[index] (excluding itself) by growing
// a radius search until enough candidates are found, then keeping the closest k.
// Also reports the distance to the nearest neighbor for spacing estimation.
std::vector<uint32_t> kNearestNeighbors(const unibn::Octree<estimator::Point>& octree, uint32_t index,
	uint32_t k, float startRadius, float& nearestDist) {
	std::vector<uint32_t> results;
	std::vector<float> sqrDists;
	float radius = startRadius;
	while (true) {
		octree.radiusNeighbors<unibn::L2Distance<estimator::Point> >(points.at(index), radius, results, sqrDists);
		// The query point itself is included in the results, hence > k
		if (results.size() > k || results.size() >= points.size()) break;
		radius *= 1.5f;
	}

	std::vector<uint32_t> order(results.size());
	std::iota(order.begin(), order.end(), 0);
	std::sort(order.begin(), order.end(), [&sqrDists](uint32_t a, uint32_t b) {
		return sqrDists.at(a) < sqrDists.at(b);
	});

	std::vector<uint32_t> neighbors;
	neighbors.reserve(k);
	nearestDist = 0.0f;
	for (size_t i = 0; i < order.size() && neighbors.size() < k; i++) {
		uint32_t idx = results.at(order.at(i));
		if (idx == index) continue;
		if (neighbors.empty()) nearestDist = std::sqrt(sqrDists.at(order.at(i)));
		neighbors.push_back(idx);
	}
	return neighbors;
}

// Eigenvalues of a symmetric 3x3 matrix via the trigonometric closed form,
// returned in ascending order.
void symmetricEigenvalues(double xx, double xy, double xz, double yy, double yz, double zz, double eig[3]) {
	const double p1 = xy * xy + xz * xz + yz * yz;
	if (p1 == 0.0) {
		// Already diagonal
		eig[0] = xx; eig[1] = yy; eig[2] = zz;
	} else {
		const double q = (xx + yy + zz) / 3.0;
		const double p2 = (xx - q) * (xx - q) + (yy - q) * (yy - q) + (zz - q) * (zz - q) + 2.0 * p1;
		const double p = std::sqrt(p2 / 6.0);
		// r = det((A - qI) / p) / 2, clamped so float drift can't push acos out of range
		const double bxx = (xx - q) / p, byy = (yy - q) / p, bzz = (zz - q) / p;
		const double bxy = xy / p, bxz = xz / p, byz = yz / p;
		double r = (bxx * (byy * bzz - byz * byz)
			- bxy * (bxy * bzz - byz * bxz)
			+ bxz * (bxy * byz - byy * bxz)) / 2.0;
		if (r > 1.0) r = 1.0;
		if (r < -1.0) r = -1.0;
		const double PI = 3.14159265358979323846;
		const double phi = std::acos(r) / 3.0;
		eig[2] = q + 2.0 * p * std::cos(phi);
		eig[0] = q + 2.0 * p * std::cos(phi + 2.0 * PI / 3.0);
		eig[1] = 3.0 * q - eig[0] - eig[2];
	}
	std::sort(eig, eig + 3);
}

// Fit a tangent plane to the neighborhood of points[index] via weighted PCA.
// Neighbors are Gaussian-weighted by distance (Pauly et al. 2002) so outliers
// at the edge of the neighborhood barely influence the fit. Covariance is
// accumulated in double to survive large coordinates. Returns false when the
// neighborhood is degenerate (collinear or coincident points).
bool fitTangentPlane(uint32_t index, const std::vector<uint32_t>& neighbors,
	estimator::Point& normalOut, float& variationOut) {
	if (neighbors.size() < 2) return false;

	const estimator::Point query = points.at(index);

	double maxSqrDist = 0.0;
	for (size_t i = 0; i < neighbors.size(); i++) {
		estimator::Point p = points.at(neighbors.at(i));
		const double dx = p.x - query.x, dy = p.y - query.y, dz = p.z - query.z;
		const double d2 = dx * dx + dy * dy + dz * dz;
		if (d2 > maxSqrDist) maxSqrDist = d2;
	}
	if (maxSqrDist <= 0.0) return false;
	// Gaussian falloff with sigma at half the neighborhood radius
	const double sigma2 = maxSqrDist / 4.0;

	// Weighted centroid over the neighborhood plus the query point itself (weight 1, distance 0)
	double wSum = 1.0;
	double cx = query.x, cy = query.y, cz = query.z;
	std::vector<double> weights(neighbors.size());
	for (size_t i = 0; i < neighbors.size(); i++) {
		estimator::Point p = points.at(neighbors.at(i));
		const double dx = p.x - query.x, dy = p.y - query.y, dz = p.z - query.z;
		const double w = std::exp(-(dx * dx + dy * dy + dz * dz) / (2.0 * sigma2));
		weights[i] = w;
		wSum += w;
		cx += w * p.x;
		cy += w * p.y;
		cz += w * p.z;
	}
	cx /= wSum; cy /= wSum; cz /= wSum;

	double xx = 0.0, xy = 0.0, xz = 0.0, yy = 0.0, yz = 0.0, zz = 0.0;
	{
		const double rx = query.x - cx, ry = query.y - cy, rz = query.z - cz;
		xx = rx * rx; xy = rx * ry; xz = rx * rz; yy = ry * ry; yz = ry * rz; zz = rz * rz;
	}
	for (size_t i = 0; i < neighbors.size(); i++) {
		estimator::Point p = points.at(neighbors.at(i));
		const double w = weights[i];
		const double rx = p.x - cx, ry = p.y - cy, rz = p.z - cz;
		xx += w * rx * rx;
		xy += w * rx * ry;
		xz += w * rx * rz;
		yy += w * ry * ry;
		yz += w * ry * rz;
		zz += w * rz * rz;
	}

	double eig[3];
	symmetricEigenvalues(xx, xy, xz, yy, yz, zz, eig);
	const double eigSum = eig[0] + eig[1] + eig[2];
	// Collinear (or coincident) neighborhoods have two vanishing eigenvalues
	if (eigSum <= 0.0 || eig[1] <= 1e-9 * eigSum) return false;

	// The normal is the eigenvector of the smallest eigenvalue: the null space
	// of M = C - lambda0*I, found as the largest cross product of M's rows
	const double m00 = xx - eig[0], m11 = yy - eig[0], m22 = zz - eig[0];
	const double r0[3] = { m00, xy, xz };
	const double r1[3] = { xy, m11, yz };
	const double r2[3] = { xz, yz, m22 };
	double best[3] = { 0.0, 0.0, 0.0 };
	double bestNorm = 0.0;
	const double* rows[3][2] = { { r0, r1 }, { r0, r2 }, { r1, r2 } };
	for (int i = 0; i < 3; i++) {
		const double* a = rows[i][0];
		const double* b = rows[i][1];
		const double c[3] = {
			a[1] * b[2] - a[2] * b[1],
			a[2] * b[0] - a[0] * b[2],
			a[0] * b[1] - a[1] * b[0]
		};
		const double n = c[0] * c[0] + c[1] * c[1] + c[2] * c[2];
		if (n > bestNorm) {
			bestNorm = n;
			best[0] = c[0]; best[1] = c[1]; best[2] = c[2];
		}
	}
	if (bestNorm <= 0.0) return false;

	normalOut = estimator::Point((float)best[0], (float)best[1], (float)best[2]);
	normalOut.normalize();
	variationOut = (float)(eig[0] / eigSum);
	return true;
}

// Riemannian graph edge cost. Hoppe et al. use 1 - |n_i . n_j|; the absolute
// value is essential because normals are still unoriented here, so their signs
// carry no information. The distance term discourages propagating orientation
// across gaps between close-but-separate surface sheets, and the curvature
// term routes propagation through flat regions before creases.
float propagationCost(uint32_t a, uint32_t b, float avgSpacing) {
	float alignment = 1.0f - std::fabs(normals.at(a).dot(normals.at(b)));
	estimator::Point pa = points.at(a);
	estimator::Point pb = points.at(b);
	const float dx = pa.x - pb.x, dy = pa.y - pb.y, dz = pa.z - pb.z;
	const float dist = std::sqrt(dx * dx + dy * dy + dz * dz);
	float cost = alignment;
	if (avgSpacing > 0.0f) cost += DISTANCE_WEIGHT * (dist / avgSpacing);
	cost += CURVATURE_WEIGHT * (variations.at(a) + variations.at(b));
	return cost;
}

// Flip the normal at nextIndex when it points away from the reference direction
void propagateOrientation(estimator::Point reference, uint32_t nextIndex) {
	if (normals.at(nextIndex).dot(reference) < 0) {
		normals.at(nextIndex).flipSign();
	}
}

class UnionFind {
public:
	UnionFind(uint32_t n) : parent(n) {
		std::iota(parent.begin(), parent.end(), 0);
	}
	uint32_t find(uint32_t x) {
		while (parent[x] != x) {
			parent[x] = parent[parent[x]];
			x = parent[x];
		}
		return x;
	}
	bool unite(uint32_t a, uint32_t b) {
		a = find(a);
		b = find(b);
		if (a == b) return false;
		parent[a] = b;
		return true;
	}
private:
	std::vector<uint32_t> parent;
};

struct HeapEdge {
	float cost;
	uint32_t from, to;
	bool operator>(const HeapEdge& other) const { return cost > other.cost; }
};

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

	// Read points in format X Y Z per line from STDIN. Input order is
	// preserved through to the output.
	float p[3];
	while (scanf("%f %f %f", &p[0], &p[1], &p[2]) == 3) {
		points.push_back(estimator::Point(p[0], p[1], p[2]));
	}

	const uint32_t n = (uint32_t)points.size();
	if (n < 3) {
		std::cerr << "at least 3 points required" << std::endl;
		return 1;
	}

	// Build an octree to search on the points
	unibn::Octree<estimator::Point> octree;
	octree.initialize(points);

	// Estimate the data scale from the bounding box so search radii adapt to
	// the input instead of assuming fixed world units
	estimator::Point bbMin = points.at(0), bbMax = points.at(0);
	for (uint32_t i = 1; i < n; i++) {
		estimator::Point pt = points.at(i);
		bbMin.x = std::min(bbMin.x, pt.x); bbMax.x = std::max(bbMax.x, pt.x);
		bbMin.y = std::min(bbMin.y, pt.y); bbMax.y = std::max(bbMax.y, pt.y);
		bbMin.z = std::min(bbMin.z, pt.z); bbMax.z = std::max(bbMax.z, pt.z);
	}
	const float dx = bbMax.x - bbMin.x, dy = bbMax.y - bbMin.y, dz = bbMax.z - bbMin.z;
	const float bbDiag = std::sqrt(dx * dx + dy * dy + dz * dz);
	const float startRadius = std::max(bbDiag / std::cbrt((float)n), FLT_EPSILON * std::max(bbDiag, 1.0f));

	const uint32_t k = std::min(K_NEIGHBORS, n - 1);

	// Estimate a tangent plane for each point via weighted PCA on its k nearest
	// neighbors. Degenerate neighborhoods retry with a doubled k. The loop is
	// embarrassingly parallel; enable OpenMP to spread it across cores.
	normals.resize(n);
	variations.resize(n);
	std::vector<std::vector<uint32_t> > neighborhoods(n);
	std::vector<float> nearestDists(n, 0.0f);
	std::atomic<bool> degenerate(false);
	#pragma omp parallel for schedule(dynamic, 256)
	for (int i = 0; i < (int)n; i++) {
		if (degenerate.load(std::memory_order_relaxed)) continue;
		uint32_t kq = k;
		while (true) {
			float nearestDist = 0.0f;
			std::vector<uint32_t> neighbors = kNearestNeighbors(octree, (uint32_t)i, kq, startRadius, nearestDist);
			if (fitTangentPlane((uint32_t)i, neighbors, normals.at(i), variations.at(i))) {
				neighborhoods.at(i).swap(neighbors);
				nearestDists.at(i) = nearestDist;
				break;
			}
			if (neighbors.size() >= n - 1) {
				// Even the full cloud doesn't span a plane
				degenerate = true;
				break;
			}
			kq *= 2;
		}
	}
	if (degenerate) {
		std::cerr << "points dont span a plane" << std::endl;
		return 1;
	}

	// Average nearest neighbor distance, used to normalize the distance term
	// of the propagation cost
	const float avgSpacing = std::accumulate(nearestDists.begin(), nearestDists.end(), 0.0f) / n;

	// Build the Riemannian graph from the symmetrized k-NN graph. Hoppe et al.
	// splice an EMST with k-NN edges purely to guarantee connectivity; the k-NN
	// graph is a superset of those neighborhood relations, and any residual
	// disconnected components are bridged explicitly below.
	estimator::Graph riemann(n);
	UnionFind components(n);
	uint32_t componentCount = n;
	{
		std::unordered_set<uint64_t> seen;
		seen.reserve((size_t)n * k);
		for (uint32_t i = 0; i < n; i++) {
			const std::vector<uint32_t>& neighbors = neighborhoods.at(i);
			for (size_t j = 0; j < neighbors.size(); j++) {
				const uint32_t idx = neighbors.at(j);
				const uint64_t key = ((uint64_t)std::min(i, idx) << 32) | std::max(i, idx);
				if (seen.insert(key).second) {
					riemann.addEdge(i, idx, propagationCost(i, idx, avgSpacing));
					if (components.unite(i, idx)) componentCount--;
				}
			}
		}
	}

	// Bridge any disconnected components with their shortest outgoing link so
	// orientation can propagate everywhere
	while (componentCount > 1) {
		const uint32_t mainRoot = components.find(0);
		uint32_t stray = n;
		for (uint32_t i = 0; i < n; i++) {
			if (components.find(i) != mainRoot) {
				stray = i;
				break;
			}
		}
		const uint32_t strayRoot = components.find(stray);
		float radius = startRadius;
		bool bridged = false;
		while (!bridged) {
			std::vector<uint32_t> results;
			std::vector<float> sqrDists;
			octree.radiusNeighbors<unibn::L2Distance<estimator::Point> >(points.at(stray), radius, results, sqrDists);
			float bestDist = FLT_MAX;
			uint32_t bestIdx = n;
			for (size_t i = 0; i < results.size(); i++) {
				const uint32_t idx = results.at(i);
				if (components.find(idx) != strayRoot && sqrDists.at(i) < bestDist) {
					bestDist = sqrDists.at(i);
					bestIdx = idx;
				}
			}
			if (bestIdx < n) {
				riemann.addEdge(stray, bestIdx, propagationCost(stray, bestIdx, avgSpacing));
				components.unite(stray, bestIdx);
				componentCount--;
				bridged = true;
			} else {
				radius *= 2.0f;
			}
		}
	}

	// Seed orientation at the highest point, whose normal must face Z+ if it
	// lies on the outer surface
	uint32_t seed = 0;
	for (uint32_t i = 1; i < n; i++) {
		if (points.at(i).z > points.at(seed).z) seed = i;
	}
	propagateOrientation(estimator::Point(0.0f, 0.0f, 1.0f), seed);

	// Traverse the minimal spanning tree of the Riemannian graph (Prim with a
	// lazy-deletion priority queue) and propagate the seed's sign along it
	std::vector<bool> oriented(n, false);
	std::priority_queue<HeapEdge, std::vector<HeapEdge>, std::greater<HeapEdge> > frontier;
	oriented.at(seed) = true;
	{
		const std::vector<estimator::Edge>& seedEdges = riemann.getNeighbors(seed);
		for (size_t i = 0; i < seedEdges.size(); i++) {
			frontier.push(HeapEdge{ seedEdges.at(i).cost, seedEdges.at(i).from, seedEdges.at(i).to });
		}
	}
	while (!frontier.empty()) {
		const HeapEdge edge = frontier.top();
		frontier.pop();
		if (oriented.at(edge.to)) continue;
		oriented.at(edge.to) = true;
		propagateOrientation(normals.at(edge.from), edge.to);
		const std::vector<estimator::Edge>& next = riemann.getNeighbors(edge.to);
		for (size_t i = 0; i < next.size(); i++) {
			if (!oriented.at(next.at(i).to)) {
				frontier.push(HeapEdge{ next.at(i).cost, next.at(i).from, next.at(i).to });
			}
		}
	}

	// Sanity-check the global sign against all six bounding box extremes: on a
	// closed surface each extreme point's normal must face outward along its
	// axis. A majority vote here recovers from a bad seed (e.g. the highest
	// point sitting on a vertical wall), which would otherwise invert the
	// entire cloud.
	{
		uint32_t extremes[6] = { 0, 0, 0, 0, 0, 0 };
		for (uint32_t i = 1; i < n; i++) {
			estimator::Point pt = points.at(i);
			if (pt.x < points.at(extremes[0]).x) extremes[0] = i;
			if (pt.x > points.at(extremes[1]).x) extremes[1] = i;
			if (pt.y < points.at(extremes[2]).y) extremes[2] = i;
			if (pt.y > points.at(extremes[3]).y) extremes[3] = i;
			if (pt.z < points.at(extremes[4]).z) extremes[4] = i;
			if (pt.z > points.at(extremes[5]).z) extremes[5] = i;
		}
		const estimator::Point outward[6] = {
			estimator::Point(-1.0f, 0.0f, 0.0f), estimator::Point(1.0f, 0.0f, 0.0f),
			estimator::Point(0.0f, -1.0f, 0.0f), estimator::Point(0.0f, 1.0f, 0.0f),
			estimator::Point(0.0f, 0.0f, -1.0f), estimator::Point(0.0f, 0.0f, 1.0f)
		};
		float score = 0.0f;
		for (int i = 0; i < 6; i++) {
			score += normals.at(extremes[i]).dot(outward[i]);
		}
		if (score < 0.0f) {
			for (uint32_t i = 0; i < n; i++) {
				normals.at(i).flipSign();
			}
		}
	}

	// Output results to STDOUT
	for (uint32_t i = 0; i < n; ++i) {
		estimator::Point point = points.at(i);
		estimator::Point normal = normals.at(i);
		std::cout << point.x << " " << point.y << " " << point.z << " ";
		std::cout << normal.x << " " << normal.y << " " << normal.z << std::endl;
	}

	return 0;
}