import pandas as pd
import numpy as np
from jqdata import gta
import matplotlib.pyplot as plt
import datetime 
import statsmodels.api as sm

steel_fd_data = pd.read_excel('钢铁基本面数据.xls',index_col=0)
###由于数据是wind直接下来的，所以不太干净，先预处理一下、去掉一些注释行
steel_fd_data = steel_fd_data[1:-2]
##数据钢材表观消费量为主，截断此数据开始之前的日期，此数据从2004年1月开始
steel_fd_data = steel_fd_data[steel_fd_data.index>=datetime.datetime(2004,1,1)]

need_adjust_columns = ['房屋新开工面积:累计值','房屋竣工面积:累计值','商品房销售面积:累计值','固定资产投资完成额:累计值',
                       '固定资产投资完成额:基础设施建设投资:累计值','房地产开发投资完成额:累计值']

steel_fd_data['month'] = steel_fd_data.index.map(lambda x:x.month)
#这些列都是一月份缺失数据的
for column in need_adjust_columns:
    if column != '固定资产投资完成额:基础设施建设投资:累计值':
        steel_fd_data[column.split(':')[0]] = steel_fd_data[column] - steel_fd_data[column].shift(1)
        steel_fd_data.ix[steel_fd_data['month']<=2,column.split(':')[0]] = steel_fd_data[steel_fd_data['month']<=2][column]
    else:
        steel_fd_data[column.split(':')[1]] = steel_fd_data[column] - steel_fd_data[column].shift(1)
        steel_fd_data.ix[steel_fd_data['month']<=2,column.split(':')[0]] = steel_fd_data[steel_fd_data['month']<=2][column]

###钢铁行业收入一月份没有缺失
column_ = '钢铁行业:大中型企业:销售收入:累计值'
steel_fd_data['钢铁行业销售收入'] = steel_fd_data[column_] - steel_fd_data[column_].shift(1)
steel_fd_data.ix[steel_fd_data['month'] == 1,'钢铁行业销售收入'] = steel_fd_data[steel_fd_data['month']==1][column_]
        
#取出接下来需要用的主要数据
data_columns_list = steel_fd_data.columns.tolist()
clean_data = steel_fd_data[data_columns_list[:2]+data_columns_list[-7:]+['产量:汽车:当月值']]
#去除第一行
clean_data = clean_data[1:]

clean_data.fillna(0,inplace=True)
TTM_columns = ['表观消费量:钢材:当月值','房屋新开工面积','房屋竣工面积','固定资产投资完成额','基础设施建设投资',\
               '商品房销售面积','房地产开发投资完成额','钢铁行业销售收入','产量:汽车:当月值']
for column in TTM_columns:
    clean_data[column+'TTM'] = pd.rolling_sum(clean_data[column],12)
    clean_data[column+'TTM_环比'] = clean_data[column+'TTM']/clean_data[column+'TTM'].shift(1)-1

clean_data['year'] = clean_data.index.map(lambda x:x.year)
clean_data['month'] = clean_data.index.map(lambda x:x.month)

clean_data.ix[(clean_data['year'].isin([2016,2017]))&(clean_data['month'].isin([1,2])),'表观消费量:钢材:当月值TTM_环比'] = 0

clean_data.ix[(clean_data['year'].isin([2014,2015,2016,2017]))&(clean_data['month'].isin([1,2])),'产量:汽车:当月值TTM_环比'] = 0

clean_data = clean_data[clean_data['year']>=2006]
clean_data.ix[(clean_data['year'] == 2006)&(clean_data['month']==8),'商品房销售面积TTM_环比'] = 0

use_columns = [i+'TTM_环比' for i in TTM_columns]+['M2:同比']
use_data = clean_data[use_columns]
use_data.dropna(inplace=True)

###由于列名太长，这里做一个简写,但其值是表示的为TTM_环比值(但M2为同比值)
use_data.columns = ['钢材消费量TTM环比','新开工房TTM环比','竣工房TTM环比','固投TTM环比',\
                    '基建TTM环比','房销售面积TTM环比','房地产投资TTM环比','钢铁行业收入TTM环比','汽车产量TTM环比','M2']

