import java.awt.Dimension;
import javax.swing.JFrame;
import org.jfree.chart.ChartFactory;
import org.jfree.chart.ChartPanel;
import org.jfree.chart.JFreeChart;
import org.jfree.data.xy.XYDataset;
import org.jfree.data.xy.XYSeries;
import org.jfree.data.xy.XYSeriesCollection;

/**
 * Description:
 * Attempt to convolve two 1D functions (original signal with response function ) using FFT, and then restore origin signal.
 *
 *
 *
 * Dependencies:
 * FFT.java ( source: https://introcs.cs.princeton.edu/java/97data/FFT.java.html )
 * Complex.java ( source: https://introcs.cs.princeton.edu/java/97data/Complex.java.html )
 * JFreeChart (for visualizing results)
 *
 *
 * **/

public class FFTConvolutionTest {

public static final int nfft = 1024; //FFT length

public static void main(String[] args) {

    //A first function for convolution : original signal is a sine curve
    double samplingFrequency = 150; //Hz
    double timeLength = 4; //sec
    double signalFrequency = 5; //Hz

    int samples = (int) (timeLength * samplingFrequency + 1);

    double[] signal = new double[samples];
    for (int i = 0; i < signal.length; i++) {
        signal[i] = Math.sin(2 * Math.PI * ((double) i / samplingFrequency) * signalFrequency);
    }
    visualizeArray(signal, 1d/samplingFrequency, "signal");

    //A second function for convolution : response function is a gaussian distribution
    double length = 1;
    double step = length/samplingFrequency;


    double[] response = new double[samples];
    double center = (int) (samples / 2);
    double stdDev = 0.3;
    double factor = 1/(Math.sqrt(2 * Math.PI * Math.pow(stdDev, 2)));
    double exponent;
    for (int i = 0; i < response.length; i++) {
        exponent = (double) - 1/2 * ( Math.pow(((double) i - center) * step /stdDev, 2));
        response[i] = Math.pow(Math.E, exponent) * factor;
    }
    visualizeArray(response, step, "response");


    //Making convolution
    double[] conv = convolution(signal, response);
    visualizeArray(conv, 1d/samplingFrequency, "convolved");

}

public static double[] convolution(double[] xd, double[] yd) {

    // fft length must be larger than both arrays lengths
    if ( (xd.length > nfft) || (yd.length > nfft)) {
        throw new IllegalArgumentException("nnft must be larger: " + xd.length + ", " + yd.length);
    }

    // Forming complex versions of arrays
    Complex[] x = new Complex[nfft];
    Complex[] y = new Complex[nfft];

    // Extending both x and y arrays to fft length, completing rests with zeros.
    for (int i=0; i< nfft; i++) {

        if (i < xd.length) {
            x[i] = new Complex(xd[i], 0);
        } else {
            x[i] = new Complex(0, 0);
        }

        if (i < yd.length) {
            y[i] = new Complex(yd[i], 0);
        } else {
            y[i] = new Complex(0, 0);
        }

    }

    // Doing the convolution
    Complex[] c = FFT.convolve(x, y);

    //returning real part of convolution
    double[] cd = new double[nfft];
    for (int i=0; i < cd.length; i++) {
        cd[i] = c[i].re();
    }
    return cd;
}

public static void visualizeArray(double[] array, double step, String name) {

    double[][] array2D = new double[array.length][2];
    for (int i=0; i<array2D.length; i++){
        array2D[i][0] = (double) i * step;
        array2D[i][1] = array[i];
    }

    final XYSeriesCollection xyCollection = new XYSeriesCollection();

    XYDataset xyDataset = xyCollection;
    XYSeries series = createSeries(array2D, name);
    xyCollection.addSeries(series);
    JFreeChart chart = ChartFactory.createXYLineChart(name, "x", "y", xyDataset);
    ChartPanel chartPanel = new ChartPanel(chart, true, true, true, false, true);

    JFrame frame = new JFrame();
    frame.add(chartPanel);
    frame.setPreferredSize(new Dimension(700, 500));
    frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
    frame.pack();
    frame.setLocationRelativeTo(null); // center the frame
    frame.setVisible(true);
}

public static XYSeries createSeries(double[][] array, String name){

    XYSeries series = new XYSeries(name);
    int[][] xyValues = new int[array.length][2];
    for (int i = 0; i < xyValues.length; i++) {
        series.add(array[i][0], array[i][1]);
    }
    return series;
}

}