package iron.culling;

class SoftwareOcclusionRasterizer {

	public var hiz: HiZPyramid;
	public var width: Int;
	public var height: Int;

	public var nearClip: Float = 0.0001;

	static inline var TILE_SIZE = 32;

	public var trianglesRasterized: Int = 0;

	public var collectMode = false;
	var screenTris: Array<Float> = [];
	var screenTriCount = 0;

	static final BOX_FACES = [
		[0, 1, 3, 2], // -Z
		[4, 5, 7, 6], // +Z
		[0, 1, 5, 4], // -Y
		[2, 3, 7, 6], // +Y
		[0, 2, 6, 4], // -X
		[1, 3, 7, 5], // +X
	];

	var inX: Array<Float> = [0,0,0];
	var inY: Array<Float> = [0,0,0];
	var inZ: Array<Float> = [0,0,0];
	var inW: Array<Float> = [0,0,0];
	var outX: Array<Float> = [0,0,0,0];
	var outY: Array<Float> = [0,0,0,0];
	var outZ: Array<Float> = [0,0,0,0];
	var outW: Array<Float> = [0,0,0,0];
	var projSx: Array<Float> = [0,0,0,0];
	var projSy: Array<Float> = [0,0,0,0];
	var projSz: Array<Float> = [0,0,0,0];

	var boxClipX: Array<Float> = [0,0,0,0,0,0,0,0];
	var boxClipY: Array<Float> = [0,0,0,0,0,0,0,0];
	var boxClipZ: Array<Float> = [0,0,0,0,0,0,0,0];
	var boxClipW: Array<Float> = [0,0,0,0,0,0,0,0];

	public function new(width: Int = 256, height: Int = 128) {
		this.width = width;
		this.height = height;
		hiz = new HiZPyramid(width, height);
	}

	public function clear() {
		hiz.clear(1e20);
		trianglesRasterized = 0;
	}

	public function rasterizeTriangle(
		v0x: Float, v0y: Float, v0z: Float,
		v1x: Float, v1y: Float, v1z: Float,
		v2x: Float, v2y: Float, v2z: Float
	) {
		var e0a = v1y - v0y; var e0b = v0x - v1x;
		var e1a = v2y - v1y; var e1b = v1x - v2x;
		var e2a = v0y - v2y; var e2b = v2x - v0x;
		var e0c = -e0a * v0x - e0b * v0y;
		var e1c = -e1a * v1x - e1b * v1y;
		var e2c = -e2a * v2x - e2b * v2y;

		var area = e0a * v2x + e0b * v2y + e0c;
		if (area < 0) {
			area = -area;
			e0a = -e0a; e0b = -e0b; e0c = -e0c;
			e1a = -e1a; e1b = -e1b; e1c = -e1c;
			e2a = -e2a; e2b = -e2b; e2c = -e2c;
		}
		if (area < 1e-8) return;

		var minX = Std.int(Math.max(0, Math.floor(Math.min(v0x, Math.min(v1x, v2x)))));
		var maxX = Std.int(Math.min(width - 1, Math.ceil(Math.max(v0x, Math.max(v1x, v2x)))));
		var minY = Std.int(Math.max(0, Math.floor(Math.min(v0y, Math.min(v1y, v2y)))));
		var maxY = Std.int(Math.min(height - 1, Math.ceil(Math.max(v0y, Math.max(v1y, v2y)))));
		if (minX > maxX || minY > maxY) return;

		var d = (v1x - v0x) * (v2y - v0y) - (v2x - v0x) * (v1y - v0y);
		var invD = 1.0 / d;
		var dzdx = ((v1z - v0z) * (v2y - v0y) - (v2z - v0z) * (v1y - v0y)) * invD;
		var dzdy = ((v2z - v0z) * (v1x - v0x) - (v1z - v0z) * (v2x - v0x)) * invD;

		var zPush = (Math.abs(dzdx) + Math.abs(dzdy)) * 0.5;
		var zMaxTri = v0z;
		if (v1z > zMaxTri) zMaxTri = v1z;
		if (v2z > zMaxTri) zMaxTri = v2z;

		var cx0 = minX + 0.5;
		var cy0 = minY + 0.5;
		var e0row = e0a * cx0 + e0b * cy0 + e0c;
		var e1row = e1a * cx0 + e1b * cy0 + e1c;
		var e2row = e2a * cx0 + e2b * cy0 + e2c;
		var zrow = v0z + dzdx * (cx0 - v0x) + dzdy * (cy0 - v0y);

		var buf = hiz.mipBuffers[0];
		var w = width;
		for (py in minY...maxY + 1) {
			var e0 = e0row; var e1 = e1row; var e2 = e2row;
			var z = zrow;
			var idx = py * w + minX;
			for (px in minX...maxX + 1) {
				if (e0 >= 0 && e1 >= 0 && e2 >= 0) {
					var zPix = z + zPush;
					if (zPix > zMaxTri) zPix = zMaxTri;
					if (zPix < 0.0) zPix = 0.0;
					if (zPix > 1.0) zPix = 1.0;
					if (zPix < buf[idx]) buf[idx] = zPix;
				}
				e0 += e0a; e1 += e1a; e2 += e2a;
				z += dzdx;
				idx++;
			}
			e0row += e0b; e1row += e1b; e2row += e2b;
			zrow += dzdy;
		}
	}

