#导入需要的程序包
from research_api import *
import pandas as pd
import seaborn as sns
import numpy as np
import math
import matplotlib.pyplot as plt

def Strategy(df,parameter):
    df[df['volume']==0]=np.nan
    #计算策略指标，标记出买卖点
    shortday = parameter[0]
    longday = parameter[1]
    close = df['close']
    ewma5 = pd.ewma(close,shortday)
    ewma10 = pd.ewma(close,longday)
    Strategy_table = pd.concat([ewma5,ewma10],keys=['ewma5','ewma10'],axis=1)
    Strategy_table['open']=df['open']
    Strategy_table['yesterday']=(Strategy_table['ewma5']<Strategy_table['ewma10']).shift(1)#计算昨天的
    Strategy_table['today']=(Strategy_table['ewma5']>Strategy_table['ewma10'])#计算今天
    Strategy_table['buypoint'] = Strategy_table['yesterday'] & Strategy_table['today']
    Strategy_table['sellpoint'] = Strategy_table['yesterday'].apply(lambda x:not x)&Strategy_table['today'].apply(lambda x:not x)
    Strategy_table = Strategy_table.dropna()
    return Strategy_table

def Buysellstrategy(init_money,buy_cost,sell_cost,Strategy_table):
    cash = [init_money]
    holdstocknum =[0]
    buysellpoint=[0]
    for i in range(1,len(Strategy_table)-1):
        cash.append(cash[i-1])
        holdstocknum.append(holdstocknum[i-1])
        buysellpoint.append('0')
        if Strategy_table['buypoint'][i]==True: #定义买入策略，根据标记出的买点以第二天开盘价买入，包括买入手续费
            current_cash = cash[i]
            buyprice = Strategy_table['open'][i+1]#明日开盘价
            number_of_shares = int((current_cash-buy_cost)/buyprice)
            if number_of_shares > 0:
                buysellpoint[i]=1
                cash[i] = cash[i]-buyprice*number_of_shares
                holdstocknum[i] = number_of_shares
        elif Strategy_table['sellpoint'][i]==True:#定义卖出策略，根据标记出的卖点以第二天开盘价卖出，包括卖出手续费
            if holdstocknum[i]>0:            
                buysellpoint[i]=-1
                curret_cash = cash[i]
                sellprice = Strategy_table['open'][i+1]#明日开盘价
                cash[i] = cash[i] + sellprice*holdstocknum[i]-sell_cost
                holdstocknum[i] = 0
    cash.insert(0,np.nan)
    holdstocknum.insert(0,np.nan)
    buysellpoint.append(np.nan)
    Strategy_table['cash']=cash
    Strategy_table['holdstocknum']=holdstocknum
    Strategy_table['buysellpoint']=buysellpoint#标记出实际买卖点
    Strategy_table['portfolio_value']=Strategy_table['cash']+Strategy_table['open']*Strategy_table['holdstocknum']
    #returns = 100*(Strategy_table['portfolio_value'][-1]/Strategy_table['portfolio_value'][1]-1)#计算收益
    returns = 100*Strategy_table['portfolio_value'].pct_change().cumsum()[-1]
    #计算最大回撤
    maxdrawdown = 100*((max(Strategy_table['portfolio_value'].dropna())-min(Strategy_table['portfolio_value'].dropna()))/max(Strategy_table['portfolio_value'].dropna()))
    return returns,maxdrawdown,Strategy_table

init_money = 10000
start_date='2013-01-01'
end_date='2015-10-18'
#days = [5,10,20,30,60,120]
parameter = [5,10]
#选择股票池，这里选择沪深300股票数据
stocklist = get_index_stocks('000300.XSHG')
buy_cost=5#0.0008
sell_cost=7#0.0013
#min_cost=5

Returns = []
Maxdrawdown = []
for j in range(0,len(stocklist)):
    stock = stocklist[j]#获取股票代码
    df = get_price(stock, start_date, end_date, frequency='daily')
    Strategy_table = Strategy(df,parameter)
    returns,maxdrawdown,Strategy_table = Buysellstrategy(init_money,buy_cost,sell_cost,Strategy_table)
    Returns.append(returns)
    Maxdrawdown.append(maxdrawdown)
Backtest = pd.DataFrame(index=stocklist)
Backtest['returns']=Returns
Backtest['maxdrawdown']=Maxdrawdown
Backtest  

Backtest.describe()

Backtest[Backtest['returns']>150]

df = get_price('600415.XSHG', start_date, end_date, frequency='daily')
Strategy_table = Strategy(df,parameter)
returns,maxdrawdown,Strategy_table = Buysellstrategy(init_money,buy_cost,sell_cost,Strategy_table)
Strategy_table['open'].plot(figsize=[18,8],label='open price')
plt.plot(Strategy_table.index[Strategy_table['buysellpoint']==1],Strategy_table['open'][Strategy_table['buysellpoint']==1],'r^',label='buypoint')
plt.plot(Strategy_table.index[Strategy_table['buysellpoint']==-1],Strategy_table['open'][Strategy_table['buysellpoint']==-1],'gv',label='sellpoint')
plt.legend(loc='best')

df = get_price('000300.XSHG', start_date, end_date, frequency='daily')
(100*df['close'].pct_change().cumsum()).plot(figsize=[18,8],label='Benchmark Returns')
df = get_price('600415.XSHG', start_date, end_date, frequency='daily')
Strategy_table = Strategy(df,parameter)
returns,maxdrawdown,Strategy_table = Buysellstrategy(init_money,buy_cost,sell_cost,Strategy_table)
(100*Strategy_table['portfolio_value'].pct_change().cumsum()).plot(label='Total Returns')
plt.legend(loc='best')

Strategy_table[400:]