# Check for PyTorch installation
try:
    import torch
    print(f"PyTorch version: {torch.__version__}")
except ImportError:
    print("PyTorch is not installed. Please install PyTorch to run this script.")
    raise

from transformers import pipeline
import gradio as gr

# Initialize models as None
model1 = None
model2 = None

# Attempt to load the models and run test predictions
try:
    model1_name = "JimminDev/jim-text-class"
    model2_name = "JimminDev/Depressive-detector"
    print("Loading models...")
    
    model1 = pipeline("text-classification", model=model1_name)
    test_output1 = model1("Testing the first model with a simple sentence.")
    print("Harassment Detector test output:", test_output1)
    
    model2 = pipeline("text-classification", model=model2_name)
    test_output2 = model2("Testing the second model with a simple sentence.")
    print("Depressive Detector test output:", test_output2)
except Exception as e:
    print(f"Failed to load or run models: {e}")

# Prediction function with model selection and error handling
def predict_sentiment(text, model_choice):
    try:
        if model_choice == "Harassment Detector":
            if model1 is None:
                raise ValueError("Harassment Detector not loaded.")
            predictions = model1(text)
        elif model_choice == "Depressive Detector":
            if model2 is None:
                raise ValueError("Depressive Detector not loaded.")
            predictions = model2(text)
        else:
            raise ValueError("Invalid model choice.")
        
        return f"Label: {predictions[0]['label']}, Score: {predictions[0]['score']:.4f}"
    except Exception as e:
        return f"Error processing input: {e}"

# Define example inputs
examples = [
    ["I absolutely love this product! It has changed my life.", "Harassment Detector"],
    ["This is the worst movie I have ever seen. Completely disappointing.", "Harassment Detector"],
    ["I'm not sure how I feel about this new update. It has some good points, but also many drawbacks.", "Depressive Detector"],
    ["The customer service was fantastic! Very helpful and polite.", "Depressive Detector"],
    ["Honestly, this was quite a mediocre experience. Nothing special.", "Harassment Detector"]
]

# Gradio interface setup
iface = gr.Interface(
    fn=predict_sentiment,
    title="Sentiment Analysis",
    description="Enter text to analyze sentiment. Powered by Hugging Face Transformers.",
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter text here..."),
        gr.Radio(choices=["Harassment Detector", "Depressive Detector"], label="Select Model")
    ],
    outputs="text",
    examples=examples
)

if __name__ == "__main__":
    iface.launch()