# 克隆自聚宽文章：https://www.joinquant.com/post/61157
# 标题：致敬市场(13)——K-means 指数增强
# 作者：Gyro^.^

import numpy as np
import pandas as pd
import datetime as dt
from sklearn.cluster import KMeans

def initialize(context):
    # setting system
    log.set_level('order', 'error')
    set_option('use_real_price', True)
    set_option('avoid_future_data', True)
    run_monthly(iUpdate, 1, 'before_open')
    run_daily(iTrader, 'every_bar')
    run_daily(iReport, 'after_close')

def iUpdate(context):
    # paramers
    nday = 243
    n_cluster = 24
    n_class = 5
    # all funds
    dt_now = context.current_dt.date()
    all_fund = get_all_securities('fund', dt_now)
    # onlist 1-year
    dt_1y = dt_now - dt.timedelta(days=365)
    funds = all_fund[all_fund.start_date < dt_1y].index.tolist()
    # filter, liquity
    hm = history(nday, '1d', 'money', funds).mean()
    funds = hm[hm > 1e6].index.tolist()
    # history return
    h = history(nday, '1d', 'close', funds).dropna(axis=1)
    r = np.log(h).diff()[1:]
    # annual return
    ar = 100*nday*r.mean()
    # K-means
    cluster = KMeans(n_clusters=n_cluster).fit(r.T)
    # labels
    c = pd.Series(cluster.labels_, r.columns)
    # choice
    df = pd.DataFrame(columns=['Cluster', 'Return', 'Name'])
    for k in range(n_cluster):
        # class
        choice = []
        for s in c.index:
            if c[s] == k:
                choice.append(s) 
        log.info(k, len(choice))
        # high return
        cr = ar[choice].sort_values(ascending=False).head(n_class)
        for s in cr.index:
            if s in context.portfolio.positions:
                df.loc[s] = [k, cr[s], get_security_info(s).display_name]
                break
        else:
            s = cr.index[0]
            df.loc[s] = [k, cr[s], get_security_info(s).display_name]
    # report
    pd.set_option('display.max_rows', None)
    log.info('Funds', len(df), '\n', df)
    # results
    g.funds = df.index.tolist()
    g.position_size = 0.95/n_cluster * context.portfolio.total_value

def iTrader(context):
    # load data
    choice  = g.funds
    position_size = g.position_size
    lm_value = 0.7*position_size
    hm_value = 1.5*position_size
    cash_size = 0.05*context.portfolio.total_value
    cdata = get_current_data()
    # sell
    for s in context.portfolio.positions:
        if cdata[s].paused or \
            cdata[s].last_price >= cdata[s].high_limit or cdata[s].last_price <= cdata[s].low_limit:
            continue # 过滤三停
        if s not in choice:
            log.info('sell', s, cdata[s].name)
            order_target(s, 0, MarketOrderStyle(0.99*cdata[s].last_price))
    # buy
    for s in choice:
        if cdata[s].paused or \
            cdata[s].last_price >= cdata[s].high_limit or cdata[s].last_price <= cdata[s].low_limit:
            continue # 过滤三停
        if context.portfolio.available_cash < cash_size:
            break # 现金耗尽
        if s not in context.portfolio.positions:
            log.info('buy', s, cdata[s].name)
            order_target_value(s, position_size, MarketOrderStyle(1.01*cdata[s].last_price))

def iReport(context):
    # table of positions
    ptable = pd.DataFrame(columns=['amount', 'value', 'weight', 'name'])
    tvalue = context.portfolio.total_value
    cdata = get_current_data()
    for s in context.portfolio.positions:
        ps = context.portfolio.positions[s]
        ptable.loc[s] = [ps.total_amount, int(ps.value), 100*ps.value/tvalue, cdata[s].name]
    ptable = ptable.sort_values(by='weight', ascending=False)
    # report portfolio
    pd.set_option('display.max_rows', None)
    log.info('  positions', len(ptable), '\n', ptable.head())
    log.info('  total value %.2f, cash %.2f', \
        context.portfolio.total_value/10000, context.portfolio.available_cash/10000)
# end