"""
@note: The split into training set and test set is completely deterministic,
	depending only on each datum's identifier, the total number of data.
	(If you use a single instance of L{bluedata} or L{bluedata_groups} to
	make several splits, then each split will be different.  But if you created
	a new instance, you'd get the exact same sequence of splits.)
"""

import math
import random
import zlib
from gmisclib import gpkmisc
from gmisclib import dictops
from gmisclib import dict_vector as DV
from gmisclib import blue_data_selector as BDS


def Hash(x):
	Mask1 = 0xffffffff
	Mask2 = 0xffffffffffff
	rv = 0
	for (s,i) in sorted(x):
		tmp = zlib.crc32(s) & Mask1
		rv = ((rv+tmp) * (17*i + 7)) & Mask2
	return rv

assert Hash([('a', 3), ('b', 4)]) == Hash([('a', 3), ('b', 4)]) 
assert Hash([('ss', 0), ('ta', 4)]) == Hash([('ss', 0), ('ta', 4)]) 


class bluedata(object):
	"""This is an extension of L{gmisclib.blue_data_selector.bluedata} which
	is aware of the classids of data.  It tries to split the
	data to take the same fraction of each class.
	"""

	def __init__(self, data, blueness = 2.0):
		"""@type data: a sequence of L{datum_tr} objects.
		"""
		r = random.Random(len(data))
		self.ntot = 0
		per_class_data = dictops.dict_of_lists()
		for datum in data:
			per_class_data.add(datum.classid, datum)
		ncltmp = [(nm, len(d)) for (nm,d) in per_class_data.items()]
		# Start the random number generator in a random, data-dependent
		# state, but one that is independent of the order of the data.
		r.jumpahead(abs(Hash(ncltmp)))
		self.r = r.getstate()

		self.selector = []
		for (cid, cldata) in sorted(per_class_data.items()):
			self.ntot += len(cldata)
			self.selector.append( (cid, BDS.bluedata(cldata, blueness, rng=r)) )
	

	def _chooseint(self, n, r):
		"""This chooses integers math.floor(n) or math.ceil(n) with
		appropriate probabilities so that the expectation value of
		_chooseint(n) is n.
		@type n: float
		@rtype: int
		@type r: L{random.Random}
		@param r: random number generator.
		"""
		f = math.floor(n)
		p = n - f
		assert 0.0 <= p <= 1.0
		i = int(f)
		if r.random() < p:
			return i+1
		return i


	def split(self, n, seed=0):
		"""Split the stored data into a group of C{n} and the remainder.
		Successive calls to split() will tend to (as nearly as possible)
		put different data into the first list.
		@type n: int or float
		@return: tuple(list of length n(or nearly), list of remainder)
		"""
		frac = float(n)/float(self.ntot)
		# First, we choose how many items to select from each class.
		r = random.Random()
		r.setstate(self.r)
		r.jumpahead(abs(seed))
		while True:
			# We pick random numbers for each class, then check to see if the
			# total is right.
			npc = {}
			ntot = 0
			for (cid, selector) in self.selector:
				tmp = self._chooseint(frac*len(selector), r)
				assert tmp >= 0 and tmp <= len(selector)
				ntot += tmp
				npc[cid] = tmp
			if n-0.5 <= ntot <= n+0.5:
				break

		# Then, we use a blue data selector for each class to pick the right number
		# of items from that class.
		o1 = []
		o2 = []
		for (cid, selector) in self.selector:
			t1, t2 = selector.split(npc[cid])
			o1.extend( t1 )
			o2.extend( t2 )
		assert len(o1) == ntot
		# Finally, we shuffle the classes back together.
		r.shuffle(o1)
		r.shuffle(o2)
		# ...and return.
		return (o1, o2)

	
	def __len__(self):
		return self.ntot


