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

# For financial data
import yfinance as yf       # Loading financial data

# 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")
import plotly.io as pio
pio.renderers.default = "notebook"
pio.templates.default = "plotly_white"

# Optimization
import cvxpy as cp                  # interface for convex optimization solvers
import riskparityportfolio as rpp

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

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

bt_ret, bt_w = backtest_single_portfolio(EWP, "1/N", stock_prices.iloc[1:], lookback=0, rebalance_every=5, cost_bps=30e-4)
# Note: We exclude the first row for a simpler comparison later with skfolio
bt_ret.head(8)

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

from skfolio.optimization import EqualWeighted

model_ewp = EqualWeighted()
model_ewp.fit(X)

print("Weights of the equally weighted portfolio:")
print(np.round(model_ewp.weights_, 4))

# Run backtest
bt_ewp = model_ewp.predict(X)

# Display first returns
bt_ewp.returns_df.head()

# Plot cumulative returns
bt_ewp.plot_cumulative_returns()

# WalkForward backtest
from skfolio.model_selection import cross_val_predict, WalkForward

model_ewp = EqualWeighted(
    portfolio_params=dict(name="EWP", transaction_costs=30e-4 / 5)
)

bt_ewp = cross_val_predict(
    model_ewp,
    X,
    cv=WalkForward(test_size=5, train_size=1),
)
# Note: The minimum value for train_size in skfolio is train_size=1

# Show last designed portfolio
bt_ewp[-1].composition

# Plot last designed portfolio
bt_ewp[-1].plot_composition()

import plotly.graph_objects as go

# Get the composition DataFrame
df = bt_ewp[-1].composition  # shape: (n_assets, 1) with portfolio name as column

# Create vertical bar chart
fig = go.Figure(go.Bar(
    x=df.index,             # asset names on x-axis
    y=df.iloc[:, 0]         # weights on y-axis
))

fig.update_layout(
    title="Portfolio Composition",
    xaxis_title="Assets",
    yaxis_title="Weight"
)

fig.show()

bt_ewp.returns_df.head(8)

# Plot cumulative returns
bt_ewp.compounded = True
fig = bt_ewp.plot_cumulative_returns()
fig.update_layout(title="Wealth", xaxis_title=None, yaxis_title="Net Asset Value (NAV)")
fig

# Plot drawdowns
bt_ewp.plot_drawdowns()

# WalkForward backtest rebalancing every 90 days
bt_ewp = cross_val_predict(
    model_ewp,
    X,
    cv=WalkForward(test_size=90, train_size=1),
)

# Plot last designed portfolio
bt_ewp[-1].plot_composition()

# Plot weights evolution over time
bt_ewp.plot_weights_per_observation()

# Get data
from pob_python import SP500_stocks_2015to2020, SP500_index_2015to2020
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 EWP(X):
    N = X.shape[1]
    return np.repeat(1/N, N)

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 using skfolio
#

# Create wrapper to backtest multiple backtests (currently not available in skfolio, but soon will be)
import time
from datetime import timedelta

def batch_cross_val_predict(optimizers, X, cv, **kwargs) -> list:
    """Run cross_val_predict on multiple optimizers with timing information"""
    bt_list = []
    n_splits = cv.get_n_splits(X)
    total_start_time = time.time()

    print(f"\n🚀 Starting batch backtesting with {len(optimizers)} portfolio designs...")
    print(f"📊 Cross-validation: {n_splits} splits")
    print("=" * 70)

    for i, optim in enumerate(optimizers, 1):
        # Handle missing or None portfolio_params (fixes the original error)
        if hasattr(optim, 'portfolio_params') and optim.portfolio_params is not None:
            name = optim.portfolio_params.get('name', f'Portfolio {i}')
        else:
            name = f'Portfolio {i} ({type(optim).__name__})'

        print(f"\n[{i}/{len(optimizers)}] 🏃 Backtesting '{name}'...")

        # Start timing for this specific backtest
        start_time = time.time()

        try:
            bt = cross_val_predict(
                optim, X, cv=cv,
                portfolio_params=getattr(optim, 'portfolio_params', {}),
                **kwargs
            )

            # Calculate elapsed time
            elapsed_time = time.time() - start_time
            elapsed_str = str(timedelta(seconds=int(elapsed_time)))

            bt_list.append(bt)
            print(f"✅ Completed '{name}' in {elapsed_str}")

        except Exception as e:
            elapsed_time = time.time() - start_time
            elapsed_str = str(timedelta(seconds=int(elapsed_time)))
            print(f"❌ Failed '{name}' after {elapsed_str}: {str(e)}")
            raise  # Re-raise the exception to maintain original behavior

    # Calculate total elapsed time
    total_elapsed = time.time() - total_start_time
    total_elapsed_str = str(timedelta(seconds=int(total_elapsed)))

    print("=" * 70)
    print(f"🎉 Batch backtesting completed!")
    print(f"⏱️  Total time: {total_elapsed_str}")
    print(f"📈 Successfully processed {len(bt_list)} portfolios")

    return bt_list

