from scipy import ravel, r_
from random import sample

from cerebrum.neuralnet.datasets.supervised import SupervisedDataSet


class EmptySequenceError(Exception): pass


class SequentialDataSet(SupervisedDataSet):
	"""A SequentialDataSet is like a SupervisedDataSet except that it can keep
	track of sequences of samples. Indices of a new sequence are stored whenever
	the method newSequence() is called. The last (open) sequence is considered
	a normal sequence even though it does not have a following "new sequence"
	marker."""

	def __init__(self, indim, targetdim):
		SupervisedDataSet.__init__(self, indim, targetdim)
		# add field that stores the beginning of a new episode
		self.addField('sequence_index', 1)
		self.append('sequence_index', 0)
		self.currentSeq = 0

	def newSequence(self):
		"""Marks the beginning of a new sequence. this function does nothing if
		called at the very start of the data set. Otherwise, it starts a new
		sequence. Empty sequences are not allowed, and an EmptySequenceError
		exception will be raised."""
		length = self.getLength()
		if length != 0:
			if ravel(self.getField('sequence_index'))[-1] == length:
				raise EmptySequenceError
			self._appendUnlinked('sequence_index', length)

	def _getSequenceField(self, index, field):
		"""Return a sequence of one single field given by `field` and indexed by
		`index`."""
		seq = ravel(self.getField('sequence_index'))
		if len(seq) == index + 1:
			# user wants to access the last sequence, return until end of data
			return self.getField(field)[seq[index]:]
		if len(seq) < index + 1:
			# sequence index beyond number of sequences. raise exception
			raise IndexError('sequence does not exist.')
		return self.getField(field)[seq[index]:seq[index + 1]]

	def getSequence(self, index):
		"""Returns the sequence given by `index`.

		A list of arrays is returned for the linked arrays. It is assumed that
		the last sequence goes until the end of the dataset."""
		return [self._getSequenceField(index, l) for l in self.link]

	def getSequenceIterator(self, index):
		"""Return an iterator over the samples of the sequence specified by
		`index`.

		Each element is a tuple."""
		return zip(*self.getSequence(index))

	def endOfSequence(self, index):
		"""Return True if the marker was moved over the last element of
		sequence `index`, False otherwise.

		Mostly used like .endOfData() with while loops."""
		seq = ravel(self.getField('sequence_index'))
		if len(seq) == index + 1:
			# user wants to access the last sequence, return until end of data
			return self.endOfData()
		if len(seq) < index + 1:
			# sequence index beyond number of sequences. raise exception
			raise IndexError('sequence does not exist.')
		else:
			return self.index >= seq[index + 1]

	def gotoSequence(self, index):
		"""Move the internal marker to the beginning of sequence `index`."""
		try:
			self.index = ravel(self.getField('sequence_index'))[index]
		except IndexError:
			raise IndexError('sequence does not exist')

	def getCurrentSequence(self):
		"""Return the current sequence, according to the marker position."""
		seq = ravel(self.getField('sequence_index'))
		return len(seq) - sum(seq > self.index) - 1

	def getNumSequences(self):
		"""Return the number of sequences. The last (open) sequence is also
		counted in, even though there is no additional 'newSequence' marker."""
		return self.getField('sequence_index').shape[0]

	def getSequenceLength(self, index):
		"""Return the length of the given sequence. If `index` is pointing
		to the last sequence, the sequence is considered to go until the end
		of the dataset."""
		seq = ravel(self.getField('sequence_index'))
		if len(seq) == index + 1:
			# user wants to access the last sequence, return until end of data
			return int(self.getLength() - seq[index])
		if len(seq) < index + 1:
			# sequence index beyond number of sequences. raise exception
			raise IndexError('sequence does not exist.')
		return int(seq[index + 1] - seq[index])

	def removeSequence(self, index):
		"""Remove the `index`'th sequence from the dataset and places the
		marker to the sample following the removed sequence."""
		if index >= self.getNumSequences():
			# sequence doesn't exist, raise exception
			raise IndexError('sequence does not exist.')
		sequences = ravel(self.getField('sequence_index'))
		seqstart = sequences[index]
		if index == self.getNumSequences() - 1:
			# last sequence is going to be removed
			lastSeqDeleted = True
			seqend = self.getLength()
		else:
			lastSeqDeleted = False
			# sequence to remove is not last one (sequence_index exists)
			seqend = sequences[index + 1]

		# cut out data from all fields
		for label in self.link:
			# concatenate rows from start to seqstart and from seqend to end
			self.data[label] = r_[self.data[label][:seqstart, :], self.data[label][seqend:, :]]
			# update endmarkers of linked fields
			self.endmarker[label] -= seqend - seqstart

		# update sequence indices
		for i, val in enumerate(sequences):
			if val > seqstart:
				self.data['sequence_index'][i, :] -= seqend - seqstart

		# remove sequence index of deleted sequence and reduce its endmarker
		self.data['sequence_index'] = r_[self.data['sequence_index'][:index, :], self.data['sequence_index'][index + 1:, :]]
		self.endmarker['sequence_index'] -= 1

		if lastSeqDeleted:
			# last sequence was removed
			# move sequence marker to last remaining sequence
			self.currentSeq = index - 1
			# move sample marker to end of dataset
			self.index = self.getLength()
			# if there was only 1 sequence left, re-initialize sequence index
			if self.getLength() == 0:
				self.clear()
		else:
			# removed sequence was not last one (sequence_index exists)
			# move sequence marker to the new sequence at position 'index'
			self.currentSeq = index
			# move sample marker to beginning of sequence at position 'index'
			self.index = ravel(self.getField('sequence_index'))[index]


	def clear(self):
		SupervisedDataSet.clear(self, True)
		self._appendUnlinked('sequence_index', [0])
		self.currentSeq = 0

	def __iter__(self):
		"""Create an iterator object over sequences which are themselves
		iterable objects."""
		for i in range(self.getNumSequences()):
			yield self.getSequenceIterator(i)

	def _provideSequences(self):
		"""Return an iterator over sequence lists."""
		return iter(map(list, iter(self)))

	def evaluateModuleMSE(self, module, averageOver=1, **args):
		"""Evaluate the predictions of a module on a sequential dataset
		and return the MSE (potentially average over a number of epochs)."""
		res = 0.
		for dummy in range(averageOver):
			ponderation = 0.
			totalError = 0
			for seq in self._provideSequences():
				module.reset()
				e, p = self._evaluateSequence(module.activate, seq, **args)
				totalError += e
				ponderation += p
			assert ponderation > 0
			res += totalError / ponderation
		return res / averageOver

	def splitWithProportion(self, proportion=0.5):
		"""Produce two new datasets, each containing a part of the sequences.

		The first dataset will have a fraction given by `proportion` of the
		dataset."""
		l = self.getNumSequences()
		leftIndices = sample(list(range(l)), int(l * proportion))
		leftDs = self.copy()
		leftDs.clear()
		rightDs = leftDs.copy()
		index = 0
		for seq in iter(self):
			if index in leftIndices:
				leftDs.newSequence()
				for sp in seq:
					leftDs.addSample(*sp)
			else:
				rightDs.newSequence()
				for sp in seq:
					rightDs.addSample(*sp)
			index += 1
		return leftDs, rightDs
