function [Gdec,Grec] = make_SWT_fatrices(varargin)
%
% function [Gdec,Grec] = make_SWT_fatrices(varargin)
%
% Create fatrices for N-D SWT (orthogonal or not) analysis/decomposition
% and synthesis/reconstruction. This code makes use of functionality in the
% MATLAB wavelet toolbox.
%
% options:
%   'im_dims' - size of image/volumes to use as input
%   'dwtdims' - dimensions of image/volume to perform DWT on (no more than three)
%   'wname' - wavelet name (should be recognized by MATLAB function wfilters()) (default: 'haar')
%   'nlevel' - number of decomposition levels (default: 2)
%
% Outputs:
%   Gdec - fatrix2 class for stationary DWT analysis/decomposition transform
%   Grec - fatrix2 class for stationary DWT synthesis/reconstruction transform
%
% Copyright 2013-02-05, Daniel Weller, University of Michigan

if nargin >= 1 && isequal(varargin{1},'test'), make_SWT_fatrices_test(); return; end

arg.class = 'fatrix2';
arg.im_dims = [];
arg.dwtdims = [];
arg.wname = 'haar'; % default wavelet filter name
arg.nlevel = 2; % wavelet transform with 2 levels 

arg = vararg_pair(arg, varargin);

ndim = length(arg.im_dims);
if isempty(arg.dwtdims), arg.dwtdims = 1:min(3,ndim); end
arg.dwtdims = unique(arg.dwtdims,'sorted');
[loD,hiD,loR,hiR] = wfilters(arg.wname); % get filters

min_sz = length(loD)*2^(arg.nlevel-1);
if any(arg.im_dims(arg.dwtdims) < min_sz)
    warn('Image is too small for DWT filter, # of levels');
end

numels = prod(arg.im_dims)*((2^length(arg.dwtdims)-1)*arg.nlevel + 1);
arg.idim = [prod(arg.im_dims),1];
arg.odim = [numels,1];

arg.lo = get_filter_ffts(loD,arg.im_dims,arg.nlevel,arg.dwtdims);
arg.hi = get_filter_ffts(hiD,arg.im_dims,arg.nlevel,arg.dwtdims);
Gdec = fatrix2('arg',arg,'forw',@SWT_forw,'back',@SWT_back);

idim = arg.odim; % swap in and out dimensions
arg.odim = arg.idim;
arg.idim = idim;
arg.lo = get_filter_ffts(loR(end:-1:1)./2,arg.im_dims,arg.nlevel,arg.dwtdims); % use recon filters
arg.hi = get_filter_ffts(hiR(end:-1:1)./2,arg.im_dims,arg.nlevel,arg.dwtdims);
Grec = fatrix2('arg',arg,'forw',@SWT_back,'back',@SWT_forw);

end

function y = SWT_forw(arg,x)

use_real = isreal(x);
ndwtdims = length(arg.dwtdims);

x = reshape(x,arg.im_dims);

% compute FFTs of x
for id = 1:ndwtdims
    dim = arg.dwtdims(id);
    x = fft(x,arg.im_dims(dim),dim);
end

y = {x};

% start decomposition
for il = 1:arg.nlevel
    x = y(1);
    
    for id = 1:ndwtdims
        dim = arg.dwtdims(id);
        xin = x;
        x = cell(length(xin)*2,1);
        
        % get filters
        lofft = arg.lo{il,id};
        hifft = arg.hi{il,id};

        % repeat filters
        lofft = repmat(lofft,[arg.im_dims(1:dim-1),1,arg.im_dims(dim+1:end)]);
        hifft = repmat(hifft,[arg.im_dims(1:dim-1),1,arg.im_dims(dim+1:end)]);
        
        for ix = 1:length(xin)
            % convolve by multiplying ffts
            x{2*ix-1} = xin{ix}.*lofft;
            x{2*ix} = xin{ix}.*hifft;
        end
    end
    
    % take inverse FFT's of output cells
    for ix = 1+(il < arg.nlevel):length(x)
        for id = 1:ndwtdims
            dim = arg.dwtdims(id);
            x{ix} = ifft(x{ix},arg.im_dims(dim),dim);
        end
        if use_real, x{ix} = real(x{ix}); end % SWT output should still be real if input is real
    end
    
    y = [x(:);y(2:end)];
end

% vectorize
y = cellfun(@(yc) yc(:),y,'UniformOutput',false);
y = cat(1,y{:});

end

function y = DWT_forw(arg,x)

use_real = isreal(x);
ndwtdims = length(arg.dwtdims);

x = reshape(x,arg.im_dims);

% compute FFTs of x
for id = 1:ndwtdims
    dim = arg.dwtdims(id);
    x = fft(x,arg.szs(1,dim),dim);
end

y = {x};

% start decomposition
for il = 1:arg.nlevel
    x = y(1);

    for id = 1:ndwtdims
        dim = arg.dwtdims(id);
        xin = x;
        x = cell(length(xin)*2,1);
        
        % get filters
        lofft = arg.lo{il,id};
        hifft = arg.hi{il,id};

        % repeat filters
        lofft = repmat(lofft,[arg.szs(il+1,1:dim-1),1,arg.szs(il,dim+1:end)]);
        hifft = repmat(hifft,[arg.szs(il+1,1:dim-1),1,arg.szs(il,dim+1:end)]);
        for ix = 1:length(xin)
            % convolve by multiplying ffts, then downsample by averaging ffts
            x{2*ix-1} = xin{ix}.*lofft;
            x{2*ix-1} = 0.5.*(subsdim(x{2*ix-1},1:arg.szs(il+1,dim),dim) + subsdim(x{2*ix-1},arg.szs(il+1,dim)+1:arg.szs(il,dim),dim));
            x{2*ix} = xin{ix}.*hifft;
            x{2*ix} = 0.5.*(subsdim(x{2*ix},1:arg.szs(il+1,dim),dim) + subsdim(x{2*ix},arg.szs(il+1,dim)+1:arg.szs(il,dim),dim));
        end
    end

    % take inverse FFT's of output cells
    for ix = 1+(il < arg.nlevel):length(x)
        for id = 1:ndwtdims
            dim = arg.dwtdims(id);
            x{ix} = ifft(x{ix},arg.szs(il+1,dim),dim);
        end
        if use_real, x{ix} = real(x{ix}); end % DWT output should still be real if input is real
    end
    
    y = [x(:);y(2:end)];
end

% vectorize
y = cellfun(@(yc) yc(:),y,'UniformOutput',false);
y = cat(1,y{:});

end

function x = SWT_back(arg,y)

use_real = isreal(y);
ndwtdims = length(arg.dwtdims);

% reassemble cell array using precomputed list of sizes
x = cell((2^ndwtdims-1)*arg.nlevel+1,1);
ico = 1; istart = 1;
nc = prod(arg.im_dims);
for il = arg.nlevel:-1:1
    for ic = 1+(il < arg.nlevel):2^ndwtdims
        x{ico} = reshape(y(istart:istart+nc-1),arg.im_dims);
        % do FFT of correct size
        for id = 1:ndwtdims
            dim = arg.dwtdims(id);
            x{ico} = fft(x{ico},arg.im_dims(dim),dim);
        end
        istart = istart + nc;
        ico = ico + 1;
    end
end

% do decomposition backwards
for il = arg.nlevel:-1:1
    y = x(1:2^ndwtdims);
    for id = ndwtdims:-1:1
        dim = arg.dwtdims(id);
        yin = y;
        y = cell(length(yin)/2,1);
        
        % get reverse filters
        lofft = conj(arg.lo{il,id});
        hifft = conj(arg.hi{il,id});

        % repeat filters
        lofft = repmat(lofft,[arg.im_dims(1:dim-1),1,arg.im_dims(dim+1:end)]);
        hifft = repmat(hifft,[arg.im_dims(1:dim-1),1,arg.im_dims(dim+1:end)]);
        
        for iy = 1:length(y)
            % convolve by multiplying DFTs, and average results together
            y{iy} = yin{2*iy-1}.*lofft + yin{2*iy}.*hifft;
        end
    end
    
    x = [y;x(2^ndwtdims+1:end)];
end

% take IDFT of final output
x = x{1};
for id = 1:ndwtdims
    dim = arg.dwtdims(id);
    x = ifft(x,arg.im_dims(dim),dim);
    x = subsdim(x,1:arg.im_dims(dim),dim); % remove zero-padding if added
