% For this to work, the following files must be stored in directory
% BASEDIR/DATADIR:
% AllLin.img - A map of the DT linearity.
% AllPlan.img - A map of the DT planarity.
% AllPD.Bdouble - A map of the DT principal direction.
% <ROINAME>.img - A binary image specifying an ROI.

% Base directory in which data files live.
BASEDIR='.';
% Name of directory containing files for current data set.
DATADIR='.';
% Name of the starting ROI before any thresholding.
ROINAME='WholeBrain';

% Information specific to this data set
XSIZE=128;
YSIZE=256;
ZSIZE=3;
COMPONENTS=372;
% Indices of the unweighted images
B0_INDICES=[1 2 3 94 95 96 187 188 189 280 281 282];


% Threshold on the DT linearity (keep voxels with higher)
linthresh = 0.6;
% Threshold on the DT linearity (keep voxels with lower)
planthresh = 0.2;
% Threshold on the b=0 intensity (keep voxels with higher)
b0_thresh = 2000;
% Threshold on the angle between principal directions (keep voxels with
% higher abs dot product with the central voxel in the smoothing kernel).
cos_thresh = 0.9;
% Basic size of smoothing kernel is (2window_size+1)^2.
window_size = 2;

% Load in all the required data files.
fid = fopen([BASEDIR, '/', DATADIR, '/', ROINAME, '.img'], 'r', 'b');
d = fread(fid, 'char');
fclose(fid);
mask = reshape(d, XSIZE, YSIZE, ZSIZE);
fid = fopen([BASEDIR, '/', DATADIR, '/AllLin.img'], 'r', 'b');
d = fread(fid, 'float');
fclose(fid);
lin = reshape(d, XSIZE, YSIZE, ZSIZE);
fid = fopen([BASEDIR, '/', DATADIR, '/AllPlan.img'], 'r', 'b');
d = fread(fid, 'float');
fclose(fid);
plan = reshape(d, XSIZE, YSIZE, ZSIZE);
fid = fopen([BASEDIR, '/', DATADIR, '/AllPD.Bdouble'], 'r', 'b');
d = fread(fid, 'double');
fclose(fid);
pds = reshape(d, 3, XSIZE, YSIZE, ZSIZE);

% Read in the full data file.
fid = fopen([BASEDIR, '/', DATADIR, '/All.Bfloat'], 'r', 'b');
% Read it in volume by volume to avoid matlab problems with
% truncation of large files.
for i=1:ZSIZE
  d = fread(fid, COMPONENTS*XSIZE*YSIZE, 'float');
  dw(:,:,:,i) = reshape(d, COMPONENTS,XSIZE, YSIZE);
end
fclose(fid);

% Remove low linearity and high planarity
mask = mask & (lin>linthresh) & (plan<planthresh);

% Also threshold on the b=0 measurement to avoid smoothing with regions of
% large b=0 intensity likely to contain significant CSF proportion.
mb0 = squeeze(mean(dw(B0_INDICES,:,:,:)));
mask = mask & (mb0<b0_thresh);

% Now do the smoothing step.
% To store the output
dw_smth = zeros(size(dw));
% Size of smoothing kernel
wid = window_size;
for i=(1+wid):(XSIZE-wid)
    for j=(1+wid):(YSIZE-wid)
        for k=1:ZSIZE
            if(mask(i,j,k)>0)
                
                % Construct a neighbourhood over which to average
                mnbhd = mask((i-wid):(i+wid), (j-wid):(j+wid), k);
                
                % Compare the principal direction in each potential
                % neighbourhood element with that of the central voxel.
                pdc = squeeze(pds(:,i,j,k));
                nbhdpds = pds(:,(i-wid):(i+wid), (j-wid):(j+wid), k);
                pdcrep = repmat(pdc, [1, 2*wid+1, 2*wid+1]);
                dps = squeeze(abs(sum(nbhdpds.*pdcrep)));
                
                % Keep only voxels with similar direction
                mnbhd = mnbhd.*(dps>cos_thresh);
                
                % Do the averaging.
                for c=1:COMPONENTS
                    nbhd = squeeze(dw(c,(i-wid):(i+wid), (j-wid):(j+wid), k));
                    toav = nbhd(find(mnbhd>0));
                    dw_smth(c,i,j,k) = mean(toav);
                end
            end
        end
    end
end

% Save the smoothed data.
FULLROINAME=sprintf('%sL%02iP%02iB%05i_D%02i_SmthW%i', ROINAME, fix(linthresh*10), fix(planthresh*10), fix(b0_thresh), fix(cos_thresh*100), window_size);
fid=fopen([BASEDIR, '/', DATADIR, '/', FULLROINAME, '_All.Bfloat'], 'w', 'b');
fwrite(fid, dw_smth, 'float');
fclose(fid);
% Save the mask
fid=fopen([BASEDIR, '/', DATADIR, '/', FULLROINAME, '_Mask.img'], 'w', 'b');
fwrite(fid, mask, 'short');
fclose(fid);