from jqdata import *
import datetime
import numpy as np
import pandas as pd
from jqdata import finance
import talib
#备注：成交量使用get_price，而财务数据获取使用的get_fundamentals函数，写成两段了

#1、取到需要取数据的具体日期
date_list = [datetime.datetime.strptime(str(i) + '-' + str(j) + '-01','%Y-%m-%d').date() for i in range(2005,2019) for j in range(1,13)]

#2、使用的财务数据种类：市值，营业额，负债，净利润,以及对应数据所属数据库中表名
factor_list = [('valuation','market_cap'),('income','total_operating_revenue'),('balance','total_liability')]

#3、取到行业成分股
def get_index_stock(industry_level,current_date,wanted):
    all_industry_code = get_industries(industry_level)
    ind_name = list(all_industry_code['name'])
    ind_codes = list(all_industry_code.index)
    ind_stock_dict = {}
    for ind_code in ind_codes:
        ind_sample = ind_name[ind_codes.index(ind_code)]
        ind_stock_dict[ind_sample] = get_industry_stocks(ind_code, date=current_date)
    if wanted == 'Name':
        return ind_name
    else:
        return ind_stock_dict

#取到想要的数据
def get_data(factor_tuple,stock_list,date):
    encoding_str = str(factor_tuple[0]) + '.' + str(factor_tuple[1])
    q = query(
        eval(encoding_str)
      ).filter(
          eval(factor_tuple[0]).code.in_(stock_list))
    factor_data = get_fundamentals(q,date)
    return factor_data.fillna(0).sum()

#4、检查2005年到现在行业分类标准是否有品类上的变更
def repeat_check(industry_level):
    ind_name = get_index_stock(industry_level,date_list[0],'Name')
    for date in date_list:
        new = get_index_stock(industry_level,date,'Name')
        if list(ind_name) != list(new):
            print(industry_level + '有变更' + str(date) )
            break
        else:
            pass
    if industry_level == 'jq_l1':
        print('聚宽1级分类无变更')

repeat_check('jq_l1')

#5、最终的画图函数
def plot_figure(industry_level):
    for factor in factor_list:
        data_dict = {}
        for date in date_list:
            current_stock_dict = get_index_stock(industry_level,date,'list')
            data_list = {}
            cal_percent_list = []
            for ind_key in current_stock_dict.keys():
                current_data = get_data(factor,current_stock_dict[ind_key],date)
                data_list[ind_key] = float(current_data)
                cal_percent_list.append(float(current_data))
            sum_factor = sum(cal_percent_list)
            for key in data_list.keys():
                data_list[key] = data_list[key] / sum_factor
            data_dict[date] = data_list
        data_df = pd.DataFrame(data_dict).T
        data_df.plot(kind = 'area',title = factor[1],figsize = (12,6))

# 6、成交量画图
vol_df = pd.DataFrame()
current_stock_dict = get_index_stock('jq_l1',date_list[-1],'list')
for ind_key in current_stock_dict.keys():
    current_data = get_price(current_stock_dict[ind_key], start_date=date_list[0], end_date=date_list[-1], frequency='daily', fields=['volume'])['volume'].fillna(0).sum(axis = 1)
    vol_df[ind_key] = current_data
vol_sum = vol_df.sum(axis = 1)
vol_sum
col_list = vol_df.columns
for item in col_list:
    vol_df[item] = vol_df[item] / vol_sum
vol_df.plot(kind = 'area',figsize = (12,6),title = '成交量百分比堆积图')

# 7、绘制财务数据图
plot_figure('jq_l1')

# 8、上市公司数量数据图
stock_num_dict = {}
for date in date_list:
    stock_num_dict[date] = {}
    stock_list = get_index_stock('jq_l1',date,'list')
    for idne_name in list(stock_list.keys()):
        stock_num_dict[date][idne_name] = len(stock_list[idne_name])
stock_num_df = pd.DataFrame(stock_num_dict)

stock_num_df = stock_num_df.T
col_sum = stock_num_df.sum(axis = 1)
for col in stock_num_df.columns:
    stock_num_df[col] = stock_num_df[col] / col_sum
stock_num_df.plot(kind = 'area',figsize = (12,6),title = '上市公司数量堆积')