function [means, variances, weights, minLogL, Si, expectation] = fitGMM(...,
    data, ...
    means, ...
    variances,...
    weights, ...
    equidistantMeans, ...
    updateMeans, ...
    sameVariance, ...
    updateVariance, ...
    maxIters, ...
    convCriterion, ...
    numBins, ...
    verbosity, ...
    gmmFig, ...
    costFig)
% Description: Fits a Gaussian mixture model to data by means of Expectation Maximization
%
% Call: [means, variances, weights, minLogL, Si, dataEstimate] = fitGMM(data, means, variances, weights, useFixedMeans,
%           useSameVariance, maxMMIters, convCriterion, numBins, verbosity, figToPlotTo)
%
% Input:    data:               Data to fit the GMM to
%           means:              Means of Gaussians in mixture
%           variances:          Variances of Gaussians in mixture
%           weights:            Weights of Gaussians in mixture
%           usedFixedMeans:     Means are updated according to the min and max value of the data (hence 2 parameters)
%           useSameVariance:    Equal variance is computed for all Gaussians
%           maxIters:           Maximum number of iterations allowed for fit
%           convCriterion:      Convergence threshold (relative change in cost)
%           numBins:            bins to plot histogram with
%           verbosity:          Verbosity level, higher value provides more output to stdout
%           figToPlotTo:        Figure handle to plot mixture model fitting to as iterations progress
%
% Output:   means:              Gaussian means that optimize likelihood of data
%           variances:          Gaussian variances that optimize likelihood of data
%           weights:            Gaussian weights that optimize likelihood of data
%           minLogL:            The negative loglikelihood at convergence
%           Si:                 The weights (or confidence) of each voxels in fit (posterior probability over variance)
%           dataEstimate:       The estimated data, given the optimized parameters
%
% Author: Christian Thode Larsen, christian.thode.larsen@gmail.com

% optional parameter initialization
if nargin < 5
    equidistantMeans = 0;
end
if nargin < 6
    updateMeans = 1;
end
if nargin < 7
    sameVariance = 0;
end
if nargin < 8
    updateVariance = 1;
end
if nargin < 9
    maxIters = 500;
end
if nargin < 10
    convCriterion = 1e-5;
end
if nargin < 11
    numBins = 1024;
end
if nargin < 12
    verbosity = 'none';
end
if nargin < 13
    gmmFig = -1;
end
if nargin < 14
    costFig = -1;
end

% initialize
numGaussians = length(means);
numVoxels = numel(data);

squaredDists = zeros(numVoxels, numGaussians);
likelihoods = zeros(numVoxels, numGaussians);

alphas = 1 - (0:numGaussians-1)' / (numGaussians-1);

% history of minLogLikelihoods - used to determine convergenge between iterations
costHist = 1/eps;
converged = false;

if ishandle(gmmFig)
    gaussianEvals = zeros(numBins, numGaussians);
    [histogram, binCenters] = hist(data, numBins);
    histogram = histogram / sum(histogram(:));
    binWidth = binCenters(2) - binCenters(1);
end

% Performing EM
for mmIter = 1:maxIters
     if ishandle(gmmFig)
        for i = 1:numGaussians
            gaussianEvals(:, i) =  (1 / sqrt(2*pi*variances(i))) * exp(-0.5 * (binCenters - means(i)).^2 / variances(i)) * weights(i) * binWidth;
        end
        figure(gmmFig)
        clf
        grid
        bar(binCenters, histogram);
        hold on
        plot(binCenters, gaussianEvals', 'g');
        plot(binCenters, sum(gaussianEvals, 2), 'r', 'lineWidth', 1);
        axis tight
        hold off
        drawnow
    end
    
    for i = 1:numGaussians
        squaredDists(:, i) = (data - means(i)).^2;
        likelihoods(:, i) = (1 / sqrt(2*pi*variances(i))) * exp(-0.5 * squaredDists(:,i) / variances(i)) * weights(i);
    end
    normalizer = sum(likelihoods, 2);
    
    minLogL = -sum(log(normalizer + eps));
    costHist(end + 1) = minLogL;%#ok
    
    relChange = (costHist(end - 1) - costHist(end)) / numVoxels;
    if abs(relChange) < convCriterion
        if ~isempty(strfind(verbosity, 'debug'))
            fprintf('\tGGM fit converged in %i iterations\n', mmIter);
        end
        converged = 1;
        break
    end
    
    if ishandle(costFig)
        figure(costFig)
        plot(costHist(2:end), 'b');
        axis tight
        title('Cost');
        xlabel('Iteration');
        ylabel('Cost');
        drawnow
    end
    
    posteriors = likelihoods ./ repmat(normalizer + eps, [1, numGaussians]);
    posteriorsSum = sum(posteriors);
    
    % classes we trust, to avoid collapses
    trusted = find(posteriorsSum >  (1e-3 * sum(posteriorsSum)));
    
    if(~all(trusted))
        warning('Classes are collapsing');
    end
    % update weights
    weights = posteriorsSum / sum(posteriorsSum);
    
    if(abs(sum(weights) - 1) > 1e-7)
        error('Weights do not sum up to one')
    end
           
    if updateMeans 
        if(equidistantMeans)
            lhs = zeros(2, 2);
            lhs(1, 1) = sum(sum(posteriors) .* (alphas').^2);
            lhs(2, 1) = sum(sum(posteriors) .* (alphas') .* (1-alphas'));
            lhs(1, 2) = lhs(2, 1);
            lhs(2, 2) = sum(sum(posteriors) .* (1-alphas').^2);
            
            rhs = [(data' * posteriors) * alphas; (data' * posteriors) * (1-alphas)];
            
            independentMeans = lhs \ rhs;
            means = ([alphas, (1-alphas)] *  independentMeans)';
        else
            newmeans = sum(posteriors .* repmat(data,[1, numGaussians])) ./ (posteriorsSum + eps);
            means(trusted) = newmeans(trusted);
        end
    end
    
    if updateVariance
        for i = 1:numGaussians
            squaredDists(:, i) = (data - means(i)).^2;
        end
        if(sameVariance)
            variance = posteriors(:)' * squaredDists(:) / sum(posteriorsSum);
            variances = repmat(variance, [1, numGaussians]);
        else
            newvariances = sum(posteriors .* squaredDists) ./ (posteriorsSum + eps);
            newvariances = max(newvariances, 0.001);
            variances(trusted)= newvariances(trusted);
        end
    end
end

if (~converged && ~isempty(strfind(verbosity, 'debug')))
    disp(['GMM stopped: Maximum number (', num2str(maxIters), ') of iterations reached'])
end

if(nargout > 4)
    Sik = posteriors ./ repmat(variances, [numVoxels, 1]);
    Si = sum(Sik, 2);
end

if(nargout > 5)
    expectation = sum(Sik .* repmat(means, [numVoxels, 1]), 2) ./ (Si + eps);
end

end