import streamlit as st
from transformers import pipeline, AutoModelForImageClassification, AutoFeatureExtractor
from PIL import Image

# =======================
# Streamlit Page Config
# =======================
st.set_page_config(
    page_title="AI-Powered Skin Cancer Detection",
    page_icon="🩺",
    layout="wide",
    initial_sidebar_state="expanded"
)

# =======================
# Load Skin Cancer Model (PyTorch)
# =======================
@st.cache_resource
def load_model():
    """
    Load the pre-trained skin cancer classification model using PyTorch.
    """
    try:
        extractor = AutoFeatureExtractor.from_pretrained("Anwarkh1/Skin_Cancer-Image_Classification")
        model = AutoModelForImageClassification.from_pretrained("Anwarkh1/Skin_Cancer-Image_Classification")
        return pipeline("image-classification", model=model, feature_extractor=extractor, framework="pt")
    except Exception as e:
        st.error(f"Error loading the model: {e}")
        return None

model = load_model()

# =======================
# Local Explanation Generator
# =======================
def generate_local_explanation(label, confidence):
    """
    Generate a simple explanation for the classification result.
    """
    explanations = {
        "Melanoma": (
            "Melanoma is a serious type of skin cancer that develops in the cells that produce melanin. "
            "If detected early, it is often treatable. You should consult a dermatologist immediately."
        ),
        "Basal Cell Carcinoma": (
            "Basal Cell Carcinoma is a common form of skin cancer that grows slowly and is typically not life-threatening. "
            "Still, it requires medical attention to prevent further complications."
        ),
        "Benign Lesion": (
            "A benign lesion is a non-cancerous growth on the skin. While it is usually harmless, "
            "consulting a dermatologist can help ensure no further treatment is needed."
        ),
        "Other": (
            "The AI could not confidently classify the lesion. It's strongly recommended to consult a dermatologist for further evaluation."
        )
    }

    explanation = explanations.get(label, explanations["Other"])
    confidence_msg = f"The model is {confidence:.2%} confident in this prediction. "
    return confidence_msg + explanation

# =======================
# Streamlit App Title and Sidebar
# =======================
st.title("🔍 AI-Powered Skin Cancer Classification and Explanation")
st.write("Upload an image of a skin lesion, and the AI model will classify it and provide a detailed explanation.")

st.sidebar.info("""
**AI Cancer Detection Platform**  
This application uses AI to classify skin lesions and generate detailed explanations for informational purposes.  
It is not intended for medical diagnosis. Always consult a healthcare professional for medical advice.
""")

# =======================
# File Upload and Prediction
# =======================
uploaded_image = st.file_uploader("Upload a skin lesion image (PNG, JPG, JPEG)", type=["png", "jpg", "jpeg"])

if uploaded_image:
    # Display uploaded image
    image = Image.open(uploaded_image).convert("RGB")
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Perform classification
    if model is None:
        st.error("Model could not be loaded. Please try again later.")
    else:
        with st.spinner("Classifying the image..."):
            try:
                results = model(image)
                label = results[0]['label']
                confidence = results[0]['score']

                # Display prediction results
                st.markdown(f"### Prediction: **{label}**")
                st.markdown(f"### Confidence: **{confidence:.2%}**")

                # Provide confidence-based insights
                if confidence >= 0.8:
                    st.success("High confidence in the prediction.")
                elif confidence >= 0.5:
                    st.warning("Moderate confidence in the prediction. Consider additional verification.")
                else:
                    st.error("Low confidence in the prediction. Results should be interpreted with caution.")

                # Generate explanation
                explanation = generate_local_explanation(label, confidence)

                st.markdown("### Explanation")
                st.write(explanation)

            except Exception as e:
                st.error(f"Error during classification: {e}")