function [X, permX] = internal_dlconv_permuted(X,permX,weights,permWeights,bias,paddingSize,stride,dilation,numSpatialDims)
% This function is undocumented and will change in a future release.

% Applies d-dimensional convolution on the CPU.
%
% Inputs:
%   X           = Some permutation of N1 x N2 x...x Nd x C x B1 x ... x Bb.
%   permX       = [dimN1 dimN2 ... dimNd dimC dimB1 dimB2 ... dimBb]
%       Note: dimN1 is the dimension of X that contains the first spatial
%             dimension N1, and similar for the others.
%   weights     = Some permutation of M1 x M2 x...x Md x CG x FG x G.
%   permWeights = [dimM1 dimM2 ... dimMd dimCG dimFG dimG].
%       Definitions:
%          - CG = C/G     number of channels per group
%          - FG           number of filters per group
%          - G            number of groups
%   bias        = Vector of length K (any orientation). Treated as zero if empty.
%   paddingSize = [L1,H1,...,Ld,Hd] where Li and Hi are the start and end
%                 padding sizes in the ith dimension.
%   stride      = [s1,...,sd] where si is stride in the ith dimension.
%   dilation    = [c1,...,cd] where ci is dilation in the ith dimension.
%   numSpatialDim = d number of dimensions
%
% Outputs:
%   Z           = Some permutation of O1 x O2 x...x Od x K x B1 x ... x Bb.
%   permZ       = Permutation vector of length d+b+1.

%   Copyright 2020 The MathWorks, Inc.

d = numSpatialDims;
b = length(permX)-d-1;

% Mixed precision is not supported by builtin.
[X, weights, bias] = nnet.internal.cnnhost.castToSuperiorFloat(X, weights, bias);

% Convert padding to the format expected by the builtin.
paddingSize = iConvertPaddingToHostAPIformat(paddingSize);

% Standardize group dimensions in weights
[weights, permWeights] = iPreProcessWeights(weights, permWeights);

% Standardize (merge) batch dimensions in X
if b == 0
    permX(end+1) = length(permX) + 1;
elseif b > 1
    [X, permX, batchPerm, batchSizes] = iPreProcessBatch(X, permX, d);
end

if ~isempty(bias)
    bias = iReshapeBias(bias, d);
end

[X, permX] = builtin('_batchconv', {X, permX}, {weights, permWeights}, paddingSize, stride, dilation, bias);

% Unmerge batch dimensions in X
if b == 0
    db = permX(d+2);
    permX(d+2) = [];
    permX(permX > db) = permX(permX > db) - 1;
    
    X = reshape(X, [size(X, [1:db-1 db+1:ndims(X)]) 1]);
elseif b > 1
    [X, permX] = iPostProcessBatch(X, permX, b, d, batchPerm, batchSizes);
end

end


function bias = iReshapeBias(bias, d)
    dimBias = d+1;
    szBias = ones(1, max(2, dimBias));
    szBias(dimBias) = numel(bias);
    bias = reshape(bias, szBias);
end


function [weights, permWeights] = iPreProcessWeights(weights, permWeights)
d1 = permWeights(end-1);
d2 = permWeights(end);
if d1+1 == d2
    % The weights are in valid format (groups-related inputs are together
    % and in correct order).
    return
end

if size(weights, d2) == 1
    % Reshaping is enough here, remove singleton d2 and add a new singleton
    % dimension after d1.
    singletonSecond = true;
    if d1 < d2
        % Need to shift intermediate dimensions upwards to make space for
        % new d1+1 dimension
        [weights, permWeights] = iAddHighSingleton(weights, permWeights, d1, d2, singletonSecond);
    else
        % Need to shift intermediate dimensions downwards to make space for
        % new d1+1 dimension.
        [weights, permWeights] = iAddLowSingleton(weights, permWeights, d2, d1, singletonSecond);
    end
elseif size(weights, d1) == 1
    % Reshaping is enough here, remove singleton d1 and add a new singleton
    % dimension before d2.
    singletonSecond = false;
    if d1 < d2
        [weights, permWeights] = iAddLowSingleton(weights, permWeights, d1, d2, singletonSecond);
    else
        [weights, permWeights] = iAddHighSingleton(weights, permWeights, d2, d1, singletonSecond);
    end
else
    % Need to permute, use format S,...,S,CG,FG,G
    weights = permute(weights, permWeights);
    permWeights = 1:length(permWeights);
end
end


