# Core numerical computing
import numpy as np
import pandas as pd

# For financial data
import yfinance as yf       # Lodaing 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
import riskparityportfolio as rpp
from scipy.optimize import minimize, LinearConstraint   # for optimization

stocks = yf.download(["AAPL", "AMD", "ADI", "A", "MOH", "CVS", "APD", "AA", "CF"], start='2021-10-01', end='2021-12-31')

stocks[['Adj Close']].head()

from pob_python import SP500_stocks_2015to2020, cryptos_2017to2021_hourly

# stock S&P500 market data
SP500_stocks_2015to2020.iloc[:, :5].head()

SP500_stocks_2015to2020[["AMD", "MGM", "AAPL"]].head()

# crypto data
cryptos_2017to2021_hourly.iloc[:, :5].head()

# Plot stock prices
stock_prices = SP500_stocks_2015to2020[["AMD", "MGM", "AAPL"]].loc["2019-10":]

fig, ax = plt.subplots(figsize=(12, 6))
stock_prices.plot(ax=ax)
ax.set_title('Prices of stocks')
ax.set_xlabel(None)
ax.legend(title="Stocks")
plt.show()

stock_returns = stock_prices.pct_change()

# Plot returns
fig, ax = plt.subplots(figsize=(12, 6))
stock_returns.plot(ax=ax)
ax.set_title('Returns of stocks')
ax.set_xlabel(None)
ax.legend(title="Stocks")
plt.show()

# Compute drawdowns
cumulative_returns = (1 + stock_returns).cumprod()
running_max = cumulative_returns.cummax()
drawdowns = (cumulative_returns - running_max) / running_max

# Plot drawdowns
fig, ax = plt.subplots(figsize=(12, 6))
drawdowns.plot(ax=ax)
ax.set_title('Drawdown of stocks')
ax.set_xlabel(None)
ax.legend(title="Stocks")
plt.show()

stock_prices = SP500_stocks_2015to2020[["AMD", "MGM", "AAPL", "AMZN", "TSCO"]].loc["2019-10":]
T, N = stock_prices.shape

# linear returns
X = stock_prices / stock_prices.shift(1) - 1
X = X.dropna()

# portfolio
w = np.repeat(1/N, N)

# naive backtest (assuming daily rebalancing and no transaction cost)
portf_ret = X @ w
portf_ret.head()

def EWP(X):
    N = X.shape[1]
    return np.repeat(1/N, N)

def backtest_single_portfolio(portfolio_func, portfolio_name, prices, lookback, rebalance_every=1, cost_bps=0):
    N = prices.shape[1]

    # Calculate returns
    X = prices.pct_change().dropna()

    # 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()

    return portf_ret, w

bt_ret, bt_w = backtest_single_portfolio(EWP, "1/N", stock_prices, lookback=0, rebalance_every=1)
bt_ret.head()

bt_ret, bt_w = backtest_single_portfolio(EWP, "1/N", stock_prices, lookback=0, rebalance_every=5, cost_bps=30e-4)
bt_ret.head()

def plot_cum_return(returns):
    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()

plot_cum_return(bt_ret)

def plot_drawdown(returns):
    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()

plot_drawdown(bt_ret)

# Run backtest
bt_ret, bt_w = backtest_single_portfolio(EWP, "1/N", stock_prices, lookback=0, rebalance_every=90)

# Filter the weights DataFrame for the desired date range
filtered_w = bt_w.loc["2020-01":"2020-08"]

# Create the plot
fig, ax = plt.subplots(figsize=(12, 6))

# Calculate bar width (in days)
bar_width = (filtered_w.index[-1] - filtered_w.index[0]).days / len(filtered_w)

# Plot stacked bars
bottom = np.zeros(len(filtered_w))
for column in filtered_w.columns:
    ax.bar(filtered_w.index, filtered_w[column], bottom=bottom, width=3*bar_width, label=column, edgecolor='none')
    bottom += filtered_w[column]

# Customize the x-axis
ax.xaxis.set_major_locator(mdates.MonthLocator())
ax.xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))
plt.xticks(rotation=45, ha='right')


# Customize other aspects of the plot
plt.title("Weight Allocation Over Time for Portfolio 1/N")
plt.xlabel(None)
plt.ylabel("Weight")
plt.legend(title="Stocks", bbox_to_anchor=(1.05, 1), loc='upper left')

# Add vertical lines for specific dates
plt.axvline(pd.to_datetime("2020-02-11"), color="black", linestyle="-", linewidth=3)
plt.axvline(pd.to_datetime("2020-06-19"), color="black", linestyle="-", linewidth=3)
plt.axhline(0.0, color="black", linestyle="--")
plt.axhline(0.2, color="black", linestyle="--")
plt.axhline(0.4, color="black", linestyle="--")
plt.axhline(0.6, color="black", linestyle="--")
plt.axhline(0.8, color="black", linestyle="--")
plt.axhline(1.0, color="black", linestyle="--")

