function lgraph = Inflated3D(numClasses,inputSize,min,max,twoDNet,rescaleWeights)
%INFLATED3D Summary of this function goes here
%   Detailed explanation goes here

if nargin == 5
    rescaleWeights = true;
end
lgraph = layerGraph(twoDNet);

if isempty(min)
    normalizationArgs = { 'Normalization', 'none'};
else
    normalizationArgs = { 'Normalization', 'rescale-symmetric',...
    'Min',min,...
    'Max',max };
end
inputLayer = image3dInputLayer(inputSize,...
    'Name','VideoInputLayer',...
    normalizationArgs{:});

numChannels = inputSize(4);
lgraph = replaceLayer(lgraph, lgraph.InputNames{1}, inputLayer);

convIdxes = findLayers(lgraph,'nnet.cnn.layer.Convolution2DLayer');
lgraph = replaceConvLayer(lgraph,lgraph.Layers(convIdxes(1)), rescaleWeights, numChannels);
for ii = 2:numel(convIdxes)
    lgraph = replaceConvLayer(lgraph,lgraph.Layers(convIdxes(ii)), rescaleWeights);
end

normIdxes = findLayers(lgraph,'nnet.cnn.layer.CrossChannelNormalizationLayer');

for ii = 1:numel(normIdxes)
    lgraph = replaceNormLayer(lgraph,lgraph.Layers(normIdxes(ii)));
end


poolIdxes = findLayers(lgraph, 'nnet.cnn.layer.MaxPooling2DLayer');

for ii = 1:numel(poolIdxes)
    lgraph = replacePoolLayer(lgraph,lgraph.Layers(poolIdxes(ii)));
end

concatIdxes = findLayers(lgraph, 'nnet.cnn.layer.DepthConcatenationLayer');


for ii = 1:numel(concatIdxes)
    lgraph = replaceConcatLayer(lgraph,lgraph.Layers(concatIdxes(ii)));
end


gPoolIdx = find(arrayfun(@(x)isa(x,'nnet.cnn.layer.GlobalAveragePooling2DLayer'), lgraph.Layers));
gPool = lgraph.Layers(gPoolIdx);
gPoolRep = globalAveragePooling3dLayer('Name',gPool.Name + "_inflated");

lgraph = replaceLayer(lgraph,gPool.Name, gPoolRep);

fcIdx = findLayers(lgraph, 'nnet.cnn.layer.FullyConnectedLayer');
fcLayer = lgraph.Layers(fcIdx);
newFc = fullyConnectedLayer(numClasses, ...
        'Name','newFullyConnectedLayer', ...
        'WeightLearnRateFactor',10, ...
        'BiasLearnRateFactor',10);
lgraph = replaceLayer(lgraph,fcLayer.Name, newFc);


classIdx = findLayers(lgraph,'nnet.cnn.layer.ClassificationOutputLayer');
if isempty(classIdx)
    classIdx = findLayers(lgraph,'nnet.layer.ClassificationLayer');    
end
classLayer = lgraph.Layers(classIdx);
lgraph = removeLayers(lgraph, classLayer.Name);
end


function lgraph = replaceConvLayer(lgraph, convLayer, rescaleWeights, numChannels)
if nargin == 3
    numChannels = [];
end
filterSize = [convLayer.FilterSize, convLayer.FilterSize(2)];
outputChannels = convLayer.NumFilters;
name = convLayer.Name + "_Inflated";
stride = [convLayer.Stride, convLayer.Stride(2)];
w = convLayer.Weights;

temporal = filterSize(end);

if ~isempty(numChannels)
    w = w(:,:,1:numChannels,:);
end

% NxN
% NxNXN/N

nw = repmat(w,1,1,1,1,temporal);
if rescaleWeights
    % To ensure convolutional filter response is the
    % same, divide the N repeated weights by N.
    nw = nw ./ temporal;
end
sz = size(w);
sz = [sz(1:2), temporal, sz(3:4)];
nw = reshape(nw, sz);

b = convLayer.Bias;
b = reshape(b,[1,size(b)]);

convLayerRep = convolution3dLayer(...
        filterSize,...
        outputChannels,... % NumFilters
        'Name', name, ...
        'Stride', stride,...
        'Weights', nw,...
        'Bias', b,...
        'Padding', 'same');
lgraph = replaceLayer(lgraph, convLayer.Name, convLayerRep);
end

function lgraph = replacePoolLayer(lgraph, poolLayer)
poolSize = [poolLayer.PoolSize, poolLayer.PoolSize(2)];
stride = [poolLayer.Stride, poolLayer.Stride(2)];

maxPool = maxPooling3dLayer(poolSize, ...
        'Name', poolLayer.Name + "_Inflated", ...
        'Stride', stride,...
        'Padding', 'same');
lgraph = replaceLayer(lgraph, poolLayer.Name, maxPool);
end

function lgraph = replaceNormLayer(lgraph, normLayer)
newNormLayer = groupNormalizationLayer(16, 'Name', normLayer.Name + "_Inflated");
lgraph = replaceLayer(lgraph, normLayer.Name, newNormLayer);
end

function lgraph = replaceConcatLayer(lgraph,depthcatLayer)
concat = concatenationLayer(4,4,'Name', depthcatLayer.Name + "_Inflated");
lgraph = replaceLayer(lgraph, depthcatLayer.Name, concat);
end

function idxes = findLayers(lgraph, name)
idxes = find(arrayfun(@(x)isa(x,name), lgraph.Layers));
end

