# coding:utf-8
from datetime import datetime as dt
import numpy as np
import os
import pandas as pd
import time
from sqlalchemy import create_engine
from jqdatasdk import *
import pymysql
import multiprocessing as mp
import math
import talib as ta
from xtquant import xtdata
import os
import traceback

pd.set_option('display.max_columns', None)  # 设置显示最大行
engine = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/qmt_stocks_whole?charset=utf8')


def err_call_back(err):
    print(f'出错啦~ error：{str(err)}')
    traceback.print_exc()


def myself_kdj(df):
    low_list = df['low_back'].rolling(9, min_periods=9).min()
    low_list.fillna(value=df['low_back'].expanding().min(), inplace=True)
    high_list = df['high_back'].rolling(9, min_periods=9).max()
    high_list.fillna(value=df['high_back'].expanding().max(), inplace=True)
    rsv = (df['close_back'] - low_list) / (high_list - low_list) * 100
    df['k'] = pd.DataFrame(rsv).ewm(com=2).mean()
    df['d'] = df['k'].ewm(com=2).mean()
    df['j'] = 3 * df['k'] - 2 * df['d']
    return df


# macd指标
def get_macd_data(data, short=0, long1=0, mid=0):
    if short == 0:
        short = 12
    if long1 == 0:
        long1 = 26
    if mid == 0:
        mid = 9
    data['sema'] = pd.Series(data['close_back']).ewm(span=short).mean()
    data['lema'] = pd.Series(data['close_back']).ewm(span=long1).mean()
    data.fillna(0, inplace=True)
    data['dif'] = data['sema'] - data['lema']
    data['dea'] = pd.Series(data['dif']).ewm(span=mid).mean()
    data['macd'] = 2 * (data['dif'] - data['dea'])
    data.fillna(0, inplace=True)
    # return data[['dif', 'dea', 'macd']]


# rsi指标
# 建议用talib库的RSI方法，亲测有用
def get_ris(data):
    data["rsi_6"] = ta.RSI(data['close_back'], timeperiod=6)
    data["rsi_12"] = ta.RSI(data['close_back'], timeperiod=12)
    data["rsi_24"] = ta.RSI(data['close_back'], timeperiod=24)


def get_bias(data):
    # 计算方法：
    # bias指标
    # N期BIAS=(当日收盘价-N期平均收盘价)/N期平均收盘价*100%
    data['bias_6'] = (data['close_back'] - data['close_back'].rolling(6, min_periods=1).mean()) / \
                     data['close_back'].rolling(6, min_periods=1).mean() * 100
    data['bias_12'] = (data['close_back'] - data['close_back'].rolling(12, min_periods=1).mean()) / \
                      data['close_back'].rolling(12, min_periods=1).mean() * 100
    data['bias_24'] = (data['close_back'] - data['close_back'].rolling(24, min_periods=1).mean()) / \
                      data['close_back'].rolling(24, min_periods=1).mean() * 100
    data['bias_6'] = round(data['bias_6'], 2)
    data['bias_12'] = round(data['bias_12'], 2)
    data['bias_24'] = round(data['bias_24'], 2)


def get_wilr(data):
    # 威廉指标
    # 建议用talib库的WILLR方法，亲测有用
    data['willr'] = ta.WILLR(data['high_back'], data['low_back'], data['close_back'], timeperiod=14)


