/**
 * Copyright (c) 2011, The University of Southampton and the individual contributors.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 *   *  Redistributions of source code must retain the above copyright notice,
 *      this list of conditions and the following disclaimer.
 *
 *   *  Redistributions in binary form must reproduce the above copyright notice,
 *      this list of conditions and the following disclaimer in the documentation
 *      and/or other materials provided with the distribution.
 *
 *   *  Neither the name of the University of Southampton nor the names of its
 *      contributors may be used to endorse or promote products derived from this
 *      software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
package org.openimaj.ml.gmm;

import java.util.Arrays;
import java.util.EnumSet;

import org.apache.commons.math.util.MathUtils;
import org.openimaj.math.matrix.MatrixUtils;
import org.openimaj.math.statistics.MeanAndCovariance;
import org.openimaj.math.statistics.distribution.AbstractMultivariateGaussian;
import org.openimaj.math.statistics.distribution.DiagonalMultivariateGaussian;
import org.openimaj.math.statistics.distribution.FullMultivariateGaussian;
import org.openimaj.math.statistics.distribution.MixtureOfGaussians;
import org.openimaj.math.statistics.distribution.MultivariateGaussian;
import org.openimaj.math.statistics.distribution.SphericalMultivariateGaussian;
import org.openimaj.ml.clustering.DoubleCentroidsResult;
import org.openimaj.ml.clustering.kmeans.DoubleKMeans;
import org.openimaj.util.array.ArrayUtils;
import org.openimaj.util.pair.IndependentPair;

import Jama.Matrix;
import gnu.trove.list.array.TDoubleArrayList;

/**
 * Gaussian mixture model learning using the EM algorithm. Supports a range of
 * different shapes Gaussian through different covariance matrix forms. An
 * initialisation step is used to learn the initial means using K-Means,
 * although this can be disabled in the constructor.
 * <p>
 * Implementation was originally inspired by the SciPy's "gmm.py".
 *
 * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
 */
