/**
 * Copyright (c) 2011, The University of Southampton and the individual contributors.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 *   *  Redistributions of source code must retain the above copyright notice,
 *      this list of conditions and the following disclaimer.
 *
 *   *  Redistributions in binary form must reproduce the above copyright notice,
 *      this list of conditions and the following disclaimer in the documentation
 *      and/or other materials provided with the distribution.
 *
 *   *  Neither the name of the University of Southampton nor the names of its
 *      contributors may be used to endorse or promote products derived from this
 *      software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
package org.openimaj.image.processing.algorithm;

import org.openimaj.image.FImage;

import edu.emory.mathcs.jtransforms.fft.FloatFFT_2D;

/**
 * Perform forward and inverse Fast Fourier Transforms on image data. This class
 * computes the result of the transform in polar form (i.e. in terms of phase
 * and magnitude). If you need the result of the transform in the raw complex
 * form, then use the {@link FourierTransformComplex} class.
 *
 * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
 *
 */
public class FourierTransform {
        private FImage phase;
        private FImage magnitude;
        private boolean centre;

        /**
         * Construct Fourier Transform by performing a forward transform on the
         * given image. If the centre option is set, the FFT will be re-ordered so
         * that the DC component is in the centre.
         *
         * @param image
         *            the image to transform
         * @param centre
         *            should the FFT be reordered so the centre is DC component
         */
        public FourierTransform(FImage image, boolean centre) {
                this.centre = centre;

                process(image);
        }

        /**
         * Construct Fourier Transform object from the given magnitude and phase
         * images in the frequency domain. The resultant object can then be used to
         * construct the image using the {@link #inverse()} method.
         *
         * @param magnitude
         *            the magnitude image
         * @param phase
         *            the phase image
         * @param centre
         *            is the DC component in the image centre?
         */
        public FourierTransform(FImage magnitude, FImage phase, boolean centre) {
                this.centre = centre;
                this.magnitude = magnitude;
                this.phase = phase;
        }

        /**
         * Prepare data for a input to the FFT, padding if necessary.
         *
         * @param input
         *            input data
         * @param rs
         *            desired number of rows
         * @param cs
         *            desired number of columns
         * @param centre
         *            if true, then the data will be prepared so that the DC
         *            component is centred.
         * @return prepared data
         */
        public static float[][] prepareData(FImage input, int rs, int cs, boolean centre) {
                return prepareData(input.pixels, rs, cs, centre);
        }

        /**
         * Prepare data for a input to the FFT, padding if necessary. The data is
         * prepared as a packed 1D array.
         *
         * @param input
         *            input data
         * @param rs
         *            desired number of rows
         * @param cs
         *            desired number of columns
         * @param centre
         *            if true, then the data will be prepared so that the DC
         *            component is centred.
         * @return prepared data
         */
        public static float[] prepareData1d(FImage input, int rs, int cs, boolean centre) {
                return prepareData1d(input.pixels, rs, cs, centre);
        }

        /**
         * Prepare data for a input to the FFT, padding if necessary.
         *
         * @param input
         *            input data
         * @param rs
         *            desired number of rows
         * @param cs
         *            desired number of columns
         * @param centre
         *            if true, then the data will be prepared so that the DC
         *            component is centered.
         * @return prepared data
         */
        public static float[][] prepareData(float[][] input, int rs, int cs, boolean centre) {
                final float[][] prepared = new float[rs][cs * 2];

                if (centre) {
                        for (int r = 0; r < Math.min(rs, input.length); r++) {
                                for (int c = 0; c < Math.min(cs, input[0].length); c++) {
                                        prepared[r][c * 2] = input[r][c] * (1 - 2 * ((r + c) % 2));
                                }
                        }
                } else {
                        for (int r = 0; r < Math.min(rs, input.length); r++) {
                                for (int c = 0; c < Math.min(cs, input[0].length); c++) {
                                        prepared[r][c * 2] = input[r][c];
                                }
                        }
                }

                return prepared;
        }

        /**
         * Prepare data for a input to the FFT, padding if necessary. The data is
         * prepared as a packed 1D array.
         *
         * @param input
         *            input data
         * @param rs
         *            desired number of rows
         * @param cs
         *            desired number of columns
         * @param centre
         *            if true, then the data will be prepared so that the DC
         *            component is centered.
         * @return prepared data
         */
        public static float[] prepareData1d(float[][] input, int rs, int cs, boolean centre) {
                final float[] prepared = new float[rs * cs * 2];

                if (centre) {
                        for (int r = 0; r < Math.min(rs, input.length); r++) {
                                for (int c = 0; c < Math.min(cs, input[0].length); c++) {
                                        prepared[r * 2 * cs + 2 * c] = input[r][c] * (1 - 2 * ((r + c) % 2));
                                }
                        }
                } else {
                        for (int r = 0; r < Math.min(rs, input.length); r++) {
                                for (int c = 0; c < Math.min(cs, input[0].length); c++) {
                                        prepared[r * 2 * cs + 2 * c] = input[r][c];
                                }
                        }
                }

                return prepared;
        }

