import numpy as np
# from matplotlib import pyplot

from pyml.utils import normalize_features
from pyml.utils import normalize_inputs

class LinearRegression(object):
    """A linear regression learning algorithm.

    Args:
        X (np.array): Matrix of features (training data).
        y (np.array): Vector of training labels.
        alpha (float): Learning rate.
        iterations (int): Number of gradient descent iterations.

    Class Attributes:
        _FEATURE_THRESHOLD (int): Min. number of features to revert to
            gradient descent.

    Properties:
        m (int): Size of training set.
        n (int): Size of feature set.
        theta (np.array): Linear regression hypothesis function coefficients.

    Usage:
        Note: The train() method defaults to using the normal equation to solve
            for the best-fit theta vector analytically; if the force_grad_desc
            arg is set to 'True' or the number of features is greater than
            _FEATURE_THRESHOLD, however,  gradient descent will be used to
            find the best-fit theta vector.

        algo = LinearRegression(csv, alpha=0.01, iterations=1000)
        (theta, costFnHistory) = algo.train()

    """

    _FEATURE_THRESHOLD = 10 * 1000

    def __init__(self, X, y, alpha, iterations):
        self._X = X
        self._y = y
        self._alpha = alpha
        self._iterations = iterations

        self._m = None
        self._n = None
        self._theta = None
        self._cost_fn_history = None

        # For feature scaling / mean normalization.
        self._features_scaled = False
        self._means = None
        self._sigmas = None

    @property
    def m(self):
        """Size of training set."""
        if self._m is None:
            X_rows, _ = self._X.shape
            Y_size = self._y.size
            if X_rows == Y_size:
                self._m = Y_size
            else:
                raise ValueError(
                    '{t_set_size} != {l_set_size}: Training set size not'\
                    ' equal to # of training labels.'.format(
                        t_set_size=X_rows,
                        l_set_size=Y_size,
                    )
                )
        return self._m

    @property
    def n(self):
        """Size of feature set."""
        if self._n is None:
            _, self._n = self._X.shape
        return self._n

    @property
    def theta(self):
        """Initialize hypothesis fn's theta vector."""
        if self._theta is None:
            # Plus one to account for theta0.
            self._theta = np.zeros(shape=(self.n + 1, 1))
        return self._theta

    @property
    def cost_fn_history(self):
        """Stores cost function (J(theta)) history for gradient descent."""
        return self._cost_fn_history

    def _compute_cost(self, X, y, theta):
        """Computes the squared-error between hypothesized outputs and
        labels given current theta weights and a labeled training set.
        """
        predictions = X.dot(theta).flatten()
        squareErrors = (predictions - y)
        cost = (1 / float(2 * self.m)) * squareErrors.T.dot(squareErrors)
        return cost

    def _use_normal_eq(self):
        return True if self.n <= self._FEATURE_THRESHOLD else False

    def _normal_eq(self):
        """Finds theta vector that minimizes the cost function analytically.
        The normal equation is as follows: (X^T X)^-1 X^T y.
        """
        temp = self._X
        y = self._y

        # Add ones feature column to X to account for theta0 coefficient.
        X = np.ones(shape=(self.m, self.n + 1))
        X[:, 1:] = temp

        theta = (np.linalg.pinv(X.T.dot(X)).dot(X.T)).dot(y)
        return theta

    def _gradient_descent(self):
        """Performs gradient descent to find the optimal theta coefficients by
        minimizing the cost function and updating the theta coefficients
        simultaneously.
        """
        m = self.m
        temp = self._X
        y = self._y
        theta = self.theta
        alpha = self._alpha
        iterations = self._iterations
        cost_fn_history = np.zeros(shape=(iterations, 1))

        # Add ones feature column to X to account for theta0 coefficient.
        X = np.ones(shape=(self.m, self.n + 1))
        X[:, 1:] = temp

        for i in range(iterations):
            # Calculate predictions.
            predictions = X.dot(theta).flatten()

            # Calculate partial derivatives of cost function, J(theta).
            partials = (1 / float(m)) * (predictions - y).dot(X).flatten()

            # Update thetas.
            theta[:, 0] = theta[:, 0] - alpha * partials

            # Update costFnHistory for plotting purposes.
            cost_fn_history[i, 0] = self._compute_cost(X, y, theta)

        return theta.flatten(), cost_fn_history

    def train(self, force_grad_desc=False, normalize=True):
        """Trains a linear regression hypothesis function using either the
        normal equation (solves for theta vector analytically) or gradient
        descent -- depending on the number of features and the force_grad_desc
        argument. Additionally, if the normalize argument is set to 'True,'
        features will be scaled and mean-normalized.
        """
        if self._use_normal_eq() and not force_grad_desc:
            # cost_fn_history is None for normal equation method.
            theta, cost_fn_history = self._normal_eq(), None
        else:
            if normalize:
                self._X, self._means, self._sigmas = normalize_features(
                    self._X
                )
                self._features_scaled = True
            theta, cost_fn_history = self._gradient_descent()

        self._theta, self._cost_fn_history = theta, cost_fn_history

    def predict(self, *args):
        """Predicts outputs based on trained theta vector and inputs; will
        automatically scale features if features were scaled during training.
        """
        temp = np.array(args)
        if self._features_scaled:
            temp = normalize_inputs(temp, self._means, self._sigmas)
        inputs = np.ones(shape=(1, self.n + 1))
        inputs[:, 1:] = temp
        return inputs.dot(self.theta).flatten()


if __name__ == "__main__":
    CSV_PATH = 'RealEstate.csv'
    csv = np.genfromtxt(CSV_PATH, delimiter = ',', skip_header=1, dtype=float, unpack=False, usecols=(4,3,2))

    normalized_features, means, stds = normalize_features(csv[:, :-1])

    algo = LinearRegression(csv[:, :-1], csv[:, -1], 0.01, 10000)
    algo.train(force_grad_desc=True, normalize=True)

    algo2 = LinearRegression(csv[:, :-1], csv[:, -1], 0.01, 10000)
    algo2.train()

    print("Theta calculated by gradient descent")
    print(algo.theta)

    print("Prediction via gradient descent:")
    print(algo.predict(3, 4))

    print("Theta calculated by the normal equation")
    print(algo2.theta)

    print("Prediction via the normal equation")
    print(algo2.predict(3, 4))

    print(algo.cost_fn_history);

    # Plot costFnHistory.
    # pyplot.figure(2)
    # pyplot.plot(algo.cost_fn_history)
    # pyplot.show()
    # it = np.ones(shape=(csv[:, -1].size, 2))
    # it[:, 1:] = csv[:, :-1]

    # values = it.dot(algo.theta)
    # pyplot.figure(1)
    # pyplot.plot(csv[:, :-1], csv[:, -1], 'bs', csv[:, :-1], values, 'r--')
    # pyplot.show()
