/*
	SuperCollider real time audio synthesis system
 Copyright (c) 2002 James McCartney. All rights reserved.
	http://www.audiosynth.com

 This program is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation; either version 2 of the License, or
 (at your option) any later version.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA
 */

//Partitioned Convolution, Nick Collins mid October 2008
//PartConv(in, fftsize, irbufnum, accumbufnum)


#include "FFT_UGens.h"

#include <stdio.h>

struct PartConv : public Unit {
	int m_counter;
	uint32 m_specbufnumcheck;
	float* m_fd_accumulate; //will be exactly fftsize*frames in size
	float * m_irspectra;
	int m_fd_accum_pos;
	int m_partitions; //number of frames impulse response is partitioned into
	int m_fullsize;

	//sliding window code
	int m_fftsize; //must be power of two
	//int m_windowsize; //also half fftsize, was partition size, just use nover2 for this
	int m_nover2;

	//int m_hopsize; //hopsize will be half fftsize
	//int m_shuntsize;
	int m_pos;
	float * m_inputbuf;
	float * m_spectrum;
	scfft* m_scfft;
	float * m_inputbuf2;
	float * m_spectrum2;
	scfft* m_scifft; //inverse
	int m_outputpos;
	float * m_output;

	//amortisation state
	int m_blocksize, m_sr;
	int m_spareblocks;
	int m_numamort; //will relate number of partitions to number of spare blocks
	int m_lastamort;
	int m_amortcount;
	int m_partitionsdone;
};

extern "C" {
	void PartConv_next(PartConv *unit, int inNumSamples);
	void PartConv_Ctor(PartConv* unit);
	void PartConv_Dtor(PartConv* unit);
}

