import pandas as pd

def KDJ_K(df,n=9):
    df['highest'] = df['high'].rolling(n).max()
    df['lowest'] = df['low'].rolling(n).min()
    df['RSV'] = (df['close']-df['lowest'])/(df['highest']-df['lowest'])*100
    df['signal'] = SMA_REC(df['RSV'],n=3,m=1, start=50)
    return df['signal']

def SMA_REC(vals, n, m , start) :
    sma=[ ]
    numb = 0
    for x in vals:
        if np.isnan(x):
            sma.append(x)
        else:
            if numb == 0:
                ret = start
            else:
                ret = (x * m + ret * (n - m)) / n
            sma.append(ret)
            numb += 1
    result = array(sma)
    return result
# #传入df的index是日期，cloumn是['high','low','close'],需要的数据列名。
def KDJ_J(df,n=9):
    df['highest'] = df['high'].rolling(n).max()
    df['lowest'] = df['low'].rolling(n).min()
    df['RSV'] = (df['close']-df['lowest'])/(df['highest']-df['lowest'])*100
    df['KDJ_K'] = SMA_REC(df['RSV'],n=3,m=1, start=50)
    df['signal'] = 3*df['KDJ_K']-2*SMA_REC(df['KDJ_K'],n=3,m=1, start=50)
    return df['signal']
def KDJ_D(df,n=9):
    df['highest'] = df['high'].rolling(n).max()
    df['lowest'] = df['low'].rolling(n).min()
    df['RSV'] = (df['close']-df['lowest'])/(df['highest']-df['lowest'])*100
    df['KDJ_K'] = SMA_REC(df['RSV'],n=3,m=1, start=50)
    df['signal'] = SMA_REC(df['KDJ_K'],n=3,m=1, start=50)
    return df['signal']
def JK_KDJ_D(security_list,end_date,count = None,start_date = None,n=9):
    if type(security_list)==str:
        security_list=[security_list]
    datalist = get_price(security=security_list,start_date=start_date,end_date=end_date,
                         frequency='daily', fields=['high','low','close'],fq='pre',count=count)
    json = {}
    d='KDJ_D'
    j='KDJ_J'
    k='KDJ_K'
    for security_code in security_list:
        df = datalist.minor_xs(security_code)
        
        df[d] = KDJ_D(df,n)
        df[j] = KDJ_J(df,n)
        df[k] = KDJ_K(df,n)
        json[security_code] = pd.DataFrame(data=df,index=df.index)  
    result = pd.Panel(data=json)
    return result
print(JK_KDJ_D('000002.XSHG','2019-01-01')['000002.XSHG'])