from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.ensemble import RandomForestClassifier
import numpy as np
import pickle


class UnbalancedSamplingClf(BaseEstimator,ClassifierMixin):
	'''
	Unbalanced Sampling Classifier: The user sepecifies a base learner (default is Random Forest).  The model building process
	takes all of the minority data points (1s) and subsamples the appropriate proportion of 0s such that it produces a number of
	different subsets of the base data, where each subset is balanced.  A number of models are iteartively trained over each subset
	and stored in the model_set list.  The models are then used as an ensemble to generate predictions. 
	Parameters:
	--------------------
	base_estimator: the model trained in each iteration of the model building process (default is RandomForestClassifier)
	majority_to_minority_ratio: the ratio of the majority to the minority class used in the sampling process
	iterations: the number of times the sampling process is ran (i.e. the total number of models built in the ensemble)
	cutoff: this is the threshold used in voting or probability calculation used by the predict method
	prediction: this is the type of prediction the model will be generating, it can either be generated by using labels of the models or using their generated
				probabilities (more sensitive to the cutoff threshold)
	random_state: used for results persistence
	--------------------
	'''

	def __init__(self, 
				 base_estimator=RandomForestClassifier(),
			     majority_to_minority_ratio=1, 
			     iterations=50,
			     cutoff = 0.5,
			     prediction='labels',
			     random_state=0):

		self.base_estimator = base_estimator
		self.majority_to_minority_ratio = majority_to_minority_ratio
		self.iterations = iterations
		self.cutoff = cutoff
		self.prediction = prediction
		self.random_state = random_state
		self.model_stack = []


	def fit(self,X,y):
		
		'''trains the model - the data is first split into a subset of the positive minority class (1s) and (0s) and then
		the 0s are subsampled in each iteration of the model according to the ratio specified upon initializing
		and array of models are then subsequently trained on various subsamples of the 1s and 0s class
		Parameters:
		--------------------
		X: feature set [n_samples,m_features]
		y: target vector [n_samples, 1]
		Returns:
		self.model_set
		--------------------
		'''
		ones_indeces = np.where(y==1)[0]
		n_ones = len(ones_indeces)
		zeros_indeces = np.where(y==0)[0]
		ones_X, ones_y = X[ones_indeces], y[ones_indeces]
		zeroes_X, zeroes_y = X[zeros_indeces], y[zeros_indeces]

		for iteration in range(self.iterations):

			#setting random seed for the iteration
			np.random.seed(iteration)

			#generating random indexes to be sampled from the majority class (done without replacement)
			indeces = np.random.choice(zeroes_X.shape[0], n_ones*self.majority_to_minority_ratio, replace= False)

			sample_X = zeroes_X[indeces]
			sample_y = zeroes_y[indeces]

			#stacking the majority class data with the minority class data
			ones_data = np.hstack([ones_y.reshape(-1,1),ones_X])
			zeroes_data = np.hstack([sample_y.reshape(-1,1),sample_X])

			#stacking vertically into a single data frame and shuffling
			all_data = np.vstack([ones_data,zeroes_data])
			np.random.shuffle(all_data)

			#training a classifier on the data
			self.base_estimator.fit(all_data[:,1:], all_data[:,0])

			#adding a picklezied version of the model to the model stack
			self.model_stack.append(pickle.dumps(self.base_estimator))

		#de-pickling all the models in the the stack
		self.model_stack = [pickle.loads(model) for model in self.model_stack]

	def predict(self,X):
			
		'''makes a class prediction using the model
		Parameters:
		-----------
		X: feature set (usually unseen data) format: [n_samples, m_features]
		Returns:
		y_pred: an array of predicted classes for the target variable'''
		
		if self.prediction == 'labels':

			predictions = np.array([list(base_estimator.predict(X)) for base_estimator in self.model_stack])

		if self.prediction == 'probabilities':

			predictions = np.array([list(base_estimator.predict_proba(X)[:,1]) for base_estimator in self.model_stack])

		y_pred_proba = predictions.mean(axis=0)

		return np.where(y_pred_proba > self.cutoff, 1, 0)

	
	def predict_proba(self,X):

		'''makes a probability prediction using the model
		Parameters:
		-----------
		X: feature set (usually unseen data) format: [n_samples, m_features]
		Returns:
		y_pred: an array of predicted probabilities for the target variable'''

		if self.prediction == 'labels':

			predictions = np.array([list(base_estimator.predict(X)) for base_estimator in self.model_stack])

		if self.prediction == 'probabilities':

			predictions = np.array([list(base_estimator.predict_proba(X)[:,1]) for base_estimator in self.model_stack])

		return predictions.mean(axis=0)
