function [LP, codistr] = hFftnImpl(codistr, LP, fftSize, sym, isFFT)
% hFftnImpl: Implement an N-dimensional FFT for codistributor1d where N >=2
% The code begins with two cases that require no communication.  These
% return when they are complete.  Otherwise, the general case is used which
% requires redistribution to a different dimension so that the transform on
% the original distribution dimension can be done locally.

%   Copyright 2014-2022 The MathWorks, Inc.

    szFftOut = codistr.Cached.GlobalSize;
    distDim = codistr.Dimension;

    % Cases where transform can be done locally without redistribution
    numberOfDimsToTransform = numel(fftSize);
    numberOfDimsOfInput = numel(szFftOut);
    is2D = (numberOfDimsToTransform == 2);
    if is2D && distDim > 2
        % Dimensions to be transformed are local to each worker
        LP = iDoLocalFft2(LP, fftSize(1), fftSize(2), isFFT, sym);
        szFftOut(1:2) = fftSize;
        codistr = hGetNewForSize(codistr, szFftOut);
        return
    end
    % Pad fftSize and szFftOut with ones to be conform with each other and
    % with distDim
    if distDim > numberOfDimsToTransform
        fftSize(numberOfDimsToTransform+1:distDim) = 1;
    end
    if numberOfDimsToTransform > numberOfDimsOfInput
        szFftOut(numberOfDimsOfInput+1:numberOfDimsToTransform) = 1;
    end
    if fftSize(distDim) <= codistr.Partition(1)
        % even though the matrix is distributed on this dimension,
        % if the part of the matrix we want to transform is all on
        % the first worker, we can do it without redistribution.
        % We return immediately because this is an edge-case.
        newPartition = zeros(1, numel(codistr.Partition));
        newPartition(1) = fftSize(distDim);
        codistr = codistributor1d(distDim, newPartition, fftSize);
        localSz = fftSize;
        localSz(distDim) = newPartition(spmdIndex);
        if is2D
            LP = iDoLocalFft2(LP, localSz(1), localSz(2), isFFT, sym);
        else
            LP = iDoLocalFftn(LP, localSz, isFFT, sym);
        end
        return
    end

    % Assert that szFftOut is at least as long as fftSize since the
    % following code is based on that assumption
    assert(numel(szFftOut)>=numel(fftSize), ...
        'At this point, szFftOut must have at least as many elements than fftSize.');
        
    % General case
    % This algorithm computes the fft on all dimensions that are not
    % distributed, then redistributes to make the former distribution 
    % dimension local and computes the fft on the that dimension.  
    % The transformed matrix is not redistributed back to its original 
    % codistributor type at the end.

    % Construct list of dimensions that are not the distribution dimension
    localDims = 1:numel(fftSize);
    localDims(distDim) = [];
    % Transform all of the non-distribution dimensions
    for ldim = localDims
        LP = iDoLocalFft(LP, fftSize(ldim), ldim, isFFT, 'nonsymmetric');
    end 
    % Modify szFftOut with the resulting fftSize in the dimensions used in
    % the above loop
    szFftOut(localDims) = fftSize(localDims);
    
    codistr = hGetNewForSize(codistr, szFftOut);

    % Transform the distribution dimension
    [maxSz, index] = max(szFftOut(localDims));
    newDistDim = localDims(index);
    fftCodistr = codistributor1d(newDistDim, ...
                codistributor1d.defaultPartition(maxSz), szFftOut);
    LP = distributedutil.Redistributor.redistribute(codistr, ...
                                     LP, fftCodistr);
    % After the redistribution, the dimension on which the array had been 
    % distributed is a local dimension, so we do a local transform on it.
    LP = iDoLocalFft(LP, fftSize(distDim), distDim, isFFT, sym);
    szFftOut(distDim) = fftSize(distDim);
    codistr = hGetNewForSize(fftCodistr, szFftOut);
end

function localX = iDoLocalFft(localX, n, dim, isFFT, sym)
% This function does an fft on the local part of a matrix
    if isFFT
        localX = fft(localX, n, dim);
    else
        localX = ifft(localX, n, dim, sym);
    end
end

function localX = iDoLocalFft2(localX, mrows, ncols, isFFT, sym)
    if isFFT
        localX = fft2(localX, mrows, ncols);
    else
        localX = ifft2(localX, mrows, ncols, sym);
    end
end

function localX = iDoLocalFftn(localX, fftSize, isFFT, sym)
    if isFFT
        localX = fftn(localX, fftSize);
    else
        localX = ifftn(localX, fftSize, sym);
    end
end
