import backtrader as bt    # 引入backtrader 框架
import datetime
import pandas as pd
import numpy as np
import os                 # 用于路径管理 数据导入时使用

# 创建策略
class TestStrategy(bt.Strategy):
    def log(self, txt, dt=None):
        ''' 策略的日志函数'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))
    def __init__(self):
        # 引用data[0]数据的收盘价数据
        self.dataclose = self.datas[0].close
        # 用于记录订单状态
        self.order = None
        self.buyprice = None
        self.buycomm = None
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # 提交给代理或者由代理接收的买/卖订单 - 不做操作
            return
        # 检查订单是否执行完毕
        # 注意：如果没有足够资金，代理会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                print("\033[35m日期：{}，BUY EXECUTED,Price:{:.2f},Cost:{:.2f},Comm:{:.2f}\033[0m"
                      .format(self.datas[0].datetime.date(0),
                              order.executed.price, order.executed.value,order.executed.comm))
                # self.log(
                #     'BUY EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                #     (order.executed.price,
                #      order.executed.value,
                #      order.executed.comm))

                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            else: # 卖
                print("\033[32m日期：{}，SELL EXECUTED,Price:{:.2f},Cost:{:.2f},Comm:{:.2f}\033[0m"
                      .format(self.datas[0].datetime.date(0),
                              order.executed.price, order.executed.value,order.executed.comm))
                # self.log('SELL EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                #          (order.executed.price,
                #           order.executed.value,
                #           order.executed.comm))
            self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            # self.log('Order Canceled/Margin/Rejected')
            print("\033[31m Order Canceled/Margin/Rejected\033[0m")
        # 订单状态处理完成，设为空
        self.order = None
        # 交易成果
    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        # 显示交易的毛利率和净利润
        print("\033[31mOPERATION PROFIT, GROSS:{:.2f}, NET:{:.2f}\033[0m"
              .format(trade.pnl,trade.pnlcomm))
        # self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' %
        #          (trade.pnl, trade.pnlcomm))
    def next(self):
        # 日志输出收盘价数据
        # self.log('Close, %.2f' % self.dataclose[0])
        # 检查是否有订单等待处理，如果是就不再进行其他下单
        if self.order:
            return
        # 检查是否已经进场
        if not self.position:
            # 还未进场，则只能进行买入
            # 当日收盘价小于前一日收盘价
            if self.dataclose[0] < self.dataclose[-1]:
                    # 前一日收盘价小于前前日收盘价
                    if self.dataclose[-1] < self.dataclose[-2]:
                        # 买买买
                        self.log('BUY CREATE, %.2f' % self.dataclose[0])
                        # 记录订单避免二次下单
                        self.order = self.buy()
        # 如果已经在场内，则可以进行卖出操作
        else:
            # 卖卖卖
            if len(self) >= (self.bar_executed + 5):
                self.log('SELL CREATE, %.2f' % self.dataclose[0])
                # 记录订单避免二次下单
                self.order = self.sell()

# 创建cerebro实例
cerebro = bt.Cerebro()
# 添加策略
cerebro.addstrategy(TestStrategy)

# 添加策略开始和结束日期
start_time = datetime.datetime(2018, 1, 1)
end_time = datetime.datetime(2019, 12, 31)

# 先找到脚本文件的位置
modepath = os.path.dirname(__file__)
# 然后根据脚本与数据的相对关系找到数据位置,这样脚本从任意地方被调用，都可以正确地访问到数据
# datapath = modepath+'/data/000998.SZ.csv'
# datapath = os.path.join(modepath, './data/000998.SZ.csv')
datapath = os.path.join(modepath, '600000qfq.csv')

"""# data feeds 创建 通过GenericCSVData方式，Backtrader里使用的GenericCSVData函数，
要求CSV文件数据是按时间升序排列的。因此，如果股票数据未按时间升序排序，需要我们在做回测前对数据重新排序，否则会导致回测结果错误。
data = bt.feeds.GenericCSVData(
        dataname=datapath,
        fromdate=start_time,
        todate=end_time,
        nullvalue=0.0,
        dtformat='%Y-%m-%d',       # 自动把index数据中的符合日期的格式变成datetime类型
        datetime=0,
        open=2,
        high=3,
        low=4,
        close=5,
        volume=9,
        openinterest=-1
)"""

# 通过PandasData 方式导入dataframe数据,此方式较为直观推荐此种方案
"""直接在数据导入时，对列表名称处理好，usecols是抽取的列，header =0 取首行为列名，names进行改列名
在PandasData时，可以自载匹配数据"""
df = pd.read_csv(datapath,usecols=[2,3,4,5,6,10],header=0,
                 names=['datetime','open','high','low','close','volume'],
                 index_col=0,parse_dates=True)
df['openinterest'] = 0
df['ma5'] = 0
df['ma20'] = 20
# data = bt.feeds.PandasData(dataname=df,
#                            fromdate=start_time,
#                            todate=end_time)

# 拓展技术数据
class NewLines_PandasData(bt.feeds.PandasData):
    # 新增加两条数据线
    lines =('ma5','ma20',)
    # 新增数据在dataframe中的位置，最好使用-1进行自动匹配
    # -1表示自动按列明匹配数据，也可以设置为线在数据源中列的位置索引 (('pe',7),('pb',8),)
    params = (
                ('ma5',-1),
                ('ma20',-1),
              )

data =NewLines_PandasData(dataname=df,
                            fromdate=start_time,
                            todate=end_time)
# 在cerebro中添加价格数据
cerebro.adddata(data)
# 设置启动资金
cerebro.broker.setcash(100000.0)
# 设置佣金率为千分之一
cerebro.broker.setcommission(commission=0.001)
# 打印开始资金
print('Starting Portfolio Value:%.2f' % cerebro.broker.getvalue())
# 遍历所有数据
cerebro.run()
# 打印最后结果
print('Final Portfolio Value:%.2f' % cerebro.broker.getvalue())
# cerebro.plot(style='candle',numfigs=4)