import numpy as np

from metanetwork.utils import create_random_state

class BatchDataset:
	"""
	Batch Dataset stores data and labels with capacity
	to shuffle, repeat, and batch data in a manner
	similar to Tensorflow's Dataset implementation.

	Parameters
	----------
	X : array-like, shape=(n_samples, n_features)
		Data.

	Y : array-like, shape=(n_samples, n_labels), default=None
		Labels.

	weights : array-like, shape=(n_samples,), default=None
		Sample weights. If None, then samples are equally weighted.

	seed : None or int or RandomState, default=None
		Initial seed for the RandomState. If seed is None,
		return the RandomState singleton. If seed is an int,
		return a RandomState with the seed set to the int.
		If seed is a RandomState, return that RandomState.

	Attributes
	----------
	available : ndarray
		List of the indices corresponding to available
		data to draw from.

	batch_size : int, range=[1, n_samples]
		Batch size.

	order : list
		Order of operations used internally.

	i : int
		Number of times data has been drawn from the BatchDataset.
	"""
	def __init__(self, X, Y=None, weights=None, seed=None):
		self.X = np.array(X)
		self.Y = np.array(Y) if Y is not None else None
		if weights is not None : self.weights = np.array(weights)
		else : self.weights = np.ones(len(X))
		self.seed = create_random_state(seed=seed)
		self.available = np.arange(len(self.X)).astype(int)
		self.batch_size = 1
		self.n_batches = self._calculate_n_batches()
		self.order = []
		self.i = 0

	def batch(self, batch_size):
		"""
		Setup the BatchDataset to batch with the
		given batch size.

		Parameters
		----------
		batch_size :  int, range=[1, n_samples]
			Batch size.
		"""
		self.batch_size = batch_size
		self.n_batches = self._calculate_n_batches()
		self.order = [op for op in self.order if op != 'batch']
		self.order.append('batch')
		self.i = 0
		return self

	def repeat(self):
		"""
		Setup the BatchDataset to repeat the data.
		"""
		self.order = [op for op in self.order if op != 'repeat']
		self.order.append('repeat')
		self.i = 0
		return self

	def shuffle(self):
		"""
		Setup the BatchDataset to shuffle the data.
		"""
		self.order = [op for op in self.order if op != 'shuffle']
		self.order.append('shuffle')
		self.i = 0
		return self

	def next(self):
		"""
		Draw a batch from the BatchDataset.
		If this is the first batch, organize the dataset.
		If this batch would cause there to be less than another
		batch, reorganize the dataset.

		Returns
		-------
		next : ndarray or tuple of ndarray
			The batched data, and if available, labels and weights.
		"""
		if self.i == 0 : self.organize()
		if len(self.available) < self.batch_size:
			batch = self.available
			self.available = np.array([])
		else:
			batch = self.available[:self.batch_size]
			self.available = self.available[self.batch_size:]
		if len(self.available) < self.batch_size:
			self.organize(prepend=self.available)
		self.i += 1
		next = [self.X[batch]]
		if self.Y is None : next.append(None)
		else : next.append(self.Y[batch])
		next.append(self.weights[batch])
		return next

	def _calculate_n_batches(self):
		"""
		Calculate the number of batches that cover
		all data in the dataset.

		Returns
		-------
		n_batches : int
			Number of batches.
		"""
		return len(self.X) // self.batch_size + 1

	def organize(self, prepend=[], append=[]):
		"""
		Organize the BatchDataset according to `order`.
		In this manner, the order of shuffle, repeat, and
		batch affect how the data is drawn.

		Parameters
		----------
		prepend : list, default=[]
			Prepend these indices to the reorganized list.
			Data drawn will first exhaust this list.

		append : list, default=[]
			Append these indices to the reorganized list.
			Data drawn will exhaust this list last.
		"""
		order = np.array(self.order)
		try : shuffle = np.argwhere(order == 'shuffle').flatten()[0]
		except : shuffle = np.inf
		try : repeat = np.argwhere(order == 'repeat').flatten()[0]
		except : repeat = np.inf
		try : batch = np.argwhere(order == 'batch').flatten()[0]
		except : batch = np.inf
		length_Y = len(self.Y)
		if shuffle < repeat and shuffle < batch:
			self.available = np.arange(length_Y)
			self.seed.shuffle(self.available)
		elif repeat < shuffle:
			if shuffle < batch:
				self.available = self.seed.choice(np.arange(length_Y), length_Y)
			elif batch < shuffle and shuffle < np.inf:
				self.available = np.arange(length_Y).reshape(-1,self.batch_size)
				n_batches = np.arange(len(self.available))
				indices = self.seed.choice(n_batches, len(self.available))
				self.available = self.available[indices].flatten()
			else:
				self.available = np.arange(length_Y)
		elif batch < shuffle < repeat:
			self.available = np.arange(length_Y).reshape(-1,self.batch_size)
			self.seed.shuffle(self.available)
			self.available.flatten()
		self.available = np.concatenate((prepend, self.available, append), axis=0).astype(int)
