#注意：需先运行下方函数代码，才能运行本段代码
df_use=cal_df('2016q3',3,4,4,0.12,0.18,0.18)
df_use

import numpy as np
import pandas as pd
from jqdata import *
from datetime import datetime
import datetime as dt
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
from matplotlib.finance import candlestick_ohlc

def back(x):
    if x[-1]=='1':
        return str(int(x[:4])-1)+'q4'
    else:
        return str(x[:5])+str(int(x[-1])-1)
    
today=datetime.strptime(str(datetime.now())[:10], '%Y-%m-%d')
def delta_days(x):
    d2=datetime.strptime(str(x)[:10], '%Y-%m-%d')
    return (today-d2).days

def cal_df(q_end,p3=4,p2=4,p1=4,r3=0.15,r2=0.15,r1=0.15,date_for_index='2017-12-31'):
    #选取上证综指和深证成指的全体股票
    g1=get_index_stocks('000001.XSHG',date=date_for_index)
    g2=get_index_stocks('399001.XSHE',date=date_for_index)
    stocks=g1+g2
    q1=query(indicator).filter(indicator.code.in_(stocks))
    q2=query(balance).filter(balance.code.in_(stocks))
    
    #计算第三段扣非净益率
    df=get_fundamentals(q1,statDate=q_end)[['code','adjusted_profit']]
    df=df.rename(columns={'adjusted_profit':'prof'+q_end})
    df.index=df['code'].values
    df_=get_fundamentals(q2,statDate=q_end)[['code','total_owner_equities']]
    df_.index=df_['code'].values
    df['equity'+q_end]=df_['total_owner_equities']
    q_date=q_end
    for i in range(p3-1):
        q_date=back(q_date)
        df_=get_fundamentals(q1,statDate=q_date)[['code','adjusted_profit']]
        df_.index=df_['code'].values
        df['prof'+q_end]=df['prof'+q_end]+df_['adjusted_profit']
    q_end3=q_end
    df['rate'+q_end3]=df['prof'+q_end3]/df['equity'+q_end3]
    
    #计算第二段扣非净益率
    q_end=back(q_date)
    df_=get_fundamentals(q1,statDate=q_end)[['code','adjusted_profit']]
    df_.index=df_['code'].values
    df['prof'+q_end]=df_['adjusted_profit']
    df_=get_fundamentals(q2,statDate=q_end)[['code','total_owner_equities']]
    df_.index=df_['code'].values
    df['equity'+q_end]=df_['total_owner_equities']
    q_date=q_end
    for i in range(p2-1):
        q_date=back(q_date)
        df_=get_fundamentals(q1,statDate=q_date)[['code','adjusted_profit']]
        df_.index=df_['code'].values
        df['prof'+q_end]=df['prof'+q_end]+df_['adjusted_profit']
    q_end2=q_end
    df['rate'+q_end2]=df['prof'+q_end2]/df['equity'+q_end2]    
    
    #计算第一段扣非净益率
    q_end=back(q_date)
    df_=get_fundamentals(q1,statDate=q_end)[['code','adjusted_profit']]
    df_.index=df_['code'].values
    df['prof'+q_end]=df_['adjusted_profit']
    df_=get_fundamentals(q2,statDate=q_end)[['code','total_owner_equities']]
    df_.index=df_['code'].values
    df['equity'+q_end]=df_['total_owner_equities']
    q_date=q_end
    for i in range(p1-1):
        q_date=back(q_date)
        df_=get_fundamentals(q1,statDate=q_date)[['code','adjusted_profit']]
        df_.index=df_['code'].values
        df['prof'+q_end]=df['prof'+q_end]+df_['adjusted_profit']
    q_end1=q_end
    df['rate'+q_end1]=df['prof'+q_end1]/df['equity'+q_end1]
    
    #选取符合条件的股票
    df=df[(df['rate'+q_end1]>r1) & (df['rate'+q_end2]>r2) & (df['rate'+q_end3]>r3) 
          &(df['prof'+q_end1]>0)&(df['prof'+q_end2]>0)&(df['prof'+q_end3]>0)]
    df=df[['rate'+q_end3,'rate'+q_end2,'rate'+q_end1]]
    
    #对符合条件的股票进行处理
    #start_date:上市日期 days:上市天数 pe_ratio：当天的市盈率 market_cap：当天的市值
    sat_list=list(df.index)
    name_list=list(map(lambda x:get_security_info(x).display_name,sat_list))
    start_list=list(map(lambda x:get_security_info(x).start_date,sat_list))
    df.insert(0,'name',name_list)
    df['start_date']=start_list
    days_list=list(map(lambda x:delta_days(x),start_list))
    df['days']=days_list
    df=df.sort(columns='rate'+q_end3,ascending=False)
    q3=query(valuation).filter(valuation.code.in_(sat_list))
    fund=get_fundamentals(q3)[['code','pe_ratio','market_cap']]
    fund.index=fund['code'].values
    df['pe_ratio']=fund['pe_ratio']
    df['market_cap']=fund['market_cap']
    return df

