function [correctedData, biasField, mask, info] = iic(data, varargin)
% Description: Estimates a multiplicative intensitity inhomogeneity field from data by means of Bayesian inference
%
% CALL: [correctedData, biasField, mask, info] = iic(data, OPTIONAL)
%
% INPUT:
%   data:                       Data to be corrected.
%
% OPTIONAL INPUT:
%   'stepSize':                 Step size of data (voxel width, height and depth) in mm. Default: [1,1,1].
%
%   'desiredStepSize':          The step size you want to estimate at (used for downsampling the input data) in mm. Default: 1.
%
%   'downSamplingType':         Method used to to downsample data (as close as possible) to 'desiredStepSize'.
%                               Available options: 'noInterp' (default), 'nearestNeighbor'. Original N3 uses 'nearestNeighbor'.
%
%   'backgroundThresholdType':  Method used to compute background treshold. Used to mask out low intensity voxels.
%                               Available options: 'none', 'mixture', 'otzu' (similar to N3).
%                               If no mask and no prior is provided, defaults to 'otzu', else 'none'.
%
%   'maxNumberIterations':      Maximum number of iterations that the method is allowed to run. Default: 1000. Original N3 default: 50.
%
%   'maxNumberGMMIterations':   Maximum number of iterations that the mixture fitting component is allowed to run. Default: 1000.
%
%   'mask':                     A mask specifying voxels to be used for intensity inhomogeneity estimation.
%                               Must be same dimensions as data and should exclude (at least) all zero intensity voxels.
%
%   'prior':                    An IxP sized matrix of P spatially smooth priors on tissue classes or structures for I image voxels.
%                               All voxel priors are checked (and normalized if \sum_p v_p^i > 1) such that \sum_p v_p^i <= 1.
%                               A garbage/background class is computed from the priors to ensure that \sum_p v_p^i = 1.
%
%   'priorThreshhold':          Percentage threshhold for masking out background voxels when a prior is provided. Default: 0.99.
%                               This mask is combined with a thresholding of zero-intensity voxels.
%
%
%   'numberOfGaussians':        Number of Gaussians to use in mixture component. 
%                               If a prior is specified, numberOfGaussians is a vector containing the number of Gaussians to use per class. 
%                               Default: 3 gaussians in a single class, original N3 uses a default value of 200 (hidden).
%
%   'equidistantMeans':         Default: false. If true, Gaussian means are restricted to equidistant placement.
%                               Original N3: true.
%
%   'updateMeans':              Default: true. If set to false, means are never updated after initialization.
%
%   'sameVariance':             Default: false. If set to true, Gaussian variances are restricted to equal value.
%                               Original N3 default: true.
%
%   'updateVariance':           Default: true. If set to false, variances are never updated after initialization
%                               (Gaussian width, i.e., squared distance between means).
%                               If varianceScale is > 1 and sameVariance is true, variance will be fixed until near convergence for speed,
%                               and then allowed to update to increase accuracy.
%                               Original N3: false.
%
%   'varianceScale':            Default: 1. Scale the Gaussian variance by a constant.
%
%   'gmmFitType':               Method used to fit the mixture model. Options: 'mixture' (default), 'wiener'.
%                               Original N3 uses 'wiener', a heuristic gaussian mixture fitting method performed by a
%                               regularized least-squares method (Wiener deconvolution).
%
%   'convergenceType':          Correction method convergence type. Options: 'cost' (default), uses change in
%                               cost per voxel relative to cost starting value (computed from initialized parameters).
%                               N3 fitting implies 'stdDifferenceField', which uses standard deviation
%                               of difference in field estimates between current and previous iteration.
%
%   'convergence':              Default: 1e-6. Convergence threshold. Original N3 uses a default value of 1e-3.
%
%   'convergenceGMM':           Default: 1e-6. Convergence threshold for Gaussian mixture. Only applies if 'gmmFitType': 'mixture'.
%
%   'smoothingType':            Type of smoothing used on intensity residual, in order to conform with the assumption of a smooth,
%                               multiplicative bias. Options: 'cosine', 'N3' (default). 'N3' is really a bspline scheme with
%                               equidistant control points and regularization on the partial second order derivatives (bending
%                               energy).
%
%   'smoothingDistance':        Default: 50. Distance in mm which controls how flexible the bias is. Implicitly controls the number
%                               of basis functions used for smoothing.
%                               Note that this parameter can be specified per dimension (e.g., [50, 60, 70]), which allows (discrete)
%                               control with how flexible the bias is along each dimension.
%
%   'smoothingRegularization':  Default value: 1e-7. A scalar that controls how heavily we regularize (penalize curvature).
%                               Must be specified when using bspline smoothing ('N3' default).
%                               Original N3 uses 1e-7.
%
%   'voxelScaleRegularization': Default: true. Scale the regularization by the number of voxels (balances reg term vs dataterm).
%                               Original N3 fitting uses a default value of 'true'.
%
%   'smoothExponentiatedBias':  Default: false. A smoothing performed by the original N3 method of the final
%                               exponentiated bias field estimate. Generally found to do very little, and it might hurt
%                               performance (e.g., CJV between WM and GM). Default in 'strict' N3, otherwise not recommended.
%
%   'fwhm':                     Default: 0.15. N3 parameter used to control variance of Gaussians in mixture. Implies
%                               'gmmFitType': 'wiener'.
%
%   'wienerFilterNoise':        Default: 0.1. N3 hidden parameter used to control noise used by wiener fitting of
%                               gaussian mixture.
%
%   'mimicN3':                  Options: 'no' (default), 'model', 'strict'.
%                               'model' implies N3 mimicking and enables all model essential N3 parameters.
%                               'strict' further mimics N3 by enabling various tweaks and hacks.
%                               Parameters that can be user controlled in the original N3 are NOT set to the original defaults.
%                               Note that you can configure N3 like behavior by setting other flags yourself.
%
%   'histogramBins':            Default: 256. Number of bins used when plotting histogram and mixture fitting.
%
%   'showPlots':                Default: false. True enables a number of plots that are useful to debug and inspect
%                               progress of correction.
%
%   'verbosity':                Options: 'none' (default), 'info' and 'debug'. Print relevant information to commandline
%                               during progress of correction.
%
% OUTPUT:
%   correctedData:              Intensity homogeneity corrected data at original resolution.
%
%   biasField:                  The estimated intensitiy inhomogeneity field at original resolution.
%
%   mask:                       The mask of voxels that are used for bias estimation at original resolution. Useful if
%                               mask is computed using background thresholding.
%
%   info:                       A struct containing various history vectors, useful for debugging and inspection of
%                               method behavior.
%
% Author:       Christian Thode Larsen (cthla@dtu.dk), christian.thode.larsen@gmail.com
% Literature:   Leemput et al 1999. Automated Model-based Bias Field Correction of MR Images of the Brain.
%               Larsen et al 2014. N3 Bias Field Correction Explained as a Bayesian Modeling Method.
%
% NOTE: IF USED FOR RESEARCH, PLEASE CONSIDER CITING THE TWO AFOREMENTIONED REFERENCES!
% NOT FOR COMMERCIAL USE


