import numpy as np
import pandas as pd
from scipy.optimize import minimize


# 风险预算优化
def calculate_portfolio_var(w, S):
    # 计算组合风险的函数
    w = np.mat(w)
    return w * S * w.T


def calculate_risk_contribution(w, S):
    # 计算单个资产对总体风险贡献度的函数
    w = np.mat(w)
    sigma = np.sqrt(calculate_portfolio_var(w, S))
    # 边际风险贡献
    MRC = S * w.T / sigma
    # 风险贡献
    RC = np.multiply(MRC, w.T)
    return RC


def risk_budget_objective(w, pars):
    # 计算组合风险
    S = pars[0]  # 协方差矩阵
    risk_target = pars[1]  # 组合中资产预期风险贡献度的目标向量
    sig_p = np.sqrt(calculate_portfolio_var(w, S))  # portfolio sigma
    risk_target = np.asmatrix(np.multiply(sig_p, risk_target))
    asset_RC = calculate_risk_contribution(w, S)
    J = sum(np.square(asset_RC - risk_target.T))[0, 0]  # sum of squared error
    return J


def total_weight_constraint(x):
    return np.sum(x) - 1.0


def long_only_constraint(x):
    return x


# 根据资产预期目标风险贡献度来计算各资产的权重
def calcu_w(w_origin, S, risk_target):
    # w0 = [0.2, 0.2, 0.2, 0.6]
    #     x_t = [0.25, 0.25, 0.25, 0.25] # 目标是让四个资产风险贡献度相等，即都为25%
    cons = ({'type': 'eq', 'fun': total_weight_constraint},)
            # {'type': 'ineq', 'fun': long_only_constraint})
    res = minimize(risk_budget_objective, w_origin, args=[S, risk_target], method='SLSQP',
                   bounds=[[0,1]] * len(w_origin),
                   constraints=cons, options={'disp': True})
    w_rb = np.asmatrix(res.x)
    return w_rb
    # return res.x


# 将各资产风险贡献度绘制成柱状图
def plot_rc(w, S):
    rc = calculate_risk_contribution(w, S)
    print(rc)


if __name__ == '__main__':
    P = pd.read_csv("tests/resources/stock_prices.csv", parse_dates=True, index_col="date")
    w_origin, S, mu = optim_drawdown(P, 0.6, [0.001, 1])
    S = np.asmatrix(S)
    w_origin = np.asarray([i for i in w_origin.values()])
    risk_target = np.asarray([1/len(w_origin)] * len(w_origin))
    print(risk_target)
    w_rb = calcu_w(w_origin, S, risk_target)
    print('各资产权重：', w_rb)
    print(w_rb @ mu)
    plot_rc(w_rb)