package org.forkalsrud.util;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Random;

public class VpTree<T> {

    public static interface Metric<T> {
        public double distance(T a, T b);
    }

    private Metric<T> metric;
    private Random random;
    private Node root;

    public VpTree(Collection<T> items, Metric<T> metric) {

        this.metric = metric;
        this.random = new Random();
        ArrayList<WithDistance> points = new ArrayList<WithDistance>(items.size());
        for (T t : items) {
            points.add(new WithDistance(t, 0d));
        }          
        this.root = buildFromPoints(points, 0, points.size());
    }

    
    public class WithDistance implements Comparable<WithDistance> {

        private T value;
        private double distance;

        public WithDistance(T value, double distance) {
            this.value = value;
            this.distance = distance;
        }
        
        public T getValue() {
            return this.value;
        }
        
        public double getDistance() {
            return this.distance;
        }

        @Override
        public int compareTo(WithDistance other) {
            double delta = other.distance - this.distance;
            return (int) Math.signum(delta);
        }
        
        @Override
        public String toString() {
            return String.valueOf(value) + " (" + distance + ")";
        }
    }

    public List<WithDistance> search(final T target, int numResults) {

        Searcher searcher = new Searcher(target, numResults);
        root.search(searcher);
        PriorityQueue<WithDistance> heap = searcher.matches;
        ArrayList<WithDistance> results = new ArrayList<WithDistance>(heap.size());
        while (!heap.isEmpty()) {
            results.add(heap.remove());
        }
        Collections.reverse(results);
        return results;
    }

    
    Node buildFromPoints(ArrayList<WithDistance> points, int left, int right) {

        WithDistance vp = points.get(left);
        Node nearChild = null;
        Node farChild = null;
        double threshold = 0d;
        
        if (right - left > 1) {
            // choose an arbitrary point and move it to the start
            int len = right - left;
            int pos = left + random.nextInt(len);
            Collections.swap(points, left, pos);
            vp = points.get(left);
            // sort the rest by distance from said point
            for (int i = left + 1; i < right; i++) {
                WithDistance wd = points.get(i);
                wd.distance = metric.distance(vp.value, wd.value);
            }
            Collections.sort(points.subList(left + 1, right));
            // Use the median distance as a split
            int midpoint = (left + 1 + right) / 2;
            threshold = points.get(midpoint).distance;
            // The leftmost points are closer, the rightmost are farher away
            nearChild = buildFromPoints(points, left + 1, midpoint);
            farChild = buildFromPoints(points, midpoint, right);
        }
        return new Node(vp.value, threshold, nearChild, farChild);
    }

    protected class Searcher {
        
        private final T query;
        private final int numMatches;
        private final PriorityQueue<WithDistance> matches;

        public Searcher(T query, int numMatches) {
            this.query = query;
            this.numMatches = numMatches;
            this.matches = new PriorityQueue<WithDistance>(numMatches);
        }

        public double queryDistance(T value) {
            return metric.distance(query, value);
        }

        /**
         * Farthest match distance from query
         * @return
         */
        public double tau() {
            return matches.peek().distance;
        }

        public void add(T value, double distance) {
            if (matches.size() < numMatches || distance <  tau()) {
                matches.add(new WithDistance(value, distance));
            }
            if (matches.size() > numMatches) {
                // drop the one with the largest distance
                matches.poll();
            }
        }
    }
    
    protected class Node {

        private T value;
        private double threshold;
        private Node near;
        private Node far;

        public Node(T value, double threshold, Node left, Node right) {
            this.value = value;
            this.threshold = threshold;
            this.near = left;
            this.far = right;
        }
        
        void search(Searcher searcher) {

            double distance = searcher.queryDistance(value);
            searcher.add(value, distance);
            if (distance < threshold) {
                // try near child first
                if (near != null && distance - searcher.tau() <= threshold) {
                    near.search(searcher);
                }
                if (far != null && distance + searcher.tau() >= threshold) {
                    far.search(searcher);
                }
            } else {
                // try far child first
                if (far != null && distance + searcher.tau() >= threshold) {
                    far.search(searcher);
                }
                if (near != null && distance - searcher.tau() <= threshold) {
                    near.search(searcher);
                }
            }
        }
    }

}
