function wss_dist= comp_wss(cleanFile, enhancedFile);
% ----------------------------------------------------------------------
%
%     Weighted Spectral Slope (WSS) Objective Speech Quality Measure
%
%     This function implements the Weighted Spectral Slope (WSS)
%     distance measure originally proposed in [1].  The algorithm
%     works by first decomposing the speech signal into a set of
%     frequency bands (this is done for both the test and reference
%     frame).  The intensities within each critical band are
%     measured.  Then, a weighted distances between the measured
%     slopes of the log-critical band spectra are computed.
%     This measure is also described in Section 2.2.9 (pages 56-58)
%     of [2].
%
%     Whereas Klatt's original measure used 36 critical-band
%     filters to estimate the smoothed short-time spectrum, this
%     implementation considers a bank of 25 filters spanning
%     the 4 kHz bandwidth.
%
%   Usage:  wss_dist=comp_wss(cleanFile.wav, enhancedFile.wav)
%
%         cleanFile.wav - clean input file in .wav format
%         enhancedFile  - enhanced output file in .wav format
%         wss_dist      - computed spectral slope distance
%
%  Example call:  ws =comp_wss('sp04.wav','enhanced.wav')
%
%  References:
%
%     [1] D. H. Klatt, "Prediction of Perceived Phonetic Distance
%	    from Critical-Band Spectra: A First Step", Proc. IEEE
%	    ICASSP'82, Volume 2, pp. 1278-1281, May, 1982.
%
%     [2] S. R. Quackenbush, T. P. Barnwell, and M. A. Clements,
%	    Objective Measures of Speech Quality.  Prentice Hall
%	    Advanced Reference Series, Englewood Cliffs, NJ, 1988,
%	    ISBN: 0-13-629056-6.
%
%  Authors: Bryan L. Pellom and John H. L. Hansen (July 1998)
%  Modified by: Philipos C. Loizou  (Oct 2006)
%
% Copyright (c) 2006 by Philipos C. Loizou
% $Revision: 0.0 $  $Date: 10/09/2006 $
%
% ----------------------------------------------------------------------
if nargin~=2
    fprintf('USAGE: WSS=comp_wss(cleanFile.wav, enhancedFile.wav)\n');
    fprintf('For more help, type: help comp_wss\n\n');
    return;
end

alpha= 0.95;

[data1, Srate1, Nbits1]= wavread(cleanFile);
[data2, Srate2, Nbits2]= wavread(enhancedFile);
if ( Srate1~= Srate2) | ( Nbits1~= Nbits2)
    error( 'The two files do not match!\n');
end

len= min( length( data1), length( data2));
data1= data1( 1: len)+eps;
data2= data2( 1: len)+eps;

wss_dist_vec= wss( data1, data2,Srate1);
wss_dist_vec= sort( wss_dist_vec);
wss_dist= mean( wss_dist_vec( 1: round( length( wss_dist_vec)*alpha)));

function distortion = wss(clean_speech, processed_speech,sample_rate)

% ----------------------------------------------------------------------
% Check the length of the clean and processed speech.  Must be the same.
% ----------------------------------------------------------------------

clean_length      = length(clean_speech);
processed_length  = length(processed_speech);

if (clean_length ~= processed_length)
  disp('Error: Files  musthave same length.');
  return
end

% ----------------------------------------------------------------------
% Global Variables
% ----------------------------------------------------------------------

winlength   = round(30*sample_rate/1000); 	   % window length in samples
skiprate    = floor(winlength/4);		   % window skip in samples
max_freq    = sample_rate/2;	   % maximum bandwidth
num_crit    = 25;		   % number of critical bands

USE_FFT_SPECTRUM = 1;		   % defaults to 10th order LP spectrum
n_fft       = 2^nextpow2(2*winlength);
n_fftby2    = n_fft/2;		   % FFT size/2
Kmax        = 20;		   % value suggested by Klatt, pg 1280
Klocmax     = 1;		   % value suggested by Klatt, pg 1280

% ----------------------------------------------------------------------
% Critical Band Filter Definitions (Center Frequency and Bandwidths in Hz)
% ----------------------------------------------------------------------

cent_freq(1)  = 50.0000;   bandwidth(1)  = 70.0000;
cent_freq(2)  = 120.000;   bandwidth(2)  = 70.0000;
cent_freq(3)  = 190.000;   bandwidth(3)  = 70.0000;
cent_freq(4)  = 260.000;   bandwidth(4)  = 70.0000;
cent_freq(5)  = 330.000;   bandwidth(5)  = 70.0000;
cent_freq(6)  = 400.000;   bandwidth(6)  = 70.0000;
cent_freq(7)  = 470.000;   bandwidth(7)  = 70.0000;
cent_freq(8)  = 540.000;   bandwidth(8)  = 77.3724;
cent_freq(9)  = 617.372;   bandwidth(9)  = 86.0056;
cent_freq(10) = 703.378;   bandwidth(10) = 95.3398;
cent_freq(11) = 798.717;   bandwidth(11) = 105.411;
cent_freq(12) = 904.128;   bandwidth(12) = 116.256;
cent_freq(13) = 1020.38;   bandwidth(13) = 127.914;
cent_freq(14) = 1148.30;   bandwidth(14) = 140.423;
cent_freq(15) = 1288.72;   bandwidth(15) = 153.823;
cent_freq(16) = 1442.54;   bandwidth(16) = 168.154;
cent_freq(17) = 1610.70;   bandwidth(17) = 183.457;
cent_freq(18) = 1794.16;   bandwidth(18) = 199.776;
cent_freq(19) = 1993.93;   bandwidth(19) = 217.153;
cent_freq(20) = 2211.08;   bandwidth(20) = 235.631;
cent_freq(21) = 2446.71;   bandwidth(21) = 255.255;
cent_freq(22) = 2701.97;   bandwidth(22) = 276.072;
cent_freq(23) = 2978.04;   bandwidth(23) = 298.126;
cent_freq(24) = 3276.17;   bandwidth(24) = 321.465;
cent_freq(25) = 3597.63;   bandwidth(25) = 346.136;

bw_min      = bandwidth (1);	   % minimum critical bandwidth

% ----------------------------------------------------------------------
% Set up the critical band filters.  Note here that Gaussianly shaped
% filters are used.  Also, the sum of the filter weights are equivalent
% for each critical band filter.  Filter less than -30 dB and set to
% zero.
% ----------------------------------------------------------------------

min_factor = exp (-30.0 / (2.0 * 2.303));       % -30 dB point of filter

for i = 1:num_crit
  f0 = (cent_freq (i) / max_freq) * (n_fftby2);
  all_f0(i) = floor(f0);
  bw = (bandwidth (i) / max_freq) * (n_fftby2);
  norm_factor = log(bw_min) - log(bandwidth(i));
  j = 0:1:n_fftby2-1;
  crit_filter(i,:) = exp (-11 *(((j - floor(f0)) ./bw).^2) + norm_factor);
  crit_filter(i,:) = crit_filter(i,:).*(crit_filter(i,:) > min_factor);
end

% ----------------------------------------------------------------------
% For each frame of input speech, calculate the Weighted Spectral
% Slope Measure
% ----------------------------------------------------------------------

