# 多周期
# 买入条件：价格回踩2级支撑或突破1级阻力
# 卖出条件：价格较最高收盘价回撤5%卖出
import backtrader as bt
import backtrader.feeds as btfeeds
import backtrader.indicators as btind
import datetime
import pandas as pd

stk_list = ['minutes1-JD2201']


class PivotMultiTF(bt.Strategy):
    params = (
        ('lowestperiod', 5),
        ('trailamount', 0.0),
        ('trailpercent', 0.05),
    )

    def __init__(self):
        # 存储不同数据的技术指标
        self.inds = dict()
        # 存储特定股票的订单，key为股票的代码
        self.orders = dict()
        # 遍历所有数据
        for i, d in enumerate(self.datas):
            self.orders[d._name] = None
            # 为每个数据定义字典，存储技术指标
            self.inds[d] = dict()
            # 判断d是否为日线数据
            if 0 == i % 2:
                self.inds[d]['lowest'] = btind.Lowest(d, period=self.p.lowestperiod)
            # 判断d是否为月线数据
            else:
                # 定义pivot point指标
                self.inds[d]['pp'] = btind.PivotPoint(d)

            # # 跳过第一只股票data，第一只股票data作为主图数据
            # if i > 0:
            #     d.plotinfo.plotmaster = self.datas[0]

    def next(self):
        for i, d in enumerate(self.datas):
            # 如果处理月线数据则跳过买卖条件，因为已在日线数据判断处理过
            if 1 == i % 2:
                continue
            pos = self.getposition(d)
            # 不在场内，则可以买入
            if not len(pos):
                # 达到买入条件
                month_pp = self.inds[self.datas[i + 1]]['pp']
                if (self.inds[d]['lowest'] <= month_pp.s2 and d.close > month_pp.s2) or (
                        self.inds[d]['lowest'] <= month_pp.r1 and d.close > month_pp.r1):
                    # 买入手数
                    stake = int(self.broker.cash / len(stk_list) // (d.close[0] * 100)) * 100
                    # 买买买
                    self.buy(data=d, size=stake)
            elif not self.orders[d._name]:
                # 下保护点卖单
                self.orders[d._name] = self.close(data=d, exectype=bt.Order.StopTrail,
                                                  trailamount=self.p.trailamount,
                                                  trailpercent=self.p.trailpercent)

    def notify_order(self, order):
        if order.status in [order.Completed]:
            if order.isbuy():
                print('{} BUY {} EXECUTED, Price: {:.2f}'.format(self.datetime.datetime(), order.data._name,
                                                                 order.executed.price))
            else:  # Sell
                self.orders[order.data._name] = None
                print('{} SELL {} EXECUTED, Price: {:.2f}'.format(self.datetime.datetime(), order.data._name,
                                                                  order.executed.price))

# 加载数据
def load_data(stk_code, fromdate, todate):
    datapath = stk_code +'.csv'
    return bt.feeds.GenericCSVData(
        dataname=datapath,
        fromdate=fromdate,
        todate=todate+datetime.timedelta(days=1),
        nullvalue=0.0,
        dtformat='%Y-%m-%d %H:%M:%S',
        timeframe=bt.TimeFrame.Minutes,
        datetime=1,
        open=2,
        high=3,
        low=4,
        close=5,
        volume=6,
        openinterest=-1
    )

cerebro = bt.Cerebro()
cerebro.broker.setcash(1000000.0)
cerebro.addstrategy(PivotMultiTF)
fromdate = datetime.datetime(2021, 11, 2)
todate = datetime.datetime(2021, 12, 2)
for stk_code in stk_list:
    data = load_data(stk_code, fromdate, todate)
    cerebro.adddata(data)
    cerebro.resampledata(data, timeframe=bt.TimeFrame.Minutes, compression=15)
cerebro.addwriter(bt.WriterFile, out='log.csv', csv=True)
cerebro.run()
print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
# Plot the result绘制结果
cerebro.plot(start=datetime.date(2021, 12, 1), end=datetime.date(2021, 12, 3),
             volume=False, style='candle',
             barup='red', bardown='green')