# Core numerical computing
import numpy as np
import pandas as pd
from typing import Dict, Tuple

# For financial data
import yfinance as yf       # Loading financial data
import empyrical as emp     # Performance metrics

# Book data (pip install "git+https://github.com/dppalomar/pob.git#subdirectory=python")
from pob_python import SP500_stocks_2015to2020, SP500_index_2015to2020

# Plotting
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import seaborn as sns
sns.set_theme(style="darkgrid")

# Optimization
import cvxpy as cp                  # interface for convex optimization solvers
from scipy.optimize import minimize, LinearConstraint   # for optimization
from scipy.linalg import cholesky

def backtest(
        portfolio_funcs: Dict[str, callable],
        prices: pd.DataFrame,
        lookback: int,
        rebalance_every: int = 1,
        cost_bps: float = 0
) -> Tuple[pd.DataFrame, Dict[str, pd.DataFrame]]:
    N = prices.shape[1]
    X = prices.pct_change().dropna()

    portf_rets = {}
    ws = {}

    for portfolio_name, portfolio_func in portfolio_funcs.items():
        print(f"Backtesting portfolio {portfolio_name}...")

        # Initialize variables
        current_w = np.repeat(0, N)
        w = pd.DataFrame(index=X.index, columns=X.columns)
        portf_ret = pd.Series(index=X.index)
        portf_ret.name = portfolio_name

        for t in range(lookback, len(X)):
            # Rebalance portfolio if necessary
            if (t - lookback) % rebalance_every == 0:
                new_w = portfolio_func(X.iloc[t-lookback:t])  # Note that the row at time t is not included
                turnover = np.abs(new_w - current_w).sum()
                transaction_cost = turnover * cost_bps / 1e4
                current_w = new_w

            # Store weights
            w.iloc[t] = current_w

            # Calculate portfolio return for this period
            period_return = (X.iloc[t] * current_w).sum()
            portf_ret.iloc[t] = period_return - transaction_cost

            # Update normalized weights based on asset performance
            current_w = current_w * (1 + X.iloc[t])
            current_w = current_w / current_w.sum()

        # Keep a record (remove inital NaNs)
        portf_rets[portfolio_name] = portf_ret[lookback:]
        ws[portfolio_name] = w[lookback:]

    # Combine all portfolio returns into a single DataFrame
    portf_rets_df = pd.DataFrame(portf_rets)

    return portf_rets_df, ws

def print_table_performance_metrics(returns: pd.DataFrame):
    risk_free_rate = 0.0  # Risk-free rate (e.g., 0 for simplicity)
    trading_days = 252    # Number of trading days in a year

    # Calculate Annualized Return
    annualized_return = (1 + returns.mean())**trading_days - 1

    # Calculate Annualized Volatility
    annualized_volatility = returns.std() * np.sqrt(trading_days)

    # Calculate Sharpe Ratio
    sharpe_ratio = (returns.mean() - risk_free_rate) / returns.std()
    sharpe_ratio *= np.sqrt(trading_days)  # Annualize Sharpe ratio

    # Calculate Maximum Drawdown
    def calculate_max_drawdown(returns):
        cumulative_returns = (1 + returns).cumprod()
        running_max = cumulative_returns.cummax()
        drawdown = (cumulative_returns - running_max) / running_max
        max_drawdown = drawdown.min()
        return max_drawdown

    max_drawdown = returns.apply(calculate_max_drawdown)

    # Combine all metrics into a single DataFrame
    performance_table = pd.DataFrame({
        'Annualized Return': annualized_return,
        'Annualized Volatility': annualized_volatility,
        'Sharpe Ratio': sharpe_ratio,
        'Maximum Drawdown': max_drawdown
    })

    # Display the performance table
    performance_table_styled = performance_table.style.format("{:.2f}")
    display(performance_table_styled)

def plot_cum_return(returns: pd.DataFrame) -> None:
    cumulative_returns = (1 + returns).cumprod() - 1
    fig, ax = plt.subplots(figsize=(12, 6))
    cumulative_returns.plot(ax=ax)
    ax.set_title('Cumulative Return')
    ax.set_xlabel(None)
    ax.legend(title="portfolio")
    plt.show()

def plot_drawdown(returns: pd.DataFrame) -> None:
    cumulative_returns = (1 + returns).cumprod()
    running_max = cumulative_returns.cummax()
    drawdowns = (cumulative_returns - running_max) / running_max

    fig, ax = plt.subplots(figsize=(12, 6))
    drawdowns.plot(ax=ax)
    ax.set_title('Drawdown')
    ax.set_xlabel(None)
    ax.legend(title="portfolio")
    plt.show()

def plot_initial_allocations(bt_w: Dict[str, pd.DataFrame]) -> None:
    # Extract the first row of each portfolio's weights and prepare for plotting
    first_rows = {
        portfolio_name: weights.iloc[0]  # Get the first row of each DataFrame
        for portfolio_name, weights in bt_w.items()
    }

    # Convert the dictionary of first rows into a DataFrame
    first_rows_df = pd.DataFrame(first_rows).T.reset_index()  # Transpose and reset index
    first_rows_df.columns = ['portfolio'] + list(first_rows_df.columns[1:])

    # Barplot of weight allocation
    melted_df = first_rows_df.melt(id_vars='portfolio', var_name='stocks', value_name='weights')

    fig, ax = plt.subplots(figsize=(12, 4))
    sns.barplot(x='stocks', y='weights', hue='portfolio', ax=ax, data=melted_df)
    ax.set_title("Weight Allocation for Heuristic Portfolios")
    plt.xticks(rotation=45)
    plt.show()

# Get data
stock_prices = SP500_stocks_2015to2020[
                   ["AAPL", "AMZN", "AMD", "GM", "GOOGL", "MGM", "MSFT", "QCOM", "TSCO", "UPS"]
               ].loc["2019":]
T = stock_prices.shape[0]
T_trn = round(0.70*T)

stock_prices.head()

#
# Define portfolios
#
def EWP(X: pd.DataFrame) -> np.ndarray:
    N = X.shape[1]
    return np.repeat(1/N, N)

def design_MVP(mu: np.ndarray, Sigma: np.ndarray, lmd: float = 1) -> np.ndarray:
    N = len(Sigma)
    w = cp.Variable(N)
    prob = cp.Problem(cp.Maximize(mu @ w - (lmd/2) * cp.quad_form(w, Sigma)),
                      [cp.sum(w) == 1, w >= 0])
    prob.solve()
    return w.value

def MVP_lmd1(X_lin: pd.DataFrame) -> np.ndarray:
    X_log = np.log(1 + X_lin)
    mu = X_log.mean().values
    Sigma = X_log.cov().values
    return design_MVP(mu, Sigma, lmd=1)

def MVP_lmd4(X_lin: pd.DataFrame) -> np.ndarray:
    X_log = np.log(1 + X_lin)
    mu = X_log.mean().values
    Sigma = X_log.cov().values
    return design_MVP(mu, Sigma, lmd=4)

def MVP_lmd16(X_lin: pd.DataFrame) -> np.ndarray:
    X_log = np.log(1 + X_lin)
    mu = X_log.mean().values
    Sigma = X_log.cov().values
    return design_MVP(mu, Sigma, lmd=16)

def MVP_lmd64(X_lin: pd.DataFrame) -> np.ndarray:
    X_log = np.log(1 + X_lin)
    mu = X_log.mean().values
    Sigma = X_log.cov().values
    return design_MVP(mu, Sigma, lmd=64)

# Run backtest
bt_ret, bt_w = backtest(
    {
        "1/N": EWP,
        "MVP (lmd = 1)": MVP_lmd1,
        "MVP (lmd = 4)": MVP_lmd4,
        "MVP (lmd = 16)": MVP_lmd16,
        "MVP (lmd = 64)": MVP_lmd64,
    },
    stock_prices, lookback=T_trn, rebalance_every=30)

plot_initial_allocations(bt_w)

plot_cum_return(bt_ret)

plot_drawdown(bt_ret)

print_table_performance_metrics(bt_ret)

#
# Define portfolios
#
def GMVP(X_lin: pd.DataFrame) -> np.ndarray:
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    Sigma = X_log.cov().values
    # Optimize portfolio
    N = len(Sigma)
    w = cp.Variable(N)
    prob = cp.Problem(cp.Minimize(cp.quad_form(w, Sigma)),
                      [cp.sum(w) == 1, w >= 0])
    prob.solve()
    return w.value

def MVP_vanilla(X_lin: pd.DataFrame) -> np.ndarray:
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    mu = X_log.mean().values
    Sigma = X_log.cov().values
    # Optimize portfolio
    lmd = 1
    N = len(Sigma)
    w = cp.Variable(N)
    prob = cp.Problem(cp.Maximize(mu.T @ w - (lmd/2) * cp.quad_form(w, Sigma)),
                      [cp.sum(w) == 1, w >= 0])
    prob.solve()
    return w.value

