///////////////////////////////////////////////////////////////////////////
//                                                                       //
// Program file name: Fourier2d.java                                     //
//                                                                       //
// © Tao Pang 2006                                                       //
//                                                                       //
// Last modified: January 18, 2006                                       //
//                                                                       //
// (1) This Java program is part of the book, "An Introduction to        //
//     Computational Physics, 2nd Edition," written by Tao Pang and      //
//     published by Cambridge University Press on January 19, 2006.      //
//                                                                       //
// (2) No warranties, express or implied, are made for this program.     //
//                                                                       //
///////////////////////////////////////////////////////////////////////////

// An example of performing the fast Fourier transform
// for a 2-dimensional array.

import java.lang.*;
import java.util.Random;
public class Fourier2d {
  static final int n1 = 128, n2 = 64, l1 = 4, l2 = 4;
  public static void main(String argv[]) {
    double fr[][] = new double[n1][n2];
    double fi[][] = new double[n1][n2];
    Random r = new Random();
    int m1 = 7, m2 = 6;

 // Assign the data with a random-number generator
    for (int i=0; i<n1; ++i) {
      for (int j=0; j<n2; ++j) {
        fr[i][j] = r.nextDouble();
        fi[i][j] = r.nextDouble();
      }
    }

 // Perform fft on the 2-dimensional data
    fft2d(fr, fi, m1, m2);

 // Output the result in every l1/l2 data steps
    for (int i=0; i<n1; i+=l1) {
      for (int j=0; j<n2; j+=l2) {
        System.out.println(fr[i][j] + "  " + fi[i][j]);
      }
    }
  }

// Method to carry out the fast Fourier transform for a 
// 2-dimenisonal array.  Here fr[][] and fi[][] are real 
// and imaginary parts in both the input and output.

  public static void fft2d(double fr[][], double fi[][],
    int m1, int m2) {
    int n1 = fr.length;
    int n2 = fr[0].length;
    double hr[] = new double[n2];
    double hi[] = new double[n2];
    double pr[] = new double[n1];
    double pi[] = new double[n1];

 // Perform fft on the 2nd index
    for (int i=0; i<n1; ++i) {
      for (int j=0; j<n2; ++j) {
        hr[j] = fr[i][j];
        hi[j] = fi[i][j];
      }
      fft(hr, hi, m2);
      for (int j=0; j<n2; ++j) {
        fr[i][j] = hr[j];
        fi[i][j] = hi[j];
      }
    }

 // Perform fft on the 1st index
    for (int j=0; j<n2; ++j) {
      for (int i=0; i<n1; ++i) {
        pr[i] = fr[i][j];
        pi[i] = fi[i][j];
      }
      fft(pr, pi, m1);
      for (int i=0; i<n1; ++i) {
        fr[i][j] = pr[i];
        fi[i][j] = pi[i];
      }
    }
  }

// Method to perform the fast Foruier transform.  Here
// fr[] and fi[] are the real and imaginary parts of 
// both the input and output data. 

  public static void fft(double fr[], double fi[],
    int m) {
    int n = fr.length;
    int nh = n/2;
    int np = (int) Math.pow(2, m);

 // Stop the program if the indices do not match
    if (np != n) {
      System.out.println("Index mismtch detected");
      System.exit(1);
    }

 // Rearrange the data to the bit-reversed order
    int k = 1;
    for (int i=0; i<n-1; ++i) {
      if (i < k-1) {
        double f1 = fr[k-1];
        double f2 = fi[k-1];
        fr[k-1] = fr[i];
        fi[k-1] = fi[i];
        fr[i] = f1;
        fi[i] = f2;
      }
      int j = nh;
      while (j < k) {
        k -= j;
        j /= 2;
      }
      k += j;
    }

 // Sum up the reordered data at all levels
    k = 1;
    for (int i=0; i<m; ++i) {
      double w = 0;
      int j = k;
      k = 2*j;
      for (int p=0; p<j; ++p) {
        double u =  Math.cos(w);
        double v = -Math.sin(w);
        w += Math.PI/j;
        for (int q=p; q<n; q+=k) {
          int r = q+j;
          double f1 = fr[r]*u - fi[r]*v;
          double f2 = fr[r]*v + fi[r]*u;
          fr[r] = fr[q] - f1;
          fr[q] += f1;
          fi[r] = fi[q] - f2;
          fi[q] += f2;
        }
      }
    }
  }
}