import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# Set random seed for reproducibility
np.random.seed(42)

def generate_price_data(num_points=1000, initial_price=100, trend_strength=0.1, noise_level=0.5):
    """
    Generates synthetic price data.

    Inputs:
    - num_points (int): The number of data points to generate.
    - initial_price (float): The starting price of the series.
    - trend_strength (float): The strength of the underlying trend.
    - noise_level (float): The magnitude of random fluctuations.

    Outputs:
    - pd.Series: A pandas Series containing the simulated price data.
    """
    time = np.arange(num_points)
    # Introduce a changing trend over time
    trend = initial_price + trend_strength * time + 0.0005 * (time ** 2)
    noise = np.random.normal(0, noise_level, num_points).cumsum()
    prices = trend + noise
    return pd.Series(prices, name='Price')

# Generate simulated data
data = generate_price_data(num_points=1000, initial_price=100, trend_strength=0.05, noise_level=1)
dates = pd.date_range(start='2020-01-01', periods=1000, freq='D')
df = pd.DataFrame({'Date': dates, 'Price': data})
df = df.set_index('Date')

print("Simulated Price Data Head:")
display(df.head())

print("\nSimulated Price Data Tail:")
display(df.tail())

# Visualize the simulated data
plt.figure(figsize=(12, 6))
plt.plot(df.index, df['Price'])
plt.title('Simulated Price Data')
plt.xlabel('Date')
plt.ylabel('Price')
plt.grid(True)
plt.show()

def simple_moving_average_strategy(data, fast_ma_len, slow_ma_len):
    """
    Applies a simple moving average crossover strategy to price data.

    Inputs:
    - data (pd.Series): The price data (e.g., 'Price' column from DataFrame).
    - fast_ma_len (int): The length of the fast moving average.
    - slow_ma_len (int): The length of the slow moving average.

    Outputs:
    - pd.Series: A series of daily returns from the strategy.
    """
    if fast_ma_len >= slow_ma_len:
        return pd.Series(0, index=data.index) # Invalid lengths, return no profit

    fast_ma = data.rolling(window=fast_ma_len).mean()
    slow_ma = data.rolling(window=slow_ma_len).mean()

    # Generate signals: 1 for buy, -1 for sell, 0 for hold
    signals = pd.Series(0, index=data.index)
    signals[fast_ma > slow_ma] = 1
    signals[fast_ma < slow_ma] = -1

    # Shift signals to ensure we trade on the next day's open after a signal
    positions = signals.shift(1).fillna(0) # 0 means no position

    # Calculate daily returns: (change in price / previous price) * position
    daily_returns = data.pct_change() * positions
    return daily_returns

def walk_forward_backtest(
    df,
    price_col='Price',
    in_sample_len=200,
    out_of_sample_len=50,
    fast_ma_range=(10, 30),
    slow_ma_range=(40, 80)
):
    """
    Performs a walk-forward backtest on a given DataFrame using a simple MA crossover strategy.

    Inputs:
    - df (pd.DataFrame): DataFrame containing price data with a datetime index.
    - price_col (str): The name of the column containing the price data.
    - in_sample_len (int): Length of the in-sample (optimization) period.
    - out_of_sample_len (int): Length of the out-of-sample (testing) period.
    - fast_ma_range (tuple): (min_len, max_len) for the fast moving average.
    - slow_ma_range (tuple): (min_len, max_len) for the slow moving average.

    Outputs:
    - pd.DataFrame: A DataFrame containing the daily returns for each out-of-sample period.
    - list: A list of dictionaries, each containing optimized parameters for a window.
    """
    all_oos_returns = pd.Series(dtype=float)
    optimized_params_history = []

    total_len = len(df)
    current_start_idx = 0

    while current_start_idx + in_sample_len + out_of_sample_len <= total_len:
        in_sample_end_idx = current_start_idx + in_sample_len
        oos_end_idx = in_sample_end_idx + out_of_sample_len

        in_sample_data = df.iloc[current_start_idx:in_sample_end_idx][price_col]
        oos_data = df.iloc[in_sample_end_idx:oos_end_idx][price_col]

        best_fast_ma = -1
        best_slow_ma = -1
        max_in_sample_return = -np.inf

        # --- Optimization (In-Sample) ---
        for fast_len in range(fast_ma_range[0], fast_ma_range[1] + 1):
            for slow_len in range(slow_ma_range[0], slow_ma_range[1] + 1):
                if fast_len < slow_len:
                    strategy_returns = simple_moving_average_strategy(in_sample_data, fast_len, slow_len)
                    total_return = strategy_returns.sum()
                    if total_return > max_in_sample_return:
                        max_in_sample_return = total_return
                        best_fast_ma = fast_len
                        best_slow_ma = slow_len

        # Store optimized parameters
        optimized_params_history.append({
            'in_sample_start': in_sample_data.index.min(),
            'in_sample_end': in_sample_data.index.max(),
            'oos_start': oos_data.index.min(),
            'oos_end': oos_data.index.max(),
            'best_fast_ma': best_fast_ma,
            'best_slow_ma': best_slow_ma,
            'max_in_sample_return': max_in_sample_return
        })

        # --- Out-of-Sample Testing ---
        if best_fast_ma != -1 and best_slow_ma != -1:
            oos_strategy_returns = simple_moving_average_strategy(oos_data, best_fast_ma, best_slow_ma)
            all_oos_returns = pd.concat([all_oos_returns, oos_strategy_returns])
        else:
            # If no valid parameters found, append zeros for the OOS period
            oos_zero_returns = pd.Series(0, index=oos_data.index)
            all_oos_returns = pd.concat([all_oos_returns, oos_zero_returns])

        # Move the window forward
        current_start_idx += out_of_sample_len # Walk forward by OOS length

    return all_oos_returns, optimized_params_history


