/**
 * 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.linear.projection;

import gnu.trove.list.array.TDoubleArrayList;

import org.openimaj.math.matrix.MatrixUtils;
import org.openimaj.math.matrix.algorithm.pca.ThinSvdPrincipalComponentAnalysis;

import Jama.Matrix;

/**
 * {@link LargeMarginDimensionalityReduction} is a technique to compress high
 * dimensional features into a lower-dimension representation using a learned
 * linear projection. Supervised learning is used to learn the projection such
 * that the squared Euclidean distance between two low-dimensional vectors is
 * less than a threshold if they correspond to the same object, or greater
 * otherwise. In addition, it is imposed that this condition is satisfied with a
 * margin of at least one.
 * <p>
 * In essence, the Euclidean distance in the low dimensional space produced by
 * this technique can be seen as a low-rank Mahalanobis metric in the original
 * space; the Mahalanobis matrix would have rank equal to the number of
 * dimensions of the smaller space.
 * <p>
 * This class implements the technique using stochastic sub-gradient descent. As
 * the objective function is not convex, initialisation is important, and
 * initial conditions are generated by selecting the largest PCA dimensions, and
 * then whitening the dimensions so they have equal magnitude. In addition, the
 * projection matrix is not regularised explicitly; instead the algorithm is
 * just stopped after a fixed number of iterations.
 * 
 * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
 */
public class LargeMarginDimensionalityReduction {
        protected int ndims;
        protected double wLearnRate = 0.25; // gamma
        protected double bLearnRate = 1; // bias gamma
        protected Matrix W;
        protected double b;

        /**
         * Construct with the given target dimensionality and default values for the
         * other parameters (learning rate of 1.0 for W; learning rate of 0 for
         * bias).
         * 
         * @param ndims
         *            target number of dimensions
         */
        public LargeMarginDimensionalityReduction(int ndims) {
                this.ndims = ndims;
        }

        /**
         * Construct with the given target dimensionality learning rates.
         * 
         * @param ndims
         *            target number of dimensions
         * @param wLearnRate
         *            learning rate for the transform matrix, W.
         * @param bLearnRate
         *            learning rate for the bias, b.
         */
        public LargeMarginDimensionalityReduction(int ndims, double wLearnRate, double bLearnRate) {
                this.ndims = ndims;
        }

        /**
         * Initialise the LMDR with the given data in three parallel data arrays.
         * The first two arrays together store pairs of vectors, and the third array
         * indicate whether the vectors in each pair represent the same class or
         * different classes. All three arrays should obviously have the same
         * length, and approximately equal numbers of true and false pairs.
         * <p>
         * Internally, this method gathers together all the vectors and performs
         * {@link ThinSvdPrincipalComponentAnalysis} to get a starting estimate for
         * the transform matrix. The distance in the projected space between all
         * true pairs and false pairs is computed and used to select an optimal
         * initial threshold.
         * 
         * @param datai
         *            array of the first vectors of the pairs
         * @param dataj
         *            array of the second vectors of the pairs
         * @param same
         *            array indicating where pairs are true (from the same class) or
         *            false (from different classes)
         */
        public void initialise(double[][] datai, double[][] dataj, boolean[] same) {
                final double[][] data = new double[2 * datai.length][];
                for (int i = 0; i < datai.length; i++) {
                        data[2 * i] = datai[i];
                        data[2 * i + 1] = dataj[i];
                }

                final ThinSvdPrincipalComponentAnalysis pca = new ThinSvdPrincipalComponentAnalysis(ndims);
                pca.learnBasis(data);

                final double[] evs = pca.getEigenValues();
                final double[] invStdDev = new double[ndims];
                for (int i = 0; i < ndims; i++)
                        invStdDev[i] = 1.0 / Math.sqrt(evs[i]);

                W = MatrixUtils.diag(invStdDev).times(pca.getBasis().transpose());

                recomputeBias(datai, dataj, same);
        }

        public void recomputeBias(double[][] datai, double[][] dataj, boolean[] same) {
                final TDoubleArrayList posDistances = new TDoubleArrayList();
                final TDoubleArrayList negDistances = new TDoubleArrayList();
                for (int i = 0; i < datai.length; i++) {
                        final Matrix diff = diff(datai[i], dataj[i]);
                        final Matrix diffProj = W.times(diff);
                        final double dist = sumsq(diffProj);

                        if (same[i]) {
                                posDistances.add(dist);
                        } else {
                                negDistances.add(dist);
                        }
                }

                b = computeOptimal(posDistances, negDistances);
        }