num_frames = clean_length/skiprate-(winlength/skiprate); % number of frames
start      = 1;					% starting sample
window     = 0.5*(1 - cos(2*pi*(1:winlength)'/(winlength+1)));

for frame_count = 1:num_frames

   % ----------------------------------------------------------
   % (1) Get the Frames for the test and reference speech.
   %     Multiply by Hanning Window.
   % ----------------------------------------------------------

   clean_frame = clean_speech(start:start+winlength-1);
   processed_frame = processed_speech(start:start+winlength-1);
   clean_frame = clean_frame.*window;
   processed_frame = processed_frame.*window;

   % ----------------------------------------------------------
   % (2) Compute the Power Spectrum of Clean and Processed
   % ----------------------------------------------------------

    if (USE_FFT_SPECTRUM)
       clean_spec     = (abs(fft(clean_frame,n_fft)).^2);
       processed_spec = (abs(fft(processed_frame,n_fft)).^2);
    else
       a_vec = zeros(1,n_fft);
       a_vec(1:11) = lpc(clean_frame,10);
       clean_spec     = 1.0/(abs(fft(a_vec,n_fft)).^2)';

       a_vec = zeros(1,n_fft);
       a_vec(1:11) = lpc(processed_frame,10);
       processed_spec = 1.0/(abs(fft(a_vec,n_fft)).^2)';
    end

   % ----------------------------------------------------------
   % (3) Compute Filterbank Output Energies (in dB scale)
   % ----------------------------------------------------------

   for i = 1:num_crit
      clean_energy(i) = sum(clean_spec(1:n_fftby2) ...
		            .*crit_filter(i,:)');
      processed_energy(i) = sum(processed_spec(1:n_fftby2) ...
			        .*crit_filter(i,:)');
   end
   clean_energy = 10*log10(max(clean_energy,1E-10));
   processed_energy = 10*log10(max(processed_energy,1E-10));

   % ----------------------------------------------------------
   % (4) Compute Spectral Slope (dB[i+1]-dB[i])
   % ----------------------------------------------------------

   clean_slope     = clean_energy(2:num_crit) - ...
		     clean_energy(1:num_crit-1);
   processed_slope = processed_energy(2:num_crit) - ...
		     processed_energy(1:num_crit-1);

   % ----------------------------------------------------------
   % (5) Find the nearest peak locations in the spectra to
   %     each critical band.  If the slope is negative, we
   %     search to the left.  If positive, we search to the
   %     right.
   % ----------------------------------------------------------

   for i = 1:num_crit-1

       % find the peaks in the clean speech signal

       if (clean_slope(i)>0) 		% search to the right
	  n = i;
          while ((n<num_crit) & (clean_slope(n) > 0))
	     n = n+1;
 	  end
	  clean_loc_peak(i) = clean_energy(n-1);
       else				% search to the left
          n = i;
	  while ((n>0) & (clean_slope(n) <= 0))
	     n = n-1;
 	  end
	  clean_loc_peak(i) = clean_energy(n+1);
       end

       % find the peaks in the processed speech signal

       if (processed_slope(i)>0) 	% search to the right
	  n = i;
          while ((n<num_crit) & (processed_slope(n) > 0))
	     n = n+1;
	  end
	  processed_loc_peak(i) = processed_energy(n-1);
       else				% search to the left
          n = i;
	  while ((n>0) & (processed_slope(n) <= 0))
	     n = n-1;
 	  end
	  processed_loc_peak(i) = processed_energy(n+1);
       end

   end

   % ----------------------------------------------------------
   %  (6) Compute the WSS Measure for this frame.  This
   %      includes determination of the weighting function.
   % ----------------------------------------------------------

   dBMax_clean       = max(clean_energy);
   dBMax_processed   = max(processed_energy);

   % The weights are calculated by averaging individual
   % weighting factors from the clean and processed frame.
   % These weights W_clean and W_processed should range
   % from 0 to 1 and place more emphasis on spectral
   % peaks and less emphasis on slope differences in spectral
   % valleys.  This procedure is described on page 1280 of
   % Klatt's 1982 ICASSP paper.

   Wmax_clean        = Kmax ./ (Kmax + dBMax_clean - ...
		 	    clean_energy(1:num_crit-1));
   Wlocmax_clean     = Klocmax ./ ( Klocmax + clean_loc_peak - ...
				clean_energy(1:num_crit-1));
   W_clean           = Wmax_clean .* Wlocmax_clean;

   Wmax_processed    = Kmax ./ (Kmax + dBMax_processed - ...
			        processed_energy(1:num_crit-1));
   Wlocmax_processed = Klocmax ./ ( Klocmax + processed_loc_peak - ...
			            processed_energy(1:num_crit-1));
   W_processed       = Wmax_processed .* Wlocmax_processed;

   W = (W_clean + W_processed)./2.0;

   distortion(frame_count) = sum(W.*(clean_slope(1:num_crit-1) - ...
		       processed_slope(1:num_crit-1)).^2);

   % this normalization is not part of Klatt's paper, but helps
   % to normalize the measure.  Here we scale the measure by the
   % sum of the weights.

   distortion(frame_count) = distortion(frame_count)/sum(W);

   start = start + skiprate;

end


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% 1. use vectorization to generate face and face2elem data structure
%% this vectorization code does not work for uniform tets

% tic
%
% matlabTri = 1;
% N = size(mesh.p,1);
% node = [mesh.p;mesh.eIntP];
% allCutElem = mesh.t(mesh.tLoc<0,:);
% isCutElem = (mesh.tLoc<0);
% vSign = [mesh.pLoc;zeros(size(mesh.eIntP,1),1)];
%
% isInterfaceNode = false(size(node,1), 1);
% isInterfaceNode(allCutElem(:)) = true; % include the vertices of allCutElem
% isInterfaceNode(N+1:end) = true;    % and the cut points and the aux points
% interfaceNode = node(isInterfaceNode,:);
% matlabversion = version();
% if matlabTri == 1
%     DT = delaunayTriangulation(interfaceNode);
%     tetElem = DT.ConnectivityList;
% elseif matlabTri == 0
%     constrant1 = [mesh.e(mesh.eLoc<0,1),N-mesh.eLoc(mesh.eLoc<0)];
%     constrant2 = [N-mesh.eLoc(mesh.eLoc<0),mesh.e(mesh.eLoc<0,2)];
%     constrant = [constrant1;constrant2];
%     DT = constrainedDelaunayTetGen(interfaceNode,constrant);
%     tetElem = DT.ConnectivityList;
% end
% [tetElem, ortidx, volume] = fixorder3(interfaceNode, tetElem);
%
% localidx2globalidx = find(isInterfaceNode); % tetElem points to interfaceNode
% tetElem = localidx2globalidx(tetElem);  % map to the global index of node
% bc = (node(tetElem(:, 1), :) + node(tetElem(:, 2), :) ...
%     + node(tetElem(:, 3), :) + node(tetElem(:, 4), :))/4.0;
% idx2cube = FindElemId(bc, mesh); % Get the cube index of each tetElem
% idx2cubeBad = (idx2cube<=0);
% idx2cube = idx2cube(~idx2cubeBad);
% tetElem = tetElem(isCutElem(idx2cube), :); % Just keep the tets in cut elem
% volume = volume(isCutElem(idx2cube), :);
% idx2cube = idx2cube(isCutElem(idx2cube));  % keep cube idx in cut elements
% % volumetet0 =  accumarray(idx2cube,volume);
% % volumetet0 = volumetet0(volumetet0>0);
% %[tetElemOld, ortidxOld, volumeOld] = fixorder3(interfaceNode, mesh.t(mesh.tLoc<0,:));
%
% % There might exist some bad tet elements whose vertices are in
% % different cubes. We just need to keep the tet in every cut elem. So
% % here we get rid of them. (but do not handle different tets)
% NT = size(tetElem,1);
% X = reshape(node(tetElem,1), NT, 4);
% Y = reshape(node(tetElem,2), NT, 4);
% Z = reshape(node(tetElem,3), NT, 4);
% isBadTet =  (max(X,[],2) - min(X,[],2)) > 2*h - eps ... % d > 2h means
%           | (max(Y,[],2) - min(Y,[],2)) > 2*h - eps ... % in differnt cube
%           | (max(Z,[],2) - min(Z,[],2)) > 2*h - eps;
% tetElem = tetElem(~isBadTet,:);
%
% toc

% plot to check
%showmesh3(node,tetElem);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% this approach is based on forlopp which is slow
tic

Np = size(mesh.p,1);
node = [mesh.p;mesh.eIntP];
allCutElem = mesh.t(mesh.tLoc<0,:);
isCutElem = (mesh.tLoc<0);
vSign = [mesh.pLoc;zeros(size(mesh.eIntP,1),1)];
isInterfaceNode = false(size(node,1), 1);
isInterfaceNode(allCutElem(:)) = true; % include the vertices of allCutElem
numCutElemV1 = sum(isInterfaceNode);
isInterfaceNode(Np+1:end) = true;    % and the cut points and the aux points
allCutElemReoderTmp = zeros(size(node,1),1);
numCutElemV2 = sum(isInterfaceNode);
allCutElemReoderTmp(isInterfaceNode) = 1:numCutElemV2;
allCutElemReoder = allCutElemReoderTmp(allCutElem);
interfaceNode = node(isInterfaceNode,:);
ntI = -min(mesh.tLoc);
intID = find(mesh.tLoc<0);

tetElem = zeros(12*ntI,4);
tetElemLoc = zeros(12*ntI,1);
idx2cube = zeros(12*ntI,1);
tetcount = 0;

for i = 1:ntI
    tID = intID(i);
    t_e = mesh.t_e(tID,:);
    pLocK = mesh.pLoc(mesh.t(tID,:));
    vert0 = mesh.p(mesh.t(tID,:),:);
    nodeid = [allCutElemReoder(i,:),numCutElemV1-mesh.eLoc(t_e(mesh.eLoc(t_e)<0))'];
    vert2 = vert0(pLocK>0,:); % plus domain
    vert1 = vert0(pLocK<0,:); % minus domain
    intpt0 = mesh.eIntP(-mesh.eLoc(t_e(mesh.eLoc(t_e)<0)),:);
    id1 = [find(pLocK<0)', 5:4+size(intpt0,1)];
    id2 = [find(pLocK>0)', 5:4+size(intpt0,1)];
    p1 = [vert1;intpt0]; DT = delaunayTriangulation(p1); t1 = DT.ConnectivityList;
    p2 = [vert2;intpt0]; DT = delaunayTriangulation(p2); t2 = DT.ConnectivityList;
    tetElem(tetcount+1:tetcount+size(t1,1),:) = nodeid(id1(t1));
    idx2cube(tetcount+1:tetcount+size(t1,1)) = tID;
    tetElemLoc(tetcount+1:tetcount+size(t1,1)) = 1; % inside subdomain
    tetcount = tetcount+size(t1,1);
    tetElem(tetcount+1:tetcount+size(t2,1),:) = nodeid(id2(t2));
    idx2cube(tetcount+1:tetcount+size(t2,1)) = tID;
    tetElemLoc(tetcount+1:tetcount+size(t2,1)) = 2; % outside subdomain
    tetcount = tetcount+size(t2,1);
end
tetElem(tetcount+1:end,:) = [];
tetElemLoc(tetcount+1:end,:) = [];
idx2cube(tetcount+1:end,:) = [];
[tetElem, ortidx, volume] = fixorder3(interfaceNode, tetElem);
% there are some tets which are coplane, need to get rid of them
idPlane = (volume<=10^(-12));
tetElem(idPlane,:) = [];
volume(idPlane,:) = [];
tetElemLoc(idPlane,:) = [];
idx2cube(idPlane,:) = [];
localidx2globalidx = find(isInterfaceNode); % tetElem points to interfaceNode
tetElem = localidx2globalidx(tetElem);  % map to the global index of node
% there are some tets which are on the interface, but contained in the surrounding tets
% need to get rid of them (but the following algorithm may not be robust)
tetSign = vSign(tetElem);
idSliver = (sum(abs(tetSign),2)==0);
tetElem(idSliver,:) = [];
volume(idSliver,:) = [];
tetElemLoc(idSliver,:) = [];
idx2cube(idSliver,:) = [];
PolyVolume = accumarray([idx2cube,tetElemLoc],volume);
PolyVolume = PolyVolume(mesh.tLoc<0,:);
h = (PolyVolume(:,1) + PolyVolume(:,2)).^(1/3);
% the first component contains the volume of the polyhedron on the
% subdomain 1, the first component contains the volume of the polyhedron
% on the subdomain 2

toc

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% try another approach for vectorization

% tic
%
% Np = size(mesh.p,1);
% Ncube = size(mesh.t,1)/6;
% node = [mesh.p;mesh.eIntP];
% allCutElem = mesh.t(mesh.tLoc<0,:);
% isCutElem = (mesh.tLoc<0);
% vSign = [mesh.pLoc;zeros(size(mesh.eIntP,1),1)];
% intID = find(mesh.tLoc<0);
% isInterfaceNode = false(size(node,1), 1);
% isInterfaceNode(allCutElem(:)) = true; % include the vertices of allCutElem
% isInterfaceNode(Np+1:end) = true;    % and the cut points and the aux points
%
% for j = 1:6
%
%     Indj = (j-1)*Ncube + 1:Ncube;
%     Indj = intersect(intID,Indj);
%     allCutElemj = mesh.t(Indj,:);
%     isInterfaceNodej = false(size(node,1), 1);
%     isInterfaceNodej(allCutElemj(:)) = true; % include the vertices of allCutElem
%     tmp = mesh.eLoc(mesh.t_e(Indj,:));
%     tmp = tmp(tmp<0);
%     %isInterfaceNodej(Np-tmp) = true;    % and the cut points and the aux points
%     interfaceNodej = node(isInterfaceNodej,:);
%
%     DT = delaunayTriangulation(interfaceNodej);
%     tetElemj = DT.ConnectivityList;
%     [tetElemj, ortidx, volume] = fixorder3(interfaceNodej, tetElemj);
%
%     localidx2globalidxj = find(isInterfaceNodej); % tetElem points to interfaceNode
%     tetElemj = localidx2globalidxj(tetElemj);  % map to the global index of node
%     bcj = (node(tetElemj(:, 1), :) + node(tetElemj(:, 2), :) ...
%         + node(tetElemj(:, 3), :) + node(tetElemj(:, 4), :))/4.0;
%     idx2cubej = FindElemId(bcj, mesh); % Get the cube index of each tetElem
%     idx2cubeBadj = (idx2cubej<=0);
%     idx2cubej = idx2cubej(~idx2cubeBadj);
%     tetElemj = tetElemj(~idx2cubeBadj,:);
%     [godElemj,godElemjIDa,godElemjIDb] = intersect(idx2cubej,Indj);
%     idx2cubej = idx2cubej(godElemjIDa);
%     tetElemj = tetElemj(godElemjIDa,:);
%
% %     tetElem = tetElem(isCutElem(idx2cube), :); % Just keep the tets in cut elem
% %     volume = volume(isCutElem(idx2cube), :);
% %     idx2cube = idx2cube(isCutElem(idx2cube));  % keep cube idx in cut elements
%
% end
%
% toc

%% Get triangular faces for interrior elements and interface
tic
localFace = [2 3 4; 1 4 3; 1 2 4; 1 3 2];
NT = size(tetElem,1);
tface = zeros(4*NT, 3);
tface2elem = zeros(4*NT, 1);
iface = zeros(4*NT, 3);
iface2elem = zeros(4*NT, 1);
% find the interior tet elements
isIntTet = min(vSign(tetElem),[], 2) == -1; % can not be == -1 as there is sliver
intTet = tetElem(isIntTet,:);
% find the corresponding cube indices
intIdx2cube = idx2cube(isIntTet);
% find triangular interface
T = auxstructure3(intTet);
neighbor = T.neighbor; % if a face is on the boundary, then the neighbor element index is itself
clear T;
tmp = (1:size(intTet, 1))';
ct = 0;
ci = 0;
for i = 1:4
    face = intTet(:, localFace(i,:));
    % find the triangle faces of polyhedron
    % 1. face and its neighbor associated to different cubes, and
    % 2. face is not on the boundary of all cut elements (which are squares)
    isPolyTriFace = ((neighbor(:, i) == tmp) | (intIdx2cube ~= intIdx2cube(neighbor(:,i)))) &...
        (sum(abs(vSign(face)), 2) > 0);% & (sum(abs(vSign(face)), 2) > 0);
    c2 = ct + sum(isPolyTriFace);
    tface((ct+1):c2,:) = face(isPolyTriFace,:);
    tface2elem((ct+1):c2,:) = intIdx2cube(isPolyTriFace); % the indices of the polyhedron
    ct = c2;
    % note that only interior elements are treated

    % find the triangle faces on interface with normal points to exterior
    % 1. face is on the boundary and
    % 2. all vertices are on the interface
    isInterface = (sum(abs(vSign(face)), 2) == 0);

    % add to interface
    c4 = ci + sum(isInterface);
    iface((ci+1):c4,:) = face(isInterface,:); % the interface tri faces.
    iface2elem((ci+1):c4,:) = intIdx2cube(isInterface); % the indices of the polyhedron
    ci = c4;
end
iface((ci+1):end,:) = [];
iface2elem((ci+1):end,:) = [];
c2old = c2;
% plot to check
%trisurf(tface(1:c2,:),node(:,1),node(:,2),node(:,3))
%trisurf(iface(1:ci,:),node(:,1),node(:,2),node(:,3))

% Find the triangular faces for exterior elements
extTet = tetElem(~isIntTet, :);
extIdx2cube = idx2cube(~isIntTet);

T = auxstructure3(extTet);
neighbor = T.neighbor;
clear T;
tmp = (1:size(extTet,1))';
for i = 1:4
    face = extTet(:, localFace(i,:));
    % find the triangle faces of polyhedron
    % 1. face and its neighbor associated to different cubes, and
    % 2. face is not on the boundary of all cut elements (which are squares)
    isPolyTriFace = ((neighbor(:, i) == tmp) | (extIdx2cube ~= extIdx2cube(neighbor(:,i)))) &...
        (sum(abs(vSign(face)), 2) > 0);
    c2 = ct + sum(isPolyTriFace);
    tface((ct+1):c2,:) = face(isPolyTriFace,:);
    tface2elem((ct+1):c2,:) = extIdx2cube(isPolyTriFace);
    % index of exterior polyhedron is append to the end of elem
    ct = c2;
end
toc

% plot to check
%trisurf(tface(c2old+1:c2,:),node(:,1),node(:,2),node(:,3))
%trisurf(iface,node(:,1),node(:,2),node(:,3))

tface((ct+1):end,:) = [];    % remove empty meomory
tface2elem((ct+1):end) = [];

face2elem = tface2elem;
face2elemLoc = [ones(c2old,1);2*ones(c2-c2old,1)]; % face2elemLoc is the same as face location
face = tface;

interfaceData.tface = tface;
interfaceData.tface2elem = tface2elem;
interfaceData.iface = iface;
interfaceData.iface2elem = iface2elem;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% 2.use face and face2elem data structure to compute ife functions

vSign = [mesh.pLoc;zeros(size(mesh.eIntP,1),1)];
node = [mesh.p;mesh.eIntP];
face = interfaceData.tface;
face2elem = interfaceData.tface2elem;
N = size(node, 1); Ndof = N;
NF = size(face,1); NC = size(mesh.t, 1);

NP = max(face2elem);
isCutPoly = false(NP, 1);
isCutPoly(face2elem) = true;
% this is the same as:
% isCutPoly = (mesh.tLoc<0);
% isCutPoly = isCutPoly(1:max(face2elem));
cutPolyIdx = zeros(NP, 1);
isExtPoly = false(NP,1);
isExtPoly(NC+1:NP) = true;

NP = sum(isCutPoly); % this new NP is max(-mesh.tLoc(mesh.tLoc<0));
cutPolyIdx(isCutPoly) = 1:NP;
% this is the same as:
% cutPolyIdx = zeros(size(mesh.tLoc,1), 1);
% cutPolyIdx(mesh.tLoc<0) = -mesh.tLoc(mesh.tLoc<0);
% cutPolyIdx = cutPolyIdxNew(1:max(face2elem));
face2elem = cutPolyIdx(face2elem);
% the old face2elem is the global index of interface elements
% the new one is only the index of interface elements
isExtPoly = isExtPoly(isCutPoly);
%%%clear isCutPoly cutPolyIdx

% prepare new data structure to compute projections
poly2node = sparse(face2elem(:)*ones(1,3), face(:), 1, NP,N);
poly2node = (poly2node > 0);
NV = poly2node*ones(N, 1);
%centroid = poly2node*node./[NV, NV, NV];
%KP = pde.A(centroid);

% generate local IFE basis functions which are not associated with any DoFs
% IFEbasis(:,i,:,:) for minus subdomain (i=1) or plus domain (i=2)
iface = interfaceData.iface;
tface = interfaceData.tface;
iface2elem = interfaceData.iface2elem;
[iface2elemReduce, uniID] = unique(iface2elem);
% length(uniID) should be NP
ifaceReduce = iface(uniID,:);
[IntFnormal,IntFarea,unitNormal,unittgt1,unittgt2] = facenormaltgt(node,ifaceReduce);
[Fnormal,Farea,FunitNormal,Funittgt1,Funittgt2] = facenormaltgt(node,tface);
% IFEbasis contains the associated IFE function on each face
IFEbasis = zeros(size(face2elem,1),3,3);
% for two pieces of an ife function, only the normal vector is different
% on the minus subdomain (1) ife = n; on the plus subdomain (2) ife = n*bm/bp
IFEbasis(:,1,:) = unitNormal(face2elem,:);
IFEbasis(:,2,:) = unittgt1(face2elem,:);
IFEbasis(:,3,:) = unittgt2(face2elem,:);
piece1tmp = (face2elemLoc==1);
IFEbasis(~piece1tmp,1,:) = IFEbasis(~piece1tmp,1,:)*bm/bp;
% generate boundary integral for computing projections
BIFE =  zeros(length(face2elem),3);
% BIFE contains  (beta*grad vi.n)*|f|/3 for i=1,2,3 (vi is the test function)
BIFE(:,1) = sum(squeeze(IFEbasis(:,1,:)).*Fnormal,2)/2;
BIFE(:,2) = sum(squeeze(IFEbasis(:,2,:)).*Fnormal,2)/2;
BIFE(:,3) = sum(squeeze(IFEbasis(:,3,:)).*Fnormal,2)/2;
% devided by 2 is because Fnormal is unitnormal*(2*area)
BIFE(piece1tmp,:) = bm*BIFE(piece1tmp,:)/3.;
BIFE(~piece1tmp,:) = bp*BIFE(~piece1tmp,:)/3;
% generate the matrix for solving IFE projections which is a diagonal matrix
VdiagTmp1 = (bm*PolyVolume(:,1) + bm^2/bp*PolyVolume(:,2)).^(-1);
VdiagTmp2 = (bm*PolyVolume(:,1) + bp*PolyVolume(:,2)).^(-1);
VdiagTmp3 = (bm*PolyVolume(:,1) + bp*PolyVolume(:,2)).^(-1);
Vdiag1 = VdiagTmp1(face2elem);
Vdiag2 = VdiagTmp2(face2elem);
Vdiag3 = VdiagTmp3(face2elem);
Vdiag = zeros(3*length(face2elem),1);
Vdiag(1:3:end-2) = Vdiag1;
Vdiag(2:3:end-1) = Vdiag2;
Vdiag(3:3:end) = Vdiag3;
Mdiag = spdiags(Vdiag,0,3*length(face2elem),3*length(face2elem));
BIFE = reshape(Mdiag*reshape(BIFE',[],1),3,[])';
% this updated new BIFE(i,:) contains (old)BIFE(i,:)*M^(-1) where M is the
% coefficient matrix associated with the element for this face.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% compute projections to IFE spaces and form the matrices
nnz = sum(NV.^2);
iiP = zeros(nnz, 1);
jjP = zeros(nnz, 1);
ssP = zeros(nnz, 1);
indexP = 0;
iiF = zeros(nnz, 1);
jjF = zeros(nnz, 1);
ssF = zeros(nnz, 1);
indexF = 0;
unv = unique(NV);
hF = h(face2elem);
b = zeros(Ndof, 1);

for kk = 1:length(unv) % group polys according to their # of nodes

    tic
    %% generate IFE functions including the IFE basis functions on each interface
    %% elements which are not associated with any DoFs, and projections of gradients
    %% and the IFE functions with matching face average
    nv = unv(kk);
    isCurrentFace = (NV(face2elem) == nv);
    CurrentFace = tface(isCurrentFace,:);
    CNF = size(CurrentFace,1);
    Currentfaceloc = face2elemLoc(isCurrentFace);
    Currentface2elem = face2elem(isCurrentFace);
    CurrentFnormal = FunitNormal(isCurrentFace,:);
    isCurrentElem = false(NP, 1);
    isCurrentElem(face2elem(isCurrentFace)) = true;
    CurrentPolyVolume = PolyVolume(isCurrentElem,:);
    % Current elem to node matrix
    currentPoly2node = poly2node(isCurrentElem, :);
    CNP = size(currentPoly2node, 1);
    currentPolyLocalIdx = zeros(NP, 1);
    currentPolyLocalIdx(isCurrentElem) = 1:CNP;

    % currentElem(i,:) contains the node index of the (current) i-th element
    [I, ~] = find(currentPoly2node');
    currentElem = reshape(I, nv, [])';
    % localIdx(i,j) contains the (current) i-th element having the
    % node(dof=1,2,3,4) on the j-th location (j is the global node index)
    localIdx = sparse(repmat((1:CNP)', 1, nv), currentElem, ones(CNP, 1)*(1:nv), CNP, N);
    clear currentPoly2node

    % Deal with current triangle face cases
    %isCTFace = isCurrentFace;
    tFace = face(isCurrentFace, :);
    NN = sum(isCurrentFace);
    subs1 = currentPolyLocalIdx(face2elem(isCurrentFace));
    subs2 = [ones(NN, 1); 2*ones(NN, 1); 3*ones(NN, 1)];
    val = BIFE(isCurrentFace,:);

    % compute the projection of gradients
    GP1 = zeros(CNP,3,nv);
    % coefficient of n t1 and t2 for the subelement on the subdomain 1
    for m = 1:3
        subs3 = full(localIdx((CurrentFace(:, m) - 1)*CNP + subs1));
        % subs1: new order (range in 1:CNP) of the element index
        % subs2: three components in the vector
        % subs3: the local index (DoF) (w.r.t. element) of the m-th node of this face
        GP1 = GP1 + accumarray([repmat(subs1,3,1), subs2, repmat(subs3, 3, 1)], val(:), [CNP, 3, nv]);
    end
    GP2 = GP1; GP2(:,1,:) = bm/bp*GP2(:,1,:);
    % generate vector basis associated with interface plane
    unitNormalCurrent = unitNormal(isCurrentElem,:);
    unittgt1Current = unittgt1(isCurrentElem,:);
    unittgt2Current = unittgt2(isCurrentElem,:);
    NTCurrent = zeros(CNP,3,3);
    NTCurrent(:,:,1) = unitNormalCurrent;
    NTCurrent(:,:,2) = unittgt1Current;
    NTCurrent(:,:,3) = unittgt2Current;
    GP1 = mytimes(NTCurrent,GP1);
    GP2 = mytimes(NTCurrent,GP2);
    % the updated GP contains three component values of the projection vector
    % generate IFE vectors on each face

    % generate some data structure on faces to compute IFE functions (not gradients)
    GP1tmp = zeros(NP,3,nv); GP2tmp = zeros(NP,3,nv);
    GP1tmp(isCurrentElem,:,:) = GP1; GP2tmp(isCurrentElem,:,:) = GP2;
    Gface = zeros(CNF,3,nv);
    Gface(Currentfaceloc==1,:,:) = GP1tmp(Currentface2elem(Currentfaceloc==1),:,:);
    Gface(Currentfaceloc==2,:,:) = GP2tmp(Currentface2elem(Currentfaceloc==2),:,:);
    clear GP1tmp GP2tmp
    % compute the IFE functions with the constant such that integral on
    % faces equal the virtual element function
    VPtmp1 = zeros(CNP,nv); % sum_i DoFi(u)|f|/3 on each face f
    VPtmp2 = zeros(CNP,nv); % grad*(xf-xm)|f|
    VPtmp3 = zeros(CNP,nv); % the surface area of each element repeated nv times
    FareaCurrent = Farea(isCurrentFace); %|f|/3
    for m = 1:3
        subs3 = full(localIdx((CurrentFace(:, m) - 1)*CNP + subs1));
        VPtmp1 = VPtmp1 + accumarray([subs1, subs3], FareaCurrent(:)/3, [CNP, nv]);
    end
    Xmf = (node(ifaceReduce(Currentface2elem,1),:) + node(ifaceReduce(Currentface2elem,2),:) +...
        node(ifaceReduce(Currentface2elem,3),:))/3; % the Xm point on the approximate interface plane
    Xf = (node(CurrentFace(:,1),:) + node(CurrentFace(:,2),:) +...
        node(CurrentFace(:,3),:))/3; % the centroid on each triangular face
    VPtmp2tmp = ((Xf - Xmf)).*FareaCurrent;
    for i = 1:nv
        tmp = sum(Gface(:,:,i).*VPtmp2tmp,2);
        VPtmp2(:,i) = accumarray(subs1, tmp,[CNP, 1]);
    end
    VPtmp3 = repmat(accumarray(subs1, FareaCurrent(:)),[1, nv]);
    IFEConst = (VPtmp1 - VPtmp2)./VPtmp3;
    clear currentPolyLocalIdx localIdx
    clear subs1 subs2 subs3
    toc

    %% form the stiffness matrices and stabilization matrices
    CurrentFace2DoF = zeros(NP,nv);
    CurrentFace2DoF(isCurrentElem,:) = currentElem;
    CurrentFace2DoF = CurrentFace2DoF(Currentface2elem,:);
    FunitNormalCurrent = FunitNormal(isCurrentFace,:);
    Funittgt1Current = Funittgt1(isCurrentFace,:);
    Funittgt2Current = Funittgt2(isCurrentFace,:);
    FNTCurrent = zeros(CNF,3,3);
    FNTCurrent(:,1,:) = Funittgt1Current;
    FNTCurrent(:,2,:) = Funittgt2Current;
    FNTCurrent(:,3,:) = FunitNormalCurrent;
    GfaceP = zeros(size(Gface));
    for n = 1:nv
        GfaceP(:,:,n) =  Gface(:,:,n) - sum(Gface(:,:,n).*CurrentFnormal,2).*CurrentFnormal;
    end
    GfaceP = mytimes(FNTCurrent,GfaceP);
    GfaceP = GfaceP(:,1:2,:);
    % compute the face gradients based on DoFs
    % generate vector basis associated with each face
    Fnode  = zeros(CNF,3,3);
    for l = 1:3
        Fnode(:,:,l) = node(CurrentFace(:,l),:);
    end
    FacePCoord = mytimes(FNTCurrent,Fnode);
    gradLambdatmp = zeros(CNF,2,3);  % barycentric coordinates
    gradLambda = zeros(CNF,2,nv);
    signtmp = ones(CNF,2); signtmp(:,1) = -1;
    gradLambdatmp(:,:,1) = (FacePCoord(:,[2,1],3) - FacePCoord(:,[2,1],2)).*signtmp./(2*Farea(isCurrentFace));
    gradLambdatmp(:,:,2) = (FacePCoord(:,[2,1],1) - FacePCoord(:,[2,1],3)).*signtmp./(2*Farea(isCurrentFace));
    gradLambdatmp(:,:,3) = -gradLambdatmp(:,:,1) - gradLambdatmp(:,:,2);
    for i = 1:nv
        for j = 1:3
            IDij = (CurrentFace(:,j) == CurrentFace2DoF(:,i));
            gradLambda(IDij,:,i) = gradLambdatmp(IDij,:,j);
        end
    end

%     for n = 1:nv
%         for m = 1:nv
%             Xn = node(currentElem(:, n), :) - centroid(isCurrentElem,:);
%             IminusP(:, n, m) = - 1/nv - dot(Xn, BP(:, :, m), 2)./volume(isCurrentElem); % stabilization
%             if(n == m)
%                 IminusP(:, n, m) = 1 + IminusP(:, n, m);
%             end
%         end
%     end

    for n = 1:nv
        for m = 1:nv
            iiP(indexP+1:indexP + CNP) = currentElem(:, n);
            jjP(indexP+1:indexP + CNP) = currentElem(:, m);
            ssP(indexP+1:indexP + CNP) = bm*dot(GP1(:,:,n), GP1(:,:,m),2).*CurrentPolyVolume(:,1)+...
                bp*dot(GP2(:,:,n), GP2(:,:,m),2).*CurrentPolyVolume(:,2);
            indexP = indexP + CNP;

            iiF(indexF+1:indexF + CNF) = CurrentFace2DoF(:, n);
            jjF(indexF+1:indexF + CNF) = CurrentFace2DoF(:, m);
            ssF(indexF+1:indexF + CNF) = hF(isCurrentFace).*...
                dot((gradLambda(:,:,m) - GfaceP(:,:,m)),(gradLambda(:,:,n) - GfaceP(:,:,n)),2).*...
                Farea(isCurrentFace);
            indexF = indexF + CNF;
        end
    end

    %% compute the RHS with only evaluating the integral at the centroid
    %% approach 1 (will give the optimal order but slightly increase the error)
%     currentElemSign = vSign(currentElem);
%     currentElemSign1 = currentElemSign<=0; nv1 = sum(currentElemSign1,2);
%     currentElemSign2 = currentElemSign>=0; nv2 = sum(currentElemSign2,2);
% %     nodeX = node(:,1); currentElemNodeX = nodeX(currentElem); clear nodeX
% %     nodeY = node(:,2); currentElemNodeY = nodeY(currentElem); clear nodeY
% %     nodeZ = node(:,1); currentElemNodeZ = nodeZ(currentElem); clear nodeZ
%     currentElemNodeX = reshape(node(currentElem',1),size(currentElem,2),[])';
%     currentElemNodeY = reshape(node(currentElem',2),size(currentElem,2),[])';
%     currentElemNodeZ = reshape(node(currentElem',3),size(currentElem,2),[])';
%     centroid1 = [sum(currentElemNodeX.*currentElemSign1,2)./nv1,...
%         sum(currentElemNodeY.*currentElemSign1,2)./nv1, ...
%         sum(currentElemNodeZ.*currentElemSign1,2)./nv1];
%     centroid2 = [sum(currentElemNodeX.*currentElemSign2,2)./nv2,...
%         sum(currentElemNodeY.*currentElemSign2,2)./nv2, ...
%         sum(currentElemNodeZ.*currentElemSign2,2)./nv2];
%     ft1 = CurrentPolyVolume(:,1).*pde.f(centroid1(:,1),centroid1(:,2),centroid1(:,3));
%     ft2 = CurrentPolyVolume(:,2).*pde.f(centroid2(:,1),centroid2(:,2),centroid2(:,3));
%
%     Xm = (node(ifaceReduce(isCurrentElem,1),:) + node(ifaceReduce(isCurrentElem,2),:) +...
%         node(ifaceReduce(isCurrentElem,3),:))/3;
%     ifeEvalcen1 = zeros(CNP,nv); ifeEvalcen2 = zeros(CNP,nv);
%     for n = 1:nv
%         ifeEvalcen1(:,n) = sum(GP1(:,:,n).*(centroid1-Xm),2) + IFEConst(:,n);
%         ifeEvalcen2(:,n) = sum(GP2(:,:,n).*(centroid2-Xm),2) + IFEConst(:,n);
%     end
%
%     Currentb = ft1.*ifeEvalcen1 + ft2.*ifeEvalcen2;
%     b = b + accumarray(currentElem(:), Currentb(:), [Ndof, 1]);

    %% approach 2
    idx2cubeNew = -mesh.tLoc(idx2cube);
    TetID = isCurrentElem(idx2cubeNew);
    currentvolume = volume(TetID);
    ElemDoF = zeros(sum(TetID),nv);
    ElemDoF(isCurrentElem,:) = currentElem;
    ElemDoF = ElemDoF(idx2cubeNew(TetID),:);
    X1 = node(tetElem(TetID,1),:);
    X2 = node(tetElem(TetID,2),:);
    X3 = node(tetElem(TetID,3),:);
    X4 = node(tetElem(TetID,4),:);
    ng = 1;
    [gx, gy, gz] = gaussPtetra(X1, X2, X3, X4, ng);
    gw = gaussWtetra(ng);
    Xm = (node(ifaceReduce(isCurrentElem,1),:) + node(ifaceReduce(isCurrentElem,2),:) +...
        node(ifaceReduce(isCurrentElem,3),:))/3;

    Bas1 = zeros(size(isCurrentElem,1),3,nv); Bas2 = Bas1;
    Bas1(isCurrentElem,:,:) = GP1;
    Bas2(isCurrentElem,:,:) = GP2;
    Bas1 = Bas1(idx2cubeNew(TetID),:,:); Bas2 = Bas2(idx2cubeNew(TetID),:,:);
    Bas = zeros(sum(TetID),3,nv);
    %%%
    IFEc0 = zeros(size(isCurrentElem,1),nv);
    IFEc0(isCurrentElem,:) = IFEConst;
    IFEc0 = IFEc0(idx2cubeNew(TetID),:);
    %%%
    Xmtet = zeros(size(isCurrentElem,1),3);
    Xmtet(isCurrentElem,:) = Xm;
    Xmtet = Xmtet(idx2cubeNew(TetID),:);
    %%%
    currentTetDoF = zeros(size(isCurrentElem,1),nv);
    currentTetDoF(isCurrentElem,:) = currentElem;
    currentTetDoF = currentTetDoF(idx2cubeNew(TetID),:);
    %%%
    piecetmp = tetElemLoc(TetID);
    Bas(piecetmp==1,:,:) = Bas1(piecetmp==1,:,:);
    Bas(piecetmp==2,:,:) = Bas2(piecetmp==2,:,:);

    %%%%%%%%%%%%%%%%%%%%%%%%%%%
    % mass matrix

%     for i = 1:nv
%         for j = 1:nv
%             Basi = squeeze(Bas(:,:,i)); IFEc0i = IFEc0(:,i);
%             Basj = squeeze(Bas(:,:,j)); IFEc0j = IFEc0(:,j);
%             uhi = (gx-Xmtet(:,1)).*Basi(:,1) + (gy-Xmtet(:,2)).*Basi(:,2) +...
%                 (gz-Xmtet(:,3)).*Basi(:,3) + IFEc0i;
%             uhj = (gx-Xmtet(:,1)).*Basj(:,1) + (gy-Xmtet(:,2)).*Basj(:,2) +...
%                 (gz-Xmtet(:,3)).*Basj(:,3) + IFEc0j;
%             Currentb(:,i) = sum(gw*ft.*uhi.*uhj,2).*currentvolume;
%             bii(count+1:count+cnt) = currentTetDoF(:,i);
%             count = count+cnt;
%         end
%     end

    %%%%%%%%%%%%%%%%%%%%%%%%%%%

    ft = pde.f(gx,gy,gz);
    Currentb = zeros(size(ft));
    cnt = size(ft,1);
    bii = nv*cnt;
    count = 0 ;
    for i = 1:nv
        Basi = squeeze(Bas(:,:,i)); IFEc0i = IFEc0(:,i);
        uhi = (gx-Xmtet(:,1)).*Basi(:,1) + (gy-Xmtet(:,2)).*Basi(:,2) +...
            (gz-Xmtet(:,3)).*Basi(:,3) + IFEc0i;
        Currentb(:,i) = sum(gw*ft.*uhi,2).*currentvolume;
        bii(count+1:count+cnt) = currentTetDoF(:,i);
        count = count+cnt;
    end
    b = b + sparse(bii,1,reshape(Currentb,[],1),Ndof,1);
    %b = b + accumarray([idx2cubeNew(TetID),currentTetDoF(:)], Currentb(:), [Ndof, 1]);

end
K = sparse(iiP, jjP, ssP, Ndof, Ndof);
S = sparse(iiF, jjF, ssF, Ndof, Ndof);


function solution = solveCobraMILP(MILPproblem, varargin)
% Solves constraint-based MILP problems
% The solver is defined in the CBT_MILP_SOLVER global variable
% (set using `changeCobraSolver`). Solvers currently available are
% 'tomlab_cplex' and 'glpk'
%
% USAGE:
%
%    solution = solveCobraMILP(MILPproblem, parameters)
%
% INPUT:
%    MILPproblem:     Structure containing the following fields describing the LP problem to be solved
%
%                       * .A - LHS matrix
%                       * .b - RHS vector
%                       * .c - Objective coeff vector
%                       * .lb - Lower bound vector
%                       * .ub - Upper bound vector
%                       * .osense - Objective sense (-1 max, +1 min)
%                       * .csense - Constraint senses, a string containting the constraint sense for
%                         each row in A ('E', equality, 'G' greater than, 'L' less than).
%                       * .vartype - Variable types ('C' continuous, 'I' integer, 'B' binary)
%                       * .x0 - Initial solution
%
% Optional parameters can be entered using parameters structure or as
% parameter followed by parameter value: i.e. ,'printLevel', 3)
% Setting `parameters` = 'default' uses default setting set in
% `getCobraSolverParameters`.
%
% OPTIONAL INPUTS:
%    varargin:        Additional parameters either as parameter struct, or as
%                     parameter/value pairs. A combination is possible, if
%                     the parameter struct is either at the beginning or the
%                     end of the optional input.
%                     All fields of the struct which are not COBRA parameters
%                     (see `getCobraSolverParamsOptionsForType`) for this
%                     problem type will be passed on to the solver in a
%                     solver specific manner. Some optional parameters which
%                     can be passed to the function as parameter value pairs,
%                     or as part of the options struct are listed below:
%
%    timeLimit:       Global solver time limit
%    intTol:          Integrality tolerance
%    relMipGapTol:    Relative MIP gap tolerance
%    logFile:         Log file (for CPLEX)
%    printLevel:      Printing level
%
%                       * 0 - Silent (Default)
%                       * 1 - Warnings and Errors
%                       * 2 - Summary information
%                       * 3 - More detailed information
%                       * > 10 - Pause statements, and maximal printing (debug mode)
%    saveInput:       Saves LPproblem to filename specified in field.
%                     i.e. parameters.saveInput = 'LPproblem.mat';
%
%
% OUTPUT:
%    solution:        Structure containing the following fields describing a MILP solution
%
%                       * .cont:        Continuous solution
%                       * .int:         Integer solution
%                       * .full:        Full MILP solution vector
%                       * .obj:         Objective value
%                       * .solver:      Solver used to solve MILP problem
%                       * .stat:        Solver status in standardized form (see below)
%
%                         * 1 - Optimal solution found
%                         * 2 - Unbounded solution
%                         * 0 - Infeasible MILP
%                         * -1 - No integer solution exists
%                         * 3 - Other problem (time limit etc, but integer solution exists)
%                       * .origStat:    Original status returned by the specific solver
%                       * .time:        Solve time in seconds
%
% .. Authors:
%       - Markus Herrgard 1/23/07
%       - Tim Harrington  05/18/12 Added support for the Gurobi 5.0 solver
%       - Ronan (16/07/2013) default MPS parameters are no longer global variables
%       - Meiyappan Lakshmanan  11/14/14 Added support for the cplex_direct solver
%       - cplex_direct solver accesible through CPLEX m-file and CPLEX C-interface
%       - Thomas Pfau (12/11/2015) Added support for ibm_cplex (the IBM Matlab
%       interface) to the solvers.

[cobraParams,solverParams] = parseSolverParameters('MILP',varargin{:}); % get the solver parameters

solver = cobraParams.solver;

% Save Input if selected
if ~isempty(cobraParams.saveInput)
    fileName = cobraParams.saveInput;
    if ~find(regexp(fileName, '.mat'))
        fileName = [fileName '.mat'];
    end
    display(['Saving MILPproblem in ' fileName]);
    save(fileName, 'MILPproblem')
end

% Defaults in case the solver does not return anything
x = [];
xInt = [];
xCont = [];
f = [];

if ~isfield(MILPproblem, 'x0')
    MILPproblem.x0 = [];
end

[A, b, c, lb, ub, csense, osense, vartype, x0] = ...
    deal(MILPproblem.A, MILPproblem.b, MILPproblem.c, MILPproblem.lb, MILPproblem.ub, ...
    MILPproblem.csense, MILPproblem.osense, MILPproblem.vartype, MILPproblem.x0);

if any(~(vartype == 'C' | vartype == 'B' | vartype == 'I'))
    display('vartype not C or B or I:  Assuming C');
    vartype(vartype ~= 'C' & vartype ~= 'I'& vartype ~= 'B') = 'C';
end

t_start = clock;
switch solver

    case 'glpk'
        %% glpk

        % Set up problem
        if (isempty(csense))
            clear csense
            csense(1:length(b), 1) = 'S';
        else
            csense(csense == 'L') = 'U';
            csense(csense == 'G') = 'L';
            csense(csense == 'E') = 'S';
            csense = columnVector(csense);
        end

        if ~isfield(solverParams,'msglev')
            solverParams.msglev = cobraParams.printLevel;
        end
        if ~isfield(solverParams,'tmlim')
            solverParams.tmlim = cobraParams.timeLimit;
        end

        % whos csense vartype
        csense = char(csense);
        vartype = char(vartype);
        % whos csense vartype

        % Solve problem
        [x, f, stat, extra] = glpk(c, A, b, lb, ub, csense, vartype, osense, solverParams);
        % Handle solution status reports
        if (stat == 5)
            solStat = 1;  % optimal
        elseif(stat == 6)
            solStat = 2;  % unbounded
        elseif(stat == 4)
            solStat = 0;  % infeasible

        elseif(stat == 171)
            solStat = 1;  % Opt integer within tolerance
        elseif(stat == 173)
            solStat = 0;  % Integer infeas
        elseif(stat == 184)
            solStat = 2;  % Unbounded
        elseif(stat == 172)
            solStat = 3;  % Other problem, but integer solution exists
        else
            solStat = -1;  % No integer solution exists
        end

    case 'cplex_direct'
        %% cplex_direct

        % Set up problem
        b = full(b);
        [m_lin, n] = size(MILPproblem.A);
        if ~isempty(csense)
            Aineq = [MILPproblem.A(csense == 'L', :); - MILPproblem.A(csense == 'G', :)];
            bineq = [b(csense == 'L', :); - b(csense == 'G', :)];
            %        min      c*x
            %        st.      Aineq*x <= bineq
            %                 Aeq*x    = beq
            %                 lb <= x <= ub
            A = MILPproblem.A(csense == 'E', :);
            b = b(csense == 'E', 1);
            [x, f, exitflag, output] = cplexmilp(c, Aineq, bineq, A, b, [], [], [], lb, ub, vartype');

            % primal
            solution.obj = osense * f;
            solution.full = x;
            % this is the dual to the equality constraints but it's not the chemical potential
            %             solution.dual=lambda.eqlin;
        else
            Aineq = [];
            bineq = [];
            [x, f, exitflag, output] = cplexmilp(c, Aineq, bineq, MILPproblem.A, b, lb, ub, vartype);
            solution.obj = osense * f;
            solution.full = x;
            % this is the dual to the equality constraints but it's not the chemical potential
            solution.dual = sparse(size(MILPproblem.A, 1), 1);
            %             solution.dual(csense == 'E')=lambda.eqlin;
            % this is the dual to the inequality constraints but it's not the chemical potential
            %             solution.dual(csense == 'L')=lambda.ineqlin(1:nnz(csense == 'L'),1);
            %             solution.dual(csense == 'G')=lambda.ineqlin(nnz(csense == 'L')+1:end,1);
        end
        solution.nInfeas = [];
        solution.sumInfeas = [];
        solution.origStat = output.cplexstatus;

        Inform = solution.origStat;
        stat = Inform;
        if (stat == 101 || stat == 102)
            solStat = 1;  % Opt integer within tolerance
        elseif(stat == 103)
            solStat = 0;  % Integer infeas
        elseif(stat == 118 || stat == 119)
            solStat = 2;  % Unbounded
        elseif(stat == 106 || stat == 106 || stat == 108 || stat == 110 || stat == 112 || stat == 114 || stat == 117)
            solStat = -1;  % No integer solution exists
        else
            solStat = 3;  % Other problem, but integer solution exists
        end

    case 'gurobi_mex'
        % Free academic licenses for the Gurobi solver can be obtained from
        % http://www.gurobi.com/html/academic.html
        %
        % The code below uses Gurobi Mex to interface with Gurobi. It can be downloaded from
        % http://www.convexoptimization.com/wikimization/index.php/Gurobi_Mex:_A_MATLAB_interface_for_Gurobi

        opts = solverParams;
        if cobraParams.printLevel == 0
            % Version v1.10 of Gurobi Mex has a minor bug. For complete silence
            % Remove Line 736 of gurobi_mex.c: mexPrintf("\n");
            if ~isfield(opts,'Display')
                opts.Display = 0;
            end
            if ~isfield(opts,'DisplayInterval')
                opts.DisplayInterval = 0;
            end
        else
            if ~isfield(opts,'Display')
                opts.Display = 1;
            end
        end

        if ~isfield(opts,'TimeLimit')
            opts.TimeLimit = solverParams.timeLimit;
        end
        if ~isfield(opts,'MIPGap')
            opts.MIPGap = solverParams.relMipGapTol;
        end
        if ~isfield(opts,'IntFeasTol')
            opts.IntFeasTol = solverParams.intTol;
        end
        if ~isfield(opts,'FeasibilityTol')
            % minimum intTol for gurobi = 1e-9
            opts.FeasibilityTol = max(solverParams.feasTol,1e-9);
        end
        if ~isfield(opts,'OptimalityTol')
            opts.OptimalityTol = solverParams.optTol;
        end

        if (isempty(csense))
            clear csense
            csense(1:length(b),1) = '=';
        else
            csense(csense == 'L') = '<';
            csense(csense == 'G') = '>';
            csense(csense == 'E') = '=';
            csense = csense(:);
        end
        % gurobi_mex doesn't automatically cast logicals to doubles
        c = double(c);
        [x,f,stat,output] = gurobi_mex(c,osense,sparse(A),b, ...
            csense,lb,ub,vartype,opts);
        if stat == 2
            solStat = 1; % Optimal solutuion found
        elseif stat == 3
            solStat = 0; % Infeasible
        elseif stat == 5
            solStat = 2; % Unbounded
        elseif stat == 4
            solStat = 0; % Gurobi reports infeasible *or* unbounded
        else
            solStat = -1; % Solution not optimal or solver problem
        end

    case 'ibm_cplex'
        % Free academic licenses for the IBM CPLEX solver can be obtained from
        % https://www.ibm.com/developerworks/community/blogs/jfp/entry/CPLEX_Is_Free_For_Students?lang=en
        cplexlp = buildCplexProblemFromCOBRAStruct(MILPproblem);
        [cplexlp, logFile, logToFile] = setCplexParametersForProblem(cplexlp,cobraParams,solverParams,'MILP');

        % Solve problem
        Result = cplexlp.solve();

        if logToFile
            % Close the output file
            fclose(logFile);
        end

        % Get results
        stat = Result.status;
        if (stat == 101 || stat == 102 || stat == 1)
            solStat = 1; % Opt integer within tolerance
            % Return solution if problem is feasible, bounded and optimal
            x = Result.x;
            f = Result.objval;
        elseif (stat == 103 || stat == 3)
            solStat = 0; % Integer infeas
        elseif (stat == 118 || stat == 119 || stat == 2)
            solStat = 2; % Unbounded
        elseif (stat == 106 || stat == 106 || stat == 108 || stat == 110 || stat == 112 || stat == 114 || stat == 117)
            solStat = -1; % No integer solution exists
        else
            solStat = 3; % Other problem, but integer solution exists
        end
        if exist([pwd filesep 'clone1.log'],'file')
            delete('clone1.log')
        end
    case 'gurobi'
        %% gurobi 5
        % Free academic licenses for the Gurobi solver can be obtained from
        % http://www.gurobi.com/html/academic.html
        MILPproblem.A = deal(sparse(MILPproblem.A));

        if cobraParams.printLevel == 0
            params.OutputFlag = 0;
            params.DisplayInterval = 1;
        else
            params.OutputFlag = 1;
            params.DisplayInterval = 5;
        end

        %return solution when time limit is reached and save the log file
        if ~isempty(cobraParams.logFile)
            params.LogFile = cobraParams.logFile;
        end
        params.TimeLimit = cobraParams.timeLimit;

        % set tolerances
        params.MIPGap = cobraParams.relMipGapTol;
        if cobraParams.intTol <= 1e-09
            params.IntFeasTol = 1e-09;
        else
            params.IntFeasTol = cobraParams.intTol;
        end
        params.FeasibilityTol = cobraParams.feasTol;
        params.OptimalityTol = cobraParams.optTol;

        if (isempty(csense))
            clear csense
            csense(1:length(b),1) = '=';
        else
            csense(csense == 'L') = '<';
            csense(csense == 'G') = '>';
            csense(csense == 'E') = '=';
            MILPproblem.csense = csense(:);
        end

        if osense == -1
            MILPproblem.osense = 'max';
        else
            MILPproblem.osense = 'min';
        end

        % overwrite default params with directParams
        fieldNames = fieldnames(solverParams);
        for i = 1:size(fieldNames,1)
            params.(fieldNames{i}) = solverParams.(fieldNames{i});
        end

        MILPproblem.vtype = vartype;
        MILPproblem.modelsense = MILPproblem.osense;
        [MILPproblem.A,MILPproblem.rhs,MILPproblem.obj,MILPproblem.sense] = deal(sparse(MILPproblem.A),MILPproblem.b,double(MILPproblem.c),MILPproblem.csense);
        if ~isempty(x0)
            MILPproblem.start = x0;
        end
        resultgurobi = gurobi(MILPproblem,params);

        stat = resultgurobi.status;
        if strcmp(resultgurobi.status,'OPTIMAL')
            solStat = 1; % Optimal solution found
            [x,f] = deal(resultgurobi.x,resultgurobi.objval);
        elseif strcmp(resultgurobi.status,'INFEASIBLE')
            solStat = 0; % Infeasible
        elseif strcmp(resultgurobi.status,'UNBOUNDED')
            solStat = 2; % Unbounded
        elseif strcmp(resultgurobi.status,'INF_OR_UNBD')
            solStat = 0; % Gurobi reports infeasible *or* unbounded
        elseif strcmp(resultgurobi.status,'TIME_LIMIT')
            solStat = 3; % Time limit reached
            warning('Time limit reached, solution might not be optimal (gurobi)')
            try
                [x,f] = deal(resultgurobi.x,resultgurobi.objval);
            catch
                %x and f could not be assigned, as there is no solution
                %yet
            end
        else
            solStat = -1; % Solution not optimal or solver problem
        end

    case 'tomlab_cplex'
        %% CPLEX through tomlab
        if (~isempty(csense))
            b_L(csense == 'E') = b(csense == 'E');
            b_U(csense == 'E') = b(csense == 'E');
            b_L(csense == 'G') = b(csense == 'G');
            b_U(csense == 'G') = inf;
            b_L(csense == 'L') = -inf;
            b_U(csense == 'L') = b(csense == 'L');
        elseif isfield(MILPproblem, 'b_L') && isfield(MILPproblem, 'b_U')
            b_L = MILPproblem.b_L;
            b_U = MILPproblem.b_U;
        else
            b_L = b;
            b_U = b;
        end
        intVars = (vartype == 'B') | (vartype == 'I');
        % intVars
        % pause;
        tomlabProblem = mipAssign(osense*c,A,b_L,b_U,lb,ub,x0,'CobraMILP',[],[],intVars);

        % Set parameters for CPLEX
        tomlabProblem.MIP.cpxControl.EPINT = cobraParams.intTol;
        tomlabProblem.MIP.cpxControl.EPGAP = cobraParams.relMipGapTol;
        tomlabProblem.MIP.cpxControl.TILIM = cobraParams.timeLimit;
        tomlabProblem.CPLEX.LogFile = cobraParams.logFile;
        tomlabProblem.PriLev = cobraParams.printLevel;
        tomlabProblem.MIP.cpxControl.THREADS = 1; % by default use only one thread

        % Strict numerical tolerances
        tomlabProblem.MIP.cpxControl.EPRHS = cobraParams.feasTol;
        tomlabProblem.MIP.cpxControl.EPOPT = cobraParams.optTol;
        tomlabProblem.MIP.cpxControl.EPAGAP = cobraParams.absMipGapTol;

        %Now, replace anything that is in the solver Specific field.
        tomlabProblem.cpxControl = updateStructData(tomlabProblem.MIP.cpxControl,solverParams);

        % Set initial solution
        tomlabProblem.MIP.xIP = x0;

        % Set up callback to print out intermediate solutions
        % only set this up if you know that you actually need these
        % results.  Otherwise do not specify intSolInd and contSolInd
        global cobraIntSolInd;
        global cobraContSolInd;
        if(~isfield(MILPproblem, 'intSolInd'))
            MILPproblem.intSolInd = [];
        else
            tomlabProblem.MIP.callback(14) = 1;
        end
        cobraIntSolInd = MILPproblem.intSolInd;
        if(~isfield(MILPproblem, 'contSolInd'))
            MILPproblem.contSolInd = [];
        end
        cobraContSolInd = MILPproblem.contSolInd;
        tomlabProblem.MIP.callbacks = [];
        tomlabProblem.PriLevOpt = 0;

        % Solve problem
        Result = tomRun('cplex', tomlabProblem);

        % Get results
        x = Result.x_k;
        f = osense*Result.f_k;
        stat = Result.Inform;
        if (stat == 101 || stat == 102)
            solStat = 1; % Opt integer within tolerance
        elseif (stat == 103)
            solStat = 0; % Integer infeas
        elseif (stat == 118 || stat == 119)
            solStat = 2; % Unbounded
        elseif (stat == 106 || stat == 106 || stat == 108 || stat == 110 || stat == 112 || stat == 114 || stat == 117)
            solStat = -1; % No integer solution exists
        else
            solStat = 3; % Other problem, but integer solution exists
        end
    case 'mps'
        fprintf(' > The interface to ''mps'' from solveCobraMILP will not be supported anymore.\n -> Use >> writeCbModel(model, ''mps'');\n');
        % temporary legacy support
        solverParams = updateStructData(cobraParams,solverParams);
        writeLPProblem(MILPproblem, 'problemName','COBRAMILPProblem','fileName','MILP.mps','solverParams',solverParams);
        return
    otherwise
        if isempty(solver)
            error('There is no solver for MILP problems available');
        else
            error(['Unknown solver: ' solver]);
        end
end
t = etime(clock, t_start);

%% Store results
if ~strcmp(solver,'mps')
    if (~isempty(x))
        % xInt = x(MILPproblem.intSolInd);
        % xCont = x(MILPproblem.contSolInd);
        xInt = x(vartype == 'B' | vartype == 'I');
        xCont = x(vartype == 'C');
    end

    solution.cont = xCont;
    solution.int = xInt;
    solution.obj = f;
    solution.solver = solver;
    solution.stat = solStat;
    solution.origStat = stat;
    solution.time = t;
    solution.full = x;
    if(isfield(MILPproblem, 'intSolInd'))
        solution.intInd = MILPproblem.intSolInd;
    end
end

%% Redirection function such that cplex redirects its output to the defined outputfile.
    function redirect(l)
        % Write the line of log output
        fprintf(outputfile, '%s\n', l);
    end

end


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
% .jm          - Jacobian Modulation. If set, intensity (Jacobian)
%                deformations are included in the model. If zero,
%                intensity deformations are not considered.
% .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.
%
%__________________________________________________________________________
% Copyright (C) 2003-2011 Wellcome Trust Centre for Neuroimaging

% Jesper Andersson
% $Id: spm_uw_estimate.m 6824 2016-06-27 16:19:51Z guillaume $

% 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 centered on the voxel-centers, to a grid whith a
%    different sampling frequency.
% 4. Inclusion of a regularisation term. The speed-up has facilitated
%    using more basis-functions, which makes it neccessary 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.

SVNid = '$Rev: 6824 $';

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

%-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 1 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,'jm'),        defpar.jm = ud.jm; 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, chosing 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   = prod(size(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));
   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(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);

%
% 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])*inv(diag([ss 1]))*eye(4);
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,:));
dBy        = spm_dctmtx(P(1).dim(2),ds.order(2),y(1,:,1),'diff');
xyz        = [x(:) y(:) z(:) ones(length(x(:)),1)]; clear x y z;
def        = zeros(size(xyz,1),nof);
ddefa      = 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,:)';
   c = spm_bsplinc(ds.sfP,ds.hold);
   ds.sfield = spm_bsplins(c,txyz(:,1),txyz(:,2),txyz(:,3),ds.hold);
   ds.sfield = ds.sfield(:);
   clear c 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,ddefa,P,xyz,M,msk,beta,dispP);
   AtA       = build_AtA(ref,dx,dy,dz,Bx,By,Bz,dBy,ds,P);
   [Aty,yty] = build_Aty(ref,dx,dy,dz,Bx,By,Bz,dBy,ds,P,xyz,beta,sf,def,ddefa,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;
      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_get_def(Bx, By,Bz,beta((i-1)*prod(ds.order)+1:i*prod(ds.order)));
         ddefa(:,i) = spm_get_def(Bx,dBy,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
   cleanup(P,ds)
   spm_uw_show('FinTot');
   fprintf('Unwarp terminated abnormally.\n');
   rethrow(lasterror);
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 indicies for first order effect.
% sot       : nx2 matrix of indicies 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(inv(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 = inv(P(end).mat) * P(1).mat;
else
   warning(sprintf('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),prod(size(fot)) + size(sot,1));
for i=1:length(P)
   p = [1 1 1 180/pi 180/pi 180/pi zeros(1,6)] .*...
      spm_imatrix(inv(P(i).mat) * P(1).mat);
   for j=1:prod(size(fot))
      q(i,j) = p(fot(j)) - ep(fot(j));
   end
   for j=1:size(sot,1)
      q(i,prod(size(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;

if regorder == 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)) );
elseif regorder == 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)) );
elseif regorder == 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)) );
elseif regorder == 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)) );
else
   error('Invalid order of regularisation');
end
return
%_______________________________________________________________________

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

if ds.jm, spm_uw_show('StartAtA',10);
else spm_uw_show('StartAtA',6); end

nof   = prod(size(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);

if ds.jm
   AjtAj = uwAtA1(ref.*ref,Bx,dBy,Bz); spm_uw_show('NewAtA',7);
   AxtAj = uwAtA2( dx.*ref,Bx, By,Bz,Bx,dBy,Bz); spm_uw_show('NewAtA',8);
   AytAj = uwAtA2( dy.*ref,Bx, By,Bz,Bx,dBy,Bz); spm_uw_show('NewAtA',9);
   AztAj = uwAtA2( dz.*ref,Bx, By,Bz,Bx,dBy,Bz); spm_uw_show('NewAtA',10);
end

R = zeros(length(P),3);
for i=1:length(P)
   tmp    = inv(P(i).mat\P(1).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));

if ds.jm
   tmp  = [ds.q.*kron(tmp,ones(length(P),1))];
   AtA  = AtA + kron(tmp'*tmp,AjtAj) +...
             2*(kron(tmp1'*tmp,AxtAj) + kron(tmp2'*tmp,AytAj) + kron(tmp3'*tmp,AztAj));
end
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 AtA = uwAtA2(y,Bx1,By1,Bz1,Bx2,By2,Bz2)
% Calculating cross-term of diagonal block of AtA
% when A is a sum of the type A1+A2 and where both
% A1 and A2 are possible to express as a kronecker
% product of lesser matrices.

[nx,mx1] = size(Bx1); [nx,mx2] = size(Bx2);
[ny,my1] = size(By1); [ny,my2] = size(By2);
[nz,mz1] = size(Bz1); [nz,mz2] = size(Bz2);
AtA      = zeros(mx1*my1*mz1,mx2*my2*mz2);
for sl =1:nz
   tmp  = reshape(y((sl-1)*nx*ny+1:sl*nx*ny),nx,ny);
   spm_krutil(Bz1(sl,:)'*Bz2(sl,:),spm_krutil(tmp,Bx1,By1,Bx2,By2),AtA);
%   AtA  = AtA + kron(Bz1(sl,:)'*Bz2(sl,:),spm_krutil(tmp,Bx1,By1,Bx2,By2));
end
return
%_______________________________________________________________________

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

nof = prod(size(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))
      [idef,jac] = spm_get_image_def(scan,ds,def,ddefa);
      txyz(:,2)  = txyz(:,2)+idef;
   end;
   c    = spm_bsplinc(P(scan),ds.hold);
   y    = sf(scan) * spm_bsplins(c,txyz(:,1),txyz(:,2),txyz(:,3),ds.hold);
   if ds.jm && ~(all(beta == 0) && isempty(ds.sfield))
      y = y .* jac;
   end;

   y_diff       = (ref - y).*msk;
   indx         = find(isnan(y));
   y_diff(indx) = 0;
   iTcol        = inv(T(1:3,1:3));
   tmpAty       = spm_get_def(Bx',By',Bz',([dx dy dz]*iTcol(:,2)).*y_diff);
   if ds.jm ~= 0
      tmpAty = tmpAty + spm_get_def(Bx',dBy',Bz',ref.*y_diff);
   end
   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,ddefa,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,ddefa,P,xyz,msk);
end

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))
      [idef,jac] = spm_get_image_def(i,ds,def,ddefa);
      txyz(:,2)  = txyz(:,2)+idef;
   end;
   c     = spm_bsplinc(P(i),ds.hold);
   f     = spm_bsplins(c,txyz(:,1),txyz(:,2),txyz(:,3),ds.hold);
   if ds.jm ~= 0 && exist('jac')==1
      f = f .* jac;
   end
   indx  = find(~isnan(f));
   sf(i) = 1 / (sum(f(indx).*msk(indx)) / sum(msk(indx)));
   ref   = ref + f;

   if rem
      indx   = find(isnan(f)); f(indx) = 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;
c              = spm_bsplinc(ref,ds.hold);
txyz           = xyz*M(1:3,:)';
[ref,dx,dy,dz] = spm_bsplins(c,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 = inv(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,ddefa,P,xyz,msk)
% Get partials w.r.t. movements from first scan.

[idef,jac] = spm_get_image_def(1,ds,def,ddefa);
T    = P(1).mat\ds.M;
txyz = xyz*T';
c    = spm_bsplinc(P(1),ds.hold);
[m_ref,dx,dy,dz] = spm_bsplins(c,txyz(:,1),...
                   txyz(:,2)+idef,txyz(:,3),ds.hold);
indx    = ~isnan(m_ref);
mref_sf = 1 / (sum(m_ref(indx).*msk(indx)) / sum(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;
if ds.jm ~= 0
   m_ref = m_ref .* jac;
   dx    = dx    .* jac;
   dy    = dy    .* jac;
   dz    = dz    .* jac;
end
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,ext] = 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); clear 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    = ((inv(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));
   c           = spm_bsplinc(dispP,ds.hold);
   field       = spm_bsplins(c,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 doesnt 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
%_______________________________________________________________________


% SVARGPLVMPREDICTIONS2 This script implements the predictions for a classification task
%
% SEEALSO : demClassificationSvargplvm, demClassificationSvargplvm3, demClassificationGeneralSvargplvm, svargplvmPredictions, svargplvmPredictions3
%
% COPYRIGHT : Andreas C. Damianou, 2011

% VARGPLVM

%%%--- Script: svargplvmPredictions2 %%%%-
%%% This script is like svargplvmPredictions with a few additions
if ~exist('testOnTraining')
    testOnTraining=0;
end
if ~exist('globalOpt') || ~isfield(globalOpt, 'predictionsOptionX')
    if ~exist('predictionsOptionX')
        predictionsOptionX = 1;
    end
end
if ~exist('globalOpt') || ~isfield(globalOpt, 'predictionsGenerationOption')
    predictionsGenerationOption = 1;
end
if ~exist('testDisplayIters')
    testDisplayIters=0; %%%
end
if ~exist('predictionVariance')
    predictionVariance = false;
end
if ~exist('predictionXVariance')
    predictionXVariance = true;
end

numberTestPoints = size(Yts{1},1); % 10;
if testOnTraining
   % perm = randperm(model.N);
   % testInd = perm(1:size(Yall{1},1));
   testInd = 1:size(Yall{1},1);
else
    perm = randperm(size(Yts{obsMod},1));
    %testInd = perm(1:numberTestPoints);
    if ~exist('testInd')
        testInd = 1:size(Yts{1},1); %%%%
    end
end

ZpredMuAll = [];
XpredAll = zeros(length(testInd), model.q);  %%%
varZpredAll = zeros(length(testInd), model.q);  %%%
pb = myProgressBar(length(testInd));
for i=1:length(testInd)
    pb = myProgressBar(pb,i);
    curInd = testInd(i);
   % fprintf('# Testing indice number %d ', curInd);
    if testOnTraining
        y_star = model.comp{obsMod}.y(curInd,:);
        x_star = model.comp{obsMod}.vardist.means(curInd,:);
        varx_star = model.comp{obsMod}.vardist.covars(curInd,:);
    else
        y_star = Yts{obsMod}(curInd,:);
        z_star = Yts{infMod}(curInd,:);
        dst = dist2(y_star, model.comp{obsMod}.y);
        [mind, mini(i)] = min(dst);

        model.comp{obsMod}.mini = mini(i);
        Init(i,:) = model.vardist.means(mini(i),:);
        vardistx = vardistCreate(model.comp{obsMod}.vardist.means(mini(i),:), model.q, 'gaussian');
        vardistx.covars = model.comp{obsMod}.vardist.covars(mini(i),:);
        model.comp{obsMod}.vardistx = vardistx;
        iters = globalOpt.reconstrIters;
        % Find p(X_* | Y_*) which is approximated by q(X_*)
        [x_star, varx_star, modelUpdated] = vargplvmOptimisePoint(model.comp{obsMod}, vardistx, y_star, testDisplayIters, iters);%%%
    end
    numberOfNN = 1;
    % Now we selected a datapoint X_* by taking into account only the
    % private dimensions for Y. Now, based on the shared dimensions of
    % that, we select the closest (in a NN manner) X from the training data.
   % fprintf('# Finding the %d NN of X_* with the training X based only on the shared dims.\n', numberOfNN);
    [ind, distInd] = nn_class(model.X(:,sharedDims), x_star(:,sharedDims), numberOfNN, 'euclidean');

    ZpredMu = zeros(length(ind), size(model.comp{infMod}.y,2));
    varZstar = zeros(size(ZpredMu));
    ZpredSigma = zeros(length(ind), size(model.comp{infMod}.y,2));

    % Find p(y_*|x_*) for every x_* found from the NN
   % fprintf('# Predicting from the NN of X_* ');
    for k=1:numberOfNN
        x_cur = svargplvmPredictX(x_star, model.X, ind(k), sharedDims, predictionsOptionX);
        %[ZpredMu(k,:), ZpredSigma(k,:)] =
        %vargplvmPosteriorMeanVar(model.comp{infMod}, model.X(ind(k),:));
        if predictionsGenerationOption
            if ~predictionXVariance
                varx_star = varx_star * 0;
            end
            % NEW: varx_star
            if ~predictionVariance
                ZpredMu(k,:) = vargplvmPosteriorMeanVar(model.comp{infMod},x_cur, varx_star);
            else
                [ZpredMu(k,:) varZstar(k,:)] = vargplvmPosteriorMeanVar(model.comp{infMod},x_cur, varx_star);
            end
        else
            [ind2, distInd] = nn_class(model.X, x_cur, 1,'euclidean');
            ZpredMu(k,:) = model.comp{infMod}.y(ind2,:);
        end
    end
    ZpredMuAll = [ZpredMuAll; ZpredMu(1,:)];
    if predictionVariance
        varZpredAll = [varZpredAll; varZstar(1,:)];
    end
    XpredAll(i,:) = x_cur;
   % fprintf('\n\n');
end
fprintf(1, '\n');

ZpredMuAllOrig = ZpredMuAll;
defThresh = 0.5;
if min(min(Yts{2} == -1))
    defThresh = 0;
end
ZpredMuAll(ZpredMuAll >= defThresh) = 1;
ZpredMuAll(ZpredMuAll < defThresh) = 0;