import datetime
import pandas as pd
from jqdata import *
from jqfactor import get_factor_values

def get_dividend_rate(code_list,date=get_trade_days(start_date=None, end_date=datetime.datetime.now().date(), count=2)[-1]):
    '''date參數需傳入datetime.date，默認為最新交易日'''
    if date < datetime.datetime.now().date():
        q_date=date
    else:
        q_date=get_trade_days(start_date=None, end_date=datetime.datetime.now().date(), count=2)[0]
    if type(code_list)!=list:
        code_list=code_list.split()
    ##計算各報告期的分紅
    f=finance.STK_XR_XD
    q1=query(f.code,f.company_name,f.report_date,f.bonus_amount_rmb).filter(finance.STK_XR_XD.code.in_(code_list)).order_by(finance.STK_XR_XD.report_date.desc())
    df = finance.run_query(q1)
    df=df.fillna(0)
    #df['report_date']=df['report_date'].apply(lambda x:datetime.datetime.combine(x, datetime.time(0)))
    df.set_index(pd.to_datetime(df.report_date), inplace=True)#字符串到datetime
    df=df.groupby('code').resample('Y').sum()
    df=df.reset_index()
    df.columns=['code','year','bonus_amount_rmb']
    ##計算股票分紅使用的報告期
    y1=str(q_date.year-2)
    y2=str(q_date.year-1)
    q2 = query(
            balance.statDate,
            balance.code,
            balance.pubDate,
        ).filter(
            balance.code.in_(code_list)
        )
    ret=get_fundamentals(q2,statDate=y1)
    ret2=get_fundamentals(q2, statDate=y2)
    x=ret.append(ret2)
    x.pubDate=pd.to_datetime(x.pubDate)
    x.set_index(pd.to_datetime(x.pubDate), inplace=True)#
    pub_df=x.groupby('code').max()
    pub_df.drop('statDate',axis=1,inplace=True)
    pub_df['year']=pub_df.applymap(lambda x: pd.Timestamp(y1+'-12-31') if q_date<=x.date() else pd.Timestamp(y2+'-12-31'))
    pub_df=pub_df.reset_index()
    ##計算市值數據
    q3 = query(
                valuation.day,
                valuation.code,
                valuation.market_cap,
            ).filter(
                balance.code.in_(code_list)
            )
    cap=get_fundamentals(q3,q_date)
    ##連接各表並計算股息率
    mix=pd.merge(df,pub_df)
    div_df=pd.merge(mix,cap)
    div_df['dividend_rate']=div_df.bonus_amount_rmb/div_df.market_cap/10000
    div_df.set_index('code', inplace=True)
    return div_df.dividend_rate

stocks = get_index_stocks('000300.XSHG')
HS300_divid=get_dividend_rate(stocks,datetime.date(2018,4,21))
HS300_divid