function ds = spm_uw_estimate(P,par)
% Estimation of partial derivatives of EPI deformation fields
%
% FORMAT [ds] = spm_uw_estimate((P),(par))
%
% P            - List of file names or headers.
% par          - Structure containing parameters governing the specifics
%                of how to estimate the fields.
% .M           - When performing multi-session realignment and Unwarp we
%                want to realign everything to the space of the first
%                image in the first time-series. M defines the space of
%                that.
% .order       - Number of basis functions to use for each dimension.
%                If the third dimension is left out, the order for 
%                that dimension is calculated to yield a roughly
%                equal spatial cut-off in all directions.
%                Default: [12 12 *]
% .sfP         - Static field supplied by the user. It should be a 
%                filename or handle to a voxel-displacement map in
%                the same space as the first EPI image of the time-
%                series. If using the FieldMap toolbox, realignment
%                should (if necessary) have been performed as part of
%                the process of creating the VDM. Note also that the
%                VDM must be in undistorted space, i.e. if it is
%                calculated from an EPI based field-map sequence
%                it should have been inverted before passing it to
%                spm_uw_estimate. Again, the FieldMap toolbox will
%                do this for you.
% .regorder    - Regularisation of derivative fields is based on the
%                regorder'th (spatial) derivative of the field.
%                Default: 1
% .lambda      - Fudge factor used to decide relative weights of
%                data and regularisation.
%                Default: 1e5
% .fot         - List of indexes for first order terms to model
%                derivatives for. Order of parameters as defined
%                by spm_imatrix. 
%                Default: [4 5]
% .sot         - List of second order terms to model second 
%                derivatives of. Should be an nx2 matrix where
%                e.g. [4 4; 4 5; 5 5] means that second partial
%                derivatives of rotation around x- and y-axis
%                should be modelled.
%                Default: []
% .fwhm        - FWHM (mm) of smoothing filter applied to images prior
%                to estimation of deformation fields.
%                Default: 6
% .rem         - Re-Estimation of Movement parameters. Set to unity means
%                that movement-parameters should be re-estimated at each
%                iteration.
%                Default: 0
% .noi         - Maximum number of Iterations.
%                Default: 5
% .exp_round   - Point in position space to do Taylor expansion around.
%                'First', 'Last' or 'Average'.
%                Default: 'Average'.
% ds           - The returned structure contains the following fields
% .P           - Copy of P on input.
% .sfP         - Copy of sfP on input (if non-empty).
% .order       - Copy of order on input, or default.
% .regorder    - Copy of regorder on input, or default.
% .lambda      - Copy of lambda on input, or default.
% .fot         - Copy of fot on input, or default.
% .sot         - Copy of sot on input, or default.
% .fwhm        - Copy of fwhm on input, or default.
% .rem         - Copy of rem on input, or default.
% .p0          - Average position vector (three translations in mm
%                and three rotations in degrees) of scans in P.
% .q           - Deviations from mean position vector of modelled
%                effects. Corresponds to deviations (and deviations
%                squared) of a Taylor expansion of deformation fields.
% .beta        - Coeffeicents of DCT basis functions for partial
%                derivatives of deformation fields w.r.t. modelled
%                effects. Scaled such that resulting deformation 
%                fields have units mm^-1 or deg^-1 (and squares 
%                thereof).
% .SS          - Sum of squared errors for each iteration.
%
%__________________________________________________________________________

% Jesper Andersson
% Copyright (C) 2003-2022 Wellcome Centre for Human Neuroimaging

