# app.py
import streamlit as st
import pandas as pd
import re
from transformers import BertTokenizer, BertForSequenceClassification
import torch
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

# Load and display the WEBP image
image = Image.open("imagefintech.webp")  # Replace 'logo.webp' with your actual WEBP file path
st.image(image, caption="Sentiment Analysis App", use_column_width=True)

# Load the model and tokenizer
@st.cache_resource
def load_model():
    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
    model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=8)
    return tokenizer, model

tokenizer, model = load_model()

# Custom sentiments
sentiments = ["happy", "motivated", "growth", "optimistic", "jealousy", "frustrated", "decline", "angry"]

# Define the preprocessing function
def preprocess_text(text):
    text = re.sub(r'[^\w\s]', '', text.lower())  # Remove punctuation and lowercase
    text = re.sub(r'\d+', '', text)  # Remove numbers
    return text

# Title and instructions
st.title("Sentiment Analysis of Financial News")
st.write("Enter a sentence to analyze its sentiment across predefined categories.")

# Input text from user
text = st.text_input("Enter a sentence:", "")

if text:
    # Preprocess and tokenize
    cleaned_text = preprocess_text(text)
    inputs = tokenizer(cleaned_text, return_tensors="pt", truncation=True, padding=True)
    
    # Get model predictions
    with torch.no_grad():
        outputs = model(**inputs)
        sentiment_score = outputs.logits.softmax(dim=1)

    # Convert tensor to list for plotting
    score_list = sentiment_score.tolist()[0]

    # Display sentiment scores as a table
    st.subheader("Sentiment Scores")
    score_df = pd.DataFrame({"Sentiment": sentiments, "Score": score_list})
    st.dataframe(score_df)

    # Plot the sentiment scores
    st.subheader("Sentiment Score Chart")
    fig, ax = plt.subplots(figsize=(10, 6))
    mustard_yellow = "#FFDB58"

    # Plot bars with spacing and color
    ax.bar(np.arange(len(sentiments)) * 1.5, score_list, color=mustard_yellow, edgecolor="black", width=0.8)
    
    # Customize the plot
    ax.set_xlabel("Sentiments", color="black", fontsize=12)
    ax.set_ylabel("Scores", color="black", fontsize=12)
    ax.set_title("Sentiment Analysis of Financial News", color="black", fontsize=14)
    ax.set_xticks(np.arange(len(sentiments)) * 1.5)
    ax.set_xticklabels(sentiments, color="black", fontsize=10, rotation=45)
    ax.tick_params(axis="y", colors="black")

    # Display the plot in Streamlit
    st.pyplot(fig)