main.js

let cameras = [
	{
		id: 0,
		img_name: "00001",
		width: 1959,
		height: 1090,
		position: [
			-3.0089893469241797, -0.11086489695181866, -3.7527640949141428,
		],
		rotation: [
			[0.876134201218856, 0.06925962026449776, 0.47706599800804744],
			[-0.04747421839895102, 0.9972110940209488, -0.057586739349882114],
			[-0.4797239414934443, 0.027805376500959853, 0.8769787916452908],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
	{
		id: 1,
		img_name: "00009",
		width: 1959,
		height: 1090,
		position: [
			-2.5199776022057296, -0.09704735754873686, -3.6247725540304545,
		],
		rotation: [
			[0.9982731285632193, -0.011928707708098955, -0.05751927260507243],
			[0.0065061360949636325, 0.9955928229282383, -0.09355533724430458],
			[0.058381769258182864, 0.09301955098900708, 0.9939511719154457],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
	{
		id: 2,
		img_name: "00017",
		width: 1959,
		height: 1090,
		position: [
			-0.7737533667465242, -0.3364271945329695, -2.9358969417573753,
		],
		rotation: [
			[0.9998813418672372, 0.013742375651625236, -0.0069605529394208224],
			[-0.014268370388586709, 0.996512943252834, -0.08220929105659476],
			[0.00580653013657589, 0.08229885200307129, 0.9965907801935302],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
	{
		id: 3,
		img_name: "00025",
		width: 1959,
		height: 1090,
		position: [
			1.2198221749590001, -0.2196687861401182, -2.3183162007028453,
		],
		rotation: [
			[0.9208648867765482, 0.0012010625395201253, 0.389880004297208],
			[-0.06298204172269357, 0.987319521752825, 0.14571693239364383],
			[-0.3847611242348369, -0.1587410451475895, 0.9092635249821667],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
	{
		id: 4,
		img_name: "00033",
		width: 1959,
		height: 1090,
		position: [
			1.742387858893817, -0.13848225198886954, -2.0566370113193146,
		],
		rotation: [
			[0.24669889292141334, -0.08370189346592856, -0.9654706879349405],
			[0.11343747891376445, 0.9919082664242816, -0.05700815184573074],
			[0.9624300466054861, -0.09545671285663988, 0.2541976029815521],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
	{
		id: 5,
		img_name: "00041",
		width: 1959,
		height: 1090,
		position: [
			3.6567309419223935, -0.16470990600750707, -1.3458085590422042,
		],
		rotation: [
			[0.2341293058324528, -0.02968330457755884, -0.9717522161434825],
			[0.10270823606832301, 0.99469554638321, -0.005638106875665722],
			[0.9667649592295676, -0.09848690996657204, 0.2359360976431732],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
	{
		id: 6,
		img_name: "00049",
		width: 1959,
		height: 1090,
		position: [
			3.9013554243203497, -0.2597500978038105, -0.8106154188297828,
		],
		rotation: [
			[0.6717235545638952, -0.015718162115524837, -0.7406351366386528],
			[0.055627354673906296, 0.9980224478387622, 0.029270992841185218],
			[0.7387104058127439, -0.060861588786650656, 0.6712695459756353],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
	{
		id: 7,
		img_name: "00057",
		width: 1959,
		height: 1090,
		position: [4.742994605467533, -0.05591660945412069, 0.9500365976084458],
		rotation: [
			[-0.17042655709210375, 0.01207080756938, -0.9852964448542146],
			[0.1165090336695526, 0.9931575292530063, -0.00798543433078162],
			[0.9784581921120181, -0.1161568667478904, -0.1706667764862097],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
	{
		id: 8,
		img_name: "00065",
		width: 1959,
		height: 1090,
		position: [4.34676307626522, 0.08168160516967145, 1.0876221470355405],
		rotation: [
			[-0.003575447631888379, -0.044792503246552894, -0.9989899137764799],
			[0.10770152645126597, 0.9931680875192705, -0.04491693593046672],
			[0.9941768441149182, -0.10775333677534978, 0.0012732004866391048],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
	{
		id: 9,
		img_name: "00073",
		width: 1959,
		height: 1090,
		position: [3.264984351114202, 0.078974937336732, 1.0117200284114904],
		rotation: [
			[-0.026919994628162257, -0.1565891128261527, -0.9872968974090509],
			[0.08444552208239385, 0.983768234577625, -0.1583319754069128],
			[0.9960643893290491, -0.0876350978794554, -0.013259786205163005],
		],
		fy: 1164.6601287484507,
		fx: 1159.5880733038064,
	},
];

const camera = cameras[0];

function getProjectionMatrix(fx, fy, width, height) {
	const znear = 0.2;
	const zfar = 200;
	return [
		[(2 * fx) / width, 0, 0, 0],
		[0, -(2 * fy) / height, 0, 0],
		[0, 0, zfar / (zfar - znear), 1],
		[0, 0, -(zfar * znear) / (zfar - znear), 0],
	].flat();
}

function getViewMatrix(camera) {
	const R = camera.rotation.flat();
	const t = camera.position;
	const camToWorld = [
		[R[0], R[1], R[2], 0],
		[R[3], R[4], R[5], 0],
		[R[6], R[7], R[8], 0],
		[
			-t[0] * R[0] - t[1] * R[3] - t[2] * R[6],
			-t[0] * R[1] - t[1] * R[4] - t[2] * R[7],
			-t[0] * R[2] - t[1] * R[5] - t[2] * R[8],
			1,
		],
	].flat();
	return camToWorld;
}

function multiply4(a, b) {
	return [
		b[0] * a[0] + b[1] * a[4] + b[2] * a[8] + b[3] * a[12],
		b[0] * a[1] + b[1] * a[5] + b[2] * a[9] + b[3] * a[13],
		b[0] * a[2] + b[1] * a[6] + b[2] * a[10] + b[3] * a[14],
		b[0] * a[3] + b[1] * a[7] + b[2] * a[11] + b[3] * a[15],
		b[4] * a[0] + b[5] * a[4] + b[6] * a[8] + b[7] * a[12],
		b[4] * a[1] + b[5] * a[5] + b[6] * a[9] + b[7] * a[13],
		b[4] * a[2] + b[5] * a[6] + b[6] * a[10] + b[7] * a[14],
		b[4] * a[3] + b[5] * a[7] + b[6] * a[11] + b[7] * a[15],
		b[8] * a[0] + b[9] * a[4] + b[10] * a[8] + b[11] * a[12],
		b[8] * a[1] + b[9] * a[5] + b[10] * a[9] + b[11] * a[13],
		b[8] * a[2] + b[9] * a[6] + b[10] * a[10] + b[11] * a[14],
		b[8] * a[3] + b[9] * a[7] + b[10] * a[11] + b[11] * a[15],
		b[12] * a[0] + b[13] * a[4] + b[14] * a[8] + b[15] * a[12],
		b[12] * a[1] + b[13] * a[5] + b[14] * a[9] + b[15] * a[13],
		b[12] * a[2] + b[13] * a[6] + b[14] * a[10] + b[15] * a[14],
		b[12] * a[3] + b[13] * a[7] + b[14] * a[11] + b[15] * a[15],
	];
}

function invert4(a) {
	let b00 = a[0] * a[5] - a[1] * a[4];
	let b01 = a[0] * a[6] - a[2] * a[4];
	let b02 = a[0] * a[7] - a[3] * a[4];
	let b03 = a[1] * a[6] - a[2] * a[5];
	let b04 = a[1] * a[7] - a[3] * a[5];
	let b05 = a[2] * a[7] - a[3] * a[6];
	let b06 = a[8] * a[13] - a[9] * a[12];
	let b07 = a[8] * a[14] - a[10] * a[12];
	let b08 = a[8] * a[15] - a[11] * a[12];
	let b09 = a[9] * a[14] - a[10] * a[13];
	let b10 = a[9] * a[15] - a[11] * a[13];
	let b11 = a[10] * a[15] - a[11] * a[14];
	let det =
		b00 * b11 - b01 * b10 + b02 * b09 + b03 * b08 - b04 * b07 + b05 * b06;
	if (!det) return null;
	return [
		(a[5] * b11 - a[6] * b10 + a[7] * b09) / det,
		(a[2] * b10 - a[1] * b11 - a[3] * b09) / det,
		(a[13] * b05 - a[14] * b04 + a[15] * b03) / det,
		(a[10] * b04 - a[9] * b05 - a[11] * b03) / det,
		(a[6] * b08 - a[4] * b11 - a[7] * b07) / det,
		(a[0] * b11 - a[2] * b08 + a[3] * b07) / det,
		(a[14] * b02 - a[12] * b05 - a[15] * b01) / det,
		(a[8] * b05 - a[10] * b02 + a[11] * b01) / det,
		(a[4] * b10 - a[5] * b08 + a[7] * b06) / det,
		(a[1] * b08 - a[0] * b10 - a[3] * b06) / det,
		(a[12] * b04 - a[13] * b02 + a[15] * b00) / det,
		(a[9] * b02 - a[8] * b04 - a[11] * b00) / det,
		(a[5] * b07 - a[4] * b09 - a[6] * b06) / det,
		(a[0] * b09 - a[1] * b07 + a[2] * b06) / det,
		(a[13] * b01 - a[12] * b03 - a[14] * b00) / det,
		(a[8] * b03 - a[9] * b01 + a[10] * b00) / det,
	];
}

function rotate4(a, rad, x, y, z) {
	let len = Math.hypot(x, y, z);
	x /= len;
	y /= len;
	z /= len;
	let s = Math.sin(rad);
	let c = Math.cos(rad);
	let t = 1 - c;
	let b00 = x * x * t + c;
	let b01 = y * x * t + z * s;
	let b02 = z * x * t - y * s;
	let b10 = x * y * t - z * s;
	let b11 = y * y * t + c;
	let b12 = z * y * t + x * s;
	let b20 = x * z * t + y * s;
	let b21 = y * z * t - x * s;
	let b22 = z * z * t + c;
	return [
		a[0] * b00 + a[4] * b01 + a[8] * b02,
		a[1] * b00 + a[5] * b01 + a[9] * b02,
		a[2] * b00 + a[6] * b01 + a[10] * b02,
		a[3] * b00 + a[7] * b01 + a[11] * b02,
		a[0] * b10 + a[4] * b11 + a[8] * b12,
		a[1] * b10 + a[5] * b11 + a[9] * b12,
		a[2] * b10 + a[6] * b11 + a[10] * b12,
		a[3] * b10 + a[7] * b11 + a[11] * b12,
		a[0] * b20 + a[4] * b21 + a[8] * b22,
		a[1] * b20 + a[5] * b21 + a[9] * b22,
		a[2] * b20 + a[6] * b21 + a[10] * b22,
		a[3] * b20 + a[7] * b21 + a[11] * b22,
		...a.slice(12, 16),
	];
}

function translate4(a, x, y, z) {
	return [
		...a.slice(0, 12),
		a[0] * x + a[4] * y + a[8] * z + a[12],
		a[1] * x + a[5] * y + a[9] * z + a[13],
		a[2] * x + a[6] * y + a[10] * z + a[14],
		a[3] * x + a[7] * y + a[11] * z + a[15],
	];
}

function createWorker(self) {
	let buffer;
	let vertexCount = 0;
	let viewProj;
	// 6*4 + 4 + 4 = 8*4
	// XYZ - Position (Float32)
	// XYZ - Scale (Float32)
	// RGBA - colors (uint8)
	// IJKL - quaternion/rot (uint8)
	const rowLength = 3 * 4 + 3 * 4 + 4 + 4;
	let depthMix = new BigInt64Array();
	let lastProj = [];

	const runSort = (viewProj) => {
		if (!buffer) return;

		

		const f_buffer = new Float32Array(buffer);
		const u_buffer = new Uint8Array(buffer);

		const covA = new Float32Array(3 * vertexCount);
		const covB = new Float32Array(3 * vertexCount);

		const center = new Float32Array(3 * vertexCount);
		const color = new Float32Array(4 * vertexCount);

		if (depthMix.length !== vertexCount) {
			depthMix = new BigInt64Array(vertexCount);
			const indexMix = new Uint32Array(depthMix.buffer);
			for (let j = 0; j < vertexCount; j++) {
				indexMix[2 * j] = j;
			}
		} else {
			let dot =
				lastProj[2] * viewProj[2] +
				lastProj[6] * viewProj[6] +
				lastProj[10] * viewProj[10];
			if (Math.abs(dot - 1) < 0.01) {
				return;
			}
		}
		// console.time("sort");

		const floatMix = new Float32Array(depthMix.buffer);
		const indexMix = new Uint32Array(depthMix.buffer);

		for (let j = 0; j < vertexCount; j++) {
			let i = indexMix[2 * j];
			floatMix[2 * j + 1] =
				10000 +
				viewProj[2] * f_buffer[8 * i + 0] +
				viewProj[6] * f_buffer[8 * i + 1] +
				viewProj[10] * f_buffer[8 * i + 2];
		}

		lastProj = viewProj;

		depthMix.sort();

		for (let j = 0; j < vertexCount; j++) {
			const i = indexMix[2 * j];


			center[3 * j + 0] = f_buffer[8 * i + 0];
			center[3 * j + 1] = f_buffer[8 * i + 1];
			center[3 * j + 2] = f_buffer[8 * i + 2];

			color[4 * j + 0] = u_buffer[32 * i + 24 + 0] / 255;
			color[4 * j + 1] = u_buffer[32 * i + 24 + 1] / 255;
			color[4 * j + 2] = u_buffer[32 * i + 24 + 2] / 255;
			color[4 * j + 3] = u_buffer[32 * i + 24 + 3] / 255;

			let scale = [ f_buffer[8 * i + 3 + 0], f_buffer[8 * i + 3 + 1],  f_buffer[8 * i + 3 + 2]];
			let rot = [(u_buffer[32 * i + 28 + 0] - 128) / 128, (u_buffer[32 * i + 28 + 1] - 128) / 128, (u_buffer[32 * i + 28 + 2] - 128) / 128, (u_buffer[32 * i + 28 + 3] - 128) / 128]

			const R = [
				1.0 - 2.0 * (rot[2] * rot[2] + rot[3] * rot[3]),
				2.0 * (rot[1] * rot[2] + rot[0] * rot[3]),
				2.0 * (rot[1] * rot[3] - rot[0] * rot[2]),

				2.0 * (rot[1] * rot[2] - rot[0] * rot[3]),
				1.0 - 2.0 * (rot[1] * rot[1] + rot[3] * rot[3]),
				2.0 * (rot[2] * rot[3] + rot[0] * rot[1]),

				2.0 * (rot[1] * rot[3] + rot[0] * rot[2]),
				2.0 * (rot[2] * rot[3] - rot[0] * rot[1]),
				1.0 - 2.0 * (rot[1] * rot[1] + rot[2] * rot[2]),
			];

			// Compute the matrix product of S and R (M = S * R)
			const M = [
				scale[0] * R[0],
				scale[0] * R[1],
				scale[0] * R[2],
				scale[1] * R[3],
				scale[1] * R[4],
				scale[1] * R[5],
				scale[2] * R[6],
				scale[2] * R[7],
				scale[2] * R[8],
			];

			
			covA[3 * j + 0] = M[0] * M[0] + M[3] * M[3] + M[6] * M[6];
			covA[3 * j + 1] = M[0] * M[1] + M[3] * M[4] + M[6] * M[7];
			covA[3 * j + 2] = M[0] * M[2] + M[3] * M[5] + M[6] * M[8];
			covB[3 * j + 0] = M[1] * M[1] + M[4] * M[4] + M[7] * M[7];
			covB[3 * j + 1] = M[1] * M[2] + M[4] * M[5] + M[7] * M[8];
			covB[3 * j + 2] = M[2] * M[2] + M[5] * M[5] + M[8] * M[8];
		}

		self.postMessage({ covA, center, color, covB, viewProj }, [
			covA.buffer,
			center.buffer,
			color.buffer,
			covB.buffer,
		]);

		// console.timeEnd("sort");
	};

	function processPlyBuffer(inputBuffer) {
		const ubuf = new Uint8Array(inputBuffer);
		// 10KB ought to be enough for a header...
		const header = new TextDecoder().decode(ubuf.slice(0, 1024 * 10));
		const header_end = "end_header\n";
		const header_end_index = header.indexOf(header_end);
		if (header_end_index < 0)
			throw new Error("Unable to read .ply file header");
		const vertexCount = parseInt(/element vertex (\d+)\n/.exec(header)[1]);
		console.log("Vertex Count", vertexCount);
		let row_offset = 0,
			offsets = {},
			types = {};
		const TYPE_MAP = {
			double: "getFloat64",
			int: "getInt32",
			uint: "getUint32",
			float: "getFloat32",
			short: "getInt16",
			ushort: "getUint16",
			uchar: "getUint8",
		};
		for (let prop of header
			.slice(0, header_end_index)
			.split("\n")
			.filter((k) => k.startsWith("property "))) {
			const [p, type, name] = prop.split(" ");
			const arrayType = TYPE_MAP[type] || "getInt8";
			types[name] = arrayType;
			offsets[name] = row_offset;
			row_offset += parseInt(arrayType.replace(/[^\d]/g, "")) / 8;
		}
		console.log("Bytes per row", row_offset, types, offsets);

		let dataView = new DataView(
			inputBuffer,
			header_end_index + header_end.length,
		);
		let row = 0;
		const attrs = new Proxy(
			{},
			{
				get(target, prop) {
					if (!types[prop]) throw new Error(prop + " not found");
					return dataView[types[prop]](
						row * row_offset + offsets[prop],
						true,
					);
				},
			},
		);

		console.time("calculate importance");
		let sizeList = new Float32Array(vertexCount);
		let sizeIndex = new Uint32Array(vertexCount);
		for (row = 0; row < vertexCount; row++) {
			sizeIndex[row] = row;
			if (!types["scale_0"]) continue;
			const size =
				Math.exp(attrs.scale_0) *
				Math.exp(attrs.scale_1) *
				Math.exp(attrs.scale_2);
			const opacity = 1 / (1 + Math.exp(-attrs.opacity));
			sizeList[row] = size * opacity;
		}
		console.timeEnd("calculate importance");

		console.time("sort");
		sizeIndex.sort((b, a) => sizeList[a] - sizeList[b]);
		console.timeEnd("sort");

		// 6*4 + 4 + 4 = 8*4
		// XYZ - Position (Float32)
		// XYZ - Scale (Float32)
		// RGBA - colors (uint8)
		// IJKL - quaternion/rot (uint8)
		const rowLength = 3 * 4 + 3 * 4 + 4 + 4;
		const buffer = new ArrayBuffer(rowLength * vertexCount);

		console.time("build buffer");
		for (let j = 0; j < vertexCount; j++) {
			row = sizeIndex[j];

			const position = new Float32Array(buffer, j * rowLength, 3);
			const scales = new Float32Array(buffer, j * rowLength + 4 * 3, 3);
			const rgba = new Uint8ClampedArray(
				buffer,
				j * rowLength + 4 * 3 + 4 * 3,
				4,
			);
			const rot = new Uint8ClampedArray(
				buffer,
				j * rowLength + 4 * 3 + 4 * 3 + 4,
				4,
			);

			if (types["scale_0"]) {
				const qlen = Math.sqrt(
					attrs.rot_0 ** 2 +
						attrs.rot_1 ** 2 +
						attrs.rot_2 ** 2 +
						attrs.rot_3 ** 2,
				);

				rot[0] = (attrs.rot_0 / qlen) * 128 + 128;
				rot[1] = (attrs.rot_1 / qlen) * 128 + 128;
				rot[2] = (attrs.rot_2 / qlen) * 128 + 128;
				rot[3] = (attrs.rot_3 / qlen) * 128 + 128;

				scales[0] = Math.exp(attrs.scale_0);
				scales[1] = Math.exp(attrs.scale_1);
				scales[2] = Math.exp(attrs.scale_2);
			} else {
				scales[0] = 0.01;
				scales[1] = 0.01;
				scales[2] = 0.01;

				rot[0] = 255;
				rot[1] = 0;
				rot[2] = 0;
				rot[3] = 0;
			}

			position[0] = attrs.x;
			position[1] = attrs.y;
			position[2] = attrs.z;

			if (types["f_dc_0"]) {
				const SH_C0 = 0.28209479177387814;
				rgba[0] = (0.5 + SH_C0 * attrs.f_dc_0) * 255;
				rgba[1] = (0.5 + SH_C0 * attrs.f_dc_1) * 255;
				rgba[2] = (0.5 + SH_C0 * attrs.f_dc_2) * 255;
			} else {
				rgba[0] = attrs.red;
				rgba[1] = attrs.green;
				rgba[2] = attrs.blue;
			}
			if (types["opacity"]) {
				rgba[3] = (1 / (1 + Math.exp(-attrs.opacity))) * 255;
			} else {
				rgba[3] = 255;
			}
		}
		console.timeEnd("build buffer");
		return buffer;
	}

	const throttledSort = () => {
		if (!sortRunning) {
			sortRunning = true;
			let lastView = viewProj;
			runSort(lastView);
			setTimeout(() => {
				sortRunning = false;
				if (lastView !== viewProj) {
					throttledSort();
				}
			}, 0);
		}
	};

	let sortRunning;
	self.onmessage = (e) => {
		if (e.data.ply) {
			vertexCount = 0;
			runSort(viewProj);
			buffer = processPlyBuffer(e.data.ply);
			vertexCount = Math.floor(buffer.byteLength / rowLength);
			postMessage({ buffer: buffer });
		} else if (e.data.buffer) {
			buffer = e.data.buffer;
			vertexCount = e.data.vertexCount;
		} else if (e.data.vertexCount) {
			vertexCount = e.data.vertexCount;
		} else if (e.data.view) {
			viewProj = e.data.view;
			throttledSort();
		}
	};
}

const vertexShaderSource = `
  precision mediump float;
  attribute vec2 position;

  attribute vec4 color;
  attribute vec3 center;
  attribute vec3 covA;
  attribute vec3 covB;

  uniform mat4 projection, view;
  uniform vec2 focal;
  uniform vec2 viewport;

  varying vec4 vColor;
  varying vec2 vPosition;

  mat3 transpose(mat3 m) {
    return mat3(
        m[0][0], m[1][0], m[2][0],
        m[0][1], m[1][1], m[2][1],
        m[0][2], m[1][2], m[2][2]
    );
  }

  void main () {
    vec4 camspace = view * vec4(center, 1);
    vec4 pos2d = projection * camspace;

    float bounds = 1.2 * pos2d.w;
    if (pos2d.z < -pos2d.w || pos2d.x < -bounds || pos2d.x > bounds
		 || pos2d.y < -bounds || pos2d.y > bounds) {
        gl_Position = vec4(0.0, 0.0, 2.0, 1.0);
        return;
    }

    mat3 Vrk = mat3(
        covA.x, covA.y, covA.z, 
        covA.y, covB.x, covB.y,
        covA.z, covB.y, covB.z
    );
	
    mat3 J = mat3(
        focal.x / camspace.z, 0., -(focal.x * camspace.x) / (camspace.z * camspace.z), 
        0., -focal.y / camspace.z, (focal.y * camspace.y) / (camspace.z * camspace.z), 
        0., 0., 0.
    );

    mat3 W = transpose(mat3(view));
    mat3 T = W * J;
    mat3 cov = transpose(T) * Vrk * T;
    
    vec2 vCenter = vec2(pos2d) / pos2d.w;

    float diagonal1 = cov[0][0] + 0.3;
    float offDiagonal = cov[0][1];
    float diagonal2 = cov[1][1] + 0.3;

	float mid = 0.5 * (diagonal1 + diagonal2);
	float radius = length(vec2((diagonal1 - diagonal2) / 2.0, offDiagonal));
	float lambda1 = mid + radius;
	float lambda2 = max(mid - radius, 0.1);
	vec2 diagonalVector = normalize(vec2(offDiagonal, lambda1 - diagonal1));
	vec2 v1 = min(sqrt(2.0 * lambda1), 1024.0) * diagonalVector;
	vec2 v2 = min(sqrt(2.0 * lambda2), 1024.0) * vec2(diagonalVector.y, -diagonalVector.x);


    vColor = color;
    vPosition = position;

    gl_Position = vec4(
        vCenter 
            + position.x * v1 / viewport * 2.0 
            + position.y * v2 / viewport * 2.0, 0.0, 1.0);

  }
`;

const fragmentShaderSource = `
precision mediump float;

  varying vec4 vColor;
  varying vec2 vPosition;

  void main () {    
	  float A = -dot(vPosition, vPosition);
    if (A < -4.0) discard;
    float B = exp(A) * vColor.a;
    gl_FragColor = vec4(B * vColor.rgb, B);
  }
`;

let defaultViewMatrix = [
	0.47, 0.04, 0.88, 0, -0.11, 0.99, 0.02, 0, -0.88, -0.11, 0.47, 0, 0.07,
	0.03, 6.55, 1,
];
let viewMatrix = defaultViewMatrix;

async function main() {
	let carousel = true;
	const params = new URLSearchParams(location.search);
	try {
		viewMatrix = JSON.parse(decodeURIComponent(location.hash.slice(1)));
		carousel = false;
	} catch (err) {}
	const url = new URL(
		// "nike.splat",
		// location.href,
		params.get("url") || "train.splat",
		"https://huggingface.co/cakewalk/splat-data/resolve/main/",
	);
	const req = await fetch(url, {
		mode: "cors", // no-cors, *cors, same-origin
		credentials: "omit", // include, *same-origin, omit
	});
	console.log(req);
	if (req.status != 200)
		throw new Error(req.status + " Unable to load " + req.url);

	const rowLength = 3 * 4 + 3 * 4 + 4 + 4;
	const reader = req.body.getReader();
	let splatData = new Uint8Array(req.headers.get("content-length"));

	const downsample = splatData.length / rowLength > 500000 ? 1 : 1 / devicePixelRatio;
	// const downsample = 1 / devicePixelRatio;
	// const downsample = 1;
	console.log(splatData.length / rowLength, downsample);

	const worker = new Worker(
		URL.createObjectURL(
			new Blob(["(", createWorker.toString(), ")(self)"], {
				type: "application/javascript",
			}),
		),
	);

	const canvas = document.getElementById("canvas");
	canvas.width = innerWidth / downsample;
	canvas.height = innerHeight / downsample;

	const fps = document.getElementById("fps");

	let projectionMatrix = getProjectionMatrix(
		camera.fx / downsample,
		camera.fy / downsample,
		canvas.width,
		canvas.height,
	);

	const gl = canvas.getContext("webgl");
	const ext = gl.getExtension("ANGLE_instanced_arrays");

	const vertexShader = gl.createShader(gl.VERTEX_SHADER);
	gl.shaderSource(vertexShader, vertexShaderSource);
	gl.compileShader(vertexShader);
	if (!gl.getShaderParameter(vertexShader, gl.COMPILE_STATUS))
		console.error(gl.getShaderInfoLog(vertexShader));

	const fragmentShader = gl.createShader(gl.FRAGMENT_SHADER);
	gl.shaderSource(fragmentShader, fragmentShaderSource);
	gl.compileShader(fragmentShader);
	if (!gl.getShaderParameter(fragmentShader, gl.COMPILE_STATUS))
		console.error(gl.getShaderInfoLog(fragmentShader));

	const program = gl.createProgram();
	gl.attachShader(program, vertexShader);
	gl.attachShader(program, fragmentShader);
	gl.linkProgram(program);
	gl.useProgram(program);

	if (!gl.getProgramParameter(program, gl.LINK_STATUS))
		console.error(gl.getProgramInfoLog(program));

	gl.disable(gl.DEPTH_TEST); // Disable depth testing

	// Enable blending
	gl.enable(gl.BLEND);

	// Set blending function
	gl.blendFuncSeparate(
		gl.ONE_MINUS_DST_ALPHA,
		gl.ONE,
		gl.ONE_MINUS_DST_ALPHA,
		gl.ONE,
	);

	// Set blending equation
	gl.blendEquationSeparate(gl.FUNC_ADD, gl.FUNC_ADD);

	// projection
	const u_projection = gl.getUniformLocation(program, "projection");
	gl.uniformMatrix4fv(u_projection, false, projectionMatrix);

	// viewport
	const u_viewport = gl.getUniformLocation(program, "viewport");
	gl.uniform2fv(u_viewport, new Float32Array([canvas.width, canvas.height]));

	// focal
	const u_focal = gl.getUniformLocation(program, "focal");
	gl.uniform2fv(
		u_focal,
		new Float32Array([camera.fx / downsample, camera.fy / downsample]),
	);

	// view
	const u_view = gl.getUniformLocation(program, "view");
	gl.uniformMatrix4fv(u_view, false, viewMatrix);

	// positions
	const triangleVertices = new Float32Array([-2, -2, 2, -2, 2, 2, -2, 2]);
	const vertexBuffer = gl.createBuffer();
	gl.bindBuffer(gl.ARRAY_BUFFER, vertexBuffer);
	gl.bufferData(gl.ARRAY_BUFFER, triangleVertices, gl.STATIC_DRAW);
	const a_position = gl.getAttribLocation(program, "position");
	gl.enableVertexAttribArray(a_position);
	gl.bindBuffer(gl.ARRAY_BUFFER, vertexBuffer);
	gl.vertexAttribPointer(a_position, 2, gl.FLOAT, false, 0, 0);

	// center
	const centerBuffer = gl.createBuffer();
	// gl.bindBuffer(gl.ARRAY_BUFFER, centerBuffer);
	// gl.bufferData(gl.ARRAY_BUFFER, center, gl.STATIC_DRAW);
	const a_center = gl.getAttribLocation(program, "center");
	gl.enableVertexAttribArray(a_center);
	gl.bindBuffer(gl.ARRAY_BUFFER, centerBuffer);
	gl.vertexAttribPointer(a_center, 3, gl.FLOAT, false, 0, 0);
	ext.vertexAttribDivisorANGLE(a_center, 1); // Use the extension here

	// color
	const colorBuffer = gl.createBuffer();
	// gl.bindBuffer(gl.ARRAY_BUFFER, colorBuffer);
	// gl.bufferData(gl.ARRAY_BUFFER, color, gl.STATIC_DRAW);
	const a_color = gl.getAttribLocation(program, "color");
	gl.enableVertexAttribArray(a_color);
	gl.bindBuffer(gl.ARRAY_BUFFER, colorBuffer);
	gl.vertexAttribPointer(a_color, 4, gl.FLOAT, false, 0, 0);
	ext.vertexAttribDivisorANGLE(a_color, 1); // Use the extension here

	// cov
	const covABuffer = gl.createBuffer();
	const a_covA = gl.getAttribLocation(program, "covA");
	gl.enableVertexAttribArray(a_covA);
	gl.bindBuffer(gl.ARRAY_BUFFER, covABuffer);
	gl.vertexAttribPointer(a_covA, 3, gl.FLOAT, false, 0, 0);
	ext.vertexAttribDivisorANGLE(a_covA, 1); // Use the extension here

	const covBBuffer = gl.createBuffer();
	const a_covB = gl.getAttribLocation(program, "covB");
	gl.enableVertexAttribArray(a_covB);
	gl.bindBuffer(gl.ARRAY_BUFFER, covBBuffer);
	gl.vertexAttribPointer(a_covB, 3, gl.FLOAT, false, 0, 0);
	ext.vertexAttribDivisorANGLE(a_covB, 1); // Use the extension here

	let lastProj = []
	let lastData

	worker.onmessage = (e) => {
		if (e.data.buffer) {
			splatData = new Uint8Array(e.data.buffer);
			const blob = new Blob([splatData.buffer], {
				type: "application/octet-stream",
			});
			const link = document.createElement("a");
			link.download = "model.splat";
			link.href = URL.createObjectURL(blob);
			document.body.appendChild(link);
			link.click();
		} else {
			let { covA, covB, center, color, viewProj } = e.data;
			lastData = e.data;

			lastProj = viewProj
			vertexCount = center.length / 3;

			gl.bindBuffer(gl.ARRAY_BUFFER, centerBuffer);
			gl.bufferData(gl.ARRAY_BUFFER, center, gl.STATIC_DRAW);

			gl.bindBuffer(gl.ARRAY_BUFFER, colorBuffer);
			gl.bufferData(gl.ARRAY_BUFFER, color, gl.STATIC_DRAW);

			
			gl.bindBuffer(gl.ARRAY_BUFFER, covABuffer);
			gl.bufferData(gl.ARRAY_BUFFER, covA, gl.STATIC_DRAW);

			gl.bindBuffer(gl.ARRAY_BUFFER, covBBuffer);
			gl.bufferData(gl.ARRAY_BUFFER, covB, gl.STATIC_DRAW);

		}
	};

	let activeKeys = [];

	window.addEventListener("keydown", (e) => {
		if (document.activeElement != document.body) return;
		carousel = false;
		if (!activeKeys.includes(e.key)) activeKeys.push(e.key);
		if (/\d/.test(e.key)) {
			viewMatrix = getViewMatrix(cameras[parseInt(e.key)]);
		}
		if (e.key == "v") {
			location.hash =
				"#" +
				JSON.stringify(
					viewMatrix.map((k) => Math.round(k * 100) / 100),
				);
		} else if (e.key === "p") {
			carousel = true;
		}
	});
	window.addEventListener("keyup", (e) => {
		activeKeys = activeKeys.filter((k) => k !== e.key);
	});
	window.addEventListener("blur", () => {
		activeKeys = [];
	});

	window.addEventListener(
		"wheel",
		(e) => {
			carousel = false;
			e.preventDefault();
			const lineHeight = 10;
			const scale =
				e.deltaMode == 1
					? lineHeight
					: e.deltaMode == 2
					? innerHeight
					: 1;
			let inv = invert4(viewMatrix);
			if (e.shiftKey) {
				inv = translate4(
					inv,
					(e.deltaX * scale) / innerWidth,
					(e.deltaY * scale) / innerHeight,
					0,
				);
			} else if (e.ctrlKey || e.metaKey) {
				// inv = rotate4(inv,  (e.deltaX * scale) / innerWidth,  0, 0, 1);
				// inv = translate4(inv,  0, (e.deltaY * scale) / innerHeight, 0);
				let preY = inv[13];
				inv = translate4(
					inv,
					0,
					0,
					(-10 * (e.deltaY * scale)) / innerHeight,
				);
				inv[13] = preY;
			} else {
				let d = 4;
				inv = translate4(inv, 0, 0, d);
				inv = rotate4(inv, -(e.deltaX * scale) / innerWidth, 0, 1, 0);
				inv = rotate4(inv, (e.deltaY * scale) / innerHeight, 1, 0, 0);
				inv = translate4(inv, 0, 0, -d);
			}

			viewMatrix = invert4(inv);
		},
		{ passive: false },
	);

	let startX, startY, down;
	canvas.addEventListener("mousedown", (e) => {
		carousel = false;
		e.preventDefault();
		startX = e.clientX;
		startY = e.clientY;
		down = e.ctrlKey || e.metaKey ? 2 : 1;
	});
	canvas.addEventListener("contextmenu", (e) => {
		carousel = false;
		e.preventDefault();
		startX = e.clientX;
		startY = e.clientY;
		down = 2;
	});

	canvas.addEventListener("mousemove", (e) => {
		e.preventDefault();
		if (down == 1) {
			let inv = invert4(viewMatrix);
			let dx = (5 * (e.clientX - startX)) / innerWidth;
			let dy = (5 * (e.clientY - startY)) / innerHeight;
			let d = 4;

			inv = translate4(inv, 0, 0, d);
			inv = rotate4(inv, dx, 0, 1, 0);
			inv = rotate4(inv, -dy, 1, 0, 0);
			inv = translate4(inv, 0, 0, -d);
			// let postAngle = Math.atan2(inv[0], inv[10])
			// inv = rotate4(inv, postAngle - preAngle, 0, 0, 1)
			// console.log(postAngle)
			viewMatrix = invert4(inv);

			startX = e.clientX;
			startY = e.clientY;
		} else if (down == 2) {
			let inv = invert4(viewMatrix);
			// inv = rotateY(inv, );
			let preY = inv[13];
			inv = translate4(
				inv,
				(-10 * (e.clientX - startX)) / innerWidth,
				0,
				(10 * (e.clientY - startY)) / innerHeight,
			);
			inv[13] = preY;
			viewMatrix = invert4(inv);

			startX = e.clientX;
			startY = e.clientY;
		}
	});
	canvas.addEventListener("mouseup", (e) => {
		e.preventDefault();
		down = false;
		startX = 0;
		startY = 0;
	});

	let altX = 0,
		altY = 0;
	canvas.addEventListener(
		"touchstart",
		(e) => {
			e.preventDefault();
			if (e.touches.length === 1) {
				carousel = false;
				startX = e.touches[0].clientX;
				startY = e.touches[0].clientY;
				down = 1;
			} else if (e.touches.length === 2) {
				// console.log('beep')
				carousel = false;
				startX = e.touches[0].clientX;
				altX = e.touches[1].clientX;
				startY = e.touches[0].clientY;
				altY = e.touches[1].clientY;
				down = 1;
			}
		},
		{ passive: false },
	);
	canvas.addEventListener(
		"touchmove",
		(e) => {
			e.preventDefault();
			if (e.touches.length === 1 && down) {
				let inv = invert4(viewMatrix);
				let dx = (4 * (e.touches[0].clientX - startX)) / innerWidth;
				let dy = (4 * (e.touches[0].clientY - startY)) / innerHeight;

				let d = 4;
				inv = translate4(inv, 0, 0, d);
				// inv = translate4(inv,  -x, -y, -z);
				// inv = translate4(inv,  x, y, z);
				inv = rotate4(inv, dx, 0, 1, 0);
				inv = rotate4(inv, -dy, 1, 0, 0);
				inv = translate4(inv, 0, 0, -d);

				viewMatrix = invert4(inv);

				startX = e.touches[0].clientX;
				startY = e.touches[0].clientY;
			} else if (e.touches.length === 2) {
				// alert('beep')
				const dtheta =
					Math.atan2(startY - altY, startX - altX) -
					Math.atan2(
						e.touches[0].clientY - e.touches[1].clientY,
						e.touches[0].clientX - e.touches[1].clientX,
					);
				const dscale =
					Math.hypot(startX - altX, startY - altY) /
					Math.hypot(
						e.touches[0].clientX - e.touches[1].clientX,
						e.touches[0].clientY - e.touches[1].clientY,
					);
				const dx =
					(e.touches[0].clientX +
						e.touches[1].clientX -
						(startX + altX)) /
					2;
				const dy =
					(e.touches[0].clientY +
						e.touches[1].clientY -
						(startY + altY)) /
					2;
				let inv = invert4(viewMatrix);
				// inv = translate4(inv,  0, 0, d);
				inv = rotate4(inv, dtheta, 0, 0, 1);

				inv = translate4(inv, -dx / innerWidth, -dy / innerHeight, 0);

				let preY = inv[13];
				inv = translate4(inv, 0, 0, 3 * (1 - dscale));
				inv[13] = preY;

				viewMatrix = invert4(inv);

				startX = e.touches[0].clientX;
				altX = e.touches[1].clientX;
				startY = e.touches[0].clientY;
				altY = e.touches[1].clientY;
			}
		},
		{ passive: false },
	);
	canvas.addEventListener(
		"touchend",
		(e) => {
			e.preventDefault();
			down = false;
			startX = 0;
			startY = 0;
		},
		{ passive: false },
	);

	let jumpDelta = 0;
	let vertexCount = 0;

	let lastFrame = 0;
	let avgFps = 0;
	let start = 0;

	const frame = (now) => {
		let inv = invert4(viewMatrix);

		if (activeKeys.includes("ArrowUp")) {
			if(activeKeys.includes("Shift")){
				inv = translate4(inv, 0, -0.03, 0);
			}else{
				let preY = inv[13];
				inv = translate4(inv, 0, 0, 0.1);
				inv[13] = preY;
			}
		}
		if (activeKeys.includes("ArrowDown")) {
			if(activeKeys.includes("Shift")){
				inv = translate4(inv, 0, 0.03, 0);
			}else{
				let preY = inv[13];
				inv = translate4(inv, 0, 0, -0.1);
				inv[13] = preY;
			}
		}
		if (activeKeys.includes("ArrowLeft"))
			inv = translate4(inv, -0.03, 0, 0);
		//
		if (activeKeys.includes("ArrowRight"))
			inv = translate4(inv, 0.03, 0, 0);
		// inv = rotate4(inv, 0.01, 0, 1, 0);
		if (activeKeys.includes("a")) inv = rotate4(inv, -0.01, 0, 1, 0);
		if (activeKeys.includes("d")) inv = rotate4(inv, 0.01, 0, 1, 0);
		if (activeKeys.includes("q")) inv = rotate4(inv, 0.01, 0, 0, 1);
		if (activeKeys.includes("e")) inv = rotate4(inv, -0.01, 0, 0, 1);
		if (activeKeys.includes("w")) inv = rotate4(inv, 0.005, 1, 0, 0);
		if (activeKeys.includes("s")) inv = rotate4(inv, -0.005, 1, 0, 0);

		if (["j", "k", "l", "i"].some((k) => activeKeys.includes(k))) {
			let d = 4;
			inv = translate4(inv, 0, 0, d);
			inv = rotate4(
				inv,
				activeKeys.includes("j")
					? -0.05
					: activeKeys.includes("l")
					? 0.05
					: 0,
				0,
				1,
				0,
			);
			inv = rotate4(
				inv,
				activeKeys.includes("i")
					? 0.05
					: activeKeys.includes("k")
					? -0.05
					: 0,
				1,
				0,
				0,
			);
			inv = translate4(inv, 0, 0, -d);
		}

		// inv[13] = preY;
		viewMatrix = invert4(inv);

		if (carousel) {
			let inv = invert4(defaultViewMatrix);

			const t = Math.sin((Date.now() - start) / 5000);
			inv = translate4(inv, 2.5 * t, 0, 6 * (1 - Math.cos(t)));
			inv = rotate4(inv, -0.6 * t, 0, 1, 0);

			viewMatrix = invert4(inv);
		}

		if (activeKeys.includes(" ")) {
			jumpDelta = Math.min(1, jumpDelta + 0.05);
		} else {
			jumpDelta = Math.max(0, jumpDelta - 0.05);
		}

		let inv2 = invert4(viewMatrix);
		inv2[13] -= jumpDelta;
		inv2 = rotate4(inv2, -0.1 * jumpDelta, 1, 0, 0);
		let actualViewMatrix = invert4(inv2);

		const viewProj = multiply4(projectionMatrix, actualViewMatrix);
		worker.postMessage({ view: viewProj });

		const currentFps = 1000 / (now - lastFrame) || 0;
		avgFps = avgFps * 0.9 + currentFps * 0.1;

		if (vertexCount > 0) {
			document.getElementById("spinner").style.display = "none";
			// console.time('render')
			gl.uniformMatrix4fv(u_view, false, actualViewMatrix);
			ext.drawArraysInstancedANGLE(gl.TRIANGLE_FAN, 0, 4, vertexCount);
			// console.timeEnd('render')
		} else {
			gl.clear(gl.COLOR_BUFFER_BIT);
			document.getElementById("spinner").style.display = "";
			start = Date.now() + 2000;
		}
		const progress = (100 * vertexCount) / (splatData.length / rowLength);
		if (progress < 100) {
			document.getElementById("progress").style.width = progress + "%";
		} else {
			document.getElementById("progress").style.display = "none";	
		}
		fps.innerText = Math.round(avgFps) + " fps";
		lastFrame = now;
		requestAnimationFrame(frame);
	};

	frame();

	const selectFile = (file) => {
		const fr = new FileReader();
		if (/\.json$/i.test(file.name)) {
			fr.onload = () => {
				cameras = JSON.parse(fr.result);
				viewMatrix = getViewMatrix(cameras[0]);
				projectionMatrix = getProjectionMatrix(
					camera.fx / downsample,
					camera.fy / downsample,
					canvas.width,
					canvas.height,
				);
				gl.uniformMatrix4fv(u_projection, false, projectionMatrix);

				console.log("Loaded Cameras");
			};
			fr.readAsText(file);
		} else {
			stopLoading = true;
			fr.onload = () => {
				splatData = new Uint8Array(fr.result);
				console.log("Loaded", Math.floor(splatData.length / rowLength));

				if (
					splatData[0] == 112 &&
					splatData[1] == 108 &&
					splatData[2] == 121 &&
					splatData[3] == 10
				) {
					// ply file magic header means it should be handled differently
					worker.postMessage({ ply: splatData.buffer });
				} else {
					worker.postMessage({
						buffer: splatData.buffer,
						vertexCount: Math.floor(splatData.length / rowLength),
					});
				}
			};
			fr.readAsArrayBuffer(file);
		}
	};

	window.addEventListener("hashchange", (e) => {
		try {
			viewMatrix = JSON.parse(decodeURIComponent(location.hash.slice(1)));
			carousel = false;
		} catch (err) {}
	});

	const preventDefault = (e) => {
		e.preventDefault();
		e.stopPropagation();
	};
	document.addEventListener("dragenter", preventDefault);
	document.addEventListener("dragover", preventDefault);
	document.addEventListener("dragleave", preventDefault);
	document.addEventListener("drop", (e) => {
		e.preventDefault();
		e.stopPropagation();
		selectFile(e.dataTransfer.files[0]);
	});

	let bytesRead = 0;
	let lastVertexCount = -1;
	let stopLoading = false;

	while (true) {
		const { done, value } = await reader.read();
		if (done || stopLoading) break;

		splatData.set(value, bytesRead);
		bytesRead += value.length;

		if (vertexCount > lastVertexCount) {
			worker.postMessage({
				buffer: splatData.buffer,
				vertexCount: Math.floor(bytesRead / rowLength),
			});
			lastVertexCount = vertexCount;
		}
	}
	if (!stopLoading)
		worker.postMessage({
			buffer: splatData.buffer,
			vertexCount: Math.floor(bytesRead / rowLength),
		});
}

main().catch((err) => {
	document.getElementById("spinner").style.display = "none";
	document.getElementById("message").innerText = err.toString();
});