% This is a major rewrite which uses some new ideas to speed up
% the estimation of the field. The time consuming part is the
% evaluation of A'*A where A is a matrix with the partial
% derivatives for each scan with respect to the parameters
% describing the warp-fields. If we denote the derivative
% matrix for a single scan by Ai, then the estimation of A'*A
% is A'*A = A1'*A1 + A2'*A2 + ... +An'*An where n is the number
% of scans in the time-series. If we model the partial-derivative
% fields w.r.t. two movement parameters (e.g. pitch and roll), each
% by [8 8 8] basis-functions and the image dimensions are
% 64x64x64 then each Ai is a 262144x1024 matrix and we need to
% evaluate and add n of these 1024x1024 matrices. It takes a while.
%
% The new idea is based on the realisation that each of these
% matrices is the kroneceker-product of the relevant movement
% parameters, our basis set and a linear combination (given by
% one column of the inverse of the rotation matrix) of the image 
% gradients in the x-, y- and z-directions. This means that 
% they really aren't all that unique, and that the amount of
% information in these matrices doesn't really warrant all 
% those calculations. After a lot of head-scratching I arrived
% at the following
%
% First some definitions
%
% n: no. of voxels
% m: no. of scans
% l: no. of effects to model
% order: [xorder yorder zorder] no. of basis functions
% q: mxl matrix of scaled realignment parameters for effects to model.
% T{i}: inv(inv(P(i).mat)*P(1).mat);
% t(i,:) = T{i}(1:3,2)';
% B: kron(Bz,By,Bx);
% Ax = repmat(dx,1,prod(order)).*B;
% Ay = repmat(dy,1,prod(order)).*B;
% Az = repmat(dz,1,prod(order)).*B;
%
% Now, my hypothesis is that A'*A is given by
%
% AtA = kron((q.*kron(ones(1,l),t(:,1)))'*(q.*kron(ones(1,l),t(:,1))),Ax'*Ax) +...
% kron((q.*kron(ones(1,l),t(:,2)))'*(q.*kron(ones(1,l),t(:,2))),Ay'*Ay) +...
% kron((q.*kron(ones(1,l),t(:,3)))'*(q.*kron(ones(1,l),t(:,3))),Az'*Az) +...
% 2*kron((q.*kron(ones(1,l),t(:,1)))'*(q.*kron(ones(1,l),t(:,2))),Ax'*Ay) +...
% 2*kron((q.*kron(ones(1,l),t(:,1)))'*(q.*kron(ones(1,l),t(:,3))),Ax'*Az) +...
% 2*kron((q.*kron(ones(1,l),t(:,2)))'*(q.*kron(ones(1,l),t(:,3))),Ay'*Az);
%
% Which turns out to be true. This means that regardless of how many
% scans we have we will always be able to create AtA as a sum of
% six individual AtAs. It has been tested against the previous
% implementation and yields exactly identical results.
%
% I know this isn't much of a derivation, but there will be a paper soon. Sorry.
%
% Other things that are new for the rewrite is
% 1. We have removed the possibility to try and estimate the "static" field.
%    There simply isn't enough information about that field, and for a
%    given time series the estimation is just too likely to fail.
% 2. New option to pass a static field (e.g. estimated from a dual
%    echo-time measurement) as an input parameter.
% 3. Re-estimation of the movement parameters at each iteration.
%    Let us say we have a nodding motion in the time series, and
%    at each nod the brain appear to shrink in the phase-encode
%    direction. The realignment will find the nods, but might in
%    addition mistake the "shrinks" for translations, leading to
%    biased estimates. We attempt to correct that by, at iteration,
%    updating both parameters pertaining to distortions and to
%    movement. In order to do so in an unbiased fashion we have
%    also switched from sampling of images (and deformation fields)
%    on a grid centred on the voxel-centres, to a grid with a 
%    different sampling frequency.   
% 4. Inclusion of a regularisation term. The speed-up has facilitated
%    using more basis-functions, which makes it necessary to impose
%    regularisation (i.e. punisihing some order derivative of the
%    deformation field).
% 5. Change of interpolation model from tri-linear/Sinc to 
%    tri-linear/B-spline.
% 6. Option to include the Jacobian compression/stretching effects
%    in the model for the estimation of the displacement fields.
%    Our tests have indicated that this is NOT a good idea though. 


%-Say hello
%--------------------------------------------------------------------------
spm('FnBanner',mfilename);

%-Parameters
%--------------------------------------------------------------------------
if nargin < 1 || isempty(P), P = spm_select(Inf,'image'); end
if ~isstruct(P), P = spm_vol(P); end

%-Hardcoded default input parameters.
%--------------------------------------------------------------------------
defpar = struct('sfP',             [],...
                'M',               P(1).mat,...
                'fot',             [4 5],...
                'sot',             [],...
                'hold',            [1 1 1 0 0 0]);

%-Merge hardcoded defaults and spm_defaults. Translate spm_defaults
% settings to internal defaults.
%--------------------------------------------------------------------------
ud = spm_get_defaults('unwarp.estimate');
if isfield(ud,'basfcn'),    defpar.order = ud.basfcn; end
if isfield(ud,'regorder'),  defpar.regorder = ud.regorder; end
if isfield(ud,'regwgt'),    defpar.lambda = ud.regwgt; end
if isfield(ud,'fwhm'),      defpar.fwhm = ud.fwhm; end
if isfield(ud,'rem'),       defpar.rem = ud.rem; end
if isfield(ud,'noi'),       defpar.noi = ud.noi; end
if isfield(ud,'expround'),  defpar.exp_round = ud.expround; end

defnames = fieldnames(defpar);

%-Go through input parameters, choosing the default
% for any parameters that are missing, warning the 
% user if there are "unknown" parameters (probably
% reflecting a misspelling).
%--------------------------------------------------------------------------
if nargin < 2 || isempty(par)
    par = defpar;
