Untitled

import pandas as pd
import numpy as np
from dateutil.relativedelta import relativedelta
import math

def get_price_on_date(prices_df, target_date, ticker):
    """
    Returns the last available price on or before target_date for a given ticker.
    """
    df_filtered = prices_df[prices_df['Date'] <= target_date]
    if df_filtered.empty:
        return np.nan
    return df_filtered.iloc[-1][ticker]

def compute_momentum_for_asset(prices_df, current_date, ticker, lookback_months):
    """
    Computes three momentum signals for the given asset and lookback period.
    Returns:
      - total_return (price-based momentum)
      - price_minus_sma (distance from moving average)
      - risk_adjusted (risk-adjusted momentum)
    """
    start_date = current_date - relativedelta(months=lookback_months)
    price_start = get_price_on_date(prices_df, start_date, ticker)
    price_current = get_price_on_date(prices_df, current_date, ticker)
    
    if pd.isna(price_start) or pd.isna(price_current):
        return np.nan, np.nan, np.nan
    
    total_return = (price_current / price_start) - 1
    
    prices_window = prices_df[(prices_df['Date'] >= start_date) & (prices_df['Date'] <= current_date)][ticker]
    sma = prices_window.mean() if not prices_window.empty else np.nan
    price_minus_sma = (price_current / sma) - 1 if (not pd.isna(sma) and sma != 0) else np.nan
    
    numerator = np.log(price_current / price_start)
    prices_period = prices_df[(prices_df['Date'] >= start_date) & (prices_df['Date'] <= current_date)][ticker].sort_index()
    if len(prices_period) < 2:
        risk_adjusted = np.nan
    else:
        log_returns = np.log(prices_period / prices_period.shift(1)).dropna()
        denominator = log_returns.abs().sum()
        risk_adjusted = numerator / denominator if denominator != 0 else np.nan
    
    return total_return, price_minus_sma, risk_adjusted

def compute_aggregated_momentum(prices_df, current_date, ticker, lookback_periods):
    """
    Computes aggregated momentum for an asset over multiple lookback periods.
    Returns:
      - aggregated_score: the average of all momentum signals (across all lookbacks)
      - is_positive: True only if every individual signal is positive.
      - signals: list of individual signals (for inspection)
    """
    signals = []
    for lb in lookback_periods:
        tr, pma, ra = compute_momentum_for_asset(prices_df, current_date, ticker, lb)
        if pd.isna(tr) or pd.isna(pma) or pd.isna(ra):
            return None, None, None  # Incomplete data; skip asset.
        signals.extend([tr, pma, ra])
    agg_score = np.mean(signals)
    # Flag as positive only if every signal is > 0.
    is_positive = all(x > 0 for x in signals)
    return agg_score, is_positive, signals

def backtest_momentum_strategy(prices_df, start_date, end_date, rebalance_frequency, lookback_periods, aum, top_n, risk_on_list, risk_off_list):
    """
    Backtests the long-only momentum strategy with adjusted selection rules.
    """
    prices_df['Date'] = pd.to_datetime(prices_df['Date'])
    prices_df.sort_values('Date', inplace=True)
    
    # Build rebalancing dates
    rebalance_dates = pd.date_range(start=start_date, end=end_date, freq=rebalance_frequency)
    current_aum = aum
    result_records = []
    current_portfolio = {}  # {ticker: (quantity, entry_price)}
    
    for i, reb_date in enumerate(rebalance_dates):
        if i > 0 and current_portfolio:
            portfolio_value = 0
            for ticker, (qty, entry_price) in current_portfolio.items():
                price_today = get_price_on_date(prices_df, reb_date, ticker)
                portfolio_value += qty * price_today
            current_aum = portfolio_value
        
        # Compute momentum scores for risk_on assets
        all_risk_on = []
        for ticker in risk_on_list:
            agg_score, is_positive, _ = compute_aggregated_momentum(prices_df, reb_date, ticker, lookback_periods)
            if agg_score is not None:
                all_risk_on.append((ticker, agg_score, is_positive))
        
        # Separate assets with strictly positive momentum
        positive_risk_on = [(ticker, score) for ticker, score, pos in all_risk_on if pos]
        # Sort lists by aggregated score descending
        positive_risk_on_sorted = sorted(positive_risk_on, key=lambda x: x[1], reverse=True)
        all_risk_on_sorted = sorted(all_risk_on, key=lambda x: x[1], reverse=True)
        
        # Selection logic based on your relaxed conditions:
        if len(positive_risk_on_sorted) >= 6:
            # Sufficient positive risk_on assets; take top 6.
            final_selection = [ticker for ticker, score in positive_risk_on_sorted[:6]]
        elif len(positive_risk_on_sorted) >= 4:
            # Between 4 and 5 positive risk_on; include all of them plus both risk_off assets.
            final_selection = [ticker for ticker, score in positive_risk_on_sorted] + risk_off_list
        else:
            # Fewer than 4 positive risk_on assets: ignore sign and select the top 4 risk_on.
            if len(all_risk_on_sorted) < 4:
                selected = [ticker for ticker, score, pos in all_risk_on_sorted]
                needed = 4 - len(selected)
                final_selection = selected + risk_off_list[:needed]
            else:
                final_selection = [ticker for ticker, score, pos in all_risk_on_sorted[:4]]
        
        # Allocate equal weight among the selected assets
        allocation = current_aum / len(final_selection) if final_selection else 0
        positions = {}
        entry_prices = {}
        for ticker in final_selection:
            price_at_entry = get_price_on_date(prices_df, reb_date, ticker)
            qty = allocation / price_at_entry if price_at_entry != 0 else 0
            positions[ticker] = (qty, price_at_entry)
            entry_prices[ticker] = price_at_entry
        
        record = {
            'Rebalance Date': reb_date,
            'Final AUM': current_aum,
            'Selected Assets': final_selection,
            'Quantities': [positions[ticker][0] for ticker in final_selection],
            'Entry Prices': [positions[ticker][1] for ticker in final_selection]
        }
        result_records.append(record)
        current_portfolio = positions.copy()
    
    result_df = pd.DataFrame(result_records)
    return result_df