import streamlit as st
from transformers import pipeline
from utils import read_poems_from_directory
import matplotlib.pyplot as plt

sentiment_classifier = pipeline("sentiment-analysis")
emotion_classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)

def analyze_poem(poem):
    sentiment = sentiment_classifier(poem)[0]
    emotion_scores = emotion_classifier(poem)[0]
    emotion = max(emotion_scores, key=lambda x: x['score'])['label']
    return sentiment, emotion, emotion_scores

def plot_emotion_scores(emotion_scores):
    emotions = [score['label'] for score in emotion_scores]
    scores = [score['score'] for score in emotion_scores]
    
    fig, ax = plt.subplots()
    ax.bar(emotions, scores)
    ax.set_xlabel("Emotion")
    ax.set_ylabel("Score")
    ax.set_title("Emotion Scores")
    plt.xticks(rotation=45)
    plt.tight_layout()
    return fig

def analyze_individual_page():
    st.header("Analyze Poems")
    poems_directory = "poems"
    
    poems = read_poems_from_directory(poems_directory)
    
    if poems:
        for i, poem in enumerate(poems, start=1):
            st.subheader(f"Poem {i}")
            st.text(poem)
            sentiment, emotion, emotion_scores = analyze_poem(poem)
            st.write(f"Sentiment: {sentiment['label']} (score: {sentiment['score']:.2f})")
            st.write(f"Emotion: {emotion}")
            
            # Plot emotion scores
            fig = plot_emotion_scores(emotion_scores)
            st.pyplot(fig)
    else:
        st.warning("No poems found in the 'poems' directory.")

def analyze_sentiment(poems):
    sentiment_labels = []
    for poem in poems:
        sentiment = sentiment_classifier(poem)[0]
        sentiment_labels.append(sentiment['label'])
    return sentiment_labels