from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

# 中文文档 http://backtrader.com.cn/docu/#1
# 英文文档 https://www.backtrader.com/docu/

import datetime  # 日期时间类
import os.path  # 路径管理
import sys  # 获取当前运行脚本的路径

# 导入backtrader框架
import backtrader as bt

import pandas as pd

# 创建自己的策略
class MyStrategy(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):
        # 保存收盘价的引用
        self.dataclose = self.datas[0].close

    # 打印出行情数据
    def next(self):
        # 记录收盘价
        self.log('Close, %.2f' % self.dataclose[0])

        # 添加买入逻辑 如果k线收盘出现三连跌 则买入
        # 今天的收盘价 < 昨天收盘价
        if self.dataclose[0] < self.dataclose[-1]:
            # 昨天收盘价 < 前天收盘价
            if self.dataclose[-1] < self.dataclose[-2]:
                # 买入
                self.log('BUY CREATE, %.2f' % self.dataclose[0])
                self.buy()

if __name__ == '__main__':
    # 创建cerebro引擎
    cerebro = bt.Cerebro()
    cerebro.broker.set_cash(1000000.00)  # 设置初始资金金额
    cerebro = bt.Cerebro(stdstats=False)
    # 账户信息
    cerebro.addobserver(bt.observers.Broker)
    # 买卖点
    cerebro.addobserver(bt.observers.BuySell)
    # 计算回撤
    cerebro.addobserver(bt.observers.DrawDown)
    # 资金曲线
    cerebro.addobserver(bt.observers.Value)
    # 期间收益
    cerebro.addobserver(bt.observers.TimeReturn)

    # cerebro引擎在后台创建broker(经纪人) 系统默认资金量为10000
    init_fund = cerebro.broker.getvalue()
    print(f'初始资金:{init_fund}')
    print('组合期初资金: %.2f' % cerebro.broker.getvalue())
    filename = 'sz000002.csv'
    print(filename)
    data = pd.read_csv(filename, index_col="date", parse_dates=True)
    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}')
    cerebro.plot(style="candle")