#!/usr/bin/env python
# -*- coding: utf-8 -*-
import cross_order as order
import time


def StochRSI(close, lengthRSI=14, lengthStoch=14, smoothK=3, smoothD=3):
    """
    计算StochRSI(close),有5个参数,第1个为数据源
    @param close: 数据源
    @param lengthRSI: RSI参数
    @param lengthStoch: Stoch参数
    @param smoothK: K参数
    @param smoothD: D参数
    @return:
    """
    # 计算RSI
    lc = close.shift(1)
    diff = close - lc
    up = diff.where(diff > 0, 0)
    down = -diff.where(diff < 0, 0)
    ema_up = up.ewm(alpha=1 / lengthRSI, adjust=False).mean()
    ema_down = down.ewm(alpha=1 / lengthRSI, adjust=False).mean()
    rs = ema_up / ema_down
    rsi = 100 - 100 / (1 + rs)
    # 计算Stochastic
    stoch = (rsi - rsi.rolling(window=lengthStoch).min()) / (rsi.rolling(window=lengthStoch).max() -
                                                             rsi.rolling(window=lengthStoch).min())
    k = stoch.rolling(window=smoothK).mean()
    d = k.rolling(window=smoothD).mean()
    # 添加到data中
    STOCHRSI = k * 100
    MASTOCHRSI = d * 100
    return STOCHRSI, MASTOCHRSI


def main():
    print("任务开始时间:", time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())))
    for symbol in order.symbol_pool:
        # 设置杠杆倍数
        order.set_leverage(symbol=symbol, leverage='25')
        # 15分钟
        df = order.get_candlesticks(symbol=symbol, interval='15m', limit=str(181))
        fastk, fastd = StochRSI(df['close'], 14, 14, 3, 3)
        if (fastk.iloc[-2] or fastd.iloc[-2] < 20) and (
                fastk.iloc[-1] or fastd.iloc[-1] >= 20):
            order.up_cross_order(symbol=symbol, ordtype='market', message='StochRSI 策略做多')
        elif (fastk.iloc[-2] or fastd.iloc[-2] < 80) and (
                fastk.iloc[-1] or fastd.iloc[-1] >= 80):
            order.down_cross_order(symbol=symbol, ordtype='market', message='StochRSI 策略做空')

        time.sleep(2)

    print("任务结束时间:", time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())))


if __name__ == '__main__':
    print("程序运行时间:", time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())))
    main()