from jqdata import *
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.linear_model import LinearRegression as LR
from sklearn.model_selection import train_test_split
from sklearn.model_selection import cross_val_score
import warnings
import time
warnings.filterwarnings('ignore') 

#获取个股所属申万行业代码，因为目前只有申万行业行情
def secu_industry_code(code):
    a=[]
    aa=get_industry(code, date=None)
    for i in aa[code].keys():
        if 'sw' in i:
            a.append(aa[code][i]['industry_code'])
    return a
code="600008.XSHG"
#测试
bb=secu_industry_code(code)
bb

#获取个股对应大盘的指数代码
def Correspondence_index(code):
    a=code[:2]
    if a=='60':
        b='000001.XSHG'
    elif a=='30':
        b='399006.XSHE'
    elif a=='00':
        if code[:3]=='000':
            b='399001.XSHE'
        else:
            b='399005.XSHE'
    return b           

def handle_table(code,num,a='g'):
    #num为获取数据的长度
    #a为修改列名称的标签，无特别意义
    secu=attribute_history(code, num, unit='1d',
            fields=[ 'close', 'volume',],
            skip_paused=True, df=True, fq='pre')
    #计算价格的涨跌幅，并删除空行
    secu_change = secu.close.pct_change()
    #新增涨跌幅列，乘以100转换成百分比
    secu['chg'+a]=secu_change*100
    returns=secu
    #returns['close']=secu.close
    returns.rename(columns={'close': 'close'+a},inplace=True)
    returns.rename(columns={'volume': 'volume'+a},inplace=True)
    #增加‘date’列，值为原来的索引
    returns['date']=returns.index
    #按日期降序排列
    returns=returns.sort_values('date', ascending=False)
    #重置索引，丢掉原来的索引
    returns=returns.reset_index(drop = True)
    return returns

#获取个股和行业以及对应大盘行情数据
def secu_industry_quotation(code,limit=None,date=None):
    #limit获取数据的行数，默认全部
    #date获取某日期之前的数据，默认当天
    #获取行业板块行情
    bb=secu_industry_code(code)
    a=1
    #如果时间没有指定，那么取当前时间
    if date:
        pass
    else:
        date=time.strftime('%Y-%m-%d',time.localtime())
    for i in bb:
        df1=finance.run_query(query(
            finance.SW1_DAILY_PRICE.date,
            finance.SW1_DAILY_PRICE.close,
            finance.SW1_DAILY_PRICE.volume,
            finance.SW1_DAILY_PRICE.change_pct
            ).filter(
            finance.SW1_DAILY_PRICE.code==i,
            finance.SW1_DAILY_PRICE.date<=date).order_by(
            finance.SW1_DAILY_PRICE.date.desc()
        ).limit(limit))
        #批量修改列名称
        df1=df1.rename(columns=lambda x: x + str(a))
        #把date列改回来，作为合并轴
        df1=df1.rename(columns={'date'+str(a):'date'})
        if a==1:
            df=df1
        else:
            df=df.merge(df1,on='date')
        a+=1
    #把‘date’列的数据的object类型，转换成datetime64[ns]类型
    df['date']=df['date'].astype('datetime64[ns]')
    num=len(df)
    #获取个股行情数据
    df1=handle_table(code,num,a='g')
    #个股行情和行业行情表合并
    df=df.merge(df1,on='date')
    #获取对应大盘行情
    c=Correspondence_index(code)
    df2=handle_table(c,num,a='d')
    df=df.merge(df2,on='date')
    #个股行情和行业行情表合并
    return df
df=secu_industry_quotation(code,date='2018-08-22')
df.head()

#把成交量都变成变化率
def change_vol(df):
    #获取columns列表
    listdf=list(df)
    for i in listdf:
        if 'volume' in i:
            #计算变化率
            dfs1 = df[i].pct_change(-1)
            a='vol'+i[-1]
            df[a]=dfs1
    #删掉nan行
    df.dropna(inplace=True)
    return df
df1=change_vol(df)
df1

#删除指定的列
def del_columns(df,col):
    dfx=df
    for i in df.columns:
        if col in i:
            dfx.drop(i,axis=1,inplace=True)
    return dfx
df2=del_columns(df1,'close')
df2=del_columns(df2,'volume')
df2

