# -*- 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
import utils


class PeriodHQData(bt.feeds.PandasData):
    lines = ('open', 'close', 'high', 'low', 'p_change', 'volume', 'vol_sma10', 'vol_sma5')
    params = (
        ('datetime', None),
        ('open', -1),
        ('high', -1),
        ('close', -1),
        ('low', -1),
        ('p_change', -1),
        ('volume', -1),
        ('vol_sma10', -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):
    df['vol_sma10'] = talib.SMA(df['volume'].values, timeperiod=10)
    df['vol_sma5'] = talib.SMA(df['volume'].values, timeperiod=5)
    return df


class Strategy(bt.Strategy):
    def __init__(self, ratio=2.5):
        super(Strategy, self).__init__()
        self.ratio = ratio

    def short_v1(self):
        # 收盘卖出，隔日买入策略，卖出条件：
        # 1、交易量是过去平均周期平均值 expand_ratio 倍数以上
        # 2、涨幅为正
        # 3、不达到涨停价格
        # 4、T+1 限制
        if len(self.data0) == self.data0.buflen():
            return
        if math.isnan(self.data0.vol_sma10[-1]):
            return
        if self.data0.volume[0] >= self.data0.vol_sma10[-1] * self.ratio and \
                        self.data0.p_change[0] > 0 and \
                        self.data0.close[0] < utils.calc_limit_up_price(self.data0.close[-1]):
            return self.sell(data=self.data0,
                             size=int(self.broker.cash / self.data0.close[0]),
                             price=self.data0.close[0],
                             exectype=bt.Order.StopLimit)

    def short_v2(self):
        # 收盘卖出，隔日买入策略，卖出条件：
        # 1、昨日量比 > 1.5，且昨日涨幅> 5%
        # 2、今日成交量>昨日成交量 * 90%
        # 3、今日涨幅 < 昨日 * 50%，且 没有跌停
        if len(self.data0) == self.data0.buflen():
            return
        if math.isnan(self.data0.vol_sma5[-2]):
            return
        if (self.data0.volume[-1] / self.data0.vol_sma5[-2]) > 1.5 and \
                        self.data0.p_change[-1] > 5 and \
                        self.data0.volume[0] > self.data0.volume[-1] * 0.9 and \
                                -10 < self.data0.p_change[0] < self.data0.p_change[-1] * 0.5:
            return self.sell(data=self.data0,
                             size=int(self.broker.cash / self.data0.close[0]),
                             price=self.data0.close[0],
                             exectype=bt.Order.StopLimit)

    def short_v3(self):
        if not self.position:
            if len(self.data0) == self.data0.buflen():
                return
            if math.isnan(self.data0.close[-1]):
                return
            # 当日大涨后被砸盘
            if self.data0.high[0] >= round(self.data0.close[-1] * 1.098 + 0.0000001, 2) and \
                            self.data0.close[0] < self.data0.high[0] and \
                            self.data0.close[0] > round(self.data0.close[-1] * 0.902 + 0.0000001, 2):
                return self.sell(data=self.data0,
                                 size=int(self.broker.cash / self.data0.close[0]),
                                 price=self.data0.close[0],
                                 exectype=bt.Order.StopLimit)

    def next_open(self):
        if not self.position:
            if self.short_v1() is None:
                if self.short_v2() is None:
                    self.short_v3()
        else:
            self.close(price=self.data0.open[0], exectype=bt.Order.Limit)
