# Core numerical computing
import numpy as np
import pandas as pd
from typing import Dict, Tuple, List, Callable
import warnings
warnings.filterwarnings('ignore')

# 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 qpsolvers import solve_qp

def backtest(
        portfolio_funcs: Dict[str, callable],
        dataset: Dict[str, 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]]:
    """
    Backtest portfolio strategies for index tracking.

    Parameters:
    portfolio_funcs: Dictionary mapping strategy names to portfolio weight functions
    dataset: Dictionary containing 'prices' (asset prices) and 'index' (benchmark prices)
    lookback: Number of periods to use for estimation
    rebalance_every: Number of periods between rebalancing
    optimize_every: Number of periods between optimization
    cost_bps: Transaction costs in basis points
    verbose: Whether to print progress messages

    Returns:
    Tuple of (returns DataFrame, weights dictionary)
    """
    stock_prices = dataset['prices']
    index_prices = dataset['index'] if dataset['index'] is not None else None
    T, N = stock_prices.shape
    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.zeros(N)
        w = pd.DataFrame(index=stock_prices.index, columns=stock_prices.columns)
        portf_ret = pd.Series(index=stock_prices.index)
        portf_ret.name = portfolio_name

        for t in range(lookback, T):
            # Rebalance portfolio if necessary
            if (t - lookback) % rebalance_every == 0:
                # Reoptimize portfolio if necessary
                if (t - lookback) % optimize_every == 0:
                    new_w = portfolio_func({
                        'prices': stock_prices.iloc[t-lookback:t],
                        'index': index_prices.iloc[t-lookback:t] if index_prices is not None else None
                    })
                    if verbose:
                        print(f"  Reoptimizing portfolio at time {t} (cardinality={sum(new_w > 1e-5)})")
                # Calculate turnover and transaction costs
                turnover = np.abs(new_w - current_w).sum()
                current_w = new_w
                transaction_cost = turnover * cost_bps / 10000
            else:
                transaction_cost = 0

            # Store weights
            w.iloc[t] = current_w

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

            # Update normalized weights based on asset performance
            current_w = current_w * (1 + Xt)
            current_w = current_w / current_w.sum() if current_w.sum() != 0 else current_w

        # Keep a record (remove initial 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

# Download S&P 500 index and SPY ETF data from Yahoo Finance
SP500_prices = yf.download("^GSPC", start="2007-01-01", end=None)['Close']
SPY_prices = yf.download("SPY", start="2007-01-01", end=None)['Close']

# Combine into a single DataFrame
SP500_SPY_prices = SP500_prices.merge(SPY_prices, left_index=True, right_index=True, how='outer')
SP500_SPY_prices.columns = ['S&P 500', 'SPY']

# Plot the data
plt.figure(figsize=(12, 6))
for col in SP500_SPY_prices.columns:
    plt.plot(SP500_SPY_prices.index, SP500_SPY_prices[col], linewidth=1.0, label=col)

plt.yscale('log')
plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
plt.gca().xaxis.set_major_locator(mdates.YearLocator())
plt.gca().xaxis.set_minor_locator(mdates.MonthLocator())
plt.title("Tracking of S&P 500 index")
plt.ylabel("dollars")
plt.xlabel(None)
plt.legend(title=None)
plt.tight_layout()
plt.show()

