/*
        AUTOMATICALLY GENERATED BY jTemp FROM
        /Users/jsh2/Work/openimaj/target/checkout/machine-learning/clustering/src/main/jtemp/org/openimaj/ml/clustering/kmeans/#T#KMeans.jtemp
*/
/**
 * 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.clustering.kmeans;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;

import org.openimaj.data.DataSource;
import org.openimaj.data.ShortArrayBackedDataSource;
import org.openimaj.ml.clustering.IndexClusters;
import org.openimaj.ml.clustering.SpatialClusterer;
import org.openimaj.ml.clustering.assignment.HardAssigner;
import org.openimaj.ml.clustering.assignment.hard.KDTreeShortEuclideanAssigner;
import org.openimaj.ml.clustering.assignment.hard.ExactShortAssigner;
import org.openimaj.ml.clustering.ShortCentroidsResult;
import org.openimaj.knn.ShortNearestNeighbours;
import org.openimaj.knn.ShortNearestNeighboursExact;
import org.openimaj.knn.ShortNearestNeighboursProvider;
import org.openimaj.knn.NearestNeighboursFactory;
import org.openimaj.knn.approximate.ShortNearestNeighboursKDTree;
import org.openimaj.util.pair.IntFloatPair;

/**
 * Fast, parallel implementation of the K-Means algorithm with support for
 * bigger-than-memory data. Various flavors of K-Means are supported through the
 * selection of different subclasses of {@link ShortNearestNeighbours}; for
 * example, approximate K-Means can be achieved using a
 * {@link ShortNearestNeighboursKDTree} whilst exact K-Means can be achieved
 * using an {@link ShortNearestNeighboursExact}. The specific choice of
 * nearest-neighbour object is controlled through the
 * {@link NearestNeighboursFactory} provided to the {@link KMeansConfiguration}
 * used to construct instances of this class. The choice of
 * {@link ShortNearestNeighbours} affects the speed of clustering; using
 * approximate nearest-neighbours algorithms for the K-Means can produces
 * comparable results to the exact KMeans algorithm in much shorter time.
 * The choice and configuration of {@link ShortNearestNeighbours} can also
 * control the type of distance function being used in the clustering.
 * <p>
 * The algorithm is implemented as follows: Clustering is initiated using a
 * {@link ShortKMeansInit} and is iterative. In each round, batches of
 * samples are assigned to centroids in parallel. The centroid assignment is
 * performed using the pre-configured {@link ShortNearestNeighbours} instances
 * created from the {@link KMeansConfiguration}. Once all samples are assigned
 * new centroids are calculated and the next round started. Data point pushing
 * is performed using the same techniques as center point assignment.
 * <p>
 * This implementation is able to deal with larger-than-memory datasets by
 * streaming the samples from disk using an appropriate {@link DataSource}. The
 * only requirement is that there is enough memory to hold all the centroids
 * plus working memory for the batches of samples being assigned.
 * 
 * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
 * @author Sina Samangooei (ss@ecs.soton.ac.uk)
 */
 public class ShortKMeans implements SpatialClusterer<ShortCentroidsResult, short[]> {
        private static class CentroidAssignmentJob implements Callable<Boolean> {
                private final DataSource<short[]> ds;
                private final int startRow;
                private final int stopRow;
                private final ShortNearestNeighbours nno;
                private final float [][] centroids_accum;
                private final int [] counts;

                public CentroidAssignmentJob(DataSource<short[]> ds, int startRow, int stopRow, ShortNearestNeighbours nno, float [][] centroids_accum, int [] counts) {
                        this.ds = ds; 
                        this.startRow = startRow;
                        this.stopRow = stopRow;
                        this.nno = nno;
                        this.centroids_accum = centroids_accum;
                        this.counts = counts;
                }
                
                @Override
                public Boolean call() {
                        try {
                                int D = nno.numDimensions();

                                short [][] points = new short[stopRow-startRow][D]; 
                                ds.getData(startRow, stopRow, points);

                                int [] argmins = new int[points.length];
                                float [] mins = new float[points.length];

                                nno.searchNN(points, argmins, mins);

                                synchronized(centroids_accum){
                                        for (int i=0; i < points.length; ++i) {
                                                int k = argmins[i];
                                                for (int d=0; d < D; ++d) {
                                                        centroids_accum[k][d] += points[i][d];
                                                }
                                                counts[k] += 1;
                                        }
                                }
                        } catch(Exception e) {
                                e.printStackTrace();
                        }
                        return true;
                }
        }
        
        /**
         * Result object for ShortKMeans, extending ShortCentroidsResult and ShortNearestNeighboursProvider,
         * as well as giving access to state information from the operation of the K-Means algorithm  
         * (i.e. number of iterations, and convergence state).
         *
         * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
         */
        public static class Result extends ShortCentroidsResult implements ShortNearestNeighboursProvider {
                protected ShortNearestNeighbours nn;
                protected int iterations;
                protected int changedCentroidCount;
                 
                @Override
                public HardAssigner<short[], float[], IntFloatPair> defaultHardAssigner() {
                        if (nn instanceof ShortNearestNeighboursExact)
                                return new ExactShortAssigner(this, ((ShortNearestNeighboursExact)nn).distanceComparator());
                
                        return new KDTreeShortEuclideanAssigner(this);
                }
                
                @Override
                public ShortNearestNeighbours getNearestNeighbours() {
                        return nn;
                }
                
                /**
                 * Get the number of K-Means iterations that produced this result.
                 * @return the number of iterations
                 */
                public int numIterations() {
                        return iterations;
                }
                
                /**
                 * Get the number of changed centroids in the last iteration. This is 
                 * an indicator of convergence as over time this should reduce to 0.
                 * @return the number of changed centroids 
                 */
                public int numChangedCentroids() {
                        return changedCentroidCount;
                }
        }
        
        private ShortKMeansInit init = new ShortKMeansInit.RANDOM(); 
        private KMeansConfiguration<ShortNearestNeighbours, short[]> conf;
        private Random rng = new Random();
        
        /**
         * Construct the clusterer with the the given configuration.
         * 
         * @param conf The configuration.
         */
        public ShortKMeans(KMeansConfiguration<ShortNearestNeighbours, short[]> conf) {
                this.conf = conf;
        }
        
        /**
         * A completely default {@link ShortKMeans} used primarily as a convenience function for reading.
         */
        protected ShortKMeans() {
                this(new KMeansConfiguration<ShortNearestNeighbours, short[]>());
        }
        
        /**
         * Get the current initialisation algorithm
         *
         * @return the init algorithm being used
         */
        public ShortKMeansInit getInit() {
                return init;
        }

        /**
         * Set the current initialisation algorithm
         *
         * @param init the init algorithm to be used
         */
        public void setInit(ShortKMeansInit init) {
                this.init = init;
        }
        
        /**
         * Set the seed for the internal random number generator.
         *
         * @param seed the random seed for init random sample selection, no seed if seed < -1
         */
        public void seed(long seed) {
                if(seed < 0)
                        this.rng = new Random();
                else
                        this.rng = new Random(seed);
        }
                
        @Override
        public Result cluster(short[][] data) {
                DataSource<short[]> ds = new ShortArrayBackedDataSource(data, rng);
                
                try {
                        Result result = cluster(ds, conf.K);
                        result.nn = conf.factory.create(result.centroids);
                                                
                        return result;
                } catch (Exception e) {
                        throw new RuntimeException(e);
                }
        }
        
        @Override
        public int[][] performClustering(short[][] data) {
                ShortCentroidsResult clusters = this.cluster(data);
                return new IndexClusters(clusters.defaultHardAssigner().assign(data)).clusters();
        }
        
        /**
         * Initiate clustering with the given data and number of clusters.
         * Internally this method constructs the array to hold the centroids 
         * and calls {@link #cluster(DataSource, short [][])}.
         *
         * @param data data source to cluster with
         * @param K number of clusters to find
         * @return cluster centroids
         */
        protected Result cluster(DataSource<short[]> data, int K) throws Exception {
                int D = data.numDimensions();
                
                Result result = new Result();
                result.centroids = new short[K][D];
        
                init.initKMeans(data, result.centroids);
        
                cluster(data, result);

                return result;
        }
        
        /**
         * Main clustering algorithm. A number of threads as specified are 
         * started each containing an assignment job and a reference to
         * the same set of ShortNearestNeighbours object (i.e. Exact or KDTree). 
         * Each thread is added to a job pool and started in parallel. 
         * A single accumulator is shared between all threads and locked on update.
         * <br/>
         * This methods expects that the initial centroids have already been set in
         * the <code>result</code> object and as such <strong>ignores</strong> the
         * init object. <strong>In normal operation you should call one of the other <code>cluster</code>
         * cluster methods instead of this one.</strong> However, if you wish to resume clustering
         * iterations from a result that you've already generated this is the method
         * to use.
         *
         * @param data the data to be clustered
         * @param result the results object to be populated
         * @throws InterruptedException if interrupted while waiting, in
     *         which case unfinished tasks are cancelled.
         */
        public void cluster(short[][] data, Result result) throws InterruptedException {
                DataSource<short[]> ds = new ShortArrayBackedDataSource(data, rng);
                
                cluster(ds, result);
        }
        
        /**
         * Main clustering algorithm. A number of threads as specified are 
         * started each containing an assignment job and a reference to
         * the same set of ShortNearestNeighbours object (i.e. Exact or KDTree). 
         * Each thread is added to a job pool and started in parallel. 
         * A single accumulator is shared between all threads and locked on update.
         * <br/>
         * This methods expects that the initial centroids have already been set in
         * the <code>result</code> object and as such <strong>ignores</strong> the
         * init object. In normal operation you should call one of the other <code>cluster</code>
         * cluster methods instead of this one. However, if you wish to resume clustering
         * iterations from a result that you've already generated this is the method
         * to use.
         *
         * @param data the data to be clustered
         * @param result the results object to be populated
         * @throws InterruptedException if interrupted while waiting, in
     *         which case unfinished tasks are cancelled.
         */
        public void cluster(DataSource<short[]> data, Result result) throws InterruptedException {
                final short[][] centroids = result.centroids;
                final int K = centroids.length;
                final int D = centroids[0].length;
                final int N = data.size();
                float [][] centroids_accum = new float[K][D];
                int [] new_counts = new int[K];

                ExecutorService service = conf.threadpool;

                for (int i=0; i<conf.niters; i++) {
                        result.iterations++;
                        
                        for (int j=0; j<K; j++) 
                                Arrays.fill(centroids_accum[j], 0);
                        Arrays.fill(new_counts, 0);

                        ShortNearestNeighbours nno = conf.factory.create(centroids);
                        
                        List<CentroidAssignmentJob> jobs = new ArrayList<CentroidAssignmentJob>();
                        for (int bl = 0; bl < N; bl += conf.blockSize) {
                                int br = Math.min(bl + conf.blockSize, N);
                                jobs.add(new CentroidAssignmentJob(data, bl, br, nno, centroids_accum, new_counts));
                        }

                        service.invokeAll(jobs);

                        result.changedCentroidCount = 0;
                        for (int k=0; k < K; ++k) {
                                float ssd = 0;
                                if (new_counts[k] == 0) {
                                        // If there's an empty cluster we replace it with a random point.
                                        new_counts[k] = 1;

                                        short [][] rnd = new short[][] {centroids[k]};
                                        data.getRandomRows(rnd);
                                        result.changedCentroidCount++;
                                } else {
                                        for (int d=0; d < D; ++d) {
                                                short newValue = (short)((float)roundFloat((double)centroids_accum[k][d] / (double)new_counts[k]));
                                                
                                                // we're going to accumulate the SSD of the old vs new centroids
                                                // as a way of determining if this centroid has changed
                                                float diff = newValue - centroids[k][d]; 
                                                ssd += diff*diff;
                                                
                                                //update to new centroid
                                                centroids[k][d] = newValue;
                                        }
                                        
                                        if (ssd != 0)
                                                result.changedCentroidCount++;
                                }
                        }
                         
                        if (result.changedCentroidCount == 0)
                                break; // convergence
                }
        }
        
        protected float roundFloat(double value) { return (float) value; }
        protected double roundDouble(double value) { return value; }
        protected long roundLong(double value) { return (long)Math.round(value); }
        protected int roundInt(double value) { return (int)Math.round(value); }
        
        @Override
        public Result cluster(DataSource<short[]> ds) {
                try {
                        Result result = cluster(ds, conf.K);
                        result.nn = conf.factory.create(result.centroids);
                        
                        return result;
                } catch (Exception e) {
                        throw new RuntimeException(e);
                }
        }

    /**
         * Get the configuration
         * 
         * @return the configuration
         */
    public KMeansConfiguration<ShortNearestNeighbours, short[]> getConfiguration() {
        return conf;
    }
    
    /**
         * Set the configuration
         * 
         * @param conf
         *            the configuration to set
         */
    public void setConfiguration(KMeansConfiguration<ShortNearestNeighbours, short[]> conf) {
        this.conf = conf;
    }
        
        /**
         * Convenience method to quickly create an exact {@link ShortKMeans}. All
         * parameters other than the number of clusters are set
         * at their defaults, but can be manipulated through the configuration
         * returned by {@link #getConfiguration()}.
         * <p>
         * Euclidean distance is used to measure the distance between points.
         * 
         * @param K
         *            the number of clusters
         * @return a {@link ShortKMeans} instance configured for exact k-means
         */
        public static ShortKMeans createExact(int K) {
                final KMeansConfiguration<ShortNearestNeighbours, short[]> conf =
                                new KMeansConfiguration<ShortNearestNeighbours, short[]>(K, new ShortNearestNeighboursExact.Factory());

                return new ShortKMeans(conf);
        }

        /**
         * Convenience method to quickly create an exact {@link ShortKMeans}. All
         * parameters other than the number of clusters and number
         * of iterations are set at their defaults, but can be manipulated through
         * the configuration returned by {@link #getConfiguration()}.
         * <p>
         * Euclidean distance is used to measure the distance between points.
         * 
         * @param K
         *            the number of clusters
         * @param niters
         *            maximum number of iterations
         * @return a {@link ShortKMeans} instance configured for exact k-means
         */
        public static ShortKMeans createExact(int K, int niters) {
                final KMeansConfiguration<ShortNearestNeighbours, short[]> conf =
                                new KMeansConfiguration<ShortNearestNeighbours, short[]>(K, new ShortNearestNeighboursExact.Factory(), niters);

                return new ShortKMeans(conf);
        }
        
        /**
         * Convenience method to quickly create an approximate {@link ShortKMeans}
         * using an ensemble of KD-Trees to perform nearest-neighbour lookup. All
         * parameters other than the number of clusters are set
         * at their defaults, but can be manipulated through the configuration
         * returned by {@link #getConfiguration()}. 
         * <p>
         * Euclidean distance is used to measure the distance between points.
         * 
         * @param K
         *            the number of clusters
         * @return a {@link ShortKMeans} instance configured for approximate k-means 
         *              using an ensemble of KD-Trees
         */
        public static ShortKMeans createKDTreeEnsemble(int K) {
                final KMeansConfiguration<ShortNearestNeighbours, short[]> conf =
                                new KMeansConfiguration<ShortNearestNeighbours, short[]>(K, new ShortNearestNeighboursKDTree.Factory());

                return new ShortKMeans(conf);
        }
        
        @Override
        public String toString() {
                return String.format("%s: {K=%d, NN=%s}", this.getClass().getSimpleName(), this.conf.K, this.conf.getNearestNeighbourFactory().getClass().getSimpleName());
        }
}