%% Prep steps


% Add in the appropriate paths
main_path = '/home/ktsirlis/ABCD_preprocessing/';
spm_path  = '/home/ktsirlis/ABCD_preprocessing/spm12/';
data_path = '/SAN/medic/ABCD_project/T1MinProcessed/fmriresults01/';


% Fill in SPM's path. And make sure MB toolbox is inside SPM12's
% toolbox folder with the name "mb". Finally you will also need all the
% files inside the MB_github folder.
addpath(spm_path, [main_path 'MB_github']);


% Initialize SPM12 without interface
% This is needed as the MB algorithm is using SPM12 to run.
spm('defaults', 'fmri');
spm_jobman('initcfg');
spm_get_defaults('cmdline',true);


% Create subject list
% You also need to have a subject list array with all the subject paths so
% that the MB algo can run them. We are providing our subject list as an
% example so that you know how to build yours.




%% MB Algo

% Path to T1w MRI
pth_img = {subj_list.Scan}';

% MB RUN module (fits the model)
run              = struct;
run.mu.exist     = {[main_path 'Mikael' filesep 'mu_T1.nii']};            % Template file
run.gmm.pr.file  = {[main_path 'Mikael' filesep 'prior_T1.mat']};         % Intensity prior file
run.odir         = {['/SAN/medic/ABCD_project/ABCD_exp1/' 'output_' num2str(output_folder_idx) '_vx_2_0']};  % Output directory (to write model files)
run.onam         = num2str(mb_fit_idx);                   % Sets post-fix of output files
if use_parfor
    run.nworker      = size(pth_img,1);
end
run.gmm(1).chan(1).images = pth_img;
run.gmm(1).chan(1).inu.inu_reg = 1e5;

