from jqdatasdk import *
import seaborn as sns
import math
import datetime
from tqdm import tqdm
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

1.1、财务数据——总市值和PE

def load_fundamentals_data(context):
    '''
    加载股票的财务数据，包括总市值和PE
    '''
    df = get_fundamentals(query(valuation,indicator), context.current_dt.strftime("%Y-%m-%d"))
    raw_data = []
    for index in range(len(df['code'])):
        raw_data_item = {
            'code'      :df['code'][index],
            'market_cap':df['market_cap'][index],
            'pe_ratio'  :df['pe_ratio'][index]
            }
        raw_data.append(raw_data_item)
    return raw_data

df=get_fundamentals(query(valuation))

df = get_fundamentals(query(valuation,indicator)) 
raw_data = [] 
for index in range(len(df['code'])): 
    raw_data_item = { 
    'code' :df['code'][index], 
    'market_cap':df['market_cap'][index], 
    'pe_ratio' :df['pe_ratio'][index] 
    } 
    raw_data.append(raw_data_item)

# 按照财务信息中的总市值降序排序
raw_data = sorted(raw_data,key = lambda item:item['market_cap'], reverse=True) 
# 剔除总市值排名最小的10%的股票
fitered_market_cap = raw_data[:int(len(raw_data) * 0.9)] 
# 剔除PE TTM 小于0或大于100
filtered_pe = [] 
for stock in fitered_market_cap: 
    if stock['pe_ratio'] == None or math.isnan(stock['pe_ratio']) or float(stock['pe_ratio']) < 0 or float(stock['pe_ratio']) > 100: 
        continue 
    filtered_pe.append(stock['code'])

#PE数据集
df_pe=pd.DataFrame(fitered_market_cap)

1.2、行情数据——25日涨跌幅

current_dt=datetime.datetime.now().strftime('%Y-%m-%d')
CHANGE_PCT_DAY_NUMBER=25

def load_change_pct_data(current_dt,codes):
    change_pct_dict_list = []
    # 计算涨跌幅需要用到前一日收盘价，所以需要多加载一天的数据，
    # 而这里在第二日的开盘前运行，计算前一个交易日收盘后的数据，所以需要再多加载一天
    # 使用固定的25个交易日，而非25个bar计算涨跌幅
    count = CHANGE_PCT_DAY_NUMBER + 1
    # 获取25个交易日的日期
    pre_25_dates = get_trade_days(start_date=None, end_date=current_dt, count=count)
    pre_25_date = pre_25_dates[0]
    pre_1_date = pre_25_dates[-1]
    for code in codes:
        pre_25_data =  get_price(code, start_date=None, end_date=pre_25_date, frequency='daily', fields=['close'], skip_paused=True, fq='post', count=1)
        pre_1_data =  get_price(code, start_date=None, end_date=pre_1_date, frequency='daily', fields=['close'], skip_paused=True, fq='post', count=1)
        pre_25_close = None
        pre_1_close = None
        if str(pre_25_date) == str(pre_25_data.index[0])[:10]:
            pre_25_close = pre_25_data['close'][0]
        if str(pre_1_date) == str(pre_1_data.index[0])[:10]:
            pre_1_close = pre_1_data['close'][0]

        if pre_25_close != None and pre_1_close != None and not math.isnan(pre_25_close) and not math.isnan(pre_1_close):
            change_pct = (pre_1_close - pre_25_close) / pre_25_close
            item = {'code':code, 'change_pct': change_pct}
            change_pct_dict_list.append(item)
    return change_pct_dict_list

def load_change_pct_data_1(current_dt,code):
    change_pct_dict_list = []
    # 计算涨跌幅需要用到前一日收盘价，所以需要多加载一天的数据，
    # 而这里在第二日的开盘前运行，计算前一个交易日收盘后的数据，所以需要再多加载一天
    # 使用固定的25个交易日，而非25个bar计算涨跌幅
    count = CHANGE_PCT_DAY_NUMBER + 1
    # 获取25个交易日的日期
    pre_25_dates = get_trade_days(start_date=None, end_date=current_dt, count=count)
    pre_25_date = pre_25_dates[0]
    pre_1_date = pre_25_dates[-1]
    pre_25_data =  get_price(code, start_date=None, end_date=pre_25_date, frequency='daily', fields=['close'], skip_paused=True, fq='post', count=1)
    pre_1_data =  get_price(code, start_date=None, end_date=pre_1_date, frequency='daily', fields=['close'], skip_paused=True, fq='post', count=1)
    pre_25_close = None
    pre_1_close = None
    if str(pre_25_date) == str(pre_25_data.index[0])[:10]:
        pre_25_close = pre_25_data['close'][0]
    if str(pre_1_date) == str(pre_1_data.index[0])[:10]:
        pre_1_close = pre_1_data['close'][0]

    if pre_25_close != None and pre_1_close != None and not math.isnan(pre_25_close) and not math.isnan(pre_1_close):
        change_pct = (pre_1_close - pre_25_close) / pre_25_close
        item = {'code':code, 'change_pct': change_pct}
        change_pct_dict_list.append(item)
    return change_pct_dict_list

