# Import required libraries
import yfinance as yf
import pandas as pd
import numpy as np
import tensorflow as tf
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
import gradio as gr
from datetime import datetime

# Step 1: Fetch stock data from yfinance
def fetch_stock_data(ticker, start_date, end_date):
    stock_data = yf.download(ticker, start=start_date, end=end_date)
    stock_data.reset_index(inplace=True)
    return stock_data

# Step 2: Prepare data for the LSTM model
def prepare_data(df):
    scaler = MinMaxScaler(feature_range=(0, 1))
    close_prices = df['Close'].values.reshape(-1, 1)
    scaled_data = scaler.fit_transform(close_prices)
    return scaled_data, scaler

# Step 3: Build the LSTM model
def build_model(input_shape):
    model = tf.keras.Sequential([
        tf.keras.layers.LSTM(50, return_sequences=True, input_shape=input_shape),
        tf.keras.layers.LSTM(50, return_sequences=False),
        tf.keras.layers.Dense(25),
        tf.keras.layers.Dense(1)
    ])
    model.compile(optimizer='adam', loss='mean_squared_error')
    return model

# Step 4: Train the LSTM model
def train_model(model, train_data, epochs=5):
    X_train, y_train = [], []
    for i in range(60, len(train_data)):
        X_train.append(train_data[i-60:i, 0])
        y_train.append(train_data[i, 0])
    
    X_train, y_train = np.array(X_train), np.array(y_train)
    X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], 1))
    
    model.fit(X_train, y_train, epochs=epochs, batch_size=32, verbose=0)
    return model

# Step 5: Predict stock prices
def predict_future(model, last_data, scaler, steps=1):
    predictions = []
    input_data = last_data[-60:].reshape(1, -1)

    # Generate predictions for the specified number of future days
    for _ in range(steps):
        input_reshaped = input_data.reshape(1, 60, 1)
        predicted_price = model.predict(input_reshaped, verbose=0)
        predictions.append(predicted_price[0][0])
        input_data = np.append(input_data[0][1:], predicted_price)
    
    predicted_prices = np.array(predictions).reshape(-1, 1)
    predicted_prices = scaler.inverse_transform(predicted_prices)  # Reverse scaling
    return predicted_prices

# Step 6: Plot historical and predicted stock prices
def plot_predictions(data, predicted_prices):
    # Create a figure for the plot
    plt.figure(figsize=(14, 6))
    
    # Plot historical stock prices
    plt.plot(data['Date'], data['Close'], label="Historical Prices", color='blue')

    # Generate future date for predicted prices
    future_date = pd.to_datetime(data['Date'].iloc[-1]) + pd.Timedelta(days=1)  # Tomorrow's date
    plt.scatter(future_date, predicted_prices[0][0], color='red', label="Predicted Price for Tomorrow")  # Red dot for prediction

    # Adding labels and title to the graph
    plt.title("Stock Price Prediction for Tomorrow")
    plt.xlabel("Date")
    plt.ylabel("Stock Price (USD)")
    plt.axhline(y=predicted_prices[0][0], color='red', linestyle='--')  # Horizontal line for predicted price
    plt.legend()
    plt.grid(True)

    # Save the plot as an image file for Gradio to display
    plt.savefig("stock_prediction.png")
    plt.close()

    return "stock_prediction.png"  # Return the path to the saved image

# Step 7: Gradio interface function
def stock_prediction_app(ticker, start_date_str, end_date_str):
    # Convert input strings to datetime objects
    start_date = datetime.strptime(start_date_str, "%Y-%m-%d").date()
    end_date = datetime.strptime(end_date_str, "%Y-%m-%d").date()

    # Fetch stock data from yfinance
    data = fetch_stock_data(ticker, start_date, end_date)
    
    # Prepare data for LSTM model
    scaled_data, scaler = prepare_data(data)
    
    # Build and train the LSTM model
    model = build_model((60, 1))
    model = train_model(model, scaled_data)
    
    # Predict stock price for tomorrow (1 day)
    predicted_price = predict_future(model, scaled_data, scaler, steps=1)

    # Current, highest, lowest prices, and percentage change
    current_price = data['Close'].iloc[-1]
    highest_price = data['Close'].max()
    lowest_price = data['Close'].min()
    
    # Calculate percentage change
    percentage_change = ((predicted_price[0][0] - current_price) / current_price) * 100

    # Buy/Sell recommendation
    recommendation = "Buy" if predicted_price[0][0] > current_price else "Sell"

    # Generate and return the plot with historical and predicted prices
    plot_path = plot_predictions(data, predicted_price)
    
    return (
        plot_path, 
        predicted_price[0][0],  # Tomorrow's predicted price
        current_price,           # Current price
        highest_price,           # Highest price
        lowest_price,            # Lowest price
        percentage_change,       # Percentage change
        recommendation            # Buy/Sell recommendation
    )

# Step 8: Gradio UI setup
tickers = ["AAPL", "GOOGL", "MSFT", "AMZN", "TSLA", "META", "NFLX", "NVDA", "BABA", "BA"]

# Define the Gradio interface
ui = gr.Interface(
    fn=stock_prediction_app,
    inputs=[
        gr.Dropdown(tickers, label="Select Stock Ticker"),
        gr.Textbox(label="Start Date (YYYY-MM-DD)"),
        gr.Textbox(label="End Date (YYYY-MM-DD)")
    ],
    outputs=[
        gr.Image(type="filepath"),  # Return the file path for the generated graph
        gr.Number(label="Predicted Price for Tomorrow (USD)"),
        gr.Number(label="Current Price (USD)"),
        gr.Number(label="Highest Price (USD)"),
        gr.Number(label="Lowest Price (USD)"),
        gr.Number(label="Percentage Change (%)"),
        gr.Textbox(label="Recommendation")
    ],
    title="Stock Price Prediction App",
    description="Predict future stock price for tomorrow based on historical data."
)

# Launch the Gradio app
ui.launch()