%% Path setting
% figDir = 
% figName


%% Figures for Paper

addpath('GridSphere/GridSphere');
load('realdata0.mat');
grid = bvecss{1,1};
rad  = 2;
centers  = sphlattice(10,rad);
xs       = ste_tan0(centers,grid);
biggrid  = biggridsphere(1000);
thetas   = (1:100)./100 .* 2*pi;
circgrid = [sin(thetas);cos(thetas);0.*thetas];
pp=size(centers,2);
% utility functions
applyToGivenRow = @(func, matrix) @(row) func(matrix(row, :));
newApplyToRows = @(func, matrix) arrayfun(applyToGivenRow(func, matrix), 1:size(matrix,1), 'UniformOutput', false)';
takeAll = @(x) reshape([x{:}], size(x{1},2), size(x,1))';
genericApplyToRows = @(func, matrix) takeAll(newApplyToRows(func, matrix)); % from http://stackoverflow.com/users/376454

%% Generate the data 
rng(0);
% The number of fibers
true_k= 3;

% Partial volume fractions for the fibers
true_amp  = [.4;.3;.3];
tb        = sum(true_amp);
true_tamp = sum(true_amp);
% We simulate three fibers with slightly different directions
%In this parameterization we combine directionality of the fiber with
% its axial diffusivity
%
% If v is the parameter of the fiber, sqrt(norm(v))=b(lambda1-lambda2)
true_pos=[0.4657    1.2493    0.4667;
          -1.6289    0.6602   -1.1222;
          -0.8092    0.9064    1.2146];
%true_pos = [[1.52;-0.153;0.1],[.8;1.2;0.111],[.8;0.23;1.3]];
%true_pos = [ 0.4657    1.2493    0.4667];

%s Noise.
sigma=0.005;
noise=sigma;

%% Simulate data. 
[data_orig,data_noiseless] = generate_data(true_pos,true_amp,grid,noise);
y = data_orig; % signal with noise
y0 = data_noiseless; % noiseless signal


%% Fit data using  EBP
% (We will plot the steps of this algorithm in Figure no 1.)
xfinder = @(r) xfinder_random2(grid,2,100,3,r); % the oracle, a function for finding the kernel most coorrelated with the residual
r = data_orig; % the residual in the 0th ieration is the data
[f1,dir1] = xfinder(r); %the first kernel and parameter returned by the oracle
betas1 = lsqnonneg(f1,data_orig);
% Compute the residual signal after removing the first Oracle function.
r   = data_orig - f1 * betas1;
[f2,dir2,corr] = xfinder(r);
intermediateBetas = lsqnonneg(f2, r);
fullModel         = [f1, f2];
betasFinal        = lsqnonneg(fullModel,data_orig);

%% Simulate diffusion signal
% Here after we first simulate data with certain bvecs (circgrid, a great
% circle, equator for this example). To simulate we use a tensor with a
% specific direction (true_pos) amplitude (true_amp), the tensor has a
% shape defined by l1 (roughly axial diffusivity), l2 (roughly radial
% diffusivity). Signma is the level of noise added tot the signal.
cy  = generate_data(true_pos, true_amp, circgrid,sigma);
cx1 = generate_data(dir1, 1, circgrid, 0);
cx2 = generate_data(dir2, 1, circgrid, 0);

