import yfinance as yf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr
from sklearn.linear_model import LinearRegression
from transformers import pipeline

# Load Hugging Face summarization pipeline
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")


# Step 1: Fetch historical stock data
def get_stock_data(ticker):
    df = yf.download(ticker, start="2020-01-01", end="2024-12-31")
    return df[['Close']]


# Step 2: Prepare training data
def prepare_data(df):
    df['Target'] = df['Close'].shift(-1)
    df = df.dropna()
    X = df[['Close']]
    y = df['Target']
    return X, y


# Step 3: Train a simple regression model
def train_model(X, y):
    model = LinearRegression()
    model.fit(X, y)
    return model


# Step 4: Predict future prices (next 7 days)
def predict_future(model, last_price, days=7):
    predictions = []
    current = last_price
    for _ in range(days):
        next_price = model.predict([[current]])[0]
        predictions.append(next_price)
        current = next_price
    return predictions


# Step 5: Generate AI explanation using Hugging Face summarizer
def generate_explanation(df_tail):
    text_data = df_tail.to_string()
    summary = summarizer(text_data, max_length=100, min_length=30, do_sample=False)
    return summary[0]['summary_text']


# Step 6: Combine all into a Gradio interface function
def stock_predictor(ticker):
    try:
        df = get_stock_data(ticker)
        X, y = prepare_data(df)
        model = train_model(X, y)
        last_price = X.iloc[-1][0]
        predictions = predict_future(model, last_price)

        # Convert predictions to a plot
        plt.figure()
        plt.plot(range(1, 8), predictions, marker='o')
        plt.title(f"{ticker} Stock Price Predictions (Next 7 Days)")
        plt.xlabel("Days")
        plt.ylabel("Predicted Price")
        plt.grid(True)
        plot_path = "prediction_plot.png"
        plt.savefig(plot_path)
        plt.close()

        # Generate explanation using last 5 days of data
        explanation = generate_explanation(df.tail())

        return plot_path, predictions, explanation

    except Exception as e:
        return None, [], f"Error: {str(e)}"


# Step 7: Gradio Interface
interface = gr.Interface(
    fn=stock_predictor,
    inputs=gr.Textbox(label="Enter Stock Ticker (e.g., AAPL)"),
    outputs=[
        gr.Image(label="📈 Predicted Price Plot"),
        gr.Textbox(label="📊 Predicted Prices (Next 7 Days)"),
        gr.Textbox(label="🧠 AI Explanation of Recent Trends")
    ],
    title="Stock Price Predictor + AI Insight Generator",
    description="Enter a stock ticker (like AAPL or GOOG) to see predicted prices and a natural language explanation of recent trends using Hugging Face models."
)

interface.launch()