% MB OUT module (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     = 2;         % template space voxel size
out1.bb      = [-80 -116 -62; 80 80 98];  % template space bounding box
out1.fwhm    = [0, 0, 0];              % template space smoothing (FWHM of Gaussian kernel (in mm))
out1.proc_zn = {@(x) clean_gwc(x)};

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




%% Make a Mask - ABCD specific

% First build the GM_2_0 cell with all the names of the nifti files.


for i = 1:length(GM_2_0)  
    i
    V1 = spm_vol(GM_2_0{i});    
    if i == 1
        GMx_2_0 = spm_read_vols(V1);
    else
        GMx_2_0 = GMx_2_0 + spm_read_vols(V1);  
    end    
end
GMx_2_0 = GMx_2_0/length(GM_2_0);
A_GM_2_0 = reshape(GMx_2_0,1,[]);
GM_2_0_mask = A_GM_2_0;
save('GM_2_0_mask.mat','GM_2_0_mask');




%% Compute TBVs/TIVs - ABCD Specific

% Using the native space images (out.c[1,2,3]) you can compute the TBV/TIV,
% if you want to use them as confounding variables as we did.

% In this first part we are building a path list c0_all_2_0 with the files of the 3
% different segmented tissue classes we used to compute the TBV/TIV.

name_idx_1 = 1;
name_idx_2 = 1;
name_idx_3 = 1;

for i = 1:10
    i
    output_var = ['output_' num2str(i) '_vx_2_0'];
    output_var = eval(output_var);        
    for j = 1:length(output_var)
        if startsWith(output_var(j,1).name,'c01')
            c01_2_0{name_idx_1,1} = [output_var(j,1).folder '/' output_var(j,1).name];
            name_idx_1 = name_idx_1 + 1;
        end
    end
    for j = 1:length(output_var)
        if startsWith(output_var(j,1).name,'c02')
            c02_2_0{name_idx_2,1} = [output_var(j,1).folder '/' output_var(j,1).name];
            name_idx_2 = name_idx_2 + 1;
        end
    end
    for j = 1:length(output_var)
        if startsWith(output_var(j,1).name,'c03')
            c03_2_0{name_idx_3,1} = [output_var(j,1).folder '/' output_var(j,1).name];
            name_idx_3 = name_idx_3 + 1;
        end
    end
end

c0_all_2_0(:,1) = c01_2_0;
c0_all_2_0(:,2) = c02_2_0;
c0_all_2_0(:,3) = c03_2_0;


parfor i = 1:length(c0_all_2_0)
    i
Nii = nifti(c0_all_2_0(i,:));
vx = sqrt(sum(Nii(1).mat(1:3,1:3).^2));
tbv_dummy = 0;
tiv_dummy = 0;
for k=1:3
    vk = sum(Nii(k).dat(:) > 0.5);
    if k < 3
        tbv_dummy = tbv_dummy + vk;
    end
    tiv_dummy = tiv_dummy + vk;
end
tbv_dummy = prod(vx)*tbv_dummy;
tiv_dummy = prod(vx)*tiv_dummy;

TBV(i,1) = tbv_dummy;
TIV(i,1) = tiv_dummy;

end

save('TBV.mat','TBV');
save('TIV.mat','TIV');




%% Smooth images - ABCD

fwhm = [10 10 10];

% As in the previous step the first thing you need to do is create the
% GM_2_0 path lists with all the subject paths. In this case you need the
% GM-segmented unsmoothed warped image file paths ("wc01") from all your
% subjects.

name_idx = 1;
for i = 1:10    
    output_var = ['output_' num2str(i) '_vx_2_0'];
    output_var = eval(output_var);        
    for j = 1:length(output_var)
        if startsWith(output_var(j,1).name,'wc01')
            GM_2_0{name_idx,1} = [output_var(j,1).folder '/' output_var(j,1).name];
            name_idx = name_idx + 1;
        end
    end
end

for i = 1:length(GM_2_0)
    new_name_2_0 = GM_2_0{i,1};
    new_name_2_0 = strrep(new_name_2_0,'vx_2_0/wc01','vx_2_0/swc01');
    spm_smooth(GM_2_0{i,1},new_name_2_0,fwhm./2);    
end




%% Mask swc images and create the X image vector

%%%%%%%% 2.0 voxel size

load('GM_2_0_mask.mat')
GM_2_0_mask = GM_2_0_mask>0.1;

% Same as before we start by creating the swc_GM_2_0 path lists with all the
% now GM-segmented smoothed warped image file paths ("swc01") from all your
% subjects.

name_idx = 1;
for i = 1:10    
    output_var = ['output_' num2str(i) '_vx_2_0'];
    output_var = eval(output_var);        
    for j = 1:length(output_var)
        if startsWith(output_var(j,1).name,'swc01')
            swc_GM_2_0{name_idx,1} = [output_var(j,1).folder '/' output_var(j,1).name];
            name_idx = name_idx + 1;
        end
    end
end


% Due to memory issues we had to split the X matrix in 3 parts while saving
% it, but feel free to build it in 1 part if your RAM can support this.


for i = 1:length(swc_GM_2_0)/3
    i
    V = spm_vol(swc_GM_2_0{i});
    V = spm_read_vols(V);
    V = reshape(V,1,[]);
    V = V(GM_2_0_mask);    
    X_2_0_1(i,:) = V;       
end

save('imageVars_2_0_1.mat','X_2_0_1','-v7.3');
clear X_2_0_1;



for i = length(swc_GM_2_0)/3+1:length(swc_GM_2_0)*2/3
    i
    V = spm_vol(swc_GM_2_0{i});
    V = spm_read_vols(V);
    V = reshape(V,1,[]);
    V = V(GM_2_0_mask);
    X_2_0_2(i-length(swc_GM_2_0)/3,:) = V;
end

save('imageVars_2_0_2.mat','X_2_0_2','-v7.3');
clear X_2_0_2;



for i = length(swc_GM_2_0)*2/3+1:length(swc_GM_2_0)
    i
    V = spm_vol(swc_GM_2_0{i});
    V = spm_read_vols(V);
    V = reshape(V,1,[]);
    V = V(GM_2_0_mask);
    X_2_0_3(i-length(swc_GM_2_0)*2/3,:) = V;       
end

save('imageVars_2_0_3.mat','X_2_0_3','-v7.3');
clear X_2_0_3;