# app.py - Audio & Text Sentiment Analyzer
# Uses exact model: google-bert/bert-base-uncased (Masked LM)
# Runs locally with Gradio interface

import gradio as gr
import whisper
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForMaskedLM
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

print("Loading models... This may take a moment.")

# Load Whisper for audio transcription
whisper_model = whisper.load_model("base")  # Fast and works well; use "small" for better accuracy

# Load exact requested BERT model
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
model = AutoModelForMaskedLM.from_pretrained("google-bert/bert-base-uncased")
model.eval()

print("Models loaded successfully!")

# Function to get [CLS] embedding
def get_cls_embedding(text):
    inputs = tokenizer(text[:512], return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs, output_hidden_states=True)
        cls_embedding = outputs.hidden_states[-1][:, 0, :].cpu().numpy()
    return cls_embedding.flatten()

# Training examples for simple sentiment classifier
example_texts = [
    "I love this, it's absolutely amazing", "Best thing ever", "Fantastic experience",
    "Highly recommend", "Super happy with it", "This is terrible", "Worst product ever",
    "Very disappointed", "Complete waste", "Poor quality", "It's okay", "Nothing special",
    "Arrived on time", "Works as expected", "Average"
]

example_labels = [1,1,1,1,1, -1,-1,-1,-1,-1, 0,0,0,0,0]  # 1=Positive, -1=Negative, 0=Neutral

# Train classifier
X_train = np.array([get_cls_embedding(t) for t in example_texts])
y_train = np.array(example_labels)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)

clf = LogisticRegression(multi_class='ovr', class_weight='balanced')
clf.fit(X_train_scaled, y_train)

print("Sentiment classifier trained!")

# Predict sentiment with stars and confidence
def predict_sentiment(text):
    if not text.strip():
        return "⭐⭐⭐ Neutral", "0%"
    
    embedding = get_cls_embedding(text)
    embedding_scaled = scaler.transform([embedding])
    
    probabilities = clf.predict_proba(embedding_scaled)[0]
    pred = clf.predict(embedding_scaled)[0]
    confidence = np.max(probabilities) * 100
    conf_str = f"{confidence:.1f}%"
    
    if pred == 1:
        level = "⭐⭐⭐⭐⭐ Very Positive"
    elif pred == -1:
        level = "⭐ Very Negative"
    else:
        level = "⭐⭐⭐ Neutral"
    
    return level, conf_str

# Main analysis function
def analyze_input(audio_path, input_text):
    # Prefer 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:
        try:
            print("Transcribing audio...")
            result = whisper_model.transcribe(audio_path)
            final_text = result["text"].strip()
            if not final_text:
                return "No speech detected in the audio.", "", "", "Please speak clearly and try again."
        except Exception as e:
            return "Error transcribing audio.", "", "", f"Error: {str(e)}"
    
    else:
        return "No input provided.", "", "", "Please type text or record/upload audio."
    
    # Perform sentiment analysis
    level, confidence = predict_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 + Sentiment Analyzer")
    gr.Markdown("""
    - Record or upload audio → **Automatically transcribed**  
    - Or type text directly  
    - Analyzes sentiment using **google-bert/bert-base-uncased** ([CLS] embedding)
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="Record or Upload Audio"
            )
            
            gr.Markdown("**OR**")
            
            text_input = gr.Textbox(
                label="Type or Paste Text",
                placeholder="Enter your review or feedback here...",
                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
    - Works with any language (Whisper handles transcription)
    - Uses raw BERT base model → educational demo
    - Run locally, no data leaves your machine
    - Made with ❤️ in Accra by Chris (@chrisbekor99)
    """)


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