# --- Demonstrate the walk-forward backtest ---
print("Performing walk-forward backtest...")
# For demonstration, using shorter ranges and lengths
oos_returns, params_history = walk_forward_backtest(
    df,
    in_sample_len=150,
    out_of_sample_len=30,
    fast_ma_range=(5, 15),
    slow_ma_range=(20, 40)
)

print("\nWalk-Forward Out-of-Sample Returns Head:")
display(oos_returns.head())

print("\nWalk-Forward Out-of-Sample Returns Tail:")
display(oos_returns.tail())

print("\nOptimized Parameters History (first 5 windows):")
display(pd.DataFrame(params_history).head())

plt.figure(figsize=(14, 7))
(1 + oos_returns).cumprod().plot()
plt.title('Walk-Forward Backtest: Cumulative Out-of-Sample Returns')
plt.xlabel('Date')
plt.ylabel('Cumulative Return')
plt.grid(True)
plt.axhline(1.0, color='red', linestyle='--', linewidth=0.8, label='Break-even (1.0)')
plt.legend()
plt.show()

params_df = pd.DataFrame(params_history)
params_df['oos_start'] = pd.to_datetime(params_df['oos_start'])
params_df = params_df.set_index('oos_start')

plt.figure(figsize=(14, 7))
plt.plot(params_df.index, params_df['best_fast_ma'], label='Optimized Fast MA Length', marker='o', linestyle='--')
plt.plot(params_df.index, params_df['best_slow_ma'], label='Optimized Slow MA Length', marker='x', linestyle='-')
plt.title('Optimized Moving Average Lengths Across Walk-Forward Windows')
plt.xlabel('Out-of-Sample Start Date')
plt.ylabel('MA Length')
plt.grid(True)
plt.legend()
plt.show()

def traditional_backtest(df, price_col='Price', fast_ma_range=(10, 30), slow_ma_range=(40, 80)):
    """
    Performs a traditional backtest by optimizing parameters once over the entire dataset.

    Inputs:
    - df (pd.DataFrame): DataFrame containing price data with a datetime index.
    - price_col (str): The name of the column containing the price data.
    - fast_ma_range (tuple): (min_len, max_len) for the fast moving average.
    - slow_ma_range (tuple): (min_len, max_len) for the slow moving average.

    Outputs:
    - pd.Series: Daily returns from the strategy applied with the best parameters.
    - dict: Best fast MA, best slow MA, and max total return.
    """
    best_fast_ma = -1
    best_slow_ma = -1
    max_total_return = -np.inf

    data_to_optimize = df[price_col]

    for fast_len in range(fast_ma_range[0], fast_ma_range[1] + 1):
        for slow_len in range(slow_ma_range[0], slow_ma_range[1] + 1):
            if fast_len < slow_len:
                strategy_returns = simple_moving_average_strategy(data_to_optimize, fast_len, slow_len)
                total_return = strategy_returns.sum()
                if total_return > max_total_return:
                    max_total_return = total_return
                    best_fast_ma = fast_len
                    best_slow_ma = slow_len

    if best_fast_ma != -1 and best_slow_ma != -1:
        final_returns = simple_moving_average_strategy(data_to_optimize, best_fast_ma, best_slow_ma)
    else:
        final_returns = pd.Series(0, index=df.index)

    return final_returns, {'best_fast_ma': best_fast_ma, 'best_slow_ma': best_slow_ma, 'max_total_return': max_total_return}

# Perform traditional backtest
trad_returns, trad_params = traditional_backtest(
    df,
    fast_ma_range=(5, 15),
    slow_ma_range=(20, 40)
)

print("Traditional Backtest Optimized Parameters:")
display(trad_params)

plt.figure(figsize=(14, 7))
(1 + trad_returns).cumprod().plot(label='Traditional Backtest Cumulative Returns')
(1 + oos_returns).cumprod().plot(label='Walk-Forward Backtest Cumulative OOS Returns', linestyle='--')
plt.title('Walk-Forward vs. Traditional Backtest Cumulative Returns')
plt.xlabel('Date')
plt.ylabel('Cumulative Return')
plt.grid(True)
plt.legend()
plt.show()