# Adjust layout and display
plt.tight_layout()
plt.show()

def backtest(portfolio_funcs, prices, lookback, rebalance_every=1, cost_bps=0):
    N = prices.shape[1]
    X = prices.pct_change().dropna()

    portf_rets = {}
    ws = {}

    for portfolio_name, portfolio_func in portfolio_funcs.items():
        # 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

# 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 heuristic portfolios
def QuintP_mu(X_lin):
    X_log = np.log(1 + X_lin)
    mu = X_log.mean()
    idx = mu.argsort()[::-1]
    w = np.zeros(len(mu))
    w[idx[:len(mu)//5]] = 1 / (len(mu)//5)
    return w

def QuintP_mu_over_sigma(X_lin):
    X_log = np.log(1 + X_lin)
    mu = X_log.mean()
    sigma = X_log.std()
    idx = (mu / sigma).argsort()[::-1]
    w = np.zeros(len(mu))
    w[idx[:len(mu)//5]] = 1 / (len(mu)//5)
    return w

def QuintP_mu_over_sigma2(X_lin):
    X_log = np.log(1 + X_lin)
    mu = X_log.mean()
    sigma = X_log.std()
    idx = (mu / sigma**2).argsort()[::-1]
    w = np.zeros(len(mu))
    w[idx[:len(mu)//5]] = 1 / (len(mu)//5)
    return w

def GMRP(X_lin):
    X_log = np.log(1 + X_lin)
    mu = X_log.mean()
    w = np.zeros(len(mu))
    w[mu.argmax()] = 1
    return w

# Run backtest
bt_ret, bt_w = backtest(
    {
        "1/N": EWP,
        "GMRP": GMRP,
        "QuintP (sorted by mu)": QuintP_mu,
        "QuintP (sorted by mu/sigma)": QuintP_mu_over_sigma,
        "QuintP (sorted by mu/sigma2)": QuintP_mu_over_sigma2,
    },
    stock_prices, lookback=T_trn, rebalance_every=30)

bt_w["1/N"].head()

# 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:])

first_rows_df

# 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()

# Plot of cumulative return
plot_cum_return(bt_ret)

# Plot of drawdown
plot_drawdown(bt_ret)

#
# Table of performance metrics
#
def print_table_performance_metrics(bt_ret):
    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 + bt_ret.mean())**trading_days - 1

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

    # Calculate Sharpe Ratio
    sharpe_ratio = (bt_ret.mean() - risk_free_rate) / bt_ret.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 = bt_ret.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)

print_table_performance_metrics(bt_ret)

# Define risk-based portfolios
def GMVP(X_lin):
    # 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 IVP(X_lin):
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    sigma = X_log.std()
    # Design portfolio
    w = 1 / sigma
    return w / w.sum()

def RPP(X_lin):
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    Sigma = X_log.cov().values
    # Optimize portfolio
    N = len(Sigma)
    my_portfolio = rpp.RiskParityPortfolio(covariance=Sigma, weights=np.ones(N) / N, budget=np.ones(N) / N)
    my_portfolio.design()
    return my_portfolio.weights

def MDP(X_lin):
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    Sigma = X_log.cov().values
    mu = np.sqrt(np.diag(Sigma))
    # Optimize portfolio
    N = len(Sigma)
    w = cp.Variable(N)
    prob = cp.Problem(cp.Minimize(cp.quad_form(w, Sigma)),
                      [mu @ w == 1, w >= 0])
    prob.solve()
    return w.value / w.value.sum()

def MDCP(X_lin):
    # Estimate parameters
    X_log = np.log(1 + X_lin)
    Sigma = X_log.cov().values
    D = np.diag(1 / np.sqrt(np.diag(Sigma)))
    C = D @ Sigma @ D
    # Optimize portfolio
    N = len(C)
    w = cp.Variable(N)
    prob = cp.Problem(cp.Minimize(cp.quad_form(w, C)),
                      [cp.sum(w) == 1, w >= 0])
    prob.solve()
    return w.value

# Run backtest
bt_ret, bt_w = backtest(
    {
        "GMVP": GMVP,
        "IVP": IVP,
        "RPP": RPP,
        "MDP": MDP,
        "MDCP": MDCP,
    },
    stock_prices, lookback=T_trn, rebalance_every=30)

#
# Barplot of weight allocation
#
first_rows = {
    portfolio_name: weights.iloc[0]  # Get the first row of each DataFrame
    for portfolio_name, weights in bt_w.items()
}
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:])
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()

# Plot of cumulative return
plot_cum_return(bt_ret)

# Plot of drawdown
plot_drawdown(bt_ret)

# Table of performance metrics
print_table_performance_metrics(bt_ret)