public class GaussianMixtureModelEM {
        /**
         * Different forms of covariance matrix supported by the
         * {@link GaussianMixtureModelEM}.
         *
         * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
         */
        public static enum CovarianceType {
                /**
                 * Spherical Gaussians: variance is the same along all axes and zero
                 * across-axes.
                 */
                Spherical {
                        @Override
                        protected void setCovariances(MultivariateGaussian[] gaussians, Matrix cv) {
                                double mean = 0;

                                for (int i = 0; i < cv.getRowDimension(); i++)
                                        for (int j = 0; j < cv.getColumnDimension(); j++)
                                                mean += cv.get(i, j);
                                mean /= (cv.getColumnDimension() * cv.getRowDimension());

                                for (final MultivariateGaussian mg : gaussians) {
                                        ((SphericalMultivariateGaussian) mg).variance = mean;
                                }
                        }

                        @Override
                        protected MultivariateGaussian[] createGaussians(int ngauss, int ndims) {
                                final MultivariateGaussian[] arr = new MultivariateGaussian[ngauss];
                                for (int i = 0; i < ngauss; i++) {
                                        arr[i] = new SphericalMultivariateGaussian(ndims);
                                }

                                return arr;
                        }

                        @Override
                        protected void mstep(EMGMM gmm, GaussianMixtureModelEM learner, Matrix X, Matrix responsibilities,
                                        Matrix weightedXsum,
                                        double[] norm)
                {
                                final Matrix avgX2uw = responsibilities.transpose().times(X.arrayTimes(X));

                                for (int i = 0; i < gmm.gaussians.length; i++) {
                                        final Matrix weightedXsumi = new Matrix(new double[][] { weightedXsum.getArray()[i] });
                                        final Matrix avgX2uwi = new Matrix(new double[][] { avgX2uw.getArray()[i] });

                                        final Matrix avgX2 = avgX2uwi.times(norm[i]);
                                        final Matrix mu = ((AbstractMultivariateGaussian) gmm.gaussians[i]).mean;
                                        final Matrix avgMeans2 = MatrixUtils.pow(mu, 2);
                                        final Matrix avgXmeans = mu.arrayTimes(weightedXsumi).times(norm[i]);
                                        final Matrix covar = MatrixUtils.plus(avgX2.minus(avgXmeans.times(2)).plus(avgMeans2),
                                                        learner.minCovar);

                                        ((SphericalMultivariateGaussian) gmm.gaussians[i]).variance = MatrixUtils.sum(covar)
                                                        / X.getColumnDimension();
                                }
                        }
                },
                /**
                 * Gaussians with diagonal covariance matrices.
                 */
                Diagonal {
                        @Override
                        protected void setCovariances(MultivariateGaussian[] gaussians, Matrix cv) {
                                for (final MultivariateGaussian mg : gaussians) {
                                        ((DiagonalMultivariateGaussian) mg).variance = MatrixUtils.diagVector(cv);
                                }
                        }

                        @Override
                        protected MultivariateGaussian[] createGaussians(int ngauss, int ndims) {
                                final MultivariateGaussian[] arr = new MultivariateGaussian[ngauss];
                                for (int i = 0; i < ngauss; i++) {
                                        arr[i] = new DiagonalMultivariateGaussian(ndims);
                                }

                                return arr;
                        }

                        @Override
                        protected void mstep(EMGMM gmm, GaussianMixtureModelEM learner, Matrix X, Matrix responsibilities,
                                        Matrix weightedXsum,
                                        double[] norm)
                {
                                final Matrix avgX2uw = responsibilities.transpose().times(X.arrayTimes(X));

                                for (int i = 0; i < gmm.gaussians.length; i++) {
                                        final Matrix weightedXsumi = new Matrix(new double[][] { weightedXsum.getArray()[i] });
                                        final Matrix avgX2uwi = new Matrix(new double[][] { avgX2uw.getArray()[i] });

                                        final Matrix avgX2 = avgX2uwi.times(norm[i]);
                                        final Matrix mu = ((AbstractMultivariateGaussian) gmm.gaussians[i]).mean;
                                        final Matrix avgMeans2 = MatrixUtils.pow(mu, 2);
                                        final Matrix avgXmeans = mu.arrayTimes(weightedXsumi).times(norm[i]);

                                        final Matrix covar = MatrixUtils.plus(avgX2.minus(avgXmeans.times(2)).plus(avgMeans2),
                                                        learner.minCovar);

                                        ((DiagonalMultivariateGaussian) gmm.gaussians[i]).variance = covar.getArray()[0];
                                }
                        }
                },
                /**
                 * Gaussians with full covariance
                 */
                Full {
                        @Override
                        protected MultivariateGaussian[] createGaussians(int ngauss, int ndims) {
                                final MultivariateGaussian[] arr = new MultivariateGaussian[ngauss];
                                for (int i = 0; i < ngauss; i++) {
                                        arr[i] = new FullMultivariateGaussian(ndims);
                                }

                                return arr;
                        }

                        @Override
                        protected void setCovariances(MultivariateGaussian[] gaussians, Matrix cv) {
                                for (final MultivariateGaussian mg : gaussians) {
                                        ((FullMultivariateGaussian) mg).covar = cv.copy();
                                }
                        }

                        @Override
                        protected void mstep(EMGMM gmm, GaussianMixtureModelEM learner, Matrix X, Matrix responsibilities,
                                        Matrix weightedXsum,
                                        double[] norm)
                {
                                // Eq. 12 from K. Murphy,
                                // "Fitting a Conditional Linear Gaussian Distribution"
                                final int nfeatures = X.getColumnDimension();
                                for (int c = 0; c < learner.nComponents; c++) {
                                        final Matrix post = responsibilities.getMatrix(0, X.getRowDimension() - 1, c, c).transpose();

                                        final double factor = 1.0 / (ArrayUtils.sumValues(post.getArray()) + 10 * MathUtils.EPSILON);

                                        final Matrix pXt = X.transpose();
                                        for (int i = 0; i < pXt.getRowDimension(); i++)
                                                for (int j = 0; j < pXt.getColumnDimension(); j++)
                                                        pXt.set(i, j, pXt.get(i, j) * post.get(0, j));

                                        final Matrix argcv = pXt.times(X).times(factor);
                                        final Matrix mu = ((FullMultivariateGaussian) gmm.gaussians[c]).mean;

                                        ((FullMultivariateGaussian) gmm.gaussians[c]).covar = argcv.minusEquals(mu.transpose().times(mu))
                                                        .plusEquals(Matrix.identity(nfeatures, nfeatures).times(learner.minCovar));
                                }
                        }
                },
                /**
                 * Gaussians with a tied covariance matrix; the same covariance matrix
                 * is shared by all the gaussians.
                 */
                Tied {
                        // @Override
                        // protected double[][] logProbability(double[][] x,
                        // MultivariateGaussian[] gaussians)
                        // {
                        // final int ndim = x[0].length;
                        // final int nmix = gaussians.length;
                        // final int nsamples = x.length;
                        // final Matrix X = new Matrix(x);
                        //
                        // final double[][] logProb = new double[nsamples][nmix];
                        // final Matrix cv = ((FullMultivariateGaussian)
                        // gaussians[0]).covar;
                        //
                        // final CholeskyDecomposition chol = cv.chol();
                        // Matrix cvChol;
                        // if (chol.isSPD()) {
                        // cvChol = chol.getL();
                        // } else {
                        // // covar probably doesn't have enough samples, so
                        // // recondition it
                        // final Matrix m = cv.plus(Matrix.identity(ndim, ndim).timesEquals(
                        // MixtureOfGaussians.MIN_COVAR_RECONDITION));
                        // cvChol = m.chol().getL();
                        // }
                        //
                        // double cvLogDet = 0;
                        // final double[][] cvCholD = cvChol.getArray();
                        // for (int j = 0; j < ndim; j++) {
                        // cvLogDet += Math.log(cvCholD[j][j]);
                        // }
                        // cvLogDet *= 2;
                        //
                        // for (int i = 0; i < nmix; i++) {
                        // final Matrix mu = ((FullMultivariateGaussian) gaussians[i]).mean;
                        // final Matrix cvSol = cvChol.solve(MatrixUtils.minusRow(X,
                        // mu.getArray()[0]).transpose())
                        // .transpose();
                        // for (int k = 0; k < nsamples; k++) {
                        // double sum = 0;
                        // for (int j = 0; j < ndim; j++) {
                        // sum += cvSol.get(k, j) * cvSol.get(k, j);
                        // }
                        //
                        // logProb[k][i] = -0.5 * (sum + cvLogDet + ndim * Math.log(2 *
                        // Math.PI));
                        // }
                        // }
                        //
                        // return logProb;
                        // }

                        @Override
                        protected void setCovariances(MultivariateGaussian[] gaussians,
                                        Matrix cv)
                {
                                for (final MultivariateGaussian mg : gaussians) {
                                        ((FullMultivariateGaussian) mg).covar = cv;
                                }
                        }

                        @Override
                        protected MultivariateGaussian[] createGaussians(int ngauss, int ndims) {
                                final MultivariateGaussian[] arr = new MultivariateGaussian[ngauss];
                                final Matrix covar = new Matrix(ndims, ndims);

                                for (int i = 0; i < ngauss; i++) {
                                        arr[i] = new FullMultivariateGaussian(new Matrix(1, ndims), covar);
                                }

                                return arr;
                        }

                        @Override
                        protected void mstep(EMGMM gmm, GaussianMixtureModelEM learner, Matrix X, Matrix responsibilities,
                                        Matrix weightedXsum, double[] norm)
                {
                                // Eq. 15 from K. Murphy, "Fitting a Conditional Linear Gaussian
                                final int nfeatures = X.getColumnDimension();

                                final Matrix avgX2 = X.transpose().times(X);
                                final double[][] mudata = new double[gmm.gaussians.length][];
                                for (int i = 0; i < mudata.length; i++)
                                        mudata[i] = ((FullMultivariateGaussian) gmm.gaussians[i]).mean.getArray()[0];
                                final Matrix mu = new Matrix(mudata);

                                final Matrix avgMeans2 = mu.transpose().times(weightedXsum);
                                final Matrix covar = avgX2.minus(avgMeans2)
                                                .plus(Matrix.identity(nfeatures, nfeatures).times(learner.minCovar))
                                                .times(1.0 / X.getRowDimension());

                                for (int i = 0; i < learner.nComponents; i++)
                                        ((FullMultivariateGaussian) gmm.gaussians[i]).covar = covar;
                        }
                };

