#!/usr/bin/env python
import psycopg2
import json
import sys

# Create test tables with following statements.
#
# create table p as (select point(random(), random()) from generate_series(1,10000000));
# create index p_idx on p using gist(v);
# create table g as (select polygon(3 + (random()*5)::int, circle(point(random(), random()), 0.001)) v from generate_series(1,10000000));
# create index g_idx on g using gist(v);

dbconn = psycopg2.connect("dbname='postgres' user='smagen' host='/tmp' password='' port=5431")

points = []
pointsCount = 16
tableName = None

def generatePoints(n):
	global points
	m = 1
	d = 0.5
	points.append((0, 0))
	while m <= n:
		for i in range(0, m):
			points.append((points[i][0] + d, points[i][1]    ))
			points.append((points[i][0]    , points[i][1] + d))
			points.append((points[i][0] + d, points[i][1] + d))
		d /= 2.0
		m *= 4

generatePoints(pointsCount)

def runKnn(point, limit):
	cursor = dbconn.cursor()
	cursor.execute("EXPLAIN (ANALYZE, FORMAT JSON) SELECT * FROM " + tableName + " ORDER BY v <-> %s::point LIMIT %s;", (point, limit))
	plan = cursor.fetchone()[0][0]
	cursor.close()
	return (plan['Planning Time'], plan['Execution Time'])

def makeTests(n, limit):
	planningTime = 0
	executionTime = 0
	for i in range(0, n):
		for j in range(0, pointsCount):
			point = '(' + str(points[j][0]) + ',' + str(points[j][1]) + ')'
			result = runKnn(point, limit)
			planningTime += result[0]
			executionTime += result[1]
	planningTime /= n * pointsCount
	executionTime /= n * pointsCount
	return (planningTime, executionTime)

if (len(sys.argv) < 2):
	print "Usage: %s table_name" % sys.argv[0]
	sys.exit(2)

tableName = sys.argv[1]

for limit in [10, 100, 1000]:
	result = makeTests(10, limit)
	print "limit: %s\nplanning:  %s\nexecution: %s" % (limit, result[0], result[1])

dbconn.close()