clear; close all;

% load eigenfaces using training YaleB images
load eigenfaces.mat
%[eigenfaces,meanFace] = calculate_eigen_faces;
normScale = 1; 
eigenfaces = 128*eigenfaces;
meanFace = meanFace';

testFolder = './data/testFaces';
D_ir = dir(testFolder);
D_ir(1:2) = [];

% initializing variables
L = length(D_ir);
mestDiff = zeros(1,L);
ransacDiff = zeros(1,L);
msacDiff = zeros(1,L);
rpcaDiff = zeros(1,L);
torreDiff = zeros(1,L);
LrsRRDiff = zeros(1,L);


mestLabel = cell(1,L);
ransacLabel = cell(1,L);
msacLabel = cell(1,L);
rpcaLabel = cell(1,L);
torreLabel = cell(1,L);
LrsRRLabel = cell(1,L);


F = cell(1,L);
mestF = cell(1,L);
ransacF = cell(1,L);
msacF = cell(1,L);
rpcaF = cell(1,L);
torreF = cell(1,L);
LrsRRF = cell(1,L);

mestTime = zeros(1,L);
ransacTime = zeros(1,L);
msacTime = zeros(1,L);
rpcaTime = zeros(1,L);
torreTime = zeros(1,L);
LrsRRTime = zeros(1,L);

% save occluded image in first column, mask in second column, randomSample in
% third column, ground truth recovery in fourth column
imageBlock = cell(L,4);

mestImage = cell(L,1);
ransacImage = cell(L,1);
msacImage = cell(L,1);
rpcaImage = cell(L,1);
torreImage = cell(L,1);
LrsRRImage = cell(L,1);

for id = 1:length(D_ir)
    %% using clean image to calculate ground truth model
    Im_bound = imread([testFolder,'/',D_ir(id).name]);
    [imM,imN] = size(Im_bound);
    
    y_bound = reshape(Im_bound',[],1)';
    y_bound = double(y_bound);
    y_bound = y_bound-meanFace;
    y_bound = y_bound/normScale;
    faces = [eigenfaces;y_bound];
    % generate random samples
%     randomSample = randperm(size(faces,2));
%     randomSample(401:end) = [];

    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % to redo the experiments, we use the sample position
    load('experiment_image_data'); randomSample = imageBlock{id,3};
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    
    X = faces(:,randomSample);
    imageBlock{id,3} = randomSample;
    
    dim = size(X,1);
    fix = size(X,1); const = 1;
    
    % calculate the ground truth model using un-blocked images
    cvx_begin
    cvx_solver sedumi;
    variable f_inf(dim,1);
    variable res_bound(1);
    minimize res_bound
    subject to
    X'*f_inf >= -res_bound;
    X'*f_inf <= res_bound;
    f_inf(fix) == const;
    cvx_end
    
    disp(res_bound);
    epsilon = res_bound;
    dist_norm = epsilon;
    
    f_inf = f_inf/f_inf(fix);
    F{id} = f_inf;
    
    % save the best possible recovery as goal (ground truth)
    imageBlock{id,4} = -f_inf(1:end-1)'*eigenfaces*normScale+meanFace;

    %% read image and generate blocks
    Im = imread([testFolder,'/',D_ir(id).name]);
    
    N = 10; % number of blocks
    B = 30; % size of blockes
    I = Im;
    J = I;
    for n=1:N
        i = randi(imM-B);
        j = randi(imN-B);
        J(i:i+B-1,j:j+B-1) = ones(B,B)*255;
    end
    
    M = ones(imM,imN)*128;
    M = uint8(M);
    M(I~=J) = J((I~=J));
    Im = J;
    % save the blocked image and mask information
%     imageBlock{id,1} = Im;
%     imageBlock{id,2} = M;
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % to redo the experiments with saved mask and sample position
    Im = imageBlock{id,1};
    M = imageBlock{id,2};
    randomSample = imageBlock{id,3};
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

    y = reshape(Im',[],1)';
    y = double(y);
    y = y-meanFace;
    y = y/normScale;
    faces = [eigenfaces;y];
    X = faces(:,randomSample);

    

    %% run OLS
    A = X';
    b = -A(:,fix);
    A(:,fix) = [];
    tic;
    f_mest = regress(b,A);
    mestTime(id) = toc;
    f_mest = [f_mest(1:fix-1);1;f_mest(fix:end)]*const;
    mestLabel{id} = abs(X'*f_mest)<=epsilon;
    f_mest = f_mest/f_mest(fix);
    mestF{id} = f_mest;
    mestDiff(id) = regressorDiff(f_inf,f_mest);
    mestImage{id} = -f_mest(1:end-1)'*eigenfaces*normScale+meanFace;
    

    %% run RANsac
    tic;
    [ransacLabel{id},f_ransac] = RANSAC(X,@nptalg,dim-1,@residual_rr,1000,dist_norm);
    ransacTime(id) = toc;
    f_ransac = f_ransac'/f_ransac(fix);
    ransacF{id} = f_ransac;
    ransacDiff(id) = regressorDiff(f_inf,f_ransac);
    ransacImage{id} = -f_ransac(1:end-1)'*eigenfaces*normScale+meanFace;
    
    %% run Msac
    tic;
    [msacLabel{id},f_msac] = MSAC(X,@nptalg,dim-1,@residual_rr,1000,dist_norm);
    msacTime(id) = toc;
    f_msac = f_msac'/f_msac(fix);
    msacF{id} = f_msac;
    msacDiff(id) = regressorDiff(f_inf,f_msac);
    msacImage{id} = -f_msac(1:end-1)'*eigenfaces*normScale+meanFace;
    
    
    
    
    %% run cRPCA+LSR
    tau = sqrt(dist_norm^2*size(X,2));
    tic;
    [f_rpca,~,~,~,~] = rpca_admm(X,(0.1:0.1:11),tau,1.01,0);
    rpcaTime(id) = toc;
    f_rpca = f_rpca/f_rpca(fix)*const;
    rpcaLabel{id} = abs(X'*f_rpca)<=epsilon;
    f_rpca = f_rpca/f_rpca(fix);
    rpcaF{id} = f_rpca;
    rpcaDiff(id) = regressorDiff(f_inf,f_rpca);
    rpcaImage{id} = -f_rpca(1:end-1)'*eigenfaces*normScale+meanFace;
    
    %% run Torre Robust Regression
    tic;
    [f_torre,D,E,rr_Errs] = RR_torre(A',b',1,2,1.01,1e-3);
    torreTime(id) = toc;
    f_torre = [f_torre(1:fix-1),1,f_torre(fix:end)]'*const;
    torreLabel{id} = abs(X'*f_torre)<=epsilon;
    f_torre = f_torre/f_torre(fix);
    torreF{id} = f_torre;
    torreDiff(id) = regressorDiff(f_inf,f_torre);
    torreImage{id} = -f_torre(1:end-1)'*eigenfaces*normScale+meanFace;
    
     %% run LRS Robust Regression
    tic;
    [f_LrsRR,Lrs_D,~,Lrs_E,Lrs_Errs] = LRSR_RR_A_NoB(A',b',1,2,1.01,1e-3);
    LrsRRTime(id) = toc;
    f_LrsRR = [f_LrsRR(1:fix-1),1,f_LrsRR(fix:end)]'*const;
    %f_LrsRR = f_LrsRR'*const;
    torreLabel{id} = abs(X'*f_LrsRR)<=epsilon;
    f_LrsRR = f_LrsRR/f_LrsRR(fix);
    LrsRRF{id} = f_LrsRR;
    LrsRRDiff(id) = regressorDiff(f_inf,f_LrsRR);
    LrsRRImage{id} = -f_LrsRR(1:end-1)'*eigenfaces*normScale+meanFace;
    
    
end

%% calculate and display model error with running time
n = sqrt(length(D_ir));
mestMean = mean(mestDiff,2); mestSem = std(mestDiff,0,2)/n;
ransacMean = mean(ransacDiff,2); ransacSem = std(ransacDiff,0,2)/n;
msacMean = mean(msacDiff,2); msacSem = std(msacDiff,0,2)/n;
rpcaMean = mean(rpcaDiff,2); rpcaSem = std(rpcaDiff,0,2)/n;
torreMean = mean(torreDiff,2); torreSem = std(torreDiff,0,2)/n;
LrsRRMean = mean(LrsRRDiff,2); LrsRRSem = std(LrsRRDiff,0,2)/n;
summary = [torreMean,torreSem;LrsRRMean,LrsRRSem;rpcaMean,rpcaSem;msacMean,msacSem;...
    ransacMean,ransacSem;mestMean,mestSem]';
runtime = [mean(torreTime),mean(LrsRRTime),mean(rpcaTime),mean(msacTime),mean(ransacTime),mean(mestTime)];
summary = [summary;runtime];
disp('Model Error:');
disp([{'','LR-RR','LRS-RR','cRPCA','MSAC','RANSAC','OLS'};
    [{'Mean RMS';'stdev';'run time'},num2cell(summary)]]);

%% calculate and display fitting error to the original image with running time
Imreal = cell(L,1);
for id = 1:length(D_ir)
    Im = imread([testFolder,'/',D_ir(id).name]);
    Imreal{id} = double(reshape(Im',[],1)');
end

mestImMean = cellfun(@fitError,Imreal,mestImage);
mestImSem = std(mestImMean,0,1)/n;  mestImMean = mean(mestImMean);

ransacImMean = cellfun(@fitError,Imreal,ransacImage);
ransacImSem = std(ransacImMean,0,1)/n;  ransacImMean = mean(ransacImMean);
msacImMean = cellfun(@fitError,Imreal,msacImage);
msacImSem = std(msacImMean,0,1)/n;  msacImMean = mean(msacImMean);

rpcaImMean = cellfun(@fitError,Imreal,rpcaImage);
rpcaImSem = std(rpcaImMean,0,1)/n;  rpcaImMean = mean(rpcaImMean);

torreImMean = cellfun(@fitError,Imreal,torreImage);
torreImSem = std(torreImMean,0,1)/n;    torreImMean = mean(torreImMean);

LrsRRImMean = cellfun(@fitError,Imreal,LrsRRImage);
LrsRRImSem = std(LrsRRImMean,0,1)/n;    LrsRRImMean = mean(LrsRRImMean);

summaryIm = [torreImMean,torreImSem;LrsRRImMean,LrsRRImSem;rpcaImMean,rpcaImSem;...
    msacImMean,msacImSem;ransacImMean,ransacImSem;mestImMean,mestImSem]';
summaryIm = [summaryIm;runtime];
disp('Image Fitting Error:');
disp([{'','LS-RR','LRS-RR','cRPCA','MSAC','RANSAC','OLS'};
    [{'Mean RMS';'stdev';'run time'},num2cell(summaryIm)]]);

%% visualize results
testFolder = './data/testFaces';
D_ir = dir([testFolder,'/','yale*.pgm']);
im = imread([testFolder,'/',D_ir(1).name]);
[imM,imN] = size(im);

goods = [5,26,33];
for i = 1:length(goods)
    id = goods(i);
    figure(1);
    subplot(3,6,(i-1)*6+1), imshow(imageBlock{id,1},[]); 
    if i == 1
    title('Corrupted');
    end
    axis off; 
    subplot(3,6,(i-1)*6+2), imshow(reshape(imageBlock{id,4},imN,imM)',[]);
    if i == 1
    title('Goal');
    end
       axis off; 
    subplot(3,6,(i-1)*6+3), imshow(reshape(LrsRRImage{id},imN,imM)',[]); 
    if i == 1
    title('LRS-RR');
    end
       axis off; 
    subplot(3,6,(i-1)*6+4), imshow(reshape(torreImage{id},imN,imM)',[]); 
     if i == 1
    title('LR-RR');
    end
       axis off; 
    subplot(3,6,(i-1)*6+5), imshow(reshape(rpcaImage{id},imN,imM)',[]);
    if i == 1
    title('cRPCA');
    end
       axis off; 
    subplot(3,6,(i-1)*6+6), imshow(reshape(ransacImage{id},imN,imM)',[]); 
    if i == 1
    title('RANSAC');
    end
       axis off; 
    if id == 1
        %tilefigs;
    end
    %pause();
end