import numpy as np
import pandas as pd
import math
import matplotlib.pyplot as plt

#模拟生成两个的权重
def weight(n):
    w = np.random.random(n)
    return w/sum(w)
#定义计算收益、风险、夏普率的函数
def portfolio_test(w):
    r1,r2 = 0.08,0.12
    sigma1,sigma2 = 0.12,0.25
    rho1,rho2,rho3,rho4,rho5 = -1,-0.5,0,0.5,1
    p_mean = r1*w[0]+r2*w[1]
    p_var1 = (w[0]**2)*(sigma1**2)+(w[1]**2)*(sigma2**2)+2*w[0]*w[1]*rho1*sigma1*sigma2
    p_var2 = (w[0]**2)*(sigma1**2)+(w[1]**2)*(sigma2**2)+2*w[0]*w[1]*rho2*sigma1*sigma2
    p_var3 = (w[0]**2)*(sigma1**2)+(w[1]**2)*(sigma2**2)+2*w[0]*w[1]*rho3*sigma1*sigma2
    p_var4 = (w[0]**2)*(sigma1**2)+(w[1]**2)*(sigma2**2)+2*w[0]*w[1]*rho4*sigma1*sigma2
    p_var5 = (w[0]**2)*(sigma1**2)+(w[1]**2)*(sigma2**2)+2*w[0]*w[1]*rho5*sigma1*sigma2
    p_sigma1 = np.sqrt(p_var1)
    p_sigma2= np.sqrt(p_var2)
    p_sigma3 = np.sqrt(p_var3)
    p_sigma4 = np.sqrt(p_var4)
    p_sigma5 = np.sqrt(p_var5)
    return w[0],w[1],p_mean,p_sigma1,p_sigma2,p_sigma3,p_sigma4,p_sigma5

df_value =np.column_stack([portfolio_test(weight(2)) for i in range(10)]) #产生随机组合
df = pd.DataFrame(df_value.T,columns = ['A资产比例','B资产比例','r_期望收益','p_rho1','p_rho2','p_rho3','p_rho4','p_rho5'])

df

df[['p_rho1','p_rho2','p_rho3','p_rho4','p_rho5']].plot()

df_value1 =np.column_stack([portfolio_test(weight(2)) for i in range(200)]) #产生随机组合
df1= pd.DataFrame(df_value1.T,columns = ['A资产比例','B资产比例','p_mean','p_rho1','p_rho2','p_rho3','p_rho4','p_rho5'])
plt.scatter(df1['p_rho1'],df1['p_mean'],marker = 'o')
plt.scatter(df1['p_rho2'],df1['p_mean'],marker = 'o')
plt.scatter(df1['p_rho3'],df1['p_mean'],marker = 'o')
plt.scatter(df1['p_rho4'],df1['p_mean'],marker = 'o')
plt.scatter(df1['p_rho5'],df1['p_mean'],marker = 'o')

#选出几个标的数据，进行演示计算
#这里选取HS300里面近一年涨幅大于4%的5只股票
stock_list = get_index_stocks('000300.XSHG')
pool = []
for i in stock_list:
    price_df = get_price(i,start_date='2017-9-1',end_date='2018-9-1',fields=['close'])
    pct = price_df['close'].values[-1]/price_df['close'].values[0]
    if pct>1.04:
        pool.append(i)
    if len(pool)==5:
        break
pool

#数据准备
#计算股票池的涨跌幅、回报率
def get_return(pool):
    prices = get_price(list(pool), end_date='2017-9-1', frequency='1d', fields=['close'],count=252).close
    returns = (prices/prices.shift(1)).iloc[1:]
    rets = returns.dropna(axis=1)
    return rets
#随机权重设置
def weight(n):
    w = np.random.random(n)
    return w/sum(w)
#获取回报值
returns = get_return(pool)     

np.log(returns).mean()*252

#输入收益returns,权重
#返回年化收益、标准差、夏普率
def portfolio(returns,w):
    #无风险收益年化设置为0.04
    r_b = 0.04
    returns = np.log(returns)
    r_mean = returns.mean()*252
    p_mean = np.sum(r_mean*w)
    r_cov = returns.cov()*252
    p_var = np.dot(w.T,np.dot(r_cov,w))
    p_std = np.sqrt(p_var)
    p_sharpe = (p_mean-r_b)/p_std
    return p_mean,p_std,p_sharpe

p_mean,p_std,p_sharpe = np.column_stack([portfolio(returns,weight(5)) for i in range(2000)]) #产生随机组合
plt.figure(figsize = (14,8))
plt.scatter(p_std, p_mean, c=p_sharpe, marker = 'o')
plt.grid(True)   #显示网格
plt.xlabel('std')
plt.ylabel('mean')
plt.colorbar(label = 'Sharpe')
plt.title('Mean and Std of Returns')

