#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# File: data_download.py
# Date: 2024/7/4
# Author: 抖音、快手、视频号：东哥策略交易，微信：dongtrader
# Description: 东哥量化，带你走上量化之路。

import os
from MyTT import *

'''
# 重要的放开头，实盘交易策略管理在stock_strategy.csv文件，用金山或excel打开，像excel表。
# 里面有state、amount_std、strategy、success、earnings、trades这几个内容。
# state设置为0即正常买卖，设置为1只卖出不买入，设置为2时清仓；
# amount_std为每一个股票可以买入金额，请根据实际情况填写，程序将根据实际股票价格，在设置金额的80%~120%的范围买入。
# strategy设置为交易的策略，‘aa_trade’，交易将按策略信号买卖，请确保策略名称与数据策略名称一致。
# success为交易的策略胜率，大于设置的胜率，则认为策略胜率足够，可以进行交易。
# earnings为交易的策略平均盈利率，大于设置的盈利率，则认为策略盈利率足够，可以进行交易。
# trades为交易的策略的交易次数，大于设置的次数，则认为交易次数足够，策略可信，可以进行交易。
'''

# 选择获取数据的市场，'sh'沪市A股包含科创板股票，'sz' 深市A股包含创业板股票，'bj'北交所。
# 如果想单独下载创业板（'cy'）或科创板（'kc'）。如果想要沪市A股但不想要科创板，则设置为['sh','-kc']。
markets =['sh', 'sz']

# 是否删除ST股票，删除为True，保留为False。
del_st_stock = True

# 是否删除已经退市股票，删除为True，保留为False。
del_stop_stock = True

# 数据下载接口设置，目前只支持AKShare，如果要使用其他接口，请自行修改。
api = 'ak'

# 历史数据周期设置，{'daily', 'weekly', 'monthly', '1', '5', '15', '30', '60'}分别对应日线、周线、月线、1分钟、5分钟、15分钟、30分钟、60分钟。
period = 'daily'

# 可用日线数量约束,上市少于多少天的不下载数据，设置为0则只要上市就下载。250约等1年。
g_available_days_limit = 250

# 历史数据复权格式，''：空或默认不复权; qfq: 前复权; hfq: 后复权
adjustflag = 'hfq'

# 任务执行多线程数
process_num = 15

# 数据保存的文件夹路径，如：'E:/feather'，注意是使用反斜杆‘/’，而不是正斜杆‘\’，如果将'E:/'改为'./'将保存到当前程序的目录下。
# 数据表的文件夹路径
stockpath = 'E:/feather/stock'
# 交易数据文件夹路径
tradepath = 'E:/feather/trade'
for path in [stockpath, tradepath]:
    if not os.path.exists(path):
        # 如果文件夹不存在，创建它
        os.makedirs(path)


'''
设置下个交易日开盘价、最高价、最低价、收盘价的涨幅比例，用于计算是否存在买入信号的可能。
如果可能则加入关注列表，从而在交易时只对关注列表进行观测，加快运算速度。
科创板或创业板的股票，涨幅设置-20~20，其它为-10~10。
如果确定影响策略关键的指标，即可对关键指标的数值进行设置，其它设置为0。
另外，要注意的是，最高价不能低于最低价，开盘价与收盘价要在最高价和最低价之间。
'''
# 沪市A股、深圳A股
nextday_open_10 = 0
nextday_high_10 = 0
nextday_low_10 = 0
nextday_close_10 = 10
# 以下为科创板或创业板设置
nextday_open_20 = 0
nextday_high_20 = 0
nextday_low_20 = 0
nextday_close_20 = 20



'''
下面为每个策略都具有的10个因子字段，分别为，
1.b2smarket多空市场状态，2.trade买卖因子，3.holddate策略买入后持有天数，4.buyprice策略买入当天收盘价格，5.profit策略卖出后的利润百分比，
6.profits策略历史累计利润率，7.earnings策略累平均收益，8.trades策略交易次数，9.win策略盈利次数，10.success策略胜率。
每增加计算一种策略，则需要增加一个策略的因子字段。采用2个大写字母，如aa_trade、bb_trade等对策略进行区分。
'''
# 计算的用于交易的因子列表,策略名称必须与factor_columns中的策略名称一致。
factorlist = ['ek', ]