void PartConv_Ctor( PartConv* unit )
{
	unit->m_fftsize= (int) ZIN0(1);
	unit->m_nover2= unit->m_fftsize>>1;

	unit->m_inputbuf= (float*)RTAlloc(unit->mWorld, unit->m_fftsize * sizeof(float));
	unit->m_spectrum= (float*)RTAlloc(unit->mWorld, unit->m_fftsize * sizeof(float));

	SCWorld_Allocator alloc(ft, unit->mWorld);
	unit->m_scfft = scfft_create(unit->m_fftsize, unit->m_fftsize, kRectWindow, unit->m_inputbuf, unit->m_spectrum, kForward, alloc);

	//inverse
	unit->m_inputbuf2= (float*)RTAlloc(unit->mWorld, unit->m_fftsize * sizeof(float));
	unit->m_spectrum2= (float*)RTAlloc(unit->mWorld, unit->m_fftsize * sizeof(float));
	//in place this time
	unit->m_scifft = scfft_create(unit->m_fftsize, unit->m_fftsize, kRectWindow, unit->m_inputbuf2, unit->m_spectrum2, kBackward, alloc);

	//debug test: changing scale factors in case amplitude summation is a problem
	//unit->m_scfft->scalefac=1.0/45.254833995939;
	//unit->m_scifft->scalefac=1.0/45.254833995939;

	unit->m_output= (float*)RTAlloc(unit->mWorld, unit->m_fftsize * sizeof(float));
	unit->m_outputpos=0;

	memset(unit->m_output, 0, unit->m_fftsize * sizeof(float));
	memset(unit->m_inputbuf, 0, unit->m_fftsize * sizeof(float));
	unit->m_pos=0;

	//get passed in buffer
	unit->m_fd_accumulate=NULL;

	uint32 bufnum = (uint32)ZIN0(2);
	SndBuf *buf;

	if (bufnum >= unit->mWorld->mNumSndBufs) {
		int localBufNum = bufnum - unit->mWorld->mNumSndBufs;
		Graph *parent = unit->mParent;
		if(localBufNum <= parent->localMaxBufNum) {
			buf = parent->mLocalSndBufs + localBufNum;
		} else {
			printf("PartConv Error: Invalid Spectral data bufnum %d \n", bufnum);
			SETCALC(*ClearUnitOutputs);
			unit->mDone = true;
			return;
		}
	}

	buf = unit->mWorld->mSndBufs + bufnum;

	unit->m_specbufnumcheck = bufnum;

	if (!buf->data) {
		//unit->mDone = true;
		printf("PartConv Error: Spectral data buffer not allocated \n");
		SETCALC(*ClearUnitOutputs);
		unit->mDone = true;
		return;
	}

	unit->m_irspectra = buf->data;
	unit->m_fullsize = buf->samples;
	unit->m_partitions=buf->samples/(unit->m_fftsize); //should be exact
	//printf("another check partitions %d irspecsize %d fftsize %d \n", unit->m_partitions,unit->m_fullsize, unit->m_fftsize);

	if((buf->samples % unit->m_fftsize !=0) || (buf->samples==0)) {
		printf("PartConv Error: fftsize doesn't divide spectral data buffer size or spectral data buffer size is zero\n");
		SETCALC(*ClearUnitOutputs);
		unit->mDone = true;
		return;
	}

	//CHECK SAMPLING RATE AND BUFFER SIZE
	unit->m_blocksize = unit->mWorld->mFullRate.mBufLength;
	//if(unit->m_blocksize!=64) printf("TPV complains: block size not 64, you have %d\n", unit->m_blocksize);
	unit->m_sr = unit->mWorld->mSampleRate;
	//if(unit->m_sr!=44100) printf("TPV complains: sample rate not 44100, you have %d\n", unit->m_sr);

	OUT0(0) = 0.f;

	if(unit->m_nover2 % unit->m_blocksize !=0) {
		printf("PartConv Error: block size doesn't divide partition size\n");
		SETCALC(*ClearUnitOutputs);
		unit->mDone = true;
		return;
	} else {

		//must be exact divisor
		int blocksperpartition = unit->m_nover2/unit->m_blocksize;

		unit->m_spareblocks = blocksperpartition-1;

		if(unit->m_spareblocks<1) {
			printf("PartConv Error: no spareblocks, amortisation not possible! \n");
			SETCALC(*ClearUnitOutputs);
			unit->mDone = true;
			return;
		}

		//won't be exact
		unit->m_numamort = (unit->m_partitions-1)/unit->m_spareblocks; //will relate number of partitions to number of spare blocks
		unit->m_lastamort= (unit->m_partitions-1)- ((unit->m_spareblocks-1)*(unit->m_numamort)); //allow for error on last one
		unit->m_amortcount= -1; //starts as flag to avoid any amortisation before have first fft done
		unit->m_partitionsdone=1;

		//printf("Amortisation stats partitions %d nover2 %d blocksize %d spareblocks %d numamort %d lastamort %d \n", unit->m_partitions,unit->m_nover2, unit->m_blocksize, unit->m_spareblocks, unit->m_numamort, unit->m_lastamort);

		unit->m_fd_accumulate= (float*)RTAlloc(unit->mWorld, unit->m_fullsize * sizeof(float));
		memset(unit->m_fd_accumulate, 0, unit->m_fullsize * sizeof(float));
		unit->m_fd_accum_pos=0;

		SETCALC(PartConv_next);
	}
}

void PartConv_Dtor(PartConv *unit)
{
	RTFree(unit->mWorld, unit->m_inputbuf);
	RTFree(unit->mWorld, unit->m_inputbuf2);
	RTFree(unit->mWorld, unit->m_spectrum);
	RTFree(unit->mWorld, unit->m_spectrum2);
	RTFree(unit->mWorld, unit->m_output);
	if (unit->m_fd_accumulate) RTFree(unit->mWorld, unit->m_fd_accumulate);

	SCWorld_Allocator alloc(ft, unit->mWorld);
	if(unit->m_scfft)  scfft_destroy(unit->m_scfft, alloc);
	if(unit->m_scifft) scfft_destroy(unit->m_scifft, alloc);
}

