import streamlit as st
from transformers import pipeline
from bs4 import BeautifulSoup
import requests

# Set up models
ner_model = "Cassie-0121/fin-bert-finetuned-ner"
sentiment_model = "yiyanghkust/finbert-tone"
text_gen_model = "gpt2"

ner = pipeline("ner", model=ner_model)
sentiment_analyzer = pipeline("sentiment-analysis", model=sentiment_model)
text_generator = pipeline("text-generation", model=text_gen_model)

# Function to scrape recent stock news
def get_stock_news(stock_symbol):
    url = f"https://finance.yahoo.com/quote/{stock_symbol}/news?p={stock_symbol}"
    headers = {"User-Agent": "Mozilla/5.0"}
    response = requests.get(url, headers=headers)
    soup = BeautifulSoup(response.text, 'html.parser')
    articles = soup.find_all('h3', class_="Mb(5px)")

    news_list = []
    for article in articles[:5]:  # Limit to the top 5 news articles
        headline = article.get_text()
        news_list.append(headline)
    return news_list

# App title
st.title("AI-Powered Financial Analysis App")

# Sidebar with stock data examples
st.sidebar.header("Stock Data & Analysis")
examples = {
    "Apple Inc. (AAPL)": "AAPL",
    "Tesla Inc. (TSLA)": "TSLA",
    "Amazon.com Inc. (AMZN)": "AMZN",
    "Microsoft Corp. (MSFT)": "MSFT",
    "Alphabet Inc. (GOOGL)": "GOOGL",
    "Other": ""  # Placeholder for custom input
}

selected_example = st.sidebar.selectbox("Select a stock symbol or choose 'Other' to enter custom text:", list(examples.keys()))

# If "Other" is selected, provide an input box for custom text
if selected_example == "Other":
    input_text = st.sidebar.text_area("Enter your own stock data for analysis:")
    stock_symbol = None
else:
    stock_symbol = examples[selected_example]
    input_text = f"Latest news for {selected_example}"

# Display selected or inputted text for review
st.subheader("Stock Data & Analysis")
st.write(input_text if input_text else "Please select a stock or enter custom text for analysis.")

# Fetch and display news for the selected stock
if stock_symbol:
    st.subheader("Latest Stock News")
    news_articles = get_stock_news(stock_symbol)
    for article in news_articles:
        st.write(f"- {article}")

# Key Financial Entities extraction with filtering
st.subheader("Extracted Key Financial Entities")
if input_text:
    entities = ner(input_text)
    filtered_entities = [entity for entity in entities if entity['score'] > 0.7 and entity['word'].isalpha()]  # Filter low-score and non-alphabetic tokens
    for entity in filtered_entities:
        st.write(f"Entity: {entity['word']}, Label: {entity.get('entity', 'N/A')}, Score: {entity['score']:.2f}")

# Sentiment Analysis
st.subheader("Sentiment Analysis")
if input_text:
    sentiment = sentiment_analyzer(input_text)
    for result in sentiment:
        st.write(f"Sentiment: {result['label']}, Score: {result['score']:.2f}")

# Investment Advice or Strategy Generation with better prompt handling
st.subheader("Investment Advice or Strategy")
if input_text:
    prompt = f"Provide a clear and concise investment strategy for {selected_example if selected_example != 'Other' else 'the selected stock'} based on recent news and financial performance. "
    advice = text_generator(prompt + input_text, max_length=80, num_return_sequences=1)
    st.write(advice[0]['generated_text'])
else:
    st.write("No investment advice generated. Please select a stock or enter custom text.")

# Make the app visually more attractive
st.markdown(
    """
    <style>
        .stApp {
            background-color: #F5F5F5;
        }
        .stSidebar {
            background-color: #333333;
            color: white;
        }
    </style>
    """,
    unsafe_allow_html=True
)

# Footer
st.sidebar.write("Powered by Hugging Face and Streamlit")