# Linear regression Project
#
# Chris Pecau

import numpy as np
from asserts import assert_Xy, assert_Xbeta


def prepend_ones_column(A):
    '''
    Add a ones column to the left side of an array

    Inputs: 
        A: a numpy array

    Output: a numpy array
    '''
    ones_col = np.ones((A.shape[0], 1))
    return np.hstack([ones_col, A])


def linear_regression(X, y):
    '''
    Compute linear regression. Finds model, beta, that minimizes
    X*beta - Y in a least squared sense.

    Accepts inputs with type array
    Returns beta, which is used only by apply_beta

    Examples
    --------
    >>> X = np.array([[5, 2], [3, 2], [6, 2.1], [7, 3]]) # predictors
    >>> y = np.array([5, 2, 6, 6]) # dependent
    >>> beta = linear_regression(X, y)  # compute the coefficients
    >>> beta
    array([ 1.20104895,  1.41083916, -1.6958042 ])
    >>> apply_beta(beta, X) # apply the function defined by beta
    array([ 4.86363636,  2.04195804,  6.1048951 ,  5.98951049])
    '''
    assert_Xy(X, y, fname='linear_regression')

    X_with_ones = prepend_ones_column(X)

    # Do actual computation
    beta = np.linalg.lstsq(X_with_ones, y)[0]

    return beta


def apply_beta(beta, X):
    '''
    Apply beta, the function generated by linear_regression, to the
    specified values

    Inputs:
        model: beta as returned by linear_regression
        Xs: 2D array of floats

    Returns:
        result of applying beta to the data, as an array.

        Given:
            beta = array([B0, B1, B2,...BK])
            Xs = array([[x11, x12, ..., x0K],
                        [x21, x22, ..., x1K],
                        ...
                        [xN1, xN2, ..., xNK]])

            result will be:
            array([B0+B1*x11+B2*x12+...+BK*x1K,
                   B0+B1*x21+B2*x22+...+BK*x2K,
                   ...
                   B0+B1*xN1+B2*xN2+...+BK*xNK])
    '''
    assert_Xbeta(X, beta, fname='apply_beta')

    # Add a column of ones
    X_incl_ones = prepend_ones_column(X)

    # Calculate X*beta
    yhat = np.dot(X_incl_ones, beta)
    return yhat


def read_file(filename):
    '''
    Read data from the specified file.  Split the lines and convert
    float strings into floats.  Assumes the first row contains labels
    for the columns.

    Inputs:
      filename: name of the file to be read

    Returns:
      (list of strings, 2D array)
    '''
    with open(filename) as f:
        labels = f.readline().strip().split(',')
        data = np.loadtxt(f, delimiter=',', dtype=np.float64)
        return labels, data


def load_files(data_set):
    col_names, data = read_file("data/" + data_set + "/training.csv")
    return col_names, data

def load_testing_files(data_set):
    col_names, data = read_file("data/" + data_set + "/testing.csv")
    return col_names, data   


def rsquared(data, X, y, beta=[], verbose=False):
    '''
    Generate R^2 for a given list of columns.

    Inputs:
        data: loaded data_set array 
        X: list of columns to calculate R2 for
        y: column integer of the dependent variable
        beta: (optional) run R2 on another model's betas

    Returns:
        R2, list of betas
    '''

    dependent = data[:,y]
    x_values = data[:,X]
    if beta == []:
        beta = linear_regression(x_values, dependent)

    yhat = apply_beta(beta, x_values)
    mu = np.mean(dependent)

    sum_errors_sq = np.sum(np.power(dependent - yhat, 2))
    sum_dist_Y_mean_sq = np.sum(np.power(dependent - mu, 2))

    R2 = 1 - (sum_errors_sq / sum_dist_Y_mean_sq)

    if verbose:
        print("Regressing columns {} on {}: R2 = {}".format(X, y, R2))
    
    return R2, beta


def maximize_r2(data, X, y):
    '''
    Find highest bivariate R2 by generating each combination
    and selecting the highest value.

    Inputs:
        data: loaded data_set array
        X: list of columns of predictors
        y: column integer of the dependent variable

    Returns:
        highest_bivariate: dictionary with highest R2 pair
    '''

    bivariates = {}
    highest_bivariate = {}
    highest_bivariate['R2'] = 0

    for i in X:
        for j in X:
            if i != j:  # Avoid identical predictors
            ### GRADER COMMENT
            ### PENALTY -2
            # Because the order of the variables doesn't matter, j should loop starting
            # from i+1. This is more efficient because it only needs to check half as 
            # many values. 
                pair = [i] + [j] 
                pair.sort()   # Avoid duplicate, reversed predictors
                pair_key = str(pair) # Use pair string as key

                if pair_key not in bivariates:

                    # Generate R2 for new pair key
                    bivariates[pair_key] = {}
                    R2_value, beta = rsquared(data, pair, y)
                    bivariates[pair_key]['R2'] = R2_value

                    # If highest, record as highest R2
                    if R2_value > highest_bivariate['R2']:
                        highest_bivariate['R2'] = R2_value
                        highest_bivariate['x_columns'] = pair

    return highest_bivariate