# 加载标的的涨跌幅信息
change_pct_dict_list = load_change_pct_data(current_dt,filtered_pe) 
# 按照涨跌幅升序排序
change_pct_dict_list = sorted(change_pct_dict_list,key = lambda item:item['change_pct'], reverse=False) 
# 取跌幅前10%的股票
change_pct_dict_list = change_pct_dict_list[0:(int(len(change_pct_dict_list)*0.1))]

change_pct_dict_list=[]
for i in tqdm(filtered_pe):
    change_pct_dict_list.append(load_change_pct_data_1(current_dt,i))

change_pct_dict_df1=[]
for i in change_pct_dict_list:
    if i!=[]:
        change_pct_dict_df1.append(i[0])

change_pct_dict_df=pd.DataFrame(change_pct_dict_df1)

1.3财务数据-资产负债率

##资产负债率ratio_liability_vs_assets=总负债total_liability/总资产total_assets   
##合并资产负债表

q=query(finance.STK_BALANCE_SHEET.company_name,
        finance.STK_BALANCE_SHEET.code,
        finance.STK_BALANCE_SHEET.pub_date,
        finance.STK_BALANCE_SHEET.start_date,
        finance.STK_BALANCE_SHEET.end_date,
        finance.STK_BALANCE_SHEET.cash_equivalents,
        finance.STK_BALANCE_SHEET.total_assets,
        finance.STK_BALANCE_SHEET.total_liability
).filter(finance.STK_BALANCE_SHEET.pub_date>='2019-01-01',finance.STK_BALANCE_SHEET.report_type==0)
df=finance.run_query(q)

df['ratio_liability_vs_assets']=df['total_liability']/df['total_assets']

df['rank_1']=df.sort_values(['pub_date', 'end_date'], ascending=[ False, False]).groupby(['code']).cumcount() + 1

df2=df[df['rank_1']==1]

df2['ratio_liability_vs_assets_seg']=pd.cut(df2['ratio_liability_vs_assets'],[0,0.3,0.4,0.6,0.8,1,10])

pd.value_counts(df2['ratio_liability_vs_assets_seg'])

(0.0, 0.3]     996
(0.4, 0.6]     920
(0.3, 0.4]     517
(0.6, 0.8]     470
(0.8, 1.0]     137
(1.0, 10.0]      8
Name: ratio_liability_vs_assets_seg, dtype: int64

df2['ratio_liability_vs_assets'].describe()

count    3048.000000
mean        0.422930
std         0.266916
min         0.009660
25%         0.251521
50%         0.402005
75%         0.562507
max         8.257912
Name: ratio_liability_vs_assets, dtype: float64

df3=df2[(df2['ratio_liability_vs_assets']<=1)]['ratio_liability_vs_assets']

sns.distplot(df3)

#格式转换

df2['zs1']=df2.ratio_liability_vs_assets_seg.astype(str)

#取资产负债率和25日跌幅

df3=pd.merge(df2,change_pct_dict_df,on='code',how='inner')

df4=df3[['company_name','code','ratio_liability_vs_assets','change_pct']]

# 资产负债率与25日跌幅相关性分析
def plot_corr(df,col_list,threshold=None,plt_size=None,is_annot=True):
    """
    df:数据集
    col_list:变量list集合
    threshold: 相关性设定的阈值
    plt_size:图纸尺寸
    is_annot:是否显示相关系数值
    
    return :相关性热力图
    """
    corr_df = df.loc[:,col_list].corr()
    plt.figure(figsize=plt_size)
    sns.heatmap(corr_df,annot=is_annot,cmap='rainbow',vmax=1,vmin=-1,mask=np.abs(corr_df)<=threshold)
    return plt.show()

# 剔除PE TTM 小于0或大于100
df_pe=df_pe[(df_pe.pe_ratio>=0)&(df_pe.pe_ratio<=100)]

#df4加入pe_ratio
df5=pd.merge(df4,df_pe,on='code',how='inner')

1.4财务数据——流动比率

##流动比率 current_ratio =流动资产合计 total_current_assets/流动负债合计 total_current_liability
##合并资产负债表

q=query(finance.STK_BALANCE_SHEET.company_name,
        finance.STK_BALANCE_SHEET.code,
        finance.STK_BALANCE_SHEET.pub_date,
        finance.STK_BALANCE_SHEET.start_date,
        finance.STK_BALANCE_SHEET.end_date,
        finance.STK_BALANCE_SHEET.total_current_assets,
        finance.STK_BALANCE_SHEET.total_current_liability
).filter(finance.STK_BALANCE_SHEET.pub_date>='2018-01-01',finance.STK_BALANCE_SHEET.report_type==0)
df_current=finance.run_query(q)

df_current['current_ratio']=df_current['total_current_assets']/df_current['total_current_liability']

#取最新一份财报数据
df_current['rank_1']=df_current.sort_values(['pub_date', 'end_date'], ascending=[ False, False]).groupby(['code']).cumcount() + 1

df_current=df_current[df_current.rank_1==1]

# df_current=df_current[['code','current_ratio']].dropna()
df_current=df_current[['code','current_ratio']]

#合并到特征表 
df6=pd.merge(df5,df_current,on='code',how='left')

#相关性分析
plot_corr(df6,['ratio_liability_vs_assets','change_pct','market_cap','pe_ratio','current_ratio'])

df_dist=df6[df6.change_pct>0.2]['current_ratio'].dropna()