function demo_mb
% Requires SPM12 and the diffeo-segment toolbox on MATLAB's path:
% * https://www.fil.ion.ucl.ac.uk/spm/software/download/
% * https://github.com/WTCN-computational-anatomy-group/diffeo-segment
%__________________________________________________________________________

% Get T1w scans
dir_data = './data';
pth_img  = spm_select('FPList',dir_data,'^.*\.nii$'); 
% pth_img  = pth_img(1,:); % uncomment for testing on a single subject

% MB RUN module (fits  the model)
run              = struct;
run.mu.exist     = {'./model/mu_T1.nii'};     % Template file
run.gmm.pr.file  = {'./model/prior_T1.mat'};  % Intensity prior file
run.odir         = {'./output-vx1'};          % Output directory (to write model files)
run.onam         = 'test1';                   % Sets post-fix of output files
run.nworker      = size(pth_img,1);
run.gmm(1).chan(1).images      = cellstr(pth_img); 
run.gmm(1).chan(1).inu.inu_reg = 1e5;

% MB OUT modules (writes different outputs)
out1         = struct;
out1.result  = {fullfile(run.odir{1}, ['mb_fit_' run.onam '.mat'])};
out1.c       = 1:3;                       % write classes in native space
out1.wc      = 1:3;                       % write classes in template space
out1.mwc     = 1:3;                       % write classes in modulated template space
out1.sm      = 1:3;                       % write scalar momentum
out1.vox     = 1.0;                       % template space voxel size
out1.bb      = [-80 -116 -62; 80 80 98];  % template space bounding box
out1.fwhm    = [10, 10, 10];              % template space smoothing (FWHM of Gaussian kernel (in mm))
out1.proc_zn = {@(x) clean_gwc(x)};
% and also for voxel size 2
out2      = out1;
out2.c    = [];
out2.vox  = 2.0;    
out2.odir = {'./output-vx2'};

% Run modules
jobs = {};
jobs{1}.spm.tools.mb.run = run;
jobs{2}.spm.tools.mb.out = out1;
jobs{3}.spm.tools.mb.out = out2;
spm_jobman('run', jobs);
%==========================================================================

%==========================================================================
function zn = clean_gwc(zn,ixt,level)
if nargin < 2 || isempty(ixt)
    % Default SPM12 template ordering
    ixt = struct('gm',1,'wm',2,'csf',3);
end
if nargin < 3, level = 1; end

b = sum(zn(:,:,:,ixt.wm),4);

% Build a 3x3x3 seperable smoothing kernel
kx=[0.75 1 0.75];
ky=[0.75 1 0.75];
kz=[0.75 1 0.75];
sm=sum(kron(kron(kz,ky),kx))^(1/3);
kx=kx/sm; ky=ky/sm; kz=kz/sm;

% Erosions and conditional dilations
th1 = 0.15;
if level==2, th1 = 0.2; end
niter  = 32;
niter2 = 32;
for j=1:niter
    if j>2
        th       = th1;
    else
        th       = 0.6;
    end  % Dilate after two its of erosion
    for i=1:size(b,3)
        gp       = double(sum(zn(:,:,i,ixt.gm),4));
        wp       = double(sum(zn(:,:,i,ixt.wm),4));
        bp       = double(b(:,:,i));
        bp       = (bp>th).*(wp+gp);
        b(:,:,i) = bp;
    end
    spm_conv_vol(b,b,kx,ky,kz,-[1 1 1]);
end

% Also clean up the CSF.
if niter2 > 0
    c = b;
    for j=1:niter2
        for i=1:size(b,3)
            gp       = double(sum(zn(:,:,i,ixt.gm),4));
            wp       = double(sum(zn(:,:,i,ixt.wm),4));
            cp       = double(sum(zn(:,:,i,ixt.csf),4));
            bp       = double(c(:,:,i));
            bp       = (bp>th).*(wp+gp+cp);
            c(:,:,i) = bp;
        end
        spm_conv_vol(c,c,kx,ky,kz,-[1 1 1]);
    end
end

th = 0.05;
for i=1:size(b,3)
    slices = cell(1,size(zn,4));
    for k1=1:size(zn,4)
        slices{k1} = double(zn(:,:,i,k1));
    end
    bp           = double(b(:,:,i));
    bp           = ((bp>th).*(sum(cat(3,slices{ixt.gm}),3)+sum(cat(3,slices{ixt.wm}),3)))>th;
    for i1=1:numel(ixt.gm)
        slices{ixt.gm(i1)} = slices{ixt.gm(i1)}.*bp;
    end
    for i1=1:numel(ixt.wm)
        slices{ixt.wm(i1)} = slices{ixt.wm(i1)}.*bp;
    end

    if niter2>0
        cp           = double(c(:,:,i));
        cp           = ((cp>th).*(sum(cat(3,slices{ixt.gm}),3)+sum(cat(3,slices{ixt.wm}),3)+sum(cat(3,slices{ixt.csf}),3)))>th;

        for i1=1:numel(ixt.csf)
            slices{ixt.csf(i1)} = slices{ixt.csf(i1)}.*cp;
        end
    end
    tot       = zeros(size(bp))+eps;
    for k1=1:size(zn,4)
        tot   = tot + slices{k1};
    end
    for k1=1:size(zn,4)
        zn(:,:,i,k1) = slices{k1}./tot;
    end
end
%==========================================================================