from jqdatasdk import *
from sqlalchemy import create_engine
from jqdatasdk.technical_analysis import *


auth('18616891214', 'Ea?*7f68nD.dafcW34d!')

import numpy as np
import pandas as pd


def calculateEMA(period, closeArray, ema, emaArray=[]):
    """计算指数移动平均"""
    length = len(closeArray)
    nanCounter = np.count_nonzero(np.isnan(closeArray))
    print('def ema', ema)
    if not emaArray:
        if not ema and ema:
            firstema = ema
        else:
            firstema = np.mean(closeArray[nanCounter:nanCounter + period - 1])
        # emaArray.extend(np.tile([np.nan], (nanCounter + period - 1)))
        print('fe', firstema)
        emaArray.append(firstema)

        for i in range(nanCounter, length):
            ema_a = (2 * closeArray[i] + (period - 1) * emaArray[-1]) / (period + 1)
            emaArray.append(ema_a)
    return np.array(emaArray[1:])


def calculateMACD(closeArray,ema, shortPeriod=12, longPeriod=26, signalPeriod=9):
    ema12 = calculateEMA(shortPeriod, closeArray, ema, [])
    print(ema12)
    ema26 = calculateEMA(longPeriod, closeArray, ema, [])
    print(ema26)
    diff = ema12 - ema26

    dea = calculateEMA(signalPeriod, diff, 0, [])
    print(diff,len(diff))
    print('dea=', dea, type(dea))
    macd = 2 * (diff - dea)
    return macd, diff, dea


stock = '300114.XSHE'
fre = '1d'
check_date = '2010-01-04'
emaArray = []
engine_stock = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/stocks?charset=utf8')
df = pd.read_sql_query('select date,open,close,high,low,volume,money from `stk%s_%s`' % (stock, fre), engine_stock)
df_close = df['close']
ema = EMA(stock, check_date=f'2009-12-31', timeperiod=30)[stock]
print('ddd', ema)
df4 = calculateMACD(df_close, ema)
# df2 = calculateMACD(emaArray, df_close)
# print('df2=', df2)
# print(len(df), len(df2[0]))
# df3 = pd.concat([df, pd.Series(df2[0]).rename('macd'), pd.Series(df2[1]).rename('diff'), pd.Series(df2[2]).rename('dea')], axis=1)

print('df=', df)
df4 = pd.concat([df, pd.Series(df4[0]).rename('macd'), pd.Series(df4[1]).rename('diff'), pd.Series(df4[2]).rename('dea')], axis=1)
print('df=', df4)
print(df4.loc[df.date== '2010-09-10',:])

x_macd_dif, x_macd_dea, x_macd_macd = MACD(stock, check_date='2010-09-13 00:00:00', SHORT=12, LONG=26, MID=9, unit=fre)
print(x_macd_macd, x_macd_dif, x_macd_dea)