if nargin < 1
    error('Call: [correctedData, OPTIONAL] = icc(data, OPTIONAL)')
end

versionString = strsplit(version, '.');
%if (str2double([versionString{1}, '.', versionString{2}]) < 8.2)
%    error('You need Matlab version 8.1 (R2013b) or higher for this program to run');
%end


% =============================================================
%
% Parse inputs
%
% =============================================================
parser = inputParser;
addRequired(parser, 'data', @isnumeric);

addParameter(parser, 'mask', [], ...
    @(x) isnumeric(x) && all(size(x) == size(data)));
addParameter(parser, 'stepSize', [1,1,1], @isnumeric);
addParameter(parser, 'desiredStepSize', 1, @isnumeric);
addParameter(parser, 'downSamplingType', 'noInterp', ...
    @(x) any(validatestring(x, {'noInterp', 'noInterpCentered', 'nearestNeigbor'})));
addParameter(parser, 'backgroundThresholdType', 'none', ...
    @(x) any(validatestring(x, {'mixture', 'otzu', 'N3'})));
addParameter(parser, 'maxNumberIterations', 1000, ...
    @(x) isinteger(x) && x > 0);
addParameter(parser, 'maxNumberGMMIterations', 500, ...
    @(x) isinteger(x) && x > 0);
addParameter(parser, 'prior', [], ...
    @(x) isnumeric(x) );
addParameter(parser, 'priorThreshhold', 0.99, ...
    @(x) isnumeric(x) );
addParameter(parser, 'numberOfGaussians', 3, ...
    @(x) isnumeric(x) && all(x > 0));
addParameter(parser, 'equidistantMeans', false, @islogical);
addParameter(parser, 'updateMeans', true, @islogical);
addParameter(parser, 'varianceScale', 1, @isnumeric);
addParameter(parser, 'sameVariance', false, @islogical);
addParameter(parser, 'updateVariance', true, @islogical);
addParameter(parser, 'gmmFitType', 'mixture', ...
    @(x) any(validatestring(x, {'mixture', 'wiener'})));
addParameter(parser, 'convergenceType', 'cost', ...
    @(x) any(validatestring(x, {'cost', 'stdDifferenceField'})));
addParameter(parser, 'convergence', 1e-5, ...
    @(x) isnumeric(x) && x > 0);
addParameter(parser, 'convergenceGMM', 1e-5, ...
    @(x) isnumeric(x) && x > 0);

% todo: make 'bspline' option that uses proper spline scaling + regularization
addParameter(parser, 'smoothingType', 'bspline', ...
    @(x) any(validatestring(x, {'bspline', 'spm'})));
addParameter(parser, 'smoothingDistance', 50, ...
    @(x) isnumeric(x) && x > 0);
addParameter(parser, 'smoothingRegularization', 1e-7, ...
    @(x) isnumeric(x) && x > 0);
addParameter(parser, 'voxelScaleRegularization', true, @islogical);
addParameter(parser, 'smoothExponentiatedBias', false, @islogical);

addParameter(parser, 'fwhm', 0.15, ...
    @(x) isnumeric(x) && x > 0);
addParameter(parser, 'wienerFilterNoise', 0.1, ...
    @(x) isnumeric(x) && x > 0);
addParameter(parser, 'mimicN3', 'no',  ...
    @(x) any(validatestring(x, {'no', 'model', 'strict'})));

addParameter(parser, 'histogramBins', 200, ...
    @(x) isnumeric(x) && x > 0);
addParameter(parser, 'showPlots', false, @islogical);
addParameter(parser, 'verbosity', 'none', ...
    @(x) any(validatestring(x, {'none', 'info', 'debug'})));



parse(parser, data, varargin{:});

% data handling
mask = parser.Results.mask;
stepSize = parser.Results.stepSize;
desiredStepSize = parser.Results.desiredStepSize;
downSamplingType = parser.Results.downSamplingType;
bkgThresType = parser.Results.backgroundThresholdType;

% model config
prior = parser.Results.prior;
priorThreshhold = parser.Results.priorThreshhold;
numGaussians = parser.Results.numberOfGaussians;
equidistantMeans = parser.Results.equidistantMeans;
updateMeans = parser.Results.updateMeans;
sameVariance = parser.Results.sameVariance;
updateVariance =  parser.Results.updateVariance;
varianceScale = parser.Results.varianceScale;
gmmFitType = lower(parser.Results.gmmFitType);

% convergence settings
convergenceType = lower(parser.Results.convergenceType);
convergence = parser.Results.convergence;
convergenceGMM = parser.Results.convergenceGMM;
maxIters = parser.Results.maxNumberIterations;
maxMMParamIters = parser.Results.maxNumberGMMIterations;

% bias shape
smoothingType = lower(parser.Results.smoothingType);
smoothingDist = parser.Results.smoothingDistance;
lambda = parser.Results.smoothingRegularization;
voxelScaleReg = parser.Results.voxelScaleRegularization;
smoothExpBias = parser.Results.smoothExponentiatedBias;

% output
histogramBins = parser.Results.histogramBins;
showPlots = parser.Results.showPlots;
verbosity = lower(parser.Results.verbosity);

% N3 specific
fwhm = parser.Results.fwhm;
wienerFiltNoise = parser.Results.wienerFilterNoise;
mimicN3 = parser.Results.mimicN3;

% save original data for final correction
origData = data;
origDataSize = size(origData);
origStepSize = stepSize;

% =================================================================
%
% Hard coded parameters - Settings we do not want to set outside the function right now
%
% =================================================================

% for stepping through bias field smoothness using the regularizer
useMultiResolutionLambda = false;
lambdaStepValues = [9*lambda, 3*lambda, lambda];
currentLambda = lambda;
lambdaStep = 1;

% We assume values below this are background
bkgIntensityCap = 1e-7;

% figure configs
estimatedBiasFieldFig = -1;
correctedDataFig = -1;
historyFig = -1;
gmmFig = -1;
gmmCostFig = -1;

% =============================================================
%
% Relevant checks and parameter dependent configs
%
% =============================================================

% if the user didn't provide priors, masks or configured the background thresholding, set to default otzu
if isempty(prior) && isempty(mask) && ~isempty(strfind(bkgThresType, 'none'))
    warning('No prior, mask or background Threshold type specified, Setting background Thresholding type to otzu.');
    bkgThresType = 'otzu';
end

% unless using priors, we assume one class consisting of multiple gaussians (or alternatively, several classes of one gaussian each)
numClasses = 1;

