import numpy as np
import pandas as pd
import datetime
import seaborn as sns
from jqdata import *
sns.set(font='serif')
plt.rcParams['axes.unicode_minus'] = False 

trade_date = get_trade_days(end_date=datetime.datetime.now(), count=50)
index_codes = ['000300.XSHG', '000905.XSHG']  # 沪深300，上证500
start_date = '2013-03-15'
end_date = trade_date[-2]

'''计算指数的平均市场与市盈率'''
def func(index):
    q = query(valuation.code, 
              valuation.market_cap,
              valuation.pe_ratio).filter(valuation.code.in_(get_index_stocks(index)))
    df = get_fundamentals(q, end_date)
    return {'market_cap':df['market_cap'].mean(),'pe_ratio':df['pe_ratio'].mean()}
dict_index = {get_security_info(index).display_name: func(index) for index in index_codes}
pd.DataFrame(dict_index).T

'''指数市值分布情况'''
df_index = pd.DataFrame(dict_index).T
df_index['market_cap'].plot(kind='bar')
plt.show()

'''指数市盈率分布情况'''
df_index = pd.DataFrame(dict_index).T
df_index['pe_ratio'].plot(kind='bar')
plt.show()

'''指数的历史价格趋势'''
dict_close = {get_security_info(code).display_name:\
              get_price(code, start_date, end_date)['close']\
              for code in index_codes}
pd.DataFrame(dict_close).plot()
plt.show()

'''指数20日前的价格与昨日价格获取'''
price = {get_security_info(index).display_name:\
         get_price(index, end_date=trade_date[-2], count=50)['close'].values[[-20, -1]]\
         for index in index_codes}
price

'''计算前20日动量'''
ratio = {key: {'ratio':(value[1] - value[0]) / value[0]} for key, value in price.items()}
ratio

'''动量可视化对比'''
df_ratio = pd.DataFrame(ratio).T
df_ratio['ratio'].plot(kind='bar')
plt.show()

'''相关ETF的价格趋势'''
etf_list = ['510300.XSHG', '510500.XSHG']
dict_close = {get_security_info(code).display_name:\
              get_price(code, start_date, end_date)['close']\
              for code in etf_list}
pd.DataFrame(dict_close).plot()
plt.show()

'''数据相关程度计算'''
a=[0.6557,0.0357,0.8491,0.9340,0.6787];
b=[0.7315,0.1100,0.8884,0.9995,0.6959];
print('计算一\n',mean(multiply((a-mean(a)),(b-mean(b))))/(std(b)*std(a)))
print('计算二\n', corrcoef(a,b))
"""
corrcoef得到相关系数矩阵（向量的相似程度）
第一个1是a与a的相关系数,左边第一个0.9976是a与b相关系数,第二个0.9976是b与a相关系数,第二个1是b与b的相关系数
"""

'''计算指数之间的相关程度'''
stock_base = '000300.XSHG'
index_list = get_all_securities(types=['fund'], date=end_date).index
func = lambda index_a, index_b: np.corrcoef(get_price(index_a, start_date, end_date)['close'].values, 
                                            get_price(index_b, start_date, end_date)['close'].values
                                           )[0, 1]
df_corr = pd.DataFrame({index : {'corrcoef': func(stock_base, index)}\
                        for index in index_list}).T.dropna()
df_corr = df_corr.sort_values(by=['corrcoef'], ascending=False).head()
df_corr

'''正当关指数与ETF趋势对比'''
compare_codes = (stock_base, df_corr.index[0])
df_test = pd.DataFrame()
for index in compare_codes:
    df_test[get_security_info(index).display_name] = get_price(index, start_date, end_date)['close']
df_test.plot(subplots=True)
plt.show()

'''找出与指数相关程度最高的ETF'''
def func(code):
    etfs = get_all_securities(types=['fund'], date=end_date).index
    func = lambda index_a, index_b: np.corrcoef(get_price(index_a, start_date, end_date)['close'].values, 
                                                get_price(index_b, start_date, end_date)['close'].values
                                               )[0, 1]
    df_corr = pd.DataFrame({etf : {'corrcoef': func(code, etf)}\
                            for etf in etfs}).T.dropna()
    df_corr = df_corr.sort_values(by=['corrcoef'], ascending=False).head()
    return df_corr.index[0], df_corr.values[0]
eft_list = map(func, index_codes)
result = list(eft_list)
result

'''找出与ETF负相关程度最高的ETF'''
def func(code):
    etfs = get_all_securities(types=['fund'], date=end_date).index
    func = lambda index_a, index_b: np.corrcoef(get_price(index_a, start_date, end_date)['close'].values, 
                                                get_price(index_b, start_date, end_date)['close'].values
                                               )[0, 1]
    df_corr = pd.DataFrame({etf : {'corrcoef': func(code, etf)}\
                            for etf in etfs}).T.dropna()
    df_corr = df_corr.sort_values(by=['corrcoef'], ascending=True).head()
    return df_corr.index[0], df_corr.values[0]
eft_list = map(func, etf_list)
result = list(eft_list)
result

'''大盘与国债指数趋势对比'''
compare_codes = (stock_base, '511010.XSHG')
df_test = pd.DataFrame()
for index in compare_codes:
    df_test[get_security_info(index).display_name] = get_price(index, start_date, end_date)['close']
df_test.plot(subplots=True)
plt.show()

'''负相关ETF趋势对比'''
compare_codes = (stock_base, result[0][0])
df_test = pd.DataFrame()
for index in compare_codes:
    df_test[get_security_info(index).display_name] = get_price(index, start_date, end_date)['close']
df_test.plot(subplots=True)
plt.show()