import Shader 			from "../../framework/WebGpu.js"

import Matrix4 			from "../../framework/Matrix4.js"

import Vector3 			from "../../framework/Vector3.js"

import Camera 			from "../../framework/Camera.js";

import EventManager 	from "../../framework/eventManager.js";

import ShaderInpector	from "../../framework/ShaderInpector.js";


export class ParticleSimulation {

	canvas;

	device;

	camera;

	useLocalSort		= true;

	eventManager		= new EventManager();

	frameCount 			= 0;

	setCanvas( canvas ) {

		this.canvas = canvas;

		this.eventManager.setCanvas( canvas );

	}

	async setup( offscreenCanvas, width, height ) {

		offscreenCanvas.width	= width;

		offscreenCanvas.height	= height;

		this.canvas				= offscreenCanvas;

		//console.log( this.canvas.width, this.canvas.height );

		const context 			= offscreenCanvas.getContext("webgpu");




		this.camera 			= new Camera( [0, 0, 5], [0, -.3, 0], [0, 1, 0] );

		this.eventManager.setup( offscreenCanvas, this.camera );

		//this.eventManager.registerEventListenersNode();

		const adapter 			= await self.navigator.gpu.requestAdapter();

		if ( !adapter ) {

			throw new Error("Failed to get GPU adapter");

		}

		this.device 			= await adapter.requestDevice();

		this.particleCount 		= 8192 * 2;

		const presentationFormat = navigator.gpu.getPreferredCanvasFormat();

		context.configure({
			device: this.device,
			format: presentationFormat,
			alphaMode: "opaque"
		});

		let boundsMinimum 		= new Vector3( -2, -2, -2 );

		let boundsMaximum		= new Vector3( 2, 2, 2 );

		let cellsPerDimension 	= 12;

		let gravity				= -2.3;

		this.workgroupSize 		= 64;

		this.numWorkgroups 		= Math.ceil( this.particleCount / this.workgroupSize );


		var jArray 				= new Array();

		var kArray 				= new Array();

		if( this.useLocalSort ) {

			var startK = this.workgroupSize * 2; // 2 * 256; 256 = this.workgroupSize

		} else {

			var startK = 2;
			
		}

		for (let k = startK; k <= this.particleCount; k <<= 1 ) {

			for (let j = k >> 1; j > 0; j >>= 1 ) {

				jArray.push( j );

				kArray.push( k );	

			}

		}


		this.gravityShader 		= new Shader( this.device );

		await this.gravityShader.setup( "../../shaders/gravity.wgsl" );

		this.gravityShader.setVariable("gravity", gravity );

		this.gravityShader.setVariable("updateDistancesAndIndices", 1);


		this.resetParticlePositions();


		this.copyIndicesShader 	= new Shader( this.device );

		await this.copyIndicesShader.setup( "../../shaders/copyBuffer.wgsl" );


		this.renderShader 		= new Shader( this.device );


		this.renderShader.setCanvas( this.canvas );

		await this.renderShader.setup("../../shaders/points.wgsl");



		const quadOffsets = new Float32Array( [
			// Triangle 1
			-1, -1,  // bottom-left
			1,  -1,  // bottom-right
			1,   1,  // top-right

			// Triangle 2
			-1, -1,  // bottom-left
			1,   1,  // top-right
			-1,  1   // top-left
		]);

		this.renderShader.setAttribute( "quadOffset", quadOffsets );


		this.findGridHashShader = new Shader( this.device );

		await this.findGridHashShader.setup( "../../shaders/findGridHash.wgsl"  );

		//console.log("this.gravityShader.getBuffer(\"positions\")", this.gravityShader.getBuffer("positions"));

		this.findGridHashShader.setBuffer( "positions", this.gravityShader.getBuffer("positions") );

		this.findGridHashShader.setVariable( "cellCount", cellsPerDimension );

		this.findGridHashShader.setVariable( "gridMin", boundsMinimum );

		this.findGridHashShader.setVariable( "gridMax", boundsMaximum );


		if( this.useLocalSort ) {

			this.localSortShader = new Shader( this.device );

			await this.localSortShader.setup( "../../shaders/localSort.wgsl"  );

			this.localSortShader.setBuffer( "gridHashes", this.findGridHashShader.getBuffer("gridHashes" ) );

			this.localSortShader.setBuffer( "indices", this.findGridHashShader.getBuffer("indices" ) );

			this.localSortShader.setVariable( "totalCount", this.particleCount );

		}

		this.bitonicSortGridHashShader = new Shader( this.device );

		await this.bitonicSortGridHashShader.setup( "../../shaders/bitonicSortUIntMultiPass.wgsl" );

		if( this.useLocalSort ) {

			this.bitonicSortGridHashShader.setBuffer( "gridHashes", this.localSortShader.getBuffer("gridHashes")  );

			this.bitonicSortGridHashShader.setBuffer( "indices", this.localSortShader.getBuffer("indices") );

		} else {

			this.bitonicSortGridHashShader.setBuffer( "gridHashes", this.findGridHashShader.getBuffer("gridHashes" ) );

			this.bitonicSortGridHashShader.setBuffer( "indices", this.findGridHashShader.getBuffer("indices") );

		}

		this.bitonicSortGridHashShader.setVariable( "totalCount", this.particleCount );

		this.bitonicSortGridHashShader.setVariable( "jArray", jArray );

		this.bitonicSortGridHashShader.setVariable( "kArray", kArray );


		this.findGridHashRangeShader = new Shader( this.device );

		await this.findGridHashRangeShader.setup( "../../shaders/findGridHashRanges.wgsl" );


		this.collisionDetectionShader = new Shader( this.device );

		await this.collisionDetectionShader.setup( "../../shaders/collisionDetection.wgsl" );

		this.collisionDetectionShader.setBuffer( "positions", this.gravityShader.getBuffer("positions" ) );

		this.collisionDetectionShader.setBuffer( "velocities", this.gravityShader.getBuffer("velocities") );

		this.collisionDetectionShader.setBuffer( "positions", this.gravityShader.getBuffer("positions" ) );

		this.collisionDetectionShader.setBuffer( "velocities", this.gravityShader.getBuffer("velocities") );

		this.collisionDetectionShader.setBuffer( "gridHashes", this.findGridHashShader.getBuffer("gridHashes") );


		if( this.useLocalSort ) {

			this.collisionDetectionShader.setBuffer( "hashSortedIndices", this.localSortShader.getBuffer("indices") );

		} else {

			this.collisionDetectionShader.setBuffer( "hashSortedIndices", this.bitonicSortGridHashShader.getBuffer("indices") );

		}


		this.collisionDetectionShader.setVariable( "cellCount", cellsPerDimension );

		this.collisionDetectionShader.setVariable( "gridMin", boundsMinimum );

		this.collisionDetectionShader.setVariable( "gridMax", boundsMaximum );

		this.collisionDetectionShader.setBuffer( "startIndices", this.findGridHashRangeShader.getBuffer("startIndices") );

		this.collisionDetectionShader.setBuffer( "endIndices", this.findGridHashRangeShader.getBuffer("endIndices") );

		this.collisionDetectionShader.setVariable( "collisionRadius", 0.06 );


		this.bitonicSortShader = new Shader(this.device, "../../shaders/bitonicSort.wgsl");

		await this.bitonicSortShader.addStage("main", GPUShaderStage.COMPUTE );


		this.bitonicSortShader.setBuffer( "compare", this.gravityShader.getBuffer("distances") );

		this.bitonicSortShader.setBuffer( "indices", this.gravityShader.getBuffer("indices") );

		this.bitonicSortShader.setVariable( "totalCount", this.particleCount );


		this.findGridHashRangeShader.setBuffer( "gridHashes", this.bitonicSortGridHashShader.getBuffer("gridHashes") );

		this.findGridHashRangeShader.setBuffer( "indices", this.bitonicSortGridHashShader.getBuffer("indices") );

		this.findGridHashRangeShader.setVariable( "totalCount", this.particleCount );


		this.copyIndicesShader.setBuffer( "indices", this.bitonicSortShader.getBuffer("indices") );


		this.renderShader.setBuffer( "positions", this.gravityShader.getBuffer("positions") );

		this.renderShader.setBuffer( "sortedIndices", this.copyIndicesShader.getBuffer("sortedIndices") );

		this.renderShader.setCanvas( this.canvas );

		//var simulation = new particleSimulation();

		this.render();


		return;

		//document.getElementById("resetParticlesButton").addEventListener("click", resetParticlePositions);


	}