% Full EM with prior
if  ~isempty(prior)
    if isempty(strfind(gmmFitType, 'mixture'))
        error('Priors are only supported with EM, please change mode to mixture');
    end
    gmmFitType = 'mixtureWithPrior';
    
    % The numclasses parameter defines the number of classes we want to fit + a garbage class
    % verify that priors are only supported for the real classes
    numClasses = length(numGaussians);
    numPriors = size(prior, 2);
    if numClasses ~= numPriors + 1
        error('The vector numGaussians should have length equal to the number of priors + 1 (the background)');
    end
    
    % We need a map of which gaussians belong to what class, since we support multiple gaussians per class
    priorMap = zeros(1,sum(numGaussians));
    priorClass = 1;
    count = 1;
    for i = 1:numClasses
        for j = 1:numGaussians(i)
            priorMap(count) = priorClass;
            count = count + 1;
        end
        priorClass = priorClass + 1;
    end
    numGaussians = sum(numGaussians);
    
    % interpolation when downsampling may mess up our priors, so we don't support it (and interpolation is bad for you in general)
    if isempty(strfind(downSamplingType, 'noInterp'))
        fprintf('Priors are only supported with noInterp downsampling, changing');
        downSamplingType = 'noInterp';
    end
    
    % compute background prior + ensure they all sum up to one
    totalProbability = sum(prior, 2);
    idx = find(totalProbability > 1);
    if(~isempty(idx))
        warning('Probabilities sum to over one, probably a resampling artifact. Normalizing these entries to ensure background probability does not become negative');
        for i = 1:numPriors
            prior(idx, i) = prior(idx, i) ./ totalProbability(idx);
        end
        totalProbability(idx) = 1;
    end
    prior(:,numPriors+1) = 1 - totalProbability;
end

% if n3 flag is model or strict, ensure these options are set, since they cannot be changed in N3
if isempty(strfind(mimicN3, 'no'))
    if isempty(strfind(gmmFitType, 'wiener'))
        fprintf('Changing gaussian mixture fit mode to wiener\n');
        gmmFitType = 'wiener';
    end
    if isempty(strfind(smoothingType, 'bspline'))
        fprintf('Changing residual smoothing to N3 (bsplines)\n');
        smoothingType = 'bspline';
    end
    if numGaussians ~= 200
        fprintf('Changing numGaussians to 200\n');
        numGaussians = 200;
    end
    
    if isempty(strfind(convergenceType, 'stddifferencefield'))
        fprintf('Changing conversion type to standard deviation of difference (bias) field\n');
        convergenceType = 'stddifferencefield';
    end
    
    if ~voxelScaleReg
        fprintf('Changing scaling of regularization by number of voxels to true\n');
        voxelScaleReg = true;
    end
    
    % used in select cases
    n3Variance = 1 / (8 * log(2)) * fwhm^2;
    lambda = lambda/(2*n3Variance);   
    currentLambda = lambda;
    numLambdaSteps = 1;   
end

% further set strict options for mimicN3
if ~isempty(strfind(mimicN3, 'strict'))
    if ~smoothExpBias
        fprintf('Changing smoothing of exponentiated bias estimate to true\n');
        smoothExpBias = true;
    end
    
    if isempty(strfind(downSamplingType, 'nearestneighbor'))
        fprintf('Changing downsampling type to nearestNeighbor\n');
        downSamplingType = 'nearestneighbor';
    end
    
    if isempty(mask) && isempty(strfind(bkgThresType, 'N3'));
        fprintf('Changing background thresholding type to N3\n');
        bkgThresType = 'N3';
    end
end

% wiener filtering does have a few dependencies...
if(~isempty(strfind(gmmFitType, 'wiener')))
    sameVariance = true;
    equidistantMeans = true;
    useParzenWindowHist = 1;
    useInterpTable = 1;
    doPad = 1;
end

% for plotting stuff
if showPlots
    gmmFig = figure('name', 'gmm');
    gmmCostFig = figure('name', 'gmm cost');
    estimatedBiasFieldFig = figure( 'name', 'bias' );
    correctedDataFig = figure( 'name', 'corrected' );
    historyFig = figure( 'name', 'History' );
    bFig = figure('name', 'basis functions and regularizer');
    if ~isempty(prior)
        priorFig = figure('name', 'priors');
        posteriorFig = figure('name', 'posteriors');
    end
end

% =================================================================
%
% Preprocessing
%
% =================================================================

% Time it
timeFitData = 0;
timeEstBias = 0;
timeFitWeights = 0;
startTime = tic;
startTimePreProcessing = tic;

% =================================================================
%
% Masking + downsampling
%
% =================================================================

% downsample according to desired resolution and sampling type
[data, start, stepSize, sampleIDX] = downSampleData(origData, origStepSize, desiredStepSize, downSamplingType);
dataSize = size(data);