class bluedata_groups(object):
	def __init__(self, d, gr):
		"""A 'grouper' function takes a DUID (a unique i.d. string
		for a datum) and returns the name of the data group to which
		it belongs.    This group name is used in constructing the training
		and test sets.
		@param d: a list of data
		@type d: list(L{datum_c})
		@return: A dictionary full of lists; each list corresponds to one group.
		@rtype: L{dictops.dict_of_lists<dictops>}
		"""
		r = random.Random(len(d))
		grps = dictops.dict_of_lists()
		self.ngc = {}
		ncl = dictops.dict_of_accums()
		self.ntot = len(d)
		self.dused = {}
		self.gmap = {}
		for t in d:
			g = gr(t)
			grps.add(g, t)
			ncl.add(t.classid, 1)
			self.gmap[id(t)] = g
		self.grps = sorted(grps.items())

		# Start the random number generator in a random, data-dependent
		# state, but one that is independent of the order of the data.
		ncltmp = [(nm, n) for (nm,n) in ncl.items()]
		ngtmp = [(nm, len(t)) for (nm,t) in self.grps]
		r.jumpahead(abs(Hash(ncltmp)^Hash(ngtmp)))

		for t in d:
			self.dused[id(t)] = r.random()
		self.gused = {}
		n_in_c = DV.dict_vector()
		for (g, pgd) in self.grps:
			self.gused[g] = 0.0
			nc = dictops.dict_of_accums()
			for t in pgd:
				nc.add(t.classid, 1)
			self.ngc[g] = DV.dict_vector(nc.items())
			n_in_c += self.ngc[g]
			assert isinstance(self.ngc[g], DV.dict_vector)
		assert (n_in_c > 0).all(), "Weird - zero total examples of a class we know about???"
		self.ncl = DV.dict_vector(ncl.items())
		assert isinstance(self.ncl, DV.dict_vector)
		self.r = r.getstate()


	def __len__(self):
		return self.ntot

	HUGE = 1e30

	def split(self, n, seed=0):
		"""Take groups until you have accumulated approximately n data.
		The trick is to make all classifiers as identical as possible
		by getting the same balance of classes in the training set.
		We are willing to sacrifice some data to do that.
		@param n: how many data in the testing set
		@type n: float or int
		@return: (testing set, training set, names in training set)
		@rtype: tuple(list(L{datum_tr}), list(L{datum_tr}), list(str))
		"""
		r = random.Random()
		r.setstate(self.r)
		r.jumpahead(abs(seed))
		assert 1 <= n <= self.ntot
		ftr = 1.0 - float(n)/float(self.ntot)
		assert isinstance(self.ncl, DV.dict_vector)
		#: How many data from each class should the training set have?
		tgt = self.ncl * ftr
		assert isinstance(tgt, DV.dict_vector)
		#: This is an accumulator of how many we've found in each class:
		ta = DV.dict_vector()

		#First, we select which groups to use in the training set.
		myscore = self.HUGE
		groups = set()
		# Never put all your groups into the training set!
		while len(groups)<len(self.grps)-1:
			minused = min([self.gused[k] for (k,v) in self.grps
					if k not in groups])
			q = [k for (k,v) in self.grps
				if self.gused[k]<minused+r.random()**2
					and k not in groups
					]
			assert q
			r.shuffle(q)
			scores = DV.dict_vector()
			for g in q:
				scores[g] = self._score(ta + self.ngc[g] - tgt)
			gfound, newscore = DV.min_within(scores)
			if newscore < myscore:
				myscore = newscore
				groups.add(gfound)
				ta += self.ngc[gfound]
			else:
				break
		assert len(groups) > 0
		# Then, we go through and select data from the previously selected
		# groups.
		tgt *= (ta/self.ncl.float()).median()
		o2 = []
		o1 = []
		for (g, d) in self.grps:
			if g in groups:
				assert (ta >= self.ngc[g]).all()
				o2.extend( self._sample(d, tgt, ta, self.ngc[g]) )
			else:
				o1.extend(d)
		r.shuffle(o1)
		r.shuffle(o2)
		return (o1, o2, groups)


	def _score(self, future_tgt):
		return abs(future_tgt).sum()


	def _sample(self, d, tgt, ta, ngc):
		"""Here, we attempt to sample items from the list of data d,
		taking care to sample less-used data first.   We also
		increment the group used counter (for later use in choosing
		groups, in C{split()}.
		Data items will be used at most once, so you may not get all
		that you ask for.

		@param d: list of data
		@type d: list(data_tr)
		@param tgt: the desired remaining number of data to select for the training set.
			This is how many we will try to get out of the groups that have
			not yet been C{_select()}ed.   It is decremented at the end of the call.
		@param ta: the total remaining number of data in each class.   This
			looks at the data within groups we have selected but where
			we on which we have not yet called C{_sample()}.
			It is decremented at the end of the call.
		@param ngc: the number of data in each class, within our particular group. 
		@type tgt: L{dict_vector.dict_vector}
		@type ta: L{dict_vector.dict_vector}
		@type ngc: L{dict_vector.dict_vector}
		"""
		assert isinstance(ngc, DV.dict_vector), "ngc=%s" % str(ngc)
		n = DV.round(tgt * (ngc/ta.float()))
		n0 = n.copy()
		tmp = [ (self.dused[id(t)], t) for t in d ]
		tmp.sort()
		f = 1.0/len(d)
		o = []
		for (u, t) in tmp:
			if n[t.classid] >= 0.5:
				self.dused[id(t)] += 1
				self.gused[self.gmap[id(t)]] += f
				o.append(t)
				n[t.classid] -= 1
		tgt -= n0 - n
		ta -= ngc
		assert (ta >= 0).all()
		ta.idropzero()
		return o
