importScripts( "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs/dist/tf.min.js" )
importScripts( "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-webgpu/dist/tf-backend-webgpu.js" )
importScripts( "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js" )
//importScripts( "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.webgpu.min.js" )

const create_hann = size                                             =>{
	return tf.tidy(() => {
		const indices = tf.range(0, size, 1);
		return tf.cos(indices.mul(2 * Math.PI / (size - 1))).mul(0.5).sub(0.5).mul(-2);

		/*
		const window = tf.cos(indices.mul(2 * Math.PI / (size - 1))).mul(0.5).sub(0.5).mul(-2);
		const normalizationFactor = window.sum();
		return window.div(normalizationFactor);
		*/
	});
}
const stft        =( signal, frameSize, hopSize, windowFunc = 'hann')=>{
  return tf.tidy(() => {
    // Get signal length and calculate number of frames
    const signalLength = signal.shape[0];
    const numFrames = Math.floor((signalLength - frameSize) / hopSize) + 1;

    // Generate window
    const window = create_hann( frameSize )

    // Prepare output array
    const frames = [];

    // Process each frame
    for (let i = 0; i < numFrames; i++) {
      const start = i * hopSize;
      const frame = signal.slice([start], [frameSize]);
      
      // Apply window and compute RFFT
      const windowedFrame = frame.mul(window);
      const fft = tf.spectral.rfft(windowedFrame);
      
      frames.push(fft);
    }

    // Stack frames into a 2D tensor [numFrames, fftLength]
    const spectrogram = tf.stack(frames);
    
    return spectrogram;
  });
}

const segment_size =  256
const fftSize      = 7680
const hop          = 1024

const chunk_size     = hop * ( segment_size - 1 )
const trim           = Math.floor( fftSize / 2 )
const generated_size = chunk_size - 2 * trim
const step           = chunk_size - fftSize

const init      = async args =>{
	await tf.setBackend( 'webgpu' )
	await tf.ready()

	self.postMessage ({ selector : "init" })
}
const get_spec  = async args =>{
	const spec = tf.tidy(()=>{
		const tensor_l        = tf.tensor( args.audio[ 0 ])
		const padded_tensor_l = tensor_l.pad([[ trim, trim ]])
		const complex_l       = stft( padded_tensor_l, fftSize, hop )
		const reals_l         = tf.real( complex_l )
		const imags_l         = tf.imag( complex_l )

		const tensor_r        = tf.tensor( args.audio[ 1 ])
		const padded_tensor_r = tensor_r.pad([[ trim, trim ]])
		const complex_r       = stft( padded_tensor_r, fftSize, hop )
		const reals_r         = tf.real( complex_r )
		const imags_r         = tf.imag( complex_r )
		
		real_l = reals_l.slice([ 0, 0 ], [ 256, 3072 ]).transpose([ 1, 0 ])
		imag_l = imags_l.slice([ 0, 0 ], [ 256, 3072 ]).transpose([ 1, 0 ])
		real_r = reals_r.slice([ 0, 0 ], [ 256, 3072 ]).transpose([ 1, 0 ])
		imag_r = imags_r.slice([ 0, 0 ], [ 256, 3072 ]).transpose([ 1, 0 ])
		
		const stacked    = tf.stack([ real_l, imag_l, real_r, imag_r ], 1 )
		const expanded   = stacked.expandDims( 0 )
		const transposed = expanded.transpose([ 0, 2, 1, 3 ])
		
		return transposed.dataSync()
	})
	self.postMessage({ selector : "audio_2_spec", index : args.index, spec })
}
const get_audio = async args =>{
	const audio = tf.tidy(()=>{
		const output      = tf.tensor( args.spec, [ 1, 4, 3072, 256 ] )
		const source1Real = output.slice([ 0, 0, 0, 0 ], [ 1, 1, 3072, 256 ]).squeeze([ 0 ])
		const source1Imag = output.slice([ 0, 1, 0, 0 ], [ 1, 1, 3072, 256 ]).squeeze([ 0 ])
		const source2Real = output.slice([ 0, 2, 0, 0 ], [ 1, 1, 3072, 256 ]).squeeze([ 0 ])
		const source2Imag = output.slice([ 0, 3, 0, 0 ], [ 1, 1, 3072, 256 ]).squeeze([ 0 ])

		//const totalLength  = mix[ 0 ].length - trim
		const chunkLength   = ( 256 - 1 ) * hop + fftSize
		const chunkSignal   = new Float32Array( chunkLength )
		const vocals_l      = new Float32Array( chunkLength )
		const vocals_r      = new Float32Array( chunkLength )

		const hannWindow = create_hann( fftSize )
		
		for( let source of [
			{ real : source1Real, imag : source1Imag, buffer : vocals_l }, 
			{ real : source2Real, imag : source2Imag, buffer : vocals_r }
		]){
			const realFull    = source.real.pad([[ 0, 0 ], [ 0, 3841 - 3072 ], [ 0, 0 ]])
			const imagFull    = source.imag.pad([[ 0, 0 ], [ 0, 3841 - 3072 ], [ 0, 0 ]])
			const realFrames  = realFull.transpose([ 0, 2, 1 ])
			const imagFrames  = imagFull.transpose([ 0, 2, 1 ])

			for( let t = 0; t < 256; ++t ){
				const realFrame = realFrames.gather( t, 1 )
				const imagFrame = imagFrames.gather( t, 1 )
				const stftFrame = tf.complex( realFrame, imagFrame )
				const timeFrame = tf.spectral.irfft( stftFrame )
				
				// Apply Hann window
				const windowedFrame = timeFrame.mul(hannWindow);
				const frameData = windowedFrame.dataSync();				
				
				// Calculate the offset and accumulate the frame data into chunkSignal
				const offset = t * hop
				for( let j = 0; j < 7680; ++j )
					chunkSignal[ offset + j ] += frameData[ j ]
				
				// Clean up
				realFrame.dispose();
				imagFrame.dispose();
				stftFrame.dispose();
				timeFrame.dispose();
				windowedFrame.dispose();
			}

			// Trim the chunkSignal to remove unwanted edges
			const trimmed = chunkSignal.slice( 3840, chunkLength - 3840 )
			
			// Accumulate into source.buffer
			//for( let j = 0; j < trimmed.length && /*i +*/ j < totalLength; ++j )
			for( let j = 0; j < trimmed.length && /*i +*/ j < chunkLength; ++j )
				source.buffer[ /*i +*/ j ] += trimmed[ j ]
		}
		return [ vocals_l, vocals_r ]
	})
	self.postMessage({ selector : "spec_2_audio", index : args.index, audio })
}

self.onmessage    = function( event ){
	switch( event.data.selector ){
		case "init"         : init({ model_array_buffer : event.data.model_array_buffer }); break	
		case "audio_2_spec": get_spec({ index : event.data.index, audio : event.data.audio }); break	
		case "spec_2_audio": get_audio({ index : event.data.index, spec : event.data.spec }); break
	}
}