#生成标签列，Y=secu_change的下一列值列。
def generate_label(df):
    #生成新列“y”等于‘secu_change’的下一个值
    df['y']=df['chgg'].shift(1)
    #生成第二天涨幅大于2的（0,1）标签
    df['y2']=np.where(df.y>2,1,0)
    #删除包含‘nan’的行
    df=df.dropna()
    df=df.reset_index(drop=True).dropna()
    return df  
df3=generate_label(df2)
df3

#获取数据X,Y数据集
def current_data(code):
    #获取个股、行业、大盘行情数据列表
    df=secu_industry_quotation(code)
    #把成交量转化成变化率
    df1=change_vol(df)
    #删除多余的列
    df2=del_columns(df1,'close')
    df2=del_columns(df2,'volume')
    #添加y标签
    df3=generate_label(df2)
    df3=df3.dropna()
    X,Y=df3.iloc[:,1:11],df3.y2
    return X,Y   

#运用随机森林进行预测
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
code="600519.XSHG"
X,Y=current_data(code)
cfc_1=[]
#进行十次十组交叉验证
for i in range(10):
    cfc=RandomForestClassifier(n_estimators=25)
    cfc_s=cross_val_score(cfc,X,Y,cv=10).mean()
    cfc_1.append(cfc_s)
plt.plot(range(1,11),cfc_1,label='RondomForst')
plt.legend(loc='best')
plt.show()

#学习曲线,确定n_estimators参数，先range(0,200,10),确定范围，然后，再在最大值上下十个点，进行精确定位n_estimators
superpa=[]
for i in range(130,150):
    cfc=RandomForestClassifier(n_estimators=i+1,n_jobs=-1)
    cfc_s=cross_val_score(cfc,X,Y,cv=10).mean()
    superpa.append(cfc_s)
print(max(superpa),superpa.index(max(superpa)))
plt.figure(figsize=[20,5])
plt.plot(range(130,150),superpa)
plt.show()

#训练调试好的参数的模型
#初始模型我们以600519的数据进行模型训练。
code="600519.XSHG"
cfc=RandomForestClassifier(n_estimators=140,n_jobs=-1)
X,Y=current_data(code)
Xtrain,Xtest,Ytrain,Ytest=train_test_split(X,Y,test_size=0.3,random_state=10)
cfc.fit(Xtrain,Ytrain)
cfc.score(Xtest,Ytest)

#保存训练好的模型
from sklearn.externals import joblib
save='tch'+'cfc'+'.pkl'
joblib.dump(cfc,save)

from sklearn.externals import joblib 
code="600835.XSHG"
X,Y=current_data(code)
model=joblib.load('tchcfc.pkl')
model.score(X,Y)

#使用上证50中个股，对模型进行有效性测试
b=get_index_stocks('000016.XSHG', date=None)
score=[]
for code in b:
    X,Y=current_data(code)
    model=joblib.load('tchcfc.pkl')
    try:
        b=model.score(X,Y)
    except ValueError:
        pass
    score.append(b)
plt.plot(range(len(score)),score)
plt.scatter(range(len(score)),score,marker='o',label='成功率')
plt.legend(loc='best')
plt.show()

np.array(score).mean()

np.array(score).max()

np.array(score).min()

b=np.array(score)
c=np.sort(b)
plt.plot(range(len(c)),c)
plt.scatter(range(len(c)),c,marker='o',label='成功率由小到大')
plt.legend(loc='best')
plt.show()

c

#密度曲线
import seaborn as sns
sns.kdeplot(c, shade=True,label='密度曲线')
plt.legend()
plt.show()

from sklearn.externals import joblib
#预测模型
def predict_data(code,date):
    df=secu_industry_quotation(code,limit=2,date=date)
    df1=change_vol(df)
    df1=del_columns(df,'close')
    df1=del_columns(df,'volume')
    x=df.iloc[:,1:11]
    model=joblib.load('tchcfc.pkl')
    a=model.predict(x)
    return a

#进行选股
def check_stock(date=None):
    #进行选股
    a=[]
    security=get_index_stocks('000016.XSHG', date=date)
    for code in security:
        try:
            if predict_data(code,date)==1:
                g.a.append(code)
        except :
            pass
    return a

a=get_trade_days(start_date='2010-8-1', end_date=None, count=None)
for i in a:
    a=check_stock(date=i)
    if len(a)>0:
        print(a)