        private double computeOptimal(TDoubleArrayList posDistances, TDoubleArrayList negDistances) {
                double bestAcc = 0;
                double bestThresh = -Double.MAX_VALUE;
                for (int i = 0; i < posDistances.size(); i++) {
                        final double thresh = posDistances.get(i);

                        final double acc = computeAccuracy(posDistances, negDistances, thresh);

                        if (acc > bestAcc) {
                                bestAcc = acc;
                                bestThresh = thresh;
                        }
                }

                for (int i = 0; i < negDistances.size(); i++) {
                        final double thresh = negDistances.get(i);

                        final double acc = computeAccuracy(posDistances, negDistances, thresh);

                        if (acc > bestAcc) {
                                bestAcc = acc;
                                bestThresh = thresh;
                        }
                }

                return bestThresh;
        }

        private double computeAccuracy(TDoubleArrayList posDistances, TDoubleArrayList negDistances, double thresh) {
                int correct = 0;
                for (int i = 0; i < posDistances.size(); i++) {
                        if (posDistances.get(i) < thresh)
                                correct++;
                }

                for (int i = 0; i < negDistances.size(); i++) {
                        if (negDistances.get(i) >= thresh)
                                correct++;
                }

                return (double) correct / (double) (posDistances.size() + negDistances.size());
        }

        private Matrix diff(double[] phii, double[] phij) {
                final Matrix diff = new Matrix(phii.length, 1);
                final double[][] diffv = diff.getArray();

                for (int i = 0; i < phii.length; i++) {
                        diffv[i][0] = phii[i] - phij[i];
                }
                return diff;
        }

        private double sumsq(Matrix diffProj) {
                final double[][] v = diffProj.getArray();

                double sumsq = 0;
                for (int i = 0; i < v.length; i++) {
                        sumsq += v[i][0] * v[i][0];
                }

                return sumsq;
        }

        /**
         * Perform a single update step of the SGD optimisation. Alternate calls to
         * this method should swap between true and false pairs.
         * 
         * @param phii
         *            first vector
         * @param phij
         *            second vector
         * @param same
         *            true if the vectors are from the same class; false otherwise
         * @return true if the transform matrix was changed; false otherwise
         */
        public boolean step(double[] phii, double[] phij, boolean same) {
                final int yij = same ? 1 : -1;

                final Matrix diff = diff(phii, phij);
                final Matrix diffProj = W.times(diff);
                final double sumsq = sumsq(diffProj);

                if (yij * (b - sumsq) > 1)
                        return false;

                // final Matrix updateW = diffProj.times(wLearnRate *
                // yij).times(diff.transpose());
                // W.minusEquals(updateW);
                fastUpdate(diffProj, wLearnRate * yij, diff);

                b += yij * bLearnRate;

                return true;
        }

        /**
         * This efficiently computes the update in place without creating loads of
         * temporary matrices, and does so in a single pass!
         * 
         * @param diffProj
         * @param weight
         * @param diff
         */
        private void fastUpdate(Matrix diffProj, double weight, Matrix diff) {
                // final Matrix updateW = diffProj.times(wLearnRate *
                // yij).times(diff.transpose());
                // W.minusEquals(updateW);

                final double[][] dp = diffProj.getArray();
                final double[][] d = diff.getArray();
                final double[][] Wdata = W.getArray();
                for (int r = 0; r < Wdata.length; r++) {
                        for (int c = 0; c < Wdata.length; c++) {
                                Wdata[r][c] -= weight * dp[r][0] * d[c][0];
                        }
                }
        }

        /**
         * Get the transform matrix W
         * 
         * @return the transform matrix
         */
        public Matrix getTransform() {
                return W;
        }

        /**
         * Get the bias, b
         * 
         * @return the bias
         */
        public double getBias() {
                return b;
        }

        /**
         * Compute the matching score between a pair of (high dimensional) features.
         * Scores >=0 indicate a matching pair.
         * 
         * @param phii
         *            first vector
         * @param phij
         *            second vector
         * @return the matching score
         */
        public double score(double[] phii, double[] phij) {
                final Matrix diff = diff(phii, phij);
                final Matrix diffProj = W.times(diff);

                return b - sumsq(diffProj);
        }

        /**
         * Determine if two features are from the same class or different classes.
         * 
         * @param phii
         *            first vector
         * @param phij
         *            second vector
         * @return the classification (true if same class; false if different)
         */
        public boolean classify(double[] phii, double[] phij) {
                return score(phii, phij) >= 0;
        }

        /**
         * Compute the low rank estimate of the given vector
         * 
         * @param in
         *            the vector
         * @return the low-rank projection of the vector
         */
        public double[] project(double[] in) {
                return W.times(new Matrix(new double[][] { in }).transpose()).getColumnPackedCopy();
        }

        public void setBias(double d) {
                this.b = d;
        }

        public void setTransform(Matrix proj) {
                this.W = proj;
        }
}