# Create wrapper to plot portfolio compositions as a barplot
def plot_composition_sidebyside(population, asset_names=None):
    """
    Plot portfolio compositions side-by-side for a Population of portfolios.

    Parameters
    ----------
    population : Population
        A skfolio Population object containing multiple portfolios.
    asset_names : list of str, optional
        List of asset names in the desired order. If provided, all assets
        will be displayed (with zero weights for missing assets).

    Returns
    -------
    fig : plotly.graph_objects.Figure
        The plotly figure object.
    """
    import plotly.graph_objects as go

    # Determine the asset order
    if asset_names is None:
        asset_names = population[0].composition.index.tolist()

    fig = go.Figure()

    # Add a bar trace for each portfolio
    for portfolio in population:
        df = portfolio.composition  # shape: (n_assets, 1)

        # Reindex to include all assets in the specified order, filling missing with 0
        weights = df.iloc[:, 0].reindex(asset_names, fill_value=0)

        fig.add_trace(go.Bar(
            x=asset_names,           # asset names in specified order
            y=weights,               # weights (0 for missing assets)
            name=portfolio.name      # legend label (portfolio name)
        ))

    fig.update_layout(
        title="Portfolio Compositions",
        xaxis_title="Assets",
        yaxis_title="Weight",
        barmode='group',                        # bars side-by-side (not stacked)
        xaxis={'categoryorder': 'array',        # enforce asset order
               'categoryarray': asset_names}
    )

    return fig

# Create a class to be able to specify portfolios direclty via a CVXPY code
from skfolio.optimization import ConvexOptimization

class Portfolio_via_CVXPY(ConvexOptimization):
    """Portfolio based on custom CVXPY function."""

    def __init__(self, cvxpy_fun=None, **kwargs):
        super().__init__(**kwargs)
        self.cvxpy_fun = cvxpy_fun

    def fit(self, X, y=None):
        """Fit the custom optimization problem using the provided function."""

        if self.cvxpy_fun is None:
            raise ValueError("cvxpy_fun must be provided at initialization")

        # Call the custom function to compute weights
        weights = self.cvxpy_fun(X)

        # Store results
        self.weights_ = np.asarray(weights)

        return self

from skfolio.optimization import MeanRisk, EqualWeighted, ConvexOptimization
from skfolio import Population
from skfolio.model_selection import cross_val_predict, WalkForward
from skfolio.preprocessing import prices_to_returns

# Wrap each portfolio within skfolio
portf_EWP_skfolio = EqualWeighted(
    portfolio_params=dict(name="1/N (skfolio)")
)
portf_EWP = Portfolio_via_CVXPY(
    cvxpy_fun=EWP,
    portfolio_params=dict(name="1/N")
)
portf_GMRP = Portfolio_via_CVXPY(
    cvxpy_fun=GMRP,
    portfolio_params=dict(name="GMRP")
)
portf_QuintP_mu = Portfolio_via_CVXPY(
    cvxpy_fun=QuintP_mu,
    portfolio_params=dict(name="QuintP (sorted by mu)")
)
portf_QuintP_mu_over_sigma = Portfolio_via_CVXPY(
    cvxpy_fun=QuintP_mu_over_sigma,
    portfolio_params=dict(name="QuintP (sorted by mu/sigma)")
)
portf_QuintP_mu_over_sigma2 = Portfolio_via_CVXPY(
    cvxpy_fun=QuintP_mu_over_sigma2,
    portfolio_params=dict(name="QuintP (sorted by mu/sigma2)")
)

