import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import bisect
from jqdata import *
import math
import datetime
from sklearn import linear_model        #表示，可以调用sklearn中的linear_model模块进行线性回归。
#指定日期的板块PE（等权重）
def get_industry_pe_date(industry,date):
    stocks = get_industry_stocks(industry, date=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_industry_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_industry_pe_date(index_code,d))
    return pd.Series(pes, index=dates)


def count_industry():
    #"sw_l1": 申万一级行业，"sw_l2": 申万二级行业，"sw_l3": 申万三级行业
    all_industry = get_industries(name='sw_l1')
    print all_industry
    industry_choose =['801170',
                   '801030',
                    '801150',
                      '801740',
                      '801110',
                      '801710',
                      '801720',
                      '801180',
                      '801890',
                      '801880',
                      '801080',
                      '801730',
                      '801750',
                      '801790',
                      '801120'
                  ]
    df_pe = pd.DataFrame()
    #分板块计算pe值
    for code in industry_choose:
        print u'正在处理: ',code
        #获取历史时间code下的所用pe
        df_pe[code]=get_industry_pe(code)

    today= pd.datetime.today()
    results=[]
    #进行分位数计算
    for code in industry_choose:
        #获取今天的板块pe值
        pe = get_industry_pe_date(code,today)
        
        #计算分位数,也就是今天的pe值在历史值上的分位数
        q_pes = [df_pe.quantile(i/10.0)[code]  for i in range(11)]
        
        #bisect其目的在于查找该数值将会插入的位置并返回，而不会插入
        #找到本code位于整体分位的次序,比如1，就是历史上的10%（因为总共切了10个分位）
        idx = bisect.bisect(q_pes,pe)
        #计算同一分位段内本code,quantile*10为当前的分位点，处于历史的哪个位置
        quantile = idx-(q_pes[idx]-pe)/(q_pes[idx]-q_pes[idx-1])
        index_name = all_industry.loc[code,"name"]
        index_roe = calculate_roe_ttm(code, today,type='industry')
        results.append([index_name,index_roe,'%.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'当前roe',u'当前PE',u'分位点%',u'最小PE']\
    +['%d%%'% (i*10) for i in range(1,10)]\
    +[u'最大PE' , u"数据个数"]
    df = pd.DataFrame(data=results,index=industry_choose,columns=columns)
    return df
    

##粗略计算过去一年板块的roe情况 industry为板块代码
def calculate_roe_ttm(industry, date,type='industry'):
    #date=xxxx-03-31/.../xxxx-12-31,不支持个性化日期，例如xxxx-01-23
    industry_stock = []
    if(type == "industry"):
        industry_stock = get_industry_stocks(industry, date=date)
    else:
        industry_stock.append(industry)
    qt_str = [str(qt) for qt in pd.period_range(end=date, periods=4, freq='Q')]
    #此处与研究模块的区别 jq.+
    q = query(
        income.statDate,
        income.code,
        indicator.adjusted_profit,
        balance.equities_parent_company_owners
    ).filter(
        income.code.in_(industry_stock)
        #历史统计分析，不考虑未来函数问题
    )
    ##df为四个季度的数据df的数组  单元素:203*4
    df = [get_fundamentals(q, statDate=qt) for qt in qt_str]
    ##四个季度数据组合 812*4
    df = pd.concat(df)
    ##相当于group by,statDate第一列,code第二列 812*2
    df.set_index(['statDate', 'code'], inplace=True)
    ##按code展开 4*406
    df = df.unstack('code')
    ##过滤缺失数据 4 * 405
    df = df.dropna(axis=1)

    #按code合计四个季度的adjusted_profit的值 ----扣除非经常损益后的净利润(元）
    adj_net_profit = df['adjusted_profit'].sum()
    #按code合计四个季度的equities_parent_company_owners的均值 ----归属于母公司股东权益合计(元)
    ##所有者权益合计＝母公司股东权益合计母＋少数股东权益合计。
    equity_avg= df['equities_parent_company_owners'].mean()
    data = pd.concat([adj_net_profit, equity_avg], axis=1, join='inner')
    
    data.columns = ['adj_net_profit', 'equity_avg']

    ##归属于母公司股东的净利润*2/（期初归属于母公司股东的净资产+期末归属于母公司股东的净资产）
    roe_ttm = data['adj_net_profit'].sum()/data['equity_avg'].sum() * 100
    return round(roe_ttm, 4)

df = count_industry()
df