%% Define the Network Architecture of RaNet
% Define the convolutional neural network architecture.

layers = [
    image3dInputLayer([12 12 6 5], 'Normalization', 'none', 'Name', 'data')
    convolution3dLayer([3 3 3], 128, 'Padding', 'same', 'Name', 'conv1') 
    batchNormalizationLayer('Name', 'norm1')
    leakyReluLayer('Name', 'relu1')
    maxPooling3dLayer([3 3 3], 'Stride', [2 2 2], 'Padding', 'same', 'Name', 'pool1')

    convolution3dLayer([3 3 3], 256, 'Padding', 'same', 'Name', 'conv2')
    batchNormalizationLayer('Name', 'norm2')
    leakyReluLayer('Name', 'relu2')
    maxPooling3dLayer([3 3 3], 'Stride', [2 2 2], 'Padding', 'same', 'Name', 'pool2')

    convolution3dLayer([2 2 2], 512, 'Padding', 'same', 'Name', 'conv3')
    leakyReluLayer('Name', 'relu3')
    globalAveragePooling3dLayer('Name', 'pool3')
    
    fullyConnectedLayer(2048, 'Name', 'fc4')
    leakyReluLayer('Name', 'relu4')
    dropoutLayer('Name', 'drop4')

    fullyConnectedLayer(2048, 'Name', 'fc5')
    leakyReluLayer('Name', 'relu5')
    dropoutLayer('Name', 'drop5')

    fullyConnectedLayer(256, 'Name', 'fc6')
    leakyReluLayer('Name', 'relu6')
    dropoutLayer('Name', 'drop6')

    fullyConnectedLayer(32, 'Name', 'fc7')
    leakyReluLayer('Name', 'relu7')

    fullyConnectedLayer(2, 'Name', 'fc8')
    
    softmaxLayer('Name', 'prob')
    weightedClassificationLayer([0.05 0.95], 'output')];

analyzeNetwork(layers)


%% Create the Layer Graph of RxNet
% Create the layer graph variable to contain the network layers.

lgraph = layerGraph();
%% Add Layer Branches
% Add the branches of the network to the layer graph. Each branch is a linear 
% array of layers.

xLayers = [
    imageInputLayer([12 12 30],"Name","input_1","Normalization","rescale-symmetric")
    convolution2dLayer([3 3],128,"Name","block1_conv1","BiasLearnRateFactor",0,"Stride",[2 2],"Padding","same")
    batchNormalizationLayer("Name","block1_conv1_bn","Epsilon",0.001)
    reluLayer("Name","block1_conv1_act")
    convolution2dLayer([3 3],256,"Padding", "same","Name","block1_conv2","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block1_conv2_bn","Epsilon",0.001)
    reluLayer("Name","block1_conv2_act")];
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    convolution2dLayer([1 1],512,"Name","conv2d_1","BiasLearnRateFactor",0,"Padding","same")
    batchNormalizationLayer("Name","batch_normalization_1","Epsilon",0.001)];
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    groupedConvolution2dLayer([3 3],1,256,"Name","block2_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],512,"Padding","same","Name","block2_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block2_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block2_sepconv2_act")
    groupedConvolution2dLayer([3 3],1,512,"Name","block2_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],512,"Name","block2_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block2_sepconv2_bn","Epsilon",0.001)
    maxPooling2dLayer([3 3],"Name","block2_pool","Padding","same")];
lgraph = addLayers(lgraph,xLayers);

xLayers = additionLayer(2,"Name","add_1");
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    convolution2dLayer([1 1],728,"Name","conv2d_2","BiasLearnRateFactor",0,"Padding","same")
    batchNormalizationLayer("Name","batch_normalization_2","Epsilon",0.001)];
lgraph = addLayers(lgraph,xLayers);
 
xLayers = [
    reluLayer("Name","block3_sepconv1_act")
    groupedConvolution2dLayer([3 3],1,512,"Name","block3_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block3_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block3_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block3_sepconv2_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block3_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block3_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block3_sepconv2_bn","Epsilon",0.001)
    maxPooling2dLayer([3 3],"Name","block3_pool","Padding","same")];
lgraph = addLayers(lgraph,xLayers);
% 
xLayers = additionLayer(2,"Name","add_2");
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    reluLayer("Name","block4_sepconv1_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block4_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block4_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block4_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block4_sepconv2_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block4_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block4_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block4_sepconv2_bn","Epsilon",0.001)
    reluLayer("Name","block4_sepconv3_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block4_sepconv3_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block4_sepconv3_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block4_sepconv3_bn","Epsilon",0.001)];
