#@ImagePlus imp
#@Integer(label="Number of Neighbors") k
#@Boolean(label="Give Each Neighbor's Distance") is_show_each

'''
Simple 2D KNN script
Olivier Burri, BioImaging & Optics Platform
Ecole Polytechnique Fédérale de Lausanne
July 12th 2016

Code provided as-is in reply to an ImageJ mailing list question
http://imagej.1557.x6.nabble.com/distance-between-adjacent-particles-td3699485.html#a5016864
'''

from ij.gui import Overlay,  Line, OvalRoi, Roi
from ij import IJ
from ij.measure import ResultsTable
from ij.measure import Calibration
from java.awt import Color
from ij.plugin.frame import RoiManager
import math

### Some functions ###

# Brute-Force KNN
def knn(data, k):
	# Just take the data, find the K nearest neighbors to each point
	
	the_knn=[]
	
	for i in range(len(data)):
		d = []
		for j in range(len(data)):
			d.append([j, dist(data[i], data[j])])

		# Sort
		d_sort = sorted(d, key=lambda thed: thed[1])
		# Keep only the k nearest
		the_knn.append(d_sort[1:(k+1)])
	return the_knn

# Euclidean Distance 2D
def dist(p0, p1):
    return math.sqrt((p0[0] - p1[0])**2 + (p0[1] - p1[1])**2)

# Convenience function to go from a ROI to points
def roiToPointList(roi):
	p = [roi.getFloatPolygon().xpoints, roi.getFloatPolygon().ypoints]
	points = map(list, zip(*p))
	return points

def showKNNResults(frt, the_knn, label):

	# Overlay
	ov = Overlay()
	the_avg = []
	
	# Draw each neighbor as a line and draw the average as a circle
	for i in range(len(the_knn)):
		avg =  sum([d[1] for d in the_knn[i]]) / k
		the_avg.append(avg)
		
		o = OvalRoi(points[i][0] - avg, points[i][1] - avg, 2*avg, 2*avg)
		o.setStrokeColor(Color.decode("#00FFFF"))
		ov.add(o)
		
		for j in range(k):
			# XY Coordinates of current point
			a = points[i];
			b = points[the_knn[i][j][0]]
			# XY Coordinates of neighbor k
			l = Line(a[0], a[1], b[0], b[1])
			l.setStrokeColor(Color.decode("#FF00FF"))
			ov.add(l)
	
	
	# Build Results Table
	cal = imp.getCalibration()
	
	for i in range(len(the_knn)):
		frt.incrementCounter()
		frt.addValue("Label", label);
		frt.addValue("Point", i)
		frt.addValue("Average distance [px]", the_avg[i])
		if( cal.scaled() ):
			frt.addValue("Average distance ["+cal.getXUnit()+"]", cal.getX(the_avg[i]))
		if is_show_each:
			for j in range(k):
				if( cal.scaled() ):
					frt.addValue("Distance Neighbor "+str(j+1)+" ["+cal.getXUnit()+"]", cal.getX(the_knn[i][j][1]))
				else:
					frt.addValue("Distance Neighbor "+str(j+1)+" [px]", the_knn[i][j][1])

	frt.show(str(k)+" Nearest Neightbors Average Distances")
	imp.setOverlay(ov)



### Starting the script ###

# Get a new Results Table
rt = ResultsTable()

# Get the ROI manager
rm = RoiManager().getInstance2();

n_rois = rm.getCount();

# Work either on the image alone
if n_rois == 0:
	roi = imp.getRoi()
	label = imp.getTitle()
	# Get Coordinates
	points = roiToPointList(imp.getRoi())
	
	
	# Compute K Nearest Neighbors
	the_knn = knn(points,k)

	# Display the result table and an overlay
	showKNNResults(rt,the_knn,label)

# Or on each ROI of the ROI manager
else:
	for i in range(n_rois):
		roi = rm.getRoi(i)
		label = roi.getName()
		# Get Coordinates
		points = roiToPointList(roi)
	
		# Compute K Nearest Neighbors
		the_knn = knn(points,k)
	
		# Display the result table and an overlay
		showKNNResults(rt,the_knn,label)