#导入需要的库
import pandas as pd
import numpy as np
from jqdata import jy
from jqdata import *
import seaborn as sns
import matplotlib as mpl
from cvxopt import solvers, matrix
import datetime as dt
from scipy.optimize import minimize

#获取所需的数据
s_date,e_date = '2012-11-10','2018-12-10'
moneyfund_price = get_price('511880.XSHG',start_date=s_date,end_date=e_date)
#通过聚源数据获取申万行业指数行情数据（输入行业指数代码）
def get_SW_index(SW_index = 801010,start_date = '2017-01-31',end_date = '2018-01-31'):
    index_list = ['PrevClosePrice','OpenPrice','HighPrice','LowPrice','ClosePrice','TurnoverVolume','TurnoverValue','TurnoverDeals','ChangePCT','UpdateTime']
    jydf = jy.run_query(query(jy.SecuMain).filter(jy.SecuMain.SecuCode==str(SW_index)))
    link=jydf[jydf.SecuCode==str(SW_index)]
    rows=jydf[jydf.SecuCode==str(SW_index)].index.tolist()
    result=link['InnerCode'][rows]

    df = jy.run_query(query(jy.QT_SYWGIndexQuote).filter(jy.QT_SYWGIndexQuote.InnerCode==str(result[0]),                                                   jy.QT_SYWGIndexQuote.TradingDay>=start_date,                                                         jy.QT_SYWGIndexQuote.TradingDay<=end_date
                                                        ))
    df.index = df['TradingDay']
    df = df[index_list]
    return df

hy_df = get_industries(name='sw_l1')

hy_price_dict = {}
for hy in hy_df.index[:]:
    hy_price =get_SW_index(SW_index = hy,start_date = s_date,end_date = e_date)
    hy_price_dict[hy] = hy_price['ClosePrice']
all_price_df = pd.concat(hy_price_dict.values(),axis=1)
all_price_df.columns = hy_df.index
#all_price_df['moneyfund'] = moneyfund_price['close']
all_price_df.head(5)

mpl.pyplot.style.use('ggplot')
all_price_df.columns = list(hy_df['name'].values)#+['货币基金']
mpl.rcParams['font.family']='serif'
mpl.rcParams['axes.unicode_minus']=False # 处理负号
prices=all_price_df
(prices/prices.iloc[0]).plot(figsize=(18,12),grid='on')
pct_daily = prices.pct_change()

#计算各行业指数相关矩阵
fig = plt.figure(figsize= (15,10))
ax = fig.add_subplot(111)
ax = sns.heatmap(pct_daily.corr(),annot=True,annot_kws={'size':9,'weight':'bold'})

#获取日期列表
#按月、季度、半年方式获取交易日列表
def get_tradeday_list(start,end,frequency=None,count=None):
    if count != None:
        df = get_price('000001.XSHG',end_date=end,count=count)
    else:
        df = get_price('000001.XSHG',start_date=start,end_date=end)
    if frequency == None or frequency =='day':
        return df.index
    else:
        df['year-month'] = [str(i)[0:7] for i in df.index]
        if frequency == 'month':
            return df.drop_duplicates('year-month').index
        elif frequency == 'quarter':
            df['month'] = [str(i)[5:7] for i in df.index]
            df = df[(df['month']=='01') | (df['month']=='04') | (df['month']=='07') | (df['month']=='10') ]
            return df.drop_duplicates('year-month').index
        elif frequency =='halfyear':
            df['month'] = [str(i)[5:7] for i in df.index]
            df = df[(df['month']=='01') | (df['month']=='06')]
            return df.drop_duplicates('year-month').index 

get_tradeday_list(start=s_date,end=e_date,frequency="quarter",count=None)

s_date,e_date = '2013-12-10','2018-12-10'
tradedays = get_tradeday_list(start=s_date,end=e_date,frequency="quarter",count=None)
tradedays