%DKE2MNI_IMG: Convert tracts from DKE Format to MNI Image Space
%
%To run this script, you will need: 
%
%    spm8             (not 12 as some defaults have been changed) - add to matlab path with subdirectories
%    vistasoft-master (http://white.stanford.edu/newlm/index.php/Software#VISTASOFT) - add to matlab path with subdirectories
%    trk_read / trk_write (https://githum.com/johncolby/along-tract-stats) - add to matlab path
%
%The non-linear transformation routine was adapted from AFQ_SegmentFiberGroups.m,
%which uses spm's normalization functions
%
%The output trk file is in the mm image space of the template image used for
%normalization, with the bottom of the image volume at [0,0,0]. To find the
%coresponding voxel indices for the ith tract do: 
%
%   [hdr trk] = trk_read(path_to_FT_DKI_mni_img.trk); 
%   ind = ceil(bsxfun(@rdivide,trk(i).matrix,hdr.voxel_size));
%
%Author: Russell Glenn
%Medical University of South Carolina

%--------------------------------------------------------------------------
%USER INPUT
%--------------------------------------------------------------------------
    %define subj_dir directory with dke / dke_ft output files

    subj_dir = ''; 
    
    %Make sure these file names are correct...
    FT_S = load(fullfile(subj_dir,'FT_struct.mat')); %Load FT_Struct
    fn_trk = fullfile(subj_dir,'FT_DKI.trk'); %get name to FT_DKI.trk
    fn_fa = fullfile(subj_dir,'fa.nii'); 

    verify_trks = 0; %Plot tracts in subject and template image space to see if they align
    save_trks = 1;  %Write tracts in mni_img space
    
%--------------------------------------------------------------------------
%INITIALIZE SUBJECT THINGS
%--------------------------------------------------------------------------

    if ~isfield(FT_S,'vox_to_ras'); 
        FT_S.vox_to_ras = FT_S.hdr.mat([FT_S.permute_img 4],:)*...
        [diag(FT_S.invert_img) double(FT_S.invert_img==-1)'.*...
        FT_S.hdr.dim(FT_S.permute_img)'-0.5*[1;1;1];0 0 0 1];
    end
    
    %Source image to align
    alignImg = spm_read_vols(spm_vol(fn_fa)); 
    
    %Subject tracts in LPS (mm) space
    [hdr trk] = trk_read(fn_trk); 
        
%--------------------------------------------------------------------------
%INITIALIZE TEMPLATE THINGS
%--------------------------------------------------------------------------
%See AFQ_SegmentFiberGroups.m for documentation     
%NOTE: Another option is to use the T2 template and alignImg = b0; 

    spm_defaults; global defaults; params = defaults.normalise.estimate;
    tdir = fullfile(fileparts(which('mrDiffusion.m')), 'templates');
    template = fullfile(tdir,'MNI_JHU_FA.nii.gz'); 
    %     template = fullfile(tdir,'MNI_JHU_T2.nii.gz');

%--------------------------------------------------------------------------
%VERIFY TRKS IN SUBJECT IMAGE SPACE
%--------------------------------------------------------------------------
if verify_trks
    figure
    imagesc(alignImg(:,:,round(size(alignImg,3)/2)))
    hold on 
    colormap gray

    for i = 1:1000
        x = bsxfun(@rdivide,trk(i).matrix,hdr.voxel_size); 
        for j = find(FT_S.invert_img==-1)
            x(:,j)=size(alignImg,j)-x(:,j); 
        end
        plot3(x(:,2),x(:,1),x(:,3),'g')
    end
end

%--------------------------------------------------------------------------
%CREATE FG STRUCTURE for dtiXformFiberCoords
%--------------------------------------------------------------------------
%Change from TRK to Vista Lab Format

    fibers = cell(length(trk),1); 
    for i = 1:length(trk); 
        fi = FT_S.vox_to_ras*[bsxfun(@rdivide,double(trk(i).matrix),hdr.voxel_size)';ones(1,trk(i).nPoints)];  
        fibers{i} = fi(1:3,:); 
    end

    params = {'faThresh',FT_S.fa_threshold,'lengthThreshMm',[FT_S.trk_length 500],...
        'stepSizeMm',FT_S.step_size}; 

%     seeds = FT_S.vox_to_ras*[FT_S.SEED';ones(1,size(FT_S.SEED,1))];
    seeds = [0;0;0];
    
    fg.name = 'wholeBrain'; 
    fg.colorRgb = [20 90 200]; 
    fg.thickness = -0.5;
    fg.visible = 1; 
    fg.seeds = seeds(1:3,:)';   
    fg.seedRadius = 0; 
    fg.seedVoxelOffsets = 'random'; 
    fg.params = params; 
    fg.fibers = fibers;
    fg.query_id = -1; 

%--------------------------------------------------------------------------
%GET SOME TRANSOFMATIONS AND PUT THE TRACTS IN MNI SPACE
%--------------------------------------------------------------------------

    %Transformations
    A = FT_S.vox_to_ras; %LPS (img) to RAS (mm)
    B = [[diag(FT_S.invert_img),...
        double(FT_S.invert_img==-1)'.*FT_S.hdr.dim(FT_S.permute_img)'];...
        [0 0 0 1]]; %LPS (img) to RAS (img); 
    C = [[eye(3),0.5.*[1;1;1]];[0 0 0 1]]; %Shift origin
    
    %Shift > RAS (img) to LPS (img) > LPS (img) to RAS (mm); 
    xformToAcpc = A*B^-1*C^-1; 
    
    %SPM NORMALIZATION
    [sn, Vtemplate, invDef] = mrAnatComputeSpmSpatialNorm(alignImg, xformToAcpc, template, params);
    
    %Tracts in MNI Space
    fg_mni = dtiXformFiberCoords(fg, invDef);
    tempImg = readFileNifti(template);
    
    %Convert from MNI to template image space
    clear fg_mni_img; 
    for fgi = 1:length(fg_mni);
        fx = fg_mni(fgi).fibers;
        A = tempImg.qto_xyz; 

        for i = 1:length(fx)
            x = fx{i}; 
            x = (A^-1*[x;ones(1,size(x,2))]);
            fx{i} = x(1:3,:);
        end

        fg_mni_img(fgi).fibers = fx; 
    end       

    for i = 1:length(fg_mni_img.fibers)
        trk_mni_img(i).matrix = fg_mni_img.fibers{i}'; 
        trk_mni_img(i).nPoints = size(fg_mni_img.fibers{i},2); 
    end
    
%--------------------------------------------------------------------------
%VERIFY TRKS IN TEMPLATE IMAGE SPACE
%--------------------------------------------------------------------------
if verify_trks
    figure
    imagesc(tempImg.data(:,:,round(tempImg.dim(3)/2)))
    hold on 
    colormap gray

    for i = 1:1000
        x = bsxfun(@rdivide,trk_mni_img(i).matrix,tempImg.pixdim); 
        plot3(x(:,2),x(:,1),x(:,3),'g')
    end
    
    figure
    cor = flipdim(permute(tempImg.data,[3 1 2]),1);
    imagesc(cor(:,:,round(size(cor,3)/2)))
    colormap gray
    hold on 
    for i = 1:1000
        x = bsxfun(@rdivide,trk_mni_img(i).matrix,tempImg.pixdim); 
        x(:,3) = size(cor,1)-x(:,3); 
        plot3(x(:,1),x(:,3),x(:,2),'g')
    end

    figure
    sag = flipdim(permute(tempImg.data,[3 2 1]),1);
    imagesc(sag(:,:,round(size(sag,3)/2)))
    colormap gray
    hold on 
    for i = 1:1000
        x = bsxfun(@rdivide,trk_mni_img(i).matrix,tempImg.pixdim); 
        x(:,3) = size(cor,1)-x(:,3); 
        plot3(x(:,2),x(:,3),x(:,1),'g')
    end
end

%--------------------------------------------------------------------------
%Save TRK_MNI_IMG
%--------------------------------------------------------------------------
if save_trks
    [d f e] = fileparts(fn_trk); 
    hdr.dim = tempImg.dim;
    hdr.voxel_size = tempImg.pixdim;
    hdr.image_orientation_patient = [1 0 0 0 -1 0]; %Temp image is probably in LAS
    trk_write(hdr,trk_mni_img,fullfile(d,[f '_mni_img' e]))
end

clear FT_S alignImg hdr trks fibers fg invDef fg_mni tempImg fg_mni_img