%% Plot the steps of the algorithm.
figure('name','Steps of EBP algorithm','color','w');
subplot(1,4,1);
polar([thetas,thetas(1)],[cy;cy(1)]','k-');

subplot(1,4,2);
polar([thetas,thetas(1)],[cy;cy(1)]','k-');
hold on;
polar([thetas,thetas(1)],betas1.*[cx1;cx1(1)]','b-');

% subplot(1,4,3);
% polar([thetas,thetas(1)],[cy;cy(1)]','k-');
% %polar([thetas,thetas(1)],betas1.*[cx1;cx1(1)]','b-');
% polar([thetas,thetas(1)],betas1.*[cx1;cx1(1)]'+intermediateBetas.*[cx2;cx2(1)]','b-');

subplot(1,4,3)
residuals = [cy;cy(1)]' - betas1.*[cx1;cx1(1)]';
residuals(residuals < 0) = 0;
sigh = polar([thetas,thetas(1)],[cy;cy(1)]','k-');
hold on;
polar([thetas,thetas(1)],residuals,'k-')
polar([thetas,thetas(1)],intermediateBetas.*[cx2;cx2(1)]','b-');
delete(sigh)

subplot(1,4,4);
polar([thetas,thetas(1)],[cy;cy(1)]','k-');
hold on;
%polar([thetas,thetas(1)],betasFinal(1).*[cx1;cx1(1)]','b-');
%polar([thetas,thetas(1)],betasFinal(2).*[cx2;cx2(1)]','m-');
polar([thetas,thetas(1)],betasFinal(1).*[cx1;cx1(1)]' + betasFinal(2).*[cx2;cx2(1)]','b-');

%% Print the figure
% printCommand = ...
% sprintf('print(%s, ''-cmyk'', ''-painters'',''-depsc2'',''-tiff'',''-r500'' , ''-noui'', ''%s'')', ...
% gcf,fullfile(figDir,figName));
% eval(printCommand);


%% Fit DTI model
[est_d,~,estQ,estQr] = dti_fit(data_noiseless,grid);
[u,v]=eig(estQ);
[vv,o] = sort(diag(v));
uu=u(:,o);
v_dti = u(:,3) .* sqrt(vv(3)-vv(2));
if v_dti(1) < 0;
    v_dti = -v_dti;
end
w_dti = exp(-vv(2));

%% Fit NNLS model
l1p=0.01;
betaa_0 = lsqnonneg([xs;ones(1,pp)],[data_orig;1]);
est_0 = centers(:,betaa_0 > 0);
mags_0 = betaa_0(betaa_0 > 0);
spar_0 = sum(betaa_0 > 0);
xs_0 = xs(:,betaa_0 > 0);
em_0 = symearthmover2(true_pos,true_amp,est_0,mags_0,1e-3);
yh_0 = xs * betaa_0;
mse_0 = norm(y-yh_0)^2;

%% Fit EBP
xfinder = @(r) xfinder_random2(grid,2,100,3,r);
[mags_e, est_e, yh_e, sses_e, mses_e, ems_e, spars_e] = dwi_ebp(y, 1, zeros(3,1), ones(151,1), mean(y), xfinder, 1000, y0, true_pos, true_amp);
spar_e = size(est_e,2);
em_e = symearthmover2(true_pos,true_amp,est_e,mags_e,1e-3);

figure;
plot(mses_e);
title('EBP algorithm : prediction error');

figure;
plot(spars_e);
title('EBP algorithm : Size of active set');

figure;
plot(sses_e);
title('EBP algorithm : objective function');

figure;
plot(ems_e);
title('EBP algorithm : true earthmover distance');


%%
[X01,Y01,Z01] = cyl2(true_pos(:,1),true_amp(1));
[X02,Y02,Z02] = cyl2(true_pos(:,2),true_amp(2));
[X03,Y03,Z03] = cyl2(true_pos(:,3),true_amp(3));

[Xd,Yd,Zd] = cyl2(v_dti,w_dti);


%% plot params

fitshade = [0.3 0.3 0.3];
view1=[40 27];
xl = [-1.5 1.5];
yl = [-1.5 1.5];
zl = [0 1.5];

%% Plot DTI fit
fig = figure;

ax11 = surf(X01,Y01,Z01); set(ax11,'FaceColor',[1,1,1]);alpha(0.5); hold on;
ax12 = surf(X02,Y02,Z02); set(ax12,'FaceColor',[1,1,1]);alpha(0.5); hold on;
ax13 = surf(X03,Y03,Z03); set(ax13,'FaceColor',[1,1,1]);alpha(0.5); hold on;

axd = surf(Xd,Yd,Zd);set(axd,'FaceColor',fitshade);
xlim(xl);
ylim(yl);
zlim(zl);
view(view1(1),view1(2));
pbaspect([1 1 1]);
%%
% set(gcf,'PaperUnits','inches','PaperSize',[10,10],'PaperPosition',[0 0 8 8])
% print -painters -dpdf -r300 figure_3dDTI.pdf

%% Plot NNLS fit
fig = figure;

ax11 = surf(X01,Y01,Z01); set(ax11,'FaceColor',[1,1,1]);alpha(0.5); hold on;
ax12 = surf(X02,Y02,Z02); set(ax12,'FaceColor',[1,1,1]);alpha(0.5); hold on;
ax13 = surf(X03,Y03,Z03); set(ax13,'FaceColor',[1,1,1]);alpha(0.5); hold on;

for i=1:spar_0;
    if mags_0(i) > 0.01;
        v = est_0(:,i);
        if v(1) < 0;
            v = -v;
        end
        [X,Y,Z] = cyl2(v,mags_0(i));
        ax = surf(X,Y,Z); set(ax,'FaceColor',fitshade); hold on;
    end
end

xlim(xl);
ylim(yl);
zlim(zl);
view(view1(1),view1(2));
pbaspect([1 1 1]);

%%
% set(gcf,'PaperUnits','inches','PaperSize',[10,10],'PaperPosition',[0 0 8 8])
% print -painters -dpdf -r300 figure_3dNNLS.pdf

%% Plot EBP fit

fig = figure;

ax11 = surf(X01,Y01,Z01); set(ax11,'FaceColor',[1,1,1]); alpha(0.5);hold on;
ax12 = surf(X02,Y02,Z02); set(ax12,'FaceColor',[1,1,1]); alpha(0.5);hold on;
ax13 = surf(X03,Y03,Z03); set(ax13,'FaceColor',[1,1,1]); alpha(0.5);hold on;

for i=1:spar_e;
    if mags_e(i) > 0.01;
        v = est_e(:,i);
        if v(1) < 0;
            v = -v;
        end
        [X,Y,Z] = cyl2(v,mags_e(i));
        ax = surf(X,Y,Z); set(ax,'FaceColor',fitshade); hold on;
    end
end

xlim(xl);
ylim(yl);
zlim(zl);
view(view1(1),view1(2));
pbaspect([1 1 1]);

%%
% set(gcf,'PaperUnits','inches','PaperSize',[10,10],'PaperPosition',[0 0 8 8])
% print -painters -dpdf -r300 figure_3dTCBP.pdf

%% Real data
realdata6_var_ker
output(1,:)=1.5;
output(2,:)=0;

rmse_dti = reshape(output(:,3),[31 31]);
rmse_nnls = reshape(output(:,4),[31 31]);
rmse_tcbp = reshape(output(:,6),[31 31]);


figure; imagesc(flipud(rmse_dti));colormap('bone');
figure; imagesc(flipud(rmse_nnls));colormap('bone');
figure; imagesc(flipud(rmse_tcbp));colormap('bone');
