import disco
 
def init_map(e, params):
	import random
	return [(random.randint(0,params.k-1),e[1])]
 
 
def estimate_map(e, params):
	return[(min([(params.dist(c,map(float,e[1].split(' '))),i) for (i,c) in enumerate(params.centers)])[1],e[1])]
 
 
def estimate_combiner(k, v, centers, done, params):
        if done:
                return [(i,' '.join(map(repr,c))) for (i,c) in centers.iteritems()]
        else:
		v=map(float,v.split(' '))
                if not centers.has_key(k): centers[k]=[0.0]*len(v) + [0]
                for i in range(len(v)): centers[k][i]+=v[i]
 
		centers[k][len(v)]+=1
 
 
def estimate_reduce(iter, out, params):
        x={}
        for k,v in iter:
                y=map(float,v.split(' '))
                if not x.has_key(k): 
			x[k]=y
                else:
                        for i in y: x[k][i]+=y[i]
 
        for k,v in x.iteritems():
                for i in range(len(v)-1): v[i]/=v[-1]
                out.add(k,' '.join(map(repr,v)))
 
 
def predict_map(e, params):
	return [(e[0],min([(params.dist(c,map(float,map(float,e[1].split(' ')))),i) for (i,c) in enumerate(params.centers)])[1])]
 
 
def d2(x,y): return sum([(x[i]-y[i])**2 for i in range(len(x))])
 
 
def estimate(input, centers, k, iterations=10, host="disco://localhost", map_reader=disco.chain_reader, nr_reduces=None):
	if centers!=None: k=len(centers)
	if nr_reduces==None: nr_reduces=k
 
	results=None
	if centers==None:	
		results = disco.job(host, name = 'kmeans_init',
				    input_files = input, 
				    map_reader = map_reader, 
				    fun_map = init_map, 
				    combiner = estimate_combiner, 
				    reduce = estimate_reduce,
				    nr_reduces = nr_reduces, 
				    params = disco.Params(k=k),
				    sort = False, clean = True)
 
	for i in range(iterations):
		if results!=None:
			centers=[None]*k
			counts=[None]*k
			for key,value in disco.result_iterator(results):
				x=map(float,value.split(' '))
				centers[int(key)]=x[:-1]
				counts[int(key)]=x[-1]
 
		results = disco.job(host, name = 'kmeans_iterate_'+str(i),
				    input_files = input,
				    map_reader = map_reader,
				    fun_map = estimate_map, 
				    combiner = estimate_combiner,
				    reduce = estimate_reduce,
				    nr_reduces = nr_reduces, 
				    params = disco.Params(centers=centers,dist=d2),
				    sort = False, clean = True)
 
	return centers
 
 
def predict(input, centers, host="disco://localhost", map_reader=disco.chain_reader, nr_reduces=None):
	if nr_reduces==None: nr_reduces=len(centers)
 
	results = disco.job(host, name = 'kmeans_output', 
			    input_files = input, 
			    map_reader = map_reader, 
			    fun_map = predict_map,
			    nr_reduces = nr_reduces, 
			    params = disco.Params(centers=centers,dist=d2),
			    sort = False, clean = True)
 
	return results