def sparse_tracking_via_MM(X: pd.DataFrame, r: pd.DataFrame, lambda_val: float,
                           w0: np.ndarray = None, max_iter: int = 5,
                           thres: float = 1e-9, p: float = 1e-3, u: float = 1, reg: float = 1e-7,
                           verbose: bool = False) -> np.ndarray | List:
    """
    Sparse index tracking via MM algorithm.

    Parameters:
    X: DataFrame of asset returns
    r: Series of benchmark returns
    lambda_val: Regularization parameter for sparsity
    w0: Initial weights (optional)
    max_iter: Maximum number of iterations
    thres: Threshold for weight pruning
    p, u: Parameters for log penalty
    verbose: Whether to return objective values

    Returns:
    Optimal weights or dictionary with weights and objective values
    """
    # Convert to numpy arrays
    X_mat = X.values
    r_vec = r.values.reshape(-1, 1)
    T, N = X_mat.shape

    # Initialize weights
    if w0 is None:
        w = np.ones(N) / N
    else:
        w = w0.copy()

    # Initialize objective value tracking if verbose
    obj_values = []
    if verbose:
        obj_val = (1/T) * np.linalg.norm(r_vec - X_mat @ w.reshape(-1, 1))**2 + \
                 lambda_val * np.sum(np.log(1 + np.abs(w)/p)/np.log(1 + u/p))
        obj_values.append(obj_val)

    # MM loop
    for k in range(max_iter):
        #print(f"Iteration {k+1}")

        # Calculate weights for MM update
        d = (1 / (np.log(1 + u/p) * (p + np.abs(w)))).reshape(-1, 1)
        w_prev = w.copy()

        # # Solve the convex subproblem using CVXPY
        # w_var = cp.Variable(N)
        # objective = cp.Minimize((1/T) * cp.sum_squares(r_vec - X_mat @ w_var) +
        #                        lambda_val * cp.sum(cp.multiply(d, cp.abs(w_var))))
        # constraints = [cp.sum(w_var) == 1, w_var >= 0]
        # problem = cp.Problem(objective, constraints)

        # Solve the convex subproblem using quadprog
        P = 2 * (1/T) * X_mat.T @ X_mat
        P += reg * np.mean(np.diag(P)) * np.eye(N)  # Add small regularization
        q = -2 * (1/T) * X_mat.T @ r_vec + lambda_val * d
        G = np.vstack((-np.ones(N), -np.eye(N)))
        h = np.concatenate(([-1], np.zeros(N)))
        A = np.ones((1, N))
        b = np.array([1.0])
        w = solve_qp(P, q, G, h, A, b, solver="quadprog")
        if w is None:
            print("    QP solver failed, increasing the matrix regularization factor...")
            w = w_prev
            reg *= 2
            continue

        # Track objective value if verbose
        if verbose:
            obj_val = (1/T) * np.linalg.norm(r_vec - X_mat @ w.reshape(-1, 1))**2 + \
                     lambda_val * np.sum(np.log(1 + np.abs(w)/p)/np.log(1 + u/p))
            obj_values.append(obj_val)

        # Check convergence
        if np.linalg.norm(w - w_prev) / np.linalg.norm(w_prev) < 1e-3:
            break

    # Threshold small weights
    w[w < thres] = 0

    # Normalize weights to sum to 1
    if np.sum(w) > 0:
        w = w / np.sum(w)

    if verbose:
        return w, obj_values
    else:
        return w

X = SP500_stocks_2015to2020.iloc[:, :10].pct_change().dropna()
r = SP500_index_2015to2020.pct_change().dropna()

sparse_tracking_via_MM(X.iloc[:, :10], r, lambda_val=5e-7)

def compute_alpha(X, r):
    """
    Compute alpha values for time-varying portfolio weights.

    Parameters:
    X: DataFrame of asset returns
    r: Series of benchmark returns

    Returns:
    Dictionary containing alpha values and related calculations
    """
    X_cum = (1 + X).cumprod()
    r_cum = (1 + r).cumprod()

    # Convert r_cum to array for division
    r_cum_array = r_cum.values.reshape(-1, 1)

    # Calculate alpha
    alpha = X_cum.divide(r_cum_array)

    # Calculate lagged alpha (shift by 1)
    alpha_lagged = alpha.shift(1)
    alpha_lagged.iloc[0] = 1

    # Calculate X_tilde
    X_tilde = X.multiply(alpha_lagged.values)

    return {
        'alpha': alpha,
        'alpha_lagged': alpha_lagged,
        'last_alpha_lagged': alpha_lagged.iloc[-1].values,
        'X_tilde': X_tilde
    }