	resetParticlePositions() {

		const positions			= new Float32Array( this.particleCount * 3 );

		const velocities 		= new Float32Array( this.particleCount * 3 );


		for (var i = 0; i < this.particleCount; i++) {

			positions[ i * 3 + 0 ] 	= Math.random() * 2.0 - 1.0;

			positions[ i * 3 + 1 ] 	= Math.random() * 2.0 - 1.0;

			positions[ i * 3 + 2 ] 	= Math.random() * 2.0 - 1.0;

			velocities[ i * 3 + 0 ] = 0;

			velocities[ i * 3 + 1 ] = 0;

			velocities[ i * 3 + 2 ] = 0;

		}

		this.gravityShader.setVariable("positions", positions);

		this.gravityShader.setVariable("velocities", velocities);

	}

	

	updateTimeDelta() {

		const now				= performance.now();

		this.deltaTimeValue		= ( now - this.lastFrameTime ) / 1000;

		this.lastFrameTime		= now;

	}

	async sortGridHash() {

		const commandEncoder = this.device.createCommandEncoder();

		// Local sort 256 items
		if( this.useLocalSort ) {

			const passEncoder = commandEncoder.beginComputePass();

			passEncoder.setPipeline( this.localSortShader.pipeline );

			for ( const [ index, bindGroup ] of this.localSortShader.bindGroups.entries() ) {

				passEncoder.setBindGroup( index, bindGroup );

			}

			passEncoder.dispatchWorkgroups( this.numWorkgroups );

			passEncoder.end();

		}

		// Bitonic global merge
		{
			const threadPassIndices 	= new Uint32Array( this.particleCount );

			for (var i = 0; i < this.particleCount; i++) {

				threadPassIndices[0]		= 0;

			}

			this.bitonicSortGridHashShader.setVariable("threadPassIndices", threadPassIndices);

			const passEncoder = commandEncoder.beginComputePass();

			passEncoder.setPipeline( this.bitonicSortGridHashShader.pipeline );

			for ( const [ index, bindGroup ] of this.bitonicSortGridHashShader.bindGroups.entries() ) {

				passEncoder.setBindGroup( index, bindGroup );

			}

			if( this.useLocalSort ) {

				var startK = this.workgroupSize * 2;

			} else {

				var startK = 2;

			}

			for (let k = startK; k <= this.particleCount; k <<= 1 ) {

				for (let j = k >> 1; j > 0; j >>= 1 ) {

					passEncoder.dispatchWorkgroups( this.numWorkgroups );

				}

			}

			passEncoder.end();
			
		}

		const commandBuffer = commandEncoder.finish();

		await this.device.queue.submit( [ commandBuffer ] );

		await this.device.queue.onSubmittedWorkDone();

	}

