#!/usr/bin/env python

print "FeatureTree demonstration"

print "(graphic rendering done using PIL)"

import Image, ImageDraw

import random

import numpy

import sys

import cv

def FindFeatures( coords, k, tree_type ):

    """

    Construct a k-approximate-nearest-neighbour graph using either a kd-tree or

    a spilltree

    Parameters

    ----------

    coords : numpy array, each row is a point

    k : number of neighbours to search

    tree_type : either 'kdtree' or 'spilltree'

    Returns

    -------

    edges : numpy array of edges, each row is of the form

            [ point_1, point_2, distance ]

    """    

    nb_points = len(coords)

    cv_coords = cv.fromarray(coords)

    if tree_type == "kdtree":

        tree = cv.CreateKDTree(cv_coords)

    elif tree_type == "spilltree":

        tree = cv.CreateSpillTree( cv_coords,

                                   100,  # limit on the leaves to visit

                                   0.7,  # spilltree under this value

                                   0.1 ) # overlapping ratio

    else:

        print "Unkwnown tree_type '" + tree_type + "'"

        print "valid types are: 'kdtree' or 'spilltree'"

        sys.exit(1)

    res = cv.CreateMat(nb_points, k, cv.CV_32SC1)

    dist = cv.CreateMat(nb_points, k, cv.CV_64FC1)

    cv.FindFeatures(tree,cv_coords,res,dist,k)

    edges = []

    for i in range(0,nb_points):

        for n in range(0,k):

            edges.append([i,int(res[i,n]),dist[i,n]])

    return numpy.array(edges)

def weight_color( w, wmin, wmax ):

    # the closer, the greener

    # the further, the reder

    return ( int( (w-wmin) / (wmax-wmin) * 255 ),

             int( (wmax-w) / (wmax-wmin) * 255 ),

             0 )

def render(points,edges,filename,image_size=600,point_size=4,line_width=2):

    img = Image.new( "RGB", (size,size), (255,255,255) )

    draw = ImageDraw.Draw(img)

    wmin = min(edges[:,2])

    wmax = max(edges[:,2])

    # draw the edges

    for e in edges:

        draw.line( [ tuple(points[e[0]]), tuple(points[e[1]]) ],

                   fill=weight_color( e[2], wmin, wmax ),

                   width=line_width )

    # draw the points

    for p in points:

        box = [(p[0]-point_size,p[1]-point_size),(p[0]+point_size,p[1]+point_size)]

        draw.ellipse(box, fill=(0,0,0) )

    img.save(filename)

    img.show()

def FindBox( coords, box ):

    """

    Returns the points that approximately fall in the box

    Parameters

    ----------

    coords : numpy array, each row is a point

    box : a numpy array defining the box, for instance

          box = numpy.array([[[xmin,ymin]],[[xmax,ymax]]],dtype=float)

    Returns

    -------

    edges : numpy array of edges, each row is of the form

            [ point_1, point_2, distance ]

    """      

    nb_points = len(coords)

    cv_coords = cv.fromarray(coords)

    tree = cv.CreateKDTree(cv_coords)

    bounds_min = cv.fromarray(box[0])

    bounds_max = cv.fromarray(box[1])

    res = cv.CreateMat(1,nb_points,cv.CV_32SC1)

    cv.Set(res, -1)

    nb_found = cv.FindFeaturesBoxed(tree, bounds_min,bounds_max,res)

    found = []    

    for i in range(nb_found):

        found.append(coords[int(res[0,i])])

    return found

def render_box(points,box, found,filename,image_size=600,point_size=4,line_width=2):

    img = Image.new( "RGB", (size,size), (255,255,255) )

    draw = ImageDraw.Draw(img)

    # draw the black points

    for p in points:

        pbox = [(p[0]-point_size,p[1]-point_size),(p[0]+point_size,p[1]+point_size)]

        draw.ellipse(pbox, fill=(0,0,0) )

    # draw the points that are in the box

    for p in found:

        pbox = [(p[0]-point_size,p[1]-point_size),(p[0]+point_size,p[1]+point_size)]

        draw.ellipse(pbox, fill=(255,0,0) )

    # draw the box

    draw.rectangle( [ tuple(box[0][0]), tuple(box[1][0]) ],

                    outline=(255,10,10) )

    img.save(filename)

    img.show()

if __name__ == '__main__':

    nb_points = 800

    k = 5

    size = 600

    points = numpy.array( [ [ random.random()*float(size),

                              random.random()*float(size) ] for i in range(nb_points) ],

                          dtype=float )

    edges = FindFeatures(points,k, 'spilltree')

    render(points,edges,"spilltree_search.jpg",size)

    edges = FindFeatures(points,k, 'kdtree')

    render(points,edges,"kdtree_search.jpg",size)

    box = numpy.array([[[100,200]],[[400,400]]],dtype=float)

    found = FindBox(points,box)

    render_box(points,box,found,"box_search.jpg",size)