import numpy as np
import pandas as pd
%matplotlib inline
import matplotlib.pyplot as plt
plt.style.use('ggplot')

"""
Explanation: <h1 align="center">ADMM for Quadratic Programming</h1>

In this notebook, we implement the Alternating Direction Method of Multipliers algorithm for solving a standard-form Quadratic Program.  The ADMM approach solves convex optimization problems by breaking them into smaller pieces, each of which are then easier to handle.
The standard form QP is:
Minimize $\dfrac {1} {2} x^T P x + q^T x + r$ subject to $lb \le x \le ub$.
End of explanation
"""


from numpy.linalg import inv, norm

def objective(P, q, r, x):
    """Return the value of the Standard form QP using the current value of x."""
    return 0.5 * np.dot(x, np.dot(P, x)) + np.dot(q, x) + r


def qp_admm(P, q, r, lb, ub,
            max_iter=1000,
            rho=1.0, 
            alpha=1.2,              
            atol=1e-4, 
            rtol=1e-2):

    n = P.shape[0]
    
    x = np.zeros(n)
    z = np.zeros(n)
    u = np.zeros(n)
    
    history = []

    R = inv(P + rho * np.eye(n))
    
    for k in range(1, max_iter+1):
        x = np.dot(R, (z - u) - q)
        
        # z-update with relaxation
        z_old = z
        x_hat = alpha * x +(1 - alpha) * z_old
        z = np.minimum(ub, np.maximum(lb, x_hat + u))

        # u-update
        u = u + (x_hat - z)

        # diagnostics, and termination checks
        objval = objective(P, q, r, x)

        r_norm = norm(x - z)
        s_norm = norm(-rho * (z - z_old))
        eps_pri = np.sqrt(n) * atol + rtol * np.maximum(norm(x), norm(-z))
        eps_dual = np.sqrt(n)* atol + rtol * norm(rho*u)
        
        history.append({
            'objval'  : objval, 
            'r_norm'  : r_norm, 
            's_norm'  : s_norm,
            'eps_pri' : eps_pri,
            'eps_dual': eps_dual,
        })
        
        if r_norm < eps_pri and s_norm < eps_dual:
            print('Optimization terminated after {} iterations'.format(k))
            break;
        
    history = pd.DataFrame(history)
    return x, history


"""
Explanation: Pure-Python ADMM Implementation
The code below is a direct Python port of the reference MATLAB implementation in Reference[1].
End of explanation
"""


import cvxpy as cvx

def qp_cvxpy(P, q, r, lb, ub,
            max_iter=1000,
            atol=1e-4, 
            rtol=1e-2):
    n = P.shape[0]
    
    # The variable we want to solve for
    x = cvx.Variable(n)
    constraints = [x >= cvx.Constant(lb), x <= cvx.Constant(ub)]
    
    # Construct the QP expression using CVX Primitives
    # Note that in the CVX-meta language '*' of vectors of matrices indicates dot product, 
    # not elementwise multiplication 
    expr = cvx.Constant(0.5) * cvx.quad_form(x, cvx.Constant(P)) + cvx.Constant(q) * x + cvx.Constant(r)
    qp = cvx.Problem(cvx.Minimize(expr), constraints=constraints)
    qp.solve(max_iters=max_iter, abstol=atol, reltol=rtol, verbose=True)    
    
    # The result is a Matrix object. Make it an NDArray and drop of 2nd dimension i.e. make it a vector.
    x_opt = np.array(x.value).squeeze()
    return x_opt

"""
Explanation: QP Solver using CVXPY
For comparison, we also implement QP solver using cvxpy.
End of explanation
"""


import requests
from statsmodels.stats.moment_helpers import corr2cov
from functools import lru_cache

@lru_cache(maxsize=5)
def get_cov(filename):
    url = r'http://people.brunel.ac.uk/~mastjjb/jeb/orlib/files/{}.txt'.format(filename)
    data = requests.get(url).text
    lines = [line.strip() for line in data.split('\n')]

    # First line is the number of assets
    n_assets = int(lines[0])
    
    # Next n_assets lines contain the space separated mean and stddev. of returns for each asset
    means_and_sds = pd.DataFrame(
        data=np.nan, 
        index=range(0, n_assets), 
        columns=['ret_mean', 'ret_std'])

    # Next n_assetsC2 lines contain the 1-based row and column index and the corresponding correlation
    for i in range(0, n_assets):
        mean, sd = map(float, lines[1+i].split())
        means_and_sds.loc[i, ['ret_mean', 'ret_std']] = [mean, sd]

    n_corrs = (n_assets * (n_assets + 1)) // 2
    corrs = pd.DataFrame(index=range(n_assets), columns=range(n_assets), data=np.nan)

    for i in range(0, n_corrs):
        row, col, corr = lines[n_assets + 1 + i].split()
        row, col = int(row)-1, int(col)-1
        corr = float(corr)
        corrs.loc[row, col] = corr
        corrs.loc[col, row] = corr
    
    cov = corr2cov(corrs, means_and_sds.ret_std)
    return cov

"""
Explanation: Generate Optimal Portfolio Holdings
In this section, we define a helper function to load the one of the five asset returns datasets from OR library (Reference [2]). The data are available by requesting filenames port[1-5]. Each file contains a progressively larger set of asset returns, standard deviations of returns and correlations of returns.
End of explanation
"""


from numpy.random import RandomState
rng = RandomState(0)
P = get_cov('port1')
n = P.shape[0]
alphas = rng.uniform(-0.4, 0.4, size=n) 
q = -alphas
ub = np.ones_like(q)
lb = np.zeros_like(q)
r = 0

"""
Explanation: Set up the Portfolio Optimization problem as a QP
End of explanation
"""


%%time
x_opt_admm, history = qp_admm(P, q, r, lb, ub)

fig, ax = plt.subplots(history.shape[1], 1, figsize=(10, 8))
ax = history.plot(subplots=True, ax=ax, rot=0)

"""
Explanation: Using ADMM
End of explanation
"""


%%time
x_opt_cvxpy = qp_cvxpy(P, q, r, lb, ub)

"""
Explanation: Using CVXPY
End of explanation
"""


holdings = pd.DataFrame(np.column_stack([x_opt_admm, x_opt_cvxpy]), columns=['opt_admm', 'opt_cvxpy'])
fig, ax = plt.subplots(1, 1, figsize=(12, 4))
ax = holdings.plot(kind='bar', ax=ax, rot=0)
labels = ax.set(xlabel='Assets', ylabel='Holdings')

"""
Explanation: Optimal Holdings Comparison
End of explanation
"""