	function clipNearAndRasterize(
		cx0: Float, cy0: Float, cz0: Float, cw0: Float,
		cx1: Float, cy1: Float, cz1: Float, cw1: Float,
		cx2: Float, cy2: Float, cz2: Float, cw2: Float
	) {
		inX[0] = cx0; inY[0] = cy0; inZ[0] = cz0; inW[0] = cw0;
		inX[1] = cx1; inY[1] = cy1; inZ[1] = cz1; inW[1] = cw1;
		inX[2] = cx2; inY[2] = cy2; inZ[2] = cz2; inW[2] = cw2;

		var outCount = 0;
		for (i in 0...3) {
			var aIdx = (i + 2) % 3; 
			var bIdx = i;
			var aInside = inW[aIdx] >= nearClip;
			var bInside = inW[bIdx] >= nearClip;

			if (aInside) {
				if (bInside) {
					outX[outCount] = inX[bIdx]; outY[outCount] = inY[bIdx];
					outZ[outCount] = inZ[bIdx]; outW[outCount] = inW[bIdx];
					outCount++;
				} else {
					var t = (inW[aIdx] - nearClip) / (inW[aIdx] - inW[bIdx]);
					outX[outCount] = inX[aIdx] + t * (inX[bIdx] - inX[aIdx]);
					outY[outCount] = inY[aIdx] + t * (inY[bIdx] - inY[aIdx]);
					outZ[outCount] = inZ[aIdx] + t * (inZ[bIdx] - inZ[aIdx]);
					outW[outCount] = inW[aIdx] + t * (inW[bIdx] - inW[aIdx]);
					outCount++;
				}
			} else {
				if (bInside) {
					var t = (inW[aIdx] - nearClip) / (inW[aIdx] - inW[bIdx]);
					outX[outCount] = inX[aIdx] + t * (inX[bIdx] - inX[aIdx]);
					outY[outCount] = inY[aIdx] + t * (inY[bIdx] - inY[aIdx]);
					outZ[outCount] = inZ[aIdx] + t * (inZ[bIdx] - inZ[aIdx]);
					outW[outCount] = inW[aIdx] + t * (inW[bIdx] - inW[aIdx]);
					outCount++;
					outX[outCount] = inX[bIdx]; outY[outCount] = inY[bIdx];
					outZ[outCount] = inZ[bIdx]; outW[outCount] = inW[bIdx];
					outCount++;
				}
			}
		}

		if (outCount < 3) return;

		for (i in 0...outCount) {
			var invW = SimdOps.reciprocal(outW[i]);
			projSx[i] = (outX[i] * invW * 0.5 + 0.5) * width;
			projSy[i] = (outY[i] * invW * 0.5 + 0.5) * height;
			projSz[i] = outZ[i] * invW;
		}

		var signedArea = (projSx[1] - projSx[0]) * (projSy[2] - projSy[0]) -
		                 (projSy[1] - projSy[0]) * (projSx[2] - projSx[0]);
		if (Math.abs(signedArea) < 1e-8) return;

		trianglesRasterized++;

		if (collectMode) {
			screenTris[screenTriCount++] = projSx[0]; screenTris[screenTriCount++] = projSy[0]; screenTris[screenTriCount++] = projSz[0];
			screenTris[screenTriCount++] = projSx[1]; screenTris[screenTriCount++] = projSy[1]; screenTris[screenTriCount++] = projSz[1];
			screenTris[screenTriCount++] = projSx[2]; screenTris[screenTriCount++] = projSy[2]; screenTris[screenTriCount++] = projSz[2];
			if (outCount == 4) {
				screenTris[screenTriCount++] = projSx[0]; screenTris[screenTriCount++] = projSy[0]; screenTris[screenTriCount++] = projSz[0];
				screenTris[screenTriCount++] = projSx[2]; screenTris[screenTriCount++] = projSy[2]; screenTris[screenTriCount++] = projSz[2];
				screenTris[screenTriCount++] = projSx[3]; screenTris[screenTriCount++] = projSy[3]; screenTris[screenTriCount++] = projSz[3];
			}
		} else {
			rasterizeTriangle(
				projSx[0], projSy[0], projSz[0],
				projSx[1], projSy[1], projSz[1],
				projSx[2], projSy[2], projSz[2]
			);
			if (outCount == 4) {
				rasterizeTriangle(
					projSx[0], projSy[0], projSz[0],
					projSx[2], projSy[2], projSz[2],
					projSx[3], projSy[3], projSz[3]
				);
			}
		}
	}

	public function rasterizeProxy(
		positions: kha.arrays.Float32Array,
		indices: kha.arrays.Uint32Array,
		vp: iron.math.Mat4,
		screenW: Int, screenH: Int
	) {
		var triCount = Std.int(indices.length / 3);

		var m00 = vp._00, m10 = vp._10, m20 = vp._20, m30 = vp._30;
		var m01 = vp._01, m11 = vp._11, m21 = vp._21, m31 = vp._31;
		var m02 = vp._02, m12 = vp._12, m22 = vp._22, m32 = vp._32;
		var m03 = vp._03, m13 = vp._13, m23 = vp._23, m33 = vp._33;

		for (t in 0...triCount) {
			var i0 = indices[t * 3] * 3;
			var i1 = indices[t * 3 + 1] * 3;
			var i2 = indices[t * 3 + 2] * 3;

			var px0 = positions[i0], py0 = positions[i0 + 1], pz0 = positions[i0 + 2];
			var px1 = positions[i1], py1 = positions[i1 + 1], pz1 = positions[i1 + 2];
			var px2 = positions[i2], py2 = positions[i2 + 1], pz2 = positions[i2 + 2];

			var cw0 = m03 * px0 + m13 * py0 + m23 * pz0 + m33;
			var cw1 = m03 * px1 + m13 * py1 + m23 * pz1 + m33;
			var cw2 = m03 * px2 + m13 * py2 + m23 * pz2 + m33;

			if (cw0 < nearClip && cw1 < nearClip && cw2 < nearClip) continue;

			var cx0 = m00 * px0 + m10 * py0 + m20 * pz0 + m30;
			var cy0 = m01 * px0 + m11 * py0 + m21 * pz0 + m31;
			var cz0 = m02 * px0 + m12 * py0 + m22 * pz0 + m32;

			var cx1 = m00 * px1 + m10 * py1 + m20 * pz1 + m30;
			var cy1 = m01 * px1 + m11 * py1 + m21 * pz1 + m31;
			var cz1 = m02 * px1 + m12 * py1 + m22 * pz1 + m32;

			var cx2 = m00 * px2 + m10 * py2 + m20 * pz2 + m30;
			var cy2 = m01 * px2 + m11 * py2 + m21 * pz2 + m31;
			var cz2 = m02 * px2 + m12 * py2 + m22 * pz2 + m32;

			clipNearAndRasterize(
				cx0, cy0, cz0, cw0,
				cx1, cy1, cz1, cw1,
				cx2, cy2, cz2, cw2
			);
		}
	}

	public function rasterizeProxyTransformed(
		clipPositions: kha.arrays.Float32Array,
		indices: kha.arrays.Uint32Array,
		screenW: Int, screenH: Int
	) {
		var triCount = Std.int(indices.length / 3);

		for (t in 0...triCount) {
			var i0 = indices[t * 3] * 4;
			var i1 = indices[t * 3 + 1] * 4;
			var i2 = indices[t * 3 + 2] * 4;

			var cx0 = clipPositions[i0], cy0 = clipPositions[i0 + 1], cz0 = clipPositions[i0 + 2], cw0 = clipPositions[i0 + 3];
			var cx1 = clipPositions[i1], cy1 = clipPositions[i1 + 1], cz1 = clipPositions[i1 + 2], cw1 = clipPositions[i1 + 3];
			var cx2 = clipPositions[i2], cy2 = clipPositions[i2 + 1], cz2 = clipPositions[i2 + 2], cw2 = clipPositions[i2 + 3];

			if (cw0 < nearClip && cw1 < nearClip && cw2 < nearClip) continue;

			clipNearAndRasterize(
				cx0, cy0, cz0, cw0,
				cx1, cy1, cz1, cw1,
				cx2, cy2, cz2, cw2
			);
		}
	}

	public function rasterizeBox(
		minX: Float, minY: Float, minZ: Float,
		maxX: Float, maxY: Float, maxZ: Float,
		vp: iron.math.Mat4,
		screenW: Int, screenH: Int
	) {
		var m00 = vp._00, m10 = vp._10, m20 = vp._20, m30 = vp._30;
		var m01 = vp._01, m11 = vp._11, m21 = vp._21, m31 = vp._31;
		var m02 = vp._02, m12 = vp._12, m22 = vp._22, m32 = vp._32;
		var m03 = vp._03, m13 = vp._13, m23 = vp._23, m33 = vp._33;

		for (c in 0...8) {
			var x = (c & 1) != 0 ? maxX : minX;
			var y = (c & 2) != 0 ? maxY : minY;
			var z = (c & 4) != 0 ? maxZ : minZ;
			boxClipX[c] = m00 * x + m10 * y + m20 * z + m30;
			boxClipY[c] = m01 * x + m11 * y + m21 * z + m31;
			boxClipZ[c] = m02 * x + m12 * y + m22 * z + m32;
			boxClipW[c] = m03 * x + m13 * y + m23 * z + m33;
		}

		for (face in BOX_FACES) {
			clipNearAndRasterize(
				boxClipX[face[0]], boxClipY[face[0]], boxClipZ[face[0]], boxClipW[face[0]],
				boxClipX[face[1]], boxClipY[face[1]], boxClipZ[face[1]], boxClipW[face[1]],
				boxClipX[face[2]], boxClipY[face[2]], boxClipZ[face[2]], boxClipW[face[2]]
			);
			clipNearAndRasterize(
				boxClipX[face[0]], boxClipY[face[0]], boxClipZ[face[0]], boxClipW[face[0]],
				boxClipX[face[2]], boxClipY[face[2]], boxClipZ[face[2]], boxClipW[face[2]],
				boxClipX[face[3]], boxClipY[face[3]], boxClipZ[face[3]], boxClipW[face[3]]
			);
		}
	}

	public function buildHiZ() {
		hiz.buildPyramid();
	}

	public function beginCollect() {
		collectMode = true;
		screenTriCount = 0;
	}

	var tileTriCounts: Array<Int> = [];
	var tileTriOffsets: Array<Int> = [];
	var tileTriLists: Array<Array<Int>> = [];

	public function flushCollected() {
		collectMode = false;
		var nTris = Std.int(screenTriCount / 9);
		if (nTris == 0) return;

		var tilesX = Std.int(Math.ceil(width / TILE_SIZE));
		var tilesY = Std.int(Math.ceil(height / TILE_SIZE));
		var totalTiles = tilesX * tilesY;

		if (tileTriCounts.length < totalTiles) {
			tileTriCounts.resize(totalTiles);
			tileTriOffsets.resize(totalTiles);
			tileTriLists.resize(totalTiles);
		}
		for (t in 0...totalTiles) tileTriCounts[t] = 0;
		for (t in 0...nTris) {
			var off = t * 9;
			var tx0 = Std.int(Math.max(0, Math.floor(Math.min(
				Math.min(screenTris[off], screenTris[off + 3]), screenTris[off + 6]) / TILE_SIZE)));
			var tx1 = Std.int(Math.min(tilesX - 1, Math.ceil(Math.max(
				Math.max(screenTris[off], screenTris[off + 3]), screenTris[off + 6]) / TILE_SIZE)));
			var ty0 = Std.int(Math.max(0, Math.floor(Math.min(
				Math.min(screenTris[off + 1], screenTris[off + 4]), screenTris[off + 7]) / TILE_SIZE)));
			var ty1 = Std.int(Math.min(tilesY - 1, Math.ceil(Math.max(
				Math.max(screenTris[off + 1], screenTris[off + 4]), screenTris[off + 7]) / TILE_SIZE)));
			for (ty in ty0...ty1 + 1) {
				for (tx in tx0...tx1 + 1) {
					tileTriCounts[ty * tilesX + tx]++;
				}
			}
		}

		for (t in 0...totalTiles) {
			var l = tileTriLists[t];
			if (l == null) { l = []; tileTriLists[t] = l; }
			if (l.length < tileTriCounts[t]) l.resize(tileTriCounts[t]);
			tileTriOffsets[t] = 0;
		}
		for (t in 0...nTris) {
			var off = t * 9;
			var tx0 = Std.int(Math.max(0, Math.floor(Math.min(
				Math.min(screenTris[off], screenTris[off + 3]), screenTris[off + 6]) / TILE_SIZE)));
			var tx1 = Std.int(Math.min(tilesX - 1, Math.ceil(Math.max(
				Math.max(screenTris[off], screenTris[off + 3]), screenTris[off + 6]) / TILE_SIZE)));
			var ty0 = Std.int(Math.max(0, Math.floor(Math.min(
				Math.min(screenTris[off + 1], screenTris[off + 4]), screenTris[off + 7]) / TILE_SIZE)));
			var ty1 = Std.int(Math.min(tilesY - 1, Math.ceil(Math.max(
				Math.max(screenTris[off + 1], screenTris[off + 4]), screenTris[off + 7]) / TILE_SIZE)));
			for (ty in ty0...ty1 + 1) {
				for (tx in tx0...tx1 + 1) {
					var ti = ty * tilesX + tx;
					tileTriLists[ti][tileTriOffsets[ti]++] = t;
				}
			}
		}

		var buf = hiz.mipBuffers[0];
		var w = width;
		var tris = screenTris;

		#if (target.threaded)
		JobSystem.parallelFor(totalTiles, function(start, end) {
			for (ti in start...end) {
				var tx = ti % tilesX;
				var ty = Std.int(ti / tilesX);
				var px0 = tx * TILE_SIZE;
				var py0 = ty * TILE_SIZE;
				var px1 = Std.int(Math.min(w - 1, px0 + TILE_SIZE - 1));
				var py1 = Std.int(Math.min(height - 1, py0 + TILE_SIZE - 1));
				var list = tileTriLists[ti];
				var listLen = tileTriCounts[ti];
				for (li in 0...listLen) {
					var off = list[li] * 9;
					rasterizeTriangleBounded(
						tris[off], tris[off + 1], tris[off + 2],
						tris[off + 3], tris[off + 4], tris[off + 5],
						tris[off + 6], tris[off + 7], tris[off + 8],
						px0, py0, px1, py1, buf, w
					);
				}
			}
		});
		#else
		for (ti in 0...totalTiles) {
			var tx = ti % tilesX;
			var ty = Std.int(ti / tilesX);
			var px0 = tx * TILE_SIZE;
			var py0 = ty * TILE_SIZE;
			var px1 = Std.int(Math.min(w - 1, px0 + TILE_SIZE - 1));
			var py1 = Std.int(Math.min(height - 1, py0 + TILE_SIZE - 1));
			var list = tileTriLists[ti];
			for (li in 0...list.length) {
				var off = list[li] * 9;
				rasterizeTriangleBounded(
					tris[off], tris[off + 1], tris[off + 2],
					tris[off + 3], tris[off + 4], tris[off + 5],
					tris[off + 6], tris[off + 7], tris[off + 8],
					px0, py0, px1, py1, buf, w
				);
			}
		}
		#end
	}

