import gradio as gr
import yfinance as yf
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from datetime import datetime, timedelta
import matplotlib.pyplot as plt

def get_stock_data(ticker):
    today = datetime.today().strftime('%Y-%m-%d')
    year_ago = (datetime.today() - timedelta(days=365)).strftime('%Y-%m-%d')
    stock_data = yf.download(ticker, start=year_ago, end=today)
    return stock_data

def preprocess_data(data):
    data['Date'] = pd.to_datetime(data.index)
    data['Date_ordinal'] = data['Date'].map(datetime.toordinal)
    return data[['Date_ordinal', 'Close']]

def train_model(data):
    X = data[['Date_ordinal']]
    y = data['Close']
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    return model

def predict_price(model, date):
    date_ordinal = datetime.toordinal(pd.to_datetime(date))
    date_df = pd.DataFrame([[date_ordinal]], columns=['Date_ordinal'])
    prediction = model.predict(date_df)
    return prediction[0]

def plot_prediction(stock_data, ticker, prediction_date, predicted_price):
    plt.figure(figsize=(12, 6))
    plt.plot(stock_data.index, stock_data['Close'], label='Historical Data')
    plt.scatter(prediction_date, predicted_price, color='red', label='Prediction')
    plt.title(f'{ticker} Stock Price Prediction')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.grid(True)
    plt.savefig('prediction_plot.png')
    return 'prediction_plot.png'

def predict_stock(ticker, date):
    stock_data = get_stock_data(ticker)
    
    if stock_data.empty:
        return "No data found for the given ticker.", None
    
    latest_price = stock_data['Close'].iloc[-1]
    
    processed_data = preprocess_data(stock_data)
    model = train_model(processed_data)
    
    try:
        predicted_price = predict_price(model, date)
        plot_path = plot_prediction(stock_data, ticker, pd.to_datetime(date), predicted_price)
        return f"The predicted closing price for {ticker} on {date} is: ${predicted_price:.2f}", plot_path
    except ValueError:
        return "Invalid date format. Please enter the date in YYYY-MM-DD format.", None

# Gradio app interface

inputs = [
    gr.Textbox(label="Enter the stock ticker"),
    gr.Textbox(label="Enter the date (YYYY-MM-DD) for the prediction")
]
outputs = [
    gr.Text(label="Prediction"),
    gr.Image(label="Prediction Plot")
]

gr.Interface(fn=predict_stock, inputs=inputs, outputs=outputs, title="Stock Price Prediction").launch(share=True)