function fname = do_prop(PT,PY, tdir)
% Label Propagation
% FORMAT fname = do_prop(PT,PY, tdir)
%     PT    - names of NIfTI file encoding the grey and white matter
%             tissue maps for the individual (char array)
%     PY    - name of NIfTI file encoding the deformation field
%     tdir  - directory containing template data, label probabilities etc
%
%     fname - filename of maximum probability label image generated by the
%             procedure.
%
% Structure volumes are saved as an ascii file, named volumes_*.txt

%   This program is free software: you can redistribute it and/or modify
%   it under the terms of the GNU General Public License as published by
%   the Free Software Foundation, either version 3 of the License, or
%   (at your option) any later version.
%
%   This program is distributed in the hope that it will be useful,
%   but WITHOUT ANY WARRANTY; without even the implied warranty of
%   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%   GNU General Public License for more details.
%
%   You should have received a copy of the GNU General Public License
%   along with this program.  If not, see <http://www.gnu.org/licenses/>.


if nargin<1, PT = spm_select(2,'^c.*\.nii','Select the GM and WM maps'); end;
if nargin<2, PY = spm_select(1,'y_.*\.nii', 'Select the deformation field'); end
if nargin<3, tdir = spm_select(1,'dir','Select the directory containing training data'); end

if size(PT,1)~=2, error('Wrong number of tissue classes.'); end;
if size(PY,1)~=1, error('Wrong number of deformations.'); end

% Load the tissue maps
Nic = nifti(PT);
T{numel(Nic)+1} = 1;
for c=1:numel(Nic),
    T{c} = Nic(c).dat(:,:,:);
    T{numel(Nic)+1} = T{numel(Nic)+1} - T{c};
end

% Compute the inverse of the deformation field, and save as file
defs.comp{1}.inv.comp{1}.def = {PY};
defs.comp{1}.inv.space = {PT(1,:)};
defs.out{1}.savedef.ofname = sprintf('%.6d', round(rand(1)*1000000));
defs.out{1}.savedef.savedir.saveusr = {tempdir};
out  = spm_deformations(defs);

% Transform the mm coordinates of the deformation into voxel indices
% and delete the file coding the deformation
Niy  = nifti(out.def{1});
Y0   = squeeze(single(Niy.dat(:,:,:,:,:)));
delete(out.def{1});
M    = inv(spm_get_space(fullfile(tdir,'label000_c1.nii')));
Y    = zeros(size(Y0),'single');
for d=1:3,
    Y(:,:,:,d) = Y0(:,:,:,1)*M(d,1) + Y0(:,:,:,2)*M(d,2) + Y0(:,:,:,3)*M(d,3) + M(d,4);
end

L = zeros(size(T{1}),'uint8');  % Map of most probable labels
P = zeros(size(T{1}),'single'); % Temporary file
maxlabel = 256;                 % Largest label index
volumes  = zeros(maxlabel,1);   % Volumes

for k=0:maxlabel % Loop over label numbers

    % Check that there is information for that label number
    fn  = fullfile(tdir,sprintf('label%.3d_c%d.nii', k,1));
    if exist(fn,'file'),
        P1   = 0;

        % Compute coordinates to sample from
        Nii  = nifti(fn);
        MM   = M*Nii.mat;
        MM   = MM(1:3,4);
        Y1   = Y;
        Y1(:,:,:,1) = Y(:,:,:,1) - MM(1);
        Y1(:,:,:,2) = Y(:,:,:,2) - MM(2);
        Y1(:,:,:,3) = Y(:,:,:,3) - MM(3);

        % Loop over tissues and compute label probabilities (p(c)) from probabilities
        % of label given tissue (precomputed p(c|t.)) and tissue probabilities (p(t.)): 
        %     p(c) = p(c|t1)*p(t1)+p(c|t2)*p(t2)+p(c|t3)*p(t3)
        for c=1:(numel(Nic)+1),
            fn   = fullfile(tdir,sprintf('label%.3d_c%d.nii', k,c));
            Nii  = nifti(fn);
            Pdat = single(Nii.dat(:,:,:));
            Pdat = spm_diffeo('bsplins',Pdat,Y1,[1 1 1 0 0 0]);
            P1   = P1 + Pdat.*T{c}; 
        end
        clear Y1

        % Decide what to do when no data is available for p(c|t.)
        if k==0
            P1(~isfinite(P1)) = 1; % Background
        else
            P1(~isfinite(P1)) = 0; % Foreground
        end

        % Volume in cc
        if k>0,
            tiss_vol   = sum(P1(:)).*abs(det(Nic(1).mat(1:3,1:3)))/1000;
            volumes(k) = tiss_vol;
            fprintf('%.3d %g\n', k, tiss_vol);
        end

        % Update maximum probability label map
        [P,tmp] = max([P(:),P1(:)],[],2);
        P   = reshape(P,size(T{1}));
        L(tmp==2) = k;

    end
end

% Write NIfTI file containing maximum probability labels
[pth,nam,ext]  = fileparts(deblank(PT(1,:)));
fname          = fullfile(pth,['label_' nam(3:end) '.nii']);
dat            = file_array(fname,size(T{1}),'uint8',0,1,0);
Ni1            = nifti;
Ni1.dat        = dat;
Ni1.mat        = Nic(1).mat;
create(Ni1);
Ni1.dat(:,:,:) = L;

% Save structure volumes for this subject
save(fullfile(pth,['volumes_' nam(3:end) '.txt']),'volumes','-ascii');