def plot_line_chart(x_axis,y1_axis,y2_axis,label1,label2):
    
    Fig = plt.figure(figsize(10,8))
    Ax = Fig.add_subplot(111)
    lines = Ax.plot(x_axis,y1_axis,'r-',x_axis,y2_axis,'b-')
    lines[0].set_label(label1)
    lines[1].set_label(label2)
    Ax.legend(loc = 0)

def plot_line_chart_double_yaxis(x_axis,y1_axis,y2_axis,label1,label2):
    
    Fig = plt.figure(figsize(10,8))
    Ax = Fig.add_subplot(111)
    Ax.plot(x_axis,y1_axis,'r-')
    Ax.set_ylabel(label1)
    
    Ax_ = Ax.twinx()
    Ax_.plot(x_axis,y2_axis,'b-')
    Ax_.set_ylabel(label2)
    
def calculate_lag(data,column1,column2):
    
    regre_results=[]
    for i in range(-5,6):
        reg_data = pd.concat([data[column1].shift(i),data[column2]],axis =1)
        reg_data.dropna(inplace=True)
        y=reg_data[column1]
        x=reg_data[column2]
        x=sm.add_constant(x)   
        est=sm.OLS(y,x)
        results=est.fit()
        r2 = results.rsquared 
        coef = results.params[1]
        p = results.pvalues[1]
        corr = reg_data[column1].corr(reg_data[column2])
        
        regre_results.append([i,corr,coef,p,r2])
        
    regre_df = pd.DataFrame(regre_results,columns=['领先期','相关系数','回归系数','P值','R方'])
    
    return regre_df

    

plot_line_chart(use_data.index,use_data['钢材消费量TTM环比'],use_data['新开工房TTM环比'],'钢材表观消费量:TTM环比','房屋新开工面积:TTM环比')
chart_xkg_gxf = calculate_lag(use_data,'新开工房TTM环比','钢材消费量TTM环比') 
chart_xkg_gxf

plot_line_chart(use_data.index,use_data['新开工房TTM环比'],use_data['房销售面积TTM环比'],'房屋新开工面积:TTM环比','商品房销售面积:TTM环比')
chart_xkg_fxs = calculate_lag(use_data,'新开工房TTM环比','房销售面积TTM环比') 
chart_xkg_fxs

plot_line_chart_double_yaxis(use_data.index,use_data['房销售面积TTM环比'],use_data['钢材消费量TTM环比'],'商品房销售面积:TTM环比','钢材消费量:TTM环比')
chart_fxs_gxf = calculate_lag(use_data,'房销售面积TTM环比','钢材消费量TTM环比') 
chart_fxs_gxf

plot_line_chart_double_yaxis(use_data.index,use_data['竣工房TTM环比'],use_data['钢材消费量TTM环比'],'房屋竣工面积:TTM环比','钢材消费量:TTM环比')
chart_jgf_gxf = calculate_lag(use_data,'竣工房TTM环比','钢材消费量TTM环比') 
chart_jgf_gxf

plot_line_chart(use_data.index,use_data['固投TTM环比'],use_data['钢材消费量TTM环比'],'固定投资完成额:TTM环比','钢材消费量:TTM环比')
chart_gt_gxf = calculate_lag(use_data,'固投TTM环比','钢材消费量TTM环比') 
chart_gt_gxf

plot_line_chart(use_data.index,use_data['基建TTM环比'],use_data['钢材消费量TTM环比'],'基础设施建设投资:TTM环比','钢材消费量:TTM环比')
chart_gt_gxf = calculate_lag(use_data,'基建TTM环比','钢材消费量TTM环比') 
chart_gt_gxf

plot_line_chart(use_data.index,use_data['房地产投资TTM环比'],use_data['钢材消费量TTM环比'],'房地产投资完成额:TTM环比','钢材消费量:TTM环比')
chart_gt_gxf = calculate_lag(use_data,'房地产投资TTM环比','钢材消费量TTM环比') 
chart_gt_gxf