#注意：需先运行下方函数代码，才能运行本段代码
draw_k(df_use,'2015-01-05','2017-12-29')

def draw_k(df,st_date,ed_date,freq='5m',delta=240):
    #获取需要使用的股票代码，并标记开始和结束的日期
    value_stocks=list(df.index)
    days=get_all_trade_days()
    trade_days=list(map(lambda x:str(x),days))
    for i in range(len(trade_days)):
        if trade_days[i]==st_date:
            day_flag1=i
        if trade_days[i]==ed_date:
            day_flag2=i
    
    #获取对应股票的价格，获取投资组合的价格情况
    prices=get_price(value_stocks,start_date=trade_days[day_flag1],end_date=trade_days[day_flag2],frequency=freq)
    prices['sum']=prices.sum(axis=1)
    close_=prices['close'].sum(axis=1).values
    open_=prices['open'].sum(axis=1).values
    high_=prices['high'].sum(axis=1).values
    low_=prices['low'].sum(axis=1).values
    volum_=prices['volume'].sum(axis=1).values

    closep=list(close_[delta*i+delta-1] for i in range(int(len(close_)/delta)))
    openp=list(close_[delta*i] for i in range(int(len(close_)/delta)))
    highp=list(max(close_[delta*i:delta*i+delta-1]) for i in range(int(len(close_)/delta)))
    lowp=list(min(close_[delta*i:delta*i+delta-1]) for i in range(int(len(close_)/delta)))
    volume=list(np.mean(close_[delta*i:delta*i+delta-1]) for i in range(int(len(close_)/delta)))

    x,y,candleAr=0,len(closep),[]
    date = np.linspace(0,y,y)
    while x < y:
        appendLine = date[x],openp[x],highp[x],lowp[x],closep[x],volume[x]
        candleAr.append(appendLine)
        x += 1
        
    #提取日期，绘制K线图横坐标
    quotes2=get_price('000001.XSHE', start_date=st_date, end_date=ed_date, frequency='daily',skip_paused=False)
    t = quotes2.index.values
    t1 = t[-len(t):].astype(dt.date)/1000000000
    ti=[]
    for i in range(int(len(t)/50)):
        a=i*50
        d = dt.date.fromtimestamp(t1[a])
        t2=d.strftime('$%Y-%m-%d$')
        ti.append(t2)
    d1= dt.date.fromtimestamp(t1[len(t)-1])
    d2=d1.strftime('$%Y-%m-%d$')
    ti.append(d2)
    
    #绘制K线图
    fig = plt.figure(figsize=(16, 8))
    ax1 = plt.subplot2grid((10,4),(0,0),rowspan=10,colspan=4)
    candlestick_ohlc(ax1, candleAr, width=0.5, colorup='r', colordown='g', alpha=0.6)
    ax1.set_xticks(np.linspace(-2,len(closep)+2,len(ti))) 
    ax1.set_xticklabels(ti)
    plt.setp(plt.gca().get_xticklabels(), rotation=45, horizontalalignment='right')
    plt.show()