package iron.culling;

class ObjectBVH {

	static inline var LEAF_THRESHOLD = 4;
	static inline var NUM_BINS = 12;
	static inline var TRAVERSAL_COST = 1.0;
	static inline var INTERSECT_COST = 1.0;

	public var nodes: Array<BVHNode>;
	public var objectIndices: Array<Int>;
	public var root: Int = -1;
	public var objectCount: Int = 0;

	var stack: Array<Int> = [];

	static var binCounts: Array<Int> = [for (i in 0...NUM_BINS) 0];
	static var binMinX: Array<Float> = [for (i in 0...NUM_BINS) 0];
	static var binMinY: Array<Float> = [for (i in 0...NUM_BINS) 0];
	static var binMinZ: Array<Float> = [for (i in 0...NUM_BINS) 0];
	static var binMaxX: Array<Float> = [for (i in 0...NUM_BINS) 0];
	static var binMaxY: Array<Float> = [for (i in 0...NUM_BINS) 0];
	static var binMaxZ: Array<Float> = [for (i in 0...NUM_BINS) 0];

	public function new() {
		nodes = [];
		objectIndices = [];
	}

	public function build(cullData: CullData) {
		objectCount = cullData.count;
		nodes = [];
		objectIndices = [];

		if (objectCount == 0) {
			root = -1;
			return;
		}

		var indices = [for (i in 0...objectCount) i];
		root = buildRecursive(indices, cullData);
	}

	static inline function surfaceArea(minX: Float, minY: Float, minZ: Float,
	                                   maxX: Float, maxY: Float, maxZ: Float): Float {
		var wx = maxX - minX;
		var wy = maxY - minY;
		var wz = maxZ - minZ;
		return 2.0 * (wx * wy + wx * wz + wy * wz);
	}