end
if use_real, x = real(x); end % SWT output should still be real if input is real

% vectorize output
x = x(:);

end

function x = DWT_back(arg,y)

use_real = isreal(y);
ndwtdims = length(arg.dwtdims);

% reassemble cell array using precomputed list of sizes
x = cell((2^ndwtdims-1)*arg.nlevel+1,1);
ico = 1; istart = 1;
for il = arg.nlevel:-1:1
    szc = arg.szs(il+1,:); nc = prod(szc);
    for ic = 1+(il < arg.nlevel):2^ndwtdims
        x{ico} = reshape(y(istart:istart+nc-1),szc);
        % do FFT of correct size
        for id = 1:ndwtdims
            dim = arg.dwtdims(id);
            x{ico} = fft(x{ico},szc(dim),dim);
        end
        istart = istart + nc;
        ico = ico + 1;
    end
end

% do decomposition backwards
for il = arg.nlevel:-1:1
    y = x(1:2^ndwtdims);
    for id = ndwtdims:-1:1
        dim = arg.dwtdims(id);
        yin = y;
        y = cell(length(yin)/2,1);
        
        % get reverse filters
        lofft = conj(arg.lo{il,id});
        hifft = conj(arg.hi{il,id});

        % repeat filters
        lofft = repmat(lofft,[arg.szs(il+1,1:dim-1),1,arg.szs(il,dim+1:end)]);
        hifft = repmat(hifft,[arg.szs(il+1,1:dim-1),1,arg.szs(il,dim+1:end)]);
        
        for iy = 1:length(y)
            % upsample by copying DFT, convolve by multiplying DFTs, and
            % add results together
            y{iy} = cat(dim,yin{2*iy-1},yin{2*iy-1}).*lofft ...
                + cat(dim,yin{2*iy},yin{2*iy}).*hifft;
        end
    end
    
    x = [y;x(2^ndwtdims+1:end)];
end

% take IDFT of final output
x = x{1};
for id = 1:ndwtdims
    dim = arg.dwtdims(id);
    x = ifft(x,arg.szs(1,dim),dim);
    x = subsdim(x,1:arg.im_dims(dim),dim); % remove zero-padding if added
end
if use_real, x = real(x); end % DWT output should still be real if input is real

% vectorize output
x = x(:);

end

function szs = get_level_sizes(im_dims,nfilter,nlevel,dwtdims)

% get size that's divisible by two enough times
im_dims(dwtdims) = max(ceil(nfilter/2),ceil(im_dims(dwtdims)/2^nlevel)).*(2^nlevel);

szs = repmat(im_dims(:).',[nlevel+1,1]);
for ilvl = 1:nlevel
    szs(ilvl+1,dwtdims) = round(szs(ilvl,dwtdims)/2);
end

end

function hffts = get_filter_ffts(h,im_dims,nlevel,dwtdims)

ndwtdims = length(dwtdims);

hffts = cell(nlevel,ndwtdims);
for ilvl = 1:nlevel
    for id = 1:ndwtdims
        dim = dwtdims(id);
        hfft = fft(h(:),im_dims(dim));
        hffts{ilvl,id} = shiftdim(hfft,1-dim);
    end
    h = upsample(h,2); % insert zeros for next level
end

end


function make_SWT_fatrices_test()

Nx = 2; Ny = 256; Nz = 256; P = 12;
[SWT,ISWT] = make_SWT_fatrices('im_dims',[Nx,Ny,Nz],'dwtdims',[2,3],'nlevel',4,'wname','db4');

% test adjoints
test_adjoint(SWT,'big',1);
fprintf(1,'SWT adjoint test succeeded.\n'); drawnow;
test_adjoint(ISWT,'big',1);
fprintf(1,'ISWT adjoint test succeeded.\n'); drawnow;

% test recon
X = complex(randn([Nx*Ny*Nz,P],'double'),randn([Nx*Ny*Nz,P],'double'));

X2 = X;
N = 10; t = tic();
for n = 1:N
    W = SWT*X2;
    X2 = ISWT*W;
end
t = toc(t);

fprintf(1,'SWT error (%d trials in %g s): ||X2-X||/||X|| = %g.\n',N,t,norm(X2-X,'fro')/norm(X,'fro')); drawnow;

end

