# Fourier transform processor.   FFT's data in the physical domain and
# corrects for the peculiarities of numpy's fft algorithm.
# First we fft the data, then we take the negative sign of every other
# mode in each direction (x and y), then we
# divide the bad boy by (sqrt(res))^2 to get the correct intensity scaling,
# finally we fftshift everything.

from numpy.random import rand
from numpy import exp, sqrt, log, tan, pi, floor, zeros, ones
from scipy.special import gammaln
import numpy as np

def FFT(f):

    shape = f.shape;
    res=max(shape);

    if shape[0] == shape[1]:
      F= np.fft.fft2(f);
      F=F/(res);
      F[1:res:2,:]=-F[1:res:2,:];
      F[:,1:res:2]=-F[:,1:res:2];
    else:
      F = np.fft.fft(f,axis=0);
      F = F.ravel('F'); #create 1-d array, colum-major order (matlab style), not really nice
      F[1:res:2] =-F[1:res:2];
      F = F.reshape(shape[1],shape[0]).T; #back to original shape
      F=F/np.sqrt(res);

    return np.fft.fftshift(F);


def fft(a):
  """
  fft(a)
  Compute the one-dimensional discrete Fourier transform of a the way Matlab
  does. When a is a vector, the Fourier transform of the vector is
  returned. When a is a matrix, each column vector of a is
  transformed individually, and a new matrix containing the transformed
  column vectors of a is returned.
    
  Parameters
  ----------
  a : array_like
      1-D or 2-D input array (can be complex)
  Returns
  -------
  result : ndarray
      1-D or 2-D array of similar shape and type
      containing the discrete fourier transform of a

  See Also
  --------
  ifft 
  
  Notes
  -----
  Using the Numpy function fft on a matrix does not 
  produce results similar to what Matlab does.
  This helper function uses the Numpy functions to produce
  a resut that agrees with what Matlab does. 
  """  
  return transformVectors(a, np.fft.fft)

def ifft(a):
  """
  ifft(a)
  Compute the one-dimensional inverse discrete Fourier transform the
  way Matlab does. When a is a vector, the inverse Fourier transform
  of the vector is returned. When a is a matrix, each column vector
  of a is transformed individually, and a new matrix containing the
  transformed column vectors of a is returned.
    
  Parameters
  ----------
  a : array_like
      1-D or 2-D input array (can be complex)
      
  Returns
  -------
  out : ndarray
      1-D or 2-D array of similar shape and type
      containing the inverse discrete fourier transform of a

  See Also
  --------
  fft 
  
  Notes
  -----
  Using the Numpy function ifft on a matrix does not 
  produce results similar to what Matlab does.
  This helper function uses the Numpy functions to produce
  a resut that agrees with what Matlab does.
  """  
  return transformVectors(a, np.fft.ifft)

def transformVectors(a, transform):
  """
  transformVectors(a, transform)
  Transform a according to the given transform function. 
  When a is a vector, it applies the transform function to a
  and returns the result. When a is a matrix, each column vector
  of a is transformed individually using the given transform function,
  and a new matrix containing the transformed column vectors
  of a is returned. 
    
  Parameters
  ----------
  a : array_like
      1-D or 2-D input array
  transform : callable function
      This function takes a 1-D array as argument and returns
      a 1-D array of same size. This function is applied to a if 
      it is a vector or to the column vectors of a if a is a matrix.
  Returns
  -------
  out : ndarray
      1-D or 2-D array of similar shape and type
      containing the transformed data
  See Also
  --------
  fft, ifft 
  
  Notes
  -----
  This function is used by fft(a) and ifft(a).
  """  

  if a.ndim == 1:
    # a is a vector
    out = transform(a)
  else:
    # assume a is a matrix (2d array)
    shape = a.shape
    colCount = shape[1]
    #result = np.empty_like(a)
    out = np.zeros_like(a)
    # for each column vector in a
    for i in range(0, colCount):
      col = a[:,i]
      fft_col = transform(col)
      out[:,i] = fft_col
  
  return out