from sklearn.linear_model import LinearRegression
from sklearn.linear_model import Ridge
from sklearn.linear_model import Lasso
from sklearn.metrics import r2_score, explained_variance_score
import statsmodels as stamo



def linear_regression(data, predictors):
    """ Calculate Linear Regression for a set of predictors
    Args:
        data (DataFrame): A Pandas DataFrame
        predictors (list): A list of strings that contains all the
            predictors keys
    Returns:
        list: with the regression results
    """
    # initialize predictors:
    if not predictors:
        predictors = ['x']

    # Fit the model
    linreg = LinearRegression(normalize=True)
    linreg.fit(data[predictors], data['y'])
    y_pred = linreg.predict(data[predictors])

    # Return the result in pre-defined format
    rss = sum((y_pred - data['y']) ** 2)
    print('RSS: ', rss)
    print('INTERCEPT: {}'.format(linreg.intercept_))
    # print('COEF: {}'.format(list(zip(predictors, [float(i) for i in linreg.coef_]))))
    print('COEF: {}'.format(list(zip(predictors, linreg.coef_))))
    print('Variance score: %.2f' % r2_score(data['y'], y_pred))
    ret = [rss]
    ret.extend([linreg.intercept_])
    ret.extend(linreg.coef_)
    return ret


def ridge_regression(data, predictors, alpha):
    """ Calculate Ridge Regression for a set of predictors
    Args:
        data (DataFrame): A Pandas DataFrame
        predictors (list): A list of strings that contains all the
            predictors keys
        alpha (float): with the coefficient penalty
    Returns:
        list: with the regression results
    """
    # Fit the model
    ridgereg = Ridge(alpha=alpha, normalize=True)
    ridgereg.fit(data[predictors][:10000], data['y'][:10000])
    y_pred = ridgereg.predict(data[predictors][10000:])

    # Return the result in pre-defined format
    rss = sum((y_pred - data['y'][10000:]) ** 2)
    ret = [rss]
    print('RSS: ', rss)
    print('INTERCEPT: {}'.format(ridgereg.intercept_))
    # print('COEF: {}'.format(
    #     list(zip(predictors, [float(i) for i in ridgereg.coef_]))))
    print('COEF: {}'.format(
        list(zip(predictors, ridgereg.coef_))))
    print('Variance score: %.2f' % r2_score(data['y'][10000:], y_pred))
    ret.extend([ridgereg.intercept_])
    ret.extend(ridgereg.coef_)
    return ret


def lasso_regression(data, predictors, alpha, models_to_plot={}):
    """ Calculate Lasso Regression for a set of predictors
    Args:
        data (DataFrame): A Pandas DataFrame
        predictors (list): A list of strings that contains all the
            predictors keys
        alpha (float): with the coefficient penalty
    Returns:
        list: with the regression results
    """
    # Fit the model
    lassoreg = Lasso(alpha=alpha, normalize=True, max_iter=1e5)
    lassoreg.fit(data[predictors], data['y'])
    y_pred = lassoreg.predict(data[predictors])

    rss = sum((y_pred - data['y']) ** 2)
    ret = [rss]
    print('RSS: ', rss)
    print('INTERCEPT: {}'.format(lassoreg.intercept_))
    print('COEF: {}'.format(
        list(zip(predictors, [float(i) for i in lassoreg.coef_]))))
    print('Variance score: %.2f' % r2_score(data['y'], y_pred))
    ret.extend([lassoreg.intercept_])
    ret.extend(lassoreg.coef_)
    return ret

def ols_regularization(formula, data, l_regu, alpha):
    """ OLS regression with regularization
    Args:
        formula (string): A string with the regression formula
        data (DataFrame): A Pandas DataFrame
        l_regu (float): 0 for Ridge and 1 for Lasso Regression
        alpha (float): with the coefficient penalty
    Returns:
        object: The regression object
    """
    model = sm.ols(formula, data=data)
    unregularized_regression = model.fit()
    result = sm.ols(
        formula, data=data).fit_regularized(
        method='elastic_net', L1_wt=l_regu, alpha=alpha,
        start_params=unregularized_regression.params, profile_scale=False,
        refit=False)
    print(result.params)
    final = stamo.regression.linear_model.OLSResults(model,
                                                     result.params,
                                                     model.normalized_cov_params)
    print(final.summary())


def ols_regression(formula, data):
    """ OLS regression
    Args:
        formula (string): A string with the regression formula
        data (DataFrame): A Pandas DataFrame
    Returns:
        object: The regression object
    """
    result = sm.ols(formula, data=data).fit()
    print(result.summary())
    return result