/*
        AUTOMATICALLY GENERATED BY jTemp FROM
        /Users/jsh2/Work/openimaj/target/checkout/machine-learning/nearest-neighbour/src/main/jtemp/org/openimaj/knn/approximate/#T#KDTreeEnsemble.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.knn.approximate;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

import cern.jet.random.Uniform;
import cern.jet.random.engine.MersenneTwister;
    
import org.openimaj.knn.ShortNearestNeighbours;
import org.openimaj.util.array.IntArrayView;
import org.openimaj.util.pair.*;

import jal.objects.BinaryPredicate;
import jal.objects.Sorting;

/**
 * Ensemble of Best-Bin-First KDTrees for short data.
 * 
 * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
 * @author Sina Samangooei (ss@ecs.soton.ac.uk)
 */
public class ShortKDTreeEnsemble {
        private static final int leaf_max_points = 14;
        private static final int varest_max_points = 128;
        private static final int varest_max_randsz = 5;
        
        Uniform rng;

    /**
         * An internal node of the KDTree
         */     
        public static class ShortKDTreeNode {
                class NodeData {}
                
            class InternalNodeData extends NodeData {
                ShortKDTreeNode right;
                float disc;
                int disc_dim;
            }
            
            class LeafNodeData extends NodeData {
                int [] indices;
            }
            
            /**
             * left == null iff this node is a leaf.
             */
            ShortKDTreeNode left;

            NodeData node_data;
            
            private Uniform rng;
            
            boolean is_leaf() { 
                return left==null; 
            }

            IntFloatPair choose_split(final short [][] pnts, final IntArrayView inds) {
                int D = pnts[0].length;
                
                // Find mean & variance of each dimension.
                float [] sum_x = new float[D];
                float [] sum_xx = new float[D];
                
                int count = Math.min(inds.size(), varest_max_points);
                for (int n=0; n<count; ++n) {
                    for (int d=0; d<D; ++d) {
                        sum_x[d]  += pnts[inds.getFast(n)][d];
                        sum_xx[d] += (pnts[inds.getFast(n)][d]*pnts[inds.getFast(n)][d]);
                    }
                }

                FloatIntPair[] var_dim = new FloatIntPair[D];
                for (int d=0; d < D; ++d) {
                        var_dim[d] = new FloatIntPair();
                    if (count <= 1)
                        var_dim[d].first = 0;
                    else
                        var_dim[d].first = (sum_xx[d] - ((float)1/count)*sum_x[d]*sum_x[d])/(count - 1);
                    var_dim[d].second = d;
                }

                // Partial sort makes a BIG difference to the build time.
                int nrand = Math.min(varest_max_randsz, D);
                Sorting.partial_sort(var_dim, 0, nrand, var_dim.length, new BinaryPredicate() {
                                @Override
                                public boolean apply(Object arg0, Object arg1) {
                                        FloatIntPair p1 = (FloatIntPair) arg0;
                                        FloatIntPair p2 = (FloatIntPair) arg1;
                                        
                                        if (p1.first > p2.first) return true;
                                        if (p2.first > p1.first) return false;
                                        return (p1.second > p2.second);
                                }});
                
                int randd = var_dim[rng.nextIntFromTo(0, nrand-1)].second;
                
                return new IntFloatPair(randd, sum_x[randd]/count);
            }

            void split_points(final short [][] pnts, IntArrayView inds) {
                IntFloatPair spl = choose_split(pnts, inds);

                ((InternalNodeData)node_data).disc_dim = spl.first;
                ((InternalNodeData)node_data).disc = spl.second;

                int N = inds.size();
                int l = 0;
                int r = N;
                while (l!=r) {
                  if (pnts[inds.getFast(l)][((InternalNodeData)node_data).disc_dim] < ((InternalNodeData)node_data).disc) l++;
                  else {
                    r--;
                    int t = inds.getFast(l);
                    inds.setFast(l, inds.getFast(r));
                    inds.setFast(r, t);
                  }
                }
            
                // If either partition is empty -> vectors identical!
                if (l==0 || l==N) { l = N/2; } // The vectors are identical, so keep nlogn performance.

                left = new ShortKDTreeNode(pnts, inds.subView(0, l), rng);
                
                ((InternalNodeData)node_data).right = new ShortKDTreeNode(pnts, inds.subView(l, N), rng);
            }

                /** Construct a new node */
            public ShortKDTreeNode() { }

                /** 
                 * Construct a new node with the given data
                 *
                 * @param pnts the data for the node and its children
                 * @param inds a list of indices that point to the relevant
                 *                      parts of the pnts array that should be used
                 * @param rng the random number generator
                 */
            public ShortKDTreeNode(final short [][] pnts, IntArrayView inds, Uniform rng) {
                this.rng = rng;
                if (inds.size() > leaf_max_points) { // Internal node
                        node_data = new InternalNodeData();
                    split_points(pnts, inds);
                }
                else {
                        node_data = new LeafNodeData();
                        ((LeafNodeData)node_data).indices = inds.toArray();
                }
            }

