import cvxpy
import numpy as np
import math
from matplotlib import pyplot as plt
plt.style.use('ggplot')

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

def compute_allocation(rho_clo = 0.9, rho_cso=0.6, rho_subprime=0.2,
                       delta_clo=1.2, delta_cso=0.4, delta_subprime=0.8,
                       mu_HY=0.02, mu_clo=0.08, mu_cso=0.07, mu_subprime=0.25):
    rho = {'CLO': rho_clo,
           'CSO': rho_cso,
           'Subprime': rho_subprime}
    delta = {'CLO': delta_clo,
             'CSO': delta_cso,
             'Subprime': delta_subprime}
    assets = ['CLO', 'CSO', 'Subprime']
    mu = np.array([mu_HY, mu_clo, mu_cso, mu_subprime])
    u = volHY * np.array([delta[a] for a in assets])
    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[a] for a in assets])
    Sigma = np.vstack((v, np.c_[v[1:], Sigma]))
    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])

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

    fund_return = mu@W
    fund_vol= np.array([math.sqrt(W[:,i]@Sigma@W[:,i]) for i in range(gamma_x.size)])
    return (W, fund_return, fund_vol)

def plot_allocation(W, fund_return, fund_vol):
    gamma_x = np.linspace(0, 20, fund_return.size)
    fig, ax1 = plt.subplots()
    ax1.stackplot(fund_vol, W[1:,], labels=['CLO', 'CSO', 'Subprime'])
    ax1.set_xlabel('risk factor')
    ax1.set_ylabel('portfolio weights')
    ax1.legend()
    # ax1.text(0.3, 0.82, 'RMBS')
    # ax1.text(0.5, 0.45, 'CSO')
    # ax1.text(0.5, 0.15, 'CLO')
    ax1.set_ylim([0, 1])
    ax2 = ax1.twinx()
    ax2.plot(fund_vol, fund_return, lw=1, color="grey")
    ax2.set_ylabel('fund volatility')
    plt.show()

if __name__=="__main__":
    volHY = 0.07

    rho = {'CLO': 0.6,
           'CSO': 0.5,
           'Subprime': 0.3}
    delta = {'CLO': 0.4,
             'CSO': 0.2,
             'Subprime': 0.6}
    mu = np.array([0.01, 0.075, 0.065, 0.25])

    W, fund_return, fund_vol = compute_allocation(rho['CLO'], rho['CSO'], rho['Subprime'],
                                                  delta['CLO'], delta['CSO'], delta['Subprime'],
                                                  mu[0], mu[1], mu[2], mu[3])
    plot_allocation(W, fund_return, fund_vol)