import gradio as gr
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    pipeline
)
import pandas as pd
import numpy as np
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# Constants
MODEL_PATH = "./imdb_sentiment_final"
TITLE = "🎬 IMDB Sentiment Classifier"
DESCRIPTION = """
A binary sentiment classifier fine-tuned on IMDB movie reviews using LoRA.
The model classifies text as **Positive** or **Negative** sentiment.
"""

# Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load model and tokenizer
@gr.Cache()
def load_model():
    """Load the model and tokenizer with caching"""
    print("Loading model...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
        model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
        model.to(device)
        model.eval()
        print("Model loaded successfully!")
        return model, tokenizer
    except Exception as e:
        print(f"Error loading model: {e}")
        # Fallback to a default model
        from transformers import DistilBertForSequenceClassification
        model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)
        tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
        model.to(device)
        model.eval()
        return model, tokenizer

# Initialize model
model, tokenizer = load_model()

# Create inference pipeline
sentiment_pipeline = pipeline(
    "text-classification",
    model=model,
    tokenizer=tokenizer,
    device=0 if torch.cuda.is_available() else -1
)

def predict_sentiment(text, batch_size=1):
    """
    Predict sentiment for single or multiple texts
    """
    if not text or not text.strip():
        return {"error": "Please enter some text"}
    
    # Single text
    if isinstance(text, str):
        try:
            result = sentiment_pipeline(text)[0]
            confidence = result['score']
            label = "Positive" if result['label'] == "LABEL_1" else "Negative"
            
            return {
                "text": text[:200] + "..." if len(text) > 200 else text,
                "sentiment": label,
                "confidence": f"{confidence:.2%}",
                "icon": "😊" if label == "Positive" else "😞",
                "color": "green" if label == "Positive" else "red"
            }
        except Exception as e:
            return {"error": f"Prediction error: {str(e)}"}
    
    # Batch processing (if text is a list)
    else:
        results = []
        for t in text:
            if t.strip():
                try:
                    result = sentiment_pipeline(t)[0]
                    confidence = result['score']
                    label = "Positive" if result['label'] == "LABEL_1" else "Negative"
                    results.append({
                        "text": t[:100] + "..." if len(t) > 100 else t,
                        "sentiment": label,
                        "confidence": f"{confidence:.2%}",
                        "icon": "😊" if label == "Positive" else "😞"
                    })
                except:
                    results.append({"text": t[:100], "sentiment": "Error", "confidence": "N/A"})
        return {"batch_results": results}

def batch_predict(input_text):
    """
    Handle batch input from text area
    """
    texts = [t.strip() for t in input_text.split('\n') if t.strip()]
    if not texts:
        return {"error": "No valid text provided"}
    
    if len(texts) == 1:
        return predict_sentiment(texts[0])
    else:
        return predict_sentiment(texts)

def analyze_file(file):
    """
    Analyze sentiment from uploaded text file
    """
    if file is None:
        return {"error": "No file uploaded"}
    
    try:
        # Read file
        if file.name.endswith('.txt'):
            with open(file.name, 'r', encoding='utf-8') as f:
                content = f.read()
        elif file.name.endswith('.csv'):
            df = pd.read_csv(file.name)
            if 'review' in df.columns:
                content = "\n".join(df['review'].astype(str).tolist())
            else:
                content = "\n".join(df.iloc[:, 0].astype(str).tolist())
        else:
            return {"error": "Unsupported file format. Use .txt or .csv"}
        
        texts = [t.strip() for t in content.split('\n') if t.strip()]
        if len(texts) > 100:  # Limit for demo
            texts = texts[:100]
        
        results = []
        for text in texts:
            result = predict_sentiment(text)
            if "error" not in result:
                results.append({
                    "Text": text[:50] + "..." if len(text) > 50 else text,
                    "Sentiment": result['sentiment'],
                    "Confidence": result['confidence']
                })
        
        if results:
            df_results = pd.DataFrame(results)
            summary = df_results['Sentiment'].value_counts()
            
            return {
                "results_table": df_results,
                "summary": f"Positive: {summary.get('Positive', 0)}, Negative: {summary.get('Negative', 0)}",
                "total_reviews": len(results)
            }
        else:
            return {"error": "No valid text found in file"}
            
    except Exception as e:
        return {"error": f"File processing error: {str(e)}"}

# Example reviews for demo
EXAMPLE_REVIEWS = [
    "This movie was absolutely fantastic! The acting was superb and the storyline kept me engaged throughout.",
    "Terrible film. I regret wasting two hours of my life watching this boring nonsense.",
    "The cinematography was beautiful but the plot was weak and predictable.",
    "An emotional rollercoaster with brilliant performances from the entire cast.",
    "Not bad, but nothing special either. Average entertainment for a lazy Sunday."
]

