#!/usr/bin/env python

"""
docstring
"""

# Authors: Henrik Linusson

import abc
import numpy as np

from sklearn.base import BaseEstimator


class RegressorMixin(object):
	def __init__(self):
		super(RegressorMixin, self).__init__()

	@classmethod
	def get_problem_type(cls):
		return 'regression'


class ClassifierMixin(object):
	def __init__(self):
		super(ClassifierMixin, self).__init__()

	@classmethod
	def get_problem_type(cls):
		return 'classification'


class BaseModelAdapter(BaseEstimator):
	__metaclass__ = abc.ABCMeta

	def __init__(self, model, fit_params=None):
		super(BaseModelAdapter, self).__init__()

		self.model = model
		self.last_x, self.last_y = None, None
		self.clean = False
		self.fit_params = {} if fit_params is None else fit_params

	def fit(self, x, y):
		"""Fits the model.

		Parameters
		----------
		x : numpy array of shape [n_samples, n_features]
			Inputs of examples for fitting the model.

		y : numpy array of shape [n_samples]
			Outputs of examples for fitting the model.

		Returns
		-------
		None
		"""

		self.model.fit(x, y, **self.fit_params)
		self.clean = False

	def predict(self, x):
		"""Returns the prediction made by the underlying model.

		Parameters
		----------
		x : numpy array of shape [n_samples, n_features]
			Inputs of test examples.

		Returns
		-------
		y : numpy array of shape [n_samples]
			Predicted outputs of test examples.
		"""
		if (
			not self.clean or
			self.last_x is None or
			self.last_y is None or
			not np.array_equal(self.last_x, x)
		):
			self.last_x = x
			self.last_y = self._underlying_predict(x)
			self.clean = True

		return self.last_y.copy()

	@abc.abstractmethod
	def _underlying_predict(self, x):
		"""Produces a prediction using the encapsulated model.

		Parameters
		----------
		x : numpy array of shape [n_samples, n_features]
			Inputs of test examples.

		Returns
		-------
		y : numpy array of shape [n_samples]
			Predicted outputs of test examples.
		"""
		pass


class ClassifierAdapter(BaseModelAdapter):
	def __init__(self, model, fit_params=None):
		super(ClassifierAdapter, self).__init__(model, fit_params)

	def _underlying_predict(self, x):
		return self.model.predict_proba(x)


class RegressorAdapter(BaseModelAdapter):
	def __init__(self, model, fit_params=None):
		super(RegressorAdapter, self).__init__(model, fit_params)

	def _underlying_predict(self, x):
		return self.model.predict(x)


class OobMixin(object):
	def __init__(self, model, fit_params=None):
		super(OobMixin, self).__init__(model, fit_params)
		self.train_x = None

	def fit(self, x, y):
		super(OobMixin, self).fit(x, y)
		self.train_x = x

	def _underlying_predict(self, x):
		# TODO: sub-sampling of ensemble for test patterns
		oob = x == self.train_x

		if hasattr(oob, 'all'):
			oob = oob.all()

		if oob:
			return self._oob_prediction()
		else:
			return super(OobMixin, self)._underlying_predict(x)


class OobClassifierAdapter(OobMixin, ClassifierAdapter):
	def __init__(self, model, fit_params=None):
		super(OobClassifierAdapter, self).__init__(model, fit_params)

	def _oob_prediction(self):
		return self.model.oob_decision_function_


class OobRegressorAdapter(OobMixin, RegressorAdapter):
	def __init__(self, model, fit_params=None):
		super(OobRegressorAdapter, self).__init__(model, fit_params)

	def _oob_prediction(self):
		return self.model.oob_prediction_
