import backtrader as bt
import pandas as pd
import os
import datetime
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHer']

class Platform(bt.Indicator):
    lines = ("up", "down") # up 上轨 down 下轨
    params = (
        ("period", 5),
    )

    def __init__(self):
        self.addminperiod(self.params.period + 1)   #5天的平台

    def next(self):
        self.up[0] = max(self.data.high.get(ago=-1, size=self.params.period))
        self.down[0] = min(self.data.low.get(ago=-1, size=self.params.period))

class MyStrategy(bt.Strategy):
    params = (
        ("period", 5),
    )

    def __init__(self):
        self.up_down = Platform(self.data, period= self.params.period)
        self.up_down = self.up_down()
        self.up_down.plotinfo.plotmaster = self.data0

        self.buy_signal = bt.indicators.CrossOver(self.datas[0].close, self.up_down.up) # 买入信号
        self.sell_signal = bt.indicators.CrossDown(self.datas[0].close, self.up_down.down) # 卖出信号
        # self.order = None
        self.buy_signal.plotinfo.plot = False
        self.sell_signal.plotinfo.plot = False
        self.flag = 'sell'



    def start(self):
        pass

    def prenext(self):
        print(f'数据准备时间：{self.data0.datetime.datetime(0)}')

    def nextstart(self):
        pass

    def next(self):
        if self.buy_signal[0] == 1 and self.flag == 'sell':
            self.order = self.buy(size=1000)
            self.flag = 'buy'
            print(f"{self.datas[0].datetime.date(0)},买入！价格为{self.data.close[0]}")
        else:
            if self.sell_signal[0] == 1 and self.flag == 'buy':
                self.order = self.sell(size=1000)
                self.flag = 'sell'
                print(f"{self.datas[0].datetime.date(0)},卖出！价格为{self.data.close[0]}")

    def stop(self):
        if self.position:
            self.order = self.sell(size=1000)
            print(f"{self.datas[0].datetime.date(0)},卖出！价格为{self.data.close[0]}")


if __name__ == '__main__':
    cerebro = bt.Cerebro()
    cerebro.broker.set_cash(100000.00)  # 设置初始资金金额
    cerebro = bt.Cerebro(stdstats=False)
    cerebro.addobserver(bt.observers.Broker)
    # cerebro.addobserver(bt.observers.Trades)
    cerebro.addobserver(bt.observers.BuySell)
    # cerebro.addobserver(bt.observers.DrawDown)
    cerebro.addobserver(bt.observers.Value)
    # cerebro.addobserver(bt.observers.TimeReturn)

    init_fund = cerebro.broker.getvalue()
    print(f'初始资金:{init_fund}')

    filename = 'sz000002.csv'
    print(filename)
    data = pd.read_csv(filename, index_col="date", parse_dates=True)
    # print(data)
    # 也可以如下处理日期
    # data.index=pd.to_datetime(data.date)
    # data.drop(columns=["date"],inplace=True)
    # print(data)

    load_data = bt.feeds.PandasData(dataname=data, fromdate=datetime.datetime(2020, 1, 1), todate=datetime.datetime(2020, 10, 18))
    cerebro.adddata(load_data)
    cerebro.addstrategy(MyStrategy)

    #####################
    cerebro.run()
    end_fund = cerebro.broker.getvalue()
    print(f'期末资金:{end_fund}')

    """
    处理错误
    cannot import name 'warnings' from 'matplotlib.dates
    pip uninstall matplotlib
    pip install matplotlib==3.2.2
    """
    cerebro.plot(style="candle")