            void search(final short [] qu, PriorityQueue<FloatObjectPair<ShortKDTreeNode>> pri_branch, List<IntFloatPair> nns, boolean[] seen, short [][] pnts, float mindsq)
            {
                ShortKDTreeNode cur = this;
                ShortKDTreeNode other = null;

                while (!cur.is_leaf()) { // Follow best bin first until we hit a leaf
                        float diff = qu[((InternalNodeData)cur.node_data).disc_dim] - ((InternalNodeData)cur.node_data).disc;

                    if (diff < 0) {
                        other = ((InternalNodeData)cur.node_data).right;
                        cur = cur.left;
                    }
                    else {
                        other = cur.left;
                        cur = ((InternalNodeData)cur.node_data).right;
                    }

                    pri_branch.add(new FloatObjectPair<ShortKDTreeNode>(mindsq + diff*diff, other));
                }

                int [] cur_inds = ((LeafNodeData)cur.node_data).indices;
                int ncur_inds = cur_inds.length;
                
                int i;
                float [] dsq = new float[1];
                for (i = 0; i < ncur_inds; ++i) {
                        int ci = cur_inds[i];
                    if (!seen[ci]) {
                        ShortNearestNeighbours.distanceFunc(qu, new short[][] {pnts[ci]}, dsq);
                        
                        nns.add(new IntFloatPair(ci, dsq[0]));
                        
                        seen[ci] = true;
                    }
                }
            }
        }
        
        /** The tree roots */ 
        public final ShortKDTreeNode [] trees;
        
        /** The underlying data array */
        public final short [][] pnts;
    
    /**
     * Construct a ShortKDTreeEnsemble with the provided data,
     * using the default of 8 trees.
     * @param pnts the data array 
     */
    public ShortKDTreeEnsemble(final short [][] pnts) {
        this(pnts, 8, 42);
    }
    
    /**
     * Construct a ShortKDTreeEnsemble with the provided data and
     * number of trees.
     * @param pnts the data array 
     * @param ntrees the number of KDTrees in the ensemble 
     */
    public ShortKDTreeEnsemble(final short [][] pnts, int ntrees) {
        this(pnts, ntrees, 42);
    }
    
    /**
     * Construct a ShortKDTreeEnsemble with the provided data and
     * number of trees.
     * @param pnts the data array 
     * @param ntrees the number of KDTrees in the ensemble
     * @param seed the seed for the random number generator used in 
     *                  tree construction 
     */
    public ShortKDTreeEnsemble(final short [][] pnts, int ntrees, int seed) {
        final int N = pnts.length;
        this.pnts = pnts;
        this.rng = new Uniform(new MersenneTwister(seed));

        // Create inds.
        IntArrayView inds = new IntArrayView(N);
        for (int n=0; n<N; ++n) inds.setFast(n, n);

        // Create trees.
        trees = new ShortKDTreeNode[ntrees];
        for (int t=0; t<ntrees; ++t) {
            trees[t] = new ShortKDTreeNode(pnts, inds,rng);
        }
    }

    void search(final short [] qu, int numnn, IntFloatPair[] ret_nns, int nchecks) {
        final int N = pnts.length;
        
        if (nchecks < numnn) nchecks = numnn;
        if (nchecks > N) nchecks = N;
        
        PriorityQueue<FloatObjectPair<ShortKDTreeNode>> pri_branch = new PriorityQueue<FloatObjectPair<ShortKDTreeNode>>(
                11, 
                new Comparator<FloatObjectPair<ShortKDTreeNode>>() {
                        @Override
                        public int compare(FloatObjectPair<ShortKDTreeNode> o1, FloatObjectPair<ShortKDTreeNode> o2) {
                                if (o1.first > o2.first) return 1;
                                if (o2.first > o1.first) return -1;
                                return 0;
                        }}
        );

        List<IntFloatPair> nns = new ArrayList<IntFloatPair>((3*nchecks)/2);
        boolean [] seen = new boolean[N];

        // Search each tree at least once.
        for (int t=0; t<trees.length; ++t) {
            trees[t].search(qu, pri_branch, nns, seen, pnts, 0);
        }

        // Continue search until we've performed enough distances
        while (nns.size() < nchecks) {
                FloatObjectPair<ShortKDTreeNode> pr = pri_branch.poll();
            
            pr.second.search(qu, pri_branch, nns, seen, pnts, pr.first);
        }

        IntFloatPair [] nns_arr = nns.toArray(new IntFloatPair[nns.size()]); 
        Sorting.partial_sort(nns_arr, 0, numnn, nns_arr.length, new BinaryPredicate() {
                        @Override
                        public boolean apply(Object lhs, Object rhs) {
                                return ((IntFloatPair)lhs).second < ((IntFloatPair)rhs).second;
                        }});

        System.arraycopy(nns_arr, 0, ret_nns, 0, Math.min(numnn, nchecks));
    }
}