% if no mask or prior has been specified, we need to mask out the background
origMaskBKG = ones(origDataSize);
maskBKG = ones(dataSize);
if isempty(strfind(bkgThresType, 'none'))
    if isempty(strfind(verbosity, 'none'))
        fprintf('Creating mask by tresholding background\n');
    end
    % oook. Use a histogram that bins intensities using the range to define the number of bins
    nBins = ceil(max(data(:))-min(data(:)) + 1);
    nBinsOrig = ceil(max(origData(:))-min(origData(:)) + 1);
    % if no mask was specified, make a brutal background mask by tresholding
    % note that N3 actually does some sort of rounding on the subsampled threshing, but not in the end when threshing
    % the full res data. I was unable to reproduce the exact value they arrive at, so this is a (very) minor deviation
    % from how N3 does things
    bkgThres = biModalThreshold(data, false, bkgThresType, nBins);
    bkgThresOrig = biModalThreshold(origData, false, bkgThresType, nBinsOrig);
    
    maskBKG(data < bkgThres) = 0;
    origMaskBKG(origData < bkgThresOrig) = 0;
end

% if provided, downsample mask
if isempty(mask)
    origMask = ones(origDataSize);
    mask = ones(dataSize);
else
    origMask = mask;
    mask = downSampleData(mask, origStepSize, desiredStepSize, downSamplingType);
end

% if provided, downsample prior mask
if isempty(prior)
    origMaskPrior = ones(origDataSize);
    maskPrior = ones(dataSize);
else
    origMaskPrior = zeros(origDataSize);
    % mask out voxels which are threshhold % or more likely to be background + voxels that are virtually zero intensities
    origMaskPrior(prior(:, end) < priorThreshhold  & origData(:) > bkgIntensityCap) = 1;
    
    % ok as long as we only allow downsampling where values are not interpolated
    maskPrior = downSampleData(origMaskPrior, origStepSize, desiredStepSize, downSamplingType);
end

% now we compute the full mask
origMask = origMask .* origMaskPrior .* origMaskBKG;
mask = mask .* maskPrior .* maskBKG;


% we use logical indexing heavily, so this saves us from quite a lot of repeated boolean operations
origMaskLogical = origMask > 0.5;
maskLogical = mask > 0.5;

% now do a final check to see if the mask contains any zero voxels (must NOT happen)
if any(origData(origMaskLogical) < bkgIntensityCap)
    error(['The final mask contains zero intensity voxels. This is likely due to a mask and/or prior where voxels have zero intensities. ' ...
        'This is not allowed because the data is log-transformed during estimation. Enable background thresholding, or correct your mask.']);
end

% if we're using priors, we need to filter out all voxels that are not used, after downsampling
if ~isempty(prior)
    priorsFiltered = prior(sampleIDX, :);
    priorsFiltered = priorsFiltered(maskLogical, :);
    
    if showPlots
        figure(priorFig)
        numFigsPerDim = ceil(sqrt(numClasses));
        for i  = 1:numClasses
            subplot(numFigsPerDim,numFigsPerDim,i);
            tmp = mask;
            tmp(mask > 0.5) = priorsFiltered(:,i);
            showImage(tmp);
        end
    end
end

% This way we know exactly how many voxels we operate on when computing the minLogLikelihood
numVoxels = sum(maskLogical(:));

if(numVoxels ~= sum(mask(:)))
    error('Mask has not been properly binarized');
end

% we work in the log-domain - and we're ensured that zero intensities have been masked out
data(maskLogical) = log(data(maskLogical));
data = data .* mask;

% pre-allocate
estimatedBiasField = zeros(dataSize);
means = zeros(1,numGaussians);
variances = zeros(1,numGaussians);
weights = zeros(1,numGaussians);

% Only work on non-zero intensities
intensities = data(maskLogical);

% =================================================================
%
% Initialize parameters for one class mixture and priors
%
% =================================================================

% if using priors, do one round of parameter updates for initialization, using the prior as the posterior
if numClasses > 1
    % compute class statistics
    posteriors = priorsFiltered;
    posteriorsSum = sum(posteriors);
    
    % update weights - flat initially for gaussians within each class
    for l = 1:numClasses
        weights(priorMap == l) = 1 / sum(priorMap == l);
    end
    
    % update class means and variances
    classMeans = sum(posteriors .* repmat(intensities,[1, numClasses])) ./ (posteriorsSum + eps);
    squaredDists = zeros(numVoxels, numClasses);
    for i = 1:numClasses
        squaredDists(:, i) = (intensities - classMeans(i)).^2;
    end
    classVariances = sum(posteriors .* squaredDists) ./ (posteriorsSum + eps);
    classStd = sqrt(classVariances);
    
    % now that we have some per class statistics, we can initialise the gaussians within each
    for i = 1:numClasses
        gaussiansInClass = sum(priorMap == i);
        range = [classMeans(i)-3*classStd(i), classMeans(i)+3*classStd(i)];
        rangeSpan = range(2) - range(1);
        gaussianWidth = rangeSpan / (numGaussians - 1);
        
        means(priorMap == i) = range(1) + (0:gaussiansInClass-1) * gaussianWidth;
        variances(priorMap == i) =  repmat(gaussianWidth^2 * varianceScale, [1, gaussiansInClass]);
    end
    % one class configurations (mixture, N3)
