import streamlit as st
import os
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)

# Directory for poems
poems_directory = "./poems"
os.makedirs(poems_directory, exist_ok=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 Individual Poems")

    # Sidebar for file upload and listing files
    st.sidebar.title("Upload New Poem")
    uploaded_file = st.sidebar.file_uploader("Choose a text file", type="txt")

    if uploaded_file is not None:
        with open(os.path.join(poems_directory, uploaded_file.name), "wb") as f:
            f.write(uploaded_file.getbuffer())
        st.sidebar.success(f"Uploaded {uploaded_file.name}")

    st.sidebar.title("Available Poems")
    poem_files = [f for f in os.listdir(poems_directory) if f.endswith(".txt")]
    st.sidebar.write("\n".join(poem_files))

    poems = read_poems_from_directory(poems_directory)
    
    if poems:
        selected_poem = st.selectbox("Select a poem to analyze", range(len(poems)), format_func=lambda i: f"Poem {i+1}")
        
        st.subheader(f"Analyzing Poem {selected_poem + 1}")
        st.text(poems[selected_poem])
        sentiment, emotion, emotion_scores = analyze_poem(poems[selected_poem])
        st.write(f"Sentiment: {sentiment['label']} (score: {sentiment['score']:.2f})")
        st.write(f"Dominant 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