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()
stock_df1 = get_data(code='600419')

# 加载并读取数据源 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)
# 创建第二个数据集
data1 = bt.feeds.PandasData(dataname=stock_df1, fromdate=fromdate, todate=todate)


# 2 设计策略
class MyStrategy_SmaCross2(bt.Strategy):

    def __init__(self):

        # 过程参数
        self.index = 0 # 回测进度
        self.cache = 100 # 最少需要的K线数量

        # 策略固定参数
        self.sl_ls = [0.01,0.04]    # 止损参数
        self.tp_pt = [0.015,0.05]   # 止盈参数
        self.fast_ma_period = 7
        self.slow_ma_period = 20
        self.interval = 200

        # 策略动态参数
        self.sl_num = 0             # 止损次数
        self.tp_num = 0             # 止盈次数
        self.over_cross_num = 0          # 上传次数
        self.wait_num = 0           # 等待次数

        # 策略指标
        sma1 = bt.ind.SMA(period=self.fast_ma_period)       # 7日均线  快均线
        sma2 = bt.ind.SMA(period=self.slow_ma_period)       # 20日均线 慢均线
        self.crossover = bt.ind.CrossOver(sma1, sma2)

    # 每个bar都会执行一次,回测的每个日期都会执行一次
    def next(self):
        '''
        该策略将短周期均线二次上穿作为买入信号,采用分段止盈止损的方式离场
        :return:
        '''


        volume = self.position.size # 获取持仓情况
        price = self.position.price # 获取当前平均价
        value = self.broker.get_value() #获取当前账户净值
        qty = value/self.datas[0].close[0]
        if not self.position:       # 没有仓位
            if self.crossover[0] > 0 and self.over_cross_num == 0:
                self.over_cross_num+=1
            elif self.crossover[0]>0 and self.over_cross_num == 1:
                self.buy(size=qty*0.8)
                self.over_cross_num = 0
                print('Successful:实现两次上穿,成交')

            if self.wait_num>=self.interval and self.over_cross_num==1:
                # 一直没等到下一个上传信号,则重新计算
                self.over_cross_num = 0
                print("Fail:一直没等到上传信号,重新计算")

            if self.over_cross_num == 1:
                self.wait_num+=1
        else:           # 有仓位
            if self.sl_num + self.tp_num == 0: # 从未触发止盈止损
                if self.data.close[0] < price * (1-self.sl_ls[self.sl_num]):    # 如果触发第一次止损
                    self.close(size=volume/2)       # 减掉一半的仓位,进行一档止损
                    self.sl_num += 1
                    print("Half_Cut:触发一次止损:减一般仓")
                elif self.data.close[0] >= price*(1-self.tp_pt[self.tp_num]):    # 如果第一次触发止盈
                    self.close(size=volume / 2)  # 减掉一半的仓位,进行一档止盈
                    print("Half_Cut:触发一次止盈:减一般仓")
                    self.tp_num += 1

            elif self.sl_num==1 and self.tp_num == 0: # 仅触发一次止损
                if self.data.close[0] < price * (1-self.sl_ls[self.sl_num]):    # 如果触发第二次止损
                    self.close(size=volume)       # 清仓、也就是平仓
                    self.sl_num = 0
                    self.tp_num = 0
                    print("All_Cut:触发第二次止损:平仓")
                elif self.data.close[0] >= price*(1-self.tp_pt[self.tp_num]):    # 如果第一次触发止盈
                    self.close(size=volume)         # 清仓、也就是平仓
                    print("All_Cut:触发第一次止损后又触发止盈:平仓")
                    self.sl_num = 0
                    self.tp_num = 0

            elif self.sl_num==0 and self.tp_num == 1: # 仅触发一次止盈
                if self.data.close[0] < price * (1-self.sl_ls[self.sl_num]):    # 如果触发第一次止损
                    self.close(size=volume)       # 清仓、也就是平仓
                    print("All_Cut:触发第一次止盈后又触发止损:平仓")
                    self.tp_num = 0
                    self.sl_num = 0
                elif self.data.close[0] >= price*(1-self.tp_pt[self.tp_num]):    # 如果第二次触发止盈
                    self.close(size=volume)  # 清仓、也就是平仓
                    print("All_Cut:触发第二次止盈后:平仓")
                    self.sl_num = 0
                    self.tp_num = 0


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

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

# 添加经纪人 初始化资金为 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}")
cerebro.plot()