	function buildRecursive(indices: Array<Int>, cullData: CullData): Int {
		var nodeIdx = nodes.length;
		var node = new BVHNode();
		nodes.push(node);

		var minX = 1e20, minY = 1e20, minZ = 1e20;
		var maxX = -1e20, maxY = -1e20, maxZ = -1e20;
		var centMinX = 1e20, centMinY = 1e20, centMinZ = 1e20;
		var centMaxX = -1e20, centMaxY = -1e20, centMaxZ = -1e20;
		var cx = 0.0, cy = 0.0, cz = 0.0;

		for (idx in indices) {
			var x = cullData.getCenterX(idx);
			var y = cullData.getCenterY(idx);
			var z = cullData.getCenterZ(idx);
			if (cullData.getMinX(idx) < minX) minX = cullData.getMinX(idx);
			if (cullData.getMinY(idx) < minY) minY = cullData.getMinY(idx);
			if (cullData.getMinZ(idx) < minZ) minZ = cullData.getMinZ(idx);
			if (cullData.getMaxX(idx) > maxX) maxX = cullData.getMaxX(idx);
			if (cullData.getMaxY(idx) > maxY) maxY = cullData.getMaxY(idx);
			if (cullData.getMaxZ(idx) > maxZ) maxZ = cullData.getMaxZ(idx);
			if (x < centMinX) centMinX = x;
			if (y < centMinY) centMinY = y;
			if (z < centMinZ) centMinZ = z;
			if (x > centMaxX) centMaxX = x;
			if (y > centMaxY) centMaxY = y;
			if (z > centMaxZ) centMaxZ = z;
			cx += x;
			cy += y;
			cz += z;
		}

		var n = indices.length;
		cx /= n;
		cy /= n;
		cz /= n;

		node.minX = minX; node.minY = minY; node.minZ = minZ;
		node.maxX = maxX; node.maxY = maxY; node.maxZ = maxZ;
		node.centerX = cx;
		node.centerY = cy;
		node.centerZ = cz;

		if (n <= LEAF_THRESHOLD) {
			node.isLeaf = true;
			node.firstObject = objectIndices.length;
			for (idx in indices) objectIndices.push(idx);
			node.objectCount = n;
			return nodeIdx;
		}

		var parentSA = surfaceArea(minX, minY, minZ, maxX, maxY, maxZ);
		var leafCost = INTERSECT_COST * n;

		var bestCost = 1e20;
		var bestAxis = -1;
		var bestSplitBin = -1;

		var centBounds = [
			[centMinX, centMaxX],
			[centMinY, centMaxY],
			[centMinZ, centMaxZ]
		];

		for (axis in 0...3) {
			var lo = centBounds[axis][0];
			var hi = centBounds[axis][1];
			var range = hi - lo;
			if (range < 1e-8) continue;

			var invRange = NUM_BINS / range;

			for (b in 0...NUM_BINS) {
				binCounts[b] = 0;
				binMinX[b] = 1e20; binMinY[b] = 1e20; binMinZ[b] = 1e20;
				binMaxX[b] = -1e20; binMaxY[b] = -1e20; binMaxZ[b] = -1e20;
			}

			for (idx in indices) {
				var cv = axis == 0 ? cullData.getCenterX(idx) :
				         axis == 1 ? cullData.getCenterY(idx) :
				                     cullData.getCenterZ(idx);
				var b = Std.int((cv - lo) * invRange);
				if (b < 0) b = 0;
				if (b >= NUM_BINS) b = NUM_BINS - 1;
				binCounts[b]++;
				if (cullData.getMinX(idx) < binMinX[b]) binMinX[b] = cullData.getMinX(idx);
				if (cullData.getMinY(idx) < binMinY[b]) binMinY[b] = cullData.getMinY(idx);
				if (cullData.getMinZ(idx) < binMinZ[b]) binMinZ[b] = cullData.getMinZ(idx);
				if (cullData.getMaxX(idx) > binMaxX[b]) binMaxX[b] = cullData.getMaxX(idx);
				if (cullData.getMaxY(idx) > binMaxY[b]) binMaxY[b] = cullData.getMaxY(idx);
				if (cullData.getMaxZ(idx) > binMaxZ[b]) binMaxZ[b] = cullData.getMaxZ(idx);
			}

			var rightCount = [for (b in 0...NUM_BINS) 0];
			var rightMinX = [for (b in 0...NUM_BINS) 0.0];
			var rightMinY = [for (b in 0...NUM_BINS) 0.0];
			var rightMinZ = [for (b in 0...NUM_BINS) 0.0];
			var rightMaxX = [for (b in 0...NUM_BINS) 0.0];
			var rightMaxY = [for (b in 0...NUM_BINS) 0.0];
			var rightMaxZ = [for (b in 0...NUM_BINS) 0.0];

			var cumCount = 0;
			var cumMinX = 1e20, cumMinY = 1e20, cumMinZ = 1e20;
			var cumMaxX = -1e20, cumMaxY = -1e20, cumMaxZ = -1e20;

			for (b in 0...NUM_BINS) {
				var bi = NUM_BINS - 1 - b;
				cumCount += binCounts[bi];
				if (binMinX[bi] < cumMinX) cumMinX = binMinX[bi];
				if (binMinY[bi] < cumMinY) cumMinY = binMinY[bi];
				if (binMinZ[bi] < cumMinZ) cumMinZ = binMinZ[bi];
				if (binMaxX[bi] > cumMaxX) cumMaxX = binMaxX[bi];
				if (binMaxY[bi] > cumMaxY) cumMaxY = binMaxY[bi];
				if (binMaxZ[bi] > cumMaxZ) cumMaxZ = binMaxZ[bi];
				rightCount[bi] = cumCount;
				rightMinX[bi] = cumMinX; rightMinY[bi] = cumMinY; rightMinZ[bi] = cumMinZ;
				rightMaxX[bi] = cumMaxX; rightMaxY[bi] = cumMaxY; rightMaxZ[bi] = cumMaxZ;
			}

			var leftCount = 0;
			var leftMinX = 1e20, leftMinY = 1e20, leftMinZ = 1e20;
			var leftMaxX = -1e20, leftMaxY = -1e20, leftMaxZ = -1e20;

			for (b in 0...NUM_BINS - 1) {
				if (binCounts[b] == 0) continue;
				leftCount += binCounts[b];
				if (binMinX[b] < leftMinX) leftMinX = binMinX[b];
				if (binMinY[b] < leftMinY) leftMinY = binMinY[b];
				if (binMinZ[b] < leftMinZ) leftMinZ = binMinZ[b];
				if (binMaxX[b] > leftMaxX) leftMaxX = binMaxX[b];
				if (binMaxY[b] > leftMaxY) leftMaxY = binMaxY[b];
				if (binMaxZ[b] > leftMaxZ) leftMaxZ = binMaxZ[b];

				var rc = rightCount[b + 1];
				if (rc == 0) continue;

				var leftSA = surfaceArea(leftMinX, leftMinY, leftMinZ, leftMaxX, leftMaxY, leftMaxZ);
				var rightSA = surfaceArea(rightMinX[b + 1], rightMinY[b + 1], rightMinZ[b + 1],
				                          rightMaxX[b + 1], rightMaxY[b + 1], rightMaxZ[b + 1]);
				var cost = TRAVERSAL_COST + (leftSA * leftCount + rightSA * rc) / parentSA;

				if (cost < bestCost) {
					bestCost = cost;
					bestAxis = axis;
					bestSplitBin = b;
				}
			}
		}

		if (bestAxis < 0 || bestCost >= leafCost) {
			node.isLeaf = true;
			node.firstObject = objectIndices.length;
			for (idx in indices) objectIndices.push(idx);
			node.objectCount = n;
			return nodeIdx;
		}

		var lo = centBounds[bestAxis][0];
		var hi = centBounds[bestAxis][1];
		var invRange = NUM_BINS / (hi - lo);
		var splitBin = bestSplitBin;

		var left = [];
		var right = [];
		for (idx in indices) {
			var cv = bestAxis == 0 ? cullData.getCenterX(idx) :
			         bestAxis == 1 ? cullData.getCenterY(idx) :
			                     cullData.getCenterZ(idx);
			var b = Std.int((cv - lo) * invRange);
			if (b < 0) b = 0;
			if (b >= NUM_BINS) b = NUM_BINS - 1;
			if (b <= splitBin) left.push(idx);
			else right.push(idx);
		}

		if (left.length == 0 || right.length == 0) {
			indices.sort(function(a, b): Int {
				var av = bestAxis == 0 ? cullData.getCenterX(a) :
				         bestAxis == 1 ? cullData.getCenterY(a) :
				                     cullData.getCenterZ(a);
				var bv = bestAxis == 0 ? cullData.getCenterX(b) :
				         bestAxis == 1 ? cullData.getCenterY(b) :
				                     cullData.getCenterZ(b);
				if (av < bv) return -1;
				if (av > bv) return 1;
				return 0;
			});
			var mid = n >> 1;
			left = indices.slice(0, mid);
			right = indices.slice(mid);
		}

		node.isLeaf = false;
		node.leftChild = buildRecursive(left, cullData);
		node.rightChild = buildRecursive(right, cullData);

		return nodeIdx;
	}