lgraph = addLayers(lgraph,xLayers);

xLayers = additionLayer(2,"Name","add_3");
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    reluLayer("Name","block5_sepconv1_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block5_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block5_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block5_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block5_sepconv2_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block5_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block5_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block5_sepconv2_bn","Epsilon",0.001)
    reluLayer("Name","block5_sepconv3_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block5_sepconv3_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block5_sepconv3_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block5_sepconv3_bn","Epsilon",0.001)];
lgraph = addLayers(lgraph,xLayers);

xLayers = additionLayer(2,"Name","add_4");
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    reluLayer("Name","block6_sepconv1_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block6_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block6_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block6_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block6_sepconv2_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block6_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block6_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block6_sepconv2_bn","Epsilon",0.001)
    reluLayer("Name","block6_sepconv3_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block6_sepconv3_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block6_sepconv3_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block6_sepconv3_bn","Epsilon",0.001)];
lgraph = addLayers(lgraph,xLayers);

xLayers = additionLayer(2,"Name","add_5");
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    reluLayer("Name","block7_sepconv1_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block7_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block7_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block7_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block7_sepconv2_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block7_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block7_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block7_sepconv2_bn","Epsilon",0.001)
    reluLayer("Name","block7_sepconv3_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block7_sepconv3_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block7_sepconv3_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block7_sepconv3_bn","Epsilon",0.001)];
lgraph = addLayers(lgraph,xLayers);

xLayers = additionLayer(2,"Name","add_6");
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    reluLayer("Name","block8_sepconv1_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block8_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block8_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block8_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block8_sepconv2_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block8_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block8_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block8_sepconv2_bn","Epsilon",0.001)
    reluLayer("Name","block8_sepconv3_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block8_sepconv3_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block8_sepconv3_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block8_sepconv3_bn","Epsilon",0.001)];
lgraph = addLayers(lgraph,xLayers);

xLayers = additionLayer(2,"Name","add_7");
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    reluLayer("Name","block9_sepconv1_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block9_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block9_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block9_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block9_sepconv2_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block9_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block9_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block9_sepconv2_bn","Epsilon",0.001)
    reluLayer("Name","block9_sepconv3_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block9_sepconv3_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block9_sepconv3_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block9_sepconv3_bn","Epsilon",0.001)];
lgraph = addLayers(lgraph,xLayers);

