## Copyright (C) 2015 Lachlan Andrew <lachlanbis@gmail.com>
## Copyright (C) 2018-2020 John Donoghue <john.donoghue@ieee.org>
## Copyright (C) 2022 Andreas Bertsatos <abertsatos@biol.uoa.gr>
##
## This file is part of the statistics package for GNU Octave.
##
## This program is free software; you can redistribute it and/or modify it under
## the terms of the GNU General Public License as published by the Free Software
## Foundation; either version 3 of the License, or (at your option) any later
## version.
##
## This program is distributed in the hope that it will be useful, but WITHOUT
## ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
## FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
## details.
##
## You should have received a copy of the GNU General Public License along with
## this program; if not, see <http://www.gnu.org/licenses/>.

classdef gmdistribution
   ## -*- texinfo -*-
   ## @deftypefn  {statistics} {@var{GMdist} =} gmdistribution (@var{mu}, @var{Sigma})
   ## @deftypefnx {statistics} {@var{GMdist} =} gmdistribution (@var{mu}, @var{Sigma}, @var{p})
   ## @deftypefnx {statistics} {@var{GMdist} =} gmdistribution (@var{mu}, @var{Sigma}, @var{p}, @var{extra})
   ##
   ## Create an object of the gmdistribution class which represents a Gaussian
   ## mixture model with k components of n-dimensional Gaussians.
   ##
   ## Input @var{mu} is a k-by-n matrix specifying the n-dimensional mean of
   ## each of the k components of the distribution.
   ##
   ## Input @var{Sigma} is an array that specifies the variances of the
   ## distributions, in one of four forms depending on its dimension.
   ## @itemize
   ##   @item n-by-n-by-k: Slice @var{Sigma}(:,:,i) is the variance of the
   ##         i'th component
   ##   @item 1-by-n-by-k: Slice diag(@var{Sigma}(1,:,i)) is the variance of the
   ##         i'th component
   ##   @item n-by-n: @var{Sigma} is the variance of every component
   ##   @item 1-by-n-by-k: Slice diag(@var{Sigma}) is the variance of every
   ##         component
   ## @end itemize
   ##
   ## If @var{p} is specified, it is a vector of length k specifying the
   ## proportion of each component.  If it is omitted or empty, each component
   ## has an equal proportion.
   ##
   ## Input @var{extra} is used by fitgmdist to indicate the parameters of the
   ## fitting process.
   ## @seealso{fitgmdist}
   ## @end deftypefn

   properties
      mu                        ## means
      Sigma                     ## covariances
      ComponentProportion       ## mixing proportions
      DistributionName          ## "gaussian mixture distribution"
      NumComponents             ## Number of mixture components
      NumVariables              ## Dimension d of each Gaussian component

      CovarianceType            ## 'diagonal' if DiagonalCovariance, 'full' othw
      SharedCovariance          ## true if all components have equal covariance

      ## Set by a call to gmdistribution.fit  or  fitgmdist
      AIC                       ## Akaike Information Criterion
      BIC                       ## Bayes Information Criterion
      Converged                 ## true  if algorithm converged by MaxIter
      NegativeLogLikelihood     ## Negative of log-likelihood
      NlogL                     ## Negative of log-likelihood
      NumIterations             ## Number of iterations
      RegularizationValue       ## const added to diag of cov to make +ve def
   endproperties

   properties(Access = private)
      DiagonalCovariance        ## bool summary of "CovarianceType"
   endproperties

   methods
      ########################################
      ## Constructor
      function obj = gmdistribution (mu,sigma,p = [],extra = [])
        obj.DistributionName = 'gaussian mixture distribution';
        obj.mu = mu;
        obj.Sigma = sigma;
        obj.NumComponents = rows (mu);
        obj.NumVariables = columns (mu);
        if (isempty (p))
          obj.ComponentProportion = ones (1,obj.NumComponents) / ...
                                            obj.NumComponents;
        else
          if any (p < 0)
            error ("gmdistribution: component weights must be non-negative");
          endif
          s = sum (p);
          if (s == 0)
            error ("gmdistribution: component weights must not be all zero");
          elseif (s != 1)
            p = p / s;
          endif
          obj.ComponentProportion = p(:)';
        endif
        if (length (size (sigma)) == 3)
          obj.SharedCovariance = false;
        else
          obj.SharedCovariance = true;
        endif
        if (rows (sigma) == 1 && columns (mu) > 1)
          obj.DiagonalCovariance = true;
          obj.CovarianceType = 'diagonal';
        else
          obj.DiagonalCovariance = false;       ## full
          obj.CovarianceType = 'full';
        endif

        if (! isempty (extra))
          obj.AIC                   = extra.AIC;
          obj.BIC                   = extra.BIC;
          obj.Converged             = extra.Converged;
          obj.NegativeLogLikelihood = extra.NegativeLogLikelihood;
          obj.NlogL                 = extra.NegativeLogLikelihood;
          obj.NumIterations         = extra.NumIterations;
          obj.RegularizationValue   = extra.RegularizationValue;
        endif
      endfunction

      ########################################
      ## Cumulative distribution function for Gaussian mixture distribution
      function c = cdf (obj, X)
        X = checkX (obj, X, 'cdf');
        p_x_l = zeros (rows (X), obj.NumComponents);
        if (obj.SharedCovariance)
          if (obj.DiagonalCovariance)
            sig = diag (obj.Sigma);
          else
            sig = obj.Sigma;
          endif
        endif
        for i = 1:obj.NumComponents
          if (! obj.SharedCovariance)
            if (obj.DiagonalCovariance)
              sig = diag (obj.Sigma(:,:,i));
            else
              sig = obj.Sigma(:,:,i);
            endif
          endif
          p_x_l(:,i) = mvncdf (X,obj.mu(i,:),sig)*obj.ComponentProportion(i);
        endfor
        c = sum (p_x_l, 2);
      endfunction

      ########################################
      ## Construct clusters from Gaussian mixture distribution
      function [idx, nlogl, P, logpdf, M] = cluster (obj,X)
        X = checkX (obj, X, 'cluster');
        [p_x_l, M] = componentProb (obj, X);
        [~, idx] = max (p_x_l, [], 2);
        if (nargout >= 2)
          PDF = sum (p_x_l, 2);
          logpdf = log (PDF);
          nlogl = -sum (logpdf);
          if (nargout >= 3)
            P = bsxfun (@rdivide, p_x_l, PDF);
          endif
        endif
      endfunction

      ########################################
      ## Display Gaussian mixture distribution object
      function c = disp (obj)
        msg = ['Gaussian mixture distribution with %d ', ...
               "components in %d dimension(s)\n"];
        fprintf (msg, obj.NumComponents, columns (obj.mu));
        for i = 1:obj.NumComponents
          fprintf ("Clust %d: weight %d\n\tMean: ", ...
                          i, obj.ComponentProportion(i));
          fprintf ("%g ", obj.mu(i,:));
          fprintf ("\n");
          if (! obj.SharedCovariance)
            fprintf ("\tVariance:");
            if (! obj.DiagonalCovariance)
              if (columns (obj.mu) > 1)
                fprintf ("\n");
              endif
              disp (squeeze (obj.Sigma(:,:,i)))
            else
              fprintf (" diag(");
              fprintf ("%g ", obj.Sigma(:,:,i));
              fprintf (")\n");
            endif
          endif
        endfor
        if (obj.SharedCovariance)
          fprintf ("Shared variance\n");
          if (! obj.DiagonalCovariance)
            obj.Sigma
          else
            fprintf (" diag(");
            fprintf ("%g ", obj.Sigma);
            fprintf (")\n");
          endif
        endif
        if (! isempty (obj.AIC))
          fprintf ("AIC=%g BIC=%g NLogL=%g Iter=%d Cged=%d Reg=%g\n", ...
                    obj.AIC, obj.BIC, obj.NegativeLogLikelihood, ...
                    obj.NumIterations, obj.Converged, obj.RegularizationValue);
        endif
      endfunction

      ########################################
      ## Display Gaussian mixture distribution object
      function c = display (obj)
        disp (obj);
      endfunction

      ########################################
      ## Mahalanobis distance to component means
      function D = mahal (obj,X)
        X = checkX (obj, X, 'mahal');
        [~, D] = componentProb (obj,X);
      endfunction

      ########################################
      ## Probability density function for Gaussian mixture distribution
      function c = pdf (obj,X)
        X = checkX (obj, X, 'pdf');
        p_x_l = componentProb (obj, X);
        c = sum (p_x_l, 2);
      endfunction

      ########################################
      ## Posterior probabilities of components
      function c = posterior (obj,X)
        X = checkX (obj, X, 'posterior');
        p_x_l = componentProb (obj, X);
        c = bsxfun (@rdivide, p_x_l, sum (p_x_l, 2));
      endfunction

      ########################################
      ## Random numbers from Gaussian mixture distribution
      function c = random (obj,n)
        if nargin == 1
          n = 1;
        endif
        c = zeros (n, obj.NumVariables);
        classes = randsample (obj.NumComponents, n, true, ...
                              obj.ComponentProportion);
        if (obj.SharedCovariance)
          if (obj.DiagonalCovariance)
            sig = diag (obj.Sigma);
          else
            sig = obj.Sigma;
          endif
        endif
        for i = 1:obj.NumComponents
          idx = (classes == i);
          k = sum (idx);
          if (k > 0)
            if (! obj.SharedCovariance)
              if (obj.DiagonalCovariance)
                sig = diag (obj.Sigma(:,:,i));
              else
                sig = obj.Sigma(:,:,i);
              endif
            endif
            # [sig] forces [sig] not to have class "diagonal",
            # since mvnrnd uses automatic broadcast,
            # which fails on structured matrices
            c(idx,:) = mvnrnd ([obj.mu(i,:)], [sig], k);
          endif
        endfor
      endfunction
    endmethods

    ########################################
    methods(Static)
      ## Gaussian mixture parameter estimates
      function c = fit (X, k, varargin)
        c = fitgmdist (X, k, varargin{:});
      endfunction
    endmethods

    ########################################
    methods(Access = private)
      ## Probability density of (row of) X *and* component l
      ## Second argument is an array of the Mahalanobis distances
      function [p_x_l, M] = componentProb (obj, X)
        M     = zeros (rows (X), obj.NumComponents);
        dets  = zeros (1, obj.NumComponents);   % sqrt(determinant)
        if (obj.SharedCovariance)
          if (obj.DiagonalCovariance)
            r = diag (sqrt (obj.Sigma));
          else
            r = chol (obj.Sigma);
          endif
        endif
        for i = 1:obj.NumComponents
          dev = bsxfun (@minus, X, obj.mu(i,:));
          if (! obj.SharedCovariance)
            if (obj.DiagonalCovariance)
              r = diag (sqrt (obj.Sigma(:,:,i)));
            else
              r = chol (obj.Sigma(:,:,i));
            endif
          endif
          M(:,i) = sumsq (dev / r, 2);
          dets(i) = prod (diag (r));
        endfor
        p_x_l = exp (-M/2);
        coeff = obj.ComponentProportion ./ ...
                    ((2 * pi) ^ (obj.NumVariables / 2) .* dets);
        p_x_l = bsxfun (@times, p_x_l, coeff);
      endfunction

      ########################################
      ## Check format of argument X
      function X = checkX (obj, X, name)
        if (columns (X) != obj.NumVariables)
          if (columns (X) == 1 && rows (X) == obj.NumVariables)
            X = X';
          else
            error ("gmdistribution.%s: X has %d columns instead of %d\n", ...
                   name, columns (X), obj.NumVariables);
          endif
        endif
      endfunction
    endmethods
endclassdef

%!test
%! mu = eye (2);
%! Sigma = eye (2);
%! GM = gmdistribution (mu, Sigma);
%! density = GM.pdf ([0 0; 1 1]);
%! assert_equal (density(1) - density(2), 0, 1e-6);
%!
%! [idx, nlogl, P, logpdf,M] = cluster (GM, eye (2));
%! assert_equal (idx, [1; 2]);
%! [idx2,nlogl2,P2,logpdf2] = GM.cluster (eye (2));
%! assert_equal (nlogl - nlogl2, 0, 1e-6);
%! [idx3,nlogl3,P3] = cluster (GM, eye (2));
%! assert_equal (P - P3, zeros (2), 1e-6);
%! [idx4,nlogl4] = cluster (GM, eye (2));
%! assert_equal (size (nlogl4), [1 1]);
%! idx5 = cluster (GM, eye (2));
%! assert_equal (idx - idx5, zeros (2,1));
%!
%! D = GM.mahal ([1;0]);
%! assert_equal (D - M(1,:), zeros (1,2), 1e-6);
%!
%! P = GM.posterior ([0 1]);
%! assert_equal (P - P2(2,:), zeros (1,2), 1e-6);
%!
%! R = GM.random(20);
%! assert_equal (size (R), [20, 2]);
%!
%! R = GM.random();
%! assert_equal (size (R), [1, 2]);