from scipy.optimize import minimize 

#最小方差函数
def min_variance(w):
    return portfolio(returns,w)[1]**2
#最大夏普函数（转化为求最小值）
def min_sharpe(w):
    return -portfolio(returns,w)[2]
#约束条件，权重总和为1
cons = ({'type':'eq', 'fun':lambda x: np.sum(x)-1})
#夏普率最大
opt_sharpe = minimize(min_sharpe,weight(5),bounds=((0,1),(0,1),(0,1),(0,1),(0,1)),constraints = cons)
#方差最小
opt_var = minimize(min_variance,weight(5),bounds=((0,1),(0,1),(0,1),(0,1),(0,1)),constraints = cons)

print("Sharpe最大时的组合最优解是：")
print "组合仓位比例：",'%s' % (opt_sharpe['x'].round(3))
print('组合收益：%s' % (portfolio(returns,opt_sharpe['x'].round(3))[0]))
print('组合标准差（波动率）：%s' % (portfolio(returns,opt_sharpe['x'].round(3))[1]) )
print('组合sharpe:%s' % (portfolio(returns,opt_sharpe['x'].round(3))[2]) )

plt.figure(figsize = (14,8))
plt.scatter(p_std, p_mean, c=p_sharpe, marker = 'o')

#红星：标记最高sharpe组合
plt.plot(portfolio(returns,opt_sharpe['x'])[1],portfolio(returns,opt_sharpe['x'])[0], 'r*', markersize = 15.0)

#黄星：标记最小方差组合
plt.plot(portfolio(returns,opt_var['x'])[1], portfolio(returns,opt_var['x'])[0], 'y*', markersize = 15.0)

plt.grid(True)   #显示网格
plt.xlabel('std')
plt.ylabel('mean')
plt.colorbar(label = 'Sharpe')
plt.title('Mean and Std of Returns')

#有效前沿
#最小方差下的收益率
r_min = portfolio(returns,opt_var['x'])[0]
r_mean_list = np.linspace(r_min,r_min*1.7,20)
v_list = []
for r in r_mean_list:    
    cons1 = ({'type':'eq','fun':lambda w:portfolio(returns,w)[0]-r},{'type':'eq','fun':lambda w:np.sum(w)-1})
    opt_var_1 = minimize(min_variance,weight(5),bounds=((0,1),(0,1),(0,1),(0,1),(0,1)),constraints = cons1)
    #最小方差下的收益率
    v_min = portfolio(returns,opt_var_1['x'])[1]
    print v_min
    v_list.append(v_min)

plt.figure(figsize = (14,8))
plt.scatter(p_std, p_mean, c=p_sharpe, marker = 'o')

#红星：标记最高sharpe组合
plt.plot(portfolio(returns,opt_sharpe['x'])[1],portfolio(returns,opt_sharpe['x'])[0], 'r*', markersize = 15.0)

#黄星：标记最小方差组合
plt.plot(portfolio(returns,opt_var['x'])[1], portfolio(returns,opt_var['x'])[0], 'y*', markersize = 15.0)

#标记有效前沿
plt.plot(v_list,r_mean_list, 'b*', markersize = 5.0)

plt.grid(True)   #显示网格
plt.xlabel('std')
plt.ylabel('mean')
plt.colorbar(label = 'Sharpe')
plt.title('Mean and Std of Returns')

#最小方差下的收益率
r_min = portfolio(returns,opt_var['x'])[0]
r_min

#输入日期
#返回HS300中股票收益超过4%的前5只股票
def get_pool(date):
    stock_list = get_index_stocks('000300.XSHG',date=date)
    pool = []
    for i in stock_list:
        price_df = get_price(i,end_date=date,fields=['close'],count=252)
        pct = price_df['close'].values[-1]/price_df['close'].values[0]
        if pct>1.04:
            pool.append(i)
        if len(pool)==5:
            break
    return pool

#输入股票列表
#返回一个组合的权重配比
def opt_func(pool,date):
    prices = get_price(list(pool), end_date=date, frequency='1d', fields=['close'],count=252).close
    returns = (prices/prices.shift(1)).iloc[1:]
    returns = returns.dropna(axis=1)

    def portfolio(returns,w):
        #无风险收益年化设置为0.04
        r_b = 0.04
        returns = np.log(returns)
        r_mean = returns.mean()*252
        p_mean = np.sum(r_mean*w)
        r_cov = returns.cov()*252
        p_var = np.dot(w.T,np.dot(r_cov,w))
        p_std = np.sqrt(p_var)
        p_sharpe = (p_mean-r_b)/p_std
        return p_mean,p_std,p_sharpe
    
    def weight(n):
        w = np.random.random(n)
        return w/sum(w)
    
    #最大夏普函数（转化为求最小值）
    def min_sharpe(w):
        return -portfolio(returns,w)[2]
    #约束条件，权重总和为1
    cons = ({'type':'eq', 'fun':lambda x: np.sum(x)-1})
    #夏普率最大
    opt_sharpe = minimize(min_sharpe,weight(5),bounds=((0,1),(0,1),(0,1),(0,1),(0,1)),constraints = cons)
    return opt_sharpe['x'].round(3)

