import numpy as np
from .helpers import lls_gradient


class LinearRegressionModel(object):
    """Model for univariate multiple linear regression.

    The model supports ridge regression.
    """

    def __init__(self, data_set, features, target, lda=0, reg_type='ridge'):
        """Initialize a LinearRegressionModel object.

        Parameters
        ----------
        data_set : epysteme.DataSet
            The data set used to train and evaluate the model.
        features : iterable[str]
            Names of the features to use from the data set.
        target : str
            Name of the target feature in the data set.
        lda : optional[float]
            Lambda for ridge regression, default is 0.
        reg_type : optional[str]
            Regularization type, can be 'ridge' (default) or 'lasso'
        """
        self._data_set = data_set
        self._features = features
        self._target = target
        self.lda = lda
        self.reg_type = reg_type.lower()

        m = len(self._features) + 1
        self.beta_hat = np.zeros((m, 1))
        self.trained = False

    def train(self, subset=(), gradient_method=False):
        """Train the model.

        Training is done by finding the closed-form solution to the linear
        least squares problem (regularised if lambda is non-zero).

        Parameters
        ----------
        subset : optional[str]
            Name of the subset on which to train the model.
        gradient_method : optional[bool]
            Forces use of gradient method.
        """
        if gradient_method is True:
            self.beta_hat = self._gradient_descent(subset)
        else:
            self.beta_hat = self._closed_form(subset)
            # TODO: catch unstable/singular and propose gradient method

        self.trained = True

    def _cost(self, beta, subset):
        cost = 0
        for x, y in self._data_set.data(self._features,
                                        self._target,
                                        const=True,
                                        subsets=subset):
            cost += np.linalg.norm(np.dot(x, beta) - y)**2
        return cost

    def _closed_form(self, subset=()):
        """Compute the closed form solution.

        Parameters
        ----------
        subset : optional[str -or- iterable[str]]
            Subset or subsets on which to train the model.

        Returns
        -------
        numpy.ndarray
            Closed form solution of the linear regression problem.
        """
        m = len(self._features) + 1
        xtx = np.zeros((m, m))
        xty = np.zeros((m, 1))

        for x, y in self._data_set.data(self._features,
                                        self._target,
                                        const=True,
                                        subsets=subset):

            xtx += np.dot(x.transpose(), x)
            xty += np.dot(x.transpose(), y)

        beta_hat = np.dot(np.linalg.inv(xtx + self.lda * np.eye(m)), xty)
        return beta_hat

    def _gradient_descent(self, subset=()):
        """Train the model using gradient descent.

        Parameters
        ----------
        """
        # TODO: variable iterations and/or break condition
        beta_hat = np.zeros(self.beta_hat.shape)
        step_size = 0.01

        x, y = next(self._data_set.data(self._features,
                                        self._target,
                                        const=True,
                                        subsets=subset))
        if len(y) == self._data_set.n_observations:
            in_memory = True
        else:
            in_memory = False

        for k in range(10000):
            grad_beta = np.zeros(beta_hat.shape)
            if in_memory is not True:
                for x, y in self._data_set.data(self._features,
                                                self._target,
                                                const=True,
                                                subsets=subset):
                    grad_beta += lls_gradient(x, y, beta_hat)
            else:
                grad_beta = lls_gradient(x, y, beta_hat)

            n = np.linalg.norm(grad_beta)
            if n > 1:
                grad_beta *= (1/n)

            beta_hat -= step_size*grad_beta

        return beta_hat

    def predict(self, obs):
        """Predict value(s) of target feature for the given observation(s).

        Parameters
        ----------
        obs : array-like
            Observation(s) for which to predict the target feature.

        Returns
        -------
        float -or- numpy.ndarray
            Prediction(s) of the target feature.
        """
        return np.dot(self.beta_hat, obs)