else
    range = [min(intensities(:)), max(intensities(:))];
    % Define range of intensities that our non-zero data spans (same as DHistogram.cc in N3)
    rangeSpan = range(2) - range(1);
    gaussianWidth = rangeSpan / (numGaussians - 1);
    
    % mean init
    means = range(1) + (0:numGaussians-1)*gaussianWidth;
    
    % default variance initialization
    variances = repmat(gaussianWidth^2, [1, numGaussians]);
    variances = variances * varianceScale;


    % uniform weights - tests indicate that using the histogram for init doesn't really make any difference
    % well, apparently this is not always the case... in the same variance scenario, using the histogram seems to be pretty decent.
    % If I do not do it, the method becomes more sensitive to how I scale the variance - which MIGHT perform comparably, but not always
    weights = ones(1, numGaussians) / numGaussians;
    if sameVariance
        weights = hist(intensities, numGaussians) / sum(hist(intensities, numGaussians));
    end  
end

% =================================================================
%
% Initialize histories
%
% =================================================================

oneOverEps = 1 / eps;

costHist = oneOverEps;
dataTermHist = oneOverEps;
regTermHist = oneOverEps;

stdFieldHist  = oneOverEps;
stdDiffHist  = oneOverEps;
cvDataHist  = oneOverEps;

meansHist = means;
variancesHist = variances;
weightsHist = weights;

% if useMultiResolutionLambda
%     lambdaHist = 1;
% end

% =================================================================
%
% Initialize parameters for bias field
%
% =================================================================

% compute everything needed for basis function bias field representation
[numBFuncs, domain] = computeNumberOfSeparableBasisFunctions(origDataSize, origStepSize, smoothingType, smoothingDist);
[B, B2, indexMapB2] = computeSeparableBasisFunctions(numBFuncs, domain, smoothingDist, dataSize, smoothingType,  stepSize, start);
J = computeRegularizers(smoothingType, numBFuncs, domain, origStepSize);

% speed-up, A2 does not depend on variance
if sameVariance
    A2 = constructAtWA(mask, B2, indexMapB2);
end

% plot the basis functions + regularizer
if showPlots
    figure(bFig);
    subplot(2,2,1); plot(B{1});     axis tight;
    subplot(2,2,2); plot(B{2});     axis tight;
    subplot(2,2,3); plot(B{3});     axis tight;
    subplot(2,2,4); imagesc(J);     axis tight;
end

% initialize coefficients (flat field in log-domain)
C = zeros(prod(numBFuncs), 1);

% Balance the data term and the reg term (corresponds to dividing the data term by N)
if voxelScaleReg
    J = J * numVoxels;
end

differenceField = data(maskLogical);

% end of preprocessing
timePreProcessing = toc(startTimePreProcessing);


% print some info
if isempty(strfind(verbosity, 'none'))
    fprintf('Settings:\n');
    fprintf('==================================\n');
    fprintf('numGaussians: %d\n', numGaussians);
    fprintf('equiDistantMeans: %d\n', equidistantMeans);
    fprintf('sameVariance: %d\n', sameVariance);
    fprintf('updateVariance: %d\n', updateVariance);
    fprintf('varianceInitScale: %.3f\n', varianceScale);
    fprintf(['gmmFitType: ', gmmFitType, '\n']);
    fprintf(['convergenceScheme: ', convergenceType, '\n']);
    fprintf('stopCriterion: %g\n', convergence);
    fprintf(['SmoothingType: ', smoothingType, '\n']);
    fprintf('SmoothingDist: %d\n', smoothingDist);
    fprintf('lambda: %g\n', lambda);
    fprintf('voxelScaleRegularization: %d\n', voxelScaleReg);
    fprintf('histogramBins: %d\n', histogramBins);
    fprintf('==================================\n\n');
end

if ~isempty(strfind(verbosity, 'debug'))
    fprintf('Bias field correcting data:\n\n');
end

% =================================================================
%
% Bias field correct data
%
% =================================================================