	function rasterizeTriangleBounded(
		v0x: Float, v0y: Float, v0z: Float,
		v1x: Float, v1y: Float, v1z: Float,
		v2x: Float, v2y: Float, v2z: Float,
		tileX0: Int, tileY0: Int, tileX1: Int, tileY1: Int,
		buf: kha.arrays.Float32Array, w: Int
	) {
		var e0a = v1y - v0y; var e0b = v0x - v1x;
		var e1a = v2y - v1y; var e1b = v1x - v2x;
		var e2a = v0y - v2y; var e2b = v2x - v0x;
		var e0c = -e0a * v0x - e0b * v0y;
		var e1c = -e1a * v1x - e1b * v1y;
		var e2c = -e2a * v2x - e2b * v2y;

		var area = e0a * v2x + e0b * v2y + e0c;
		if (area < 0) {
			area = -area;
			e0a = -e0a; e0b = -e0b; e0c = -e0c;
			e1a = -e1a; e1b = -e1b; e1c = -e1c;
			e2a = -e2a; e2b = -e2b; e2c = -e2c;
		}
		if (area < 1e-8) return;

		var minX = Std.int(Math.max(tileX0, Math.floor(Math.min(v0x, Math.min(v1x, v2x)))));
		var maxX = Std.int(Math.min(tileX1, Math.ceil(Math.max(v0x, Math.max(v1x, v2x)))));
		var minY = Std.int(Math.max(tileY0, Math.floor(Math.min(v0y, Math.min(v1y, v2y)))));
		var maxY = Std.int(Math.min(tileY1, Math.ceil(Math.max(v0y, Math.max(v1y, v2y)))));
		if (minX > maxX || minY > maxY) return;

		var d = (v1x - v0x) * (v2y - v0y) - (v2x - v0x) * (v1y - v0y);
		var invD = 1.0 / d;
		var dzdx = ((v1z - v0z) * (v2y - v0y) - (v2z - v0z) * (v1y - v0y)) * invD;
		var dzdy = ((v2z - v0z) * (v1x - v0x) - (v1z - v0z) * (v2x - v0x)) * invD;

		var zPush = (Math.abs(dzdx) + Math.abs(dzdy)) * 0.5;
		var zMaxTri = v0z;
		if (v1z > zMaxTri) zMaxTri = v1z;
		if (v2z > zMaxTri) zMaxTri = v2z;

		var cx0 = minX + 0.5;
		var cy0 = minY + 0.5;
		var e0row = e0a * cx0 + e0b * cy0 + e0c;
		var e1row = e1a * cx0 + e1b * cy0 + e1c;
		var e2row = e2a * cx0 + e2b * cy0 + e2c;
		var zrow = v0z + dzdx * (cx0 - v0x) + dzdy * (cy0 - v0y);

		for (py in minY...maxY + 1) {
			var e0 = e0row; var e1 = e1row; var e2 = e2row;
			var z = zrow;
			var idx = py * w + minX;
			for (px in minX...maxX + 1) {
				if (e0 >= 0 && e1 >= 0 && e2 >= 0) {
					var zPix = z + zPush;
					if (zPix > zMaxTri) zPix = zMaxTri;
					if (zPix > 1.0) zPix = 1.0;
					if (zPix < buf[idx]) { buf[idx] = zPix; }
				}
				e0 += e0a; e1 += e1a; e2 += e2a;
				z += dzdx;
				idx++;
			}
			e0row += e0b; e1row += e1b; e2row += e2b;
			zrow += dzdy;
		}
	}

	public function isOccluded(
		minX: Float, minY: Float, minZ: Float,
		maxX: Float, maxY: Float, maxZ: Float,
		vp: iron.math.Mat4,
		screenW: Int, screenH: Int
	): Bool {
		return hiz.isOccluded(minX, minY, minZ, maxX, maxY, maxZ, vp, screenW, screenH);
	}
}