end
ds = [];
for i=1:length(defnames)
    if isfield(par,defnames{i}) && ~isempty(par.(defnames{i}))
        ds.(defnames{i}) = par.(defnames{i});
    else
        ds.(defnames{i}) = defpar.(defnames{i});
    end
end
parnames = fieldnames(par);
for i=1:length(parnames)
    if ~isfield(defpar,parnames{i})
        warning('Unknown par field %s',parnames{i});
    end
end

%-Resolve ambiguities.
%--------------------------------------------------------------------------
if length(ds.order) == 2
    mm          = sqrt(sum(P(1).mat(1:3,1:3).^2)).*P(1).dim(1:3);
    ds.order(3) = round(ds.order(1)*mm(3)/mm(1));
end
if isfield(ds,'sfP') && ~isempty(ds.sfP)
    if ~isstruct(ds.sfP)
        ds.sfP = spm_vol(ds.sfP);
    end
end
nscan = length(P);
nof   = numel(ds.fot) + size(ds.sot,1);
ds.P  = P;

%-Get matrix of 'expansion point'-corrected movement parameters 
% for which we seek partial derivative fields.
%--------------------------------------------------------------------------
[ds.q, ds.p0] = make_q(P, ds.fot, ds.sot, ds.exp_round);

%-Create matrix for regularisation of warps.
%--------------------------------------------------------------------------
H = make_H(P, ds.order, ds.regorder);

%-Create a temporary smooth time series to use for the estimation.
%--------------------------------------------------------------------------
old_P = P;
if ds.fwhm ~= 0
    spm_uw_show('SmoothStart', length(P));
    sfname = cell(length(old_P), 1);
    for i=1:length(old_P)
        spm_uw_show('SmoothUpdate',i);
        sfname{i} = [tempname spm_file_ext ',1,1'];
        to_smooth = sprintf('%s,%d,%d', old_P(i).fname, old_P(i).n);
        spm_smooth(to_smooth, sfname{i}, ds.fwhm);
    end
    P = spm_vol(char(sfname));
    spm_uw_show('SmoothEnd');
end

%-Now that we have littered the disk with smooth
% temporary files we  should use a try-catch
% block for the rest of the function, to ensure files
% get deleted in the event of an error.
%--------------------------------------------------------------------------
try        % Try block starts here

%-Initialize some stuff.
%--------------------------------------------------------------------------
beta0    = zeros(nof*prod(ds.order), 10);
beta     = zeros(nof*prod(ds.order), 1);
%old_beta = zeros(nof*prod(ds.order),1); % Unused

%-Sample images on irregular grid to avoid biasing
% re-estimated movement parameters with smoothing
% effect from voxel-interpolation (See Andersson ,
% EJNM (1998), 25:575-586.).
%--------------------------------------------------------------------------
ss         = [1.1 1.1 0.9];
xs         = 1:ss(1):P(1).dim(1); xs = xs + (P(1).dim(1)-xs(end)) / 2;
ys         = 1:ss(2):P(1).dim(2); ys = ys + (P(1).dim(2)-ys(end)) / 2;
zs         = 1:ss(3):P(1).dim(3); zs = zs + (P(1).dim(3)-zs(end)) / 2;
M          = spm_matrix([-xs(1)/ss(1)+1 -ys(1)/ss(2)+1 -zs(1)/ss(3)+1]) / diag([ss 1]);
nx         = length(xs);
ny         = length(ys);
nz         = length(zs);
[x,y,z]    = ndgrid(xs,ys,zs);
Bx         = spm_dctmtx(P(1).dim(1), ds.order(1), x(:,1,1));
By         = spm_dctmtx(P(1).dim(2), ds.order(2), y(1,:,1));
Bz         = spm_dctmtx(P(1).dim(3), ds.order(3), z(1,1,:));
xyz        = [x(:) y(:) z(:) ones(length(x(:)),1)]; clear x y z;
def        = zeros(size(xyz, 1), nof);

%-Create file struct for use with spm_orthviews to draw
% representations of the field.
%--------------------------------------------------------------------------
dispP       = P(1);
dispP       = rmfield(dispP,{'fname','descrip','n','private'});
dispP.dim   = [nx ny nz];
dispP.dt    = [64 spm_platform('bigend')];
dispP.pinfo = [1 0]';
p          = spm_imatrix(dispP.mat); p = p.*[zeros(1,6) ss 0 0 0];
p(1)       = -mean(1:nx)*p(7);
p(2)       = -mean(1:ny)*p(8);
p(3)       = -mean(1:nz)*p(9);
dispP.mat  = spm_matrix(p); clear p;

