import numpy as np
from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import MinMaxScaler, OneHotEncoder
from imblearn.over_sampling import ADASYN

class ScoreDataset:
	'''
	Basic structure implementing data access and behavior useful for training ANN in keras/tensorflow.
	'''
	def __init__(self, features, labels, oversample=None, kfold=5, norm=False):
		assert len(features) == len(labels), "Labels (Y) and features (X) do not have the same number of observations!"
		self.X = np.array(features).astype('float64')
		if norm:
			scaler = MinMaxScaler(feature_range=(-1, 1), copy=False).fit(self.X)
			scaler.transform(self.X)
		self.Y = np.array(labels).astype('int32').reshape(len(labels),-1)
		self._X, self._Y = self.X, self.Y
		if oversample:
			self._oversample(oversample)
		self.leaves = []
		self.num_leaves = 1
		self.n_batches = 1
		self.kfold = kfold
		self.enc = OneHotEncoder() ; self.enc.fit(self.Y)

	def __len__(self):
		return len(self.X)

	def __getitem__(self, i):
		return [self.X[i], self.Y[i]]

	def get_x_dims(self):
		return self.X.shape[1]

	def get_y_dims(self):
		return self.Y.shape[1]

	def _oversample(self, ratio=0.7):
		'''
		Random resampling of immunogenic peptides to equal ratio
		of non-immunogenic (training set only).
		'''
		# Might be better ways to do this than just random oversampling -- e.g. SMOTE?
		self._X, self._Y = ADASYN(sampling_strategy=ratio).fit_resample(self.X, self.Y)
		self._Y = self._Y.reshape(-1, 1)

	def get_splits(self, folds):
		'''
		Defines k training/validation leaves.
		'''
		assert type(folds) == int, "Folds must be type int."
		assert folds>1, "Number of folds must be greater than 1."
		skf = StratifiedKFold(n_splits=folds, shuffle=True)
		for train_idx, val_idx in skf.split(self._X, self._Y):
			self.leaves.append((train_idx, val_idx))
		self.num_leaves = folds

	def get_training_batch(self):
		'''
		Method generating data for training and validation leaves.
		'''
		self.get_splits(self.kfold)
		for leaf in self.leaves:
			yield self._X[leaf[0]], self.enc.transform(self._Y[leaf[0]]).toarray(), self._X[leaf[1]], self.enc.transform(self._Y[leaf[1]]).toarray()