def MVP_ub(X_lin: pd.DataFrame) -> np.ndarray:
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    mu = X_log.mean().values
    Sigma = X_log.cov().values
    # Optimize portfolio
    N = len(Sigma)
    lmd = 1
    ub = 0.25
    w = cp.Variable(N)
    prob = cp.Problem(cp.Maximize(mu.T @ w - (lmd/2) * cp.quad_form(w, Sigma)),
                      [cp.sum(w) == 1, w >= 0, w <= ub])
    prob.solve()
    return w.value

def MVP_diversification(X_lin: pd.DataFrame) -> np.ndarray:
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    mu = X_log.mean().values
    Sigma = X_log.cov().values
    # Optimize portfolio
    N = len(Sigma)
    lmd = 1
    D = 2/N
    w = cp.Variable(N)
    prob = cp.Problem(cp.Maximize(mu.T @ w - (lmd/2) * cp.quad_form(w, Sigma)),
                      [cp.sum(w) == 1, w >= 0, cp.sum_squares(w) <= D])
    prob.solve()
    return w.value

# Run backtest
bt_ret, bt_w = backtest(
    {
        "1/N": EWP,
        "GMVP": GMVP,
        "MVP": MVP_vanilla,
        "MVP with upper bound": MVP_ub,
        "MVP with diversific. const.": MVP_diversification,
    },
    stock_prices, lookback=T_trn, rebalance_every=30)

plot_initial_allocations(bt_w)

plot_cum_return(bt_ret)

plot_drawdown(bt_ret)

print_table_performance_metrics(bt_ret)

stock_prices = SP500_stocks_2015to2020[
                   ["AAPL", "AMZN", "AMD", "GM", "GOOGL", "MGM", "MSFT", "QCOM", "TSCO", "UPS"]
               ].loc["2019":]
X_lin = stock_prices.pct_change().dropna()
X_log = np.log(1 + X_lin)
T, N = X_lin.shape
mu = X_log.mean().values
Sigma = X_log.cov().values

# Define the nonconvex objective function
def fn_SR(w):
    return (w @ mu) / np.sqrt(w @ Sigma @ w)

def neg_fn_SR(w):
    return -fn_SR(w)

# Initial point
N = len(mu)
w0 = np.ones(N) / N

# Equality constraint: sum(w) = 1
def constraint_eq(w):
    return np.sum(w) - 1

# Optimization
res = minimize(neg_fn_SR, w0, method='SLSQP',
               bounds=[(0, None) for _ in range(N)],  # w >= 0
               constraints={'type': 'eq', 'fun': constraint_eq})

w_nonlinear_solver = res.x

print(res)

# Square-root of matrix Sigma
Sigma_12 = cholesky(Sigma)
print(np.max(np.abs(Sigma_12.T @ Sigma_12 - Sigma)))  # sanity check

# Create function for MVP (each iteration of bisection)
import warnings

def SOCP_bisection(t):
    w = cp.Variable(N)
    prob = cp.Problem(cp.Maximize(0),
                      constraints=[t * cp.norm(Sigma_12 @ w, 2) <= mu.T @ w,
                                   sum(w) == 1,
                                   w >= 0])
    # Solve the problem (ignore annoying warnings):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        prob.solve()
    return {"status": prob.status, "w": w.value}

# Now run the bisection algorithm
t_lb = 0   # for sure the problem is feasible in this case
t_ub = 10  # a tighter upper bound could be chosen (10 is a conservative choice)
while t_ub - t_lb > 1e-3:
    t = (t_ub + t_lb) / 2  # midpoint
    print(f"Current interval = ({t_lb:.3f}, {t_ub:.3f}) Trying midpoint t = {t:.4f}")
    if 'infeasible' in SOCP_bisection(t)["status"]:
        t_ub = t
    else:
        t_lb = t

w_bisection = SOCP_bisection(t_lb)["w"]

# Comparison between two solutions
w_df = pd.DataFrame({
    'nonlinear_solver': w_nonlinear_solver,
    'bisection': w_bisection,
})
print(w_df.round(3))

