import talib
import pandas as pd
import numpy as np
import datetime
from jqdata import*
import matplotlib.pyplot as plt

#使用成交量，5日MA，5日波动率，预测三日后价格变化

#归一化函数
def turn_one(array):
    range_num = array.max() - array.min()
    min_num = array.min()
    array = (array - min_num) / range_num
    return array

#确定基准股票池，时间，获取价格
stock_list = get_index_stocks('000906.XSHG')
sample_price = get_price(stock_list,'2018-06-30','2018-11-30','1d',['close'])['close']
sample_price = sample_price.fillna(0)

#打标签，3日后涨为1，跌为-1，无变化为0
date_list = sample_price.index
trade_days = [i for i in get_all_trade_days()]
stock_col = sample_price.columns
deal_change = pd.DataFrame()
for date in date_list[4:len(date_list)-3]:
    now_price = sample_price.loc[date,:]
    future_date = trade_days[trade_days.index(date.date()) + 3]
    future_price = sample_price.loc[future_date,:]
    price_change = future_price - now_price
    price_change[price_change > 0] = 1
    price_change[price_change < 0] = -1
    price_change[price_change == 0] = 0
    deal_change[date] = price_change
deal_change = deal_change.T

#计算MA
MA_data = pd.DataFrame(index = date_list[4:len(date_list) - 3])
for i in sample_price.columns:
    array_close = np.array(sample_price[i])
    MA_data[i] = turn_one(talib.MA(array_close,timeperiod = 5)[4:len(array_close) - 3])

#获取成交量
sample_vol = get_price(stock_list,'2018-06-30','2018-11-30','1d',['money'])['money']
VOL_df = pd.DataFrame(index = date_list[4:len(date_list) - 3])
for i in sample_vol.columns:
    array_close = np.array(sample_vol[i])
    VOL_df[i] = turn_one(array_close)[4:len(array_close) - 3]

#获取五日波动率
STD_data = pd.DataFrame(index = date_list[4:len(date_list) - 3])
for i in sample_price.columns:
    array_close = np.array(sample_price[i])
    std_array = []
    for j in range(4,len(array_close) - 3):
        std_array.append(array_close[j-4:j].std())
    std_array = turn_one(np.array(std_array))[:len(array_close) - 3]
    STD_data[i] = std_array

#每只股票时间序列储存在字典里
data_dict = {}
for stock in deal_change.columns:
    data_dict[stock] = []
    for date in deal_change.index:
        current_data = []
        current_data.append(MA_data.loc[date,stock])
        current_data.append(STD_data.loc[date,stock])
        current_data.append(VOL_df.loc[date,stock])
        current_data.append(deal_change.loc[date,stock])
        data_dict[stock].append(current_data)
    data_dict[stock] = np.array(data_dict[stock])

#清洗数据，提取特征和标签
def DataCreat(data,test_part):
    data_length = data.shape[0]
    train_length = int((1 - test_part) * data_length)
    test_length = data_length - train_length
    train_Data = np.zeros((train_length,3))
    test_Data = np.zeros((test_length,3))
    train_label = []
    test_label = []
    index = 0
    for array in data:
        if index <= train_length - 1:
            train_label.append(array[3])
            train_Data[index] = array[:3]
        else:
            test_Data[index - train_length] = array[:3]
            test_label.append(array[3])
        index += 1
    return train_Data,train_label,test_Data,test_label

#KNN L2范数计算方法
def deal_distance(train_data,test_vector,train_label,K):
    KNN = ((train_data - test_vector)**2).sum(axis = 1).argsort()[:K+1]
    predict_label = []
    count_label = []
    predict_set = set(train_label)
    for index in KNN:
        predict_label.append(train_label[index])
    for label in predict_set:
        count_label.append((predict_label.count(label)))
    max_index = count_label.index(max(count_label))
    return list(predict_set)[max_index]

#进行预测
def KNN_claassify_predict(data,K,part):
    train_data,train_label,test_data,test_label = DataCreat(data,part)
    preditc_results = []
    for test_vector in test_data:
        result = deal_distance(train_data,test_vector,train_label,K)
        preditc_results.append(result)
    return preditc_results

#不同参属下评估性能
def KNN_evaluat(data,K,part):
    train_data,train_label,test_data,test_label = DataCreat(data,part)
    predict = KNN_claassify_predict(data,K,part)
    score = 0
    for i in range(len(predict)):
        if predict[i] == test_label[i]:
            score += 1
        else:
            pass
    right_rate = score/len(predict)
    return right_rate
rate_list = []
for key in data_dict.keys():
    params_list = [KNN_evaluat(data_dict[key],i,0.2) for i in range(3,20)]
    mean_rate = sum(params_list)/len(params_list)
    rate_list.append(mean_rate)
print('全市场股票涨跌预测准确率为：' + str(sum(rate_list)/len(rate_list)*100) + '%')
print('全市场股票涨跌预测最高准确率为：' + str(max(rate_list)*100) + '%')