%% shape prior using stacked autoencoder

clc;

clear all;

close all;

addpath('functions');

%% STEP 0: 

% parameters

patchsize = 40;

visibleSize = patchsize*patchsize;   % number of input units 

hiddenSizeL1 = 100;     % number of hidden units 

hiddenSizeL2=100;

sparsityParam1 = 0.01;   % desired average activation of the hidden units.

sparsityParam2=0.01;

lambda = 3e-3;       % weight decay parameter       

beta = 3;            % weight of sparsity penalty term       

outputSize=visibleSize; % number of output units

%%======================================================================

%% STEP 1: laod training inputs and labels from mat files

load matFiles/training_dataES; 

train_input=sampleIMAGES(t_Iroi,patchsize);

train_labels=sampleIMAGES(t_yLV,patchsize);

%% train sparse Auto Encoder 1

%  Randomly initialize the parameters

saeTheta1 = initializeParameters(hiddenSizeL1, visibleSize);

%  Use minFunc to minimize the function

addpath minFunc/

options.Method = 'lbfgs'; % Here, we use L-BFGS to optimize our cost

                          % function. Generally, for minFunc to work, you

                          % need a function pointer with two outputs: the

                          % function value and the gradient. In our problem,

                          % sparseAutoencoderCost.m satisfies this.

options.maxIter = 400;    % Maximum number of iterations of L-BFGS to run 

options.display = 'on';

[sae1OptTheta, cost] = minFunc( @(p) sparseAutoencoderCost(p, ...

                                   visibleSize, hiddenSizeL1, ...

                                   lambda, sparsityParam1, ...

                                   beta, train_input), ...

                                   saeTheta1, options);

%% STEP 5: Visualization of AE1

W1 = reshape(sae1OptTheta(1:hiddenSizeL1*visibleSize), hiddenSizeL1, visibleSize);

display_network(W1', 12); 

%% compute activations from layer 1

[sae1Features] = feedForwardAutoencoder(sae1OptTheta, hiddenSizeL1, ...

                                        visibleSize, train_input);

%% train sparse Auto Encoder 2                                   

%  Randomly initialize the parameters

sae2Theta = initializeParameters(hiddenSizeL2, hiddenSizeL1);

[sae2OptTheta, costL2] = minFunc( @(p) sparseAutoencoderCost(p, ...

                                  hiddenSizeL1, hiddenSizeL2, ...

                                  lambda, sparsityParam2, ...

                                  beta, sae1Features), ...

                                  sae2Theta, options);

%% compute activation from layer 2

[sae2Features] = feedForwardAutoencoder(sae2OptTheta, hiddenSizeL2, ...

                                        hiddenSizeL1, sae1Features);

%% train multi outputs logstic regression                                    

lambda_mr=1e-4;

options_mr.maxIter = 100;

trainLabels=train_labels;

mrModel = mrTrain(hiddenSizeL2, outputSize, lambda_mr, ...

                            sae2Features, trainLabels, options_mr);

saeMultRegOptTheta = mrModel.optTheta(:);

%% fine tuning

% Initialize the stack using the parameters learned

stack = cell(2,1);

inputSize=visibleSize;

stack{1}.w = reshape(sae1OptTheta(1:hiddenSizeL1*inputSize), ...

                     hiddenSizeL1, inputSize);

stack{1}.b = sae1OptTheta(2*hiddenSizeL1*inputSize+1:2*hiddenSizeL1*inputSize+hiddenSizeL1);

stack{2}.w = reshape(sae2OptTheta(1:hiddenSizeL2*hiddenSizeL1), ...

                     hiddenSizeL2, hiddenSizeL1);

stack{2}.b = sae2OptTheta(2*hiddenSizeL2*hiddenSizeL1+1:2*hiddenSizeL2*hiddenSizeL1+hiddenSizeL2);

% Initialize the parameters for the deep model

[stackparams, netconfig] = stack2params(stack);

stackedAETheta = [ saeMultRegOptTheta ; stackparams ];

[stackedAEOptTheta, loss] = minFunc( @(x) stackedAECost(x, ...

      inputSize, hiddenSizeL2, outputSize, netconfig, ...

      lambda, train_input, train_labels), ...

      stackedAETheta, options);

%% test 

% load test data

load matFiles/validation_dataES; 

testI=t_Iroi;

test_input=sampleIMAGES(testI,patchsize);

[pred_yLV] = stackedAEPredict(stackedAEOptTheta, inputSize, hiddenSizeL2, ...

                          outputSize, netconfig, test_input);

% the final output is a mask of LV segmentation                      

yLVhr=reshape(pred_yLV,patchsize,patchsize,[]);

% scale to image size

scale=size(t_Iroi,1)/patchsize;

Mroi=100;

for k=1:size(yLVhr,3)

    y1=yLVhr(:,:,k);

    yLV_h=imresize(y1,scale);

end

%% dispaly segmentation

for k=1:30

    I1=testI(:,:,k);

    figure(1)

    subplot(5,6,k)

    imagesc(I1);

    colormap(gray);hold on

    contour(yLV_h,[0 0],'r','LineWidth',2); 

    contour(t_yLV,[0 0],'g','LineWidth',2);    

end

% one title for all subplots

set(gcf,'NextPlot','add');

axes;

h = title(['HiddenSize=',num2str(hiddenSizeL1),' sparsity=',num2str(sparsityParam1)]);

set(gca,'Visible','off');

set(h,'Visible','on');

%% save results

filename=['DLconfigure/LV_ES_V',num2str(patchsize),'_H',num2str(hiddenSizeL1),'_rho_',num2str(sparsityParam1*100),'div100'];

%save (filename);