import yfinance as yf
import tensorflow as tf
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import gradio as gr
from datetime import datetime

# Stock tickers for dropdown
stock_tickers = ['AAPL', 'MSFT', 'GOOGL', 'AMZN', 'TSLA', 'NFLX', 'FB', 'NVDA', 'JPM', 'BAC']

# Function to download stock data from Yahoo Finance
def fetch_stock_data(ticker, start_date, end_date):
    data = yf.download(ticker, start=start_date, end=end_date)
    return data

# Function to train LSTM model and predict stock prices
def train_lstm(data):
    # Preprocessing and creating training dataset
    data['Close'] = data['Close'].fillna(method='ffill')  # Forward-fill to handle NaNs
    scaled_data = data['Close'].values.reshape(-1, 1)
    training_data = scaled_data[:int(len(scaled_data) * 0.8)]  # 80% data for training
    test_data = scaled_data[int(len(scaled_data) * 0.8):]

    # Creating sequences for LSTM input
    def create_sequences(data, seq_length):
        X, y = [], []
        for i in range(len(data) - seq_length):
            X.append(data[i:i + seq_length])
            y.append(data[i + seq_length])
        return np.array(X), np.array(y)

    seq_length = 60  # Using 60 days for prediction
    X_train, y_train = create_sequences(training_data, seq_length)

    # Define the LSTM model
    model = tf.keras.Sequential([
        tf.keras.layers.LSTM(50, return_sequences=True, input_shape=(X_train.shape[1], 1)),
        tf.keras.layers.LSTM(50, return_sequences=False),
        tf.keras.layers.Dense(25),
        tf.keras.layers.Dense(1)
    ])

    # Compile and train the model
    model.compile(optimizer='adam', loss='mean_squared_error')
    model.fit(X_train, y_train, batch_size=1, epochs=1)  # Short training for demo, increase epochs for better results

    # Predict future prices
    X_test, _ = create_sequences(test_data, seq_length)
    predictions = model.predict(X_test)

    return predictions, data.index[-len(predictions):]

# Function to plot historical and predicted prices
def plot_stock_data(data, predictions, pred_dates):
    plt.figure(figsize=(10,6))
    plt.plot(data['Close'], label='Historical Data')
    plt.plot(pred_dates, predictions, label='Predicted Data')
    plt.xlabel('Date')
    plt.ylabel('Close Price (USD)')
    plt.legend()
    plt.grid()
    plt.show()

# Main function to wrap everything into Gradio interface
def stock_prediction(ticker, start_date, end_date):
    data = fetch_stock_data(ticker, start_date, end_date)
    predictions, pred_dates = train_lstm(data)
    plot_stock_data(data, predictions, pred_dates)

    # Calculate percentage change, highest and lowest values
    percentage_change = ((data['Close'][-1] - data['Close'][0]) / data['Close'][0]) * 100
    highest_value = data['Close'].max()
    lowest_value = data['Close'].min()
    
    # Buy/Sell decision based on prediction (buy if next value is higher)
    decision = 'Buy' if predictions[-1] > data['Close'][-1] else 'Sell'
    
    return f"Percentage Change: {percentage_change:.2f}%", f"Highest Value: {highest_value:.2f}", f"Lowest Value: {lowest_value:.2f}", f"Decision: {decision}"

# Gradio UI
interface = gr.Interface(
    fn=stock_prediction,
    inputs=[
        gr.inputs.Dropdown(choices=stock_tickers, label="Select Stock Ticker"),
        gr.inputs.Date(label="Start Date"),
        gr.inputs.Date(label="End Date")
    ],
    outputs=[
        gr.outputs.Textbox(label="Percentage Change"),
        gr.outputs.Textbox(label="Highest Value"),
        gr.outputs.Textbox(label="Lowest Value"),
        gr.outputs.Textbox(label="Prediction: Buy/Sell")
    ],
    title="Stock Prediction App",
    description="Predict if you should buy or sell a stock based on historical performance."
)

interface.launch()