def tracking_using_plain_returns_wfixed(dataset, lambda_val=5e-7):
    """Index tracking assuming fixed portfolio weights."""
    # Calculate returns
    X = dataset['prices'].pct_change()
    r = dataset['index'].pct_change()

    # Remove NA's
    common_valid_indices = X.dropna().index.intersection(r.dropna().index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Get optimal weights
    w = sparse_tracking_via_MM(X, r, lambda_val=lambda_val, max_iter=10)
    return w

def tracking_using_plain_returns_wt(dataset, lambda_val=5e-7):
    """Index tracking assuming time-varying portfolio weights."""
    # Calculate returns
    X = dataset['prices'].pct_change()
    r = dataset['index'].pct_change()

    # Remove NA's
    common_valid_indices = X.dropna().index.intersection(r.dropna().index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Compute alpha
    alpha = compute_alpha(X, r)

    # Get optimal weights
    w0 = sparse_tracking_via_MM(alpha['X_tilde'], r, lambda_val=lambda_val, max_iter=10)

    # Adjust weights by last alpha
    w = w0 * alpha['last_alpha_lagged']

    # Normalize weights
    if np.sum(w) > 0:
        w = w / np.sum(w)
    return w

w = tracking_using_plain_returns_wfixed(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    })
sum(w), min(w)

w = tracking_using_plain_returns_wt(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    })
sum(w), min(w)

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Run backtest
returns, weights = backtest(
    portfolio_funcs={
        "Tracking assuming fixed portfolio": tracking_using_plain_returns_wfixed,
        "Tracking assuming time-varying portfolio": tracking_using_plain_returns_wt
    },
    dataset={
        'prices': SP500_stocks_2015to2020,  #.loc["2018":]
        'index': SP500_index_2015to2020  # .loc["2018":]
    },
    lookback=T_trn,
    rebalance_every=6*21,  # Rebalance every 6 months (assuming 21 trading days per month)
    optimize_every=6*21,
)

def compute_tracking_error(returns: pd.DataFrame, benchmark: pd.DataFrame) -> pd.DataFrame:
    benchmark_rets = benchmark.pct_change().dropna()

    # Align the indices
    common_idx = returns.index.intersection(benchmark_rets.index)
    returns_aligned = returns.loc[common_idx]
    benchmark_rets_aligned = benchmark_rets.loc[common_idx]

    # Calculate tracking error
    tracking_errors = (returns_aligned.subtract(benchmark_rets_aligned.iloc[:, 0], axis=0)) ** 2

    return tracking_errors.mean()

def plot_tracking_error(returns: pd.DataFrame, benchmark: pd.DataFrame, window: int = 252) -> pd.DataFrame:
    """
    Plot tracking error over time.

    Parameters:
    returns: DataFrame of portfolio returns
    benchmark: Series of benchmark prices (will be converted to returns)
    window: Rolling window size for smoothing
    """
    benchmark_rets = benchmark.pct_change().dropna()

    # Align the indices
    common_idx = returns.index.intersection(benchmark_rets.index)
    returns_aligned = returns.loc[common_idx]
    benchmark_rets_aligned = benchmark_rets.loc[common_idx]

    # Calculate tracking error
    tracking_errors = (returns_aligned.subtract(benchmark_rets_aligned.iloc[:, 0], axis=0)) ** 2

    # Apply rolling mean
    smoothed_te = tracking_errors.rolling(window=window).mean()

    plt.figure(figsize=(12, 6))
    smoothed_te.plot()
    plt.title(f'Tracking Error (Rolling {window}-day window)')
    plt.xlabel(None)
    plt.ylabel('Tracking Error')
    plt.legend(title="Portfolio")
    plt.grid(True)
    plt.show()

    return tracking_errors

TE = plot_tracking_error(returns, benchmark=SP500_index_2015to2020, window=252)

def compute_cardinality(weights: Dict[str, pd.DataFrame]) -> pd.DataFrame:
    """Plot portfolio cardinality (number of assets with non-zero weights)."""
    cardinality = {}

    for name, w in weights.items():
        cardinality[name] = (w > 1e-5).sum(axis=1)

    cardinality_df = pd.DataFrame(cardinality)

    return cardinality_df.mean()

def plot_cardinality(weights: Dict[str, pd.DataFrame]) -> pd.DataFrame:
    """Plot portfolio cardinality (number of assets with non-zero weights)."""
    cardinality = {}

    for name, w in weights.items():
        cardinality[name] = (w > 1e-5).sum(axis=1)

    cardinality_df = pd.DataFrame(cardinality)

    plt.figure(figsize=(12, 6))
    cardinality_df.plot()
    plt.title('Portfolio Cardinality Over Time')
    plt.xlabel(None)
    plt.ylabel('Number of Assets')
    plt.legend(title="Portfolio")
    plt.grid(True)
    plt.show()

    return cardinality_df

cardinality = plot_cardinality(weights)

# Calculate column means for TE and cardinality
summary_df = pd.DataFrame({
    "TE": TE.mean(),
    "cardinality": cardinality.mean()
})