                protected abstract MultivariateGaussian[] createGaussians(int ngauss, int ndims);

                protected abstract void setCovariances(MultivariateGaussian[] gaussians, Matrix cv);

                /**
                 * Mode specific maximisation-step. Implementors should use the state to
                 * update the covariance of each of the
                 * {@link GaussianMixtureModelEM#gaussians}.
                 *
                 * @param gmm
                 *            the mixture model being learned
                 * @param X
                 *            the data
                 * @param responsibilities
                 *            matrix with the same number of rows as X where each col is
                 *            the amount that the data point belongs to each gaussian
                 * @param weightedXsum
                 *            responsibilities.T * X
                 * @param inverseWeights
                 *            1/weights
                 */
                protected abstract void mstep(EMGMM gmm, GaussianMixtureModelEM learner, Matrix X,
                                Matrix responsibilities, Matrix weightedXsum, double[] inverseWeights);
        }

        /**
         * Options for controlling what gets updated during the initialisation
         * and/or iterations.
         *
         * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
         */
        public static enum UpdateOptions {
                /**
                 * Update the means
                 */
                Means,
                /**
                 * Update the weights
                 */
                Weights,
                /**
                 * Update the covariances
                 */
                Covariances
        }

        protected static class EMGMM extends MixtureOfGaussians {
                EMGMM(int nComponents) {
                        super(null, null);

                        this.weights = new double[nComponents];
                        Arrays.fill(this.weights, 1.0 / nComponents);
                }
        }

