% Developed by Dr. Sonia Dahdouh
% Modified by Dr. Yuan-Chiao Lu
% Last Modified Date: 1/26/2022
clear
clc
close all

% Goal: Relabel the segmentation and parcellation
% This code use BrainSuite21a. If you use other versions of BrainSuite, please change the code in the funtion "reconstructAndLabelCorticalSurface"

%% Preinstalled software
% Toolboxes to download and install on your own home directory:
% (1) nifti (2) iso2mesh (3) peyre's toolbox (4) gptoolbox (5) meshlab (6) Brainsuite

%% Preparation
% Create a folder called 'Data_Step1' and put the raw scans in this folder
% There are 2 input files for each scan.
% (1) brain segmentation file: *_brain_tissue_labels_*.nii.gz
% (2) brain parcellation file: *_brain_labels_*.nii.gz

%% Please change the file path
main_dir = '/data/mril/users/all/mrdata/research/processed/CNMC/chd_r01/fetal/wip/yuan/Final_Work/Brain/Cortical_features';
data_path=fullfile(main_dir,'Attempt3/Data_Step1');

%% Main program starts here 
matlab_dir = fullfile(main_dir,'matlab_code');
app_dir = '/Applications';
cd(main_dir)
addpath(genpath(fullfile(main_dir,'matlab_code')))
addpath(genpath(fullfile(main_dir,'matlab_toolbox')))
addpath(genpath(fullfile(main_dir,'matlab_toolbox','nifti')))
addpath(genpath(fullfile(main_dir,'matlab_toolbox','iso2mesh')))
addpath(genpath(fullfile(main_dir,'matlab_toolbox','peyre')))
addpath(genpath(fullfile(main_dir,'matlab_toolbox','toolbox_brainsuite','utilities')));

cprintf('_blue','Step 1/3: Re-label Parcellation\n');

WM_seg_label = 3;
DGM_seg_label = 7;
Ventricles_seg_label=5; 
CorpusCallosum_label=48; % Corpus Callosum
HippocampusL_label=1; % Hippocampus Left
HippocampusR_label=2; % Hippocampus Right
AmygdalaL_label=3; % Amygdala Left
AmygdalaR_label=4; % Amygdala Right

Sbj=dir(fullfile(data_path,'fetus*'));

for Sbj_count=1:size(Sbj,1)

    Sbj_path=fullfile(data_path,Sbj(Sbj_count).name);
    
    nm = strfind(Sbj(Sbj_count).name,'_');
    Sbj_ID=Sbj(Sbj_count).name(nm+1:end-1);
    Scan_ID=Sbj(Sbj_count).name(end);
    cprintf('_blue',['##### Subject ID: ' num2str(Sbj_ID) ', and Scan: ' num2str(Scan_ID) ' #####\n']);
    cd(Sbj_path)
    
    % Check whether the parcellation and segmentation files are unzipped
    % Segmentation
    if ~exist(fullfile(Sbj_path,'segmentation',[Sbj_ID '-0' Scan_ID '_seg.nii']),'file')
        cmd=['find ',Sbj_path,' -name ','*_tissue_labels_*.nii.gz'];
        [~,listing]= system(cmd);
        if ~isempty(listing)
            startIndex=regexp(listing,'\n');
            startIndex=[1 startIndex];
            for ii=1:(size(startIndex,2)-1)
                gz_filename = strtrim(listing(startIndex(ii):startIndex(ii+1)));
                [path_res,name,ext] = fileparts(gz_filename);
                if name(1)~='.' 
                    gunzip(strcat(fullfile(Sbj_path,[name '.gz'])))
                    delete(strcat(fullfile(Sbj_path,[name '.gz'])))
                end
            end
        end
        cmd=['find ',Sbj_path,' -name ','*_tissue_labels_*.nii'];
        [~,listing]= system(cmd);
        mkdir('segmentation')
        movefile(strcat(listing),fullfile(Sbj_path,'segmentation',[Sbj_ID '-0' Scan_ID '_seg.nii']))
    end
    % Parcellation
    if ~exist(fullfile(Sbj_path,'parcellation',[Sbj_ID '-0' Scan_ID '_labels.nii']),'file')
        cmd=['find ',Sbj_path,' -name ','*_labels_*.nii.gz'];
        [~,listing]= system(cmd);
        if ~isempty(listing)
            startIndex=regexp(listing,'\n');
            startIndex=[1 startIndex];
            for ii=1:(size(startIndex,2)-1)
                gz_filename = strtrim(listing(startIndex(ii):startIndex(ii+1)));
                [path_res,name,ext] = fileparts(gz_filename);
                if name(1)~='.' 
                    gunzip(strcat(fullfile(Sbj_path,[name '.gz'])))
                    delete(strcat(fullfile(Sbj_path,[name '.gz'])))
                end
            end
        end
        cmd=['find ',Sbj_path,' -name ','*_labels_*.nii']; 
        [~,listing]= system(cmd);
        mkdir('parcellation')
        movefile(strcat(listing),fullfile(Sbj_path,'parcellation',[Sbj_ID '-0' Scan_ID '_labels.nii']))
    end
    
    segmentationPath = fullfile(Sbj_path,'segmentation');
    parcellationPath = fullfile(Sbj_path,'parcellation');
    
    cmd=['find ',parcellationPath,' -name ','*_labels.nii'];
    [status,listing]= system(cmd);
    startIndex=regexp(listing,'\n');
    startIndex=[1 startIndex];

    for ii=1:(size(startIndex,2)-1)

        parc_filename = strtrim(listing(startIndex(ii):startIndex(ii+1)));
        [path_res,name,ext] = fileparts(parc_filename);

        if name(1)~='.'
            nm = strfind(name,'_');
            segmentationName =strtrim(fullfile(segmentationPath,[name(1:nm(1)),'seg.nii']));
            reduced_parcellationPath =strtrim(fullfile(parcellationPath,[name(1:nm(1)),'WM_SubCort_labels.nii']));
            increased_segmentationPath =strtrim(fullfile(segmentationPath,[name(1:nm(1)),'seg_HA.nii']));
            outputLabeFileName =strtrim(fullfile(parcellationPath,[name(1:nm(1)),'labels.csv']));
            meshFileName_1 =strtrim(fullfile(parcellationPath,[name(1:nm(1)),'WM_1.off']));
            meshFileName_2 =strtrim(fullfile(parcellationPath,[name(1:nm(1)),'WM_2.off']));

            if exist(segmentationName,'file')

                [parcel] = computeWM_Parcellation(segmentationName,parc_filename,WM_seg_label,reduced_parcellationPath); % (Yuan) This gives *_WM_SubCort_labels.nii

                %Relabel the parcellation
                old_lab_1=[37,39,22,29,13,11,31,7,5,15,9,27,25,20,34,32]; % "20" is "insula right"
                new_lab = [1,2,3,4,4,4,4,4,4,4,4,4,4,5,6,7];

                old_lab_2=[36,38,23,14,28,30,12,8,6,16,10,26,24,21,35,33]; % "21" is "insula left"
                new_lab_2 = new_lab+max(new_lab);

                parcel2=parcel;
                parcel2(parcel2>0)=20;

                for iii=1:max(size(old_lab_1))            
                    parcel2(parcel==old_lab_1(iii))=new_lab(iii);
                    parcel2(parcel==old_lab_2(iii))=new_lab_2(iii);           
                end

                %Separate hemisphere
                parcel2(round(size(parcel,1)/2),:,:)=0;

                parcel3=parcel2;
                parcel3(parcel2==20)=0;

                %Add subcortical and ventricles to obtain a spherical mesh
                segmentation_nii = load_untouch_nii(segmentationName);
                segm=segmentation_nii.img;
 
                parcellation_nii = load_untouch_nii(parc_filename);
                parc=parcellation_nii.img;

                rh=zeros(size(parcel3));
	            rh(1:round(size(rh,1)/2),:,:)=1; 
                parcel3(segm==DGM_seg_label & rh==1)=15; % DGM right
                parcel3(segm==Ventricles_seg_label & rh==1)=17; % Ventricle right
                parcel3((parc==HippocampusR_label | parc==AmygdalaR_label) & rh==1)=11; % Hippocampus and Amygdala right
                parcel3(parc==CorpusCallosum_label & rh==1)=19; % Corpus Callosum right
                
                lh=zeros(size(parcel3));  
                lh((round(size(lh,1)/2)+1):end,:,:)=1;
                parcel3(segm==DGM_seg_label & lh==1)=16; % DGM left
                parcel3(segm==Ventricles_seg_label & lh==1)=18; % Ventricle left
                parcel3((parc==HippocampusL_label | parc==AmygdalaL_label) & lh==1)=4; % Hippocampus and Amygdala left
                parcel3(parc==CorpusCallosum_label & lh==1)=20; % Corpus Callosum left

                parcellation_nii.img = parcel3;
                save_untouch_nii(parcellation_nii,reduced_parcellationPath);
                
                segm_new=zeros(size(segm));
                segm_new(segm==1 & rh==1)=8;  % CSF, right
                segm_new(segm==2 & rh==1)=9;  % CGM, right
                segm_new(segm==3 & rh==1)=11; % WM, right
                segm_new(segm==4 & rh==1)=3;  % Background, right
                segm_new(segm==5 & rh==1)=5;  % Ventricles, right
                segm_new(segm==6 & rh==1)=6;  % Cerebellum, right
                segm_new(segm==7 & rh==1)=7;  % DGM, right
                segm_new(segm==8 & rh==1)=15; % Brainstem, right
                segm_new(parc==HippocampusR_label)=17;
                segm_new(parc==AmygdalaR_label)=19;

                if max(max(max(segm)))==11
                    segm_new(segm==10 & rh==1)=8; % If segm==10 (Third ventricle) exists, combine it to CSF
                    segm_new(segm==11 & rh==1)=8; % If segm==11 (Fourth ventricle) exists, combine it to CSF
                elseif max(max(max(segm)))==14
                    segm_new(segm==13 & rh==1)=8; % If segm==13 (Third ventricle) exists, combine it to CSF
                    segm_new(segm==14 & rh==1)=8; % If segm==14 (Fourth ventricle) exists, combine it to CSF
                end
    
                segm_new(segm==1 & lh==1)=1;  % CSF, left
                segm_new(segm==2 & lh==1)=2;  % CGM, left
                segm_new(segm==3 & lh==1)=4;  % WM, left
                segm_new(segm==4 & lh==1)=10; % Background, left
                segm_new(segm==5 & lh==1)=12; % Ventricles, left          
                segm_new(segm==6 & lh==1)=13; % Cerebellum, left
                segm_new(segm==7 & lh==1)=14; % DGM, left
                segm_new(segm==8 & lh==1)=16; % Brainstem, left
                segm_new(parc==HippocampusL_label)=18;
                segm_new(parc==AmygdalaL_label)=20;
                
                if max(max(max(segm)))==11
                    segm_new(segm==10 & lh==1)=1; % If segm==10 (Third ventricle) exists, combine it to CSF
                    segm_new(segm==11 & lh==1)=1; % If segm==11 (Fourth ventricle) exists, combine it to CSF 
                elseif max(max(max(segm)))==14
                    segm_new(segm==13 & lh==1)=1; % If segm==13 (Third ventricle) exists, combine it to CSF
                    segm_new(segm==14 & lh==1)=1; % If segm==14 (Fourth ventricle) exists, combine it to CSF 
                end
                
                segmentation_nii.img = segm_new;
                save_untouch_nii(segmentation_nii,increased_segmentationPath);

                try
                    % Reconstruct using brainsuite and label
                    reconstructAndLabelCorticalSurface(reduced_parcellationPath,meshFileName_1,meshFileName_2,outputLabeFileName,matlab_dir,app_dir)
                catch ME
                    disp 'reconstructAndLabelCorticalSurface.m is not working!'
                    errors{ii}=name;
                end
            end
            clear parcel segm new_img parcel3 parcel2 ix iy iz segmentation_nii tmp_2 tmp_1;   
        end

    end
    clearvars -except main_dir matlab_dir app_dir data_path Sbj Sbj_count ...
        WM_seg_label DGM_seg_label Ventricles_seg_label CorpusCallosum_label ...
        HippocampusL_label HippocampusR_label AmygdalaL_label AmygdalaR_label

end



