#!/usr/bin/env python

"""This is an implementation of the K-means algorithm.
You give it a bunch of named M{N}-dimensional vectors, and it
finds M{K} clusters so that the mean distance between the
cluster center and the points in that cluster are as small
as possible.   You can think of this as finding a set of
M{K} points that do the best possible job of approximating the
input data.

The algorithm is designed for flexibility rather than speed.
It allows you to change the function used to compute the distance
from cluster center to the points.  It also lets you ignore
wild, outlying points in the computation of the class centers
by using a trimmed mean, rather than the average.

For a script version, use C{g_kmeans} without the ".py".
"""

import math
import random
import numpy
from gmisclib import g_implements
from gmisclib import gpkmisc
from gmisclib import Numeric_gpk as NG
from gmisclib import die
import q_classifier_r

HUGE = 1e30

def find_which_cluster(datum, cpos, distfcn):
	"""Which cluster does a point belog to?
	@param cpos: a list of cluster centers.
	@return: (the index of the closest cluster and the distance)
	@rtype (int, float)
	"""
	bestdist = HUGE
	best = -1
	for (i, cp) in enumerate(cpos):
		tmp = distfcn(datum.value, cp)
		if tmp < bestdist:
			bestdist = tmp
			best = i
	return (best, bestdist)


def find_center(i, membership, data, clip):
	tmp = []
	for (mem, datum) in zip(membership, data):
		if mem == i:
			tmp.append( datum.value )
	ctr = NG.trimmed_mean_across(tmp, None, clip)
	return ctr


class cluster_descriptor(object):
	def __init__(self, center, membership, size):
		self.center = center	#: Vector position of the cluster center.
		self.membership = membership	#: a list of data in the cluster
		self.size = size	#: mean (or medan) distance of the cluster's data from cluster center



def kmeans(data, Ncl, distfcn=None, Nit=None, clip=None):
	"""Run the K-means algorithm to divide a list of points into C{Ncl} clusters.
	@param data: The data to be clustered.
	@type data: L{list}(L{q_classifier_r.datum_c}) or a list of equivalent classes.
		Each C{datum} must hae a C{value} member containing a C{numpy.ndarray} vector.
	@param Ncl: How many clusters?
	@param distfcn: A function to compute the distance between a point and a cluster center.
	@type distfcn: C{fcn(vector, vector) -> float}, where a vector is a one-dimensional C{numpy.ndarray} of floats.
	@param Nit: A limit on the number of iterations.
	@type Nit: int
	@param clip: C{0} to execute a normal K-means algorithm.   Greater than zero to use a trimmed
		mean to define cluster centers.   When C{clip > 0}, the reported error is a median
		rather than a mean, to make it robust against wild points.
	@type clip: L{float} on M{[0, 0,5)}.
	@return: (1) A list of cluster descriptors, each of which tells you where a cluster is and which points are members;
		(2) A dictionary that maps from each datum to the (integer) index of each cluster, and
		(3) the overall mean (or median if C{clip>0}) distance between points and cluster centers.
	@rtype: (L{list}(L{cluster_descriptor}), L{dict}(L{q_classifier_r.datum_c}:L{int}), L{float})
	"""
	for datum in data:
		g_implements.check(datum, q_classifier_r.datum_c, ignore="__init__")
	
	if Nit is None:
		Nit = 10 + int(round(math.sqrt(len(data)*Ncl)))
	if clip is None:
		clip = 0.0
	assert len(data) >= Ncl
	bestcp = None
	bestm = None
	besterr = HUGE
	starts = 0
	while starts < Nit:
		cpos = [data[c].value for c in random.sample(range(len(data)), Ncl) ]
		membership = numpy.zeros((len(data),), numpy.int)
		assert len(cpos) == Ncl
		passes = 0
		while passes < Nit:
			cldtmp = [ [] for x in range(Ncl) ]
			scatter = 0.0
			omem = numpy.array(membership, copy=True)
			for (i, datum) in enumerate(data):
				membership[i], dist = find_which_cluster(datum, cpos, distfcn)
				if membership[i] >= 0:
					cldtmp[membership[i]].append(dist)
				scatter += dist

			for i in range(Ncl):
				cpos[i] = find_center(i, membership, data, clip)

			if numpy.equal(omem, membership).all():
				break
			passes += 1

		if clip > 0.0:
			cldist = [ numpy.median(cldtmp[i]) for i in range(Ncl) ]
		else:
			cldist = [ numpy.mean(cldtmp[i]) for i in range(Ncl) ]

		if scatter < besterr:
			besterr = scatter
			bestcp = cpos
			bestm = membership

		starts += 1

	o = []
	for i in range(Ncl):
		m = [ d.uid for (mem, d) in zip(bestm, data) if mem == i ]
		o.append( cluster_descriptor(bestcp[i], m, cldist[i]) )
	map_to_cluster = {}
	for (datum, mem) in zip(data, bestm):
		map_to_cluster[datum] = mem
	return (o, map_to_cluster, besterr)




def read_data(fd):
	"""Reads in feature vectors, each with a uid as a comment."""
	d = []
	dim = None
	# die.info('Reading')
	ln = 0
	# for l in fd.readlines():
	for l in fd:
		ln += 1
		if l.startswith('#'):
			continue
		aa = l.split('#', 1)
		if len(aa) > 1:
			uid = aa[1].strip()
		else:
			uid = 'Line:%d' % ln
		a = aa[0].split()
		if dim is None:
			dim = len(a)
		elif len(a) != dim:
			die.die('Not all vectors have length=%d.  Problem on line %d'
					% (dim, ln)
				)
		d.append( q_classifier_r.datum_c(numpy.array([float(x) for x in a]),
				uid)
				)
	return d


def euclid(a, b):
	return math.sqrt(numpy.sum(numpy.ravel(a-b)**2))