plot_line_chart_double_yaxis(use_data.index,use_data['汽车产量TTM环比'],use_data['钢材消费量TTM环比'],'汽车产量:TTM环比','钢材消费量:TTM环比')
chart_gt_gxf = calculate_lag(use_data,'汽车产量TTM环比','钢材消费量TTM环比') 
chart_gt_gxf

plot_line_chart_double_yaxis(use_data.index,use_data['M2'],use_data['钢材消费量TTM环比'],'M2','钢材消费量:TTM环比')
chart_gt_gxf = calculate_lag(use_data,'M2','钢材消费量TTM环比') 
chart_gt_gxf

###以前一年数据作为样本进行滚动回归，得到回归系数，以此系数预测下一个月钢材表观消费量TTM环比增速
###商品房销售面积、基础设施建设投资、M2的领先期分别选为2个月、4个月、2个月
testdata = use_data[['钢材消费量TTM环比','房销售面积TTM环比','基建TTM环比','M2']]

def predict_TTM_growth(testdata,y_var,x_var,x_lag,regre_period = 12):
    
    x_var_lag = []
    for x,lag in zip(x_var,x_lag):
        testdata[x+'_lag'+str(lag)] = testdata[x].shift(lag)
        x_var_lag.append(x+'_lag'+str(lag))
        use_testdata = testdata[[y_var] + x_var_lag]
        use_testdata.dropna(inplace=True) 
        use_testdata[y_var+'_预测'] = np.NaN
    
    corr_chart = use_testdata[x_var_lag].corr()
    print (corr_chart)
    for i in range(regre_period,len(use_testdata)):
    
        train_data = use_testdata[i-regre_period:i]
        y = train_data[y_var]
        x = train_data[x_var_lag]
        x = sm.add_constant(x)
        est = sm.OLS(y,x)
        results = est.fit()
    
        predict_data = use_testdata[i:i+1]
        predict_x = predict_data[x_var_lag]
        predict_x = sm.add_constant(predict_x)

        
        use_testdata.iloc[i,-1] = results.predict(predict_x)
    
    return use_testdata
    

predict_data = predict_TTM_growth(testdata,'钢材消费量TTM环比',['房销售面积TTM环比','基建TTM环比','M2'],[2,4,2])
plot_line_chart(predict_data.index,predict_data['钢材消费量TTM环比'],predict_data['钢材消费量TTM环比_预测'],'钢材消费量TTM环比','钢材消费量TTM环比预测')

steel_price = steel_fd_data[['钢材价格综合指数']]
steel_price.fillna(method='ffill',inplace = True)
steel_price['钢材价格综合指数_lag2'] = steel_price['钢材价格综合指数'].shift(2)
steel_price['钢材价格综合指数_lag14'] = steel_price['钢材价格综合指数'].shift(14)

tdata = use_data[['钢材消费量TTM环比','房销售面积TTM环比','基建TTM环比','M2']]

###以前一年数据作为样本进行滚动回归，得到回归系数，以此系数预测下一个月钢材表观消费量TTM环比增速
testdata = pd.concat([tdata[['房销售面积TTM环比','基建TTM环比','M2']],\
                      steel_price[['钢材价格综合指数']],\
                      clean_data[['钢铁行业销售收入TTM_环比']][:-3]],axis=1)


predict_data_ = predict_TTM_growth(testdata,'钢铁行业销售收入TTM_环比',\
                                   ['房销售面积TTM环比','基建TTM环比','钢材价格综合指数','钢材价格综合指数','M2'],[2,4,2,14,2])
plot_line_chart(predict_data_.index,predict_data_['钢铁行业销售收入TTM_环比'],predict_data_['钢铁行业销售收入TTM_环比_预测'],\
                '钢铁行业销售收入TTM_环比','钢铁行业销售收入TTM_环比_预测')