date = '2018-1-1'
pool = get_pool(date)
res = opt_func(pool,date).round(3)
res

#工具函数
#获取日期列表
def get_tradeday_list(s='2017-1-1',e='2018-7-25',label='month',count=None):
    if count != None:
        df = get_price('000001.XSHG',end_date=e,count=count)
    else:
        df = get_price('000001.XSHG',start_date=s,end_date=e)
    if label==None:
        return df.index
    else:
        df['year-month'] = [str(i)[0:7] for i in df.index]
        if label == 'month':
            return df.drop_duplicates('year-month').index
        elif label == 'quarter':
            df['month'] = [str(i)[5:7] for i in df.index]
            df = df[(df['month']=='01') | (df['month']=='04') | (df['month']=='07') | (df['month']=='10') ]
            return df.drop_duplicates('year-month').index
        elif label =='halfyear':
            df['month'] = [str(i)[5:7] for i in df.index]
            df = df[(df['month']=='01') | (df['month']=='06')]
            return df.drop_duplicates('year-month').index
#获取股票池的组合收益(不进行市值加权)
def ret_se(sl,start_date='2018-6-1',end_date='2018-7-1',w=[1.0/len(sl)]*len(sl)):
    if len(sl) != 0:
        #得到股票的历史价格数据
        df = get_price(list(sl),start_date=start_date,end_date=end_date,fields=['close']).close
        df = df.dropna(axis=1)
        #获取列表中的股票流通市值对数值
        df_mkt = get_fundamentals(query(valuation.code,valuation.circulating_market_cap).filter(valuation.code.in_(df.columns)))
        df_mkt.index = df_mkt['code'].values
        fact_se =pd.Series(df_mkt['circulating_market_cap'].values,index = df_mkt['code'].values)
        fact_se = np.log(fact_se)
    else:
        df = get_price('000001.XSHG',start_date=start_date,end_date=end_date,fields=['close'])
        df['v'] = [1]*len(df)
        del df['close']
    #相当于昨天的百分比变化
    pct = df.pct_change()+1
    pct.iloc[0,:] = 1
    #df1 = pct.cumprod()
    #等权重平均收益结果
    #按权重的方式计算
    se1 = (pct*w).cumsum(axis=1).iloc[:,-1]
    se2 = pct.cumsum(axis=1).iloc[:,-1]/pct.shape[1]
    return se1,se2
def f_sum(x):
    return sum(x)

#设置回测时间
start,end = '2016-9-1','2018-9-1'
#获取交易日列表
trade_list = get_tradeday_list(start,end,'month')
num = 1

for s,e in zip(trade_list[:-1],trade_list[1:]): 
    #获取满足条件的股票列表
    pool = get_pool(s)   
    #获取优化权重
    w = opt_func(pool,s)
    if num == 0:
        df_returns1 = ret_se(pool,s,e,w)[0] 
        df_a1 = ret_se(pool,s,e,w)[1] 
        df_returns = pd.concat([df_returns,df_returns1])
        df_a = pd.concat([df_a,df_a1])
    else:
        df_returns = ret_se(pool,s,e,w)[0]
        df_a = ret_se(pool,s,e,w)[1]
        num = 0
#进行图表展示
df_returns.cumprod().plot(figsize=(12,7),legend=True,label='组合优化收益',title='图2 组合收益走势')
df_a.cumprod().plot(figsize=(12,7),legend=True,label='平均分仓收益')

#计算总收益率
total_return1,total_return2 = df_returns.cumprod().values[-1]-1,df_a.cumprod().values[-1]-1
#计算年化收益
total_a_r1,total_a_r2 = (1+total_return1)**(250.0/len(df_returns))-1,(1+total_return2)**(250.0/len(df_a))-1
#计算夏普率
sharpe1,sharpe2 = (total_a_r1-0.04)/np.std(df_returns.cumprod()),(total_a_r2-0.04)/np.std(df_a.cumprod())

df_summary = pd.DataFrame([],index=['均值方差优化组合','平均分仓组合'])
df_summary['总收益'] = [total_return1,total_return2]
df_summary['年化收益'] = [total_a_r1,total_a_r2]
df_summary['sharpe'] = [sharpe1,sharpe2]
df_summary