/**
 * 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;
	}
}