        /**
         * Extract the actual data from prepared data. The output image must have
         * the same number of rows as the prepared data, and half the number of
         * columns.
         *
         * @param prepared
         *            the prepared data
         * @param output
         *            the output
         * @param centre
         *            if true, then the data will be prepared so that the DC
         *            component is centered.
         */
        public static void unprepareData(float[][] prepared, FImage output, boolean centre) {
                unprepareData(prepared, output.pixels, centre);
        }

        /**
         * Extract the actual data from prepared data. The output image must have
         * the same number of rows as the prepared data, and half the number of
         * columns.
         *
         * @param prepared
         *            the prepared data
         * @param output
         *            the output
         * @param centre
         *            if true, then the data will be prepared so that the DC
         *            component is centered.
         */
        public static void unprepareData(float[] prepared, FImage output, boolean centre) {
                unprepareData(prepared, output.pixels, centre);
        }

        /**
         * Extract the actual data from prepared data. The output array must have
         * the same number of rows as the prepared data, and half the number of
         * columns.
         *
         * @param prepared
         *            the prepared data
         * @param output
         *            the output
         * @param centre
         *            if true, then the data will be prepared so that the DC
         *            component is centered.
         */
        public static void unprepareData(float[][] prepared, float[][] output, boolean centre) {
                final int rs = output.length;
                final int cs = output[0].length;

                if (centre) {
                        for (int r = 0; r < rs; r++) {
                                for (int c = 0; c < cs; c++) {
                                        output[r][c] = prepared[r][c * 2] * (1 - 2 * ((r + c) % 2));
                                }
                        }
                } else {
                        for (int r = 0; r < rs; r++) {
                                for (int c = 0; c < cs; c++) {
                                        output[r][c] = prepared[r][c * 2];
                                }
                        }
                }
        }

        /**
         * Extract the actual data from prepared data. The output array must have
         * the same number of rows as the prepared data, and half the number of
         * columns.
         *
         * @param prepared
         *            the prepared data
         * @param output
         *            the output
         * @param centre
         *            if true, then the data will be prepared so that the DC
         *            component is centered.
         */
        public static void unprepareData(float[] prepared, float[][] output, boolean centre) {
                final int rs = output.length;
                final int cs = output[0].length;

                if (centre) {
                        for (int r = 0; r < rs; r++) {
                                for (int c = 0; c < cs; c++) {
                                        output[r][c] = prepared[r * 2 * cs + 2 * c] * (1 - 2 * ((r + c) % 2));
                                }
                        }
                } else {
                        for (int r = 0; r < rs; r++) {
                                for (int c = 0; c < cs; c++) {
                                        output[r][c] = prepared[r * 2 * cs + 2 * c];
                                }
                        }
                }
        }

        private void process(FImage image) {
                final int cs = image.getCols();
                final int rs = image.getRows();

                phase = new FImage(cs, rs);
                magnitude = new FImage(cs, rs);

                final FloatFFT_2D fft = new FloatFFT_2D(rs, cs);
                final float[][] prepared = prepareData(image.pixels, rs, cs, centre);

                fft.complexForward(prepared);

                for (int y = 0; y < rs; y++) {
                        for (int x = 0; x < cs; x++) {
                                final float re = prepared[y][x * 2];
                                final float im = prepared[y][1 + x * 2];

                                phase.pixels[y][x] = (float) Math.atan2(im, re);
                                magnitude.pixels[y][x] = (float) Math.sqrt(re * re + im * im);
                        }
                }
        }

        /**
         * Perform the inverse FFT using the underlying magnitude and phase images.
         * The resultant reconstructed image may need normalisation.
         *
         * @return the reconstructed image
         */
        public FImage inverse() {
                final int cs = magnitude.getCols();
                final int rs = magnitude.getRows();

                final FloatFFT_2D fft = new FloatFFT_2D(rs, cs);
                final float[][] prepared = new float[rs][cs * 2];
                for (int y = 0; y < rs; y++) {
                        for (int x = 0; x < cs; x++) {
                                final float p = phase.pixels[y][x];
                                final float m = magnitude.pixels[y][x];

                                final float re = (float) (m * Math.cos(p));
                                final float im = (float) (m * Math.sin(p));

                                prepared[y][x * 2] = re;
                                prepared[y][1 + x * 2] = im;
                        }
                }

                fft.complexInverse(prepared, true);

                final FImage image = new FImage(cs, rs);
                unprepareData(prepared, image, centre);

                return image;
        }

        /**
         * Get a log-normalised copy of the magnitude image suitable for displaying.
         *
         * @return a log-normalised copy of the magnitude image
         */
        public FImage getLogNormalisedMagnitude() {
                final FImage im = magnitude.clone();

                for (int y = 0; y < im.height; y++) {
                        for (int x = 0; x < im.width; x++) {
                                im.pixels[y][x] = (float) Math.log(im.pixels[y][x] + 1);
                        }
                }

                return im.normalise();
        }

        /**
         * @return the phase image
         */
        public FImage getPhase() {
                return phase;
        }

        /**
         * @return the magnitude image
         */
        public FImage getMagnitude() {
                return magnitude;
        }

        /**
         * @return true if the DC component is in the centre; false otherwise
         */
        public boolean isCentre() {
                return centre;
        }
}