# Sharpe ratio of two solutions
sharpe_ratios_df = pd.DataFrame({
    'Method': ['nonlinear_solver', 'bisection'],
    'Sharpe Ratio': [fn_SR(w_nonlinear_solver), fn_SR(w_bisection)]
})
print(sharpe_ratios_df.round(3))

# Create function for MVP (each iteration of Dinkelbach)
def SOCP_Dinkelbach(y):
    w = cp.Variable(Sigma.shape[0])
    prob = cp.Problem(cp.Maximize(mu.T @ w - y * cp.norm(Sigma_12 @ w, 2)),
                      constraints=[cp.sum(w) == 1,
                                   w >= 0])
    prob.solve()
    return w.value

# Initial point (has to satisfy w_k.T @ mu >= 0)
i_max = np.argmax(mu)
w_k = np.zeros(N)
w_k[i_max] = 1

# Now the iterative Dinkelbach algorithm
k = 1
while k == 1 or np.max(np.abs(w_k - w_prev)) > 1e-6:
    w_prev = w_k
    y_k = (w_k.T @ mu) / np.sqrt(w_k.T @ Sigma @ w_k)
    w_k = SOCP_Dinkelbach(y_k)
    k += 1

w_Dinkelbach = w_k
print(f"Number of iterations: {k-1}")

# Comparison among three solutions
w_df = pd.DataFrame({
    'nonlinear_solver': w_nonlinear_solver,
    'bisection': w_bisection,
    'Dinkelbach': w_Dinkelbach
})
print(w_df.round(3))

# Sharpe ratio of three solutions
sharpe_ratios_df = pd.DataFrame({
    'Method': ['nonlinear_solver', 'bisection', 'Dinkelbach'],
    'Sharpe Ratio': [fn_SR(w_nonlinear_solver), fn_SR(w_bisection), fn_SR(w_Dinkelbach)]
})
print(sharpe_ratios_df.round(3))

# create function for MSRP in convex form
def MSRP_convex(mu, Sigma):
    w_ = cp.Variable(Sigma.shape[0])
    prob = cp.Problem(cp.Minimize(cp.quad_form(w_, Sigma)),
                      constraints=[w_ >= 0, mu.T @ w_ == 1])
    prob.solve()
    w = w_.value / np.sum(w_.value)
    return w

# This function can now be used as
w_MSRP = MSRP_convex(mu, Sigma)

# Comparison among solutions
comparison = pd.DataFrame({
    'w_nonlinear_solver': w_nonlinear_solver,
    'w_bisection': w_bisection,
    'w_Dinkelbach': w_Dinkelbach,
    'w_MSRP': w_MSRP
})
print(comparison.round(3))

# Sharpe ratio of different solutions
sharpe_ratios_df = pd.DataFrame({
    'Method': ['nonlinear_solver', 'bisection', 'Dinkelbach', 'Schaible'],
    'Sharpe Ratio': [fn_SR(w_nonlinear_solver), fn_SR(w_bisection), fn_SR(w_Dinkelbach), fn_SR(w_MSRP)]
})
print(sharpe_ratios_df.round(3))

# Get data
stock_prices = SP500_stocks_2015to2020[
                   ["AAPL", "AMZN", "AMD", "GM", "GOOGL", "MGM", "MSFT", "QCOM", "TSCO", "UPS"]
               ].loc["2019":]
T = stock_prices.shape[0]
T_trn = round(0.70*T)

#
# Define portfolios
#
def GMVP(X_lin: pd.DataFrame) -> np.ndarray:
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    Sigma = X_log.cov().values
    # Optimize portfolio
    N = len(Sigma)
    w = cp.Variable(N)
    prob = cp.Problem(cp.Minimize(cp.quad_form(w, Sigma)),
                      [cp.sum(w) == 1, w >= 0])
    prob.solve()
    return w.value

def MVP(X_lin: pd.DataFrame) -> np.ndarray:
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    mu = X_log.mean().values
    Sigma = X_log.cov().values
    # Optimize portfolio
    lmd = 1
    N = len(Sigma)
    w = cp.Variable(N)
    prob = cp.Problem(cp.Maximize(mu.T @ w - (lmd/2) * cp.quad_form(w, Sigma)),
                      [cp.sum(w) == 1, w >= 0])
    prob.solve()
    return w.value

def MSRP(X_lin: pd.DataFrame) -> np.ndarray:
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    mu = X_log.mean().values
    Sigma = X_log.cov().values
    # Optimize portfolio
    w = MSRP_convex(mu, Sigma)
    return w

