# -*- coding: utf-8 -*-

import backtrader
from backtrader.utils.date import num2date
import pandas as pd


class BaseStrategy(backtrader.Strategy):
    def __init__(self, *args, **kwargs):
        super(BaseStrategy, self).__init__(*args, **kwargs)
        self.order_history = []  # 交易下单记录
        self.trade_history = []  # 平仓历史记录
        self.order_history_df = None
        self.trade_history_df = None

    def log(self, txt, dt=None):
        ''' Logging function fot this strategy'''
        dt = dt or self.data.datetime.date(0)
        # print('%s  %s' % (dt.isoformat(), txt))

    def notify_order(self, order):
        # 打印Accepted
        # if order.status in [order.Accepted]:
        #     # print order
        #     if order.isbuy():
        #         self.log('BUY  ACCEPTED, %.02f * (%d)' % (order.price, order.size))
        #     elif order.issell():
        #         self.log('SELL ACCEPTED, %.02f * (%d)' % (order.price, order.size))
        #     return
        # 打印Completed
        if order.status in [order.Completed]:
            order_info = {
                'datetime': num2date(order.executed.dt),
                'price': order.executed.price,
                'size': order.executed.size,
                'comm': order.executed.comm,
            }
            if order.isbuy():
                self.log('BUY  EXECUTED, %(price).2f * %(size)d' % order_info)
            elif order.issell():
                self.log('SELL EXECUTED, %(price).2f * %(size)d' % order_info)
            self.order_history.append(order_info)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        trade_info = {
            'datetime': trade.close_datetime(),
            'open_datetime': trade.open_datetime(),
            'gross': trade.pnl,
            'net': trade.pnlcomm,
            'comm': trade.commission,
            'cash': self.broker.cash,
        }
        self.log('OPERATION PROFIT, GROSS %(gross).2f, NET %(net).2f, COMM %(comm).2f'
                 % trade_info)

        self.trade_history.append(trade_info)

    def stop(self):
        self.trade_history_df = pd.DataFrame(
            self.trade_history,
            index=pd.DatetimeIndex([trade['datetime'] for trade in self.trade_history]),
            columns=['open_datetime', 'gross', 'net', 'comm', 'cash'],
        )
        self.trade_history_df.index.name = 'datetime'
        self.order_history_df = pd.DataFrame(
            self.order_history,
            index=pd.DatetimeIndex([order['datetime'] for order in self.order_history]),
            columns=['price', 'size', 'comm'])
        self.order_history_df.index.name = 'datetime'
