package party.bami.eqfix.analysis;

import org.jtransforms.fft.FloatFFT_1D;

public class Deconvolver {
    private FloatFFT_1D fft;
    private long size;

    public Deconvolver (long size){
        this.fft = new FloatFFT_1D(size);
        this.size = size;
    }

    public float[] getKernel(float[] f, float[] h, int kernelSize) {



        float[] fftX = getMeanSpectrum(f);
//        fft.complexForward(fftX);

        float[] fftY = getMeanSpectrum(h);
//        fft.complexForward(fftY);

        float[] fftH = new float[f.length];

        for (int i = 0; i < this.size; i++) {
            float reX = fftX[2*i];
            float imX = fftX[2*i + 1];
            float reY = fftY[2*i];
            float imY = fftY[2*i + 1];

            fftH[2*i] = (reY*reX + imY*imX)/(reX*reX + imX * imX);
            fftH[2*i + 1] = (imY*reX - reY*imX)/(reX*reX + imX * imX);
        }

        fft.complexInverse(fftH, true);


        // convert impulse response to filter kernel of size kernelSize
        float[] temp = new float[fftH.length];

        for (int i = 0; i < fftH.length; i++) {
            int index = (i + kernelSize/2) % fftH.length;
            temp[index] = fftH[i];
        }

        for (int i = 0; i < fftH.length; i++) {
            fftH[i] = temp[i];
        }

        // apply  windowing function
        for (int i = 0; i < fftH.length; i++) {
            if (i <= kernelSize) {
                fftH[i] *= 0.54f - 0.46f * Math.cos(2f * Math.PI * i / kernelSize);
            } else {
                fftH[i] = 0;
            }
        }

        return trim(fftH, kernelSize + 1);
    }


    // divides the given array into mutiple slices of current size * 2 and calculates the mean
    private float[] getMeanSpectrum(float[] f){
        float[] fftF = new float[(int)this.size * 2];

        for (int i = 0; i < f.length/(this.size * 2); i++) {
            float[] temp = new float[(int)this.size * 2];

            // copy current part of f into temp
            for (int j = 0; j < this.size * 2; j++) {
                temp[j] = f[j + i * temp.length];
            }
            fft.complexForward(temp);

            // add to fftF
            for (int j = 0; j < this.size * 2; j++) {
                fftF[j] += temp[j];
            }
        }

        for (int i = 0; i < fftF.length; i++) {
            fftF[i] /= f.length/(this.size * 2);
        }

        return fftF;
    }


    // trims an array to a given size
    private float[] trim(float[] f, int size) {
        float[] output = new float[size];
        for (int i = 0; i < size; i++) {
            output[i] = f[i];
        }

        return output;
    }
}
