import cvxpy as cp
import numpy as np
import matplotlib.pyplot as plt

# d. Solve (RO.4)
np.random.seed(1)
n = 10
mu = np.abs(np.random.randn(n))     # n risky assets with the expected returns
# mu_hat = np.abs(np.random.randn(n))  # empirical expected returns
Sigma = np.random.randn(n, n)       # covariance of asset returns
Sigma = Sigma.T.dot(Sigma)
Sigma_mu = (1 / n) * np.diag(np.diag(Sigma))
lamb = cp.Parameter(nonneg=True)    # risk preference parameter of investor, non-negative
iota = np.ones((n, 1))              # vector of ones
eps = cp.Parameter(nonneg=True)
k = cp.Parameter(nonneg=True)

# w = cp.Variable((n, 1))      # 'investment proportion' or 'porfolio allocation vector'
W = cp.Variable((n, n), PSD=True)
ret = iota.T @ W @ mu
risk = cp.trace(Sigma @ W)
prob = cp.Problem(cp.Minimize(risk - lamb * ret),
                  [cp.trace(iota @ iota.T @ W) == 1,
                  cp.trace(Sigma_mu @ W) <= eps,
                  iota.T @ cp.abs(W) @ iota <= k * cp.trace(W)])

# Compute trade-off curve
kays = [3, 5, 7]
SAMPLES = 100
w_vec = np.zeros((n, SAMPLES))
risk_data = np.zeros((len(kays), SAMPLES))
ret_data = np.zeros((len(kays), SAMPLES))
lamb_vals = np.logspace(-2, 3, num=SAMPLES)

for idx, kay in enumerate(kays):
    for i in range(SAMPLES):
        eps.value = 1
        k.value = kays[idx]
        lamb.value = lamb_vals[i]
        prob.solve()
        eigval, eigvec = np.linalg.eig(W.value)
        max_idx = np.argmax(eigval)
        w_vec[:, i] = eigvec[:, max_idx]
        w_vec[:, i] = w_vec[:, i] / sum(w_vec[:, i])
        risk_data[idx, i] = cp.sqrt(cp.quad_form(w_vec[:, i], Sigma).value).value
        ret_data[idx, i] = mu.T @ w_vec[:, i]

# Plot trade-off curve
fig = plt.figure()
ax = fig.add_subplot(111)

for i in range(n):
    plt.plot(cp.sqrt(Sigma[i, i]).value, mu[i], 'ro')  # generated by randn

colors = ['r', 'g', 'b']
for idx, kay in enumerate(kays):
    plt.plot(risk_data[idx, :], ret_data[idx, :], linestyle='dashed', color=colors[idx],
             label=r"$k$ = %d, $\epsilon$ = 1" % kay)

markers_on = [40, 50, 60, 70]
markers_color = ['rs', 'gs', 'bs']

# plot lambda points
for idx in range(3):
    for marker in markers_on:
        plt.plot(risk_data[idx, marker], ret_data[idx, marker], markers_color[idx])
        ax.annotate(r"$\lambda = %.2f$" % lamb_vals[marker],
                    xy=(risk_data[idx, marker] + .08, ret_data[idx, marker] - .03))

plt.title('(RO.4)')
plt.xlabel('Risk (Standard deviation)')
plt.ylabel('Return')
plt.legend(loc='lower right')
plt.show()
