import pandas as pd
import numpy as np
import statsmodels.formula.api as smf
import statsmodels.api as sm
from R2OOS import R2OOS
from CWtest import CWtest

# =============================================================================
# Read data
# =============================================================================
ret=pd.read_excel('#Asset_Return.xlsx','Sheet1')
ret=ret[ret['YM']>=198501]# make sure all returns are after 198501
ret['RF']=ret['Rf']+1#capital RF is the gross risk free rate
ret['Return']=ret['ExRet']+ret['RF']
ret['ExRet12']=np.nan

#define the variables that we want to achieve
Table2=pd.DataFrame(columns=['Beta','T','R2','R2OOS','CW'],index=np.arange(1,56))

for i in range(1,56):
    ret_sub=ret[ret['ID']==i]
    ret_sub['ExRet12']=(pd.rolling_apply(ret_sub['Return'],12,lambda x : x.prod())-pd.rolling_apply(ret_sub['RF'],12,lambda x : x.prod()))/12
    ret_sub['ExRet_f1']=ret_sub['ExRet'].shift(-1)
    ret_sub['ExRet12_std']=(ret_sub['ExRet12']-ret_sub['ExRet12'].mean())/ret_sub['ExRet12'].std()
    result=smf.ols(formula="ExRet_f1 ~ ExRet12_std", data=ret_sub).fit(cov_type='HAC',cov_kwds={'maxlags':12})
    Table2.loc[[i],['Beta']]=result.params[1]*100
    Table2.loc[[i],['T']]=result.tvalues[1]
    Table2.loc[[i],['R2']]=result.rsquared*100

    r_hat=[];r_bar=[]    
    training_length=len(ret_sub[ret_sub['YM']<=199912])
    length=len(ret_sub)
    
    for j in range(training_length,length):
        lhs=ret_sub['ExRet'].iloc[12:j-1].reset_index(drop=True)
        rhs=ret_sub['ExRet12'].iloc[11:j-2].reset_index(drop=True)
        result=sm.OLS(lhs,sm.add_constant(rhs)).fit()
        r_hat.append(result.params[0]+result.params[1]*ret_sub['ExRet12'].iloc[j-1])
        r_bar.append(ret_sub['ExRet'].iloc[12:j-1].mean())
    
    r_real=np.array(ret_sub['ExRet'].iloc[training_length:])
    r_hat=np.array(list(r_hat))
    r_bar=np.array(list(r_bar))
    
    Table2.loc[[i],['R2OOS']]=R2OOS(r_real,r_hat,r_bar)*100
    Table2.loc[[i],['CW']]=CWtest(r_real,r_hat,r_bar,4)

writer = pd.ExcelWriter('Table2_InOutTSRegression.xlsx')
pd.DataFrame.to_excel(Table2,writer,'Table2_InOutTSRegression')