import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM, pipeline

# Load the pre-trained BERT model and tokenizer for sentiment analysis
sentiment_model_name = "nlptown/bert-base-multilingual-uncased-sentiment"
sentiment_tokenizer = AutoTokenizer.from_pretrained(sentiment_model_name)
sentiment_model = AutoModelForSequenceClassification.from_pretrained(sentiment_model_name)
nlp_pipeline = pipeline('sentiment-analysis', model=sentiment_model, tokenizer=sentiment_tokenizer)

# Load the pre-trained T5 model and tokenizer for reason extraction
reason_model_name = "google/flan-t5-large"
reason_tokenizer = AutoTokenizer.from_pretrained(reason_model_name)
reason_model = AutoModelForSeq2SeqLM.from_pretrained(reason_model_name)
reason_extraction_pipeline = pipeline("text2text-generation", model=reason_model, tokenizer=reason_tokenizer)

# Function to classify sentiment as Positive, Neutral, or Negative
def classify_sentiment(feedback):
    if isinstance(feedback, str):
        # Truncate review if it exceeds 512 tokens for sentiment analysis
        result = nlp_pipeline(feedback[:512])
        label = result[0]['label']
        
        # Convert Hugging Face star labels to "Positive", "Negative", or "Neutral"
        if label in ['4 stars', '5 stars']:
            return "Positive"
        elif label == '3 stars':
            return "Neutral"
        else:
            return "Negative"
    else:
        return "No Feedback"

# Function to generate the prompt for reason extraction
def generate_prompt(feedback_text):
    return f"""
    You are given customer feedback. Please identify the key product attribute or aspect that the customer is discussing, whether positive or negative.
    Focus on words like price, quality, smell, packaging, durability, or any other relevant product feature.
    If no clear aspect can be identified, return "Could not extract a specific reason".
    
    Feedback: '{feedback_text}'
    """

# Function to post-process reasons and filter based on known product aspects
def post_process_reason(reason):
    possible_reasons = ["price", "quality", "smell", "packaging", "durability"]
    for r in possible_reasons:
        if r in reason.lower():
            return r
    return reason

# Function to extract the reason for the sentiment
def analyze_feedback(feedback_text):
    feedback_text = feedback_text.strip()
    prompt = generate_prompt(feedback_text)
    
    try:
        result = reason_extraction_pipeline(prompt, max_length=50)
        reason = result[0]['generated_text'].strip()
        return post_process_reason(reason)
    except Exception as e:
        print(f"Error processing feedback: {feedback_text}")
        print(e)
        return "Error in processing"

# Function to integrate both sentiment classification and reason extraction
def sentiment_and_reason(review_text):
    sentiment = classify_sentiment(review_text)
    reason = analyze_feedback(review_text)
    return f"Sentiment: {sentiment}\nReason: {reason}"

# Define the Gradio interface
def main_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## Sentiment and Reason Extraction for Product Reviews")
        
        with gr.Row():
            review_input = gr.Textbox(label="Enter your review:", lines=5, placeholder="Type your product review here...")
            output_box = gr.Textbox(label="Output:", lines=5, placeholder="Sentiment and reason will be displayed here...")
        
        analyze_button = gr.Button("Analyze")
        
        analyze_button.click(fn=sentiment_and_reason, inputs=review_input, outputs=output_box)
    
    return demo

# Launch the Gradio app
main_interface().launch()