% iterative EM scheme starts here, iteration cap is given by maxIters
for iter = 1:maxIters
    if ~isempty(strfind(verbosity, 'debug'))
        fprintf('ITERATION %d\n', iter);
    end
    
    
    % bias field is always estimated with respect to original data; only in the end is the data corrected.
    biasFieldCorrectedData = data - estimatedBiasField;
    differenceField = differenceField - estimatedBiasField(maskLogical);
    
    if showPlots
        figure(correctedDataFig);
        showImage( exp(biasFieldCorrectedData) .* mask);
        figure(estimatedBiasFieldFig);
        showImage(exp(estimatedBiasField) .* mask);
    end
    
    % N3 doesn't even use the coefficient of variation for termination - it uses the standard deviation!
    % but ok, the CV is just a normalization...
    stdFieldHist(end + 1) = std(estimatedBiasField(maskLogical)); %#ok
    stdDiffHist(end + 1) = std(differenceField); %#ok
    % coefficient of variation within mask on subsampled data
    cvDataHist(end + 1) = std(biasFieldCorrectedData(maskLogical)) / mean(biasFieldCorrectedData(maskLogical)); %#ok
    
    % difference field convergence should take place prior to next mixture model fit, as it depends on the correction step
    if(~isempty(strfind(convergenceType, 'stddifferencefield')) && stdDiffHist(end) < convergence)
        if ~isempty(strfind(verbosity, 'debug'))
            fprintf('Converged in iteration %d!\n', iter);
        end
        break;
    end
    
    % pull out the current set of bias field corrected intensities for next mixture model fit
    intensities = biasFieldCorrectedData(maskLogical);
  
    startTimeFitData = tic;
    switch gmmFitType
        % N3 fit
        case 'wiener'
            [means, variances, weights, minLogL, expectation, elapsedTimeWeightsIter] = fitGMMWiener(...
                intensities, ...
                numGaussians, ...
                fwhm, ...
                wienerFiltNoise,...
                histogramBins, ...
                useParzenWindowHist, ...
                useInterpTable, ...
                doPad, ...
                gmmFig);
            timeFitWeights = timeFitWeights + elapsedTimeWeightsIter;
            
            % regular GMM fit
        case 'mixture'
            [means, variances, weights, minLogL, Si, expectation] = fitGMM(...
                intensities, ...
                means, ...
                variances, ...
                weights, ...
                equidistantMeans, ...
                updateMeans, ...
                sameVariance, ...
                updateVariance, ...
                maxMMParamIters, ...
                convergenceGMM, ...
                histogramBins, ...
                verbosity, ...
                gmmFig, ...
                gmmCostFig);
            % regular GMM fit + priors
        case 'mixtureWithPrior'
            [means, variances, weights, minLogL, Si, expectation, posteriors] = fitGMMwithPrior(...
                intensities, ...
                priorsFiltered, ...
                priorMap, ...
                means, ...
                variances, ...
                weights, ...
                maxMMParamIters, ...
                convergenceGMM, ...
                histogramBins, ...
                verbosity, ...
                gmmFig, ...
                gmmCostFig);
            
            if showPlots
                figure(posteriorFig);
                segmentation = zeros(dataSize);
                for i = 1:numClasses
                    segmentation(maskLogical) = sum(posteriors(:, priorMap == i), 2);
                    subplot(2,2,i)
                    showImage(segmentation);
                end
            end
        otherwise
            error('invalid Gaussian Mixture Model Fitting');
    end
    % Time the histogram fit
    timeFitData = timeFitData + toc(startTimeFitData);
    
    meansHist(end+1,:) = means; %#ok
    variancesHist(end+1,:) = variances; %#ok
    weightsHist(end+1,:) = weights; %#ok
    
    dataTermHist(end + 1) = minLogL; %#ok
    regTermHist(end + 1) = currentLambda * C' * J * C;%#ok
    costHist(end+1) = dataTermHist(end) + regTermHist(end);%#ok
    
    if ~isempty(strfind(verbosity, 'debug'))
        fprintf('\tCost: %.5f\n', costHist(end));
    end
    
    if showPlots
        figure(historyFig);
        %  subplot(2,2,1);
        plot(costHist(2:end), 'b');
        axis tight
        hold on
        %  plot(dataTermHist(2:end), 'r');
        %  plot(regTermHist(2:end), 'g');
        hold off
        title('Cost');
        xlabel('Iteration');
        ylabel('Cost');
    end
    
    % the relative change in cost for this iteration
    relativeChangeCost = (costHist(end-1) - costHist(end)) / numVoxels;
    
    % Convergence / multi resolution
    if ((isempty(strfind(gmmFitType, 'wiener'))) && (abs(relativeChangeCost) < (convergence * 10)))
        % Start with a very rigid field, and increase flexibility as we go along.
        % Avoid local minima when the field is too flexible compared to how well we're able to fit the data
        if useMultiResolutionLambda && (lambdaStep < numLambdaSteps)  
            % increase smoothing flexibility by decreasing lambda (regularization)
            lambdaStep = lambdaStep + 1;
            currentLambda = lambdaStepValues(lambdaStep);
            lambdaHist(end + 1) = iter; %#ok
            
            if ~isempty(strfind(verbosity, 'debug'))
                fprintf('Decreasing lambda to %f \n', currentLambda);
            end 
            % we're done!
        elseif (abs(relativeChangeCost) < convergence)
            if ~isempty(strfind(verbosity, 'debug'))
                fprintf('Converged in iteration %d!\n', iter);
            end
            break;
        end
    end
    
    % store previous fild for difference field computation in next iter
    differenceField = estimatedBiasField(maskLogical);
    
    % Estimate and time the bias field
    startTimeEstBias = tic;
    
    % the noisy residual that we want smoothed
    residualImage = zeros(dataSize);
    residualImage(maskLogical) = data(maskLogical) - expectation;
    
    % fit basis functions to data to obtain coefficients
    if(sameVariance)
        lhs = A2 + variances(1) * 2 * currentLambda * J;
        rhs = constructAtWr(residualImage, mask, B);
    else
        weightImage = zeros(dataSize);
        weightImage(maskLogical) = Si;
        
        lhs = constructAtWA(weightImage, B2, indexMapB2) + 2 * currentLambda * J;
        rhs = constructAtWr(residualImage, weightImage, B);
    end
    C = lhs \ rhs;
    
    % evaluate basis functions
    estimatedBiasField = expandSeparableBasisFunctions(C, B);
    
