# coding=utf-8
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import backtrader as bt
import datetime as datetime
import pandas as pd


# Create a Stratey
class TestStrategy(bt.Strategy):
    params = dict(
        exitbars=5,
        pfast=10,
        pslow=30
    )

    def __init__(self):
        # Keep a reference to the "close" line in the data[0] dataseries
        self.dataclose = self.datas[0].close
        # To keep track of pending orders
        self.order = None
        self.buyprice = None
        self.buycomm = None
        self.bar_executed = None
        self.sma = bt.ind.SMA(period=self.p.pfast)

    def log(self, txt, dt=None):
        """ Logging function for this strategy"""
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    # order 状态，为了收到订单通知，需要在用户自定义的Strategy子类中，重写notify_order方法，该方法默认为空
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # Buy/Sell order submitted/accepted to/by broker - Nothing to do
            return

        # Check if an order has been completed
        # Attention: broker could reject order if not enough cash
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log("买单成交，成交价：%.3f，成交额：%.2f，佣金：%.2f" % (
                    order.executed.price,
                    order.executed.value,
                    order.executed.comm))
            elif order.issell():
                self.log("卖单成交，成交价：%.3f，成交额：%.2f，佣金：%.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("订单取消/资金不足/拒绝")
        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log("毛利润 %.2f, 净利润 %.2f" % (trade.pnl, trade.pnlcomm))

    def next(self):
        # Simply log the closing price of the series from the reference
        self.log('收盘, %.3f' % self.dataclose[0])

        # Check if an order is pending ... if yes, we cannot send a 2nd one
        if self.order:
            return

        # 检查是否已有持仓
        if not self.position:
            if self.dataclose[0] > self.sma[0]:  # 今收盘 小于昨收
                self.log("BUY CREATE, %.3f" % self.dataclose[0])
                self.buy()
        else:
            # Already in the market ... we might sell
            # if len(self) >= (self.bar_executed + self.params.exitbars):
            if self.dataclose[0] < self.sma[0]:
                self.log("SELL CREATE, %.2f" % self.dataclose[0])
                self.order = self.sell()


if __name__ == '__main__':
    cerebro = bt.Cerebro()
    # Set our desired cash start

    # Add a strategy
    cerebro.addstrategy(TestStrategy)

    # Create a Data Feed
    df_feeder = bt.feeds.GenericCSVData(  # 默认 OHLC
        dataname="tdx2csv/csv/sz300033.day.csv",
        # fromdate=datetime.datetime(2019, 1, 1),
        todate=datetime.datetime(2019, 12, 20),
        dtformat="%Y%m%d",  # "%Y-%m-%d",
        openinterest=-1  # -1表示该字段不是存在于CSV data
    )
    # 添加数据
    cerebro.adddata(df_feeder, name="sz300033.day.csv")

    # 设置初始金额，佣金
    cerebro.broker.setcash(100000.0)
    cerebro.broker.setcommission(commission=0.001)  # 佣金 0.001 即是 0.1%
    # 打印初始条件
    print('Starting Portfolio Value: %.3f' % cerebro.broker.getvalue())
    # 运行回测
    cerebro.run()
    # 打印输出结果
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    # 绘图
    cerebro.plot(style='bar')