import yfinance as yf
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
import gradio as gr
import matplotlib.pyplot as plt
from datetime import datetime, timedelta

# Define stock tickers
STOCK_TICKERS = [
    "AAPL",  # Apple
    "GOOGL", # Alphabet
    "MSFT",  # Microsoft
    "AMZN",  # Amazon
    "TSLA",  # Tesla
    "META",  # Meta Platforms
    "NVDA",  # NVIDIA
    "JPM",   # JPMorgan Chase
    "V",     # Visa
    "NFLX"   # Netflix
]

def fetch_stock_data(ticker: str, start_date: str, end_date: str) -> pd.DataFrame:
    """
    Fetches historical stock data from Yahoo Finance.
    """
    stock = yf.Ticker(ticker)
    data = stock.history(start=start_date, end=end_date)
    return data

def preprocess_data(data: pd.DataFrame) -> (np.ndarray, np.ndarray):
    """
    Preprocesses the stock data for Random Forest Regressor.
    """
    data['Target'] = data['Close'].shift(-1)  # Predict next day's close price
    data = data[:-1]

    # Create features using the previous 5 days' close prices
    for i in range(1, 6):
        data[f'Close_{i}'] = data['Close'].shift(i)

    data.dropna(inplace=True)

    X = data[[f'Close_{i}' for i in range(1, 6)]].values
    y = data['Target'].values

    return X, y

def train_model(X: np.ndarray, y: np.ndarray) -> RandomForestRegressor:
    """
    Trains the Random Forest Regressor model.
    """
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)
    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)

    predictions = model.predict(X_test)
    mse = mean_squared_error(y_test, predictions)
    print(f"Model Mean Squared Error: {mse}")

    return model

def make_prediction(model: RandomForestRegressor, recent_data: pd.DataFrame, days: int) -> np.ndarray:
    """
    Makes predictions for the next 'days' closing prices.
    """
    predicted_prices = []
    last_close = recent_data['Close'].values[-5:]

    for _ in range(days):
        if len(last_close) < 5:
            raise ValueError("Not enough data to make a prediction.")

        X_new = last_close[::-1].reshape(1, -1)  # Reverse to match feature order
        predicted_price = model.predict(X_new)[0]
        predicted_prices.append(predicted_price)

        # Update the last_close with the predicted price for the next iteration
        last_close = np.roll(last_close, -1)
        last_close[-1] = predicted_price

    return predicted_prices

def buy_or_sell(current_price: float, predicted_price: float) -> str:
    """
    Determines whether to buy or sell based on price prediction.
    """
    return "Buy" if predicted_price > current_price else "Sell"

def stock_prediction_app(ticker: str, start_date: str, end_date: str):
    """
    Main function to handle stock prediction and return outputs.
    """
    data = fetch_stock_data(ticker, start_date, end_date)
    if data.empty:
        return "No data available for the selected dates.", "N/A", "N/A", "N/A", "N/A", None

    # Calculate basic stats
    current_price = data['Close'].iloc[-1]
    highest_price = data['Close'].max()
    lowest_price = data['Close'].min()
    start_price = data['Close'].iloc[0]
    percentage_change = ((current_price - start_price) / start_price) * 100

    # Preprocess and train model
    X, y = preprocess_data(data)
    if len(X) == 0:
        return "Insufficient data to train the model.", current_price, highest_price, lowest_price, "Error", None

    model = train_model(X, y)

    # Predict the next 60 days
    predicted_prices = make_prediction(model, data, days=60)

    # Create dates for predictions
    future_dates = pd.date_range(start=data.index[-1] + timedelta(days=1), periods=60, freq='B')  # Business days

    # Prepare data for plotting
    plt.figure(figsize=(14, 7))
    
    # Plot historical prices
    plt.plot(data.index, data['Close'], label='Historical Close Price', color='blue', linewidth=2)
    
    # Plot predicted prices
    plt.plot(future_dates, predicted_prices, label='Predicted Prices', color='orange', linestyle='--', linewidth=2)

    # Highlight current price
    plt.axhline(y=current_price, color='green', linestyle='--', label='Current Close Price', linewidth=1.5)

    # Adding labels and grid for clarity
    plt.title(f'{ticker} Price Prediction', fontsize=16)
    plt.xlabel('Date', fontsize=14)
    plt.ylabel('Price ($)', fontsize=14)
    plt.xticks(rotation=45)
    plt.legend()
    plt.grid()
    plt.tight_layout()
    fig = plt.gcf()
    plt.close()

    # Format outputs
    percentage_change_str = f"{percentage_change:.2f}%"
    decision = buy_or_sell(current_price, predicted_prices[0])  # Decision based on the first predicted price

    return percentage_change_str, current_price, highest_price, lowest_price, decision, fig

# Gradio interface setup
iface = gr.Interface(
    fn=stock_prediction_app,
    inputs=[
        gr.Dropdown(choices=STOCK_TICKERS, label="Select Stock Ticker"),
        gr.Textbox(label="Enter Start Date (YYYY-MM-DD)", placeholder="e.g., 2020-01-01"),
        gr.Textbox(label="Enter End Date (YYYY-MM-DD)", placeholder="e.g., 2023-12-31")
    ],
    outputs=[
        gr.Textbox(label="Percentage Change"),
        gr.Number(label="Current Closing Price"),
        gr.Number(label="Highest Closing Price"),
        gr.Number(label="Lowest Closing Price"),
        gr.Textbox(label="Decision (Buy/Sell)"),
        gr.Plot(label="Stock Performance")
    ],
    title="Stock Prediction App",
    description="Predict stock trends using historical data. Please enter dates in the format YYYY-MM-DD."
)

# Launch the app
iface.launch()