def create_indices(data, y):
    '''
    Separates out predictor columns from dependent variable column

    Inputs:
        data: loaded data_set array
        y: column integer of the dependent variable

    Returns:
        X: list of columns of predictors
        y: column integer of the dependent variable 
    '''

    # Create list of length of number of rows in data_set
    X = list(range(data.shape[1]))
    X.pop(y)    # Remove dependent variable from list

    return X, y


def print_betas(col_names, X, output=True):
    '''
    Generates a list of headers associated with a list of predictors,
    and prints them.

    Input:
        col_names: list of column headers

    Returns:
        betas_list: list of corresponding column_headers
    '''

    betas_list = []

    for column in X:
        betas_list.append(col_names[column])
    
    if output:
        names_print = ", ".join(betas_list)
        print(names_print)

    return betas_list



def generate_betas(data, X, y):
    '''
    Generates all the betas for a list of predictors

    Inputs:
        data: loaded data_set array
        X: list of columns of predictors
        y: column integer of the dependent variable

    Returns:
        betas: dictionary of betas, with predictor columns as keys
    '''
    
    betas = {}

    for x_column in X:
        betas[x_column] = rsquared(data, [x_column], y)
    return betas


def increase_k_to_threshold(col_names, data, X, y, threshold=0, verbose=False):
    '''
    Greedy algorithm that calculates the highest R2 for each
    K (total number of betas), and adds them to a list.

    Inputs:
        col_names: list of column headers
        data: loaded data_set array
        X: list of columns of predictors
        y: column integer of the dependent variable
        threshold (optional): minimum increase in R2 to advance to next K

    Returns:
        k_values: dictionary with best value at each K, including:
              'winners':  list of predictors for best R2
               'bestR2':  best R2 at each K
            'best_beta':  column of new, best R2 at this K
                'betas':  list of betas at this K [for testing.csv]
    '''


    winners = []
    k_values = {}
    k_key = 1
    prev_best_R2 = 0
    curr_best_R2 = 0
    difference = 0

    for k in range(len(X)):

        k_values[k_key] = {}

        # Run for all values remaining in each new K-value
        for column in [x for x in X if x not in winners]:

            # Initialize bestR2 for each K
            if 'bestR2' not in k_values[k_key]:
                k_values[k_key]['bestR2'] = 0

            # Add each remaining predictor, in turn, to winner list
            # then calculate the R2
            current_X = winners + [column]
            R2, beta = rsquared(data, current_X, y)

            # If R2 is best so far in this K, record as best
            # 
            # Also record---
            #       winners:  list of predictors for best R2
            #     best_beta:  column of new, best R2 at this K
            #         betas:  list of betas at this K [for testing.csv]
            #       
            if R2 > k_values[k_key]['bestR2']:
                k_values[k_key]['winners']   = current_X
                k_values[k_key]['bestR2']    = R2
                k_values[k_key]['best_beta'] = column
                k_values[k_key]['betas'] = beta
                if verbose:
                    print("Adding R2", k_values[k_key]['bestR2'])

        # if verbose:
        #     k_values[k_key]['winners'] = winners

        ### GRADER COMMENT
        # You should remove commented out code before submitting.

        # Only calculate and record difference if K > 1
        if k_key > 1:
            prev_best_R2 = k_values[k_key-1]['bestR2']
            curr_best_R2 = k_values[k_key]['bestR2']
            difference   = curr_best_R2 - prev_best_R2

        # Calculate whether threshold is reached (if included as arg)
        if threshold != 0 and k_key > 1 and difference < threshold:

            print_betas(col_names, k_values[k_key - 1]['winners'])
            print("R2: {0:.2f}".format(prev_best_R2))

            if verbose:
                print("Prev-", prev_best_R2)
                print("Best-", curr_best_R2)
                print("Dif =", difference)
                print("Thres", threshold)
                print("Previous R2", k_values[k_key - 1]['best_beta'])
                print("Hit threshold..")
                print()

            return(k_values)

        # Add best R2 from this K to winners list
        winners.append(k_values[k_key]['best_beta'])
        
        # Print results from each K, if no threshold (Task 3a)
        if threshold == 0:
            print_betas(col_names, winners)
            if k_key != 1:
                print("R2: {0:.2f}".format(curr_best_R2)) 
            else:
                print("R2: {0:.2f}".format(k_values[k_key]['bestR2'])) 

        k_key += 1

    return k_values