# Run backtest
bt_ret, bt_w = backtest(
    {
        "GMVP": GMVP,
        "MVP": MVP,
        "MSRP": MSRP,
    },
    stock_prices, lookback=T_trn, rebalance_every=30)

plot_initial_allocations(bt_w)

plot_cum_return(bt_ret)

plot_drawdown(bt_ret)

print_table_performance_metrics(bt_ret)

# Prepare data
stock_prices = SP500_stocks_2015to2020[
                   ["AAPL", "AMZN", "AMD", "GM", "GOOGL", "MGM", "MSFT", "QCOM", "TSCO", "UPS"]
               ].loc["2019":]
X_lin = stock_prices.pct_change().dropna()
X_log = np.log(1 + X_lin)
T, N = X_lin.shape
mu = X_log.mean().values
Sigma = X_log.cov().values

# Using Schaible transform + CVX to get the optimal value
w_ = cp.Variable(N)
prob = cp.Problem(cp.Minimize(cp.quad_form(w_, Sigma)),
                  constraints=[w_ >= 0, mu.T @ w_ == 1])
prob.solve()
w_cvx = w_.value / np.sum(w_.value)
obj_cvx = (w_cvx.T @ mu) / np.sqrt(w_cvx.T @ Sigma @ w_cvx)

# Using the SQP-MVP algorithm
w = np.ones(N) / N
obj_sqp = [(w.T @ mu) / np.sqrt(w.T @ Sigma @ w)]
k = 0
for k in range(21):
    lmd_k = (w.T @ mu) / (w.T @ Sigma @ w)
    w_prev = w
    w = cp.Variable(N)
    prob = cp.Problem(cp.Maximize(mu.T @ w - (lmd_k/2) * cp.quad_form(w, Sigma)),
                      constraints=[cp.sum(w) == 1, w >= 0])
    prob.solve()
    w = w.value
    obj_sqp.append((w.T @ mu) / np.sqrt(w.T @ Sigma @ w))
    if np.max(np.abs(w - w_prev)) < 1e-4:
        break

# Plot
df = pd.DataFrame({
    "k": range(len(obj_sqp)),
    "SQP iterative method": obj_sqp
})

plt.figure(figsize=(10, 6))
plt.axhline(y=obj_cvx, color='r', linestyle='-', linewidth=1.5)
plt.plot(df["k"], df["SQP iterative method"], color='b', linewidth=1.5, marker='o', markersize=2.5)
plt.title("Convergence")
plt.ylabel("objective value")
plt.xlabel("iteration")
plt.legend(["optimal value", "SQP iterative method"])
plt.grid(True)
plt.show()

kappa = 1.0

# Using CVX to get the optimal value
Sigma_12 = cholesky(Sigma)
w = cp.Variable(N)
prob = cp.Problem(cp.Maximize(mu.T @ w - kappa * cp.norm(Sigma_12 @ w, 2)),
                  constraints=[cp.sum(w) == 1, w >= 0])
result = prob.solve()
w_cvx = w.value
obj_cvx = mu.T @ w_cvx - kappa * np.sqrt(w_cvx.T @ Sigma @ w_cvx)

# Using the SQP-MVP algorithm
w = np.ones(N) / N
obj_sqp = [w.T @ mu - np.sqrt(w.T @ Sigma @ w)]
for k in range(21):
    lmd_k = kappa / np.sqrt(w.T @ Sigma @ w)
    w_prev = w
    w = cp.Variable(N)
    prob = cp.Problem(cp.Maximize(mu.T @ w - (lmd_k/2) * cp.quad_form(w, Sigma)),
                      constraints=[cp.sum(w) == 1, w >= 0])
    prob.solve()
    w = w.value
    obj_sqp.append(w.T @ mu - kappa * np.sqrt(w.T @ Sigma @ w))
    if np.max(np.abs(w - w_prev)) < 1e-4:
        break

# Plot
df = pd.DataFrame({
    "k": range(len(obj_sqp)),
    "SQP iterative method": obj_sqp
})

plt.figure(figsize=(10, 6))
plt.axhline(y=obj_cvx, color='r', linestyle='-', linewidth=1.5)
plt.plot(df["k"], df["SQP iterative method"], color='b', linewidth=1.5, marker='o', markersize=2.5)
plt.title("Convergence")
plt.ylabel("objective value")
plt.xlabel("iteration")
plt.legend(["optimal value", "SQP iterative method"])
plt.grid(True)
plt.show()