%-We will need to resample the static field (if one was supplied)
% on the same grid (given by xs, ys and zs) as we are going to use
% for the time series. We will assume that the fieldmap has been
% realigned to the space of the first EPI image in the time-series.
%--------------------------------------------------------------------------
if isfield(ds,'sfP') && ~isempty(ds.sfP)
    T         = ds.sfP.mat \ ds.M;
    txyz      = xyz*T(1:3,:)';
    ds.sfield = spm_bsplins(spm_bsplinc(ds.sfP, ds.hold),...
                            txyz(:,1), txyz(:,2), txyz(:,3), ds.hold);
    ds.sfield = ds.sfield(:);
    clear txyz;
else
    ds.sfield = [];
end

msk = get_mask(P,xyz,ds,[nx ny nz]);
ssq = [];

%-Here starts iterative search for deformation fields.
%--------------------------------------------------------------------------
for iter=1:ds.noi
    spm_uw_show('NewIter',iter);

    [ref,dx,dy,dz,P,ds,sf,dispP] = make_ref(ds,def,P,xyz,M,msk,beta,dispP);
    AtA       = build_AtA(dx,dy,dz,Bx,By,Bz,ds,P);
    [Aty,yty] = build_Aty(ref,dx,dy,dz,Bx,By,Bz,ds,P,xyz,beta,sf,def,msk);

    % Clean up a bit, cause inverting AtA may use a lot of memory
    clear ref dx dy dz

    % Check that residual error still decreases.
    if iter > 1 && yty > ssq(iter-1)
        % This means previous iteration was no good,
        % and we should go back to old_beta.
        %
        %beta = old_beta; % Unused
        break;
    else
        ssq(iter) = yty;

        spm_uw_show('StartInv',1);

        % Solve for beta
        Aty = Aty + AtA*beta;
        AtA = AtA + ds.lambda * kron(eye(nof),diag(H)) * ssq(iter)/(nscan*sum(msk));

        try % Fastest if it works
             beta0(:,iter) = AtA\Aty;
        catch % Sometimes necessary
             beta0(:,iter) = pinv(AtA)*Aty;
        end

       %old_beta = beta;
        beta     = beta0(:,iter);

        for i=1:nof
            def(:,i)   = spm_sepmul3d(Bx,By,Bz, beta((i-1)*prod(ds.order)+1:i*prod(ds.order)));
        end

        % If we are re-estimating the movement parameters, remove any DC
        % components from the deformation fields, so that they end up in
        % the movement-parameters instead. It shouldn't make any difference
        % to the variance reduction, but might potentially lead to a better
        % explanation of the variance components.
        % Note that since we sub-sample the images (and the DCT basis set)
        % it is NOT sufficient to reset beta(1) (and beta(prod(order)+1) etc,
        % instead we explicitly subtract the DC component. Note that a DC
        % component does NOT affect the Jacobian.
        if ds.rem ~= 0
            def = def - repmat(mean(def),length(def),1);
        end
        spm_uw_show('EndInv');
        tmp       = dispP.mat;
        dispP.mat = P(1).mat;
        spm_uw_show('FinIter',ssq,def,ds.fot,ds.sot,dispP,ds.q);
        dispP.mat = tmp;
    end
    clear AtA
end

ds.P = old_P;
for i=1:length(ds.P)
    ds.P(i).mat = P(i).mat;    % Save P with new movement parameters.
end
ds.beta = reshape(refit(P,dispP,ds,def),prod(ds.order),nof);
ds.SS   = ssq;
if isfield(ds,'sfield')
    ds = rmfield(ds,'sfield');
end

cleanup(P,ds)
spm_uw_show('FinTot');

%-Document outcome
%--------------------------------------------------------------------------
F = spm_figure('FindWin','Graphics');
if ~isempty(F), spm_figure('Focus', F), spm_print; end

catch exception
    cleanup(P,ds)
    spm_uw_show('FinTot');
    fprintf('Unwarp terminated abnormally.\n');
    rethrow(exception);
end

fprintf('%-40s: %30s\n','Completed',spm('time'))                        %-#


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% Utility functions.
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [q,ep] = make_q(P,fot,sot,exp_round)
%
% P         : Array of file-handles
% fot       : nx1 array of indices for first order effect.
% sot       : nx2 matrix of indices of second order effects.
% exp_round : Point in position space to perform Taylor-expansion
%             around ('First','Last' or 'Average'). 'Average' should
%             (in principle) give the best variance reduction. If a
%             field-map acquired before the time-series is supplied
%             then expansion around the 'First' MIGHT give a slightly
%             better average geometric fidelity.

if strcmpi(exp_round,'average')
    %-Get geometric mean of all transformation matrices.
    % This will be used as the zero-point in the space 
    % of object position vectors (i.e. the point at which
    % the partial derivatives are estimated). This is presumably
    % a little more correct than taking the average of the
    % parameter estimates. 
    mT = zeros(4);
    for i=1:length(P)
        mT = mT+logm(P(i).mat \ P(1).mat);
    end
    mT = real(expm(mT/length(P)));

elseif strcmpi(exp_round,'first')
    mT = eye(4);
elseif strcmpi(exp_round,'last')
    mT = P(end).mat \ P(1).mat;
else
    warning('Unknown expansion point %s',exp_round);
end

%-Rescaling to degrees makes translations and rotations
% roughly equally scaled. Since the scaling influences
% strongly the effects of regularisation, it would surely
% be nice with a more principled approach. Suggestions
% anyone?
ep = [1 1 1 180/pi 180/pi 180/pi zeros(1,6)] .* spm_imatrix(mT);

%-Now, get a nscan-by-nof matrix containing 
% mean (expansion point) corrected values for each effect we
% model.
q = zeros(length(P),numel(fot) + size(sot,1));
for i=1:length(P)
    p = [1 1 1 180/pi 180/pi 180/pi zeros(1,6)] .*...
        spm_imatrix(P(i).mat \ P(1).mat);
    for j=1:numel(fot)
         q(i,j) = p(fot(j)) - ep(fot(j));
    end
    for j=1:size(sot,1)
        q(i,numel(fot)+j) = (p(sot(j,1)) - ep(sot(j,1))) * (p(sot(j,2)) - ep(sot(j,2)));
    end
end
return
%_______________________________________________________________________

%_______________________________________________________________________
function H = make_H(P,order,regorder)
% Utility Function to create Regularisation term of AtA.

mm = sqrt(sum(P(1).mat(1:3,1:3).^2)).*P(1).dim(1:3);
kx = (pi*((1:order(1))'-1)/mm(1)).^2;
ky = (pi*((1:order(2))'-1)/mm(2)).^2;
kz = (pi*((1:order(3))'-1)/mm(3)).^2;

switch regorder
    case 0
        % Cost function based on sum of squares
        H = (1*kron(kz.^0, kron(ky.^0, kx.^0)) +...
             1*kron(kz.^0, kron(ky.^0, kx.^0)) +...
             1*kron(kz.^0, kron(ky.^0, kx.^0)) );
    case 1
        % Cost function based on sum of squared 1st derivatives
        H  = (1*kron(kz.^1, kron(ky.^0, kx.^0)) +...
              1*kron(kz.^0, kron(ky.^1, kx.^0)) +...
              1*kron(kz.^0, kron(ky.^0, kx.^1)) );
    case 2
        % Cost function based on sum of squared 2nd derivatives
        H = (1*kron(kz.^2, kron(ky.^0, kx.^0)) +...
             1*kron(kz.^0, kron(ky.^2, kx.^0)) +...
             1*kron(kz.^0, kron(ky.^0, kx.^2)) +...
             3*kron(kz.^1, kron(ky.^1, kx.^0)) +...
             3*kron(kz.^1, kron(ky.^0, kx.^1)) +...
             3*kron(kz.^0, kron(ky.^1, kx.^1)) );
    case 3
        % Cost function based on sum of squared 3rd derivatives
        H = (1*kron(kz.^3, kron(ky.^0, kx.^0)) +...
             1*kron(kz.^0, kron(ky.^3, kx.^0)) +...
             1*kron(kz.^0, kron(ky.^0, kx.^3)) +...
             3*kron(kz.^2, kron(ky.^1, kx.^0)) +...
             3*kron(kz.^2, kron(ky.^0, kx.^1)) +...
             3*kron(kz.^1, kron(ky.^2, kx.^0)) +...
             3*kron(kz.^0, kron(ky.^2, kx.^1)) +...
             3*kron(kz.^1, kron(ky.^0, kx.^2)) +...
             3*kron(kz.^0, kron(ky.^1, kx.^2)) +...
             6*kron(kz.^1, kron(ky.^1, kx.^1)) );
    otherwise
        error('Invalid order of regularisation');
end
return
%_______________________________________________________________________

%_______________________________________________________________________
function AtA = build_AtA(dx,dy,dz, Bx,By,Bz, ds, P)
% Now lets build up the design matrix A, or rather A'*A  since
% A itself would be forbiddingly large.

spm_uw_show('StartAtA',6);

nof   = numel(ds.fot) + size(ds.sot,1);

AxtAx = uwAtA1(dx.*dx, Bx,By,Bz); spm_uw_show('NewAtA',1);
AytAy = uwAtA1(dy.*dy, Bx,By,Bz); spm_uw_show('NewAtA',2);
AztAz = uwAtA1(dz.*dz, Bx,By,Bz); spm_uw_show('NewAtA',3);
AxtAy = uwAtA1(dx.*dy, Bx,By,Bz); spm_uw_show('NewAtA',4);
AxtAz = uwAtA1(dx.*dz, Bx,By,Bz); spm_uw_show('NewAtA',5);
AytAz = uwAtA1(dy.*dz, Bx,By,Bz); spm_uw_show('NewAtA',6);

R = zeros(length(P),3);
for i=1:length(P)
    tmp    = P(1).mat \ P(i).mat;
    R(i,:) = tmp(1:3,2)';
end

tmp  = ones(1,nof);
tmp1 = ds.q.*kron(tmp, R(:,1));
tmp2 = ds.q.*kron(tmp, R(:,2));
tmp3 = ds.q.*kron(tmp, R(:,3));
AtA  =    kron(tmp1'*tmp1, AxtAx) + kron(tmp2'*tmp2, AytAy) + kron(tmp3'*tmp3, AztAz) +...
       2*(kron(tmp1'*tmp2, AxtAy) + kron(tmp1'*tmp3, AxtAz) + kron(tmp2'*tmp3, AytAz));

spm_uw_show('EndAtA');
return;
%_______________________________________________________________________

%_______________________________________________________________________
function AtA = uwAtA1(y, Bx, By, Bz)
% Calculating off-diagonal block of AtA.

[nx,mx] = size(Bx);
[ny,my] = size(By);
[nz,mz] = size(Bz);
AtA     = zeros(mx*my*mz);
for sl =1:nz
     tmp  = reshape(y((sl-1)*nx*ny+1:sl*nx*ny),nx,ny);
     spm_krutil(Bz(sl,:)'*Bz(sl,:), spm_krutil(tmp, Bx, By, 1), AtA);
%    AtA  = AtA + kron(Bz(sl,:)'*Bz(sl,:),spm_krutil(tmp,Bx,By,1));
end
return
%_______________________________________________________________________

%_______________________________________________________________________
function [Aty, yty] = build_Aty(ref, dx,dy,dz, Bx,By,Bz, ds, P, xyz, beta, sf, def, msk)
% Building Aty.

nof = numel(ds.fot) + size(ds.sot,1);
Aty = zeros(nof*prod(ds.order),1);
yty = 0;

spm_uw_show('StartAty',length(P));
for scan = 1:length(P)
    spm_uw_show('NewAty',scan);
    T    = P(scan).mat \ ds.M;
    txyz = xyz*T(1:3,:)';
    if ~(all(beta == 0) && isempty(ds.sfield))
        txyz(:,2) = txyz(:,2) + spm_uw_get_image_def(scan,ds,def);
    end
    y    = sf(scan) * spm_bsplins(spm_bsplinc(P(scan), ds.hold),...
                                  txyz(:,1), txyz(:,2), txyz(:,3), ds.hold);

    y_diff       = (ref - y).*msk;
    indx         = isnan(y);
    y_diff(indx) = 0;
    iTcol        = inv(T(1:3,1:3));
    tmpAty       = spm_sepmul3d(Bx',By',Bz',([dx dy dz]*iTcol(:,2)).*y_diff);
    for i=1:nof
        rindx      = (i-1)*prod(ds.order)+1:i*prod(ds.order);
        Aty(rindx) = Aty(rindx) + ds.q(scan,i)*tmpAty;
    end
    yty = yty + y_diff'*y_diff;
end
spm_uw_show('EndAty');
return;
%_______________________________________________________________________

%_______________________________________________________________________
function [ref, dx,dy,dz, P, ds, sf, dispP] = make_ref(ds, def, P, xyz, M, msk, beta, dispP)
% First of all get the mean of all scans given their
% present deformation fields. When using this as the
% "reference" we will explicitly be minimising variance.
%
% First scan in P is still reference in terms of
% y-direction (in the scanner framework). A single set of
% image gradients is estimated from the mean of all scans,
% transformed into the space of P(1), and used for all scans.
%
spm_uw_show('StartRef',length(P));

rem = ds.rem;
if all(beta==0), rem=0; end

if rem
    [m_ref,D] = get_refD(ds,def,P,xyz,msk);
end
Dty = cell(length(P), 1);
sf  = zeros(length(P), 1);
ref = zeros(size(xyz, 1), 1);
for i=1:length(P)
    spm_uw_show('NewRef',i);
    T    = P(i).mat \ ds.M;
    txyz = xyz*T(1:3,:)';
    if ~(all(beta == 0) && isempty(ds.sfield))
        txyz(:,2) = txyz(:,2) + spm_uw_get_image_def(i, ds, def);
    end
    f     = spm_bsplins(spm_bsplinc(P(i), ds.hold),...
                        txyz(:,1), txyz(:,2), txyz(:,3), ds.hold);
    indx  = find(~isnan(f));
    sf(i) = 1 / (sum(f(indx).*msk(indx)) / sum(msk(indx)));
    ref   = ref + f;

    if rem
        f(isnan(f)) = 0;
        Dty{i}      = (((m_ref - sf(i)*f).*msk)'*D)';
    end
end
ref  = ref / length(P);
ref  = reshape(ref,dispP.dim(1:3));
indx = find(~isnan(ref));

%-Scale to roughly 100 mean intensity to ensure
% consistent weighting of regularisation.
%--------------------------------------------------------------------------
gl             = spm_global(ref);
sf             = (100 * (sum(ref(indx).*msk(indx)) / sum(msk(indx))) / gl) * sf;
ref            = (100 / gl) * ref;
dispP.dat      = ref;
txyz           = xyz*M(1:3,:)';
[ref,dx,dy,dz] = spm_bsplins(spm_bsplinc(ref, ds.hold),...
                             txyz(:,1), txyz(:,2), txyz(:,3), ds.hold);
ref            = ref.*msk;
dx             =  dx.*msk;
dy             =  dy.*msk;
dz             =  dz.*msk;

%-Re-estimate (or rather nudge) movement parameters.
%--------------------------------------------------------------------------
if rem ~= 0
    iDtD = inv(D'*D);
    for i=2:length(P)
        P(i).mat = spm_matrix((iDtD*Dty{i})') \ P(i).mat;
    end
    [ds.q,ds.p0] = make_q(P,ds.fot,ds.sot,ds.exp_round);
end
spm_uw_show('EndRef');
return
%_______________________________________________________________________

%_______________________________________________________________________
function [m_ref, D] = get_refD(ds, def, P, xyz, msk)
% Get partials w.r.t. movements from first scan.

T    = P(1).mat\ds.M;
txyz = xyz*T';
[m_ref,dx,dy,dz] = spm_bsplins(spm_bsplinc(P(1), ds.hold), ...
                               txyz(:,1),...
                               txyz(:,2) + spm_uw_get_image_def(1,ds,def),...
                               txyz(:,3), ds.hold);
indx        = ~isnan(m_ref);
mref_sf     = sum(msk(indx)) / sum(m_ref(indx).*msk(indx));
m_ref(indx) = m_ref(indx) .* mref_sf; m_ref(~indx) = 0;
dx(indx)    =    dx(indx) .* mref_sf; dx(~indx)    = 0;
dy(indx)    =    dy(indx) .* mref_sf; dy(~indx)    = 0;
dz(indx)    =    dz(indx) .* mref_sf; dz(~indx)    = 0;
D           = make_D(dx.*msk, dy.*msk, dz.*msk, txyz, P(1).mat);
return;
%_______________________________________________________________________

%_______________________________________________________________________
function D = make_D(dx, dy, dz, xyz, M)
% Utility function that creates a matrix of partial
% derivatives in units of /mm and /radian.
%
% dx,dy,dz   - Partial derivatives (/pixel) wrt x-, y- and z-translation.
% xyz        - xyz-matrix (original positions).
% M          - Current voxel->world matrix.

D = zeros(length(xyz), 6);
tiny = 0.0001;
for i=1:6
    p      = [0 0 0 0 0 0 1 1 1 0 0 0];
    p(i)   = p(i)+tiny;
    T      = M\spm_matrix(p)*M;
    dxyz   = xyz*T';
    dxyz   = dxyz(:,1:3) - xyz(:,1:3);
    D(:,i) = sum(dxyz.*[dx dy dz], 2)/tiny;
end
return
%_______________________________________________________________________

%_______________________________________________________________________
function cleanup(P,ds)
% Delete temporary smooth files.
%
if ~isempty(ds.fwhm) && ds.fwhm > 0
    for i=1:length(P)
        spm_unlink(P(i).fname);
        [fpath,fname] = fileparts(P(i).fname);
        spm_unlink(fullfile(fpath,[fname '.hdr']));
        spm_unlink(fullfile(fpath,[fname '.mat']));
    end
end
return;
%_______________________________________________________________________

%_______________________________________________________________________
function beta = refit(P,dispP,ds,def)
% We have now estimated the fields on a grid that
% does not coincide with the voxel centers. We must
% now calculate the betas that correspond to a
% a grid centered on the voxel centers.
%
spm_uw_show('StartRefit',1);
rsP     = P(1);
p       = spm_imatrix(rsP.mat);
p       = [zeros(1,6) p(7:9) 0 0 0];
p(1)    = -mean(1:rsP.dim(1))*p(7);
p(2)    = -mean(1:rsP.dim(2))*p(8);
p(3)    = -mean(1:rsP.dim(3))*p(9);
rsP.mat = spm_matrix(p);
[x,y,z] = ndgrid(1:rsP.dim(1),1:rsP.dim(2),1:rsP.dim(3));
xyz     = [x(:) y(:) z(:) ones(numel(x),1)];
txyz    = ((dispP.mat \ rsP.mat)*xyz')';
Bx      = spm_dctmtx(rsP.dim(1),ds.order(1));
By      = spm_dctmtx(rsP.dim(2),ds.order(2));
Bz      = spm_dctmtx(rsP.dim(3),ds.order(3));
nx      = rsP.dim(1); mx = ds.order(1);
ny      = rsP.dim(2); my = ds.order(2);
nz      = rsP.dim(3); mz = ds.order(3);
nof     = numel(ds.fot) + size(ds.sot,1);
for i=1:nof
    dispP.dat   = reshape(def(:,i),dispP.dim(1:3));
    field       = spm_bsplins(spm_bsplinc(dispP, ds.hold),...
                              txyz(:,1), txyz(:,2), txyz(:,3), ds.hold);
    indx        = isnan(field);
    wgt         = ones(size(field));
    wgt(indx)   = 0;
    field(indx) = 0;
    AtA         = zeros(mx*my*mz);
    Aty         = zeros(mx*my*mz, 1);
    for sl = 1:nz
        indx  = (sl-1)*nx*ny+1:sl*nx*ny;
        tmp1  = reshape(field(indx).*wgt(indx), nx, ny);
        tmp2  = reshape(wgt(indx), nx, ny);
        spm_krutil(Bz(sl,:)'*Bz(sl,:), spm_krutil(tmp2, Bx, By, 1), AtA);
%       AtA   = AtA + kron(Bz(sl,:)'*Bz(sl,:),spm_krutil(tmp2,Bx,By,1));
        Aty   = Aty + kron(Bz(sl,:)', spm_krutil(tmp1, Bx, By, 0));
    end
    beta((i-1)*prod(ds.order)+1:i*prod(ds.order)) = AtA\Aty;
end
spm_uw_show('EndRefit');
return;
%_______________________________________________________________________

%_______________________________________________________________________
function msk = get_mask(P,xyz,ds,dm)
%
% Create a mask to avoid regions where data doesn't exist
% for all scans. This mask is slightly less Q&D than that
% of version 1 of Unwarp. It checks where data exist for
% all scans with present movement parameters and given
% (optionally) the static field. It creates a mask with
% non-zero values only for those voxels, and then does a
% 3D erode to preempt effects of re-estimated movement
% parameters and movement-by-susceptibility effects.
%
spm_uw_show('MaskStart',length(P));
msk = true(length(xyz),1);
for i=1:length(P)
    txyz = xyz * (P(i).mat\ds.M)';
    tmsk = (txyz(:,1)>=1 & txyz(:,1)<=P(1).dim(1) &...
            txyz(:,2)>=1 & txyz(:,2)<=P(1).dim(2) &...
            txyz(:,3)>=1 & txyz(:,3)<=P(1).dim(3));
    msk  = msk & tmsk;
    spm_uw_show('MaskUpdate',i);
end

%-Include static field in mask estimation if one has been supplied.
%--------------------------------------------------------------------------
if isfield(ds,'sfP') && ~isempty(ds.sfP)
    txyz = xyz * (ds.sfP.mat\ds.M)';
    tmsk = (txyz(:,1)>=1 & txyz(:,1)<=ds.sfP.dim(1) &...
            txyz(:,2)>=1 & txyz(:,2)<=ds.sfP.dim(2) &...
            txyz(:,3)>=1 & txyz(:,3)<=ds.sfP.dim(3));
    msk  = msk & tmsk;
end
msk = erode_msk(msk,dm);
spm_uw_show('MaskEnd');

% maskP                 = P(1);
% [mypath,myname,myext] = fileparts(maskP.fname);
% maskP.fname           = fullfile(mypath,['mask' myext]);
% maskP.dim             = [nx ny nz];
% maskP.dt              = P(1).dt;
% maskP                 = spm_write_vol(maskP,reshape(msk,[nx ny nz]));
return;
%_______________________________________________________________________

%_______________________________________________________________________
function omsk = erode_msk(msk, dim)
omsk = zeros(dim + [4 4 4]);
omsk(3:end-2, 3:end-2, 3:end-2) = reshape(msk, dim);
omsk = spm_erode(omsk(:), dim + [4 4 4]);
omsk = reshape(omsk, dim + [4 4 4]);
omsk = omsk(3:end-2, 3:end-2, 3:end-2);
omsk = omsk(:);
return
%_______________________________________________________________________
