import sys
import os
import re
import datetime as dt

import numpy as np
import pandas as pd
import statsmodels.api as sm
%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, compute_inputs

# 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 = []
alphaToken = None
 
list_tickers = list_stocks + list_factors
if alphaToken is not None:
    data_download(list_tickers, alphaToken)  

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

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

def multiperiod_mvo(N, T, m, G, x_0, delta, a, b):
    
    with Model("multiperiod") as M:
        # Settings
        M.setLogHandler(sys.stdout)
        
        # Variable
        x = M.variable("x", [N, T], Domain.greaterThan(0.0))
        s = M.variable("s", T)
        v = M.variable("v", [N, T])
        w = M.variable("w", [N, T])
        
        # Constraint
        M.constraint("budget", Expr.sum(x, 0), Domain.equalsTo(np.ones(T)))

        # Objective
        obj_terms = []
        for t in range(T):
            xt = x.pick([[i, t] for i in range(N)])
            term_t = Expr.add(
                [
                    Expr.dot(m[t], xt), 
                    Expr.mul(-delta[t], s.index(t)), 
                    Expr.add([Expr.mul(-a[i,t], v.index(i,t)) for i in range(N)]),
                    Expr.add([Expr.mul(-b[i,t], w.index(i,t)) for i in range(N)]),
                ]
            ) 
            obj_terms.append(term_t)
        M.objective("obj", ObjectiveSense.Maximize, Expr.add(obj_terms))
        
        # Objective cones
        for t in range(T):
            xt = x.pick([[i, t] for i in range(N)])
            vt = v.pick([[i, t] for i in range(N)])
            wt = w.pick([[i, t] for i in range(N)])
            xtprev = x_0 if t == 0 else x.pick([[i, t - 1] for i in range(N)])
            xtdiff = Expr.sub(xt, xtprev)
            M.constraint(f'risk_{t}', Expr.vstack(s.index(t), 0.5, Expr.mul(G[t].T, xt)), Domain.inRotatedQCone())
            absval(M, xtdiff, vt)
            M.constraint(f'market_impact_{t}', Expr.hstack(wt, Expr.constTerm(N, 1.0), xtdiff), Domain.inPPowerCone(2 / 3))
        
        # Solve the problem
        M.solve()
        
        # Get the solution values
        x_value = x.level().reshape(N, T)
        
        return x_value

# Number of securities
N = df_prices.shape[1]

# Number of periods
T = 10

# Initial weights
x_0 = np.array([1] * N) / N
portfolio_value = 10**8

def symmat(m):
    return (m + m.T) / 2

def makepsd(m):
    mineig = np.min(np.linalg.eigvals(m))
    if mineig < 0:
        m = m - (mineig - 0.0001) * np.identity(m.shape[0])
    return m

mu, Sigma = compute_inputs(df_prices)
m = [mu + np.random.normal(0, mu/10) for i in range(T)]
S = [makepsd(Sigma + symmat(np.random.normal(0, Sigma/10))) for i in range(T)]

G = [np.linalg.cholesky(s) for s in S] 

df_lin_returns = df_prices.pct_change()
volatility = df_lin_returns.std()
volume = (df_volumes * df_prices).mean()
vty = [abs(volatility + np.random.normal(0, volatility/10)) for i in range(T)]
vol = [abs(volume + np.random.normal(0, volume/10)) for i in range(T)]

# Transaction cost
a = 0.05 * np.ones((N, T))

# Market impact
beta = 3 / 2
b = 1
rel_volume = [v / portfolio_value for v in vol] # Relative volume (the variable x is also portfolio relative).
impact_coef = np.vstack([(b * v / r**(beta - 1)).to_numpy() for v, r in zip(vty, rel_volume)]).T

# Holding cost
s = 0.01