void PartConv_next( PartConv *unit, int inNumSamples )
{
	float *in = IN(0);
	float *out = OUT(0);
	int pos = unit->m_pos;

	//safety check
	if (!(unit->mWorld->mSndBufs + unit->m_specbufnumcheck)->data) {
		printf("PartConv Error: Spectral data buffer not allocated \n");
		ClearUnitOutputs(unit, inNumSamples);
		SETCALC(*ClearUnitOutputs);
		unit->mDone = true;
		return;
	}

	float * input= unit->m_inputbuf;
	float * output= unit->m_output;
	int outputpos= unit->m_outputpos;

	//into input buffer
	memcpy(input+pos, in, inNumSamples * sizeof(float));

	pos += inNumSamples;

	//if ready for new FFT
	int nover2 = unit->m_nover2;

	//assumes that blocksize perfectly divides windowsize
	if (pos == nover2) {
		//FFT this input, second half of input always zero
		//memset(input+unit->m_nover2, 0, sizeof(float)*unit->m_nover2);
		scfft_dofft(unit->m_scfft);

		//reset pos into input buffer
		pos=0;
		//reset outputpos
		outputpos= 0;

		float * spectrum = unit->m_spectrum;
		float * spectrum2 = unit->m_spectrum2;

		//multiply spectra and accumulate for all ir spectra across storage buffer

		int fftsize = unit->m_fftsize;
		int accumpos = unit->m_fd_accum_pos;
		float * accumbuffer = unit->m_fd_accumulate;

		float * irspectrum = unit->m_irspectra;

		int fullsize = unit->m_fullsize;

		//JUST DO FIRST ONE FOR NOW, AMORTISED FOR OTHERS
		//frames
		for (int i=0; i<1; ++i) {

			int irpos= (i*fftsize);
			int posnow= (accumpos+irpos)%fullsize;
			float * target= accumbuffer+posnow;
			float * ir= irspectrum+irpos;

			//real multiply for dc and nyquist
			target[0] += ir[0]*spectrum[0];
			target[1] += ir[1]*spectrum[1];

			//complex multiply for frequency bins
			for (int j=1; j<nover2; ++j) {
				int binposr= 2*j;
				int binposi= binposr+1;
				target[binposr] += (ir[binposr]*spectrum[binposr]) - (ir[binposi]*spectrum[binposi]);
				target[binposi] += (ir[binposi]*spectrum[binposr]) + (ir[binposr]*spectrum[binposi]);
			}
		}

		//IFFT this partition
		float * input2 = unit->m_inputbuf2;
		memcpy(input2, accumbuffer+accumpos, fftsize * sizeof(float));
		scfft_doifft(unit->m_scifft);

		//shunt output data down and zero top bit
		memcpy(output, output+nover2, nover2 * sizeof(float));
		memset(output+nover2, 0, nover2 * sizeof(float));

		//sum into output
		for (int j=0; j<fftsize; ++j)
			output[j] += spectrum2[j];

		//zero this partition
		memset(accumbuffer+accumpos, 0, fftsize * sizeof(float));

		//ONLY DO AT END OF AMORTISATION???? no, amort code has extra -1 in indexing to cope
		//update partition counter
		accumpos= (accumpos+fftsize)%fullsize;
		unit->m_fd_accum_pos= accumpos;

		//set up for amortisation (calculate output for other partitions of impulse response)
		unit->m_amortcount=0;
		unit->m_partitionsdone=1;
	} else {
		//amortisation steps:
		//complex multiply of this new fft spectrum against existing irspectrums and sum to accumbuffer
		if (unit->m_amortcount>=0) {
			float * spectrum= unit->m_spectrum;

			//multiply spectra and accumulate for all ir spectra across storage buffer

			int fftsize= unit->m_fftsize;
			int nover2= unit->m_nover2;
			//int frames= unit->m_partitions;
			int accumpos= unit->m_fd_accum_pos;
			float * accumbuffer= unit->m_fd_accumulate;

			float * irspectrum= unit->m_irspectra;

			int fullsize= unit->m_fullsize;

			int starti, stopi;
			int number;

			if(unit->m_amortcount==(unit->m_spareblocks-1)) {
				number= unit->m_lastamort;
			}else{
				number= unit->m_numamort;
			}

			starti= unit->m_partitionsdone;//-1;
			stopi= starti+number-1;

			//printf("amort check count %d starti %d stopi %d number %d framesdone %d \n",unit->m_amortcount, starti, stopi, number, unit->m_partitionsdone);

			unit->m_partitionsdone += number;
			++unit->m_amortcount;

			for (int i=starti; i<=stopi; ++i) {
				int posnow= (accumpos+((i-1)*fftsize))%fullsize;
				float * target= accumbuffer+posnow;
				int irpos= (i*fftsize);
				float * ir= irspectrum+irpos;

				//real multiply for dc and nyquist
				target[0]+=  ir[0]*spectrum[0];
				target[1]+=	 ir[1]*spectrum[1];
				//complex multiply for frequency bins
				for (int j=1; j<nover2; ++j) {
					int binposr= 2*j;
					int binposi= binposr+1;
					target[binposr]+=  (ir[binposr]*spectrum[binposr]) - (ir[binposi]*spectrum[binposi]);
					target[binposi]+=  (ir[binposi]*spectrum[binposr]) + (ir[binposr]*spectrum[binposi]);
				}
			}
		}
	}

	//do this second!
	memcpy(out, output+outputpos, inNumSamples * sizeof(float));

	//debugging tests: output values tend to be fftsize too big, probably due to complex multiply and also summation over all partitions
	//	RGen& rgen = *unit->mParent->mRGen;
	//	int testindex= rgen.irand(inNumSamples-1);
	//	printf("inNumSamples %d testindex %d out %f output %f \n",inNumSamples, testindex, out[testindex], *(output+outputpos+testindex));

	outputpos+=inNumSamples;

	unit->m_outputpos= outputpos;
	unit->m_pos= pos;
}

