from jqlib.technical_analysis import *
import pandas as pd
import numpy as py
import talib
import matplotlib.pyplot as plt
import datetime
from sklearn.ensemble import RandomForestRegressor
from sklearn.svm import SVR
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

stock = '000999.XSHE'
#stock = '000001.XSHG'
start_date='2017-01-01'
end_date1='2018-08-17'

#print(get_index_stocks('000988.XSHG'))

early = datetime.datetime.strptime(start_date,'%Y-%m-%d') - datetime.timedelta(52)
#获取股票历史数据，并根据数据对股票买入点进行计算，最后对买入点好坏进行评估，给出相应因素
df = get_price(stock, early, end_date1,'daily', ['close','volume'], skip_paused=True, fq='pre', count=None)

security_list = [stock]
rsi_l=[]
v1_l=[]
v2_l=[]
K_l=[]
D_l=[]
J_l=[]


for date in df.index:

    #RSI指数
    rsi = RSI(security_list, date, N1=6)
    rsi_l.append(rsi[stock])

    #成交量
    v1,v2 = VR(security_list, date, N=26, M=6)
    v1_l.append(v1[stock])
    v2_l.append(v2[stock])

    #KDJ指标
    K,D,J = KDJ(security_list, date, N =9, M1=3, M2=3)
    K_l.append(K[stock])
    D_l.append(D[stock])
    J_l.append(J[stock])
    
df['RSI']=rsi_l
df['V1']=v1_l
df['V2']=v2_l
df['K']=K_l
df['D']=D_l
df['J']=J_l

lag=5
RR=[]
for i in range(len(df)):
    if(i<lag):
        rr=1
    else:
        rr=df.close[i]/df.close[i-lag]
    RR.append(rr)
df['RR']=RR[lag:]+[1]*lag

#df['K_D']=df.K-df.D
#df['J_D']=df.J-df.D

ff=df[df.columns.difference(['close','volume'])]

rfr=RandomForestRegressor()
#svr=SVR(kernel='rbf')

#trainx,testx,trainy,testy=train_test_split(df[df.columns.difference(['RR'])],df.RR,test_size=0.2)

size=int(len(ff)*0.5)
dd=ff[ff.columns.difference(['RR'])]
trainx=dd.iloc[:size,:]
testx=dd.iloc[size:,:]
trainy=df.RR.iloc[:size]
testy=df.RR.iloc[size:]

rfr.fit(trainx,trainy)
#svr.fit(trainx,trainy)
tt=rfr.predict(testx)
#tt=svr.predict(testx)
#plt.plot(tt)
#plt.show()

plt.plot((testy-tt)/testy)
plt.show()
plt.plot(testy)
plt.show()
plt.plot(tt)
plt.show()
plt.plot(trainy-rfr.predict(trainx))
plt.show()


sum=100.0
num=0
for i in range(len(testx)):
    if rfr.predict(testx.iloc[i].reshape(1,-1))>1.02:
        if sum>0:
            num=sum/df.iloc[size+i].close
            sum=0
            print('buy at:'+str(df.iloc[size+i].close)+" in:"+str(testx.index[i]))
    if rfr.predict(testx.iloc[i].reshape(1,-1))<0.99:
        if num>0:
            sum=num*df.iloc[size+i].close
            num=0
            print('sell at:'+str(df.iloc[size+i].close)+" in:"+str(testx.index[i]))
total=sum+num*df.iloc[size+i].close
print(total)
print(num)
print(ff.columns.difference(['RR']))
rfr.feature_importances_