#!/usr/bin/env python
# coding: utf-8

# LaTeX macros (hidden cell)
# $
# \newcommand{\Q}{\mathcal{Q}}
# \newcommand{\ECov}{\boldsymbol{\Sigma}}
# \newcommand{\EMean}{\boldsymbol{\mu}}
# \newcommand{\EAlpha}{\boldsymbol{\alpha}}
# \newcommand{\EBeta}{\boldsymbol{\beta}}
# $

# # Imports and configuration

# In[29]:


import sys
import os
import re
import glob
import datetime as dt

import numpy as np
import pandas as pd
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

from mosek.fusion import *

from notebook.services.config import ConfigManager

from portfolio_tools import data_download, DataReader


# In[30]:


# Version checks
print(sys.version)
print('matplotlib: {}'.format(matplotlib.__version__))

# Jupyter configuration
c = ConfigManager()
c.update('notebook', {"CodeCell": {"cm_config": {"autoCloseBrackets": False}}})  

# Numpy options
np.set_printoptions(precision=5, linewidth=120, suppress=True)

# Pandas options
pd.set_option('display.max_rows', None)

# Matplotlib options
plt.rcParams['figure.figsize'] = [12, 8]
plt.rcParams['figure.dpi'] = 200


# # Prepare input data

# In this example, the input data is given. It consists of the vector $\EMean$ of expected returns, and the covariance matrix $\ECov$.

# In[31]:


# Linear return statistics on the investment horizon
mu = np.array([0.07197349, 0.15518171, 0.17535435, 0.0898094 , 0.42895777, 0.39291844, 0.32170722, 0.18378628])
Sigma = np.array([
        [0.09460323, 0.03735969, 0.03488376, 0.03483838, 0.05420885, 0.03682539, 0.03209623, 0.03271886],
        [0.03735969, 0.07746293, 0.03868215, 0.03670678, 0.03816653, 0.03634422, 0.0356449 , 0.03422235],
        [0.03488376, 0.03868215, 0.06241065, 0.03364444, 0.03949475, 0.03690811, 0.03383847, 0.02433733],
        [0.03483838, 0.03670678, 0.03364444, 0.06824955, 0.04017978, 0.03348263, 0.04360484, 0.03713009],
        [0.05420885, 0.03816653, 0.03949475, 0.04017978, 0.17243352, 0.07886889, 0.06999607, 0.05010711],
        [0.03682539, 0.03634422, 0.03690811, 0.03348263, 0.07886889, 0.09093307, 0.05364518, 0.04489357],
        [0.03209623, 0.0356449 , 0.03383847, 0.04360484, 0.06999607, 0.05364518, 0.09649728, 0.04419974],
        [0.03271886, 0.03422235, 0.02433733, 0.03713009, 0.05010711, 0.04489357, 0.04419974, 0.08159633]
      ])


# # Define the optimization model

# The optimization problem we would like to solve is 
# $$
#     \begin{array}{lrcl}
#     \mbox{maximize}     & \EMean^\mathsf{T}\mathbf{x}       &          &\\
#     \mbox{subject to}   & \left(\gamma^2, \frac{1}{2}, \mathbf{G}^\mathsf{T}\mathbf{x}\right)      & \in      & \Q_\mathrm{r}^{N+2},\\
#                         & \mathbf{1}^\mathsf{T}\mathbf{x}                & =        & 1,\\
#                         & \mathbf{x}                                     & \geq     & 0.\\
#     \end{array}
# $$
# 
# Here we define this model in MOSEK Fusion.

# In[32]:


# Define function solving the optimization model
def Markowitz(N, m, G, gamma2):
    with Model("markowitz") as M:
        # Settings
        M.setLogHandler(sys.stdout) 

        # Decision variable (fraction of holdings in each security)
        # The variable x is restricted to be positive, which imposes the constraint of no short-selling.   
        x = M.variable("x", N, Domain.greaterThan(0.0)) 

        # Budget constraint
        M.constraint('budget', Expr.sum(x), Domain.equalsTo(1))

        # Objective 
        M.objective('obj', ObjectiveSense.Maximize, Expr.dot(m, x))

        # Imposes a bound on the risk
        M.constraint('risk', Expr.vstack(gamma2, 0.5, Expr.mul(G.transpose(), x)), Domain.inRotatedQCone())

        # Solve optimization
        M.solve()
        
        # Check if the solution is an optimal point
        solsta = M.getPrimalSolutionStatus()
        if (solsta != SolutionStatus.Optimal):
            # See https://docs.mosek.com/latest/pythonfusion/accessing-solution.html about handling solution statuses.
            raise Exception("Unexpected solution status!") 
        
        returns = M.primalObjValue()
        portfolio = x.level()
        
    return returns, portfolio


# # Run the optimization

# ## Define the parameters

# The problem parameters are the number of securities $N$ and the risk limit $\gamma^2$. 

# In[33]:


N = mu.shape[0]  # Number of securities
gamma2 = 0.05   # Risk limit (variance)


# ## Factorize the covariance matrix

# Here we factorize $\ECov$ because the model is defined in conic form, and it expects a matrix $G$ such that $\ECov = GG^\mathsf{T}$.

# In[34]:


G = np.linalg.cholesky(Sigma)  # Cholesky factor of S to use in conic risk constraint 


# ## Solve the optimization problem

# Next we call the function that defines the Fusion model and runs the optimization.

# In[35]:


# Run optimization 
f, x = Markowitz(N, mu, G, gamma2)
print("========================\n")
print("RESULTS:")
print(f"Optimal expected portfolio return: {f*100:.4f}%")
print(f"Optimal portfolio weights: {x}")
print(f"Sum of weights: {np.sum(x)}")


# # Test result

# In[36]:


expected_x = np.array([0., 0.09126, 0.26911, 0., 0.02531, 0.32162, 0.17652, 0.11618])
diff = np.sum(np.abs(expected_x - x))
assert diff < 1e-4, f"Resulting portfolio does not match expected one. Difference is {diff}"