# Custom CSS for better styling
custom_css = """
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}
.sentiment-positive {
    color: green;
    font-weight: bold;
}
.sentiment-negative {
    color: red;
    font-weight: bold;
}
.result-box {
    padding: 20px;
    border-radius: 10px;
    margin: 10px 0;
    background: #f5f5f5;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css, title=TITLE) as demo:
    gr.HTML(f"""
    <div style="text-align: center;">
        <h1>🎬 IMDB Sentiment Classifier</h1>
        <p>Fine-tuned on 50K movie reviews using LoRA with DistilBERT (≤ 110M params)</p>
        <p style="color: #666; font-size: 0.9em;">Model Type: Binary Sentiment Classification (Positive/Negative)</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            # Single text input
            gr.Markdown("### 📝 Enter Your Movie Review")
            input_text = gr.Textbox(
                label="Review Text",
                placeholder="Type your movie review here...",
                lines=4,
                elem_id="input_text"
            )
            
            with gr.Row():
                submit_btn = gr.Button("Analyze Sentiment 🚀", variant="primary")
                clear_btn = gr.Button("Clear", variant="secondary")
            
            # Batch input
            gr.Markdown("### 📚 Batch Analysis")
            batch_input = gr.Textbox(
                label="Enter multiple reviews (one per line)",
                placeholder="Review 1\nReview 2\nReview 3\n...",
                lines=5
            )
            batch_btn = gr.Button("Analyze All", variant="secondary")
            
            # File upload
            gr.Markdown("### 📁 Upload File")
            file_input = gr.File(
                label="Upload .txt or .csv file",
                file_types=[".txt", ".csv"]
            )
            file_btn = gr.Button("Analyze File", variant="secondary")
            
            # Examples
            gr.Markdown("### ✨ Example Reviews")
            gr.Examples(
                examples=EXAMPLE_REVIEWS,
                inputs=input_text,
                label="Click an example to try it out"
            )
        
        with gr.Column(scale=1):
            # Results display
            gr.Markdown("### 📊 Results")
            result_box = gr.JSON(label="Analysis Result", elem_classes="result-box")
            
            # Visual output
            sentiment_output = gr.Label(label="Sentiment", num_top_classes=2)
            
            # Confidence meter
            confidence_output = gr.Slider(
                minimum=0,
                maximum=100,
                label="Confidence Level",
                interactive=False
            )
            
            # Stats
            gr.Markdown("### 📈 Statistics")
            stats_output = gr.Markdown("")
            
            # Model info
            with gr.Accordion("ℹ️ Model Information", open=False):
                gr.Markdown(f"""
                **Model Details:**
                - Base Model: DistilBERT-base-uncased
                - Parameters: ~66M (with LoRA adaptation)
                - Training: Fine-tuned on IMDB dataset
                - Task: Binary sentiment classification
                - Framework: PyTorch + Hugging Face Transformers
                
                **Last Updated:** {datetime.now().strftime('%Y-%m-%d')}
                
                **Note:** This is a demo model trained on a subset of data.
                """)
    
    # Functions to update all outputs
    def update_all_outputs(result):
        """Update all output components based on prediction result"""
        if "error" in result:
            return [
                result,
                {"Positive": 0, "Negative": 0},
                50,
                f"**Error:** {result['error']}"
            ]
        elif "batch_results" in result:
            # For batch results
            positive_count = sum(1 for r in result['batch_results'] if r['sentiment'] == 'Positive')
            negative_count = sum(1 for r in result['batch_results'] if r['sentiment'] == 'Negative')
            avg_confidence = 75  # Placeholder
            summary = f"**Batch Analysis:** {len(result['batch_results'])} reviews\nPositive: {positive_count}, Negative: {negative_count}"
            
            return [
                result,
                {"Positive": positive_count, "Negative": negative_count},
                avg_confidence,
                summary
            ]
        else:
            # For single result
            confidence_percent = float(result['confidence'].strip('%')) / 100
            sentiment_dict = {"Positive": 0, "Negative": 0}
            sentiment_dict[result['sentiment']] = 1
            
            stats = f"""
            **Review:** {result['text']}
            
            **Sentiment:** <span class="sentiment-{result['sentiment'].lower()}">{result['sentiment']} {result['icon']}</span>
            
            **Confidence:** {result['confidence']}
            """
            
            return [
                result,
                sentiment_dict,
                confidence_percent * 100,
                stats
            ]
    
    # Set up event handlers
    submit_btn.click(
        fn=predict_sentiment,
        inputs=input_text,
        outputs=result_box
    ).then(
        fn=update_all_outputs,
        inputs=result_box,
        outputs=[result_box, sentiment_output, confidence_output, stats_output]
    )
    
    clear_btn.click(
        fn=lambda: [None, {"Positive": 0, "Negative": 0}, 50, ""],
        outputs=[input_text, sentiment_output, confidence_output, stats_output]
    )
    
    batch_btn.click(
        fn=batch_predict,
        inputs=batch_input,
        outputs=result_box
    ).then(
        fn=update_all_outputs,
        inputs=result_box,
        outputs=[result_box, sentiment_output, confidence_output, stats_output]
    )
    
    file_btn.click(
        fn=analyze_file,
        inputs=file_input,
        outputs=result_box
    ).then(
        fn=update_all_outputs,
        inputs=result_box,
        outputs=[result_box, sentiment_output, confidence_output, stats_output]
    )
    
    # Auto-run when example is clicked
    input_text.change(
        fn=predict_sentiment,
        inputs=input_text,
        outputs=result_box
    ).then(
        fn=update_all_outputs,
        inputs=result_box,
        outputs=[result_box, sentiment_output, confidence_output, stats_output]
    )

# Launch configuration
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )