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

# 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

stock_prices = SP500_stocks_2015to2020[
                   ["AAPL", "AMZN", "AMD", "GM", "GOOGL", "MGM", "MSFT", "QCOM", "TSCO", "UPS"]
               ].loc["2018":]

# Display first and last rows
print("First few rows of stock prices:")
display(stock_prices.head())
print("\nLast few rows of stock prices:")
display(stock_prices.tail())

# Plot log-prices
plt.figure(figsize=(12, 6))
log_prices = np.log(stock_prices)
for col in log_prices.columns:
    plt.plot(log_prices.index, log_prices[col], linewidth=1, label=col)
plt.title('Log-prices')
plt.legend(title='stocks')
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

# Define portfolio functions

def one_over_N(X: pd.DataFrame) -> np.ndarray:
    """Equal weight portfolio (1/N)"""
    N = X.shape[1]
    return np.repeat(1/N, N)

def design_GMVP(Sigma: np.ndarray) -> np.ndarray:
    """Design Global Minimum Variance Portfolio"""
    N = Sigma.shape[0]
    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 GMVP(X: pd.DataFrame) -> np.ndarray:
    """Global Minimum Variance Portfolio"""
    # Convert to log returns for covariance estimation
    X_log = np.log(1 + X)
    Sigma = X_log.cov().values
    return design_GMVP(Sigma)

def design_IVP(sigma: np.ndarray) -> np.ndarray:
    """Design Inverse Volatility Portfolio"""
    w = 1/sigma
    w = w/np.sum(w)
    return w

def IVP(X: pd.DataFrame) -> np.ndarray:
    """Inverse Volatility Portfolio"""
    # Convert to log returns for volatility estimation
    X_log = np.log(1 + X)
    sigma = X_log.std()
    return design_IVP(sigma.values)

# Compute returns
X_lin = stock_prices.pct_change().dropna()
X_log = np.log(stock_prices).diff().dropna()
# or: X_log = np.log(1 + X_lin)

# Split data into training and test
T, N = X_lin.shape
T_trn = round(0.50*T)
X_lin_trn = X_lin.iloc[:T_trn]
X_lin_tst = X_lin.iloc[T_trn:]
X_log_trn = X_log.iloc[:T_trn]
X_log_tst = X_log.iloc[T_trn:]

# Estimate mu and Sigma with training data
mu = X_log_trn.mean().values
Sigma = X_log_trn.cov().values

# Design portfolios
w_one_over_N = np.repeat(1/N, N)
# or: w_one_over_N = one_over_N(X_log_trn)
w_GMVP = design_GMVP(Sigma)
w_IVP = design_IVP(np.sqrt(np.diag(Sigma)))

# Backtest portfolios with test data (assuming daily rebalancing and no transaction cost)
ret_one_over_N = X_lin_tst @ w_one_over_N
ret_GMVP = X_lin_tst @ w_GMVP
ret_IVP = X_lin_tst @ w_IVP

# Compute cumulative returns
wealth_one_over_N = (1 + ret_one_over_N).cumprod() - 1
wealth_GMVP = (1 + ret_GMVP).cumprod() - 1
wealth_IVP = (1 + ret_IVP).cumprod() - 1

# Plot cumulative returns
plt.figure(figsize=(12, 6))
wealth_one_over_N.plot(label='1/N', linewidth=1)
wealth_GMVP.plot(label='GMVP', linewidth=1)
wealth_IVP.plot(label='IVP', linewidth=1)
plt.title('Cumulative returns')
plt.legend(title='Portfolios')
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

def backtest(
        portfolio_funcs: Dict[str, callable],
        prices: pd.DataFrame,
        lookback: int,
        rebalance_every: int = 1,
        optimize_every: int = 1,
        cost_bps: float = 0,
        verbose: bool = True
) -> 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():
        if verbose:
            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:
                # Reoptimize portfolio if necessary
                if (t - lookback) % optimize_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 calculate_performance_metrics(returns: pd.DataFrame) -> Dict[str, float]:
    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)

    return {
        'Annualized Return': annualized_return,
        'Annualized Volatility': annualized_volatility,
        'Sharpe Ratio': sharpe_ratio,
        'Maximum Drawdown': max_drawdown
    }

