import numpy as np
import pandas as pd
import statsmodels.api as sm
import seaborn as sns

# 输入是一DataFrame，每一列是一支股票在每一日的价格
def find_cointegrated_pairs(dataframe):
    # 得到DataFrame长度
    n = dataframe.shape[1]
    # 初始化p值矩阵
    pvalue_matrix = np.ones((n, n))
    # 抽取列的名称
    keys = dataframe.keys()
    # 初始化强协整组
    pairs = []
    # 对于每一个i
    for i in range(n):
        # 对于大于i的j
        for j in range(i+1, n):
            # 获取相应的两只股票的价格Series
            stock1 = dataframe[keys[i]]
            stock2 = dataframe[keys[j]]
            # 分析它们的协整关系
            result = sm.tsa.stattools.coint(stock1, stock2)
            # 取出并记录p值
            pvalue = result[1]
            pvalue_matrix[i, j] = pvalue
            # 如果p值小于0.05
            if pvalue < 0.05:
                # 记录股票对和相应的p值
                pairs.append((keys[i], keys[j], pvalue))
    # 返回结果
    return pvalue_matrix, pairs

instruments = ["002142.SZA", "600000.SHA", "600015.SHA", "600016.SHA", "600036.SHA", "601009.SHA",
              "601166.SHA", "601169.SHA", "601328.SHA", "601398.SHA", "601988.SHA", "601998.SHA"]

# 确定起始时间
start_date = '2015-01-01' 
# 确定结束时间
end_date = '2017-07-18' 
# 获取股票总市值数据，返回DataFrame数据格式
prices_temp = D.history_data(instruments,start_date,end_date,
              fields=['close'] )
prices_df=pd.pivot_table(prices_temp, values='close', index=['date'], columns=['instrument'])
pvalues, pairs = find_cointegrated_pairs(prices_df)
#画协整检验热度图，输出pvalue < 0.05的股票对
sns.heatmap(1-pvalues, xticklabels=instruments, yticklabels=instruments, cmap='RdYlGn_r', mask = (pvalues == 1))
print(pairs)

df = pd.DataFrame(pairs, index=range(0,len(pairs)), columns=list(['Name1','Name2','pvalue']))
#pvalue越小表示相关性越大,按pvalue升序排名就是获取相关性从大到小的股票对
df.sort_values(by='pvalue')

T.plot(prices_df[['601328.SHA','601998.SHA']], chart_type='line', title='Price')

# ols
x = prices_df['601328.SHA']
y = prices_df['601998.SHA']
X = sm.add_constant(x)
result = (sm.OLS(y,X)).fit()
print(result.summary())

import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(8,6))
ax.plot(x, y, 'o', label="data")
ax.plot(x, result.fittedvalues, 'r', label="OLS")
ax.legend(loc='best')

# T.plot(pd.DataFrame({'Stationary Series':0.8602*x-y, 'Mean':[np.mean(0.8602*x-y)]}), chart_type='line')
df = pd.DataFrame({'Stationary Series':y-0.8602*x, 'Mean':np.mean(y-0.8602*x)})
T.plot(df, chart_type='line', title='Stationary Series')

def zscore(series):
    return (series - series.mean()) / np.std(series)

zscore_calcu = zscore(y-0.8602*x)
T.plot(pd.DataFrame({'zscore':zscore_calcu, 'Mean':np.mean(y-0.8602*x), 'upper':1, 'lower':-1}) ,chart_type='line', title='zscore')

instrument = {'y':'601998.SHA','x':'601328.SHA'}  # 协整股票对
start_date = '2015-01-01' # 起始日期   
end_date = '2017-07-18' # 结束日期

# 初始化账户和传入需要的变量
def initialize(context):
    context.set_commission(PerDollar(0.0015)) # 手续费设置
    context.zscore = zscore_calcu # 交易信号需要根据zscore_calcu的具体数值给出
    context.ins  = instrument # 传入协整股票对
    
# 策略主题函数   
def handle_data(context, data):
    
    date = data.current_dt.strftime('%Y-%m-%d') # 运行到当根k线的日期
    zscore = context.zscore.ix[date]  # 当日的zscore
    stock_1 = context.ins['y'] # 股票y
    stock_2 = context.ins['x'] # 股票x
    
    symbol_1 = context.symbol(stock_1) # 转换成回测引擎所需要的symbol格式
    symbol_2 = context.symbol(stock_2)
    
    # 持仓
    cur_position_1 = context.portfolio.positions[symbol_1].amount
    cur_position_2 = context.portfolio.positions[symbol_2].amount
   
    
    # 交易逻辑
    # 如果zesore大于上轨（>1），则价差会向下回归均值，因此需要买入股票x，卖出股票y
    if zscore > 1 and cur_position_2 == 0 and data.can_trade(symbol_1) and data.can_trade(symbol_2):  
        context.order_target_percent(symbol_1, 0)
        context.order_target_percent(symbol_2, 1)
        print(date, '全仓买入：交通银行')
        
    # 如果zesore小于下轨（<-1），则价差会向上回归均值，因此需要买入股票y，卖出股票x
    elif zscore < -1 and cur_position_1 == 0 and data.can_trade(symbol_1) and data.can_trade(symbol_2):  
        context.order_target_percent(symbol_1, 1)
        print(date, '全仓买入：中信银行')
        context.order_target_percent(symbol_2, 0)     

# 回测启动接口
m=M.trade.v2( 
    instruments=list(instrument.values()),# 保证instrument是有字符串的股票代码组合成的列表（list）
    start_date=start_date,
    end_date=end_date,
    initialize=initialize,
    handle_data=handle_data,
    order_price_field_buy='open',
    order_price_field_sell='open',
    capital_base=10000,
    benchmark='000300.INDX',
)