	public function refit(cullData: CullData) {
		var i = nodes.length - 1;
		while (i >= 0) {
			var node = nodes[i];
			if (node.isLeaf) {
				var minX = 1e20, minY = 1e20, minZ = 1e20;
				var maxX = -1e20, maxY = -1e20, maxZ = -1e20;
				for (j in 0...node.objectCount) {
					var idx = objectIndices[node.firstObject + j];
					if (cullData.getMinX(idx) < minX) minX = cullData.getMinX(idx);
					if (cullData.getMinY(idx) < minY) minY = cullData.getMinY(idx);
					if (cullData.getMinZ(idx) < minZ) minZ = cullData.getMinZ(idx);
					if (cullData.getMaxX(idx) > maxX) maxX = cullData.getMaxX(idx);
					if (cullData.getMaxY(idx) > maxY) maxY = cullData.getMaxY(idx);
					if (cullData.getMaxZ(idx) > maxZ) maxZ = cullData.getMaxZ(idx);
				}
				node.minX = minX; node.minY = minY; node.minZ = minZ;
				node.maxX = maxX; node.maxY = maxY; node.maxZ = maxZ;
			} else {
				var l = nodes[node.leftChild];
				var r = nodes[node.rightChild];
				node.minX = l.minX < r.minX ? l.minX : r.minX;
				node.minY = l.minY < r.minY ? l.minY : r.minY;
				node.minZ = l.minZ < r.minZ ? l.minZ : r.minZ;
				node.maxX = l.maxX > r.maxX ? l.maxX : r.maxX;
				node.maxY = l.maxY > r.maxY ? l.maxY : r.maxY;
				node.maxZ = l.maxZ > r.maxZ ? l.maxZ : r.maxZ;
			}
			node.centerX = (node.minX + node.maxX) * 0.5;
			node.centerY = (node.minY + node.maxY) * 0.5;
			node.centerZ = (node.minZ + node.maxZ) * 0.5;
			i--;
		}
	}