        private static final double DEFAULT_THRESH = 1e-2;
        private static final double DEFAULT_MIN_COVAR = 1e-3;
        private static final int DEFAULT_NITERS = 100;
        private static final int DEFAULT_NINIT = 1;

        CovarianceType ctype;
        int nComponents;
        private double thresh;
        private double minCovar;
        private int nIters;
        private int nInit;

        private boolean converged = false;
        private EnumSet<UpdateOptions> initOpts;
        private EnumSet<UpdateOptions> iterOpts;

        /**
         * Construct with the given arguments.
         *
         * @param nComponents
         *            the number of gaussian components
         * @param ctype
         *            the form of the covariance matrices
         * @param thresh
         *            the threshold at which to stop iterating
         * @param minCovar
         *            the minimum value allowed in the diagonal of the estimated
         *            covariance matrices to prevent overfitting
         * @param nIters
         *            the maximum number of iterations
         * @param nInit
         *            the number of runs of the algorithm to perform; the best
         *            result will be kept.
         * @param iterOpts
         *            options controlling what is updated during iteration
         * @param initOpts
         *            options controlling what is updated during initialisation.
         *            Enabling the {@link UpdateOptions#Means} option will cause
         *            K-Means to be used to generate initial starting points for the
         *            means.
         */
        public GaussianMixtureModelEM(int nComponents, CovarianceType ctype, double thresh, double minCovar,
                        int nIters, int nInit, EnumSet<UpdateOptions> iterOpts, EnumSet<UpdateOptions> initOpts)
        {
                this.ctype = ctype;
                this.nComponents = nComponents;
                this.thresh = thresh;
                this.minCovar = minCovar;
                this.nIters = nIters;
                this.nInit = nInit;
                this.iterOpts = iterOpts;
                this.initOpts = initOpts;

                if (nInit < 1) {
                        throw new IllegalArgumentException("GMM estimation requires at least one run");
                }
                this.converged = false;
        }

        /**
         * Construct with the given arguments.
         *
         * @param nComponents
         *            the number of gaussian components
         * @param ctype
         *            the form of the covariance matrices
         */
        public GaussianMixtureModelEM(int nComponents, CovarianceType ctype) {
                this(nComponents, ctype, DEFAULT_THRESH, DEFAULT_MIN_COVAR, DEFAULT_NITERS, DEFAULT_NINIT, EnumSet
                                .allOf(UpdateOptions.class), EnumSet.allOf(UpdateOptions.class));
        }

        /**
         * Get's the convergence state of the algorithm. Will return false if
         * {@link #estimate(double[][])} has not been called, or if the last call to
         * {@link #estimate(double[][])} failed to reach convergence before running
         * out of iterations.
         *
         * @return true if the last call to {@link #estimate(double[][])} reached
         *         convergence; false otherwise
         */
        public boolean hasConverged() {
                return converged;
        }

        /**
         * Estimate a new {@link MixtureOfGaussians} from the given data. Use
         * {@link #hasConverged()} to check whether the EM algorithm reached
         * convergence in the estimation of the returned model.
         *
         * @param X
         *            the data matrix.
         * @return the generated GMM.
         */
        public MixtureOfGaussians estimate(Matrix X) {
                return estimate(X.getArray());
        }