# 针对固定盈利的因子,填写固定盈利百分比，采用字典形式，如：profit_dict ={'ek': 0.04, 'aa': 0.05}。
profit_dict ={'ek': 0.04,}


# 以下为回测设置
# 设置MIN_PERIOD用于后面剔除数据长度不足的股票
MIN_PERIOD = 60
# 设置回测启动资金
capital = 100000.0




def basic_factors(df, oldindex):
    """
    计算股票的基本因素，可以在这里计算需要的基础因子/指标、以及需要被其它因子所调用的公共因子。
    如：MACD指标、各种均线等。

    参数:
    df: pandas.DataFrame，包含股票收盘价等数据。
    oldindex: list，原股票列表的索引。

    返回:
    pandas.DataFrame，包含计算后的技术指标数据。
    """
    # 下面是计算MACD指标的示例代码，可以根据自己的需求进行修改。
    # 获取收盘价数组
    try:
        CLOSE = df['close'].values
        dif, dea, macd = np.nan_to_num(np.round(MACD(CLOSE, SHORT=12, LONG=26, M=9), 4))
        df['dif'] = dif
        df['dea'] = dea
        df['macd'] = macd
        df['ma_5'] = np.nan_to_num(np.round(MA(CLOSE, 5), 4))
        df['ma_10'] = np.nan_to_num(np.round(MA(CLOSE, 10), 4))
        df['ema_10'] = np.nan_to_num(np.round(EMA(CLOSE, 10), 4))
        df['ma_30'] = np.nan_to_num(np.round(MA(CLOSE, 30), 4))
        df['ema_30'] = np.nan_to_num(np.round(EMA(CLOSE, 30), 4))
        df['ma_100'] = np.nan_to_num(np.round(MA(CLOSE, 100), 4))

        for i in range(oldindex, len(df)):
            if i == 0:
                df.loc[i, ['trend', 'transition', 'change', 'demand', 'supply', 'percent_10', 'b']] = 0.0
                df.loc[i, ['snake', 'snake_10']] = -100

                if df['high'][i] == df['low'][i]:
                    direction = 0
                    unit = df['volume'][i]
                else:
                    direction = df['close'][i] - (df['high'][i] + df['low'][i]) / 2
                    unit = df['volume'][i] / (df['high'][i] - df['low'][i])
                df.at[i, 'SD_strength'] = round(direction * unit, 0)
                if df['SD_strength'][i] > 0:
                    df.at[i, 'demand'] = df['SD_strength'][i]
                    df.at[i, 'supply'] = 0
                elif df['SD_strength'][i] < 0:
                    df.at[i, 'demand'] = 0
                    df.at[i, 'supply'] = df['SD_strength'][i]
                else:
                    df.at[i, 'demand'] = 0
                    df.at[i, 'supply'] = 0

                if df['open'][i] == 0:
                    df.at[i, 'long'] = 0.0
                else:
                    if df['open'][i] != 0:
                        pow = round((df['close'][i] / df['open'][i] - 1) * 100, 4)
                    else:
                        pow = 0.0
                    df.at[i, 'long'] = pow

            else:

                if df['close'][i] == df['high'][i] \
                        and df['close'][i - 1] != 0 and df['close'][i] / df['close'][i - 1] > 1.099:
                    df.at[i, 'percent_10'] = 1
                else:
                    df.at[i, 'percent_10'] = 0

                if df['high'][i] == df['low'][i]:
                    if df['close'][i] == df['close'][i - 1]:
                        direction = 0
                    elif df['close'][i] > df['close'][i - 1]:
                        direction = 1
                    else:
                        direction = -1
                    unit = df['volume'][i]
                else:
                    direction = df['close'][i] - (df['high'][i] + df['low'][i]) / 2
                    unit = df['volume'][i] / (df['high'][i] - df['low'][i])
                df.at[i, 'SD_strength'] = round(direction * unit, 0)
                if df['SD_strength'][i] > 0 and df['SD_strength'][i - 1] > 0:
                    if df['demand'][i - 1] > 0:
                        df.at[i, 'demand'] = df['SD_strength'][i] + df['demand'][i - 1]
                        df.at[i, 'supply'] = 0
                    else:
                        if df['demand'][i - 1] > 0:
                            df.at[i, 'demand'] = df['SD_strength'][i] + df['demand'][i - 1]
                        else:
                            df.at[i, 'demand'] = df['SD_strength'][i]
                        df.at[i, 'supply'] = 0
                elif df['SD_strength'][i] < 0 and df['SD_strength'][i - 1] < 0:
                    if df['supply'][i - 1] > 0:
                        df.at[i, 'demand'] = 0
                        df.at[i, 'supply'] = df['SD_strength'][i] + df['supply'][i - 1]
                    else:
                        df.at[i, 'demand'] = 0
                        if df['supply'][i - 1] < 0:
                            df.at[i, 'supply'] = df['SD_strength'][i] + df['supply'][i - 1]
                        else:
                            df.at[i, 'supply'] = df['SD_strength'][i]
                elif df['SD_strength'][i] > 0 and df['SD_strength'][i - 1] <= 0:
                    if df['SD_strength'][i] > abs(df['SD_strength'][i - 1]):
                        if df['demand'][i - 1] > 0:
                            df.at[i, 'demand'] = df['SD_strength'][i] + df['demand'][i - 1]
                        else:
                            df.at[i, 'demand'] = df['SD_strength'][i]
                        df.at[i, 'supply'] = 0
                    else:
                        if df['supply'][i - 1] < 0:
                            df.at[i, 'demand'] = 0
                            df.at[i, 'supply'] = df['supply'][i - 1] + df['SD_strength'][i]
                        else:
                            if df['demand'][i - 1] > 0:
                                df.at[i, 'demand'] = df['SD_strength'][i] + df['demand'][i - 1]
                            else:
                                df.at[i, 'demand'] = df['SD_strength'][i]
                            df.at[i, 'supply'] = 0
                elif df['SD_strength'][i] < 0 and df['SD_strength'][i - 1] >= 0:
                    if abs(df['SD_strength'][i]) < df['SD_strength'][i - 1]:
                        if df['demand'][i - 1] > 0:
                            df.at[i, 'demand'] = df['SD_strength'][i] + df['demand'][i - 1]
                            df.at[i, 'supply'] = 0
                        else:
                            df.at[i, 'demand'] = 0
                            if df['supply'][i - 1] < 0:
                                df.at[i, 'supply'] = df['SD_strength'][i] + df['supply'][i - 1]
                            else:
                                df.at[i, 'supply'] = df['SD_strength'][i]
                    else:
                        df.at[i, 'demand'] = 0
                        if df['supply'][i - 1] < 0:
                            df.at[i, 'supply'] = df['SD_strength'][i] + df['supply'][i - 1]
                        else:
                            df.at[i, 'supply'] = df['SD_strength'][i]
                elif df['SD_strength'][i] == 0:
                    df.at[i, 'demand'] = df['demand'][i - 1]
                    df.at[i, 'supply'] = df['supply'][i - 1]
                else:
                    df.at[i, 'demand'] = 0
                    df.at[i, 'supply'] = 0

                if df['open'][i] == 0 or df['close'][i - 1] == 0:
                    df.at[i, 'long'] = 0.0
                else:
                    if df['open'][i] != 0 and df['close'][i - 1] != 0:
                        pow = round((df['open'][i] / df['close'][i - 1] + df['close'][i] / df['open'][i] - 2) / 2 * 100,
                                    4)
                    else:
                        pow = 0.0
                    df.at[i, 'long'] = pow
                if i < 11:
                    df.at[i, 'be_high'] = 0.0
                elif i < 21:
                    if df['high'][i - 11] == max(df['high'][0:i]):
                        df.at[i, 'be_high'] = df['high'][i - 11]
                    else:
                        df.at[i, 'be_high'] = df['be_high'][i - 1]
                else:
                    if df['high'][i - 11] == max(df['high'][i - 21:i]):
                        df.at[i, 'be_high'] = df['high'][i - 11]
                    else:
                        df.at[i, 'be_high'] = df['be_high'][i - 1]
                if df['close'][i] != 0:
                    df.at[i, 'snake'] = np.round((df['ema_30'][i] - df['close'][i]) / df['close'][i] * 100, 4)
                    df.at[i, 'snake_10'] = np.round((df['ema_10'][i] - df['close'][i]) / df['close'][i] * 100, 4)
                    df.at[i, 'b'] = np.round((1 + df['long'][i] / 100) * df['close'][i], 4)
                else:
                    df.at[i, 'snake'] = -100
                    df.at[i, 'snake_10'] = -100
                    df.at[i, 'b'] = 0

        df['energy_f'] = fast = np.nan_to_num(np.round(EMA(df['long'], 13) * 2, 4))
        df['energy_s'] = slow = np.nan_to_num(np.round(EMA(df['long'], 34) * 2, 4))
        df['energy'] = np.nan_to_num(np.round(fast - slow, 4))

        df['long_ma_3'] = np.nan_to_num(np.round(EMA(df['long'], 3), 4))
        df['long_ma_6'] = np.nan_to_num(np.round(EMA(df['long'], 6), 4))
        df['long_ma_9'] = np.nan_to_num(np.round(EMA(df['long'], 9), 4))
        df['long_ma_10'] = np.nan_to_num(np.round(EMA(df['long'], 10), 4))
        df['long_ma_30'] = np.nan_to_num(np.round(EMA(df['long'], 30), 4))

        df['snake_ma_10'] = np.nan_to_num(np.round(EMA(df['snake'], 10), 4))
        df['snake_ma_30'] = np.nan_to_num(np.round(EMA(df['snake'], 30), 4))
        df['volume_ma_10'] = np.nan_to_num(np.round(EMA(df['volume'], 10), 4))
        df['b_ma_10'] = np.nan_to_num(np.round(MA(df['b'], 10), 4))
        longhigh_big = BARSLAST(df['long'] > 0)
        longhigh_small = BARSLAST(df['long'] < 0)
        condition_big = (df['long'] > 0) & (REF(df['long'], 1) < 0)
        condition_small = (df['long'] < 0) & (REF(df['long'], 1) > 0)
        longvalue_big = BARSLAST(condition_big)
        longvalue_small = BARSLAST(condition_small)

        for i in range(oldindex, len(df)):

            if i == 0:
                df.loc[i, ['longvalue', 'longhigh']] = round(df['long'][i], 4)
                df.loc[i, ['touch', 'mixed', 'shadow', 'link', 'sumtouch', 'malink', 'ma_posture']] = 0.0

            else:
                if df['close'][i] != 0:
                    df.at[i, 'ma_posture'] = np.round((df['b_ma_10'][i] - df['ma_10'][i]) / df['close'][i] * 100, 4)
                else:
                    df.at[i, 'ma_posture'] = 0

                if df['long'][i] > 0:
                    df.at[i, 'longhigh'] = round(sum(df['long'][i - longhigh_small[i]:i]) + df['long'][i], 4)
                    df.at[i, 'longvalue'] = round(sum(df['long'][i - longvalue_small[i]:i]) + df['long'][i], 4)
                else:
                    df.at[i, 'longhigh'] = round(sum(df['long'][i - longhigh_big[i]:i]) + df['long'][i], 4)
                    df.at[i, 'longvalue'] = round(sum(df['long'][i - longvalue_big[i]:i]) + df['long'][i], 4)

                try:
                    if df['high'] > df['ma_5'] and df['high'] > df['ma_10'] \
                            and df['low'] < df['ma_5'] and df['low'] < df['ma_10']:
                        df.at[i, 'touch'] = 1
                    else:
                        df.at[i, 'touch'] = 0
                except:
                    df.at[i, 'touch'] = 0

                try:
                    df.at[i, 'mixed'] = np.nan_to_num(np.round(df['long_ma_30'][i] - df['snake_ma_30'][i], 4))
                except:
                    df.at[i, 'mixed'] = 0

                if (df['macd'][i] > 0 and df['mixed'][i] > 0) or df['mixed'][i] < 0 or (
                        df['macd'][i] < 0 and df['dif'][i] > df['dif'][i - 1]):
                    df.at[i, 'shadow'] = round(df['mixed'][i], 4)
                else:
                    df.at[i, 'shadow'] = round(df['mixed'][i - 1], 4)

                df.at[i, 'link'] = np.nan_to_num(
                    np.round(df['longvalue'][i] - df['longvalue'][i - 1] + df['mixed'][i], 4))
                if i < 3:
                    df.at[i, 'malink'] = 0
                    df.at[i, 'sumtouch'] = sum(df['touch'][:i]) + df['touch'][i]
                elif i < 10:
                    df.at[i, 'sumtouch'] = sum(df['touch'][:i]) + df['touch'][i]
                    if df['ma_100'][i] > df['ma_100'][i - 1] > df['ma_100'][i - 2]:
                        df.at[i, 'malink'] = 1
                    else:
                        df.at[i, 'malink'] = 0
                else:
                    df.at[i, 'sumtouch'] = sum(df['touch'][i - 9:i]) + df['touch'][i]
                    if df['ma_100'][i] > df['ma_100'][i - 1] > df['ma_100'][i - 2]:
                        df.at[i, 'malink'] = 1
                    else:
                        df.at[i, 'malink'] = 0

        df['longvalue_ma_10'] = np.nan_to_num(np.round(EMA(df['longvalue'], 10), 4))
        df['link_ma_10'] = np.nan_to_num(np.round(MA(df['link'], 10), 4))
        df['longhigh_ma_25'] = np.nan_to_num(np.round(MA(df['longhigh'], 25), 4))

        for i in range(oldindex, len(df)):
            if i == 0:
                df.at[i, 'money'] = 0
            else:
                df.at[i, 'money'] = round(df['longvalue_ma_10'][i] - df['longvalue_ma_10'][i - 1] + df['mixed'][i], 4)

        df['money_ma_5'] = np.nan_to_num(np.round(MA(df['money'], 5), 4))
        df['money_ma_30'] = np.nan_to_num(np.round(MA(df['money'], 30), 4))

        condition_1 = REF(df['money_ma_30'], 2) < df['money_ma_30']
        condition_2 = REF(df['money_ma_30'], 1) < df['money_ma_30']
        condition_3 = df['money_ma_30'] < df['money_ma_5']
        condition_4 = df['money_ma_5'] > 0
        condition_all = condition_1 & condition_2 & condition_3 & condition_4
        df['money_up'] = IF(condition_all, 1, 0)

    except Exception as e:
        print(e)


    return df