	public function forEachLeafInSubtree(node: BVHNode, cb: (Int) -> Void) {
		if (node.isLeaf) {
			for (i in 0...node.objectCount) {
				cb(objectIndices[node.firstObject + i]);
			}
		} else {
			forEachLeafInSubtree(nodes[node.leftChild], cb);
			forEachLeafInSubtree(nodes[node.rightChild], cb);
		}
	}

	public function traverseFrontToBack(
		camX: Float, camY: Float, camZ: Float,
		shouldSkip: (BVHNode) -> Bool,
		onLeaf: (Int, Int) -> Void
	) {
		if (root < 0) return;

		stack.resize(0);
		stack.push(root);

		while (stack.length > 0) {
			var nodeIdx = stack.pop();
			var node = nodes[nodeIdx];

			if (shouldSkip(node)) continue;

			if (node.isLeaf) {
				var start = node.firstObject;
				for (i in 0...node.objectCount) {
					onLeaf(objectIndices[start + i], i);
				}
			} else {
				var leftNode = nodes[node.leftChild];
				var rightNode = nodes[node.rightChild];
				var leftDist = distSq(leftNode.centerX, leftNode.centerY, leftNode.centerZ, camX, camY, camZ);
				var rightDist = distSq(rightNode.centerX, rightNode.centerY, rightNode.centerZ, camX, camY, camZ);

				if (leftDist > rightDist) {
					stack.push(node.leftChild);
					stack.push(node.rightChild);
				} else {
					stack.push(node.rightChild);
					stack.push(node.leftChild);
				}
			}
		}
	}

	inline function distSq(ax: Float, ay: Float, az: Float, bx: Float, by: Float, bz: Float): Float {
		var dx = ax - bx;
		var dy = ay - by;
		var dz = az - bz;
		return dx * dx + dy * dy + dz * dz;
	}

	public static function aabbInFrustumRaw(
		minX: Float, minY: Float, minZ: Float,
		maxX: Float, maxY: Float, maxZ: Float,
		planes: Array<iron.object.FrustumPlane>
	): Bool {
		for (i in 0...planes.length) {
			var plane = planes[i];
			var px = plane.normal.x >= 0 ? maxX : minX;
			var py = plane.normal.y >= 0 ? maxY : minY;
			var pz = plane.normal.z >= 0 ? maxZ : minZ;
			if (SimdOps.dot3(px, py, pz, plane.normal.x, plane.normal.y, plane.normal.z) + plane.constant < 0) {
				return false;
			}
		}
		return true;
	}

	public static function aabbInFrustum(
		node: BVHNode,
		planes: Array<iron.object.FrustumPlane>
	): Bool {
		for (i in 0...planes.length) {
			var plane = planes[i];
			var px = plane.normal.x >= 0 ? node.maxX : node.minX;
			var py = plane.normal.y >= 0 ? node.maxY : node.minY;
			var pz = plane.normal.z >= 0 ? node.maxZ : node.minZ;
			if (SimdOps.dot3(px, py, pz, plane.normal.x, plane.normal.y, plane.normal.z) + plane.constant < 0) {
				return false;
			}
		}
		return true;
	}
}