import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import itertools # For generating parameter combinations
import warnings
warnings.filterwarnings('ignore') # Ignore warnings for cleaner output

def generate_mock_data(start_date='2020-01-01', end_date='2021-12-31', initial_price=100, volatility=0.01):
    """
    Generates mock stock price data.

    Args:
        start_date (str): Start date for the data.
        end_date (str): End date for the data.
        initial_price (float): Starting price of the stock.
        volatility (float): Daily volatility for price fluctuations.

    Returns:
        pd.DataFrame: DataFrame with 'Date' and 'Adj Close' columns.
    """
    dates = pd.date_range(start=start_date, end=end_date, freq='B') # Business days
    prices = [initial_price]
    for _ in range(1, len(dates)):
        daily_return = np.random.normal(0, volatility)
        prices.append(prices[-1] * (1 + daily_return))

    df = pd.DataFrame({'Date': dates, 'Adj Close': prices})
    df.set_index('Date', inplace=True)
    return df

# Generate data for demonstration
mock_data = generate_mock_data(initial_price=100, volatility=0.008)

print("Mock Data Head:")
display(mock_data.head())
print("\nMock Data Tail:")
display(mock_data.tail())

def run_ma_crossover_backtest(data, short_window, long_window):
    """
    Runs a Moving Average Crossover backtest and calculates cumulative returns.

    Args:
        data (pd.DataFrame): DataFrame with 'Adj Close' prices.
        short_window (int): Period for the short-term moving average.
        long_window (int): Period for the long-term moving average.

    Returns:
        float: Cumulative returns of the strategy.
    """
    if short_window >= long_window:
        return -np.inf # Invalid parameters, return a very low value

    df = data.copy()

    # Calculate moving averages
    df['short_ma'] = df['Adj Close'].rolling(window=short_window, min_periods=1).mean()
    df['long_ma'] = df['Adj Close'].rolling(window=long_window, min_periods=1).mean()

    # Generate signals
    df['signal'] = 0
    # When short MA crosses above long MA, generate a buy signal (1)
    df['signal'][short_window:] = np.where(df['short_ma'][short_window:] > df['long_ma'][short_window:], 1, 0)
    # When short MA crosses below long MA, generate a sell signal (-1)
    df['signal'][short_window:] = np.where(df['short_ma'][short_window:] < df['long_ma'][short_window:], -1, df['signal'][short_window:])

    # Calculate daily returns of the asset
    df['daily_returns'] = df['Adj Close'].pct_change()

    # Calculate strategy returns
    # We assume we are in the market (holding a position) when signal is 1
    # For simplicity, we are not handling short selling (-1 signal) in returns calculation here
    # A more robust backtest would use a 'position' column and multiply by daily returns
    df['strategy_returns'] = df['daily_returns'] * df['signal'].shift(1) # Shift signal to avoid look-ahead bias

    # Calculate cumulative returns
    # Add 1 to returns to allow product calculation, then subtract 1 at the end
    cumulative_returns = (1 + df['strategy_returns'].fillna(0)).cumprod() - 1

    # Return the last cumulative return value as the performance metric
    return cumulative_returns.iloc[-1] if not cumulative_returns.empty else -np.inf

# Example usage with arbitrary parameters
example_returns = run_ma_crossover_backtest(mock_data, short_window=20, long_window=50)
print(f"Example Cumulative Returns with short_window=20, long_window=50: {example_returns:.4f}")

def perform_grid_search(data, short_windows, long_windows):
    """
    Performs grid search optimization for the MA Crossover strategy.

    Args:
        data (pd.DataFrame): DataFrame with 'Adj Close' prices.
        short_windows (list): List of integers for short moving average windows.
        long_windows (list): List of integers for long moving average windows.

    Returns:
        tuple: A tuple containing:
            - dict: Best parameters found.
            - float: Best cumulative returns achieved.
            - pd.DataFrame: Results of all backtests.
    """
    best_returns = -np.inf
    best_params = {}
    results = []

    # Iterate through all combinations of short and long windows
    for sw, lw in itertools.product(short_windows, long_windows):
        if sw < lw: # Ensure short window is always less than long window
            cumulative_returns = run_ma_crossover_backtest(data, sw, lw)
            results.append({'short_window': sw, 'long_window': lw, 'cumulative_returns': cumulative_returns})

            if cumulative_returns > best_returns:
                best_returns = cumulative_returns
                best_params = {'short_window': sw, 'long_window': lw}

    results_df = pd.DataFrame(results)
    return best_params, best_returns, results_df

