function [GDFT,GIDFT] = make_DFT_fatrices(varargin)
%
% function [GDFT,GIDFT] = make_DFT_fatrices(varargin)
%
% Create fatrices for N-D DFT (unitary or not) and IDFT, with or without
% centering.
%
% options:
%   'im_dims' - size of image/volumes to use as input
%   'dftdims' - dimensions of image/volume to perform DFT (default: 1:min(3,length(im_dims)))
%   'unitary' - use unitary or unnormalized DFT (default: false)
%   'dftshift' - recenter in DFT domain (default: false)
%   'imgshift' - recenter in image domain (default: false)
%
% Outputs:
%   GDFT - fatrix2 class for forward DFT
%   GIDFT - fatrix2 class for inverse DFT (provided for convenience)
%
% Copyright 2012-10-11, Daniel Weller, University of Michigan

arg.class = 'fatrix2';
arg.im_dims = [];
arg.dftdims = [];
arg.unitary = false;
arg.dftshift = false;
arg.imgshift = false;

arg = vararg_pair(arg, varargin);

ndim = length(arg.im_dims);
if isempty(arg.dftdims), arg.dftdims = 1:min(3,ndim); end
arg.dftdims = unique(arg.dftdims,'sorted'); % ensure specified in increasing order
arg.im_dims(end+1:arg.dftdims(end)) = 1;
N = prod(arg.im_dims);
arg.idim = [N,1];
arg.odim = [N,1];
K = prod(arg.im_dims(arg.dftdims));
if arg.unitary
    arg.scale = 1/sqrt(K);
else
    arg.scale = 1;
end
GDFT = fatrix2('arg',arg,'forw',@DFT_forw,'back',@IDFT_forw);

arg.idim = [N,1];
arg.odim = [N,1];
if ~arg.unitary
    arg.scale = 1/K;
end
GIDFT = fatrix2('arg',arg,'forw',@IDFT_forw,'back',@DFT_forw);

end

function x = DFT_forw(arg,x)

x = reshape(x,arg.im_dims); % reshape vector to image dimensions

for id = 1:length(arg.dftdims)
    dim = arg.dftdims(id);

    if arg.imgshift, x = ifftshift(x,dim); end % recenter if desired
    
    x = fft(x,arg.im_dims(dim),dim);
    
    if arg.dftshift, x = fftshift(x,dim); end % recenter if desired
end

x = arg.scale.*x(:); % expects a vector

end

function x = IDFT_forw(arg,x)

x = reshape(x,arg.im_dims);

for id = 1:length(arg.dftdims)
    dim = arg.dftdims(id);
    
    if arg.dftshift, x = ifftshift(x,dim); end % recenter if desired
    
    x = ifft(x,arg.im_dims(dim),dim);
    
    if arg.imgshift, x = fftshift(x,dim); end % recenter if desired
    
end

K = prod(arg.im_dims(arg.dftdims));
x = (arg.scale*K).*x(:);

end


