import numpy as npy
from scipy.spatial import cKDTree

def nearestneighbors(md,data,goodids,badids,knn):
	'''
	fill holes using nearest neigbors.  Arguments include:

	md:		the model
	data:		the data field to be filled (full field, including holes)
	goodids:	id's into the vertices that have good data
	badids:	id's into the vertices with missing/bad data
	knn:		integer representing the k nearest neighbors to use for filling
				holes.  The average data value over the k nearest neighbors is 
				then used to fill the hole.

	Usage:
		filleddata=nearestneighbors(md,goodids,badids,knn)

	Example:
		filledthickness=nearestneighbors(md,goodids,badids,5)
	'''

	if type(knn) != int or knn<1:
		raise TypeError('nearestneighbors error: knn should be an integer>1')

	if len(data) != md.mesh.numberofvertices:
		raise StandardError('nearestneighbors error: "data" should have length md.mesh.numberofvertices')

	filled=data
	
	XYGood=npy.dstack([md.mesh.x[goodids],md.mesh.y[goodids]])[0]
	XYBad=npy.dstack([md.mesh.x[badids],md.mesh.y[badids]])[0]
	tree=cKDTree(XYGood)
	nearest=tree.query(XYBad,k=knn)[1]
	
	if knn==1:
		filled[badids]=filled[goodids][nearest] # can add k=N to return the N nearest neighbors
	else:
		for i in range(len(badids)):
			neardat=[]
			for j in range(knn):
				neardat.append(filled[goodids][nearest[i][j]])
				filled[badids[i]]=npy.mean(neardat)
				
	return filled