	async sortDistance() {

		await this.gravityShader.execute( this.numWorkgroups );

		this.gravityShader.setVariable("updateDistancesAndIndices", 0);

		for (let k = 2; k <= this.particleCount; k <<= 1) {

			for (let j = k >> 1; j > 0; j >>= 1) {

				const commandBuffers = new Array();

				await this.bitonicSortShader.setVariable( "j", j );  // Must await uniform update
				await this.bitonicSortShader.setVariable( "k", k );

				const commandEncoder = this.device.createCommandEncoder();

				const passEncoder = commandEncoder.beginComputePass();

				passEncoder.setPipeline( this.bitonicSortShader.pipeline );

				for ( const [ index, bindGroup ] of this.bitonicSortShader.bindGroups.entries() ) {
					
					passEncoder.setBindGroup( index, bindGroup );

				}

				passEncoder.dispatchWorkgroups( this.numWorkgroups );

				passEncoder.end();

				commandBuffers.push( commandEncoder.finish() );

				this.device.queue.submit( commandBuffers );

			}
		}

		this.gravityShader.setVariable( "updateDistancesAndIndices", 1 );

		await this.copyIndicesShader.execute( this.numWorkgroups );

	}

	async findStartEndIndices( logBuffers ) {

		await this.gravityShader.execute( this.numWorkgroups );

		await this.gravityShader.setVariable("updateDistancesAndIndices", 0);

		await this.findGridHashShader.execute( this.numWorkgroups );

		await this.sortGridHash();

		await this.findGridHashRangeShader.execute( this.workgroupSize );

		if( logBuffers ) { 

			await this.bitonicSortGridHashShader.debugBuffer("indices");

			await this.bitonicSortGridHashShader.debugBuffer("gridHashes");

			await this.findGridHashRangeShader.debugBuffer("startIndices");

			await this.findGridHashRangeShader.debugBuffer("endIndices");
		}

		await this.gravityShader.setVariable("updateDistancesAndIndices", 1);

	}



	async render() {

		this.updateTimeDelta();

		const viewMatrixData 			= this.camera.getViewMatrix();

		const projectionMatrixData 		= Matrix4.createProjectionMatrix( this.camera, this.canvas )

		const viewProjectionMatrix 		= Matrix4.multiply( projectionMatrixData, viewMatrixData );

		const cameraWorldMatrix 		= Matrix4.invert( viewMatrixData );  

		const cameraRight				= Matrix4.getColumn( cameraWorldMatrix, 0 );

		const cameraUp					= Matrix4.getColumn( cameraWorldMatrix, 1 );

		const cameraPosition			= Matrix4.getColumn( cameraWorldMatrix, 3 );


		this.gravityShader.execute( this.numWorkgroups );

		this.gravityShader.setVariable( "cameraPosition", cameraPosition );

		this.gravityShader.setVariable( "deltaTimeSeconds", this.deltaTimeValue );

		if ( this.frameCount % 20 === 0 ) {

			this.sortDistance();

		}

		this.renderShader.setVariable( "viewProjectionMatrix", viewProjectionMatrix );

		this.renderShader.setVariable( "cameraRight", cameraRight );

		this.renderShader.setVariable( "cameraUp",    cameraUp );

		this.renderShader.renderToCanvas( 6, this.particleCount, 0 );


		this.frameCount++;

		if ( this.frameCount % 6 === 0  ) {

			this.findStartEndIndices();

		}

		this.collisionDetectionShader.setVariable( "deltaTimeSeconds", this.deltaTimeValue );

		this.collisionDetectionShader.execute( this.numWorkgroups );

		requestAnimationFrame( this.render.bind( this ) );

	}

}