import numpy as np
import matplotlib.pyplot as plt
import cvxopt as opt
from cvxopt import blas, solvers
import pandas as pd

def optimal_portfolio(returns):
    n = len(returns)
    returns = np.asmatrix(returns)
    
    N = 100
    mus = [10**(5.0 * t/N - 1.0) for t in range(N)]
    
    # 转化为cvxopt matrices
    S = opt.matrix(np.cov(returns))
    pbar = opt.matrix(np.mean(returns, axis=1))
    
    # 约束条件
    G = -opt.matrix(np.eye(n))   # opt默认是求最大值，因此要求最小化问题，还得乘以一个负号
    h = opt.matrix(0.0, (n ,1))
    A = opt.matrix(1.0, (1, n))
    b = opt.matrix(1.0)
    
    # 使用凸优化计算有效前沿
    portfolios = [solvers.qp(mu*S, -pbar, G, h, A, b)['x'] 
                  for mu in mus]
    ## 计算有效前沿的收益率和风险
    returns = [blas.dot(pbar, x) for x in portfolios]
    risks = [np.sqrt(blas.dot(x, S*x)) for x in portfolios]
    m1 = np.polyfit(returns, risks, 2)
    x1 = np.sqrt(m1[2] / m1[0])
    # 计算最优组合
    wt = solvers.qp(opt.matrix(x1 * S), -pbar, G, h, A, b)['x']
    return np.asarray(wt), returns, risks

#weights, returns, risks = optimal_portfolio(return_vec)

# 获取数据
start_date = '2017-02-01'
end_date = '2017-12-10'
instruments = ['600000.SHA','600016.SHA','600019.SHA','600028.SHA','600029.SHA','600030.SHA','600036.SHA','600048.SHA','600050.SHA','600104.SHA',
               '600111.SHA','600309.SHA','600340.SHA','600518.SHA','600519.SHA','600547.SHA','600606.SHA','600837.SHA','600887.SHA','600919.SHA',
               '600958.SHA','600999.SHA','601006.SHA','601088.SHA','601166.SHA','601169.SHA','601186.SHA','601211.SHA','601229.SHA','601288.SHA',
               '601318.SHA','601328.SHA','601336.SHA','601390.SHA','601398.SHA','601601.SHA','601628.SHA','601668.SHA','601669.SHA','601688.SHA',
               '601766.SHA','601800.SHA','601818.SHA','601857.SHA','601878.SHA','601881.SHA','601985.SHA','601988.SHA','601989.SHA','603993.SHA']
   
data = D.history_data(instruments,start_date,end_date,
                     fields=['close'])
# 整理数据
data = pd.pivot_table(data,values='close',index=['date'],columns=['instrument'])
#T.plot(data)

def initialize(context):
    
    context.days = 0
    context.ins = instruments
     
def handle_data(context, data):
    
    context.days += 1
    if context.days < 40:
        return
    # 每60天调仓一次
    if context.days % 30 != 0:
        return
    # 获取数据的时间窗口并计算收益率
    prices = data.history(context.symbols(context.ins[0],context.ins[1],context.ins[2],context.ins[3],context.ins[4]), 'price',40, '1d').dropna()
    returns = prices.pct_change().dropna()
    try:
        # 马科维茨组合优化
        weights, _, _ = optimal_portfolio(returns.T)
        print(weights)
        # 对持仓进行权重调整
        for stock, weight in zip(prices.columns, weights):
            if data.can_trade(stock):
                order_target_percent(stock, weight[0])
    except ValueError as e:
        pass

m = M.trade.v2(
    instruments=instruments,
    start_date=start_date,
    end_date=end_date,
    initialize=initialize,
    handle_data=handle_data,
    # 买入订单以开盘价成交
    order_price_field_buy='open',
    # 卖出订单以开盘价成交
    order_price_field_sell='open',
    capital_base=1000000,
    benchmark='000016.SHA',
)