% looks at overlap between nnUNet output and a hippocampal mask generated
% via affine to CITI168. Note most of the time these will already be affine
% registered but not when the --rigOnly flag was used (ie. in TSE data)

%% ASHSatlas (7T-TSE)
addpath(genpath([getenv('AUTOTOP_DIR') '/tools']));
ref_mask = [getenv('AUTOTOP_DIR') '/hippunfold/hippocampal_autotop/atlases/CITI/Mask_hemi-R_300umCoronalOblique.nii.gz'];
nnUNet_hippo_labels = [1 2 7 8]; % GM SRLM cysts DG
mask = load_untouch_nii(ref_mask);

workdirs = {'datasets/ASHSatlas_7T-TSE/HippUnfold/work/' ...
    'datasets/Maastricht_7T-TSE/HippUnfold/work/'...
    'datasets/HCP1200/work/' ...
    'datasets/HCPaging_T2w/work/' ...
    'datasets/HCPaging_T1w/work/'}; % note: need T1w reference
n=0;
for d = 1:length(workdirs)
    if strcmp(workdirs{d},'datasets/HCPaging_T1w/work/')
        ref_img = [getenv('AUTOTOP_DIR') '/hippunfold/hippocampal_autotop/atlases/CITI/T1w_300umCoronalOblique_hemi-R.nii.gz'];
        mod = 'T1w';
    else
        ref_img = [getenv('AUTOTOP_DIR') '/hippunfold/hippocampal_autotop/atlases/CITI/T2w_300umCoronalOblique_hemi-R.nii.gz'];
        mod = 'T2w';
    end
    subs = ls(workdirs{d});
    subs = strsplit(subs)';
    subs(end) = [];
    nsubs(d) = length(subs);
    
    for s = 1:length(subs)
        hemi = {'R' 'Lflip'};
        for LR = 1:2
            % generate f3d registration to CITI168
            cropped_img = [workdirs{d} subs{s} '/' subs{s} '_hemi-' hemi{LR} '_space-corobl_desc-cropped_modality-' mod '_autotop/img.nii.gz'];
            cropped_nnUNet = [workdirs{d} subs{s} '/' subs{s} '_hemi-' hemi{LR} '_space-corobl_desc-cropped_modality-' mod '_autotop/manual_lbl.nii.gz'];
            cpp = [workdirs{d} subs{s} '/' subs{s} '_hemi-' hemi{LR} '_space-corobl_desc-cropped_modality-' mod '_autotop/reg_to_CITI.nii.gz'];
            res = [workdirs{d} subs{s} '/' subs{s} '_hemi-' hemi{LR} '_space-corobl_desc-cropped_modality-' mod '_autotop/nnUNet_space-CITI.nii.gz'];
            if ~exist(res,'file')
            system(['reg_f3d -flo ' cropped_img ' -ref ' ref_img ' -cpp ' cpp]);
            system(['reg_resample -flo ' cropped_nnUNet ' -ref ' ref_img ' -cpp ' cpp ' -res ' res ' -inter 0']);
            end
            
            % check mask overlap
            n=n+1;
            try
            nnUNet = load_untouch_nii(res);
            f3dMask_Dice(n) = dice(mask.img==1, ismember(nnUNet.img,nnUNet_hippo_labels));
            end
        end
    end
end
save('verify_integrity_generalizationData.mat');

%% plot histogram

nsubscum = cumsum(nsubs)*2;
edges = [0.5:0.02:0.9];

tmp = f3dMask_Dice;
tmp(tmp<0.5) = 0.5; % xaxis window

figure;
hold on;
yyaxis left;
set(gca,'ycolor','k');
histogram(tmp(nsubscum(2)+1:nsubscum(3)),edges,'facecolor','w','facealpha',0.5,'edgecolor', 'k');
histogram(tmp(nsubscum(3)+1:nsubscum(4)),edges-0.005,'facecolor','g','facealpha',0.5,'edgecolor', 'none');
histogram(tmp(nsubscum(4)+1:nsubscum(5)),edges-0.01,'facecolor','b','facealpha',0.5,'edgecolor', 'none');
yyaxis right;
set(gca,'ycolor','r');
ylim([0 60]);
histogram(tmp(1:nsubscum(2)),edges-0.015,'facecolor','r','facealpha',0.5,'edgecolor', 'none');
legend({'HCP1200-T2w' 'HCPaging-T2w' 'HCPaging-T1w' '7T-TSE'},'location','northwest');
saveas(gcf,'f3d_comparison.png');

%% questionable subs

f = fopen('badsubs.txt','w');
n=0;
for d = 1:length(workdirs)
    subs = ls(workdirs{d});
    subs = strsplit(subs)';
    subs(end) = [];
    for s = 1:length(subs)
        for LR = 1:2
            n=n+1;
            if f3dMask_Dice(n) < 0.7
                fprintf(f,['Check file ' workdirs{d} subs{s} '_hemi' hemi{LR}...
                    ' Dice overlap with F3D CITI168 atlas was: ' num2str(f3dMask_Dice(n)) '\n']);
            end
        end
    end
end
fclose(f);