all_formulations = [
    portf_EWP,
    portf_GMRP,
    portf_QuintP_mu,
    portf_QuintP_mu_over_sigma,
    portf_QuintP_mu_over_sigma2,
]
n_portfolios = len(all_formulations)

# WalkForward backtest
bt_list = batch_cross_val_predict(
    all_formulations,
    X=prices_to_returns(stock_prices),  #or: stock_prices.pct_change(fill_method=None)
    cv=WalkForward(test_size=30, train_size=T_trn, purged_size=1),
)
bt_population = Population(bt_list)

# Extract last portfolio composition of each portfolio model
last_portfolios = Population([
    setattr(mpp[-1], 'name', mpp.name) or mpp[-1]
    for mpp in bt_population
])

# Plot stacked barplot
last_portfolios.plot_composition()

# Plot side-by-side barplot
plot_composition_sidebyside(last_portfolios)

# Portfolio NAV (Net Asset Value)
bt_population.set_portfolio_params(compounded=True)
nav_df = bt_population.cumulative_returns_df()
nav_df.head()

# Summary
bt_population.summary().loc[[
    'Annualized Sharpe Ratio',
    'CVaR at 95%',
    'MAX Drawdown',
    'Annualized Mean',
    'Annualized Standard Deviation'
]].T

# Plot NAV
fig = bt_population.plot_cumulative_returns()
fig.update_layout(title="Wealth", xaxis_title=None, yaxis_title="Net Asset Value (NAV)")
fig

# Plot drawdowns
bt_population.plot_drawdowns()

# 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

from skfolio.optimization import InverseVolatility, MeanRisk, ObjectiveFunction
from skfolio import RiskMeasure
from skfolio import Population
from skfolio.model_selection import cross_val_predict, WalkForward
from skfolio.preprocessing import prices_to_returns


# Wrap each portfolio within skfolio
portf_GMVP = Portfolio_via_CVXPY(
    cvxpy_fun=GMVP,
    portfolio_params=dict(name="GMVP")
)
portf_GMVP_skfolio = MeanRisk(
    objective_function=ObjectiveFunction.MINIMIZE_RISK,
    risk_measure=RiskMeasure.VARIANCE,
    portfolio_params=dict(name="GMVP (skfolio)")
)

portf_IVP = Portfolio_via_CVXPY(
    cvxpy_fun=IVP,
    portfolio_params=dict(name="IVP")
)
portf_IVP_skfolio = InverseVolatility(
    portfolio_params=dict(name="IVP (skfolio)")
)

portf_RPP = Portfolio_via_CVXPY(
    cvxpy_fun=RPP,
    portfolio_params=dict(name="RPP")
)
portf_MDP = Portfolio_via_CVXPY(
    cvxpy_fun=MDP,
    portfolio_params=dict(name="MDP")
)
portf_MDCP = Portfolio_via_CVXPY(
    cvxpy_fun=MDCP,
    portfolio_params=dict(name="MDCP")
)

all_formulations = [
    portf_GMVP,
    #portf_GMVP_skfolio,
    portf_IVP,
    #portf_IVP_skfolio,
    portf_RPP,
    portf_MDP,
    portf_MDCP,
]
n_portfolios = len(all_formulations)

# WalkForward backtest
bt_list = batch_cross_val_predict(
    all_formulations,
    X=prices_to_returns(stock_prices),  #or: stock_prices.pct_change(fill_method=None)
    cv=WalkForward(test_size=30, train_size=T_trn, purged_size=1),
)
bt_population = Population(bt_list)

# Extract last portfolio composition of each portfolio model
last_portfolios = Population([
    setattr(mpp[-1], 'name', mpp.name) or mpp[-1]
    for mpp in bt_population
])

# Plot side-by-side barplot
plot_composition_sidebyside(last_portfolios, asset_names=stock_prices.columns)

# Summary
bt_population.summary().loc[[
    'Annualized Sharpe Ratio',
    'CVaR at 95%',
    'MAX Drawdown',
    'Annualized Mean',
    'Annualized Standard Deviation'
]].T

# Plot NAV
bt_population.set_portfolio_params(compounded=True)
fig = bt_population.plot_cumulative_returns()
fig.update_layout(title="Wealth", xaxis_title=None, yaxis_title="Net Asset Value (NAV)")
fig

# Plot drawdowns
bt_population.plot_drawdowns()