"""
title:自编指数研究 v1.1 by 小宽(QQ:1026070)
description: 根据给定的股池, 按深交所指数编制算法, 计算出该股池的每日指数myindex, 再推算出ROC(20)值和是否清仓标识
    也可通过编写 getter 函数, 动态获取符合条件的股池
lastModified: 2016-12-09 12:50:00
"""
from jqdata import * # get_trade_days
import pandas as pd
import datetime as dt
#start_date = dt.datetime.strptime('2016-05-01', '%Y-%m-%d').date()
#trdays = get_all_trade_days() # 全部交易日
trdays = get_trade_days(count=365) #取最近365个交易日
#print trdays
ORIGIN_DATE = trdays[0] #起点交易日

#print '全部交易日总数: %d' % len(trdays)

# 获取多日成分股的流通市值累计值
def get_sumcaps(stock_list=None, periods=50, startdate=None, getter=None):
    #import datetime
    # 用 stock_list 传入静态成分股池, 或用 getter 函数动态获取, 两者必须有一
    if (stock_list is None or len(stock_list) == 0) and getter is None:
        print '[Error]get_mcaps: getter function or stock_list is required'
        return None
    else:
        stocks = stock_list
    result = pd.DataFrame(columns=['date','sumcap'])#, index=['date'])
    #daterange = pd.date_range(start_date, periods=100)
    #for dt in reversed(daterange):
    if startdate is None:
        #print '[Warning] param "startdate" not specified, using today\'s date instead.'
        startdate = dt.date.today()
    i = 0
    adate = startdate
    # 从当前日期开始, 逐日倒退, 获取每日累计流通市值
    while i < periods and adate >= ORIGIN_DATE: #必须有个最小日期作为限制,否则可能陷入死循环
        adate = adate + dt.timedelta(-1) # 时间倒退一日
        if adate not in trdays: #忽略非交易日
            #print '*date: %s is not a trading day, skipped' % adate
            continue
        #TODO: 动态获取成分股列表
        if getter is not None:
            stocks = getter(adate, i, periods) # 可根据 date 获取该日符合条件的成分股
        
        df = get_fundamentals(query(
                valuation.code, valuation.circulating_market_cap.label('cmcap') #获取该日流通市值
            ).filter(valuation.code.in_(stocks)), date=adate)
        
        sumcap = df['cmcap'].sum() #所有成分股的流通市值累计
        stock_count = len(df['code']) # 成分股数量, 用于归一化
        if sumcap == 0: break # 没数据了, 说明到头了, 退出循环
        result = result.append({'date': adate, 'sumcap': sumcap, 'count': stock_count}, ignore_index=True)
        i += 1
        
    return result

#startdate = dt.datetime.strptime('2016-12-01', '%Y-%m-%d').date()
#startdate = dt.date.today()
#DEMO成分股池, 也可以用 getter 函数动态获取成分股
stock_list = ['300434.XSHE', '300483.XSHE', '002743.XSHE'] 
TOTAL = 50 #至少为 ROC 周期数 2 倍, 但数据越多并不影响每日ROC涨跌幅
df = get_sumcaps(stock_list, periods=TOTAL)

# 计算指数, 附加到新的 myindex 列
def calc_myindex(df, basis=1000.0):
    df['myindex'] = basis #初始化,全置1, 也可用1000作为起始点位
    df_sumcap = df['sumcap']
    df_count = df['count']
    # 自定义指数的算法参考深交所指数算法: 当日指数=前日指数*Σ(当日成分股流通市值)/Σ(前日成分股流通市值)
    for i in reversed(range(len(df)-1)): # 从倒数第2个开始滚动计算
        prev = i + 1
        v = df['myindex'][prev]*(df_sumcap[i]/df_count[i])/(df_sumcap[prev]/df_count[prev])
        df.ix[i, 'myindex']= v # df 修改某个单元值的正确方式!

    return df

df = calc_myindex(df)
#print '---------自编指数:myindex:-----------\n', df

roc20 = df['myindex'][0] / df['myindex'][20] - 1 # 计算ROC(20)
print '自编指数:myindex 结论: 当前 roc20=%f, 清仓=%s' % (roc20, roc20<=0)

# 计算指数的ROC值和是否清仓标识clear, 附加到新的列
def calc_myindex_roc(df, days=20):
    len_df = len(df)
    df['roc'] = numpy.nan
    df['clear'] = numpy.nan
    for i in range(len_df):
        if i + days >= len_df: # 超出最大范围就跳出循环
            break
        v = df['myindex'][i] / df['myindex'][i+days] - 1
        df.ix[i, 'roc'] = v # df 修改某个单元值的正确方式! #ROC值
        df.ix[i, 'clear'] = v <= 0 #是否清仓标识
    return df
    
df = calc_myindex_roc(df, 20) # ROC(20)
print '---------自编指数:myindex with ROC20:-----------\n', df