def get_hlfx(data):
    Trading_signals = 0
    data_temp = data[['time', 'open_back', 'close_back', 'high_back', 'low_back', 'dif', 'dea', 'macd']]
    data_temp.columns = ['time', 'open', 'close', 'high', 'low', 'dif', 'dea', 'macd']
    df_day = pd.DataFrame(columns=['time', 'open', 'close', 'high', 'low', 'volume', 'money', 'HL'])
    # 先处理去包含
    for i in data_temp.index:
        if i == 0 or i == 1:
            df_day = pd.concat([df_day, data_temp.iloc[[i]]], ignore_index=True)
        # 不包含
        elif (df_day.iloc[-1, 3] > data_temp.loc[i, 'high']
              and df_day.iloc[-1, 4] > data_temp.loc[i, 'low']) \
                or (df_day.iloc[-1, 3] < data_temp.loc[i, 'high']
                    and df_day.iloc[-1, 4] < data_temp.loc[i, 'low']):
            df_day = pd.concat([df_day, data_temp.loc[[i]]], ignore_index=True)
        # 包含
        else:
            # 左高，下降
            if df_day.iloc[-2, 3] > df_day.iloc[-1, 3]:
                df_day.iloc[-1, 3] = min(df_day.iloc[-1, 3], data_temp.loc[i, 'high'])
                df_day.iloc[-1, 4] = min(df_day.iloc[-1, 4], data_temp.loc[i, 'low'])
            else:
                # 右高，上升
                df_day.iloc[-1, 3] = max(df_day.iloc[-1, 3], data_temp.loc[i, 'high'])
                df_day.iloc[-1, 4] = max(df_day.iloc[-1, 4], data_temp.loc[i, 'low'])
    # print('111', df_day, data_temp)

    if len(df_day.index) > 2:
        # 寻找顶底分型
        for x in range(2, len(df_day.index)):
            m = x - 1
            # 底
            if ((df_day.loc[x, 'high'] > df_day.loc[x - 1, 'high']) and (
                    df_day.loc[x - 2, 'high'] > df_day.loc[x - 1, 'high'])):

                df_day.loc[x, 'HL'] = 'L*'
                while m:
                    if df_day.loc[m, 'HL'] in ['H', 'HH', 'H*']:
                        if (x - m) > 3:
                            # 成笔——>L
                            df_day.loc[x, 'HL'] = 'L'
                            # 产生信号，进入hlfx_pool
                            if x == len(df_day.index) - 1:
                                Trading_signals = 1

                    elif df_day.loc[m, 'HL'] == 'L':
                        if df_day.loc[m - 1, 'low'] > df_day.loc[x - 1, 'low']:
                            # 前一个为底更高，且中间不存在更低的底
                            df_day.loc[x, 'HL'] = 'L'

                            # 产生信号，进入hlfx_pool
                            if x == len(df_day.index) - 1:
                                Trading_signals = 1

                            # 获得MACD，判断MACD判断背驰
                            x_macd_dif, x_macd_dea, x_macd_macd = data_temp.loc[x, 'dif'], data_temp.loc[x, 'dea'], data_temp.loc[x, 'macd']
                            m_macd_dif, m_macd_dea, m_macd_macd = data_temp.loc[m, 'dif'], data_temp.loc[m, 'dea'], data_temp.loc[m, 'macd']

                            # MACD底背驰
                            if m_macd_dif < x_macd_dif:
                                # 背驰底->LL
                                df_day.loc[x, 'HL'] = 'LL'
                            break
                        break
                    m = m - 1
                    if m == 0:
                        df_day.loc[x, 'HL'] = 'L'

            # 顶
            elif ((df_day.loc[x, 'high'] < df_day.loc[x - 1, 'high']) and (
                    df_day.loc[x - 2, 'high'] < df_day.loc[x - 1, 'high'])):

                df_day.loc[x, 'HL'] = 'H*'
                while m:
                    if df_day.loc[m, 'HL'] in ['L', 'LL', 'L*']:
                        if x - m > 3:
                            # 成笔->H
                            df_day.loc[x, 'HL'] = 'H'
                            # 产生信号，进入hlfx_pool
                            if x == len(df_day.index) - 1:
                                Trading_signals = 2

                    elif (df_day.loc[m, 'HL'] == 'H'):
                        if df_day.loc[x - 1, 'high'] > df_day.loc[m - 1, 'high']:
                            # 前一个为顶，且中间存在不包含 or 更高的顶
                            df_day.loc[x, 'HL'] = 'H'
                            # 产生信号，进入hlfx_pool
                            if x == len(df_day.index) - 1:
                                Trading_signals = 2

                            # 获得MACD，判断MACD判断背驰
                            x_macd_dif, x_macd_dea, x_macd_macd = data_temp.loc[x, 'dif'], data_temp.loc[x, 'dea'], data_temp.loc[x, 'macd']
                            m_macd_dif, m_macd_dea, m_macd_macd = data_temp.loc[m, 'dif'], data_temp.loc[m, 'dea'], data_temp.loc[m, 'macd']

                            # MACD顶背驰
                            if x_macd_dif < m_macd_dif:
                                df_day.loc[x, 'HL'] = 'HH'
                            break
                        break
                    m = m - 1
                    if m == 0:
                        df_day.loc[x, 'HL'] = 'H'

            else:
                df_day.loc[x, 'HL'] = '-'
    df_temp = df_day[['time', 'HL']]

    return df_temp, Trading_signals


