import streamlit as st
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
import os

@st.cache_resource
def load_model():
    """Load your trained model"""
    try:
        # Load from the same directory as the script
        model_path = os.path.dirname(__file__)
        
        st.info(f"Loading model from: {model_path}")
        
        # Load model
        model = DistilBertForSequenceClassification.from_pretrained(
            model_path,
            local_files_only=True
        )
        tokenizer = DistilBertTokenizer.from_pretrained(
            model_path,
            local_files_only=True
        )
        
        st.success("✅ Model loaded successfully!")
        return model, tokenizer
        
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None, None

def predict_text(text, model, tokenizer):
    """Make prediction"""
    inputs = tokenizer(
        text, 
        return_tensors="pt", 
        truncation=True, 
        padding=True,
        max_length=128
    )
    
    model.eval()
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        predicted_class = torch.argmax(predictions, dim=-1).item()
        confidence = predictions[0][predicted_class].item()
    
    return predicted_class, confidence

# Main App
st.title("🏥 Medical Text Classifier")
st.write("Enter text to classify as Medical or Non-Medical")

# Load model
model, tokenizer = load_model()

if model is not None:
    # Text input
    user_input = st.text_area(
        "Enter your text:",
        placeholder="Example: I have a headache and need to see a doctor...",
        height=100
    )
    
    # Classify button
    if st.button("🔍 Classify Text", type="primary"):
        if user_input.strip():
            with st.spinner("Analyzing..."):
                predicted_class, confidence = predict_text(user_input, model, tokenizer)
                
                # Show results
                labels = ["Non-Medical", "Medical"]
                result = labels[predicted_class]
                
                if predicted_class == 1:  # Medical
                    st.success(f"🏥 **{result}**")
                else:  # Non-Medical
                    st.info(f"ℹ️ **{result}**")
                
                st.write(f"**Confidence:** {confidence:.1%}")
                st.progress(confidence)
        else:
            st.warning("Please enter some text to classify!")

else:
    st.error("Failed to load the model. Please check the files.")