# Import necessary libraries
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

# Configure plot style
plt.style.use('seaborn-v0_8-darkgrid')
plt.rcParams['figure.figsize'] = (12, 7)
plt.rcParams['lines.linewidth'] = 2
plt.rcParams['axes.labelsize'] = 14
plt.rcParams['xtick.labelsize'] = 12
plt.rcParams['ytick.labelsize'] = 12
plt.rcParams['legend.fontsize'] = 12

def generate_mock_data(start_date, end_date, initial_price=100, volatility=0.01, drift=0.0005):
    """
    Generates mock daily stock price data.

    Inputs:
    - start_date (str): Start date in 'YYYY-MM-DD' format.
    - end_date (str): End date in 'YYYY-MM-DD' format.
    - initial_price (float): Starting price of the asset.
    - volatility (float): Daily volatility of the asset price.
    - drift (float): Daily drift (average daily return) of the asset price.

    Outputs:
    - pd.DataFrame: DataFrame with 'Date' and 'Close' columns.
    """
    dates = pd.date_range(start=start_date, end=end_date, freq='B') # Business days
    n_days = len(dates)

    prices = [initial_price]
    for _ in range(1, n_days):
        # Simulate daily price movement using a geometric Brownian motion-like process
        daily_return = np.random.normal(drift, volatility)
        new_price = prices[-1] * (1 + daily_return)
        prices.append(new_price)

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

# Generate 2 years of mock data
data = generate_mock_data('2022-01-01', '2023-12-31')
print(f"Generated {len(data)} daily price points.")

# Display the first few rows of the generated data
display(data.head())

def moving_average_crossover_strategy(data, short_window=40, long_window=100):
    """
    Generates trading signals based on a Moving Average Crossover strategy.

    Inputs:
    - data (pd.DataFrame): DataFrame with 'Close' prices.
    - short_window (int): Lookback period for the short-term moving average.
    - long_window (int): Lookback period for the long-term moving average.

    Outputs:
    - pd.DataFrame: Contains 'Short_MA', 'Long_MA', and 'Signal' columns.
    """
    signals = pd.DataFrame(index=data.index)
    signals['Close'] = data['Close']
    signals['Short_MA'] = data['Close'].rolling(window=short_window, min_periods=1).mean()
    signals['Long_MA'] = data['Close'].rolling(window=long_window, min_periods=1).mean()

    # Generate signals
    signals['Signal'] = 0.0
    # Use .loc for assignment to avoid FutureWarning
    signals.loc[signals.index[short_window:], 'Signal'] = np.where(
        signals['Short_MA'][short_window:] > signals['Long_MA'][short_window:], 1.0, 0.0
    )

    # Differentiate to get actual trading orders (-1 for sell, 1 for buy)
    signals['Positions'] = signals['Signal'].diff()

    return signals.dropna()

# Generate signals using the strategy
signals = moving_average_crossover_strategy(data)
print("Value counts for signals['Positions']:")
print(signals['Positions'].value_counts())
display(signals.head())

def run_backtest(signals, initial_capital=100000.0, commission=0.001):
    """
    Simulates the execution of a trading strategy.
    """
    portfolio = pd.DataFrame(index=signals.index)
    portfolio['Holdings'] = 0.0
    portfolio['Cash'] = initial_capital
    portfolio['Total Assets'] = initial_capital

    # Keep track of current position (1 = long, 0 = out)
    current_position = 0  # Start with no position
    shares = 0

    print(f"Initial Capital: {initial_capital}")

    for i in range(len(signals)):
        current_date = signals.index[i]
        current_close = signals['Close'].iloc[i]
        signal = signals['Positions'].iloc[i]  # This is the diff

        # Execute trades based on signal
        if signal == 1.0 and current_position == 0:  # Buy signal and not in position
            # Buy with all available cash
            shares = int(portfolio.loc[current_date, 'Cash'] / current_close)
            cost = shares * current_close
            commission_cost = cost * commission

            if portfolio.loc[current_date, 'Cash'] >= cost + commission_cost:
                portfolio.loc[current_date, 'Cash'] -= (cost + commission_cost)
                current_position = 1
                print(f"BUY on {current_date} at ${current_close:.2f}, Shares: {shares}")

        elif signal == -1.0 and current_position == 1:  # Sell signal and in position
            # Sell all shares
            sale_revenue = shares * current_close
            commission_cost = sale_revenue * commission

            portfolio.loc[current_date, 'Cash'] += (sale_revenue - commission_cost)
            shares = 0
            current_position = 0
            print(f"SELL on {current_date} at ${current_close:.2f}, Revenue: ${sale_revenue:.2f}")

        # Update holdings value for current day
        portfolio.loc[current_date, 'Holdings'] = shares * current_close
        portfolio.loc[current_date, 'Total Assets'] = portfolio.loc[current_date, 'Cash'] + portfolio.loc[current_date, 'Holdings']

        # Carry forward cash if no trade
        if i + 1 < len(signals):
            portfolio.loc[signals.index[i+1], 'Cash'] = portfolio.loc[current_date, 'Cash']

    # Calculate returns
    portfolio['Daily Return'] = portfolio['Total Assets'].pct_change()
    portfolio['Cumulative Return'] = (1 + portfolio['Daily Return']).cumprod() - 1

    print(f"Final Capital: ${portfolio['Total Assets'].iloc[-1]:.2f}")
    return portfolio.fillna(0)