function [weights, permWeights] = iAddHighSingleton(weights, permWeights, low, high, addSingletonAfterLow)
% Eliminate singleton dimension high, add a new singleton dimension low+1,
% increment intermediate indices to make space for this.
% Illustration:
%  Input:  [1, ..., low-1, low, XXXXX, low+1, ..., high-1, high, high+1, ..., end]
%  Output: [1, ..., low-1, low, low+1, low+2, ..., high  , XXXX, high+1, ..., end]
% If addSingletonAfterLow is true, replace nlow = size(x, low) with the two
% tuple [nlow 1]. Otherwise, replace it with [1 nlow].

for ii=1:length(permWeights)-2
    p = permWeights(ii);
    if low < p && p < high
        permWeights(ii) = p + 1;
    end
end
permWeights(end-1) = low;
permWeights(end) = low+1;

szWeights = size(weights, 1:length(permWeights));
if addSingletonAfterLow
    % Second dimension of (low, low+1) is singleton
    szWeights = [szWeights(1:low) 1 szWeights(low+1:high-1) szWeights(high+1:end)];
else
    % First dimension of (low, low+1) is singleton
    szWeights = [szWeights(1:low-1) 1 szWeights(low:high-1) szWeights(high+1:end)];
end
weights = reshape(weights, szWeights);
end


function [weights, permWeights] = iAddLowSingleton(weights, permWeights, low, high, addSingletonAfterLow)
% Eliminate singleton dimension low, add a new singleton dimension high-1,
% decrement intermediate indices to make space for this.
% Illustration:
%  Input:  [1, ..., low-1, low, low+1,..., high-1, XXXXXX, high, ..., end]
%  Output: [1, ..., low-1, XXX, low  ,..., high-2, high-1, high, ..., end]
% If addSingletonAfterLow is true, replace nlow = size(x, low) with the two
% tuple [nlow 1]. Otherwise, replace it with [1 nlow].

for ii=1:length(permWeights)-2
    p = permWeights(ii);
    if low < p && p < high
        permWeights(ii) = p - 1;
    end
end
permWeights(end-1) = high-1;
permWeights(end) = high;

szWeights = size(weights, 1:length(permWeights));
if addSingletonAfterLow
    % Second dimension of (high, high+1) is singleton
    szWeights = [szWeights(1:low-1) szWeights(low+1:high) 1 szWeights(high+1:end)];
else
    % First dimension of (high, high+1) is singleton
    szWeights = [szWeights(1:low-1) szWeights(low+1:high-1) 1 szWeights(high:end)];
end
weights = reshape(weights, szWeights);
end


function [X, permX, batchPerm, batchSizes] = iPreProcessBatch(X, permX, d)
% Need to adjust X such that all batch dimensions are together:
batchDims = permX(d+2:end);
b = length(batchDims);
minBatch = min(batchDims);
maxBatch = max(batchDims);

if maxBatch-minBatch+1 ~= b
    % Block of batch dimensions is not contiguous. Permute to
    % format 'CB...BS...S'
    X = permute(X, permX([d+1:end 1:d]));
    permX = [b+2:b+d+1 1:b+1];
    
    batchDims = permX(d+2:end);
    minBatch = min(batchDims);
    maxBatch = max(batchDims);
end

% Store for post processing
batchPerm = batchDims - minBatch + 1; % permutation vector
batchSizes = size(X, minBatch:maxBatch);

% Shorten permX to merge batch dimensions
jj = 1;
for ii=1:d+b+1
    if permX(ii) <= minBatch
        permX(jj) = permX(ii);
        jj = jj+1;
    elseif permX(ii) > maxBatch
        permX(jj) = permX(ii) - b + 1;
        jj = jj+1;
    end
end
permX(d+3:end) = [];

% Reshape X to merge batch dimensions
szX = size(X, 1:d+b+1);
szX(minBatch) = prod(szX(minBatch:maxBatch));
szX(minBatch+1:maxBatch) = [];

X = reshape(X, szX);
end


function [X, permX] = iPostProcessBatch(X, permX, b, d, batchPerm, batchSizes)
batchInX = permX(end);

szX = size(X, 1:b+d+1);
szX = [szX(1:batchInX-1) batchSizes szX(batchInX+1:b+d+1)];
X = reshape(X, szX);

for ii=1:length(permX)
    if permX(ii) > batchInX
        permX(ii) = permX(ii) + b - 1;
    end
end
permX(end:end+b-1) = batchPerm+batchInX-1;
end


function paddingSize = iConvertPaddingToHostAPIformat(paddingSize)
% Rearrang [L1,H1,L2,H2,...,Ld,Hd] to be [L1,L2,...,Ld,H1,H2,...,Hd]
paddingSize = [paddingSize(1:2:end), paddingSize(2:2:end)];
end
