import cvxpy
import numpy as np
import math

def cor2cov(Rho, vol):
    return np.diag(vol) @ Rho @ np.diag(vol)

def rho(sigma, delta, volF):
    """ computes the correlation between the asset and the factor """
    return 1/math.sqrt(1+sigma**2/(delta**2*volF**2))

def resid_vol(rho, delta, volF):
    """ computes the residual of the asset """
    return math.sqrt(delta**2*volF**2*(1/rho**2-1))

def var(rho, delta, volF):
    """ computes the variance of the asset """
    return delta**2*volF**2+resid_vol(rho, delta, volF)**2

volHY = 0.4
rho = {'CLO': 0.9,
       'CSO': 0.6,
       'Subprime': 0.4}
delta = {'CLO': 1.5,
         'CSO': 0.4,
         'Subprime': 1}

u = volHY * np.array([delta['CLO'], delta['CSO'], delta['Subprime']])
Sigma = np.outer(u, u) + np.diag([resid_vol(rho[a], delta[a], volHY)**2
                                  for a in ['CLO', 'CSO', 'Subprime']])
v = volHY**2 * np.array([1, delta['CLO'], delta['CSO'], delta['Subprime']])
Sigma = np.vstack((v, np.c_[v[1:], Sigma]))

mu = np.array([0.03, 0.07, 0.04, 0.15])
sharpe = mu/np.sqrt(np.diag(Sigma))

gamma = cvxpy.Parameter(sign='positive')
w = cvxpy.Variable(4)
ret = mu.T*w
risk = cvxpy.quad_form(w, Sigma)
prob = cvxpy.Problem(cvxpy.Maximize(ret-gamma*risk),
               [cvxpy.sum_entries(w[1:]) - 0.1*w[0] == 1,
                w[1:] >= 0,
                w[0] <= 0])

W = np.empty((4, 100))
gamma_x = np.linspace(0, 1, 100)
for i, val in enumerate(gamma_x):
    gamma.value = val
    prob.solve()
    W[:,i] = np.asarray(w.value).squeeze()

fund_return = mu@X
fund_vol= np.array([math.sqrt(X[:,i]@Sigma@X[:,i]) for i in range(100)])

from matplotlib import pyplot as plt
plt.style.use('ggplot')
fig, ax1 = plt.subplots()
ax1.stackplot(gamma_x, W[1:,])
ax1.set_xlabel('risk factor')
ax1.set_ylabel('portfolio weights')
ax1.text(0.3, 0.82, 'RMBS')
ax1.text(0.5, 0.45, 'CSO')
ax1.text(0.5, 0.15, 'CLO')
ax2 = ax1.twinx()
ax2.plot(gamma_x, fund_vol, lw=1)
ax2.set_ylabel('fund volatility')
plt.show()