import numpy as np
import pandas as pd
import yfinance as yf
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
import gradio as gr
import matplotlib.pyplot as plt
from datetime import datetime, timedelta

# Fetch historical stock data
def fetch_data(ticker, start_date, end_date):
    data = yf.download(ticker, start=start_date, end=end_date)
    return data

# Preprocess data
def preprocess_data(data):
    # Use only the 'Close' prices for prediction
    data = data[['Close']]
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data)

    # Prepare training data
    x_train, y_train = [], []
    for i in range(60, len(scaled_data)):
        x_train.append(scaled_data[i-60:i, 0])
        y_train.append(scaled_data[i, 0])
    x_train, y_train = np.array(x_train), np.array(y_train)

    # Reshape for LSTM [samples, time steps, features]
    x_train = np.reshape(x_train, (x_train.shape[0], x_train.shape[1], 1))
    
    return x_train, y_train, scaler

# Build the LSTM model
def build_model():
    model = Sequential()
    model.add(LSTM(units=50, return_sequences=True, input_shape=(60, 1)))
    model.add(Dropout(0.2))
    model.add(LSTM(units=50, return_sequences=False))
    model.add(Dropout(0.2))
    model.add(Dense(units=25))
    model.add(Dense(units=1))  # Output layer for price prediction
    model.compile(optimizer='adam', loss='mean_squared_error')
    return model

# Train the model
def train_model(ticker, start_date, end_date):
    data = fetch_data(ticker, start_date, end_date)
    x_train, y_train, scaler = preprocess_data(data)
    
    model = build_model()
    model.fit(x_train, y_train, batch_size=1, epochs=1)
    
    return model, scaler, data

# Predict the stock price for tomorrow
def predict_next_day(model, scaler, data):
    last_60_days = data['Close'][-60:].values
    last_60_days_scaled = scaler.transform(last_60_days.reshape(-1, 1))
    
    x_test = []
    x_test.append(last_60_days_scaled)
    x_test = np.array(x_test)
    x_test = np.reshape(x_test, (x_test.shape[0], x_test.shape[1], 1))

    predicted_price = model.predict(x_test)
    predicted_price = scaler.inverse_transform(predicted_price)  # Convert back to original scale
    return predicted_price[0][0]

# Create a graph of historical and predicted prices
def create_graph(data, predicted_price):
    plt.figure(figsize=(14, 5))
    plt.plot(data.index, data['Close'], label='Historical Prices', color='blue')
    tomorrow = datetime.now() + timedelta(days=1)
    plt.scatter(tomorrow, predicted_price, label='Predicted Price for Tomorrow', color='red')
    plt.title('Stock Price Prediction')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.grid()
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig('/mnt/data/stock_prediction_graph.png')  # Save the graph
    plt.close()

# Gradio Interface
def stock_prediction_app(ticker, start_date, end_date):
    model, scaler, data = train_model(ticker, start_date, end_date)
    predicted_price = predict_next_day(model, scaler, data)
    create_graph(data, predicted_price)
    
    return f'The predicted stock price for tomorrow is ${predicted_price:.2f}', '/mnt/data/stock_prediction_graph.png'

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

# Create Gradio Interface
ticker_input = gr.Dropdown(choices=stock_tickers, label="Select Stock Ticker")
start_date_input = gr.Date(label="Start Date")  # Corrected date input
end_date_input = gr.Date(label="End Date")  # Corrected date input

iface = gr.Interface(
    fn=stock_prediction_app,
    inputs=[ticker_input, start_date_input, end_date_input],
    outputs=["text", "image"],
    title="Stock Price Prediction App",
    description="Predict tomorrow's stock price based on historical data.",
)

iface.launch()