# app.py - Audio & Text Sentiment Analyzer
# Transcription: openai/whisper-base.en (official HF version)
# Sentiment: nlptown/bert-base-multilingual-uncased-sentiment (5-star accurate model)

import gradio as gr
import torch
import numpy as np
import librosa
from transformers import (
    AutoProcessor,
    AutoModelForSpeechSeq2Seq,
    AutoTokenizer,
    AutoModelForSequenceClassification
)
import torch.nn.functional as F

print("Loading models... Please wait.")

# === Load Whisper exactly as requested ===
processor = AutoProcessor.from_pretrained("openai/whisper-base.en")
whisper_model = AutoModelForSpeechSeq2Seq.from_pretrained("openai/whisper-base.en")
whisper_model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
whisper_model.to(device)

# === Load Sentiment model exactly as requested ===
sentiment_tokenizer = AutoTokenizer.from_pretrained('nlptown/bert-base-multilingual-uncased-sentiment')
sentiment_model = AutoModelForSequenceClassification.from_pretrained('nlptown/bert-base-multilingual-uncased-sentiment')
sentiment_model.eval()
sentiment_model.to(device)

print("All models loaded successfully!")

# Transcribe audio using official Whisper
def transcribe_audio(audio_path):
    if audio_path is None:
        return ""
    
    try:
        # Load and resample to 16kHz
        speech, _ = librosa.load(audio_path, sr=16000)
        
        # Process input
        input_features = processor(speech, sampling_rate=16000, return_tensors="pt").input_features
        input_features = input_features.to(device)
        
        # Generate transcription
        with torch.no_grad():
            predicted_ids = whisper_model.generate(input_features)
        
        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
        return transcription.strip()
    
    except Exception as e:
        print(f"Transcription error: {e}")
        return "[Transcription failed]"

# Sentiment analysis with 5-star rating and confidence
def analyze_sentiment(text):
    if not text.strip():
        return "⭐⭐⭐ Neutral", "0%"
    
    inputs = sentiment_tokenizer(text[:512], return_tensors="pt", truncation=True).to(device)
    
    with torch.no_grad():
        logits = sentiment_model(**inputs).logits
    
    probabilities = F.softmax(logits, dim=-1)[0]
    predicted_class = torch.argmax(probabilities).item() + 1  # 1 to 5
    confidence = probabilities[predicted_class - 1].item() * 100
    conf_str = f"{confidence:.1f}%"
    
    stars = "⭐" * predicted_class
    if predicted_class == 1:
        level = f"{stars} Very Negative"
    elif predicted_class == 2:
        level = f"{stars} Negative"
    elif predicted_class == 3:
        level = f"{stars} Neutral"
    elif predicted_class == 4:
        level = f"{stars} Positive"
    else:
        level = f"{stars} Very Positive"
    
    return level, conf_str

# Main unified function
def analyze_input(audio_path, input_text):
    # Use typed text if provided
    if input_text and input_text.strip():
        final_text = input_text.strip()
    
    # Otherwise transcribe audio
    elif audio_path is not None:
        print("Transcribing audio...")
        final_text = transcribe_audio(audio_path)
        if not final_text or "failed" in final_text.lower():
            return "Transcription failed or no speech detected.", "", "", "Please try again with clearer English audio."
    
    else:
        return "No input provided.", "", "", "Please type text or record/upload audio."
    
    # Sentiment analysis
    level, confidence = analyze_sentiment(final_text)
    final_result = f"{level} (Confidence: {confidence})"
    
    return final_text, level, confidence, final_result

# Gradio Interface
with gr.Blocks(title="Audio & Text Sentiment Analyzer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎤✍️ Audio to Text + 5-Star Sentiment Analyzer")
    gr.Markdown("""
    - **Transcription**: OpenAI Whisper-base.en (excellent English accuracy)  
    - **Sentiment**: Multilingual BERT fine-tuned on reviews → accurate **1–5 star** ratings  
    - Record/upload audio **or** type text directly
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="Record or Upload Audio (English recommended)"
            )
            
            gr.Markdown("**OR**")
            
            text_input = gr.Textbox(
                label="Type or Paste Text",
                placeholder="Enter your review, feedback, or transcribed text...",
                lines=6
            )
    
    btn = gr.Button("Transcribe & Analyze Sentiment", variant="primary", size="lg")
    
    with gr.Column():
        gr.Markdown("### 📝 Transcribed / Entered Text")
        text_display = gr.Textbox(label="Text", lines=8, interactive=False)
        
        gr.Markdown("### 🌟 Sentiment Result")
        with gr.Row():
            level_out = gr.Textbox(label="Sentiment Level", scale=2)
            conf_out = gr.Textbox(label="Confidence", scale=1)
        
        result_out = gr.Textbox(label="Final Verdict", lines=2, interactive=False)
    
    btn.click(
        fn=analyze_input,
        inputs=[audio_input, text_input],
        outputs=[text_display, level_out, conf_out, result_out]
    )
    
    # gr.Markdown("""
    # ### Notes
    # - Best performance with **clear English speech**
    # - Sentiment model excels at review-style language (opinions, experiences)
    # - Confidence >80% = very reliable prediction
    # - Runs completely locally — perfect for privacy
    # - Built with ❤️ in Accra by Chris (@chrisbekor99)
    # """)


# Run app
if __name__ == "__main__":
    demo.launch()