# Define parameter ranges for the grid search
short_ma_windows = range(10, 31, 5) # e.g., [10, 15, 20, 25, 30]
long_ma_windows = range(40, 81, 10)  # e.g., [40, 50, 60, 70, 80]

print(f"Short MA windows to test: {list(short_ma_windows)}")
print(f"Long MA windows to test: {list(long_ma_windows)}\n")

# Run the grid search
best_params, best_returns, grid_search_results = perform_grid_search(mock_data, list(short_ma_windows), list(long_ma_windows))

print("Grid Search Complete!")
print(f"Best Parameters: {best_params}")
print(f"Best Cumulative Returns: {best_returns:.4f}")

print("\nTop 5 Grid Search Results:")
display(grid_search_results.sort_values(by='cumulative_returns', ascending=False).head())

print("\nBottom 5 Grid Search Results:")
display(grid_search_results.sort_values(by='cumulative_returns', ascending=False).tail())

import matplotlib.pyplot as plt
import seaborn as sns

# Pivot the results to create a matrix suitable for a heatmap
heatmap_data = grid_search_results.pivot_table(index='short_window', columns='long_window', values='cumulative_returns')

plt.figure(figsize=(10, 8))
sns.heatmap(heatmap_data, annot=True, fmt=".2f", cmap="viridis", cbar_kws={'label': 'Cumulative Returns'})
plt.title('Strategy Cumulative Returns by MA Window Parameters (Grid Search)')
plt.xlabel('Long Moving Average Window')
plt.ylabel('Short Moving Average Window')
plt.show()

# Group by long_window and find the max returns (or mean, depending on analysis goal)
performance_by_long_window = grid_search_results.groupby('long_window')['cumulative_returns'].max().reset_index()

plt.figure(figsize=(12, 6))
sns.lineplot(x='long_window', y='cumulative_returns', data=performance_by_long_window, marker='o')
plt.title('Maximum Cumulative Returns vs. Long Moving Average Window (Across Short Windows)')
plt.xlabel('Long Moving Average Window')
plt.ylabel('Maximum Cumulative Returns')
plt.grid(True, linestyle='--', alpha=0.7)
plt.axhline(y=0, color='r', linestyle='--', label='Zero Returns Benchmark')
plt.legend()
plt.show()

def plot_strategy_cumulative_returns(data, short_window, long_window, title):
    """
    Calculates and plots the cumulative returns for a given strategy and parameters.

    Args:
        data (pd.DataFrame): DataFrame with 'Adj Close' prices.
        short_window (int): Period for the short-term moving average.
        long_window (int): Period for the long-term moving average.
        title (str): Title for the plot.
    """
    df = data.copy()

    # Calculate moving averages
    df['short_ma'] = df['Adj Close'].rolling(window=short_window, min_periods=1).mean()
    df['long_ma'] = df['Adj Close'].rolling(window=long_window, min_periods=1).mean()

    # Generate signals
    df['signal'] = 0
    df['signal'][short_window:] = np.where(df['short_ma'][short_window:] > df['long_ma'][short_window:], 1, 0)
    df['signal'][short_window:] = np.where(df['short_ma'][short_window:] < df['long_ma'][short_window:], -1, df['signal'][short_window:])

    # Calculate daily returns of the asset
    df['daily_returns'] = df['Adj Close'].pct_change()

    # Calculate strategy returns
    df['strategy_returns'] = df['daily_returns'] * df['signal'].shift(1).fillna(0)

    # Calculate cumulative returns for the strategy and the benchmark (buy & hold)
    df['cumulative_strategy_returns'] = (1 + df['strategy_returns']).cumprod() - 1
    df['cumulative_benchmark_returns'] = (1 + df['daily_returns']).cumprod() - 1

    plt.figure(figsize=(14, 7))
    plt.plot(df.index, df['cumulative_strategy_returns'], label=f'Optimized Strategy (MA {short_window}/{long_window})')
    plt.plot(df.index, df['cumulative_benchmark_returns'], label='Buy & Hold Benchmark', linestyle='--')
    plt.title(title)
    plt.xlabel('Date')
    plt.ylabel('Cumulative Returns')
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend()
    plt.show()

# Plot the performance of the best strategy
plot_strategy_cumulative_returns(
    mock_data,
    best_params['short_window'],
    best_params['long_window'],
    'Cumulative Returns of Optimized MA Crossover Strategy vs. Buy & Hold'
)