import pandas as pd
from datetime import datetime

def roe_ttm(stocklist,scandate,num = 10,period ='Q',cut = True,tb = True):
    if isinstance(scandate,datetime):
        scandate = datetime.strftime(scandate,'%Y-%m-%d')
        
    q = query(
        income.statDate,
        income.pubDate,
        income.code,
        income.np_parent_company_owners,
        indicator.adjusted_profit,
        balance.equities_parent_company_owners,
    ).filter(
        income.code.in_(stocklist)
    )
    
    if period.upper() == 'Q':
        datelist = [str(dt) for dt in pd.date_range(start = '2000-01-01',end = scandate,freq ='Q').to_period('Q')]
    elif period.upper() == 'A':
        datelist = [str(dt) for dt in pd.date_range(start = '2000-01-01',end = scandate,freq ='A').to_period('A')]

    if num + 5 < len(datelist) and period.upper() == 'Q':
        datelist = datelist[-num-5:]
    elif num+1 < len(datelist) and period.upper() == 'A':
        datelist = datelist[-num-1:]
  
    df = pd.DataFrame()
    for dt in datelist:
        df = pd.concat([df,get_fundamentals(q,statDate = dt)],axis =0)
    df.columns =['statDate','pubDate','code','net_profit','adjust_profit','equities']
    df.set_index(['statDate','code'],inplace = True)
    df = df.unstack('code')

    if period.upper() == 'Q':
        net_profit_ttm_df    = pd.rolling_sum(df['net_profit'],4)
        adjust_profit_ttm_df = pd.rolling_sum(df['adjust_profit'],4)
        equities_avg_df      = pd.rolling_apply(df['equities'],4,lambda x:x[0]/2+x[-1]/2)
        
        #tb = 摊薄  cut =扣非
        if cut and tb:
            roe_df = adjust_profit_ttm_df/df['equities']
        elif cut and (not tb):
            roe_df = adjust_profit_ttm_df/equities_avg_df
        elif (not cut) and tb:
            roe_df = net_profit_ttm_df/df['equities']
        else:
            roe_df = net_profit_ttm_df/equities_avg_df
        
    elif period.upper() == 'A':
        equities_avg_df      = pd.rolling_apply(df['equities'],2,lambda x:x[0]/2+x[-1]/2)
    
        #tb = 摊薄  cut =扣非
        if cut and tb:
            roe_df = df['adjust_profit']/df['equities']
        elif cut and (not tb):
            roe_df = df['adjust_profit']/equities_avg_df
        elif (not cut) and tb:
            roe_df = df['net_profit']/df['equities']
        else:
            roe_df = df['net_profit']/equities_avg_df
        
    
    roe_df = roe_df.stack()
    roe_df.columns = 'roe'
    df = df.stack()['pubDate']
        
    df = pd.concat([df,roe_df],axis = 1)
    if cut and tb:
        field = 'roe_cut_tb'
    elif cut and (not tb):
        field = 'roe_cut'
    elif (not cut) and tb:
        field = 'roe_tb'
    else:
        field = 'roe'   
    df.columns =['pubDate',field]
    df.reset_index(inplace = True)
    return df.dropna()
   
#example#
roe_ttm(get_index_stocks('000010.XSHG'),'2017-04-28',num = 10,period ='Q',cut = True,tb = False)