#!/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[196]:


import sys
import os
import re
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 *
import mosek.fusion.pythonic    # Requires MOSEK >= 10.2

from notebook.services.config import ConfigManager

from portfolio_tools import data_download, DataReader, compute_inputs


# In[197]:


# 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

# Here we load the raw data that will be used to compute the yearly return observation series used for the optimization. The data consists of daily stock prices of $8$ stocks from the US market, and SPY as the benchmark.

# ## Download data

# In[198]:


# Data downloading:
# If the user has an API key for alphavantage.co, then this code part will download the data. 
# The code can be modified to download from other sources. To be able to run the examples, 
# and reproduce results in the cookbook, the files have to have the following format and content:
# - File name pattern: "daily_adjusted_[TICKER].csv", where TICKER is the symbol of a stock. 
# - The file contains at least columns "timestamp", "adjusted_close", and "volume".
# - The data is daily price/volume, covering at least the period from 2016-03-18 until 2021-03-18, 
# - Files are for the stocks PM, LMT, MCD, MMM, AAPL, MSFT, TXN, CSCO.
list_stocks = ["PM", "LMT", "MCD", "MMM", "AAPL", "MSFT", "TXN", "CSCO"]
list_factors = ["SPY"]
alphaToken = None
 
list_tickers = list_stocks + list_factors
if alphaToken is not None:
    data_download(list_tickers, alphaToken)  


# ## Read data

# We load the daily stock price data from the downloaded CSV files. The data is adjusted for splits and dividends. Then a selected time period is taken from the data.

# In[199]:


investment_start = "2016-03-18"
investment_end = "2021-03-18"


# In[200]:


# The files are in "stock_data" folder, named as "daily_adjusted_[TICKER].csv"
dr = DataReader(folder_path="stock_data", symbol_list=list_tickers)
dr.read_data()
df_prices, _ = dr.get_period(start_date=investment_start, end_date=investment_end)


# # Run the optimization

# ## Define the optimization model

# Below we implement the optimization model in Fusion API. We create it inside a function so we can call it later.

# In[201]:


def absval(M, x, z):
    M.constraint(z >= x)
    M.constraint(z >= -x)
    
def norm1(M, x, t):
    z = M.variable(x.getSize(), Domain.greaterThan(0.0))
    absval(M, x, z)
    M.constraint(Expr.sum(z) == t)

def MinTrackingError(N, R, r_bm, x0, lambda_1, lambda_2, beta=1.5):

    with Model("Case study") as M:
        # Settings
        M.setLogHandler(sys.stdout)
        
        # Variables 
        # The variable x is the fraction of holdings in each security. 
        # It is restricted to be positive, which imposes the constraint of no short-selling.   
        x = M.variable("x", N, Domain.greaterThan(0.0))
        xt = x - x0
        
        # The variable t models the OLS objective function term (tracking error).
        t = M.variable("t", 1, Domain.unbounded())
        # The variables u and v model the regularization terms (transaction cost penalties).
        u = M.variable("u", 1, Domain.unbounded())
        v = M.variable("v", N, Domain.unbounded())   
    
        # Budget constraint
        M.constraint('budget', Expr.sum(x) == 1.0)
        
        # Objective 
        penalty_lin = lambda_1 * u
        penalty_32 = lambda_2 * Expr.sum(v)
        M.objective('obj', ObjectiveSense.Minimize, t + penalty_lin + penalty_32)
        
        # Constraints for the penalties
        norm1(M, xt, u)
        M.constraint('market_impact', Expr.hstack(v, Expr.constTerm(N, 1.0), xt), Domain.inPPowerCone(1.0 / beta))

        # Constraint for the tracking error 
        residual = R.T @ x - r_bm
        M.constraint('tracking_error', Expr.vstack(t, 0.5, residual), Domain.inRotatedQCone())
        
        # Create DataFrame to store the results. Last security name (the SPY) is removed.
        columns = ["track_err", "lin_tcost", "mkt_tcost"] + df_prices.columns[:N].tolist()
        df_result = pd.DataFrame(columns=columns)
        
        # 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!") 

        # Save results
        tracking_error = t.level()[0]
        linear_tcost = u.level()[0]
        market_impact_tcost = np.sum(v.level())
        row = pd.Series([tracking_error, linear_tcost, market_impact_tcost] + list(x.level()), index=columns)
        df_result = pd.concat([df_result, pd.DataFrame([row])], ignore_index=True)

        return df_result


# ## Compute optimization input variables

# Here we use the loaded daily price data to compute the corresponding yearly mean return and covariance matrix for logarithmic returns, and compute linear return observations from that. The benchmark will be the last data series, corresponding to SPY.

# In[202]:


# Number of securities and observations
N = df_prices.shape[1] - 1
T = 1000

# Mean and covariance of historical yearly log-returns.  
m_log, S_log = compute_inputs(df_prices, return_log=True)

# Generate logarithmic return observations assuming normal distribution
scenarios_log = np.random.default_rng().multivariate_normal(m_log, S_log, T)
    
# Convert logarithmic return observations to linear return observations 
scenarios_lin = np.exp(scenarios_log) - 1


# We center and normalize the data matrices.

# In[203]:


# Center the return data
centered_return = scenarios_lin - scenarios_lin.mean(axis=0)

# Security return scenarios
security_return = scenarios_lin[:, :N] / np.sqrt(T - 1)

# Benchmark return scenarios
benchmark_return = scenarios_lin[:, -1] / np.sqrt(T - 1)


# ## Call the optimizer function

# We run the optimization for the given penalty coefficients, and initial portfolio.

# In[206]:


lambda_1 = 0.0001 
lambda_2 = 0.0001 
x0 = np.ones(N) / N

df_result = MinTrackingError(N, security_return.T, benchmark_return, x0, lambda_1, lambda_2)


# ## Check the results

# In[207]:


df_result


# In[ ]: