import pandas as pd
import numpy as np
from itertools import chain, combinations

import statsmodels.formula.api as smf
import scipy.stats as scipystats
from scipy.stats import norm
from scipy.stats.stats import pearsonr
import statsmodels.api as sm
import statsmodels.stats.stattools as stools
import statsmodels.stats as stats 
from patsy import dmatrices
from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.graphics.regressionplots import *
import matplotlib.pyplot as plt
import math
mypath =  r'F:\Dropbox\ipynb_notes\ipynb_blog\\'

indata = pd.read_pickle(mypath + "data.pkl")

pd.options.display.max_rows
pd.set_option('display.max_colwidth', -1)

from IPython.display import display, HTML

def pearson_corr(indata, xvar, yvar):
    outdf = pd.DataFrame()
    for i in range(len(xvar)):
        pcorr = pearsonr(indata[xvar[i]], indata[yvar])
        corrdf = pd.DataFrame({"var_name": xvar[i], "pearson_corr":pcorr[0], "p_value":pcorr[1]}, index = [0])
        outdf = pd.concat([outdf, corrdf], axis = 0)
        outdf.sort_values('p_value')
    outdf.columns = corrdf.columns
    return outdf

def pearson_calc(indata, xvar, yvar):
    pearsoncorr = pearson_corr(indata, xvar, yvar)
    pearsoncorr = pearsoncorr.sort_values(by = ['p_value'])
    pearsoncorr = pearsoncorr.query('p_value <= 0.1')
    pearsoncorr = pearsoncorr.reset_index(drop = True)
    return pearsoncorr

xvarcol = indata.columns[1:-3]
yvarcol = indata.columns[0]
pearsoncorr = pearson_calc(indata, xvarcol, yvarcol)
pearsoncorr.head()

varsign = {'bbb_spread_p': 1, 'ted_spread_p': 1, 'unemp_p': 1, 'realgpd_yoy_p': -1, 'cap_rate_p': 1, 'crea_hp_yoy_p': -1,    \
    'corp_profit_p': -1, 'com_price_index_yoy_p': -1, 'spindex_yoy_p': -1}

def check_sign(indata, varsign):
    varsign = pd.DataFrame.from_dict(varsign, orient = "index").reset_index()
    varsign.columns = ["var_name_abb", "sign"]
    indata['var_name_abb'] = indata.var_name.map(lambda x: x[:-1])
    indata = pd.merge(indata, varsign, how = "left", on = ['var_name_abb'])
    indata['sign_true'] = np.where(indata.pearson_corr > 0, 1, -1) * indata.sign
    return indata.query('sign_true == 1')

output = check_sign(pearsoncorr, varsign)
output.sort_values('var_name')
output.head()

xvar = output.var_name
cniX = indata.ix[:, xvar]
cniY = indata.ix[:, "cum_pd_num"]    
train_flag = indata.train_flag

xtrain, xtest = cniX[train_flag], cniX[~train_flag] 
ytrain, ytest = cniY[train_flag], cniY[~train_flag] 

print xtrain.shape
print xtest.shape

(91, 28)
(39, 28)

def verify_sign(params):
    setsign = varsign
    setsign = pd.DataFrame.from_dict(setsign, orient = "index").reset_index()
    setsign.columns = ['vars', 'coefsign']
    params = params.reset_index()
    params.columns = ['vars', 'coefd']
    params['vars'] = params.vars.map(lambda x: x[:-1])
    combinedata = pd.merge(setsign, params, how = 'inner', on = 'vars')
    combinedata['coefd_'] = combinedata.coefsign * np.sign(combinedata.coefd)
    if sum(combinedata['coefd_'] > 0) == combinedata.shape[0]:
        return True
    else:
        return False


def check_dup_vars(xvar):
    xx = xvar
    xxx = set([x[:-1] for x in xx])
    if len(xxx) == len(xx):
        return True
    else:
        return False

def all_combination_withsign2(xtrain, ytrain, xtest, ytest, k_features = [1]):
    outdf = pd.DataFrame()
    features = xtrain.columns.tolist()
    n_features = xtrain.shape[1]
    subsets = chain.from_iterable(combinations(xrange(n_features), k) for k in k_features)
    best_score = -np.inf
    all_combination = None
    for subset in subsets:
        newxtrain = xtrain.iloc[:, subset]
        newxtrain = sm.add_constant(newxtrain)
        vif = dict(zip(newxtrain.columns, [variance_inflation_factor(newxtrain.values, i) for i in range(newxtrain.shape[1])])) 
        del vif['const']
        vif_max = max(vif.itervalues())
        newxtest = xtest.iloc[:, subset]
        newxtest = sm.add_constant(newxtest)        
        lin_reg = sm.OLS(ytrain, newxtrain).fit()       
        #lin_reg = sm.OLS(ytrain, newxtrain).fit(cov_type='HAC',cov_kwds={'maxlags':1})
        if check_dup_vars(lin_reg.params.index.tolist()):        
            r2train = lin_reg.rsquared
            r2test1 = np.sum((lin_reg.predict(newxtest) - np.mean(ytest))**2) / np.sum((ytest - np.mean(ytest))**2)
            r2test = 1 - np.sum((lin_reg.predict(newxtest) -  ytest)**2) / np.sum((ytest - np.mean(ytest))**2)            
            r2test2 = 1 - np.var(ytest - lin_reg.predict(newxtest)) / np.var(ytest)       
            rmse_train = math.sqrt(sum((lin_reg.predict(newxtrain) - ytrain)**2) / (len(ytrain) - len(lin_reg.params)))
            rmse_train2 = math.sqrt(lin_reg.mse_resid)
            pred = lin_reg.predict(newxtest)        
            rmse_test = math.sqrt(sum((pred - ytest)**2)/(len(ytest) - len(lin_reg.params)))
            if (verify_sign(lin_reg.params)):
                sss = " ".join([features[i] for i in subset])
                pvalues = lin_reg.pvalues.drop('const').to_dict(); pvalues_max = max(pvalues.itervalues())
                dd = {'vars':sss, 'r2train':r2train, 'r2test':r2test, 'rmse_train':rmse_train,  'rmse_test':rmse_test, 'pvalues': [pvalues], \
                    'pvalues_max': pvalues_max, "reg_coef": [dict(lin_reg.params)], "vif":[vif], "vif_max":vif_max}
                df_i = pd.DataFrame(dd, index = [0], columns = dd.keys())  
                outdf = pd.concat([df_i, outdf], axis = 0)
                score = lin_reg.rsquared
                if (score > best_score) & (vif_max < 5) & (pvalues_max < 0.1):                  
                    best_score, all_combination = score, subset
                    best_reg_result = lin_reg

    outdf = outdf[["vars", "reg_coef", "pvalues", "pvalues_max", "vif", "vif_max", "r2train", "r2test", "rmse_train", "rmse_test"]]                                       
    return all_combination, best_score, best_reg_result, outdf

all_combination, best_score, best_reg_result, outdf  = all_combination_withsign2(xtrain, ytrain, xtest, ytest, k_features = [3])

outdf.sort_values('r2train', ascending = False).query('r2train > 0.5')

all_combination

(5, 14, 18)

best_reg_result.summary()