def ek_buy2sell(df, i, profit=0.0):
    """
    根据技术指标判断买卖信号。

    参数:
    df: DataFrame, 包含股票价格及相关技术指标的数据框。
    i: int, 当前评估的索引位置。

    返回:
    int, 买卖信号:
        -1: 卖出信号,
         1: 买入信号,
         0: 无买卖信号。
    """

    if df['EN_trader'][i] == -1 \
            or (df['longhigh'][i - 1] - df['longhigh'][i]) > 10 \
            or (df['BH_buyprice'][i - 1] > 0 and df['high'][i] > df['BH_buyprice'][i - 1] * (1 + float(profit) / 100)) \
            or df['BH_holddate'][i - 1] >= 7:
        buy2sell = -1

    elif df['money_up'][i] == 1 and df['close'][i - 1] != 0 and df['close'][i] / df['close'][i - 1] < 1.0975 and \
            max(df['close'][i - 11: i - 1]) < df['be_high'][i] and df['low'][i] < df['be_high'][i - 1] and \
            df['volume'][i] * 1.2 > df['volume'][i - 1] > df['volume'][i] * 0.5:
        if df['close'][i] > df['be_high'][i] > df['close'][i - 1]:
            buy2sell = 1
        elif df['be_high'][i] == df['close'][i - 1] and df['close'][i] > df['be_high'][i] > df['open'][i - 1]:
            buy2sell = 1
        else:
            buy2sell = 0
    else:
        buy2sell = 0

    return buy2sell