def calculate_performance_metrics(portfolio, risk_free_rate=0.02):
    """
    Calculates key performance metrics for a backtest.

    Inputs:
    - portfolio (pd.DataFrame): DataFrame from run_backtest.
    - risk_free_rate (float): Annual risk-free rate (e.g., treasury bond yield).

    Outputs:
    - dict: A dictionary of performance metrics.
    """
    metrics = {}

    # Total Return
    total_return = (portfolio['Total Assets'].iloc[-1] / portfolio['Total Assets'].iloc[0]) - 1
    metrics['Total Return'] = total_return

    # Annualized Return
    num_years = (portfolio.index[-1] - portfolio.index[0]).days / 365.25
    metrics['Annualized Return'] = (1 + total_return)**(1 / num_years) - 1 if num_years > 0 else 0

    # Daily Returns (excluding the first NaN)
    daily_returns = portfolio['Daily Return'].dropna()

    # Annualized Volatility
    annualized_volatility = daily_returns.std() * np.sqrt(252) # 252 trading days in a year
    metrics['Annualized Volatility'] = annualized_volatility

    # Sharpe Ratio
    if annualized_volatility != 0:
        metrics['Sharpe Ratio'] = (metrics['Annualized Return'] - risk_free_rate) / annualized_volatility
    else:
        metrics['Sharpe Ratio'] = np.nan

    # Maximum Drawdown
    # Calculate the running maximum (peak)
    peak = portfolio['Total Assets'].cummax()
    # Calculate the daily drawdown from the peak
    drawdown = (portfolio['Total Assets'] - peak) / peak
    metrics['Maximum Drawdown'] = drawdown.min()

    return metrics

# Calculate and display performance metrics
performance_metrics = calculate_performance_metrics(portfolio)

print("--- Backtest Performance Metrics ---")
for metric, value in performance_metrics.items():
    print(f"{metric:<20}: {value:,.4f}")

fig1, ax1 = plt.subplots(figsize=(14, 7))
ax1.plot(portfolio.index, portfolio['Total Assets'], label='Portfolio Value', color='blue')
ax1.set_title('Equity Curve (Portfolio Value Over Time)', fontsize=16)
ax1.set_xlabel('Date', fontsize=14)
ax1.set_ylabel('Portfolio Value ($)', fontsize=14)
ax1.legend()
ax1.grid(True)
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

print("Interpretation: This chart shows how the portfolio's total value changed over the backtesting period. An upward sloping curve indicates profitability, while steep drops indicate drawdowns.")

# @title step_artifacts
num_fig = "1" # @param {type:"string"}
step = 'DataVisualization'  # @param ["DataLoading", "DataExploration", "DataCleaning", "DataWrangling", "DataVisualization", "DataSplitting", "ModelDevelopment", "ModelOptimization", "ModelEvaluation", "Summary"] {type:"string"}
# Assume `upload_plt_to_gcs` is defined elsewhere or not strictly required for this notebook.

fig2, ax2 = plt.subplots(figsize=(14, 8))

# Plot the closing price
ax2.plot(data.index, data['Close'], label='Close Price', alpha=0.7)

# Plot the moving averages
ax2.plot(signals.index, signals['Short_MA'], label='Short MA', color='orange', linestyle='--')
ax2.plot(signals.index, signals['Long_MA'], label='Long MA', color='red', linestyle='--')

# Plot buy signals
ax2.plot(
    signals.loc[signals['Positions'] == 1.0].index,
    signals['Short_MA'][signals['Positions'] == 1.0],
    '^', markersize=10, color='green', lw=0, label='Buy Signal'
)

# Plot sell signals
ax2.plot(
    signals.loc[signals['Positions'] == -1.0].index,
    signals['Short_MA'][signals['Positions'] == -1.0],
    'v', markersize=10, color='red', lw=0, label='Sell Signal'
)

ax2.set_title('Price Chart with Trading Signals and Moving Averages', fontsize=16)
ax2.set_xlabel('Date', fontsize=14)
ax2.set_ylabel('Price ($)', fontsize=14)
ax2.legend()
ax2.grid(True)
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

print("Interpretation: This chart visually confirms where the strategy generated buy (green up arrows) and sell (red down arrows) signals relative to the price action and moving averages. It helps in understanding the mechanics of the strategy's entry and exit points.")

# @title step_artifacts
num_fig = "2" # @param {type:"string"}
step = 'DataVisualization'  # @param ["DataLoading", "DataExploration", "DataCleaning", "DataWrangling", "DataVisualization", "DataSplitting", "ModelDevelopment", "ModelOptimization", "ModelEvaluation", "Summary"] {type:"string"}
# Assume `upload_plt_to_gcs` is defined elsewhere or not strictly required for this notebook.