package net.heartshield.filter;

import net.heartshield.signal.CVec;
import net.heartshield.signal.DVec;

import org.apache.commons.math3.complex.Complex;
import org.apache.commons.math3.transform.DftNormalization;
import org.apache.commons.math3.transform.FastFourierTransformer;
import org.apache.commons.math3.transform.TransformType;

/**
 * Realtime FIR filter using FFT convolution.
 *
 * Introduces a delay that depends on the filter parameters.
 * Also, the first 2*delay output samples are invalid (initial settling of the filter).
 *
 * @author David Madl (git@abanbytes.eu)
 * @date 2017-03-26
 */
public class FirFilter extends RFilterBlock {
    private double[] mTaps;
    private double mFps;
    private int mDelay;
    private double[] mBuf;

    public FirFilter(double fps, double[] taps) {
        mTaps = taps;
        mFps = fps;
        mDelay = taps.length / 2;
        mBuf = DVec.zeros(taps.length-1);
    }

    /** filter delay in number of samples */
    public int getDelay() {
        return mDelay;
    }

    public double getFps() {
        return mFps;
    }

    /*
    // instead, use convolve() using FFT to speed things up
    @Deprecated
    @Override
    protected double[] batch(double[] x) {
        double[] buf = DVec.stack(mBuf, x);
        int mathArraysLen = buf.length + mTaps.length - 1;
        int trimLen = mathArraysLen - buf.length;
        int trimLeft = trimLen / 2;
        int trimRight = trimLen - trimLeft;
        int trimEnd = mathArraysLen - trimRight;

        // just keep trailing bit to include it into leading boundary of next batch
        mBuf = DVec.get(buf, buf.length - mTaps.length+1, buf.length);


        return DVec.get(MathArrays.convolve(buf, mTaps), trimLeft, trimEnd);

        // to do: this has a slicing bug somewhere -- convolution using FFT yields correct results, while this convolution yields an offset between imag and real in Hilbert
        // to do: (meaning there is some indexing trouble)
    }
    */

    @Override
    protected double[] batch(double[] x) {
        double[] buf = DVec.stack(mBuf, x);

        // just keep trailing bit to include it into leading boundary of next batch
        mBuf = DVec.get(buf, buf.length - mTaps.length+1, buf.length);

        return convolveFft(buf, mTaps);
    }

    /**
     * Applies a filter to a signal, implementing the overlap-save method (chunk up the signal)
     * see https://en.wikipedia.org/wiki/Overlap%E2%80%93save_method
     *
     * This is equivalent to:  np.convolve(x, taps, mode='valid')  but more efficient for long signals
     */
    private static double[] convolveFft(double[] sig, double[] taps) {
        if(sig.length <= taps.length)
            throw new IllegalArgumentException("signal must be longer than taps");

        //double[] h = taps;
        int M = taps.length;
        int overlap = M - 1;
        int N = nextpow2(4*overlap);
        int stepSize = N - overlap;
        //double[][] dataRI = new double[2][N];
        //*   <li>{@code dataRI[0][i]} is the real part of the {@code i}-th data point,</li>
        //*   <li>{@code dataRI[1][i]} is the imaginary part of the {@code i}-th data point.</li>
        double[] h = new double[N];
        System.arraycopy(taps, 0, h, 0, taps.length);
        // rest is padded with zeros from the right

        FastFourierTransformer fft = new FastFourierTransformer(DftNormalization.STANDARD);
        Complex[] H = fft.transform(h, TransformType.FORWARD);

        double[] x = DVec.stack(sig, DVec.zeros(N)); // end padding, so the last step_size batch will surely cover the end
        double[] y = DVec.zeros(x.length);
        for(int pos = 0; pos < x.length + 1 - N; pos += stepSize) {
            Complex[] freqTransform = fft.transform(DVec.get(x, pos, pos+N), TransformType.FORWARD);
            Complex[] freqFiltered = CVec.mul(freqTransform, H);
            Complex[] yt = fft.transform(freqFiltered, TransformType.INVERSE);
            for(int i = 0; i < stepSize; i++)
                y[pos+i] = yt[M-1+i].getReal();
        }
        // cut back the end padding, and the overlap region where taps hang out of the signal
        return DVec.get(y, 0, sig.length-taps.length+1);
    }

    private static int nextpow2(int n) {
        double mf = Math.log(n) / Math.log(2);
        double mi = Math.ceil(mf);
        return (int) Math.pow(2, (int) mi);
    }
}
