# import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import bisect

#指定日期的指数PE（等权重）
def get_index_pe_date(index_code,date):
    stocks = get_index_stocks(index_code, date)
    q = query(valuation).filter(valuation.code.in_(stocks))
    df = get_fundamentals(q, date)
    if len(df)>0:
        pe = len(df)/sum([1/p if p>0 else 0 for p in df.pe_ratio])
        return pe
    else:
        return float('NaN')
    
#指数历史PE
def get_index_pe(index_code):
    start='2005-1-1'
    end = pd.datetime.today();
    dates=[]
    pes=[]
    for d in pd.date_range(start,end,freq='M'): #频率为月
        dates.append(d)
        pes.append(get_index_pe_date(index_code,d))
    return pd.Series(pes, index=dates)

all_index = get_all_securities(['index'])
index_choose =['000016.XSHG',
               '000300.XSHG',
               '000905.XSHG',
               '399001.XSHE',
               '399005.XSHE',
               '399006.XSHE',
               '000990.XSHG',
               '000991.XSHG',
               '000992.XSHG',
               '000993.XSHG',
               '000827.XSHG',
               '399812.XSHE',
              ]
df_pe = pd.DataFrame()
for code in index_choose:
    print u'正在处理: ',code
    df_pe[code]=get_index_pe(code)

today= pd.datetime.today()
results=[]
for code in index_choose:
    pe = get_index_pe_date(code,today)
    q_pes = [df_pe.quantile(i/10.0)[code]  for i in range(11)]
    idx = bisect.bisect(q_pes,pe)
    quantile = idx-(q_pes[idx]-pe)/(q_pes[idx]-q_pes[idx-1])
    index_name = all_index.ix[code].display_name
    results.append([index_name,'%.2f'% pe,'%.2f'% (quantile*10)]+['%.2f'%q  for q in q_pes]+[df_pe[code].count()])

df_pe.columns=np.array(results)[:,0]
df_pe.plot(figsize=(12,10))
columns=[u'名称',u'当前PE',u'分位点%',u'最小PE']+['%d%%'% (i*10) for i in range(1,10)]+[u'最大PE' , u"数据个数"]
df = pd.DataFrame(data=results,index=index_choose,columns=columns)
df