# Display
summary_df.style.format({
    'TE': '{:.2e}',  # Scientific notation with 6 decimal places
    'cardinality': '{:.2f}'  # Fixed-point with 2 decimal places
})

def tracking_using_cumreturns_wfixed(dataset, lambda_val=20e-7):
    """Index tracking using cumulative returns."""
    # Calculate returns
    X = dataset['prices'].pct_change()
    r = dataset['index'].pct_change()

    # Remove NA's
    common_valid_indices = X.dropna().index.intersection(r.dropna().index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Calculate cumulative returns
    X_cum = X.cumsum()
    r_cum = r.cumsum()

    # Get optimal weights
    w = sparse_tracking_via_MM(X_cum, r_cum, lambda_val=lambda_val, max_iter=10)
    return w


def tracking_using_cumreturns_wt(dataset, lambda_val=15e-7):
    """Index tracking assuming time-varying portfolio weights."""
    # Calculate returns
    X = dataset['prices'].pct_change()
    r = dataset['index'].pct_change()

    # Remove NA's
    common_valid_indices = X.dropna().index.intersection(r.dropna().index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Compute alpha
    alpha = compute_alpha(X, r)

    # Calculate cumulative returns
    X_cum = alpha['X_tilde'].cumsum()
    r_cum = r.cumsum()

    # Get optimal weights
    w0 = sparse_tracking_via_MM(X_cum, r_cum, lambda_val=lambda_val, max_iter=10)

    # Adjust weights by last alpha
    w = w0 * alpha['last_alpha_lagged']

    # Normalize weights
    if np.sum(w) > 0:
        w = w / np.sum(w)
    return w

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Run backtest
returns, weights = backtest(
    portfolio_funcs={
        "Tracking using plain returns": tracking_using_plain_returns_wfixed,
        "Tracking using cumulative returns": tracking_using_cumreturns_wfixed
    },
    dataset={
        'prices': SP500_stocks_2015to2020,  #.loc["2018":]
        'index': SP500_index_2015to2020  # .loc["2018":]
    },
    lookback=T_trn,
    rebalance_every=6*21,  # Rebalance every 6 months (assuming 21 trading days per month)
    optimize_every=6*21,
)

def plot_prices(returns: pd.DataFrame, benchmark: pd.DataFrame) -> None:
    NAV = (1 + returns).cumprod()
    benchmark = benchmark.loc[NAV.index]
    B0 = benchmark.iloc[0,0]
    all_prices = pd.concat([benchmark, B0 * NAV], axis=1)

    # Plot
    plt.figure(figsize=(12, 6))
    all_prices.plot()
    plt.title(f'Tracking of S&P 500 index')
    plt.xlabel(None)
    plt.ylabel('dollars')
    plt.legend(title="Portfolio")
    plt.grid(True)
    plt.show()

plot_prices(returns, benchmark=SP500_index_2015to2020)

# Prep data
X = SP500_stocks_2015to2020.pct_change().dropna()
r = SP500_index_2015to2020.pct_change().dropna()

# Run index tracking MM algorithmn
w, obj_values = sparse_tracking_via_MM(X, r, lambda_val=5e-7, max_iter=10, verbose=True)

# Plot convergence over iterations
last_iter = len(obj_values) - 1
iterations = list(range(last_iter + 1))

plt.figure(figsize=(10, 6))
plt.plot(iterations, obj_values, linewidth=1.2, color="blue")
plt.title("Convergence over iterations")
plt.xlabel("k")
plt.ylabel("Objective Value")
plt.xticks(ticks=iterations)
plt.grid(True)
plt.show()

def full_tracking(dataset):
    """Full index tracking (dense portfolio)."""
    # Calculate returns
    X = dataset['prices'].pct_change().dropna()
    r = dataset['index'].pct_change().dropna()

    # Remove NA's
    common_valid_indices = X.index.intersection(r.index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Convert to numpy arrays
    X_mat = X.values
    r_vec = r.values.reshape(-1, 1)
    T, N = X_mat.shape

    # # Solve the convex subproblem using CVXPY
    # N = X_mat.shape[1]
    # w = cp.Variable(N)
    # objective = cp.Minimize(cp.sum_squares(X_mat.values @ w - r_vec.values))
    # constraints = [cp.sum(w) == 1, w >= 0]

    # Solve the convex subproblem using quadprog
    P = 2 * (1 / T) * X_mat.T @ X_mat  # Quadratic term matrix
    P += 1e-7 * np.mean(np.diag(P)) * np.eye(N)  # Regularization to ensure positive definiteness
    q = -2 * (1 / T) * X_mat.T @ r_vec  # Linear term vector
    G = -np.eye(N)  # Non-negativity constraint (w >= 0)
    h = np.zeros(N)
    A = np.ones((1, N))  # Equality constraint (sum(w) == 1)
    b = np.array([1.0])
    w = solve_qp(P, q.flatten(), G, h, A, b, solver="quadprog")
    return w

w = full_tracking(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    })
sum(w), min(w)

def two_stage_naive_tracking(dataset, K=20):
    """Two-stage tracking with proportional weights."""
    # Calculate returns
    X = dataset['prices'].pct_change().dropna()
    r = dataset['index'].pct_change().dropna()

    # Remove NA's
    common_valid_indices = X.index.intersection(r.index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Convert to numpy arrays
    X_mat = X.values
    r_vec = r.values.reshape(-1, 1)
    T, N = X_mat.shape

    # Step 1: Get dense portfolio
    P = 2 * (1 / T) * X_mat.T @ X_mat  # Quadratic term matrix
    P += 1e-7 * np.mean(np.diag(P)) * np.eye(N)  # Regularization to ensure positive definiteness
    q = -2 * (1 / T) * X_mat.T @ r_vec  # Linear term vector
    G = -np.eye(N)  # Non-negativity constraint (w >= 0)
    h = np.zeros(N)
    A = np.ones((1, N))  # Equality constraint (sum(w) == 1)
    b = np.array([1.0])
    b = solve_qp(P, q.flatten(), G, h, A, b, solver="quadprog")

    # Step 2: Keep K largest weights
    idx_Klargest = np.argsort(b)[-K:]

    w_sparse = np.zeros(N)
    w_sparse[idx_Klargest] = b[idx_Klargest]

    # Normalize weights
    if np.sum(w_sparse) > 0:
        w_sparse = w_sparse / np.sum(w_sparse)

    return w_sparse

w = two_stage_naive_tracking(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    }, K=10)
sum(w), min(w), sum(w > 1e-5)

def two_stage_refitted_tracking(dataset, K=20):
    """Two-stage tracking with refitted weights."""
    # Calculate returns
    X = dataset['prices'].pct_change().dropna()
    r = dataset['index'].pct_change().dropna()

    # Remove NA's
    common_valid_indices = X.index.intersection(r.index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Convert to numpy arrays
    X_mat = X.values
    r_vec = r.values.reshape(-1, 1)
    T, N = X_mat.shape

    # Step 1: Get dense portfolio
    P = 2 * (1 / T) * X_mat.T @ X_mat  # Quadratic term matrix
    P += 1e-7 * np.mean(np.diag(P)) * np.eye(N)  # Regularization to ensure positive definiteness
    q = -2 * (1 / T) * X_mat.T @ r_vec  # Linear term vector
    G = -np.eye(N)  # Non-negativity constraint (w >= 0)
    h = np.zeros(N)
    A = np.ones((1, N))  # Equality constraint (sum(w) == 1)
    b = np.array([1.0])
    b = solve_qp(P, q.flatten(), G, h, A, b, solver="quadprog")

    # Step 2: Keep K largest weights and reoptimize
    idx_Klargest = np.argsort(b)[-K:]
    X_selected = X_mat[:, idx_Klargest]
    P = 2 * (1 / T) * X_selected.T @ X_selected  # Quadratic term matrix
    P += 1e-7 * np.mean(np.diag(P)) * np.eye(K)  # Regularization to ensure positive definiteness
    q = -2 * (1 / T) * X_selected.T @ r_vec  # Linear term vector
    G = -np.eye(K)  # Non-negativity constraint (w >= 0)
    h = np.zeros(K)
    A = np.ones((1, K))  # Equality constraint (sum(w) == 1)
    b = np.array([1.0])
    wK = solve_qp(P, q.flatten(), G, h, A, b, solver="quadprog")

    # Create sparse weight vector
    w_sparse = np.zeros(N)
    w_sparse[idx_Klargest] = wK

    return w_sparse

w = two_stage_refitted_tracking(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    }, K=15)
sum(w), min(w), sum(w > 1e-5)

def iterative_reweighted_l1_norm_tracking(dataset, lambda_val=5e-7):
    """Index tracking assuming fixed portfolio weights."""
    # Calculate returns
    X = dataset['prices'].pct_change().dropna()
    r = dataset['index'].pct_change().dropna()

    # Keep only overlapping indices
    common_valid_indices = X.index.intersection(r.index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Get optimal weights
    w = sparse_tracking_via_MM(X, r, lambda_val=lambda_val, max_iter=10)
    return w

w = iterative_reweighted_l1_norm_tracking(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    })
sum(w), min(w), sum(w > 1e-5)

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Parameters for the desired sparsity
K = 20
lambda_val=5e-7

# Run backtest
returns, weights = backtest(
    portfolio_funcs={
        "Full tracking": full_tracking,
        "Two-step design with proportional weights": two_stage_naive_tracking,
        "Two-step design with refitted weights": two_stage_refitted_tracking,
        "Iterative reweighted l1-norm method": iterative_reweighted_l1_norm_tracking
    },
    dataset={
        'prices': SP500_stocks_2015to2020,  #.loc["2018":]
        'index': SP500_index_2015to2020  # .loc["2018":]
    },
    lookback=T_trn,
    rebalance_every=6*21,  # Rebalance every 6 months (assuming 21 trading days per month)
    optimize_every=6*21,
)

plot_prices(returns, benchmark=SP500_index_2015to2020)

methods = ["Two-step design with proportional weights",
           "Two-step design with refitted weights",
           "Iterative reweighted l1-norm method"]

TE = plot_tracking_error(returns[methods], benchmark=SP500_index_2015to2020, window=252)

cardinality = plot_cardinality({key: weights[key] for key in methods})

# Calculate column means for TE and cardinality
summary_df = pd.DataFrame({
    "TE": TE.mean(),
    "cardinality": cardinality.mean()
})

# Display
summary_df.style.format({
    'TE': '{:.2e}',  # Scientific notation with 6 decimal places
    'cardinality': '{:.2f}'  # Fixed-point with 2 decimal places
})

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Sweeping loop
lmd_sweep = 1e-5 * np.array([1, 2, 5, 10, 20, 50, 100, 200, 1000]) / 300
df_list = []
for lmd in lmd_sweep:
    print(f"Backtesting for lambda: {lmd:.2e}")

    # Run backtest
    returns, weights = backtest(
        portfolio_funcs={
            "Assuming fixed portfolio": lambda dataset: tracking_using_plain_returns_wfixed(dataset, lambda_val=lmd),
            "Assuming time-varying portfolio": lambda dataset: tracking_using_plain_returns_wt(dataset, lambda_val=lmd)
        },
        dataset={
            'prices': SP500_stocks_2015to2020,
            'index': SP500_index_2015to2020
        },
        lookback=T_trn,
        rebalance_every=6*21,  # Rebalance every 6 months (assuming 21 trading days per month)
        optimize_every=6*21,
        verbose=False
    )

    # Calculate tracking error and cardinality
    TE = compute_tracking_error(returns, benchmark=SP500_index_2015to2020)
    cardinality = compute_cardinality(weights)
    df_list.append(pd.DataFrame({
        "method": TE.index,
        "TE": TE.values,
        "K": cardinality.values,
    }))

# Combine all rows into a single DataFrame
df = pd.concat(df_list, ignore_index=True)

# Plot tracking error versus sparsity level
plt.figure(figsize=(10, 6))
sns.lineplot(data=df, x="K", y="TE", hue="method", style="method", linewidth=1.2)
plt.title("Tracking error versus sparsity level")
plt.xlabel("K (active assets)")
plt.ylabel("Tracking Error")
plt.legend(title=None)
plt.grid(True)
plt.show()

def TE_index_tracking(dataset):
    """Index tracking using cumulative returns."""
    # Calculate returns
    X = dataset['prices'].pct_change().dropna()
    r = dataset['index'].pct_change().dropna()

    # Keep only overlapping indices
    common_valid_indices = X.index.intersection(r.index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Get optimal weights
    w = sparse_tracking_via_MM(X.cumsum(), r.cumsum(), lambda_val=150*6e-7, u=0.5, max_iter=10)
    return w

def DR_index_tracking(dataset):
    """
    Index tracking using downside risk measure with cumulative returns.

    Parameters:
    dataset: Dictionary containing 'prices' (asset prices) and 'index' (benchmark prices)

    Returns:
    Optimal portfolio weights
    """
    # Calculate returns
    X = dataset['prices'].pct_change().dropna()
    r = dataset['index'].pct_change().dropna()

    # Keep only overlapping indices
    common_valid_indices = X.index.intersection(r.index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]
    T, N = X.shape

    # Cumulative returns
    X_cum = X.cumsum()
    r_cum = r.cumsum()

    # Initial weights using cumulative returns
    w = sparse_tracking_via_MM(X_cum, r_cum, lambda_val=150*6e-7, u=0.5, max_iter=10)

    # Iterative reweighting to approximate L1 norm
    for k in range(5):
        # Calculate modified benchmark
        r_cum_tilde = r_cum + np.maximum(0, (X_cum.values @ w).reshape(-1, 1) - r_cum.values)

        # Solve the modified problem
        w = sparse_tracking_via_MM(X_cum, r_cum_tilde, lambda_val=150*6e-7, u=0.5, max_iter=10)  #190*1.8e-7

    return w

def TE1_index_tracking(dataset):
    """L1-norm index tracking using cumulative returns with iterative reweighting."""
    # Calculate returns
    X = dataset['prices'].pct_change().dropna()
    r = dataset['index'].pct_change().dropna()

    # Keep overlap of indices
    common_valid_indices = X.index.intersection(r.index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Cumulative returns
    X_cum = X.cumsum()
    r_cum = r.cumsum()

    # Initial weights using cumulative returns
    w = sparse_tracking_via_MM(X_cum, r_cum, lambda_val=100*6e-7, u=0.5, max_iter=10)

    # Iterative reweighting to approximate L1 norm
    for k in range(5):
        # Calculate weights based on current tracking error
        alpha_sqrt = np.sqrt(1/(2*np.abs((X_cum.values @ w).reshape(-1, 1) - r_cum.values)))

        # Weight the returns by alpha
        X_cum_weighted = X_cum.multiply(alpha_sqrt.reshape(-1, 1))
        r_cum_weighted = r_cum * alpha_sqrt

        # Solve the weighted problem using cumulative returns
        w = sparse_tracking_via_MM(X_cum_weighted, r_cum_weighted, lambda_val=150*6e-7, u=0.5, max_iter=10)

    return w

def HubTE_index_tracking(dataset):
    """L1-norm index tracking using cumulative returns with iterative reweighting."""
    # Calculate returns
    X = dataset['prices'].pct_change().dropna()
    r = dataset['index'].pct_change().dropna()

    # Keep overlap of indices
    common_valid_indices = X.index.intersection(r.index)
    X = X.loc[common_valid_indices]
    r = r.loc[common_valid_indices]

    # Cumulative returns
    X_cum = X.cumsum()
    r_cum = r.cumsum()

    # Initial weights using cumulative returns
    w = sparse_tracking_via_MM(X_cum, r_cum, lambda_val=50*5e-8, u=0.5, max_iter=10) #50*5e-8

    # Iterative reweighting to approximate L1 norm
    M = 1e-4
    for k in range(5):
        # Calculate weights based on current tracking error
        alpha_sqrt = np.sqrt(np.minimum(1, M/np.abs((X_cum.values @ w).reshape(-1, 1) - r_cum.values)))

        # Weight the returns by alpha
        X_cum_weighted = X_cum.multiply(alpha_sqrt.reshape(-1, 1))
        r_cum_weighted = r_cum * alpha_sqrt

        # Solve the weighted problem using cumulative returns
        w = sparse_tracking_via_MM(X_cum_weighted, r_cum_weighted, lambda_val=50*5e-8, u=0.5, max_iter=10)

    return w

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Run backtest
returns, weights = backtest(
    portfolio_funcs={
        "TE based design": TE_index_tracking,
        "Hub-TE based design": HubTE_index_tracking,
        "TE_1 based design": TE1_index_tracking,
        "DR based design": DR_index_tracking
    },
    dataset={
        'prices': SP500_stocks_2015to2020,
        'index': SP500_index_2015to2020
    },
    lookback=T_trn,
    rebalance_every=6*21,
    optimize_every=6*21,
    verbose=True
)

plot_prices(returns, benchmark=SP500_index_2015to2020)

#
# T-Rex index tracking in Pyhon under development.
#
# See R version in: https://portfoliooptimizationbook.com/R-code/R-index-tracking.html
#