import pandas as pd
from datetime import datetime
import backtrader as bt
import matplotlib.pyplot as plt
import tushare as ts

plt.rcParams['font.sans-serif'] = ['SimHei']        # 设置画图时的中文显示
plt.rcParams['axes.unicode_minus'] = False          # 设置画图时的负号显示

# 1.数据加载
def get_data(code='600519',startTime='2017-01-01',endTime='2020-01-01'):
    df = ts.get_k_data(code,start=startTime,end = endTime)
    ## 设置为日期格式
    df.index = pd.to_datetime(df.date)
    # df['op']
    # print(df)
    df['openinterest'] = 0
    # print(df)
    df = df[['open','high','low','close','volume','openinterest']]
    # print(df)
    return df

stock_df = get_data()

# 加载并读取数据源 dataname:数据来源 fromdate(date格式):开始时间 todate:截至时间
fromdate = datetime(2017,1,1)
todate = datetime(2020,1,1)
data = bt.feeds.PandasData(dataname = stock_df,fromdate=fromdate,todate=todate)

# print(data)
# 2、构建自己的策略
# 上穿20日均线买入,跌穿20日均线卖出
class MyStrategy(bt.Strategy):
    params=(
        # 20日均线
        ('maperiod',20),
    )
    def __init__(self):
        self.order = None
        self.ma = bt.indicators.SimpleMovingAverage(self.datas[0],period=self.params.maperiod)

    # 每个bar都会执行一次,回测的每个日期都会执行一次
    def next(self):
        if self.order:
            # 如果有交易正在进行就返回
            return
        '''
        close[0]  是当天的日期收盘价
        close[-1] 是昨天的日期收盘价
        close[-2] 是前天的日期收盘价
        :return:
        '''
        # 如果现在是空仓状态
        if(not self.position):
            # 如果大于20日均线就买入
            if self.datas[0].close[0] > self.ma[0]:
                self.order = self.buy(size=200)
            else:
                if self.datas[0].close[0] <self.ma[0]:
                    self.order = self.sell(size=200)


# 3.策略设置
cerebro = bt.Cerebro()      # 创建大脑
# 将数据加入回测系统
cerebro.adddata(data)
# cerebro.adddata(data1)        # 可以加很多个不同的品种
# cerebro.adddata(data2)
# cerebro.adddata(data3)

# 加入自己的策略
cerebro.addstrategy(MyStrategy)

# 添加经纪人 初始化资金为 100000
start_cash = 100000
cerebro.broker.setcash(start_cash)
# 设置手续费 万分之2
cerebro.broker.setcommission(0.0002)

# 执行回测
s = fromdate.strftime("%Y-%m-%d")
t = todate.strftime("%Y-%m-%d")
print(f"初始资金:{start_cash}\n回测时间:{s}   {t}")
cerebro.run()
portval = cerebro.broker.getvalue()
print(f"策略执行完之后的资金:{portval}\n回测时间:{s}   {t}")