%     if ~isempty(strfind(verbosity, 'debug'))
%         fprintf('\tResidual min: %f, max: %f\n',  exp(min(residualImage(maskLogical))), exp(max(residualImage(maskLogical))));
%         fprintf('\tField    min: %f, max: %f\n',  exp(min(estimatedBiasField(maskLogical))), exp(max(estimatedBiasField(maskLogical))));
%     end
    timeEstBias = timeEstBias + toc(startTimeEstBias);
end

startTimePostProcessing = tic;

% Reconstruct the bias field at the original resolution
if ~isempty(strfind(verbosity, 'debug'))
    disp('Reconstructing the bias field at the original resolution')
end

% wacko N3 smoothing of bias in original domain.
if smoothExpBias
    residualImage(maskLogical) = exp(estimatedBiasField(maskLogical));
    
    % fit basis functions to data to obtain coefficients
    if(sameVariance)
        lhs = A2 + variances(1) * 2 * lambda * J;
        rhs = constructAtWr(residualImage, mask, B);
    else
        lhs = constructAtWA(weights, B2, indexMapB2) + 2 * lambda * J;
        rhs = constructAtWr(residualImage, weightImage, B);
    end
    
    C = lhs \ rhs;
end

% evaluate basis functions obtain bias at original resolution
B = computeSeparableBasisFunctions(numBFuncs, domain, smoothingDist, origDataSize, smoothingType,  origStepSize, start);
biasField = expandSeparableBasisFunctions(C, B);


% Exponentiate back into original intensity domain.
if ~smoothExpBias
    biasField(origMaskLogical) = exp(biasField(origMaskLogical));
end

% mask it
biasField = biasField .* origMask;

% Correct data at original resolution.
correctedData = zeros(origDataSize);
correctedData(origMaskLogical) = origData(origMaskLogical) ./ (biasField(origMaskLogical) + eps);
mask = origMask;

timePostProcessing = toc(startTimePostProcessing);

% Report on computation time
elapsedTime = toc(startTime);
oneOverElapsedTime100 = 100 / elapsedTime;
timePerIter = (timeFitData + timeEstBias) / (iter-1);

if nargout > 3
    info.iterations = iter;
    info.timePerIteration = timePerIter;
    info.costHistory = costHist;
    info.dataTermCostHistory = dataTermHist;
    info.regularizerTermCostHist = regTermHist;
    info.cvDataHistory = cvDataHist;
    info.stdFieldHistory = stdFieldHist;
    info.stdDifferenceFieldHistory = stdDiffHist;
    info.meanHistory = meansHist;
    info.varianceHistory = variancesHist;
    info.weightHistory = weightsHist;
end

if isempty(strfind(verbosity, 'none'))
    disp('Done computing bias fields, displaying statistics');
    disp(['   Iterations used           : ', num2str(iter-1)]);
    disp(['   Time taken                : ', num2str(elapsedTime), ' seconds']);
    disp(['   Time per iteration        : ', num2str(timePerIter), ' seconds']);
    disp(['   Pre processing            : ', num2str(timePreProcessing * oneOverElapsedTime100) '%']);
    disp(['   Fitting models            : ', num2str(timeFitData * oneOverElapsedTime100) '%']);
    disp(['   Updating weights          : ', num2str(timeFitWeights * oneOverElapsedTime100) '%']);
    disp(['   Computing bias fields     : ', num2str(timeEstBias * oneOverElapsedTime100) '%']);
    disp(['   Post processing           : ', num2str(timePostProcessing * oneOverElapsedTime100) '%']);
end

% end of iic
end