//buffer preparation
void PreparePartConv(World *world, struct SndBuf *buf, struct sc_msg_iter *msg)
{
	// 'channels' not used- should just be mono, num frames= num samples
	float *data1 = buf->data;
	uint32 frombufnum = msg->geti();
	int fftsize = msg->geti();

	//output size must be frombuf->frames*2
	if (frombufnum >= world->mNumSndBufs) frombufnum = 0;
	SndBuf* frombuf = world->mSndBufs + frombufnum;
	int frames2 = frombuf->frames;
	float *data2 = frombuf->data;

	//scfft
	int nover2= fftsize>>1;

	int numpartitions;
	if(frames2 % nover2 == 0){
		numpartitions= frames2/nover2;
	}else{
		numpartitions= (frames2/nover2)+1;
	}

	//printf("reality check numpartitions %d fftsize %d product %d numinputframes %d \n", numpartitions, fftsize, numpartitions*fftsize, frames2);

	float * inputbuf= (float*)RTAlloc(world, fftsize * sizeof(float));
	float * spectrum= (float*)RTAlloc(world, fftsize * sizeof(float));

	SCWorld_Allocator alloc(ft, world);
	scfft* m_scfft = scfft_create(fftsize, fftsize, kRectWindow, inputbuf, spectrum, kForward, alloc);

	memset(inputbuf, 0, sizeof(float)*fftsize); // for zero padding

	//run through input data buffer, taking nover2 chunks, zero padding each
	for (int i=0; i<numpartitions; ++i) {
		int indexnow= nover2*i;
		int indexout= fftsize*i;

		if(i<(numpartitions-1)){
			memcpy(inputbuf, data2+indexnow, nover2 * sizeof(float));
		}else{
			int takenow = frames2 % nover2;
			if(takenow == 0){
				takenow = nover2;
			}

			memcpy(inputbuf, data2+indexnow, takenow * sizeof(float));
			if(takenow<nover2){
				memset(inputbuf+takenow, 0, (nover2-takenow)*sizeof(float));
			}
		}
		scfft_dofft(m_scfft);
		memcpy(data1+indexout, spectrum, fftsize * sizeof(float));
	}
	//clean up
	RTFree(world, inputbuf);
	RTFree(world, spectrum);

	if(m_scfft) scfft_destroy(m_scfft, alloc);
}

void initPartConv(InterfaceTable *inTable)
{
	ft = inTable;
	DefineDtorCantAliasUnit(PartConv);
	DefineBufGen("PreparePartConv", PreparePartConv);
}