def print_table_performance_metrics(returns: pd.DataFrame):
    performance_table = pd.DataFrame(calculate_performance_metrics(returns))

    # 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_weights_over_time(weights: pd.DataFrame, portfolio_name: str):
    """Plot the evolution of portfolio weights over time as stacked area"""

    # Convert index to datetime if it's not already
    weights.index = pd.to_datetime(weights.index)

    # Convert data to numeric, replacing any non-numeric values with NaN
    weights = weights.apply(pd.to_numeric, errors='coerce')

    # Fill NaN values with 0 (or another appropriate value)
    weights = weights.fillna(0)

    # Create figure
    fig, ax = plt.subplots(figsize=(14, 6))

    # Get dates for x-axis
    dates = weights.index

    # Create bottom offset for stacking
    bottom = np.zeros(len(dates))

    # Use a colormap
    colors = plt.colormaps['tab20'](np.linspace(0, 1, len(weights.columns)))

    # Plot each asset as a segment of the stacked area
    for i, asset in enumerate(weights.columns):
        ax.fill_between(dates, bottom, bottom + weights[asset].values,
                        label=asset, color=colors[i])
        bottom += weights[asset].values

    # Format plot
    ax.set_title(f"Weight allocation over time for portfolio {portfolio_name}")
    ax.set_ylabel("weight")
    ax.set_ylim(0, 1.0)  # Set y-axis from 0 to 1
    ax.legend(title="stocks", bbox_to_anchor=(1.05, 1), loc='upper left')

    # Format x-axis
    ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2))
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))
    ax.xaxis.set_minor_locator(mdates.WeekdayLocator())
    plt.xticks(rotation=45)

    plt.tight_layout()
    plt.show()

# This will perform the backtest with daily rebalancing and no transaction costs
bt_ret_daily, bt_w_daily = backtest(
    {
        "1/N": one_over_N,
        "GMVP": GMVP,
        "IVP": IVP
    },
    stock_prices,
    lookback=T_trn,
    rebalance_every=1,
    optimize_every=10000,  # only optimize once at the beginning
    cost_bps=0
)

# Sanity check - compare with manual calculation
print("Checking if results match manual calculation:")
print(f"1/N portfolio: {np.allclose(ret_one_over_N.values, bt_ret_daily['1/N'].values, rtol=1e-5)}")
print(f"GMVP portfolio: {np.allclose(ret_GMVP.values, bt_ret_daily['GMVP'].values, rtol=1e-5)}")
print(f"IVP portfolio: {np.allclose(ret_IVP.values, bt_ret_daily['IVP'].values, rtol=1e-5)}")

# Plot cumulative returns
plot_cum_return(bt_ret_daily)

# Display performance metrics
print("Performance metrics with daily rebalancing and no transaction costs:")
print_table_performance_metrics(bt_ret_daily)

# Plot drawdowns
plot_drawdown(bt_ret_daily)

# Realistic rebalancing including transaction costs
# Rebalance every 5 days with 60 bps transaction costs

bt_ret_realistic, bt_w_realistic = backtest(
    {
        "1/N": one_over_N,
        "GMVP": GMVP,
        "IVP": IVP
    },
    stock_prices,
    lookback=T_trn,
    rebalance_every=5,
    optimize_every=10000,  # only optimize once at the beginning
    cost_bps=60
)

# Display performance metrics
print("Performance metrics with realistic rebalancing:")
print_table_performance_metrics(bt_ret_realistic)

# Plot cumulative returns
plot_cum_return(bt_ret_realistic)

# Plot drawdowns
plot_drawdown(bt_ret_realistic)

# Walk-forward backtesting
# Reoptimize every 20 days, rebalance every 5 days with transaction costs

bt_ret_wf, bt_w_wf = backtest(
    {
        "1/N": one_over_N,
        "GMVP": GMVP,
        "IVP": IVP
    },
    stock_prices,
    lookback=T_trn,
    rebalance_every=20,
    optimize_every=20,
    cost_bps=60
)

# Display performance metrics
print("Performance metrics with realistic rebalancing:")
print_table_performance_metrics(bt_ret_realistic)

# Plot cumulative returns
plot_cum_return(bt_ret_realistic)

# Plot drawdowns
plot_drawdown(bt_ret_realistic)

# Plot weights for each portfolio
plot_weights_over_time(bt_w_wf["1/N"], portfolio_name="1/N")
plot_weights_over_time(bt_w_wf["GMVP"], portfolio_name="GMVP")
plot_weights_over_time(bt_w_wf["IVP"], portfolio_name="IVP")