xLayers = additionLayer(2,"Name","add_8");
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    convolution2dLayer([1 1],1024,"Name","conv2d_3","BiasLearnRateFactor",0,"Padding","same","Stride",[2 2])
    batchNormalizationLayer("Name","batch_normalization_3","Epsilon",0.001)];
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    reluLayer("Name","block10_sepconv1_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block10_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],728,"Name","block10_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block10_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block10_sepconv2_act")
    groupedConvolution2dLayer([3 3],1,728,"Name","block10_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],1024,"Name","block10_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block10_sepconv2_bn","Epsilon",0.001)
    maxPooling2dLayer([3 3],"Name","block10_pool","Padding","same","Stride",[2 2])];
lgraph = addLayers(lgraph,xLayers);

xLayers = [
    additionLayer(2,"Name","add_9")
    groupedConvolution2dLayer([3 3],1,1024,"Name","block11_sepconv1_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],1536,"Name","block11_sepconv1_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block11_sepconv1_bn","Epsilon",0.001)
    reluLayer("Name","block11_sepconv1_act")
    groupedConvolution2dLayer([3 3],1,1536,"Name","block11_sepconv2_channel-wise","BiasLearnRateFactor",0,"Padding","same")
    convolution2dLayer([1 1],2048,"Name","block11_sepconv2_point-wise","BiasLearnRateFactor",0)
    batchNormalizationLayer("Name","block11_sepconv2_bn","Epsilon",0.001)
    reluLayer("Name","block11_sepconv2_act")
    globalAveragePooling2dLayer("Name","avg_pool")   
    fullyConnectedLayer(1024,"Name","fc1")
    reluLayer('Name','fc1_act')
    fullyConnectedLayer(128,"Name","fc2")
    reluLayer('Name','fc2_act')
    fullyConnectedLayer(16,"Name","fc3")
    reluLayer('Name','fc3_act')
    fullyConnectedLayer(2,"Name","predictions")            
    softmaxLayer("Name","predictions_softmax")
    weightedClassificationLayer([1.0, 19.0]/20.0,"weightedClassificationLayer_predictions")];
  
lgraph = addLayers(lgraph,xLayers);

% clean up helper variable
clear xLayers;

%% Connect Layer Branches
% Connect all the branches of the network to create the network graph.

lgraph = connectLayers(lgraph,"block1_conv2_act","conv2d_1");
lgraph = connectLayers(lgraph,"block1_conv2_act","block2_sepconv1_channel-wise");
lgraph = connectLayers(lgraph,"batch_normalization_1","add_1/in2");
lgraph = connectLayers(lgraph,"block2_pool","add_1/in1");
lgraph = connectLayers(lgraph,"add_1","block3_sepconv1_act");
lgraph = connectLayers(lgraph,"add_1","conv2d_2");
lgraph = connectLayers(lgraph,"batch_normalization_2","add_2/in2");
lgraph = connectLayers(lgraph,"block3_pool","add_2/in1");
lgraph = connectLayers(lgraph,"add_2","block4_sepconv1_act");
lgraph = connectLayers(lgraph,"block4_sepconv3_bn","add_3/in1");
lgraph = connectLayers(lgraph,"add_2","add_3/in2");
lgraph = connectLayers(lgraph,"add_3","block5_sepconv1_act");
lgraph = connectLayers(lgraph,"add_3","add_4/in2");
lgraph = connectLayers(lgraph,"block5_sepconv3_bn","add_4/in1");
lgraph = connectLayers(lgraph,"add_4","block6_sepconv1_act");
lgraph = connectLayers(lgraph,"add_4","add_5/in2");
lgraph = connectLayers(lgraph,"block6_sepconv3_bn","add_5/in1");
lgraph = connectLayers(lgraph,"add_5","block7_sepconv1_act");
lgraph = connectLayers(lgraph,"add_5","add_6/in2");
lgraph = connectLayers(lgraph,"block7_sepconv3_bn","add_6/in1");
lgraph = connectLayers(lgraph,"add_6","block8_sepconv1_act");
lgraph = connectLayers(lgraph,"add_6","add_7/in2");
lgraph = connectLayers(lgraph,"block8_sepconv3_bn","add_7/in1");
lgraph = connectLayers(lgraph,"add_7","block9_sepconv1_act");
lgraph = connectLayers(lgraph,"add_7","add_8/in2");
lgraph = connectLayers(lgraph,"block9_sepconv3_bn","add_8/in1");
lgraph = connectLayers(lgraph,"add_8","conv2d_3");
lgraph = connectLayers(lgraph,"add_8","block10_sepconv1_act");
lgraph = connectLayers(lgraph,"batch_normalization_3","add_9/in2");
lgraph = connectLayers(lgraph,"block10_pool","add_9/in1");
% Plot Layers

analyzeNetwork(lgraph);

%% The following needs to be copied to a separate file named as "weightedClassificationLayer.m"
% 
% classdef weightedClassificationLayer < nnet.layer.ClassificationLayer
%                
%     properties
%         % Vector of weights corresponding to the classes in the training
%         % data
%         ClassWeights
%     end
% 
%     methods
%         function layer = weightedClassificationLayer(classWeights, name)
%             % layer = weightedClassificationLayer(classWeights) creates a
%             % weighted cross entropy loss layer. classWeights is a row
%             % vector of weights corresponding to the classes in the order
%             % that they appear in the training data.
%             % 
%             % layer = weightedClassificationLayer(classWeights, name)
%             % additionally specifies the layer name. 
% 
%             % Set class weights
%             layer.ClassWeights = classWeights;
% 
%             % Set layer name
%             if nargin == 2
%                 layer.Name = name;
%             end
% 
%             % Set layer description
%             layer.Description = 'Weighted cross entropy';
%         end
%         
%         function loss = forwardLoss(layer, Y, T)
%             % loss = forwardLoss(layer, Y, T) returns the weighted cross
%             % entropy loss between the predictions Y and the training
%             % targets T.
% 
%             N = size(Y, 4);
%             Y = squeeze(Y);
%             T = squeeze(T);
%             W = layer.ClassWeights;
%     
%             loss = -sum(W*(T.*log(Y)))/N;
%         end
%     end
% end