def tech_anal(stocks, hlfx_pool,err_list):
    print(f'{dt.now()}开始循环计算！ MyPid is {os.getpid()},池子长度为{len(stocks)}')

    engine_tech = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/qmt_stocks_tech?charset=utf8')
    m = 0

    for stock in stocks:
        # print(stock)
        try:
            df = pd.read_sql_table('%s_1d' % stock, con=engine)
            df.dropna(axis=0, how='any')
        except BaseException:
            print(f'{stock}读取有问题')
            traceback.print_exc()
            pass
        else:
            if len(df) != 0:
                try:
                    get_macd_data(df)
                    get_ris(df)
                    get_bias(df)
                    get_wilr(df)
                    df_temp, T_signals = get_hlfx(df)
                    df = pd.merge(df, df_temp, on='time', how='left')
                    df['HL'].fillna(value='-', inplace=True)
                    df = df.reset_index(drop=True)
                    # print(stock, '\n', df[['open_front', 'HL']])
                    df = df.replace([np.inf, -np.inf], np.nan)
                    df.to_sql('%s_1d' % stock, con=engine_tech, index=False, if_exists='replace')
                # with engine.connect() as con:
                #     con.execute("ALTER TABLE `%s_1d` ADD PRIMARY KEY (`time`);" % stock)
                except BaseException:
                    print(f'{stock}存储有问题')
                    traceback.print_exc()
                    err_list.append(stock)
                    pass
                else:
                    # print(f"{stock} 成功！")
                    m += 1
            else:
                err_list.append(stock)
                print(f'{stock}数据为空')

            if stock in hlfx_pool and T_signals == 2:
                hlfx_pool.remove(stock)
            elif stock not in hlfx_pool and T_signals == 1:
                hlfx_pool.append(stock)


    print(f'Pid：{os.getpid()}已经完工了，应处理{len(stocks)},共计算{m}支个股')


if __name__ == '__main__':
    sttime = dt.now()

    stocks = xtdata.get_stock_list_in_sector('沪深A股')
    print(len(stocks))
    stocks.sort()

    err_list = mp.Manager().list()
    fre = '1d'
    engine_hlfx_pool = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/hlfx_pool?charset=utf8')
    hlfx_pool = mp.Manager().list()
    hlfx_pool.extend(pd.read_sql_query(
        'select value from `%s`' % fre, engine_hlfx_pool).iloc[-1, 0].split(","))

    pool = mp.Pool(processes=mp.cpu_count())
    step = math.ceil(len(stocks) / mp.cpu_count())
    # step = 10000
    x = 1
    # tech_anal(stocks, hlfx_pool)
    for i in range(0, len(stocks), step):
        print(x)
        pool.apply_async(func=tech_anal, args=(stocks[i:i + step], hlfx_pool, err_list,), error_callback=err_call_back)
        x += 1
    time.sleep(5)
    pool.close()
    pool.join()

    print(hlfx_pool)
    print(len(err_list), err_list)

    # 存档入库
    db_pool = pymysql.connect(host='localhost',
                              user='root',
                              port=3307,
                              password='r6kEwqWU9!v3',
                              database='hlfx_pool')
    cursor_pool = db_pool.cursor()
    results_list = ','.join(set(hlfx_pool))

    sql = "INSERT INTO %s (date,value) VALUES('%s','%s')" % (fre, dt.now().strftime('%Y-%m-%d %H:%M:%S'), results_list)
    cursor_pool.execute(sql)
    db_pool.commit()
    edtime = dt.now()


    print(edtime - sttime)