from jqdatasdk import *
from datetime import datetime as dt
import pandas as pd
from sqlalchemy import create_engine
import numpy as np
from jqdatasdk.technical_analysis import *

def calculateEMA(period, closeArray, ema, emaArray=[]):
    """计算指数移动平均"""
    length = len(closeArray)
    nanCounter = np.count_nonzero(np.isnan(closeArray))
    if not emaArray:
        if ema and (ema !=0):
            firstema = ema
            emaArray.append(firstema)
        else:
            print('走这里了')
            emaArray.extend(np.tile([np.nan], (nanCounter + period - 1)))
            firstema = np.mean(closeArray[nanCounter:nanCounter + period - 1])
            emaArray.append(firstema)
        for i in range(nanCounter+period, length):
            ema_a = (2 * closeArray[i] + (period - 1) * emaArray[-1]) / (period + 1)
            emaArray.append(ema_a)
    return np.array(emaArray)


def calculateMACD(closeArray, ema, shortPeriod=12, longPeriod=26, signalPeriod=9):
    ema12 = calculateEMA(shortPeriod, closeArray, ema, [])
    ema26 = calculateEMA(longPeriod, closeArray, ema, [])
    diff = ema12 - ema26
    dea = calculateEMA(signalPeriod, diff, 0, [])
    macd = 2 * (diff - dea)
    return macd, diff, dea

# auth('18616891214', 'Ea?*7f68nD.dafcW34d!')
auth('18521506014', 'Abc123!@#')
stocks = list(get_all_securities(['stock'], date=dt.today().strftime('%Y-%m-%d')).index)
engine = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/stocks?charset=utf8')
engine_data = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/stocks_data?charset=utf8')

fre = '1d'
print('ready to write to mysql %s' % fre)
for stock in stocks[2500:2501]:
    print(stock, fre)
    starttime ='2022-01-01'
    # endtime = pd.read_sql_table('stk%s_%s' % (stock, fre), con=engine).iloc[-1, 1]
    df_stock2 = get_price(stock, start_date=starttime, end_date=dt.today().strftime('%Y-%m-%d %H:%M:%S'),
                         frequency=fre, fields=['open', 'close', 'high', 'low', 'volume', 'money'],
                         skip_paused=False,
                         fq='pre', count=None, panel=False)

    df_stock = pd.read_sql_query('select date,open,close,high,low,volume,money from `stk%s_%s`' % (stock, fre), engine)
    # df_stock.index.name = 'date'
    df_money = get_money_flow(stock, start_date=starttime, end_date=dt.today().strftime('%Y-%m-%d %H:%M:%S'),
                              fields=None, count=None)
    df_money = df_money.drop(columns=['sec_code'])
    df_stock = pd.merge(df_stock, df_money, how='outer', left_index=False , on='date')
    # df_stock.to_csv('/Users/daniel/Downloads/Result.csv')
    df_stock = df_stock.dropna(axis=0)
    df_stock2=df_stock2.dropna(axis=0)
    df_stock2.reset_index(inplace=True)
    df_stock2.rename(columns={'index': 'date'}, inplace=True)
    print(df_stock2)
    df_close = df_stock2['close']

    if starttime != df_stock2.loc[0, 'date'].strftime('%Y-%m-%d'):
        ema = 0
    else:
        ema = EMA(stock, check_date=starttime, timeperiod=30)[stock]
    df_macd = calculateMACD(df_close, ema)

    df_stock = pd.concat([df_stock2, pd.Series(df_macd[0]).rename('macd'), pd.Series(df_macd[1]).rename('diff'), pd.Series(df_macd[2]).rename('dea')], axis=1)

    x_macd_dif, x_macd_dea, x_macd_macd = MACD(stock, check_date=dt.today().strftime('%Y-%m-%d %H:%M:%S'), SHORT=12, LONG=26, MID=9,
                                               unit=fre)
    print(x_macd_macd, x_macd_dif, x_macd_dea)
    print(df_stock)
    # df_stock.to_sql('stk%s_%s' % (stock, fre), con=engine_data, index=True, if_exists='append')
    # with engine.connect() as con:
    #     con.execute("ALTER TABLE `stk%s_%s` ADD PRIMARY KEY (`date`);" % (stock, fre))
    # print(df_stock)