# Function so resample data
def financial_data_resample(
    prices: pd.DataFrame,
    num_datasets: int = 200,
    N_sample: int = 8,
    T_sample: int = 252*2
) -> List[pd.DataFrame]:
    """
    Resample financial data multiple times

    Args:
        prices: DataFrame of stock prices
        num_datasets: Number of resampled datasets to create
        N_sample: Number of stocks to sample in each dataset
        T_sample: Number of time periods to sample in each dataset

    Returns:
        List of resampled price DataFrames
    """
    resampled_datasets = []

    for i in range(num_datasets):
        # Randomly select stocks
        selected_stocks = np.random.choice(prices.columns, size=N_sample, replace=False)

        # Randomly select a starting point
        max_start = len(prices) - T_sample
        if max_start <= 0:
            start_idx = 0
        else:
            start_idx = np.random.randint(0, max_start)

        # Extract the subset
        subset = prices.iloc[start_idx:start_idx+T_sample, :][selected_stocks]
        resampled_datasets.append(subset)

    return resampled_datasets

# Resample data
stock_prices_2017 = SP500_stocks_2015to2020.loc["2017":]
# stock_prices_2017 = SP500_stocks_2015to2020[
#     ["AAPL", "AMZN", "AMD", "GM", "GOOGL", "MGM", "MSFT", "QCOM", "TSCO", "UPS"]
# ].loc["2017":]

stock_prices_resampled = financial_data_resample(
    stock_prices_2017,
    num_datasets=200,
    N_sample=8,
    T_sample=252*2
)

len(stock_prices_resampled)

stock_prices_resampled[0].shape

stock_prices_resampled[0].head()

# Function to run multiple backtests
def run_multiple_backtests(
    portfolio_funcs: Dict[str, callable],
    price_datasets: List[pd.DataFrame],
    lookback: int = 252,
    optimize_every: int = 20,
    rebalance_every: int = 5,
    cost_bps: float = 60
) -> Dict[str, List[pd.DataFrame]]:
    """
    Run backtests on multiple datasets

    Returns:
        Dictionary with portfolio names as keys and lists of return DataFrames as values
    """
    results = {name: [] for name in portfolio_funcs.keys()}

    for i, prices in enumerate(price_datasets):
        if i % 10 == 0:
            print(f"Processing dataset {i+1}/{len(price_datasets)}...")

        # Run backtest on this dataset
        bt_ret, _ = backtest(
            portfolio_funcs,
            prices,
            lookback=lookback,
            optimize_every=optimize_every,
            rebalance_every=rebalance_every,
            cost_bps=cost_bps,
            verbose=False
        )

        # Store results
        for portfolio_name in portfolio_funcs.keys():
            results[portfolio_name].append(bt_ret[portfolio_name])

    return results

multiple_bt_results = run_multiple_backtests(
    {
        "1/N": one_over_N,
        "GMVP": GMVP,
        "IVP": IVP
    },
    stock_prices_resampled,
    lookback=252,
    optimize_every=20,
    rebalance_every=5,
    cost_bps=60
)

def summarize_multiple_backtests(results: Dict[str, List[pd.Series]]) -> pd.DataFrame:
    """Summarize results from multiple backtests"""
    summary = {}

    for portfolio_name, returns_list in results.items():
        portfolio_metrics = []

        for returns in returns_list:
            metrics = calculate_performance_metrics(pd.DataFrame(returns))
            portfolio_metrics.append(metrics)

        # Calculate mean of metrics across all datasets
        mean_metrics = {}
        for metric in portfolio_metrics[0].keys():
            values = [metrics[metric] for metrics in portfolio_metrics]
            mean_metrics[metric] = np.mean(values)

        summary[portfolio_name] = mean_metrics

    return pd.DataFrame(summary).T

# Generate summary table
summary_table = summarize_multiple_backtests(multiple_bt_results)
display(summary_table.style.format("{:.4f}"))

# Create bar plots for selected metrics
plt.figure(figsize=(12, 6))
summary_table[["Maximum Drawdown", "Annualized Volatility"]].abs().plot(kind='bar')
plt.title("Risk Metrics Comparison")
plt.ylabel("Value")
plt.xticks(rotation=0)
plt.tight_layout()
plt.show()