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

import math
import backtrader as bt
import t1_short.t1_short_runner
import pandas as pd
import numpy as np
import talib


class PeriodHQData(bt.feeds.PandasData):
    lines = ('open', 'close', 'high', 'low', 'p_change', 'volume', 'vol_sma5')
    params = (
        ('datetime', None),
        ('open', -1),
        ('high', -1),
        ('close', -1),
        ('low', -1),
        ('p_change', -1),
        ('volume', -1),
        ('vol_sma5', -1),
    )
    datafields = [x[0] for x in params]

    def __init__(self):
        super(PeriodHQData, self).__init__()


def prepare_full_period_df_func(df, stock):
    df['vol_sma5'] = talib.SMA(df['volume'].values, timeperiod=5)
    return df


class Strategy(bt.Strategy):
    def __init__(self, p_chg_pct=9.8):
        super(Strategy, self).__init__()
        self.last_buy_date = None
        #         self.size = None
        self.p_chg_pct = p_chg_pct

    def next_open(self):
        if not self.position:
            if len(self.data0) == self.data0.buflen():
                return
            if math.isnan(self.data0.close[-1]):
                return
            # 当日被砸跌停，收盘价涨回来，但总体依然下跌
            if self.data.low[0] < round(self.data0.close[-1] * 0.903 + 0.0000001, 2) and self.data.p_change[0] > -9.7:
                # and self.data.p_change[0] < 0:
                self.sell(data=self.data0,
                          size=int(self.broker.cash / self.data0.close[0]),
                          price=self.data0.close[0],
                          exectype=bt.Order.StopLimit)
        else:
            self.close(price=self.data0.open[0], exectype=bt.Order.Limit)
