import bt

def compare_momentum_strategies(data: pd.DataFrame, momentum_days_list: list,
                                strategy: Callable[[pd.DataFrame, int], bt.Backtest],
                                baselines: bt.Backtest | list[bt.Backtest],
                                top_percent: float = 0.1) -> None:
    """
    Compares multiple momentum strategies with different momentum windows.

    Args:
        data (pd.DataFrame): Historical price data for the tickers.
        momentum_days_list (list): List of integers representing different momentum windows (in trading days).
        strategy (Callable): A function that takes the data and momentum_days and returns a bt.Backtest object.
        baselines (bt.Backtest | list[bt.Backtest]): A backtest to choose as a baseline
        top_percent (float): The percentage of top momentum stocks to select. Defaults to 0.1 (top 10%).

    Returns:
        None
    """
    # Create backtests for each momentum window using the provided strategy function
    backtests = [strategy(data=data, momentum_days=days, top_percent=top_percent) for days in momentum_days_list]

    match baselines:
        case bt.Backtest():
            backtests.append(baselines)
        case list():
            for baseline in baselines:
                backtests.append(baseline)
        case _:
            print("You passed the wrong datatype for baselines, ya dingus!")

    # Run all backtests
    res = bt.run(*backtests)

    # Plot the performance
    res.plot(title=f"Comparison of Strategies vs. Baselines")

    # Display performance statistics
    res.display()

    return res

import json
    from plotting import backtest_plot

    # Example usage:
    tickers = "aapl,msft,goog,amzn,tsla,nflx,meta,spy"
    start_date = "2020-01-01"
    momentum_days_list = [1, 5, 10, 21, 42]

    # Fetch the data
    data = get_data(tickers, start_date)
    backtest_baseline = baseline_strategy()

    # Compare different momentum strategies
    for strat, title in zip(
        [momentum_strategy, volatility_adjusted_momentum_strategy],
        ["Simple Momentum", "Volatility-Adjusted Momentum"],
    ):
        res = compare_momentum_strategies(
            data, momentum_days_list, strategy=strat, baselines=backtest_baseline
        )
        fig = backtest_plot(res, f"Cumulative Returns: {title}")

def momentum_strategy(tickers: str, start_date: str, momentum_days: int, top_percent: float = 0.1) -> bt.Backtest:
    """
    Creates a momentum strategy backtest with a specified momentum window.

    Args:
        tickers (str): Comma-separated list of tickers to include in the backtest.
        start_date (str): The start date for downloading stock data (YYYY-MM-DD).
        momentum_days (int): The number of trading days to calculate momentum over.
        top_percent (float): The percentage of top momentum stocks to select. Defaults to 0.1 (top 10%).

    Returns:
        bt.Backtest: The backtest for the momentum strategy.
    """
    # Download historical price data
    data = bt.get(tickers, start=start_date)

    # Calculate momentum (percentage change over the specified window)
    momentum = data.pct_change(momentum_days).shift(1)

    # Select top momentum stocks based on the specified top percentage
    signal = momentum.rank(axis=1, pct=True) > (1 - top_percent)

    # Strategy: select top momentum stocks, weigh them equally, and rebalance monthly
    strategy = bt.Strategy(
        f'{momentum_days}_day_momentum',
        [
            bt.algos.SelectWhere(signal),
            bt.algos.WeighEqually(),
            bt.algos.Rebalance()
        ]
    )

    return bt.Backtest(strategy, data)

def volatility_adjusted_momentum_strategy(data: pd.DataFrame, momentum_days: int, volatility_days: int = 21, top_percent: float = 0.1) -> bt.Backtest:
    """
    Creates a volatility-adjusted momentum strategy backtest with a specified momentum window.

    Args:
        data (pd.DataFrame): Historical price data for the tickers.
        momentum_days (int): The number of trading days to calculate momentum over.
        volatility_days (int): The number of days to calculate volatility over. Defaults to 21 (approximately 1 month).
        top_percent (float): The percentage of top momentum stocks to select. Defaults to 0.1 (top 10%).

    Returns:
        bt.Backtest: The backtest for the volatility-adjusted momentum strategy.
    """
    # Ensure volatility_days is a valid integer
    if not isinstance(volatility_days, int) or volatility_days <= 0:
        raise ValueError("volatility_days must be a positive integer.")

    # Calculate momentum (percentage change over the specified window)
    momentum = data.pct_change(momentum_days).shift(1)

    # Calculate rolling volatility (standard deviation of percentage changes)
    volatility = data.pct_change().rolling(volatility_days).std()

    # Adjust momentum by volatility
    vol_adjusted_momentum = momentum / volatility

    # Select top momentum stocks based on the specified top percentage
    signal = vol_adjusted_momentum.rank(axis=1, pct=True) > (1 - top_percent)

    # Strategy: select top momentum stocks, weigh them equally, and rebalance monthly
    strategy = bt.Strategy(
        f'{momentum_days}_day_vol_adj_momentum',
        [
            bt.algos.SelectWhere(signal),
            bt.algos.WeighEqually(),
            bt.algos.Rebalance()
        ]
    )

    return bt.Backtest(strategy, data)