        /**
         * Estimate a new {@link MixtureOfGaussians} from the given data. Use
         * {@link #hasConverged()} to check whether the EM algorithm reached
         * convergence in the estimation of the returned model.
         *
         * @param X
         *            the data array.
         * @return the generated GMM.
         */
        public MixtureOfGaussians estimate(double[][] X) {
                final EMGMM gmm = new EMGMM(nComponents);

                if (X.length < nComponents)
                        throw new IllegalArgumentException(String.format(
                                        "GMM estimation with %d components, but got only %d samples", nComponents, X.length));

                double max_log_prob = Double.NEGATIVE_INFINITY;

                for (int j = 0; j < nInit; j++) {
                        gmm.gaussians = ctype.createGaussians(nComponents, X[0].length);

                        if (initOpts.contains(UpdateOptions.Means)) {
                                // initialise using k-means
                                final DoubleKMeans km = DoubleKMeans.createExact(nComponents);
                                final DoubleCentroidsResult means = km.cluster(X);

                                for (int i = 0; i < nComponents; i++) {
                                        ((AbstractMultivariateGaussian) gmm.gaussians[i]).mean.getArray()[0] = means.centroids[i];
                                }
                        }

                        if (initOpts.contains(UpdateOptions.Weights)) {
                                gmm.weights = new double[nComponents];
                                Arrays.fill(gmm.weights, 1.0 / nComponents);
                        }

                        if (initOpts.contains(UpdateOptions.Covariances)) {
                                // cv = np.cov(X.T) + self.min_covar * np.eye(X.shape[1])
                                final Matrix cv = MeanAndCovariance.computeCovariance(X);

                                ctype.setCovariances(gmm.gaussians, cv);
                        }

                        // EM algorithm
                        final TDoubleArrayList log_likelihood = new TDoubleArrayList();

                        // reset converged to false
                        converged = false;
                        double[] bestWeights = null;
                        MultivariateGaussian[] bestMixture = null;
                        for (int i = 0; i < nIters; i++) {
                                // Expectation step
                                final IndependentPair<double[], double[][]> score = gmm.scoreSamples(X);
                                final double[] curr_log_likelihood = score.firstObject();
                                final double[][] responsibilities = score.secondObject();
                                log_likelihood.add(ArrayUtils.sumValues(curr_log_likelihood));

                                // Check for convergence.
                                if (i > 0 && Math.abs(log_likelihood.get(i) - log_likelihood.get(i - 1)) < thresh) {
                                        converged = true;
                                        break;
                                }

                                // Perform the maximisation step
                                mstep(gmm, X, responsibilities);

                                // if the results are better, keep it
                                if (nIters > 0) {
                                        if (log_likelihood.getQuick(i) > max_log_prob) {
                                                max_log_prob = log_likelihood.getQuick(i);
                                                bestWeights = gmm.weights;
                                                bestMixture = gmm.gaussians;
                                        }
                                }

                                // check the existence of an init param that was not subject to
                                // likelihood computation issue.
                                if (Double.isInfinite(max_log_prob) && nIters > 0) {
                                        throw new RuntimeException(
                                                        "EM algorithm was never able to compute a valid likelihood given initial " +
                                                                        "parameters. Try different init parameters (or increasing n_init) or " +
                                                                        "check for degenerate data.");
                                }

                                if (nIters > 0) {
                                        gmm.gaussians = bestMixture;
                                        gmm.weights = bestWeights;
                                }
                        }
                }

                return gmm;
        }

        protected void mstep(EMGMM gmm, double[][] X, double[][] responsibilities) {
                final double[] weights = ArrayUtils.colSum(responsibilities);
                final Matrix resMat = new Matrix(responsibilities);
                final Matrix Xmat = new Matrix(X);

                final Matrix weighted_X_sum = resMat.transpose().times(Xmat);
                final double[] inverse_weights = new double[weights.length];
                for (int i = 0; i < inverse_weights.length; i++)
                        inverse_weights[i] = 1.0 / (weights[i] + 10 * MathUtils.EPSILON);

                if (iterOpts.contains(UpdateOptions.Weights)) {
                        final double sum = ArrayUtils.sumValues(weights);
                        for (int i = 0; i < weights.length; i++) {
                                gmm.weights[i] = (weights[i] / (sum + 10 * MathUtils.EPSILON) + MathUtils.EPSILON);
                        }
                }

                if (iterOpts.contains(UpdateOptions.Means)) {
                        // self.means_ = weighted_X_sum * inverse_weights
                        final double[][] wx = weighted_X_sum.getArray();

                        for (int i = 0; i < nComponents; i++) {
                                final double[][] m = ((AbstractMultivariateGaussian) gmm.gaussians[i]).mean.getArray();

                                for (int j = 0; j < m[0].length; j++) {
                                        m[0][j] = wx[i][j] * inverse_weights[i];
                                }
                        }
                }

                if (iterOpts.contains(UpdateOptions.Covariances)) {
                        ctype.mstep(gmm, this, Xmat, resMat, weighted_X_sum, inverse_weights);
                }
        }

        @Override
        public GaussianMixtureModelEM clone() {
                try {
                        return (GaussianMixtureModelEM) super.clone();
                } catch (final CloneNotSupportedException e) {
                        throw new RuntimeException(e);
                }
        }
}