# Import necessary libraries
import pandas as pd
import yfinance as yf
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr

def sma_crossover_strategy(initial_budget, start_date, end_date, ticker):
    # Fetch the selected stock data using yfinance
    df = yf.download(ticker, start=start_date, end=end_date, progress=False)
    df = df[['Close']]  # Use only the 'Close' price for SMA calculations

    # Calculate SMAs
    df['SMA_50'] = df['Close'].rolling(window=50).mean()
    df['SMA_150'] = df['Close'].rolling(window=150).mean()

    # Determine buy and sell signals
    df['Signal'] = 0  # Initialize a column for signals
    df['Signal'][df['SMA_50'] > df['SMA_150']] = 1  # Buy signal
    df['Signal'][df['SMA_50'] < df['SMA_150']] = -1  # Sell signal
    df['Position'] = df['Signal'].diff()  # Capture points where the signal changes

    # Initialize investment simulation
    cash = initial_budget  # Start with user-specified cash
    shares = 0  # No shares initially
    portfolio_values = []  # Store portfolio values

    # Iterate over the dataframe to simulate trading
    for index, row in df.iterrows():
        # Buy signal
        if row['Position'] == 1 and cash > 0:
            shares = cash / row['Close']
            cash = 0  # All money is invested in shares

        # Sell signal
        elif row['Position'] == -1 and shares > 0:
            cash = shares * row['Close']
            shares = 0  # All shares are sold

        # Calculate current portfolio value
        portfolio_value = cash + (shares * row['Close'])
        portfolio_values.append(portfolio_value)

    # Add portfolio values to the dataframe for visualization
    df = df.iloc[149:]  # Ignore rows with NaN values due to SMA calculations
    df['Portfolio Value'] = portfolio_values[149:]  # Align portfolio values

    # Plot Portfolio Value over time
    plt.figure(figsize=(14, 8))
    plt.plot(df['Portfolio Value'], label='Portfolio Value', color='purple')
    plt.xlabel('Date')
    plt.ylabel('Portfolio Value ($)')
    plt.title(f'Portfolio Value Over Time with 50/150 SMA Crossover Strategy ({ticker})')
    plt.legend()
    plt.grid()
    plt.tight_layout()

    # Save plot to a file for display
    plot_file = "portfolio_value_plot.png"
    plt.savefig(plot_file)
    plt.close()

    # Final results
    final_value = portfolio_values[-1]
    profit_loss = final_value - initial_budget
    percentage_return = (profit_loss / initial_budget) * 100

    # Create summary text
    results = f"""
    Ticker: {ticker}
    Trading Period: {start_date} to {end_date}
    Initial Investment: ${initial_budget}
    Final Portfolio Value: ${final_value:.2f}
    Total Profit/Loss: ${profit_loss:.2f}
    Percentage Return: {percentage_return:.2f}%
    """

    # Save results to a text file
    results_file = "simulation_results.txt"
    with open(results_file, "w") as f:
        f.write(results)

    return plot_file, results, results_file

# Define Gradio interface components
with gr.Blocks() as app:
    gr.Markdown("# SMA Crossover Trading Strategy Simulator")

    with gr.Tabs():
        # Tab for SMA Strategy Simulation
        with gr.Tab("SMA Strategy Simulator"):
            with gr.Row():
                initial_budget = gr.Number(label="Initial Investment ($)", value=100, interactive=True)
                start_date = gr.Text(label="Start Date (YYYY-MM-DD)", value="1993-01-01", interactive=True)
                end_date = gr.Text(label="End Date (YYYY-MM-DD)", value="2023-12-31", interactive=True)
                ticker = gr.Dropdown(
                    label="Stock Ticker Symbol",
                    choices=["SPY", "TSLA", "GOOGL", "AAPL", "MSFT"],
                    value="SPY",
                )
            run_button = gr.Button("Run Simulation")
            with gr.Row():
                portfolio_graph = gr.Image(label="Portfolio Value Over Time")
                summary_text = gr.Textbox(label="Simulation Summary", lines=8)
                download_button = gr.File(label="Download Results (.txt)")

        # Tab for Instructions
        with gr.Tab("Instructions"):
            gr.Markdown("""
            ## How to Use:
            1. Enter your initial investment amount.
            2. Specify the trading period (start and end dates).
            3. Select a stock ticker symbol (e.g., SPY, TSLA, GOOGL).
            4. Click "Run Simulation" to visualize the portfolio value over time and view a summary of results.
            5. Download the results as a `.txt` file using the download button.
            
            ### Notes:
            - The 50-day and 150-day SMAs are used for buy and sell signals.
            - Ensure the trading period is valid for the selected ticker symbol.
            """)

    # Link simulation function to UI
    run_button.click(
        sma_crossover_strategy,
        inputs=[initial_budget, start_date, end_date, ticker],
        outputs=[portfolio_graph, summary_text, download_button],
    )

# Launch the app
app.launch()